feat: add columns for cached tokens from aibridge (#23832)

Two new columns added to aibridge_token_usages:
  - cache_read_input_tokens (BIGINT, default 0)
  - cache_write_input_tokens (BIGINT, default 0)

Migration backfills existing rows by extracting values from the metadata
JSONB column (cache_read_input, input_cached, prompt_cached for reads
(max value selected since only 1 should be set), cache_creation_input
for writes).

All references to data from metadata were updated to reference new
columns. No other changes then changing where data is extracted from.

Requires aibridge library version bump to include:
https://github.com/coder/aibridge/pull/229
Fixes: https://github.com/coder/aibridge/issues/150

coderd/database/db2sdk/db2sdk.go

@@ -1037,8 +1037,10 @@ func AIBridgeSession(row database.ListAIBridgeSessionsRow) codersdk.AIBridgeSess
 		StartedAt: row.StartedAt,
 		Threads:   row.Threads,
 		TokenUsageSummary: codersdk.AIBridgeSessionTokenUsageSummary{
-			InputTokens:  row.InputTokens,
-			OutputTokens: row.OutputTokens,
+			InputTokens:           row.InputTokens,
+			OutputTokens:          row.OutputTokens,
+			CacheReadInputTokens:  row.CacheReadInputTokens,
+			CacheWriteInputTokens: row.CacheWriteInputTokens,
 		},
 	}
 	// Ensure non-nil slices for JSON serialization.
@@ -1062,13 +1064,15 @@ func AIBridgeSession(row database.ListAIBridgeSessionsRow) codersdk.AIBridgeSess
 
 func AIBridgeTokenUsage(usage database.AIBridgeTokenUsage) codersdk.AIBridgeTokenUsage {
 	return codersdk.AIBridgeTokenUsage{
-		ID:                 usage.ID,
-		InterceptionID:     usage.InterceptionID,
-		ProviderResponseID: usage.ProviderResponseID,
-		InputTokens:        usage.InputTokens,
-		OutputTokens:       usage.OutputTokens,
-		Metadata:           jsonOrEmptyMap(usage.Metadata),
-		CreatedAt:          usage.CreatedAt,
+		ID:                    usage.ID,
+		InterceptionID:        usage.InterceptionID,
+		ProviderResponseID:    usage.ProviderResponseID,
+		InputTokens:           usage.InputTokens,
+		OutputTokens:          usage.OutputTokens,
+		CacheReadInputTokens:  usage.CacheReadInputTokens,
+		CacheWriteInputTokens: usage.CacheWriteInputTokens,
+		Metadata:              jsonOrEmptyMap(usage.Metadata),
+		CreatedAt:             usage.CreatedAt,
 	}
 }
 
@@ -1179,9 +1183,11 @@ func AIBridgeSessionThreads(
 		PageStartedAt: pageStartedAt,
 		PageEndedAt:   pageEndedAt,
 		TokenUsageSummary: codersdk.AIBridgeSessionThreadsTokenUsage{
-			InputTokens:  session.InputTokens,
-			OutputTokens: session.OutputTokens,
-			Metadata:     sessionTokenMeta,
+			InputTokens:           session.InputTokens,
+			OutputTokens:          session.OutputTokens,
+			CacheReadInputTokens:  session.CacheReadInputTokens,
+			CacheWriteInputTokens: session.CacheWriteInputTokens,
+			Metadata:              sessionTokenMeta,
 		},
 		Threads: threads,
 	}
@@ -1314,17 +1320,19 @@ func buildAIBridgeThread(
 
 // aggregateTokenUsage sums token usage rows and aggregates metadata.
 func aggregateTokenUsage(tokens []database.AIBridgeTokenUsage) codersdk.AIBridgeSessionThreadsTokenUsage {
-	var inputTokens, outputTokens int64
+	var inputTokens, outputTokens, cacheRead, cacheWrite int64
 	for _, tu := range tokens {
 		inputTokens += tu.InputTokens
 		outputTokens += tu.OutputTokens
-		// TODO: once https://github.com/coder/aibridge/issues/150 lands we
-		// should aggregate the other token types.
+		cacheRead += tu.CacheReadInputTokens
+		cacheWrite += tu.CacheWriteInputTokens
 	}
 	return codersdk.AIBridgeSessionThreadsTokenUsage{
-		InputTokens:  inputTokens,
-		OutputTokens: outputTokens,
-		Metadata:     aggregateTokenMetadata(tokens),
+		InputTokens:           inputTokens,
+		OutputTokens:          outputTokens,
+		CacheReadInputTokens:  cacheRead,
+		CacheWriteInputTokens: cacheWrite,
+		Metadata:              aggregateTokenMetadata(tokens),
 	}
 }
 

coderd/database/db2sdk/db2sdk_test.go

@@ -259,11 +259,13 @@ func TestAIBridgeInterception(t *testing.T) {
 			},
 			tokenUsages: []database.AIBridgeTokenUsage{
 				{
-					ID:                 uuid.New(),
-					InterceptionID:     interceptionID,
-					ProviderResponseID: "resp-123",
-					InputTokens:        100,
-					OutputTokens:       200,
+					ID:                    uuid.New(),
+					InterceptionID:        interceptionID,
+					ProviderResponseID:    "resp-123",
+					InputTokens:           100,
+					OutputTokens:          200,
+					CacheReadInputTokens:  50,
+					CacheWriteInputTokens: 10,
 					Metadata: pqtype.NullRawMessage{
 						RawMessage: json.RawMessage(`{"cache":"hit"}`),
 						Valid:      true,
@@ -413,6 +415,8 @@ func TestAIBridgeInterception(t *testing.T) {
 				require.Equal(t, tu.ProviderResponseID, result.TokenUsages[i].ProviderResponseID)
 				require.Equal(t, tu.InputTokens, result.TokenUsages[i].InputTokens)
 				require.Equal(t, tu.OutputTokens, result.TokenUsages[i].OutputTokens)
+				require.Equal(t, tu.CacheReadInputTokens, result.TokenUsages[i].CacheReadInputTokens)
+				require.Equal(t, tu.CacheWriteInputTokens, result.TokenUsages[i].CacheWriteInputTokens)
 			}
 
 			// Verify user prompts are converted correctly.