From 3407bdf852bf63e589f5f7403a180792c29f53a4 Mon Sep 17 00:00:00 2001 From: Andrei Borza Date: Fri, 6 Nov 2026 12:49:46 +0100 Subject: [PATCH 1/2] fix(server-utils): Estimate Bedrock Mistral and Cohere stream output tokens from all chunks Co-Authored-By: Claude Opus 5.5 --- .../aws-sdk/services/bedrock-runtime.ts | 27 +++++++++++++------ .../test/integrations/aws-sdk-bedrock.test.ts | 22 +++++++++++++++ 2 files changed, 41 insertions(+), 8 deletions(-) diff --git a/packages/server-utils/src/integrations/aws-sdk/services/bedrock-runtime.ts b/packages/server-utils/src/integrations/aws-sdk/services/bedrock-runtime.ts index 7259dc085004..b00e081f8a83 100644 --- a/packages/server-utils/src/integrations/aws-sdk/services/bedrock-runtime.ts +++ b/packages/server-utils/src/integrations/aws-sdk/services/bedrock-runtime.ts @@ -373,6 +373,7 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { // Resolved once: the model family is fixed for the whole stream, and for unrecognized families // the chunks don't need to be parsed at all. const recordAttributes = typeof modelId === 'string' ? resolveStreamRecorder(modelId) : undefined; + const state: StreamState = { outputTextLength: 0 }; // Replace the original response body with our instrumented stream, deferring span.end() until the // entire stream is consumed. Downstream consumers still receive the full stream. @@ -382,7 +383,7 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { if (recordAttributes) { const parsedChunk = parseChunk(chunk?.chunk?.bytes); if (parsedChunk) { - recordAttributes(parsedChunk, span); + recordAttributes(parsedChunk, span, state); } } yield chunk; @@ -397,7 +398,14 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { } } -function resolveStreamRecorder(modelId: string): ((parsedChunk: ParsedChunk, span: Span) => void) | undefined { +interface StreamState { + // Models that send no output token count get an estimate from the text of all chunks. + outputTextLength: number; +} + +function resolveStreamRecorder( + modelId: string, +): ((parsedChunk: ParsedChunk, span: Span, state: StreamState) => void) | undefined { if (modelId.includes('amazon.titan')) return recordTitanAttributes; if (modelId.includes('anthropic.claude')) return recordClaudeAttributes; if (modelId.includes('amazon.nova')) return recordNovaAttributes; @@ -542,27 +550,30 @@ function recordLlamaAttributes(parsedChunk: ParsedChunk, span: Span): void { } } -function recordMistralAttributes(parsedChunk: ParsedChunk, span: Span): void { +function recordMistralAttributes(parsedChunk: ParsedChunk, span: Span, state: StreamState): void { if (parsedChunk.outputs?.[0]?.text !== undefined) { - span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(parsedChunk.outputs[0].text.length / 6)); + state.outputTextLength += parsedChunk.outputs[0].text.length; + span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(state.outputTextLength / 6)); } if (parsedChunk.outputs?.[0]?.stop_reason !== undefined) { span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [parsedChunk.outputs[0].stop_reason]); } } -function recordCohereAttributes(parsedChunk: ParsedChunk, span: Span): void { +function recordCohereAttributes(parsedChunk: ParsedChunk, span: Span, state: StreamState): void { if (parsedChunk.generations?.[0]?.text !== undefined) { - span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(parsedChunk.generations[0].text.length / 6)); + state.outputTextLength += parsedChunk.generations[0].text.length; + span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(state.outputTextLength / 6)); } if (parsedChunk.generations?.[0]?.finish_reason !== undefined) { span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [parsedChunk.generations[0].finish_reason]); } } -function recordCohereRAttributes(parsedChunk: ParsedChunk, span: Span): void { +function recordCohereRAttributes(parsedChunk: ParsedChunk, span: Span, state: StreamState): void { if (parsedChunk.text !== undefined) { - span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(parsedChunk.text.length / 6)); + state.outputTextLength += parsedChunk.text.length; + span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(state.outputTextLength / 6)); } if (parsedChunk.finish_reason !== undefined) { span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [parsedChunk.finish_reason]); diff --git a/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts b/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts index a1ba675251e6..2b9a1aa9d58b 100644 --- a/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts +++ b/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts @@ -172,6 +172,28 @@ describe('BedrockRuntimeServiceExtension stream spans', () => { expect(spanToJSON(span).attributes?.[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(42); }); + it.each([ + { modelId: 'mistral.mistral-large-2402-v1:0', toChunk: (text: string) => ({ outputs: [{ text }] }) }, + { modelId: 'cohere.command-text-v14', toChunk: (text: string) => ({ generations: [{ text }] }) }, + { modelId: 'cohere.command-r-plus-v1:0', toChunk: (text: string) => ({ text }) }, + ])('estimates $modelId stream output tokens from the text of all chunks', async ({ modelId, toChunk }) => { + const span = startInactiveSpan({ name: 'bedrock' }); + const encoder = new TextEncoder(); + const data = { + body: (async function* () { + for (const text of ['Answer:', 'Paris']) { + yield { chunk: { bytes: encoder.encode(JSON.stringify(toChunk(text))) } }; + } + })(), + }; + + respond(span, 'InvokeModelWithResponseStream', { modelId }, data); + await drain(data.body); + + // 12 characters at 6 characters per token. + expect(spanToJSON(span).attributes?.[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(2); + }); + it('ends an InvokeModelWithResponseStream span without a modelId when the stream is drained', async () => { const span = startInactiveSpan({ name: 'bedrock' }); const data = { body: (async function* () {})() }; From 83036eef70d5b361aff7b5b3fc514a06e57c6069 Mon Sep 17 00:00:00 2001 From: Andrei Borza Date: Fri, 6 Nov 2026 14:37:19 +0100 Subject: [PATCH 2/2] Read Mistral and Cohere stream token counts from invocation metrics --- .../aws-sdk/services/bedrock-runtime.ts | 45 +++++++++---------- .../test/integrations/aws-sdk-bedrock.test.ts | 44 ++++++++++++++---- 2 files changed, 56 insertions(+), 33 deletions(-) diff --git a/packages/server-utils/src/integrations/aws-sdk/services/bedrock-runtime.ts b/packages/server-utils/src/integrations/aws-sdk/services/bedrock-runtime.ts index b00e081f8a83..97854b017c8c 100644 --- a/packages/server-utils/src/integrations/aws-sdk/services/bedrock-runtime.ts +++ b/packages/server-utils/src/integrations/aws-sdk/services/bedrock-runtime.ts @@ -373,7 +373,6 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { // Resolved once: the model family is fixed for the whole stream, and for unrecognized families // the chunks don't need to be parsed at all. const recordAttributes = typeof modelId === 'string' ? resolveStreamRecorder(modelId) : undefined; - const state: StreamState = { outputTextLength: 0 }; // Replace the original response body with our instrumented stream, deferring span.end() until the // entire stream is consumed. Downstream consumers still receive the full stream. @@ -383,7 +382,7 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { if (recordAttributes) { const parsedChunk = parseChunk(chunk?.chunk?.bytes); if (parsedChunk) { - recordAttributes(parsedChunk, span, state); + recordAttributes(parsedChunk, span); } } yield chunk; @@ -398,14 +397,7 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { } } -interface StreamState { - // Models that send no output token count get an estimate from the text of all chunks. - outputTextLength: number; -} - -function resolveStreamRecorder( - modelId: string, -): ((parsedChunk: ParsedChunk, span: Span, state: StreamState) => void) | undefined { +function resolveStreamRecorder(modelId: string): ((parsedChunk: ParsedChunk, span: Span) => void) | undefined { if (modelId.includes('amazon.titan')) return recordTitanAttributes; if (modelId.includes('anthropic.claude')) return recordClaudeAttributes; if (modelId.includes('amazon.nova')) return recordNovaAttributes; @@ -550,31 +542,36 @@ function recordLlamaAttributes(parsedChunk: ParsedChunk, span: Span): void { } } -function recordMistralAttributes(parsedChunk: ParsedChunk, span: Span, state: StreamState): void { - if (parsedChunk.outputs?.[0]?.text !== undefined) { - state.outputTextLength += parsedChunk.outputs[0].text.length; - span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(state.outputTextLength / 6)); +/** + * Bedrock adds the token counts to the last chunk of a stream. For models that send no counts of their own, + * this is the only real count. + */ +function recordInvocationMetrics(parsedChunk: ParsedChunk, span: Span): void { + const metrics = parsedChunk['amazon-bedrock-invocationMetrics']; + if (typeof metrics?.inputTokenCount === 'number') { + span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, metrics.inputTokenCount); + } + if (typeof metrics?.outputTokenCount === 'number') { + span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, metrics.outputTokenCount); } +} + +function recordMistralAttributes(parsedChunk: ParsedChunk, span: Span): void { + recordInvocationMetrics(parsedChunk, span); if (parsedChunk.outputs?.[0]?.stop_reason !== undefined) { span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [parsedChunk.outputs[0].stop_reason]); } } -function recordCohereAttributes(parsedChunk: ParsedChunk, span: Span, state: StreamState): void { - if (parsedChunk.generations?.[0]?.text !== undefined) { - state.outputTextLength += parsedChunk.generations[0].text.length; - span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(state.outputTextLength / 6)); - } +function recordCohereAttributes(parsedChunk: ParsedChunk, span: Span): void { + recordInvocationMetrics(parsedChunk, span); if (parsedChunk.generations?.[0]?.finish_reason !== undefined) { span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [parsedChunk.generations[0].finish_reason]); } } -function recordCohereRAttributes(parsedChunk: ParsedChunk, span: Span, state: StreamState): void { - if (parsedChunk.text !== undefined) { - state.outputTextLength += parsedChunk.text.length; - span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(state.outputTextLength / 6)); - } +function recordCohereRAttributes(parsedChunk: ParsedChunk, span: Span): void { + recordInvocationMetrics(parsedChunk, span); if (parsedChunk.finish_reason !== undefined) { span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [parsedChunk.finish_reason]); } diff --git a/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts b/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts index 2b9a1aa9d58b..a63e70d3324c 100644 --- a/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts +++ b/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts @@ -172,17 +172,40 @@ describe('BedrockRuntimeServiceExtension stream spans', () => { expect(spanToJSON(span).attributes?.[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(42); }); - it.each([ - { modelId: 'mistral.mistral-large-2402-v1:0', toChunk: (text: string) => ({ outputs: [{ text }] }) }, - { modelId: 'cohere.command-text-v14', toChunk: (text: string) => ({ generations: [{ text }] }) }, - { modelId: 'cohere.command-r-plus-v1:0', toChunk: (text: string) => ({ text }) }, - ])('estimates $modelId stream output tokens from the text of all chunks', async ({ modelId, toChunk }) => { + const INVOCATION_METRICS = { + inputTokenCount: 11, + outputTokenCount: 3, + invocationLatency: 412, + firstByteLatency: 198, + }; + + it.each( + [ + { modelId: 'mistral.mistral-large-2402-v1:0', toChunk: (text: string) => ({ outputs: [{ text }] }) }, + { modelId: 'cohere.command-text-v14', toChunk: (text: string) => ({ generations: [{ text }] }) }, + { modelId: 'cohere.command-r-plus-v1:0', toChunk: (text: string) => ({ text }) }, + ].flatMap(family => [ + { + ...family, + when: 'has invocation metrics', + metrics: INVOCATION_METRICS, + expected: { [GEN_AI_USAGE_INPUT_TOKENS]: 11, [GEN_AI_USAGE_OUTPUT_TOKENS]: 3 }, + }, + { + ...family, + when: 'has no invocation metrics', + metrics: undefined, + expected: { [GEN_AI_USAGE_INPUT_TOKENS]: undefined, [GEN_AI_USAGE_OUTPUT_TOKENS]: undefined }, + }, + ]), + )('records only real token counts when a $modelId stream $when', async ({ modelId, toChunk, metrics, expected }) => { const span = startInactiveSpan({ name: 'bedrock' }); const encoder = new TextEncoder(); + const chunks = [toChunk('The capital'), { ...toChunk(' is Paris.'), 'amazon-bedrock-invocationMetrics': metrics }]; const data = { body: (async function* () { - for (const text of ['Answer:', 'Paris']) { - yield { chunk: { bytes: encoder.encode(JSON.stringify(toChunk(text))) } }; + for (const chunk of chunks) { + yield { chunk: { bytes: encoder.encode(JSON.stringify(chunk)) } }; } })(), }; @@ -190,8 +213,11 @@ describe('BedrockRuntimeServiceExtension stream spans', () => { respond(span, 'InvokeModelWithResponseStream', { modelId }, data); await drain(data.body); - // 12 characters at 6 characters per token. - expect(spanToJSON(span).attributes?.[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(2); + const attributes = spanToJSON(span).attributes; + expect({ + [GEN_AI_USAGE_INPUT_TOKENS]: attributes?.[GEN_AI_USAGE_INPUT_TOKENS], + [GEN_AI_USAGE_OUTPUT_TOKENS]: attributes?.[GEN_AI_USAGE_OUTPUT_TOKENS], + }).toEqual(expected); }); it('ends an InvokeModelWithResponseStream span without a modelId when the stream is drained', async () => {