fix(server-utils): Count Bedrock prompt-cache tokens as input tokens

packages/server-utils/src/integrations/aws-sdk/services/bedrock-runtime.ts

@@ -9,10 +9,13 @@ import {
   GEN_AI_REQUEST_TEMPERATURE,
   GEN_AI_REQUEST_TOP_P,
   GEN_AI_RESPONSE_FINISH_REASONS,
+  GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS,
+  GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS,
   GEN_AI_USAGE_INPUT_TOKENS,
   GEN_AI_USAGE_OUTPUT_TOKENS,
 } from '@sentry/conventions/attributes';
 import { GEN_AI_CHAT, GEN_AI_GENERATE_CONTENT } from '@sentry/conventions/op';
+import { getTokenUsageAttributes } from '../../../ai/core/utils';
 import { DEBUG_BUILD } from '../../../debug-build';
 import {
   GEN_AI_OPERATION_NAME_VALUE_CHAT,
@@ -28,6 +31,8 @@ interface TokenUsage {
   inputTokens: number | undefined;
   outputTokens: number | undefined;
   totalTokens: number | undefined;
+  cacheReadInputTokens?: number;
+  cacheWriteInputTokens?: number;
 }
 
 interface ConverseStreamOutput {
@@ -254,7 +259,17 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
     const { stopReason, usage } = response.data;
 
     setStopReason(span, stopReason);
-    setUsage(span, usage);
+    span.setAttributes(
+      getTokenUsageAttributes(
+        // With prompt caching, `inputTokens` leaves out cache reads and writes, which the conventions count as input.
+        typeof usage?.inputTokens === 'number'
+          ? usage.inputTokens + (usage.cacheReadInputTokens ?? 0) + (usage.cacheWriteInputTokens ?? 0)
+          : undefined,
+        usage?.outputTokens,
+        usage?.cacheWriteInputTokens,
+        usage?.cacheReadInputTokens,
+      ),
+    );
   }
 
   private _responseHookConverseStream(response: NormalizedResponse, span: Span): void {
@@ -296,12 +311,18 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
           span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.stopReason]);
         }
       } else if (currentModelId.includes('anthropic.claude')) {
-        if (responseBody.usage?.input_tokens !== undefined) {
-          span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, responseBody.usage.input_tokens);
-        }
-        if (responseBody.usage?.output_tokens !== undefined) {
-          span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.usage.output_tokens);
-        }
+        const usage = responseBody.usage;
+        span.setAttributes(
+          getTokenUsageAttributes(
+            // Anthropic's `input_tokens` leaves out cache reads and writes, which the conventions count as input.
+            typeof usage?.input_tokens === 'number'
+              ? usage.input_tokens + (usage.cache_read_input_tokens ?? 0) + (usage.cache_creation_input_tokens ?? 0)
+              : undefined,
+            usage?.output_tokens,
+            usage?.cache_creation_input_tokens,
+            usage?.cache_read_input_tokens,
+          ),
+        );
         if (responseBody.stop_reason !== undefined) {
           span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.stop_reason]);
         }
@@ -399,7 +420,17 @@ async function* wrapConverseStreamResponse(
       usage = item.metadata?.usage;
       yield item;
     }
-    setUsage(span, usage);
+    span.setAttributes(
+      getTokenUsageAttributes(
+        // With prompt caching, `inputTokens` leaves out cache reads and writes, which the conventions count as input.
+        typeof usage?.inputTokens === 'number'
+          ? usage.inputTokens + (usage.cacheReadInputTokens ?? 0) + (usage.cacheWriteInputTokens ?? 0)
+          : undefined,
+        usage?.outputTokens,
+        usage?.cacheWriteInputTokens,
+        usage?.cacheReadInputTokens,
+      ),
+    );
   } catch (error) {
     setStreamErrorStatus(span, abortSignal);
     throw error;
@@ -418,15 +449,21 @@ function setStopReason(span: Span, stopReason: string | undefined): void {
   }
 }
 
-function setUsage(span: Span, usage: TokenUsage | undefined): void {
-  if (usage) {
-    const { inputTokens, outputTokens } = usage;
-    if (inputTokens !== undefined) {
-      span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, inputTokens);
-    }
-    if (outputTokens !== undefined) {
-      span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, outputTokens);
-    }
+/**
+ * `getTokenUsageAttributes` is not used here: `message_start` carries only a placeholder output token count,
+ * so a total computed from it would be wrong.
+ */
+function setClaudeStreamInputTokens(span: Span, usage: ParsedChunk): void {
+  const cacheReadTokens = usage?.cache_read_input_tokens;
+  const cacheWriteTokens = usage?.cache_creation_input_tokens;
+  if (typeof usage?.input_tokens === 'number') {
+    span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, usage.input_tokens + (cacheReadTokens ?? 0) + (cacheWriteTokens ?? 0));
+  }
+  if (typeof cacheReadTokens === 'number') {
+    span.setAttribute(GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, cacheReadTokens);
+  }
+  if (typeof cacheWriteTokens === 'number') {
+    span.setAttribute(GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, cacheWriteTokens);
   }
 }
 
