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 5dfe2b0f55b6..295c0201301e 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,9 +258,14 @@ 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); + response.data.stream = wrapConverseStreamResponse(response.data.stream, span, response.request.abortSignal); } private _responseHookInvokeModel(response: NormalizedResponse, span: Span): void { @@ -338,13 +343,15 @@ 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) { + const abortSignal = response.request.abortSignal; + 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 +366,9 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { } yield chunk; } + } catch (error) { + setStreamErrorStatus(span, abortSignal); + throw error; } finally { span.end(); } @@ -380,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; @@ -389,11 +400,18 @@ async function* wrapConverseStreamResponse( yield item; } setUsage(span, usage); + } catch (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 3cbcc70e0db1..827bd1b60ff0 100644 --- a/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts +++ b/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts @@ -6,9 +6,12 @@ import { GEN_AI_REQUEST_TEMPERATURE, } from '@sentry/conventions/attributes'; import { GEN_AI_GENERATE_CONTENT } from '@sentry/conventions/op'; -import { describe, expect, it } from 'vitest'; +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 { RequestMetadata } from '../../src/integrations/aws-sdk/types'; +import type { NormalizedResponse, RequestMetadata } from '../../src/integrations/aws-sdk/types'; +import { getDefaultTestClientOptions, TestClient } from '../mocks/client'; const CLAUDE_MODEL_ID = 'anthropic.claude-3-5-sonnet-20240620-v1:0'; @@ -71,3 +74,85 @@ describe('BedrockRuntimeServiceExtension InvokeModel request', () => { }); }); }); + +const STREAM_COMMANDS = [ + { commandName: 'ConverseStream', streamField: 'stream' }, + { commandName: 'InvokeModelWithResponseStream', streamField: 'body' }, +]; + +function respond( + span: Span, + commandName: string, + commandInput: Record, + data: unknown, + abortSignal?: AbortSignal, +): void { + const response: NormalizedResponse = { + data, + request: { serviceName: 'BedrockRuntime', commandName, commandInput, abortSignal }, + 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.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 = { + [streamField]: (async function* () { + yield { messageStart: { role: 'assistant' } }; + throw error; + })(), + }; + + respond(span, commandName, { modelId: CLAUDE_MODEL_ID }, data, abortSignal); + + await expect(drain(data[streamField] as AsyncIterable)).rejects.toBe(error); + expect(spanToJSON(span).status).toBe(status); + 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(); + }); +});