From 4ac20c7b82366c6367dc3c76cec1f4798ffee7d4 Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Wed, 21 Oct 2026 09:38:34 -0400 Subject: [PATCH 1/2] fix(anthropic-ai): use cumulative input and cache token counts from streamed message_delta Co-Authored-By: Claude Sonnet 5.5 --- .../src/ai/anthropic-ai/streaming.ts | 46 +++++--- .../server-utils/src/ai/anthropic-ai/types.ts | 4 + .../ai/lib/tracing/anthropic-ai-usage.test.ts | 105 ++++++++++++++++++ 3 files changed, 142 insertions(+), 13 deletions(-) diff --git a/packages/server-utils/src/ai/anthropic-ai/streaming.ts b/packages/server-utils/src/ai/anthropic-ai/streaming.ts index 20507ed4f3bf..1c3ed58674d6 100644 --- a/packages/server-utils/src/ai/anthropic-ai/streaming.ts +++ b/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 finishReasons: [], responseId: '', responseModel: '', + inputTokens: undefined, promptTokens: undefined, completionTokens: undefined, cacheCreationInputTokens: undefined, diff --git a/packages/server-utils/src/ai/anthropic-ai/types.ts b/packages/server-utils/src/ai/anthropic-ai/types.ts index bb768e63063e..ddee39046d57 100644 --- a/packages/server-utils/src/ai/anthropic-ai/types.ts +++ b/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 */ diff --git a/packages/server-utils/test/ai/lib/tracing/anthropic-ai-usage.test.ts b/packages/server-utils/test/ai/lib/tracing/anthropic-ai-usage.test.ts index aa62ff8d78ba..1a3cf3629cdf 100644 --- a/packages/server-utils/test/ai/lib/tracing/anthropic-ai-usage.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/anthropic-ai-usage.test.ts @@ -124,3 +124,108 @@ 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, + }); + }); + }, +); From dfd9154eb1dccc9f96ed63f2852320e7ea4d47dc Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Thu, 22 Oct 2026 16:06:51 -0400 Subject: [PATCH 2/2] test(anthropic-ai): cover a zero input count at message_start filled in by message_delta Co-Authored-By: Claude Sonnet 5.5 --- .../ai/lib/tracing/anthropic-ai-usage.test.ts | 42 +++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/packages/server-utils/test/ai/lib/tracing/anthropic-ai-usage.test.ts b/packages/server-utils/test/ai/lib/tracing/anthropic-ai-usage.test.ts index 1a3cf3629cdf..c506fbcb9387 100644 --- a/packages/server-utils/test/ai/lib/tracing/anthropic-ai-usage.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/anthropic-ai-usage.test.ts @@ -227,5 +227,47 @@ describe.each(['async iterable', 'message stream'] as const)( [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, + }); + }); }, );