diff --git a/internal/tracing/snapshot_worker.go b/internal/tracing/snapshot_worker.go index e74b98af..2e9eed7e 100644 --- a/internal/tracing/snapshot_worker.go +++ b/internal/tracing/snapshot_worker.go @@ -340,7 +340,7 @@ func querySpanAggregates(ctx context.Context, db *sql.DB, from, to time.Time) ([ COALESCE(t.channel, '') as channel, COALESCE(s.provider, '') as provider, COALESCE(s.model, '') as model, - COUNT(*) FILTER (WHERE s.span_type = 'llm_call') as llm_call_count, + COUNT(*) as llm_call_count, COALESCE(SUM(s.input_tokens), 0) as span_input_tokens, COALESCE(SUM(s.output_tokens), 0) as span_output_tokens, COALESCE(SUM(s.total_cost), 0) as span_cost, @@ -348,7 +348,7 @@ func querySpanAggregates(ctx context.Context, db *sql.DB, from, to time.Time) ([ COALESCE(SUM(CAST(s.metadata->>'cache_creation_tokens' AS INTEGER)), 0) as cache_create_tokens, COALESCE(SUM(CAST(s.metadata->>'thinking_tokens' AS INTEGER)), 0) as thinking_tokens FROM traces t - JOIN spans s ON s.trace_id = t.id AND s.span_type IN ('llm_call', 'tool_call') + JOIN spans s ON s.trace_id = t.id AND s.span_type = 'llm_call' WHERE t.start_time >= $1 AND t.start_time < $2 AND t.parent_trace_id IS NULL GROUP BY t.agent_id, t.channel, s.provider, s.model`, from, to) @@ -480,6 +480,27 @@ type agentChannelKey struct { Channel string } +// findTotalsSnapshotIndex returns the index of the existing totals row +// (Provider=="" && Model=="") in snapshots matching key, or -1 if none found. +func findTotalsSnapshotIndex(snapshots []store.UsageSnapshot, key agentChannelKey) int { + for i, snap := range snapshots { + if snap.Provider != "" || snap.Model != "" { + continue + } + if snap.Channel != key.Channel { + continue + } + var agentID uuid.UUID + if snap.AgentID != nil { + agentID = *snap.AgentID + } + if agentID == key.AgentID { + return i + } + } + return -1 +} + func mergeTraceAndSpanRows( bucketStart time.Time, traceRows []traceAggregate, @@ -526,8 +547,29 @@ func mergeTraceAndSpanRows( snapshots = append(snapshots, snap) } - // 2. Create detail rows from span data (with actual provider/model) + // 2. Create detail rows from span data (with actual provider/model). + // Defense in depth: a span row with empty provider/model would collide + // with the totals row's conflict-target key (agent_id, provider='', + // model='', channel). Instead of appending a duplicate row, merge its + // metrics into the existing totals row for that (agent_id, channel). for _, sp := range spanRows { + if sp.Provider == "" && sp.Model == "" { + key := agentChannelKey{Channel: sp.Channel} + if sp.AgentID != nil { + key.AgentID = *sp.AgentID + } + if idx := findTotalsSnapshotIndex(snapshots, key); idx >= 0 { + snapshots[idx].LLMCallCount += sp.LLMCallCount + snapshots[idx].InputTokens += sp.InputTokens + snapshots[idx].OutputTokens += sp.OutputTokens + snapshots[idx].TotalCost += sp.TotalCost + snapshots[idx].CacheReadTokens += sp.CacheReadTokens + snapshots[idx].CacheCreateTokens += sp.CacheCreateTokens + snapshots[idx].ThinkingTokens += sp.ThinkingTokens + continue + } + } + snapshots = append(snapshots, store.UsageSnapshot{ BucketHour: bucketStart, AgentID: sp.AgentID, diff --git a/internal/tracing/snapshot_worker_test.go b/internal/tracing/snapshot_worker_test.go index 71a47f2d..65b76db7 100644 --- a/internal/tracing/snapshot_worker_test.go +++ b/internal/tracing/snapshot_worker_test.go @@ -3,6 +3,8 @@ package tracing import ( "testing" "time" + + "github.com/google/uuid" ) func TestUsageCatchUpStartHourRefreshesLatestClosedBucket(t *testing.T) { @@ -35,3 +37,112 @@ func TestUsageCatchUpStartHourWithoutSnapshotsComputesPreviousHourOnly(t *testin t.Fatalf("start hour = %s, want %s", got, target) } } + +// TestMergeTraceAndSpanRowsDedupesToolCallProviderModelRow guards against +// SQLSTATE 21000 ("ON CONFLICT DO UPDATE command cannot affect row a second +// time"): a span row that resolves to Provider=="" && Model=="" (e.g. from +// tool_call spans, which never carry provider/model) must be merged into the +// existing totals row for that (agent_id, channel) instead of appended as a +// second row sharing the same upsert conflict-target key. +func TestMergeTraceAndSpanRowsDedupesToolCallProviderModelRow(t *testing.T) { + bucketStart := time.Date(2026, 7, 3, 1, 0, 0, 0, time.UTC) + agentID := uuid.New() + + traceRows := []traceAggregate{ + { + AgentID: &agentID, + Channel: "web", + RequestCount: 10, + ErrorCount: 1, + UniqueUsers: 3, + InputTokens: 100, + OutputTokens: 200, + TotalCost: 1.5, + ToolCallCount: 4, + AvgDurationMS: 500, + }, + } + + spanRows := []spanAggregate{ + // Simulates leftover/edge-case span row with empty provider/model + // (e.g. tool_call-only spans) that would otherwise collide with the + // totals row's conflict-target key. + { + AgentID: &agentID, + Channel: "web", + Provider: "", + Model: "", + LLMCallCount: 2, + InputTokens: 10, + OutputTokens: 20, + TotalCost: 0.1, + CacheReadTokens: 1, + CacheCreateTokens: 2, + ThinkingTokens: 3, + }, + { + AgentID: &agentID, + Channel: "web", + Provider: "anthropic", + Model: "claude-sonnet-4", + LLMCallCount: 5, + InputTokens: 50, + OutputTokens: 60, + TotalCost: 0.9, + CacheReadTokens: 4, + CacheCreateTokens: 5, + ThinkingTokens: 6, + }, + } + + snapshots := mergeTraceAndSpanRows(bucketStart, traceRows, spanRows, nil, nil) + + var totalsRows []int + for i, snap := range snapshots { + if snap.AgentID != nil && *snap.AgentID == agentID && snap.Channel == "web" && + snap.Provider == "" && snap.Model == "" { + totalsRows = append(totalsRows, i) + } + } + + if len(totalsRows) != 1 { + t.Fatalf("expected exactly 1 totals row with Provider==\"\" && Model==\"\" for (agentID, web), got %d", len(totalsRows)) + } + + totals := snapshots[totalsRows[0]] + if want := 2; totals.LLMCallCount != want { + t.Fatalf("LLMCallCount = %d, want %d", totals.LLMCallCount, want) + } + if want := int64(10); totals.InputTokens != want { + t.Fatalf("InputTokens = %d, want %d", totals.InputTokens, want) + } + if want := int64(20); totals.OutputTokens != want { + t.Fatalf("OutputTokens = %d, want %d", totals.OutputTokens, want) + } + if want := 0.1; totals.TotalCost != want { + t.Fatalf("TotalCost = %v, want %v", totals.TotalCost, want) + } + if want := int64(1); totals.CacheReadTokens != want { + t.Fatalf("CacheReadTokens = %d, want %d", totals.CacheReadTokens, want) + } + if want := int64(2); totals.CacheCreateTokens != want { + t.Fatalf("CacheCreateTokens = %d, want %d", totals.CacheCreateTokens, want) + } + if want := int64(3); totals.ThinkingTokens != want { + t.Fatalf("ThinkingTokens = %d, want %d", totals.ThinkingTokens, want) + } + + // The distinct provider/model detail row must still exist separately. + var detailFound bool + for _, snap := range snapshots { + if snap.Provider == "anthropic" && snap.Model == "claude-sonnet-4" { + detailFound = true + if snap.LLMCallCount != 5 { + t.Fatalf("detail row LLMCallCount = %d, want 5", snap.LLMCallCount) + } + } + } + if !detailFound { + t.Fatal("expected detail row for provider=anthropic model=claude-sonnet-4") + } +}