You've already forked agentic-coding-workflow
bugfix
This commit is contained in:
@@ -0,0 +1,317 @@
|
||||
package acp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"math/big"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgtype"
|
||||
|
||||
acpsqlc "github.com/yan1h/agent-coding-workflow/internal/acp/sqlc"
|
||||
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
|
||||
)
|
||||
|
||||
func isNoRows(err error) bool { return errors.Is(err, pgx.ErrNoRows) }
|
||||
|
||||
// ===== NUMERIC <-> float64 转换 helper（复制 chat/repository.go 的实现）=====
|
||||
|
||||
func numericToFloat64(n pgtype.Numeric) float64 {
|
||||
if !n.Valid {
|
||||
return 0
|
||||
}
|
||||
f8, err := n.Float64Value()
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return f8.Float64
|
||||
}
|
||||
|
||||
func numericToFloat64Ptr(n pgtype.Numeric) *float64 {
|
||||
if !n.Valid {
|
||||
return nil
|
||||
}
|
||||
v := numericToFloat64(n)
|
||||
return &v
|
||||
}
|
||||
|
||||
func float64ToNumeric(f float64) pgtype.Numeric {
|
||||
var n pgtype.Numeric
|
||||
if err := n.ScanScientific(fmt.Sprintf("%.10e", f)); err != nil {
|
||||
bi := new(big.Int)
|
||||
bi.SetInt64(int64(f))
|
||||
return pgtype.Numeric{Int: bi, Exp: 0, Valid: true}
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
func float64PtrToNumeric(f *float64) pgtype.Numeric {
|
||||
if f == nil {
|
||||
return pgtype.Numeric{}
|
||||
}
|
||||
return float64ToNumeric(*f)
|
||||
}
|
||||
|
||||
// ===== 成本核算 / 预算 / reaper 仓储方法 =====
|
||||
|
||||
func (r *pgRepo) InsertSessionUsage(ctx context.Context, rec *SessionUsageRecord) (bool, error) {
|
||||
var srcID *int64
|
||||
if rec.SourceEventID != nil {
|
||||
v := *rec.SourceEventID
|
||||
srcID = &v
|
||||
}
|
||||
id, err := r.q.InsertSessionUsage(ctx, acpsqlc.InsertSessionUsageParams{
|
||||
SessionID: toPgUUID(rec.SessionID),
|
||||
UserID: toPgUUID(rec.UserID),
|
||||
ProjectID: toPgUUID(rec.ProjectID),
|
||||
AgentKindID: toPgUUID(rec.AgentKindID),
|
||||
ModelID: toPgUUIDPtr(rec.ModelID),
|
||||
PromptTokens: rec.PromptTokens,
|
||||
CompletionTokens: rec.CompletionTokens,
|
||||
ThinkingTokens: rec.ThinkingTokens,
|
||||
CostUsd: float64ToNumeric(rec.CostUSD),
|
||||
SourceEventID: srcID,
|
||||
})
|
||||
if err != nil {
|
||||
// ON CONFLICT DO NOTHING + no row -> not inserted (idempotent de-dup).
|
||||
if isNoRows(err) {
|
||||
return false, nil
|
||||
}
|
||||
return false, errs.Wrap(err, errs.CodeInternal, "insert session usage")
|
||||
}
|
||||
rec.ID = id
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (r *pgRepo) AddSessionUsageTotals(ctx context.Context, sid uuid.UUID, dp, dc, dt int64, dCost float64) (int64, int64, int64, float64, error) {
|
||||
row, err := r.q.AddSessionUsageTotals(ctx, acpsqlc.AddSessionUsageTotalsParams{
|
||||
ID: toPgUUID(sid),
|
||||
PromptTokens: dp,
|
||||
CompletionTokens: dc,
|
||||
ThinkingTokens: dt,
|
||||
TotalCostUsd: float64ToNumeric(dCost),
|
||||
})
|
||||
if err != nil {
|
||||
return 0, 0, 0, 0, errs.Wrap(err, errs.CodeInternal, "add session usage totals")
|
||||
}
|
||||
return row.PromptTokens, row.CompletionTokens, row.ThinkingTokens, numericToFloat64(row.TotalCostUsd), nil
|
||||
}
|
||||
|
||||
func (r *pgRepo) SumProjectCostUSD(ctx context.Context, projectID uuid.UUID, since time.Time) (float64, error) {
|
||||
n, err := r.q.SumProjectCostUSD(ctx, acpsqlc.SumProjectCostUSDParams{
|
||||
ProjectID: toPgUUID(projectID),
|
||||
CreatedAt: pgtype.Timestamptz{Time: since, Valid: true},
|
||||
})
|
||||
if err != nil {
|
||||
return 0, errs.Wrap(err, errs.CodeInternal, "sum project cost")
|
||||
}
|
||||
return numericToFloat64(n), nil
|
||||
}
|
||||
|
||||
func (r *pgRepo) MarkSessionTerminatedReason(ctx context.Context, sid uuid.UUID, reason string) error {
|
||||
if err := r.q.MarkSessionTerminatedReason(ctx, acpsqlc.MarkSessionTerminatedReasonParams{
|
||||
ID: toPgUUID(sid),
|
||||
TerminatedReason: &reason,
|
||||
}); err != nil {
|
||||
return errs.Wrap(err, errs.CodeInternal, "mark session terminated reason")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *pgRepo) ListSessionsForReaper(ctx context.Context, wallClockStartedBefore, idleSince time.Time) ([]ReaperSession, error) {
|
||||
var wallTS pgtype.Timestamptz
|
||||
if !wallClockStartedBefore.IsZero() {
|
||||
wallTS = pgtype.Timestamptz{Time: wallClockStartedBefore, Valid: true}
|
||||
}
|
||||
rows, err := r.q.ListSessionsForReaper(ctx, acpsqlc.ListSessionsForReaperParams{
|
||||
Column1: wallTS,
|
||||
Column2: pgtype.Timestamptz{Time: idleSince, Valid: true},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "list sessions for reaper")
|
||||
}
|
||||
out := make([]ReaperSession, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
out = append(out, ReaperSession{
|
||||
ID: fromPgUUID(row.ID),
|
||||
UserID: fromPgUUID(row.UserID),
|
||||
StartedAt: row.StartedAt.Time,
|
||||
LastActivityAt: row.LastActivityAt.Time,
|
||||
MaxWallClockSeconds: row.BudgetMaxWallClockSeconds,
|
||||
})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *pgRepo) ListSessionUsage(ctx context.Context, sessionID uuid.UUID, limit int32) ([]*SessionUsageRecord, error) {
|
||||
if limit <= 0 {
|
||||
limit = 50
|
||||
}
|
||||
rows, err := r.q.ListSessionUsageBySession(ctx, acpsqlc.ListSessionUsageBySessionParams{
|
||||
SessionID: toPgUUID(sessionID),
|
||||
Limit: limit,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "list session usage")
|
||||
}
|
||||
out := make([]*SessionUsageRecord, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
out = append(out, rowToSessionUsage(row))
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func rowToSessionUsage(row acpsqlc.AcpSessionUsage) *SessionUsageRecord {
|
||||
var srcID *int64
|
||||
if row.SourceEventID != nil {
|
||||
v := *row.SourceEventID
|
||||
srcID = &v
|
||||
}
|
||||
return &SessionUsageRecord{
|
||||
ID: row.ID,
|
||||
SessionID: fromPgUUID(row.SessionID),
|
||||
UserID: fromPgUUID(row.UserID),
|
||||
ProjectID: fromPgUUID(row.ProjectID),
|
||||
AgentKindID: fromPgUUID(row.AgentKindID),
|
||||
ModelID: fromPgUUIDPtr(row.ModelID),
|
||||
PromptTokens: row.PromptTokens,
|
||||
CompletionTokens: row.CompletionTokens,
|
||||
ThinkingTokens: row.ThinkingTokens,
|
||||
CostUSD: numericToFloat64(row.CostUsd),
|
||||
SourceEventID: srcID,
|
||||
CreatedAt: row.CreatedAt.Time,
|
||||
}
|
||||
}
|
||||
|
||||
func (r *pgRepo) SummarizeAcpUsageByUser(ctx context.Context, from, to time.Time) ([]AcpUsageUserRow, error) {
|
||||
rows, err := r.q.SummarizeAcpUsageByUser(ctx, acpsqlc.SummarizeAcpUsageByUserParams{
|
||||
CreatedAt: pgtype.Timestamptz{Time: from, Valid: true},
|
||||
CreatedAt_2: pgtype.Timestamptz{Time: to, Valid: true},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "summarize acp usage by user")
|
||||
}
|
||||
out := make([]AcpUsageUserRow, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
out = append(out, AcpUsageUserRow{
|
||||
UserID: fromPgUUID(row.UserID),
|
||||
PromptTokens: row.PromptTokens,
|
||||
CompletionTokens: row.CompletionTokens,
|
||||
ThinkingTokens: row.ThinkingTokens,
|
||||
CostUSD: numericToFloat64(row.CostUsd),
|
||||
})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *pgRepo) SummarizeAcpUsageByModel(ctx context.Context, from, to time.Time) ([]AcpUsageModelRow, error) {
|
||||
rows, err := r.q.SummarizeAcpUsageByModel(ctx, acpsqlc.SummarizeAcpUsageByModelParams{
|
||||
CreatedAt: pgtype.Timestamptz{Time: from, Valid: true},
|
||||
CreatedAt_2: pgtype.Timestamptz{Time: to, Valid: true},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "summarize acp usage by model")
|
||||
}
|
||||
out := make([]AcpUsageModelRow, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
out = append(out, AcpUsageModelRow{
|
||||
ModelID: fromPgUUID(row.ModelID),
|
||||
PromptTokens: row.PromptTokens,
|
||||
CompletionTokens: row.CompletionTokens,
|
||||
ThinkingTokens: row.ThinkingTokens,
|
||||
CostUSD: numericToFloat64(row.CostUsd),
|
||||
})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *pgRepo) SummarizeAcpUsageByDay(ctx context.Context, from, to time.Time) ([]AcpUsageDayRow, error) {
|
||||
rows, err := r.q.SummarizeAcpUsageByDay(ctx, acpsqlc.SummarizeAcpUsageByDayParams{
|
||||
CreatedAt: pgtype.Timestamptz{Time: from, Valid: true},
|
||||
CreatedAt_2: pgtype.Timestamptz{Time: to, Valid: true},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "summarize acp usage by day")
|
||||
}
|
||||
out := make([]AcpUsageDayRow, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
out = append(out, AcpUsageDayRow{
|
||||
Day: row.Day.Time,
|
||||
PromptTokens: row.PromptTokens,
|
||||
CompletionTokens: row.CompletionTokens,
|
||||
ThinkingTokens: row.ThinkingTokens,
|
||||
CostUSD: numericToFloat64(row.CostUsd),
|
||||
})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// ===== run-history dashboard 聚合 =====
|
||||
|
||||
func (r *pgRepo) SessionRollup(ctx context.Context, f DashboardFilter) (SessionRollup, error) {
|
||||
row, err := r.q.SessionRollup(ctx, acpsqlc.SessionRollupParams{
|
||||
Column1: toPgUUIDPtr(f.UserID),
|
||||
Column2: toPgUUIDPtr(f.ProjectID),
|
||||
Column3: toPgUUIDPtr(f.RequirementID),
|
||||
StartedAt: pgtype.Timestamptz{Time: f.From, Valid: true},
|
||||
StartedAt_2: pgtype.Timestamptz{Time: f.To, Valid: true},
|
||||
})
|
||||
if err != nil {
|
||||
return SessionRollup{}, errs.Wrap(err, errs.CodeInternal, "session rollup")
|
||||
}
|
||||
return SessionRollup{
|
||||
Total: row.Total,
|
||||
Succeeded: row.Succeeded,
|
||||
Crashed: row.Crashed,
|
||||
Active: row.Active,
|
||||
TotalCostUSD: numericToFloat64(row.TotalCostUsd),
|
||||
TotalTokens: row.TotalTokens,
|
||||
AvgDurationSeconds: row.AvgDurationSeconds,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (r *pgRepo) SessionsByDay(ctx context.Context, f DashboardFilter) ([]SessionDayRow, error) {
|
||||
rows, err := r.q.SessionsByDay(ctx, acpsqlc.SessionsByDayParams{
|
||||
Column1: toPgUUIDPtr(f.UserID),
|
||||
Column2: toPgUUIDPtr(f.ProjectID),
|
||||
Column3: toPgUUIDPtr(f.RequirementID),
|
||||
StartedAt: pgtype.Timestamptz{Time: f.From, Valid: true},
|
||||
StartedAt_2: pgtype.Timestamptz{Time: f.To, Valid: true},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "sessions by day")
|
||||
}
|
||||
out := make([]SessionDayRow, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
out = append(out, SessionDayRow{
|
||||
Day: row.Day.Time,
|
||||
Total: row.Total,
|
||||
Succeeded: row.Succeeded,
|
||||
Crashed: row.Crashed,
|
||||
TotalCostUSD: numericToFloat64(row.TotalCostUsd),
|
||||
})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *pgRepo) SessionTimeline(ctx context.Context, sessionID uuid.UUID) ([]SessionTimelineBucket, error) {
|
||||
rows, err := r.q.SessionTimeline(ctx, toPgUUID(sessionID))
|
||||
if err != nil {
|
||||
return nil, errs.Wrap(err, errs.CodeInternal, "session timeline")
|
||||
}
|
||||
out := make([]SessionTimelineBucket, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
out = append(out, SessionTimelineBucket{
|
||||
Direction: row.Direction,
|
||||
RPCKind: row.RpcKind,
|
||||
Method: row.Method,
|
||||
EventCount: row.EventCount,
|
||||
FirstAt: row.FirstAt.Time,
|
||||
LastAt: row.LastAt.Time,
|
||||
})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
Reference in New Issue
Block a user