fix: report per-request Anthropic usage in chat token accounting (#27966)

coderd/x/chatd/chatd_test.go

@@ -6872,6 +6872,45 @@ func TestActiveServer_AnthropicUsageMatchesFinalDelta(t *testing.T) {
 	require.Equal(t, sql.NullInt64{Int64: 150, Valid: true}, last.CacheReadTokens)
 }
 
+func TestActiveServer_AnthropicPersistsPerRequestUsage(t *testing.T) {
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	db, ps := dbtestutil.NewDB(t)
+	anthropicURL := chattest.NewAnthropic(t, func(_ *chattest.AnthropicRequest) chattest.AnthropicResponse {
+		return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunksWithMessageUsages(
+			chattest.AnthropicUsage{
+				InputTokens:              2,
+				CacheReadInputTokens:     139956,
+				CacheCreationInputTokens: 7770,
+			},
+			chattest.AnthropicUsage{
+				InputTokens:              4,
+				CacheReadInputTokens:     287682,
+				CacheCreationInputTokens: 15556,
+				OutputTokens:             4996,
+			},
+			"cached response",
+		)...)
+	})
+	user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL)
+
+	server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) {
+		cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, anthropicURL, chattest.WithPreservePath()))
+	})
+	chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "hello")
+	waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting)
+
+	messages := chatMessages(ctx, t, db, chat.ID)
+	last := messages[len(messages)-1]
+	require.Equal(t, database.ChatMessageRoleAssistant, last.Role)
+	require.Equal(t, sql.NullInt64{Int64: 139956, Valid: true}, last.CacheReadTokens)
+	require.Equal(t, sql.NullInt64{Int64: 7770, Valid: true}, last.CacheCreationTokens)
+	require.Equal(t, sql.NullInt64{Int64: 2, Valid: true}, last.InputTokens)
+	require.Equal(t, sql.NullInt64{Int64: 4996, Valid: true}, last.OutputTokens)
+	require.Equal(t, sql.NullInt64{Int64: 4998, Valid: true}, last.TotalTokens)
+}
+
 func TestActiveServer_ChatTurnDebugRunRecordsStreamStep(t *testing.T) {
 	t.Parallel()
 

coderd/x/chatd/chattest/anthropic.go

@@ -375,6 +375,16 @@ func AnthropicTextChunks(deltas ...string) []AnthropicChunk {
 // the initial input and cache token counts, and the final message_delta
 // carries the output token count.
 func AnthropicTextChunksWithCacheUsage(usage AnthropicUsage, deltas ...string) []AnthropicChunk {
+	return AnthropicTextChunksWithMessageUsages(
+		usage,
+		AnthropicUsage{OutputTokens: usage.OutputTokens},
+		deltas...,
+	)
+}
+
+// AnthropicTextChunksWithMessageUsages creates a streaming response with
+// independent message_start and message_delta usage.
+func AnthropicTextChunksWithMessageUsages(messageStartUsage, messageDeltaUsage AnthropicUsage, deltas ...string) []AnthropicChunk {
 	if len(deltas) == 0 {
 		return nil
 	}
@@ -383,13 +393,26 @@ func AnthropicTextChunksWithCacheUsage(usage AnthropicUsage, deltas ...string) [
 	model := "claude-3-opus-20240229"
 
 	messageUsage := map[string]int{
-		"input_tokens": usage.InputTokens,
+		"input_tokens": messageStartUsage.InputTokens,
 	}
-	if usage.CacheCreationInputTokens != 0 {
-		messageUsage["cache_creation_input_tokens"] = usage.CacheCreationInputTokens
+	if messageStartUsage.CacheCreationInputTokens != 0 {
+		messageUsage["cache_creation_input_tokens"] = messageStartUsage.CacheCreationInputTokens
 	}
-	if usage.CacheReadInputTokens != 0 {
-		messageUsage["cache_read_input_tokens"] = usage.CacheReadInputTokens
+	if messageStartUsage.CacheReadInputTokens != 0 {
+		messageUsage["cache_read_input_tokens"] = messageStartUsage.CacheReadInputTokens
+	}
+
+	deltaUsage := map[string]int{
+		"output_tokens": messageDeltaUsage.OutputTokens,
+	}
+	if messageDeltaUsage.InputTokens != 0 {
+		deltaUsage["input_tokens"] = messageDeltaUsage.InputTokens
+	}
+	if messageDeltaUsage.CacheCreationInputTokens != 0 {
+		deltaUsage["cache_creation_input_tokens"] = messageDeltaUsage.CacheCreationInputTokens
+	}
+	if messageDeltaUsage.CacheReadInputTokens != 0 {
+		deltaUsage["cache_read_input_tokens"] = messageDeltaUsage.CacheReadInputTokens
 	}
 
 	chunks := []AnthropicChunk{
@@ -432,9 +455,7 @@ func AnthropicTextChunksWithCacheUsage(usage AnthropicUsage, deltas ...string) [
 		AnthropicChunk{
 			Type:       "message_delta",
 			StopReason: "end_turn",
-			UsageMap: map[string]int{
-				"output_tokens": usage.OutputTokens,
-			},
+			UsageMap:   deltaUsage,
 		},
 		AnthropicChunk{
 			Type: "message_stop",
@@ -556,6 +577,7 @@ func AnthropicToolCallChunks(toolName string, inputJSONDeltas ...string) []Anthr
 				Type:  "message",
 				Role:  "assistant",
 				Model: model,
+				Usage: map[string]int{"input_tokens": 10},
 			},
 		},
 		{
@@ -589,9 +611,8 @@ func AnthropicToolCallChunks(toolName string, inputJSONDeltas ...string) []Anthr
 		AnthropicChunk{
 			Type:       "message_delta",
 			StopReason: "tool_use",
-			Usage: AnthropicUsage{
-				InputTokens:  10,
-				OutputTokens: 5,
+			UsageMap: map[string]int{
+				"output_tokens": 5,
 			},
 		},
 		AnthropicChunk{