fix(usage): align session totals with streaming snapshots

Claude streaming turns can persist partial and complete usage rows. Subtracting discarded partial output keeps the session total consistent with its cost and breakdown while preserving stored output that has no priceable row.

Author: mjacobs

internal/db/usage.go

@@ -3232,9 +3232,11 @@ func (db *DB) GetSessionUsage(
 		}
 	}
 	snapshotMask := activity.ClaudeSnapshotSurvivorMask(snapshotRows)
+	deduplicatedOutputTokens := 0
 	seen := make(map[usageDedupToken]struct{})
 	for i, r := range usageRows {
 		if !snapshotMask[i] {
+			deduplicatedOutputTokens += snapshotRows[i].OutputTokens
 			continue
 		}
 		if key, ok := usageDedupTokenForRow(
@@ -3301,7 +3303,7 @@ func (db *DB) GetSessionUsage(
 		SessionID:         sess.ID,
 		Agent:             sess.Agent,
 		Project:           sess.Project,
-		TotalOutputTokens: sess.TotalOutputTokens,
+		TotalOutputTokens: max(sess.TotalOutputTokens-deduplicatedOutputTokens, 0),
 		PeakContextTokens: sess.PeakContextTokens,
 		HasTokenData:      sess.HasTotalOutputTokens || sess.HasPeakContextTokens,
 		Models:            sortedSetKeys(modelsSet),

internal/db/usage_test.go

@@ -3889,6 +3889,8 @@ func TestGetSessionUsage_PrefersCompleteClaudeSnapshot(t *testing.T) {
 	insertSession(t, d, "claude:streamed", "proj", func(s *Session) {
 		s.Agent = "claude-code"
 		s.StartedAt = new("2026-05-20T10:00:00Z")
+		s.TotalOutputTokens = 636
+		s.HasTotalOutputTokens = true
 	})
 	insertMessages(t, d,
 		Message{
@@ -3907,6 +3909,7 @@ func TestGetSessionUsage_PrefersCompleteClaudeSnapshot(t *testing.T) {
 
 	u, err := d.GetSessionUsage(ctx, "claude:streamed", true)
 	requireNoError(t, err, "GetSessionUsage")
+	assert.Equal(t, 631, u.TotalOutputTokens)
 	assert.Equal(t, money.MustParseDollars("0.020775"), u.Cost)
 	require.Len(t, u.Breakdown, 1)
 	assert.Equal(t, 631, u.Breakdown[0].OutputTokens)

internal/duckdb/analytics_usage.go

@@ -3516,7 +3516,16 @@ func duckUsageCTEFromRaw(
 								COALESCE(message_ordinal, -1) ASC
 						)
 					ELSE 1
-				END AS snapshot_rank
+				END AS snapshot_rank,
+				CASE
+					WHEN claude_message_id != '' AND claude_request_id != ''
+						THEN SUM(output_tokens_norm) OVER (
+							PARTITION BY session_id, claude_message_id, claude_request_id
+						) - MAX(output_tokens_norm) OVER (
+							PARTITION BY session_id, claude_message_id, claude_request_id
+						)
+					ELSE 0
+				END AS snapshot_deduplicated_output_tokens
 			FROM usage_windowed
 		),
 		usage_snapshot_survivors AS (
@@ -3640,6 +3649,7 @@ type duckUsageAggregateRow struct {
 	reportedCostRows      int
 	authoritativeCost     int64
 	authoritativeCostRows int
+	snapshotDedupOutput   int
 }
 
 type duckSessionUsageRow struct {
@@ -4279,6 +4289,7 @@ func (s *Store) forEachSessionUsageAggregateRow(
 			display_name, started_at,
 			input_tokens_norm AS input_tokens,
 			output_tokens_norm AS output_tokens,
+			snapshot_deduplicated_output_tokens,
 			cache_create_norm AS cache_creation_tokens,
 			cache_read_norm AS cache_read_tokens,
 			CASE WHEN cost_microdollars IS NULL OR cost_source = 'copilot-reported' THEN input_tokens_norm ELSE 0 END AS billable_input_tokens,
@@ -4310,7 +4321,8 @@ func (s *Store) forEachSessionUsageAggregateRow(
 			&r.sessionID, &r.project, &r.agent, &r.model,
 			&r.priceModel, &r.source, &r.messageOrdinal,
 			&r.displayName, &startedAt,
-			&r.inputTok, &r.outputTok, &r.cacheCr, &r.cacheRd,
+			&r.inputTok, &r.outputTok, &r.snapshotDedupOutput,
+			&r.cacheCr, &r.cacheRd,
 			&r.billableInput, &r.billableOutput, &r.billableReason,
 			&r.billableCacheCr, &r.billableCacheRd,
 			&r.billableWebSearch,
@@ -4617,10 +4629,12 @@ func (s *Store) GetSessionUsage(
 	var totalCost money.Money
 	var authoritativeCost *money.Money
 	var hasComputedCost, hasReportedCost bool
+	deduplicatedOutputTokens := 0
 	hasRows := false
 	err = s.forEachSessionUsageAggregateRow(
 		ctx, db.UsageFilter{}, sessionID,
 		func(r duckUsageAggregateRow) error {
+			deduplicatedOutputTokens += r.snapshotDedupOutput
 			if r.authoritativeCostRows > 0 {
 				v := money.Money{Microdollars: r.authoritativeCost}
 				authoritativeCost = &v
@@ -4693,7 +4707,7 @@ func (s *Store) GetSessionUsage(
 	}
 	out := &db.SessionUsage{
 		SessionID: sessionID, Agent: sess.Agent, Project: sess.Project,
-		TotalOutputTokens: sess.TotalOutputTokens,
+		TotalOutputTokens: max(sess.TotalOutputTokens-deduplicatedOutputTokens, 0),
 		PeakContextTokens: sess.PeakContextTokens,
 		HasTokenData:      hasRows || sess.HasTotalOutputTokens || sess.HasPeakContextTokens,
 		Models:            sortedBoolKeys(models),

internal/duckdb/store_test.go

@@ -2688,10 +2688,13 @@ func TestSessionUsagePrefersCompleteClaudeSnapshot(t *testing.T) {
 		`{"input_tokens":1000,"output_tokens":631}`)
 	second.OutputTokens = 631
 
+	session := syncSession(
+		"duck-streamed", "alpha", "streamed",
+		"2026-01-13T00:00:00.000Z", 2)
+	session.TotalOutputTokens = 636
+	session.HasTotalOutputTokens = true
 	_, err := local.WriteSessionBatchAtomic([]db.SessionBatchWrite{{
-		Session: syncSession(
-			"duck-streamed", "alpha", "streamed",
-			"2026-01-13T00:00:00.000Z", 2),
+		Session:         session,
 		Messages:        []db.Message{first, second},
 		DataVersion:     1,
 		ReplaceMessages: true,
@@ -2707,6 +2710,7 @@ func TestSessionUsagePrefersCompleteClaudeSnapshot(t *testing.T) {
 	got, err := store.GetSessionUsage(ctx, "duck-streamed", true)
 	require.NoError(t, err)
 	require.NotNil(t, got)
+	assert.Equal(t, 631, got.TotalOutputTokens)
 	assert.Equal(t, money.MustParseDollars("0.020775"), got.Cost)
 	require.Len(t, got.Breakdown, 1)
 	assert.Equal(t, 631, got.Breakdown[0].OutputTokens)

internal/postgres/usage.go

@@ -1363,9 +1363,11 @@ func (s *Store) GetSessionUsage(
 		}
 	}
 	snapshotMask := activity.ClaudeSnapshotSurvivorMask(snapshotRows)
+	deduplicatedOutputTokens := 0
 	seen := make(map[pgUsageDedupToken]struct{})
 	for i, r := range usageRows {
 		if !snapshotMask[i] {
+			deduplicatedOutputTokens += snapshotRows[i].OutputTokens
 			continue
 		}
 		if key, ok := pgUsageDedupTokenForRow(
@@ -1432,7 +1434,7 @@ func (s *Store) GetSessionUsage(
 		SessionID:         sess.ID,
 		Agent:             sess.Agent,
 		Project:           sess.Project,
-		TotalOutputTokens: sess.TotalOutputTokens,
+		TotalOutputTokens: max(sess.TotalOutputTokens-deduplicatedOutputTokens, 0),
 		PeakContextTokens: sess.PeakContextTokens,
 		HasTokenData: sess.HasTotalOutputTokens ||
 			sess.HasPeakContextTokens,

internal/postgres/usage_pgtest_test.go

@@ -561,10 +561,11 @@ func TestStoreGetSessionUsagePrefersCompleteClaudeSnapshot(t *testing.T) {
 		) VALUES ('claude-opus-4-6', 5000000, 25000000, 6250000, 500000, 'seed');
 		INSERT INTO sessions (
 			id, machine, project, agent, started_at,
-			message_count, user_message_count
+			message_count, user_message_count,
+			total_output_tokens, has_total_output_tokens
 		) VALUES (
 			'claude:streamed', 'test-machine', 'proj', 'claude-code',
-			'2026-03-12T10:00:00Z'::timestamptz, 2, 1
+			'2026-03-12T10:00:00Z'::timestamptz, 2, 1, 636, TRUE
 		);
 		INSERT INTO messages (
 			session_id, ordinal, role, content, timestamp, content_length,
@@ -583,6 +584,7 @@ func TestStoreGetSessionUsagePrefersCompleteClaudeSnapshot(t *testing.T) {
 	got, err := store.GetSessionUsage(ctx, "claude:streamed", true)
 	require.NoError(t, err)
 	require.NotNil(t, got)
+	assert.Equal(t, 631, got.TotalOutputTokens)
 	assert.Equal(t, money.MustParseDollars("0.020775"), got.Cost)
 	require.Len(t, got.Breakdown, 1)
 	assert.Equal(t, 631, got.Breakdown[0].OutputTokens)