From 84a0aa9e91fc08c22692a368b11438572bc74bfc Mon Sep 17 00:00:00 2001 From: Andrei Borza Date: Thu, 9 Jul 2026 23:19:29 +0200 Subject: [PATCH 1/3] feat(server-utils): Port Bedrock aws-sdk extension and add integration tests Ports the BedrockRuntime gen_ai service extension from the OTel aws-sdk integration to the orchestrion channel integration: chat span name and request parameters for `Converse`/`ConverseStream`, per-model-family request body parsing for `InvokeModel`/`InvokeModelWithResponseStream` (titan, nova, claude, llama, cohere, mistral), token usage and finish reasons from responses, and deferred span end for the streaming commands (the span ends when the wrapped stream is consumed). Adds nock-mocked integration tests for the non-streaming `Converse` and `InvokeModel` commands, asserted against both the OTel and diagnostics-channel paths. The streaming commands use the AWS binary eventstream framing which nock cannot mock, so they stay untested here. This completes orchestrion parity with the OTel aws-sdk integration. Fixes #20946 Co-Authored-By: Claude Fable 5 --- .../aws-serverless/bedrock/instrument.mjs | 9 + .../aws-serverless/bedrock/scenario.mjs | 84 +++ .../suites/aws-serverless/bedrock/test.ts | 75 +++ .../tracing-channel/aws-sdk/constants.ts | 5 + .../aws-sdk/services/ServicesExtensions.ts | 2 + .../aws-sdk/services/bedrock-runtime.ts | 480 ++++++++++++++++++ 6 files changed, 655 insertions(+) create mode 100644 dev-packages/node-integration-tests/suites/aws-serverless/bedrock/instrument.mjs create mode 100644 dev-packages/node-integration-tests/suites/aws-serverless/bedrock/scenario.mjs create mode 100644 dev-packages/node-integration-tests/suites/aws-serverless/bedrock/test.ts create mode 100644 packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/bedrock-runtime.ts diff --git a/dev-packages/node-integration-tests/suites/aws-serverless/bedrock/instrument.mjs b/dev-packages/node-integration-tests/suites/aws-serverless/bedrock/instrument.mjs new file mode 100644 index 000000000000..fe1c5c47983f --- /dev/null +++ b/dev-packages/node-integration-tests/suites/aws-serverless/bedrock/instrument.mjs @@ -0,0 +1,9 @@ +import * as Sentry from '@sentry/aws-serverless'; +import { loggingTransport } from '@sentry-internal/node-integration-tests'; + +Sentry.init({ + dsn: 'https://public@dsn.ingest.sentry.io/1337', + release: '1.0', + tracesSampleRate: 1.0, + transport: loggingTransport, +}); diff --git a/dev-packages/node-integration-tests/suites/aws-serverless/bedrock/scenario.mjs b/dev-packages/node-integration-tests/suites/aws-serverless/bedrock/scenario.mjs new file mode 100644 index 000000000000..c1705ff0d403 --- /dev/null +++ b/dev-packages/node-integration-tests/suites/aws-serverless/bedrock/scenario.mjs @@ -0,0 +1,84 @@ +import * as Sentry from '@sentry/aws-serverless'; +import { BedrockRuntimeClient, ConverseCommand, InvokeModelCommand } from '@aws-sdk/client-bedrock-runtime'; +// Force the HTTP/1 handler so `nock` can intercept the request instead of hitting real AWS. +import { NodeHttpHandler } from '@smithy/node-http-handler'; +import nock from 'nock'; + +nock.disableNetConnect(); + +const region = 'us-east-1'; +const credentials = { accessKeyId: 'aws-test-key', secretAccessKey: 'aws-test-secret' }; +const host = `https://bedrock-runtime.${region}.amazonaws.com`; + +async function converse() { + const client = new BedrockRuntimeClient({ + region, + credentials, + maxAttempts: 1, + requestHandler: new NodeHttpHandler(), + }); + + nock(host) + .post(/\/model\/.*\/converse$/) + .reply( + 200, + JSON.stringify({ + output: { message: { role: 'assistant', content: [{ text: 'Hello from Bedrock' }] } }, + stopReason: 'end_turn', + usage: { inputTokens: 12, outputTokens: 8, totalTokens: 20 }, + }), + { 'content-type': 'application/json' }, + ); + + await client.send( + new ConverseCommand({ + modelId: 'anthropic.claude-3-5-sonnet-20240620-v1:0', + messages: [{ role: 'user', content: [{ text: 'Hello' }] }], + inferenceConfig: { maxTokens: 100, temperature: 0.5, topP: 0.9 }, + }), + ); +} + +async function invokeModel() { + const client = new BedrockRuntimeClient({ + region, + credentials, + maxAttempts: 1, + requestHandler: new NodeHttpHandler(), + }); + + nock(host) + .post(/\/model\/.*\/invoke$/) + .reply( + 200, + JSON.stringify({ + content: [{ type: 'text', text: 'Hello from Bedrock' }], + stop_reason: 'end_turn', + usage: { input_tokens: 15, output_tokens: 9 }, + }), + { 'content-type': 'application/json' }, + ); + + await client.send( + new InvokeModelCommand({ + modelId: 'anthropic.claude-3-5-sonnet-20240620-v1:0', + contentType: 'application/json', + body: JSON.stringify({ + anthropic_version: 'bedrock-2023-05-31', + max_tokens: 100, + temperature: 0.5, + top_p: 0.9, + messages: [{ role: 'user', content: 'Hello' }], + }), + }), + ); +} + +async function run() { + await Sentry.startSpan({ name: 'Test Transaction' }, async () => { + await converse(); + await invokeModel(); + }); +} + +run(); diff --git a/dev-packages/node-integration-tests/suites/aws-serverless/bedrock/test.ts b/dev-packages/node-integration-tests/suites/aws-serverless/bedrock/test.ts new file mode 100644 index 000000000000..e8e6e9c337e5 --- /dev/null +++ b/dev-packages/node-integration-tests/suites/aws-serverless/bedrock/test.ts @@ -0,0 +1,75 @@ +import type { TransactionEvent } from '@sentry/core'; +import { afterAll, describe, expect } from 'vitest'; +import { isOrchestrionEnabled } from '../../../utils'; +import { cleanupChildProcesses, createEsmAndCjsTests } from '../../../utils/runner'; + +// The suite runs twice on CI: once with the OTel `Aws` integration (default) and once with the +// orchestrion diagnostics-channel integration auto-injected (`INJECT_ORCHESTRION`). Both emit the +// same gen_ai spans; only the origin differs. +const ORIGIN = isOrchestrionEnabled() ? 'auto.aws.orchestrion.aws-sdk' : 'auto.otel.aws'; + +const MODEL_ID = 'anthropic.claude-3-5-sonnet-20240620-v1:0'; + +function assertBedrockSpans(transaction: TransactionEvent): void { + const spans = transaction.spans ?? []; + + expect(transaction.transaction).toBe('Test Transaction'); + + // Converse (non-streaming) + expect(spans, 'expected a Bedrock Converse span').toContainEqual( + expect.objectContaining({ + description: `chat ${MODEL_ID}`, + origin: ORIGIN, + status: 'ok', + data: expect.objectContaining({ + 'sentry.origin': ORIGIN, + 'gen_ai.system': 'aws.bedrock', + 'gen_ai.operation.name': 'chat', + 'gen_ai.request.model': MODEL_ID, + 'gen_ai.request.max_tokens': 100, + 'gen_ai.request.temperature': 0.5, + 'gen_ai.request.top_p': 0.9, + 'gen_ai.usage.input_tokens': 12, + 'gen_ai.usage.output_tokens': 8, + 'gen_ai.response.finish_reasons': ['end_turn'], + }), + }), + ); + + // InvokeModel (non-streaming, anthropic.claude request/response body) + expect(spans, 'expected a Bedrock InvokeModel span').toContainEqual( + expect.objectContaining({ + origin: ORIGIN, + status: 'ok', + data: expect.objectContaining({ + 'sentry.origin': ORIGIN, + 'gen_ai.system': 'aws.bedrock', + 'gen_ai.request.model': MODEL_ID, + 'gen_ai.request.max_tokens': 100, + 'gen_ai.request.temperature': 0.5, + 'gen_ai.request.top_p': 0.9, + 'gen_ai.usage.input_tokens': 15, + 'gen_ai.usage.output_tokens': 9, + 'gen_ai.response.finish_reasons': ['end_turn'], + }), + }), + ); +} + +describe('awsIntegration - Bedrock', () => { + afterAll(() => { + cleanupChildProcesses(); + }); + + createEsmAndCjsTests( + __dirname, + 'scenario.mjs', + 'instrument.mjs', + (createTestRunner, test) => { + test('auto-instruments Bedrock Converse and InvokeModel', { timeout: 90_000 }, async () => { + await createTestRunner().ignore('event').expect({ transaction: assertBedrockSpans }).start().completed(); + }); + }, + { additionalDependencies: { '@aws-sdk/client-bedrock-runtime': '^3.1046.0' } }, + ); +}); diff --git a/packages/server-utils/src/integrations/tracing-channel/aws-sdk/constants.ts b/packages/server-utils/src/integrations/tracing-channel/aws-sdk/constants.ts index b466c61437dc..e57a4d1e153a 100644 --- a/packages/server-utils/src/integrations/tracing-channel/aws-sdk/constants.ts +++ b/packages/server-utils/src/integrations/tracing-channel/aws-sdk/constants.ts @@ -60,3 +60,8 @@ export const ATTR_FAAS_EXECUTION = 'faas.execution'; export const ATTR_MESSAGING_DESTINATION = 'messaging.destination'; export const ATTR_MESSAGING_DESTINATION_KIND = 'messaging.destination_kind'; export const MESSAGING_DESTINATION_KIND_VALUE_TOPIC = 'topic'; + +// Bedrock (gen_ai) +export const ATTR_GEN_AI_REQUEST_STOP_SEQUENCES = 'gen_ai.request.stop_sequences'; +export const GEN_AI_OPERATION_NAME_VALUE_CHAT = 'chat'; +export const GEN_AI_SYSTEM_VALUE_AWS_BEDROCK = 'aws.bedrock'; diff --git a/packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/ServicesExtensions.ts b/packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/ServicesExtensions.ts index ac6bb90fa71e..4d4a18142791 100644 --- a/packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/ServicesExtensions.ts +++ b/packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/ServicesExtensions.ts @@ -1,5 +1,6 @@ import type { Span } from '@sentry/core'; import type { NormalizedRequest, NormalizedResponse, RequestMetadata } from '../types'; +import { BedrockRuntimeServiceExtension } from './bedrock-runtime'; import { DynamodbServiceExtension } from './dynamodb'; import { KinesisServiceExtension } from './kinesis'; import { LambdaServiceExtension } from './lambda'; @@ -25,6 +26,7 @@ export class ServicesExtensions implements ServiceExtension { this._services.set('Lambda', new LambdaServiceExtension()); this._services.set('S3', new S3ServiceExtension()); this._services.set('Kinesis', new KinesisServiceExtension()); + this._services.set('BedrockRuntime', new BedrockRuntimeServiceExtension()); } public requestPreSpanHook(request: NormalizedRequest): RequestMetadata { diff --git a/packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/bedrock-runtime.ts b/packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/bedrock-runtime.ts new file mode 100644 index 000000000000..60de948a0fbc --- /dev/null +++ b/packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/bedrock-runtime.ts @@ -0,0 +1,480 @@ +import type { Span } from '@sentry/core'; +import { debug } from '@sentry/core'; +import { + GEN_AI_OPERATION_NAME, + GEN_AI_REQUEST_MAX_TOKENS, + GEN_AI_REQUEST_MODEL, + GEN_AI_REQUEST_TEMPERATURE, + GEN_AI_REQUEST_TOP_P, + GEN_AI_RESPONSE_FINISH_REASONS, + GEN_AI_SYSTEM, + GEN_AI_USAGE_INPUT_TOKENS, + GEN_AI_USAGE_OUTPUT_TOKENS, +} from '@sentry/conventions/attributes'; +import { DEBUG_BUILD } from '../../../../debug-build'; +import { + ATTR_GEN_AI_REQUEST_STOP_SEQUENCES, + GEN_AI_OPERATION_NAME_VALUE_CHAT, + GEN_AI_SYSTEM_VALUE_AWS_BEDROCK, +} from '../constants'; +import type { NormalizedRequest, NormalizedResponse } from '../types'; +import type { RequestMetadata, ServiceExtension } from './ServiceExtension'; + +// Simplified types inlined from @aws-sdk/client-bedrock-runtime +// Only the fields accessed by this instrumentation are included +interface TokenUsage { + inputTokens: number | undefined; + outputTokens: number | undefined; + totalTokens: number | undefined; +} + +interface ConverseStreamOutput { + messageStop?: { stopReason?: string }; + metadata?: { usage?: TokenUsage }; + [key: string]: any; +} + +export class BedrockRuntimeServiceExtension implements ServiceExtension { + public requestPreSpanHook(request: NormalizedRequest): RequestMetadata { + switch (request.commandName) { + case 'Converse': + return this._requestPreSpanHookConverse(request, false); + case 'ConverseStream': + return this._requestPreSpanHookConverse(request, true); + case 'InvokeModel': + return this._requestPreSpanHookInvokeModel(request, false); + case 'InvokeModelWithResponseStream': + return this._requestPreSpanHookInvokeModel(request, true); + } + + return {}; + } + + public responseHook(response: NormalizedResponse, span: Span): void { + if (!span.isRecording()) { + return; + } + + switch (response.request.commandName) { + case 'Converse': + return this._responseHookConverse(response, span); + case 'ConverseStream': + return this._responseHookConverseStream(response, span); + case 'InvokeModel': + return this._responseHookInvokeModel(response, span); + case 'InvokeModelWithResponseStream': + return this._responseHookInvokeModelWithResponseStream(response, span); + } + } + + private _requestPreSpanHookConverse(request: NormalizedRequest, isStream: boolean): RequestMetadata { + let spanName = GEN_AI_OPERATION_NAME_VALUE_CHAT; + const spanAttributes: Record = { + // oxlint-disable-next-line typescript/no-deprecated + [GEN_AI_SYSTEM]: GEN_AI_SYSTEM_VALUE_AWS_BEDROCK, + [GEN_AI_OPERATION_NAME]: GEN_AI_OPERATION_NAME_VALUE_CHAT, + }; + + const modelId = request.commandInput.modelId; + if (modelId) { + spanAttributes[GEN_AI_REQUEST_MODEL] = modelId; + if (spanName) { + spanName += ` ${modelId}`; + } + } + + const inferenceConfig = request.commandInput.inferenceConfig; + if (inferenceConfig) { + const { maxTokens, temperature, topP, stopSequences } = inferenceConfig; + if (maxTokens !== undefined) { + spanAttributes[GEN_AI_REQUEST_MAX_TOKENS] = maxTokens; + } + if (temperature !== undefined) { + spanAttributes[GEN_AI_REQUEST_TEMPERATURE] = temperature; + } + if (topP !== undefined) { + spanAttributes[GEN_AI_REQUEST_TOP_P] = topP; + } + if (stopSequences !== undefined) { + spanAttributes[ATTR_GEN_AI_REQUEST_STOP_SEQUENCES] = stopSequences; + } + } + + return { + spanName, + isStream, + spanAttributes, + }; + } + + private _requestPreSpanHookInvokeModel(request: NormalizedRequest, isStream: boolean): RequestMetadata { + const spanAttributes: Record = { + // oxlint-disable-next-line typescript/no-deprecated + [GEN_AI_SYSTEM]: GEN_AI_SYSTEM_VALUE_AWS_BEDROCK, + }; + + const modelId = request.commandInput?.modelId; + if (modelId) { + spanAttributes[GEN_AI_REQUEST_MODEL] = modelId; + } + + if (request.commandInput?.body) { + const requestBody = JSON.parse(request.commandInput.body); + if (modelId.includes('amazon.titan')) { + if (requestBody.textGenerationConfig?.temperature !== undefined) { + spanAttributes[GEN_AI_REQUEST_TEMPERATURE] = requestBody.textGenerationConfig.temperature; + } + if (requestBody.textGenerationConfig?.topP !== undefined) { + spanAttributes[GEN_AI_REQUEST_TOP_P] = requestBody.textGenerationConfig.topP; + } + if (requestBody.textGenerationConfig?.maxTokenCount !== undefined) { + spanAttributes[GEN_AI_REQUEST_MAX_TOKENS] = requestBody.textGenerationConfig.maxTokenCount; + } + if (requestBody.textGenerationConfig?.stopSequences !== undefined) { + spanAttributes[ATTR_GEN_AI_REQUEST_STOP_SEQUENCES] = requestBody.textGenerationConfig.stopSequences; + } + } else if (modelId.includes('amazon.nova')) { + if (requestBody.inferenceConfig?.temperature !== undefined) { + spanAttributes[GEN_AI_REQUEST_TEMPERATURE] = requestBody.inferenceConfig.temperature; + } + if (requestBody.inferenceConfig?.top_p !== undefined) { + spanAttributes[GEN_AI_REQUEST_TOP_P] = requestBody.inferenceConfig.top_p; + } + if (requestBody.inferenceConfig?.max_new_tokens !== undefined) { + spanAttributes[GEN_AI_REQUEST_MAX_TOKENS] = requestBody.inferenceConfig.max_new_tokens; + } + if (requestBody.inferenceConfig?.stopSequences !== undefined) { + spanAttributes[ATTR_GEN_AI_REQUEST_STOP_SEQUENCES] = requestBody.inferenceConfig.stopSequences; + } + } else if (modelId.includes('anthropic.claude')) { + if (requestBody.max_tokens !== undefined) { + spanAttributes[GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens; + } + if (requestBody.temperature !== undefined) { + spanAttributes[GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature; + } + if (requestBody.top_p !== undefined) { + spanAttributes[GEN_AI_REQUEST_TOP_P] = requestBody.top_p; + } + if (requestBody.stop_sequences !== undefined) { + spanAttributes[ATTR_GEN_AI_REQUEST_STOP_SEQUENCES] = requestBody.stop_sequences; + } + } else if (modelId.includes('meta.llama')) { + if (requestBody.max_gen_len !== undefined) { + spanAttributes[GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_gen_len; + } + if (requestBody.temperature !== undefined) { + spanAttributes[GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature; + } + if (requestBody.top_p !== undefined) { + spanAttributes[GEN_AI_REQUEST_TOP_P] = requestBody.top_p; + } + // request for meta llama models does not contain stop_sequences field + } else if (modelId.includes('cohere.command-r')) { + if (requestBody.max_tokens !== undefined) { + spanAttributes[GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens; + } + if (requestBody.temperature !== undefined) { + spanAttributes[GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature; + } + if (requestBody.p !== undefined) { + spanAttributes[GEN_AI_REQUEST_TOP_P] = requestBody.p; + } + if (requestBody.message !== undefined) { + // NOTE: We approximate the token count since this value is not directly available in the body. + // According to Bedrock docs they use (total_chars / 6) to approximate token count for pricing. + spanAttributes[GEN_AI_USAGE_INPUT_TOKENS] = Math.ceil(requestBody.message.length / 6); + } + if (requestBody.stop_sequences !== undefined) { + spanAttributes[ATTR_GEN_AI_REQUEST_STOP_SEQUENCES] = requestBody.stop_sequences; + } + } else if (modelId.includes('cohere.command')) { + if (requestBody.max_tokens !== undefined) { + spanAttributes[GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens; + } + if (requestBody.temperature !== undefined) { + spanAttributes[GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature; + } + if (requestBody.p !== undefined) { + spanAttributes[GEN_AI_REQUEST_TOP_P] = requestBody.p; + } + if (requestBody.prompt !== undefined) { + spanAttributes[GEN_AI_USAGE_INPUT_TOKENS] = Math.ceil(requestBody.prompt.length / 6); + } + if (requestBody.stop_sequences !== undefined) { + spanAttributes[ATTR_GEN_AI_REQUEST_STOP_SEQUENCES] = requestBody.stop_sequences; + } + } else if (modelId.includes('mistral')) { + if (requestBody.prompt !== undefined) { + spanAttributes[GEN_AI_USAGE_INPUT_TOKENS] = Math.ceil(requestBody.prompt.length / 6); + } + if (requestBody.max_tokens !== undefined) { + spanAttributes[GEN_AI_REQUEST_MAX_TOKENS] = requestBody.max_tokens; + } + if (requestBody.temperature !== undefined) { + spanAttributes[GEN_AI_REQUEST_TEMPERATURE] = requestBody.temperature; + } + if (requestBody.top_p !== undefined) { + spanAttributes[GEN_AI_REQUEST_TOP_P] = requestBody.top_p; + } + if (requestBody.stop !== undefined) { + spanAttributes[ATTR_GEN_AI_REQUEST_STOP_SEQUENCES] = requestBody.stop; + } + } + } + + return { + isStream, + spanAttributes, + }; + } + + private _responseHookConverse(response: NormalizedResponse, span: Span): void { + const { stopReason, usage } = response.data; + + setStopReason(span, stopReason); + setUsage(span, usage); + } + + private _responseHookConverseStream(response: NormalizedResponse, span: Span): void { + // Wrap and replace the response stream in place to process events into telemetry + // before yielding to the user. + response.data.stream = wrapConverseStreamResponse(response.data.stream, span); + } + + private _responseHookInvokeModel(response: NormalizedResponse, span: Span): void { + const currentModelId = response.request.commandInput?.modelId; + if (response.data?.body) { + const decodedResponseBody = new TextDecoder().decode(response.data.body); + const responseBody = JSON.parse(decodedResponseBody); + if (currentModelId.includes('amazon.titan')) { + if (responseBody.inputTextTokenCount !== undefined) { + span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, responseBody.inputTextTokenCount); + } + if (responseBody.results?.[0]?.tokenCount !== undefined) { + span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.results[0].tokenCount); + } + if (responseBody.results?.[0]?.completionReason !== undefined) { + span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.results[0].completionReason]); + } + } else if (currentModelId.includes('amazon.nova')) { + if (responseBody.usage !== undefined) { + if (responseBody.usage.inputTokens !== undefined) { + span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, responseBody.usage.inputTokens); + } + if (responseBody.usage.outputTokens !== undefined) { + span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.usage.outputTokens); + } + } + if (responseBody.stopReason !== undefined) { + 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); + } + if (responseBody.stop_reason !== undefined) { + span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.stop_reason]); + } + } else if (currentModelId.includes('meta.llama')) { + if (responseBody.prompt_token_count !== undefined) { + span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, responseBody.prompt_token_count); + } + if (responseBody.generation_token_count !== undefined) { + span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.generation_token_count); + } + if (responseBody.stop_reason !== undefined) { + span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.stop_reason]); + } + } else if (currentModelId.includes('cohere.command-r')) { + if (responseBody.text !== undefined) { + span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(responseBody.text.length / 6)); + } + if (responseBody.finish_reason !== undefined) { + span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.finish_reason]); + } + } else if (currentModelId.includes('cohere.command')) { + if (responseBody.generations?.[0]?.text !== undefined) { + span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(responseBody.generations[0].text.length / 6)); + } + if (responseBody.generations?.[0]?.finish_reason !== undefined) { + span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.generations[0].finish_reason]); + } + } else if (currentModelId.includes('mistral')) { + if (responseBody.outputs?.[0]?.text !== undefined) { + span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(responseBody.outputs[0].text.length / 6)); + } + if (responseBody.outputs?.[0]?.stop_reason !== undefined) { + span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [responseBody.outputs[0].stop_reason]); + } + } + } + } + + private _responseHookInvokeModelWithResponseStream(response: NormalizedResponse, span: Span): void { + const stream = response.data?.body; + const modelId = response.request.commandInput?.modelId; + if (!stream || !modelId) { + return; + } + + // 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. + response.data.body = (async function* () { + try { + for await (const chunk of stream) { + const parsedChunk = parseChunk(chunk?.chunk?.bytes); + + if (!parsedChunk) { + // pass through + } else if (modelId.includes('amazon.titan')) { + recordTitanAttributes(parsedChunk, span); + } else if (modelId.includes('anthropic.claude')) { + recordClaudeAttributes(parsedChunk, span); + } else if (modelId.includes('amazon.nova')) { + recordNovaAttributes(parsedChunk, span); + } else if (modelId.includes('meta.llama')) { + recordLlamaAttributes(parsedChunk, span); + } else if (modelId.includes('cohere.command-r')) { + recordCohereRAttributes(parsedChunk, span); + } else if (modelId.includes('cohere.command')) { + recordCohereAttributes(parsedChunk, span); + } else if (modelId.includes('mistral')) { + recordMistralAttributes(parsedChunk, span); + } + yield chunk; + } + } finally { + span.end(); + } + })(); + } +} + +async function* wrapConverseStreamResponse( + stream: AsyncIterable, + span: Span, +): AsyncGenerator { + try { + let usage: TokenUsage | undefined; + for await (const item of stream) { + setStopReason(span, item.messageStop?.stopReason); + usage = item.metadata?.usage; + yield item; + } + setUsage(span, usage); + } finally { + span.end(); + } +} + +function setStopReason(span: Span, stopReason: string | undefined): void { + if (stopReason !== undefined) { + span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [stopReason]); + } +} + +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); + } + } +} + +function parseChunk(bytes?: Uint8Array): any { + if (!bytes || !(bytes instanceof Uint8Array)) { + return null; + } + try { + const str = Buffer.from(bytes).toString('utf-8'); + return JSON.parse(str); + } catch (err) { + DEBUG_BUILD && debug.warn('[orchestrion:aws-sdk] failed to parse streamed bedrock chunk', err); + return null; + } +} + +function recordNovaAttributes(parsedChunk: any, span: Span): void { + if (parsedChunk.metadata?.usage !== undefined) { + if (parsedChunk.metadata?.usage.inputTokens !== undefined) { + span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, parsedChunk.metadata.usage.inputTokens); + } + if (parsedChunk.metadata?.usage.outputTokens !== undefined) { + span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, parsedChunk.metadata.usage.outputTokens); + } + } + if (parsedChunk.messageStop?.stopReason !== undefined) { + span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [parsedChunk.messageStop.stopReason]); + } +} + +function recordClaudeAttributes(parsedChunk: any, span: Span): void { + if (parsedChunk.message?.usage?.input_tokens !== undefined) { + span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, parsedChunk.message.usage.input_tokens); + } + if (parsedChunk.message?.usage?.output_tokens !== undefined) { + span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, parsedChunk.message.usage.output_tokens); + } + if (parsedChunk.delta?.stop_reason !== undefined) { + span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [parsedChunk.delta.stop_reason]); + } +} + +function recordTitanAttributes(parsedChunk: any, span: Span): void { + if (parsedChunk.inputTextTokenCount !== undefined) { + span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, parsedChunk.inputTextTokenCount); + } + if (parsedChunk.totalOutputTextTokenCount !== undefined) { + span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, parsedChunk.totalOutputTextTokenCount); + } + if (parsedChunk.completionReason !== undefined) { + span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [parsedChunk.completionReason]); + } +} + +function recordLlamaAttributes(parsedChunk: any, span: Span): void { + if (parsedChunk.prompt_token_count !== undefined) { + span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, parsedChunk.prompt_token_count); + } + if (parsedChunk.generation_token_count !== undefined) { + span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, parsedChunk.generation_token_count); + } + if (parsedChunk.stop_reason !== undefined) { + span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [parsedChunk.stop_reason]); + } +} + +function recordMistralAttributes(parsedChunk: any, span: Span): void { + if (parsedChunk.outputs?.[0]?.text !== undefined) { + span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(parsedChunk.outputs[0].text.length / 6)); + } + if (parsedChunk.outputs?.[0]?.stop_reason !== undefined) { + span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [parsedChunk.outputs[0].stop_reason]); + } +} + +function recordCohereAttributes(parsedChunk: any, span: Span): void { + if (parsedChunk.generations?.[0]?.text !== undefined) { + span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(parsedChunk.generations[0].text.length / 6)); + } + if (parsedChunk.generations?.[0]?.finish_reason !== undefined) { + span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [parsedChunk.generations[0].finish_reason]); + } +} + +function recordCohereRAttributes(parsedChunk: any, span: Span): void { + if (parsedChunk.text !== undefined) { + span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(parsedChunk.text.length / 6)); + } + if (parsedChunk.finish_reason !== undefined) { + span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [parsedChunk.finish_reason]); + } +} From 0b01d2d30f45ead50ec73bb91856008d0f3490d1 Mon Sep 17 00:00:00 2001 From: Andrei Borza Date: Fri, 10 Jul 2026 00:44:48 +0200 Subject: [PATCH 2/3] Update expected origin to auto.aws.orchestrion.aws_sdk --- .../suites/aws-serverless/bedrock/test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dev-packages/node-integration-tests/suites/aws-serverless/bedrock/test.ts b/dev-packages/node-integration-tests/suites/aws-serverless/bedrock/test.ts index e8e6e9c337e5..d04c2cf1d7fa 100644 --- a/dev-packages/node-integration-tests/suites/aws-serverless/bedrock/test.ts +++ b/dev-packages/node-integration-tests/suites/aws-serverless/bedrock/test.ts @@ -6,7 +6,7 @@ import { cleanupChildProcesses, createEsmAndCjsTests } from '../../../utils/runn // The suite runs twice on CI: once with the OTel `Aws` integration (default) and once with the // orchestrion diagnostics-channel integration auto-injected (`INJECT_ORCHESTRION`). Both emit the // same gen_ai spans; only the origin differs. -const ORIGIN = isOrchestrionEnabled() ? 'auto.aws.orchestrion.aws-sdk' : 'auto.otel.aws'; +const ORIGIN = isOrchestrionEnabled() ? 'auto.aws.orchestrion.aws_sdk' : 'auto.otel.aws'; const MODEL_ID = 'anthropic.claude-3-5-sonnet-20240620-v1:0'; From 187e0a3885edab32b03d591ee7e3e8d88f07bd6e Mon Sep 17 00:00:00 2001 From: Andrei Borza Date: Fri, 10 Jul 2026 09:40:15 +0200 Subject: [PATCH 3/3] Consolidate bedrock chunk any behind documented ParsedChunk alias --- .../aws-sdk/services/bedrock-runtime.ts | 23 +++++++++++-------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/bedrock-runtime.ts b/packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/bedrock-runtime.ts index 60de948a0fbc..1956a36d6688 100644 --- a/packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/bedrock-runtime.ts +++ b/packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/bedrock-runtime.ts @@ -31,9 +31,14 @@ interface TokenUsage { interface ConverseStreamOutput { messageStop?: { stopReason?: string }; metadata?: { usage?: TokenUsage }; - [key: string]: any; + [key: string]: unknown; } +// Streamed `InvokeModel` chunks and response bodies are model-family-specific JSON (titan, nova, +// claude, llama, cohere, mistral); the record helpers probe the shapes defensively, so `any` instead +// of one structural type per family. +type ParsedChunk = any; + export class BedrockRuntimeServiceExtension implements ServiceExtension { public requestPreSpanHook(request: NormalizedRequest): RequestMetadata { switch (request.commandName) { @@ -389,7 +394,7 @@ function setUsage(span: Span, usage: TokenUsage | undefined): void { } } -function parseChunk(bytes?: Uint8Array): any { +function parseChunk(bytes?: Uint8Array): ParsedChunk { if (!bytes || !(bytes instanceof Uint8Array)) { return null; } @@ -402,7 +407,7 @@ function parseChunk(bytes?: Uint8Array): any { } } -function recordNovaAttributes(parsedChunk: any, span: Span): void { +function recordNovaAttributes(parsedChunk: ParsedChunk, span: Span): void { if (parsedChunk.metadata?.usage !== undefined) { if (parsedChunk.metadata?.usage.inputTokens !== undefined) { span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, parsedChunk.metadata.usage.inputTokens); @@ -416,7 +421,7 @@ function recordNovaAttributes(parsedChunk: any, span: Span): void { } } -function recordClaudeAttributes(parsedChunk: any, 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); } @@ -428,7 +433,7 @@ function recordClaudeAttributes(parsedChunk: any, span: Span): void { } } -function recordTitanAttributes(parsedChunk: any, span: Span): void { +function recordTitanAttributes(parsedChunk: ParsedChunk, span: Span): void { if (parsedChunk.inputTextTokenCount !== undefined) { span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, parsedChunk.inputTextTokenCount); } @@ -440,7 +445,7 @@ function recordTitanAttributes(parsedChunk: any, span: Span): void { } } -function recordLlamaAttributes(parsedChunk: any, span: Span): void { +function recordLlamaAttributes(parsedChunk: ParsedChunk, span: Span): void { if (parsedChunk.prompt_token_count !== undefined) { span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, parsedChunk.prompt_token_count); } @@ -452,7 +457,7 @@ function recordLlamaAttributes(parsedChunk: any, span: Span): void { } } -function recordMistralAttributes(parsedChunk: any, span: Span): void { +function recordMistralAttributes(parsedChunk: ParsedChunk, span: Span): void { if (parsedChunk.outputs?.[0]?.text !== undefined) { span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(parsedChunk.outputs[0].text.length / 6)); } @@ -461,7 +466,7 @@ function recordMistralAttributes(parsedChunk: any, span: Span): void { } } -function recordCohereAttributes(parsedChunk: any, span: Span): void { +function recordCohereAttributes(parsedChunk: ParsedChunk, span: Span): void { if (parsedChunk.generations?.[0]?.text !== undefined) { span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(parsedChunk.generations[0].text.length / 6)); } @@ -470,7 +475,7 @@ function recordCohereAttributes(parsedChunk: any, span: Span): void { } } -function recordCohereRAttributes(parsedChunk: any, span: Span): void { +function recordCohereRAttributes(parsedChunk: ParsedChunk, span: Span): void { if (parsedChunk.text !== undefined) { span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(parsedChunk.text.length / 6)); }