fix(server-utils): Use cumulative input and cache token counts from streamed message_delta

packages/server-utils/src/ai/anthropic-ai/streaming.ts

@@ -16,7 +16,9 @@ interface StreamingState {
   responseId: string;
   /** The model name. */
   responseModel: string;
-  /** Number of prompt/input tokens used. */
+  /** Number of input tokens, not counting cache reads and writes. */
+  inputTokens: number | undefined;
+  /** Number of prompt/input tokens used, including cache reads and writes. */
   promptTokens: number | undefined;
   /** Number of completion/output tokens used. */
   completionTokens: number | undefined;
@@ -68,8 +70,12 @@ function handleMessageMetadata(event: AnthropicAiStreamingEvent, state: Streamin
   // Cumulative token counts and the final stop reason both arrive on the message_delta event.
   // @see https://docs.anthropic.com/en/docs/build-with-claude/streaming#event-types
   if (event.type === 'message_delta') {
-    if (event.usage && typeof event.usage.output_tokens === 'number') {
-      state.completionTokens = event.usage.output_tokens;
+    if (event.usage) {
+      // Input and cache counts can grow after message_start (e.g. with server tools), so the last value wins.
+      applyInputUsage(event.usage, state);
+      if (typeof event.usage.output_tokens === 'number') {
+        state.completionTokens = event.usage.output_tokens;
+      }
     }
     if (event.delta?.stop_reason) {
       state.finishReasons.push(event.delta.stop_reason);
@@ -83,20 +89,32 @@ function handleMessageMetadata(event: AnthropicAiStreamingEvent, state: Streamin
     if (message.model) state.responseModel = message.model;
 
     if (message.usage) {
-      if (typeof message.usage.input_tokens === 'number') {
-        state.promptTokens =
-          message.usage.input_tokens +
-          (message.usage.cache_creation_input_tokens ?? 0) +
-          (message.usage.cache_read_input_tokens ?? 0);
-      }
-      if (typeof message.usage.cache_creation_input_tokens === 'number')
-        state.cacheCreationInputTokens = message.usage.cache_creation_input_tokens;
-      if (typeof message.usage.cache_read_input_tokens === 'number')
-        state.cacheReadInputTokens = message.usage.cache_read_input_tokens;
+      applyInputUsage(message.usage, state);
     }
   }
 }
 
+/**
+ * Records the input and cache token counts of a usage object that are numbers.
+ */
+function applyInputUsage(
+  usage: {
+    input_tokens?: number | null;
+    cache_creation_input_tokens?: number | null;
+    cache_read_input_tokens?: number | null;
+  },
+  state: StreamingState,
+): void {
+  if (typeof usage.input_tokens === 'number') state.inputTokens = usage.input_tokens;
+  if (typeof usage.cache_creation_input_tokens === 'number')
+    state.cacheCreationInputTokens = usage.cache_creation_input_tokens;
+  if (typeof usage.cache_read_input_tokens === 'number') state.cacheReadInputTokens = usage.cache_read_input_tokens;
+
+  if (typeof state.inputTokens === 'number') {
+    state.promptTokens = state.inputTokens + (state.cacheCreationInputTokens ?? 0) + (state.cacheReadInputTokens ?? 0);
+  }
+}
+
 /**
  * Handle start of a content block (e.g., tool_use)
  */
@@ -215,6 +233,7 @@ export async function* instrumentAsyncIterableStream(
     finishReasons: [],
     responseId: '',
     responseModel: '',
+    inputTokens: undefined,
     promptTokens: undefined,
     completionTokens: undefined,
     cacheCreationInputTokens: undefined,
@@ -246,6 +265,7 @@ export function instrumentMessageStream<R extends { on: (...args: unknown[]) =>
     finishReasons: [],
     responseId: '',
     responseModel: '',
+    inputTokens: undefined,
     promptTokens: undefined,
     completionTokens: undefined,
     cacheCreationInputTokens: undefined,

packages/server-utils/src/ai/anthropic-ai/types.ts

@@ -96,6 +96,10 @@ export type AnthropicAiStreamingEvent = {
   };
   usage?: {
     output_tokens: number; // Final total output tokens; emitted on the last `message_delta` event
+    // Cumulative input counts, which can be larger than at `message_start`
+    input_tokens?: number | null;
+    cache_creation_input_tokens?: number | null;
+    cache_read_input_tokens?: number | null;
   };
   message?: AnthropicAiMessage;
   /** Present for fine-grained tool streaming */

packages/server-utils/test/ai/lib/tracing/anthropic-ai-usage.test.ts

@@ -124,3 +124,150 @@ describe.each(['response', 'async iterable', 'message stream'] as const)('Anthro
     });
   });
 });
+
+describe.each(['async iterable', 'message stream'] as const)(
+  'Anthropic streamed input usage from message_delta (%s)',
+  mode => {
+    it('uses the cumulative input and cache counts of the final message_delta', async () => {
+      const { span, attributes } = createMockSpan();
+      // With server tools, message_delta reports input and cache counts that grew past message_start's.
+      const events: AnthropicAiStreamingEvent[] = [
+        {
+          type: 'message_start',
+          message: {
+            id: 'msg_server_tool',
+            type: 'message',
+            model: 'claude-sonnet-4-6',
+            role: 'assistant',
+            content: [],
+            stop_reason: null,
+            stop_sequence: null,
+            usage: {
+              input_tokens: 2285,
+              output_tokens: 25,
+              cache_creation_input_tokens: 0,
+              cache_read_input_tokens: 0,
+            },
+          },
+        },
+        {
+          type: 'message_delta',
+          usage: {
+            input_tokens: 2292,
+            output_tokens: 150,
+            cache_creation_input_tokens: 12032,
+            cache_read_input_tokens: 0,
+          },
+        },
+        { type: 'message_stop' },
+      ];
+
+      if (mode === 'async iterable') {
+        const stream = (async function* () {
+          yield* events;
+        })();
+        for await (const _ of instrumentAsyncIterableStream(stream, span, false)) {
+          void _;
+        }
+      } else {
+        const stream = instrumentMessageStream(new EventEmitter(), span, false);
+        for (const event of events) {
+          stream.emit('streamEvent', event);
+        }
+        stream.emit('message', {});
+      }
+
+      expect(attributes).toMatchObject({
+        [GEN_AI_USAGE_INPUT_TOKENS]: 2292 + 12032,
+        [GEN_AI_USAGE_OUTPUT_TOKENS]: 150,
+        [GEN_AI_USAGE_TOTAL_TOKENS]: 2292 + 12032 + 150,
+        [GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS]: 12032,
+        [GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]: 0,
+      });
+    });
+
+    it('keeps the message_start input counts that message_delta omits or sends as null', async () => {
+      const { span, attributes } = createMockSpan();
+      const events: AnthropicAiStreamingEvent[] = [
+        {
+          type: 'message_start',
+          message: {
+            id: 'msg_keep',
+            type: 'message',
+            model: 'claude-sonnet-4-6',
+            role: 'assistant',
+            content: [],
+            stop_reason: null,
+            stop_sequence: null,
+            usage: { input_tokens: 10, output_tokens: 1, cache_read_input_tokens: 40 },
+          },
+        },
+        { type: 'message_delta', usage: { input_tokens: null, cache_read_input_tokens: null, output_tokens: 12 } },
+        { type: 'message_stop' },
+      ];
+
+      if (mode === 'async iterable') {
+        const stream = (async function* () {
+          yield* events;
+        })();
+        for await (const _ of instrumentAsyncIterableStream(stream, span, false)) {
+          void _;
+        }
+      } else {
+        const stream = instrumentMessageStream(new EventEmitter(), span, false);
+        for (const event of events) {
+          stream.emit('streamEvent', event);
+        }
+        stream.emit('message', {});
+      }
+
+      expect(attributes).toMatchObject({
+        [GEN_AI_USAGE_INPUT_TOKENS]: 50,
+        [GEN_AI_USAGE_OUTPUT_TOKENS]: 12,
+        [GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS]: 40,
+      });
+    });
+
+    it('takes the input count from message_delta when message_start reports zero', async () => {
+      const { span, attributes } = createMockSpan();
+      const events: AnthropicAiStreamingEvent[] = [
+        {
+          type: 'message_start',
+          message: {
+            id: 'msg_zero_start',
+            type: 'message',
+            model: 'claude-sonnet-4-6',
+            role: 'assistant',
+            content: [],
+            stop_reason: null,
+            stop_sequence: null,
+            usage: { input_tokens: 0, output_tokens: 1 },
+          },
+        },
+        { type: 'message_delta', usage: { input_tokens: 1500, output_tokens: 40 } },
+        { type: 'message_stop' },
+      ];
+
+      if (mode === 'async iterable') {
+        const stream = (async function* () {
+          yield* events;
+        })();
+        for await (const _ of instrumentAsyncIterableStream(stream, span, false)) {
+          void _;
+        }
+      } else {
+        const stream = instrumentMessageStream(new EventEmitter(), span, false);
+        for (const event of events) {
+          stream.emit('streamEvent', event);
+        }
+        stream.emit('message', {});
+      }
+
+      expect(attributes).toMatchObject({
+        [GEN_AI_USAGE_INPUT_TOKENS]: 1500,
+        [GEN_AI_USAGE_OUTPUT_TOKENS]: 40,
+        [GEN_AI_USAGE_TOTAL_TOKENS]: 1540,
+      });
+    });
+  },
+);