@@ -471,9 +508,7 @@ function recordNovaAttributes(parsedChunk: ParsedChunk, span: Span): void {
 }
 
 function recordClaudeAttributes(parsedChunk: ParsedChunk, span: Span): void {
-  if (parsedChunk.message?.usage?.input_tokens !== undefined) {
-    span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, parsedChunk.message.usage.input_tokens);
-  }
+  setClaudeStreamInputTokens(span, parsedChunk.message?.usage);
   if (parsedChunk.message?.usage?.output_tokens !== undefined) {
     span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, parsedChunk.message.usage.output_tokens);
   }

packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts

@@ -4,6 +4,11 @@ import {
   GEN_AI_REQUEST_MAX_TOKENS,
   GEN_AI_REQUEST_MODEL,
   GEN_AI_REQUEST_TEMPERATURE,
+  GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS,
+  GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS,
+  GEN_AI_USAGE_INPUT_TOKENS,
+  GEN_AI_USAGE_OUTPUT_TOKENS,
+  GEN_AI_USAGE_TOTAL_TOKENS,
 } from '@sentry/conventions/attributes';
 import { GEN_AI_GENERATE_CONTENT } from '@sentry/conventions/op';
 import type { Span } from '@sentry/core';
@@ -156,3 +161,102 @@ describe('BedrockRuntimeServiceExtension stream spans', () => {
     expect(spanToJSON(span).end_timestamp).toBeDefined();
   });
 });
+
+describe('BedrockRuntimeServiceExtension token usage with prompt caching', () => {
+  const encoder = new TextEncoder();
+
+  // 20 uncached prompt tokens, 1800 read from the cache, and 300 written to it.
+  const EXPECTED_INPUT_USAGE = {
+    [GEN_AI_USAGE_INPUT_TOKENS]: 2120,
+    [GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]: 1800,
+    [GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS]: 300,
+  };
+
+  const EXPECTED_USAGE = {
+    ...EXPECTED_INPUT_USAGE,
+    [GEN_AI_USAGE_OUTPUT_TOKENS]: 50,
+    [GEN_AI_USAGE_TOTAL_TOKENS]: 2170,
+  };
+
+  const CLAUDE_USAGE = {
+    input_tokens: 20,
+    cache_read_input_tokens: 1800,
+    cache_creation_input_tokens: 300,
+    output_tokens: 50,
+  };
+
+  function usageAttributes(span: Span, expected: Record<string, unknown>): Record<string, unknown> {
+    const attributes = spanToJSON(span).attributes ?? {};
+    return Object.fromEntries(Object.keys(expected).map(key => [key, attributes[key]]));
+  }
+
+  beforeEach(() => {
+    const client = new TestClient(getDefaultTestClientOptions({ tracesSampleRate: 1 }));
+    setCurrentClient(client);
+    client.init();
+  });
+
+  afterEach(() => {
+    setCurrentClient(undefined as never);
+  });
+
+  it('counts cache reads and writes as input tokens for Converse', () => {
+    const span = startInactiveSpan({ name: 'bedrock' });
+
+    respond(
+      span,
+      'Converse',
+      { modelId: CLAUDE_MODEL_ID },
+      {
+        stopReason: 'end_turn',
+        usage: {
+          inputTokens: 20,
+          outputTokens: 50,
+          totalTokens: 2170,
+          cacheReadInputTokens: 1800,
+          cacheWriteInputTokens: 300,
+        },
+      },
+    );
+
+    expect(usageAttributes(span, EXPECTED_USAGE)).toEqual(EXPECTED_USAGE);
+  });
+
+  it('counts cache reads and writes as input tokens for Claude InvokeModel', () => {
+    const span = startInactiveSpan({ name: 'bedrock' });
+    const body = encoder.encode(
+      JSON.stringify({
+        id: 'msg_bdrk_01Fv2yWBqLQ8Pp4q9kJbDxXz',
+        type: 'message',
+        role: 'assistant',
+        content: [{ type: 'text', text: 'Paris.' }],
+        stop_reason: 'end_turn',
+        usage: CLAUDE_USAGE,
+      }),
+    );
+
+    respond(span, 'InvokeModel', { modelId: CLAUDE_MODEL_ID }, { body });
+
+    expect(usageAttributes(span, EXPECTED_USAGE)).toEqual(EXPECTED_USAGE);
+  });
+
+  it('counts cache reads and writes as input tokens for a Claude stream', async () => {
+    const span = startInactiveSpan({ name: 'bedrock' });
+    const messageStart = {
+      type: 'message_start',
+      message: { id: 'msg_bdrk_01Fv2yWBqLQ8Pp4q9kJbDxXz', role: 'assistant', usage: CLAUDE_USAGE },
+    };
+    const data = {
+      body: (async function* () {
+        yield { chunk: { bytes: encoder.encode(JSON.stringify(messageStart)) } };
+      })(),
+    };
+
+    respond(span, 'InvokeModelWithResponseStream', { modelId: CLAUDE_MODEL_ID }, data);
+    await drain(data.body);
+
+    // The total waits for the final output count, which `message_start` does not carry.
+    const expected = { ...EXPECTED_INPUT_USAGE, [GEN_AI_USAGE_TOTAL_TOKENS]: undefined };
+    expect(usageAttributes(span, expected)).toEqual(expected);
+  });
+});