From 4f98836dafa2da4d48f40c3a447593cf65935c03 Mon Sep 17 00:00:00 2001 From: Andrei Borza Date: Thu, 5 Nov 2026 14:04:38 +0100 Subject: [PATCH 1/2] fix(server-utils): End Bedrock stream spans with an error status on failure Co-Authored-By: Claude Opus 5.5 --- .../aws-sdk/services/bedrock-runtime.ts | 18 ++++- .../test/integrations/aws-sdk-bedrock.test.ts | 78 +++++++++++++++++++ 2 files changed, 93 insertions(+), 3 deletions(-) create mode 100644 packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts 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 98a97d408765..c0a3c89f8a4f 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 @@ -1,5 +1,5 @@ import type { Span } from '@sentry/core'; -import { debug } from '@sentry/core'; +import { debug, SPAN_STATUS_ERROR } from '@sentry/core'; import { GEN_AI_OPERATION_NAME, GEN_AI_PROVIDER_NAME, @@ -258,6 +258,11 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { } private _responseHookConverseStream(response: NormalizedResponse, span: Span): void { + if (!response.data?.stream) { + span.end(); + return; + } + // 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); @@ -338,13 +343,14 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { private _responseHookInvokeModelWithResponseStream(response: NormalizedResponse, span: Span): void { const stream = response.data?.body; const modelId = response.request.commandInput?.modelId; - if (!stream || !modelId) { + if (!stream) { + span.end(); return; } // 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 = resolveStreamRecorder(modelId); + const recordAttributes = typeof modelId === 'string' ? resolveStreamRecorder(modelId) : undefined; // 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. @@ -359,6 +365,9 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { } yield chunk; } + } catch (error) { + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + throw error; } finally { span.end(); } @@ -389,6 +398,9 @@ async function* wrapConverseStreamResponse( yield item; } setUsage(span, usage); + } catch (error) { + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + throw error; } finally { span.end(); } diff --git a/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts b/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts new file mode 100644 index 000000000000..947af0962b5e --- /dev/null +++ b/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts @@ -0,0 +1,78 @@ +import type { Span } from '@sentry/core'; +import { setCurrentClient, spanToJSON, startInactiveSpan } from '@sentry/core'; +import { afterEach, beforeEach, describe, expect, it } from 'vitest'; +import { BedrockRuntimeServiceExtension } from '../../src/integrations/aws-sdk/services/bedrock-runtime'; +import type { NormalizedResponse } from '../../src/integrations/aws-sdk/types'; +import { getDefaultTestClientOptions, TestClient } from '../mocks/client'; + +const CLAUDE_MODEL_ID = 'anthropic.claude-3-5-sonnet-20240620-v1:0'; + +const STREAM_COMMANDS = [ + { commandName: 'ConverseStream', streamField: 'stream' }, + { commandName: 'InvokeModelWithResponseStream', streamField: 'body' }, +]; + +function respond(span: Span, commandName: string, commandInput: Record, data: unknown): void { + const response: NormalizedResponse = { + data, + request: { serviceName: 'BedrockRuntime', commandName, commandInput }, + requestId: 'b0a5e5a2-2c2b-4f56-9d0e-8c7c2d3f4a10', + }; + new BedrockRuntimeServiceExtension().responseHook(response, span); +} + +async function drain(stream: AsyncIterable): Promise { + for await (const _ of stream) { + // consume + } +} + +describe('BedrockRuntimeServiceExtension stream spans', () => { + beforeEach(() => { + const client = new TestClient(getDefaultTestClientOptions({ tracesSampleRate: 1 })); + setCurrentClient(client); + client.init(); + }); + + afterEach(() => { + setCurrentClient(undefined as never); + }); + + it.each(STREAM_COMMANDS)( + 'marks a $commandName span as failed when the stream fails', + async ({ commandName, streamField }) => { + const span = startInactiveSpan({ name: 'bedrock' }); + const error = new Error('ThrottlingException: Too many tokens, please wait before trying again.'); + const data: Record = { + [streamField]: (async function* () { + yield { messageStart: { role: 'assistant' } }; + throw error; + })(), + }; + + respond(span, commandName, { modelId: CLAUDE_MODEL_ID }, data); + + await expect(drain(data[streamField] as AsyncIterable)).rejects.toBe(error); + expect(spanToJSON(span).status).toBe('error'); + expect(spanToJSON(span).end_timestamp).toBeDefined(); + }, + ); + + it.each(STREAM_COMMANDS)('ends a $commandName span when the response has no stream', ({ commandName }) => { + const span = startInactiveSpan({ name: 'bedrock' }); + + respond(span, commandName, { modelId: CLAUDE_MODEL_ID }, {}); + + expect(spanToJSON(span).end_timestamp).toBeDefined(); + }); + + it('ends an InvokeModelWithResponseStream span without a modelId when the stream is drained', async () => { + const span = startInactiveSpan({ name: 'bedrock' }); + const data = { body: (async function* () {})() }; + + respond(span, 'InvokeModelWithResponseStream', {}, data); + await drain(data.body); + + expect(spanToJSON(span).end_timestamp).toBeDefined(); + }); +}); From 130afe4ba15747d7c5b2cc5cd9c39bc31e7deb18 Mon Sep 17 00:00:00 2001 From: Andrei Borza Date: Thu, 5 Nov 2026 14:29:00 +0100 Subject: [PATCH 2/2] Record aborted Bedrock streams as cancelled --- .../src/integrations/aws-sdk/index.ts | 1 + .../aws-sdk/services/bedrock-runtime.ts | 12 ++++++--- .../src/integrations/aws-sdk/types.ts | 2 ++ .../test/integrations/aws-sdk-bedrock.test.ts | 26 ++++++++++++++----- 4 files changed, 31 insertions(+), 10 deletions(-) diff --git a/packages/server-utils/src/integrations/aws-sdk/index.ts b/packages/server-utils/src/integrations/aws-sdk/index.ts index 10951179dd04..e98c00a6ee6f 100644 --- a/packages/server-utils/src/integrations/aws-sdk/index.ts +++ b/packages/server-utils/src/integrations/aws-sdk/index.ts @@ -108,6 +108,7 @@ function instrumentAwsSdk(servicesExtensions: ServicesExtensions): void { command.input = {}; } const normalizedRequest = normalizeV3Request(serviceName, commandName, command.input, undefined); + normalizedRequest.abortSignal = (data.arguments[1] as { abortSignal?: AbortSignal } | undefined)?.abortSignal; const requestMetadata = servicesExtensions.requestPreSpanHook(normalizedRequest); // `rpc` matches what the exporter infers from `rpc.service` for the OTel aws-sdk spans; 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 c0a3c89f8a4f..f3c7194acb48 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 @@ -265,7 +265,7 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { // 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); + response.data.stream = wrapConverseStreamResponse(response.data.stream, span, response.request.abortSignal); } private _responseHookInvokeModel(response: NormalizedResponse, span: Span): void { @@ -343,6 +343,7 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { private _responseHookInvokeModelWithResponseStream(response: NormalizedResponse, span: Span): void { const stream = response.data?.body; const modelId = response.request.commandInput?.modelId; + const abortSignal = response.request.abortSignal; if (!stream) { span.end(); return; @@ -366,7 +367,7 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { yield chunk; } } catch (error) { - span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + setStreamErrorStatus(span, abortSignal); throw error; } finally { span.end(); @@ -389,6 +390,7 @@ function resolveStreamRecorder(modelId: string): ((parsedChunk: ParsedChunk, spa async function* wrapConverseStreamResponse( stream: AsyncIterable, span: Span, + abortSignal: AbortSignal | undefined, ): AsyncGenerator { try { let usage: TokenUsage | undefined; @@ -399,13 +401,17 @@ async function* wrapConverseStreamResponse( } setUsage(span, usage); } catch (error) { - span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + setStreamErrorStatus(span, abortSignal); throw error; } finally { span.end(); } } +function setStreamErrorStatus(span: Span, abortSignal: AbortSignal | undefined): void { + span.setStatus({ code: SPAN_STATUS_ERROR, message: abortSignal?.aborted ? 'cancelled' : 'internal_error' }); +} + function setStopReason(span: Span, stopReason: string | undefined): void { if (stopReason !== undefined) { span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [stopReason]); diff --git a/packages/server-utils/src/integrations/aws-sdk/types.ts b/packages/server-utils/src/integrations/aws-sdk/types.ts index da6e978cfc45..206b10242e50 100644 --- a/packages/server-utils/src/integrations/aws-sdk/types.ts +++ b/packages/server-utils/src/integrations/aws-sdk/types.ts @@ -11,6 +11,8 @@ export interface NormalizedRequest { commandName: string; commandInput: CommandInput; region?: string; + /** The `abortSignal` the caller passed to `send`, if any. */ + abortSignal?: AbortSignal; } export interface NormalizedResponse { 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 947af0962b5e..ad2cd19ab53d 100644 --- a/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts +++ b/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts @@ -12,10 +12,16 @@ const STREAM_COMMANDS = [ { commandName: 'InvokeModelWithResponseStream', streamField: 'body' }, ]; -function respond(span: Span, commandName: string, commandInput: Record, data: unknown): void { +function respond( + span: Span, + commandName: string, + commandInput: Record, + data: unknown, + abortSignal?: AbortSignal, +): void { const response: NormalizedResponse = { data, - request: { serviceName: 'BedrockRuntime', commandName, commandInput }, + request: { serviceName: 'BedrockRuntime', commandName, commandInput, abortSignal }, requestId: 'b0a5e5a2-2c2b-4f56-9d0e-8c7c2d3f4a10', }; new BedrockRuntimeServiceExtension().responseHook(response, span); @@ -38,9 +44,15 @@ describe('BedrockRuntimeServiceExtension stream spans', () => { setCurrentClient(undefined as never); }); - it.each(STREAM_COMMANDS)( - 'marks a $commandName span as failed when the stream fails', - async ({ commandName, streamField }) => { + it.each( + STREAM_COMMANDS.flatMap(command => [ + { ...command, when: 'fails', abortSignal: undefined, status: 'error' }, + // `cancelled` is sent as `ok` under span streaming. + { ...command, when: 'is aborted by the caller', abortSignal: AbortSignal.abort(), status: 'ok' }, + ]), + )( + 'sets status $status when a $commandName stream $when', + async ({ commandName, streamField, abortSignal, status }) => { const span = startInactiveSpan({ name: 'bedrock' }); const error = new Error('ThrottlingException: Too many tokens, please wait before trying again.'); const data: Record = { @@ -50,10 +62,10 @@ describe('BedrockRuntimeServiceExtension stream spans', () => { })(), }; - respond(span, commandName, { modelId: CLAUDE_MODEL_ID }, data); + respond(span, commandName, { modelId: CLAUDE_MODEL_ID }, data, abortSignal); await expect(drain(data[streamField] as AsyncIterable)).rejects.toBe(error); - expect(spanToJSON(span).status).toBe('error'); + expect(spanToJSON(span).status).toBe(status); expect(spanToJSON(span).end_timestamp).toBeDefined(); }, );