fix(server-utils): End Bedrock stream spans with an error status on failure

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;

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<ConverseStreamOutput>,
   span: Span,
+  abortSignal: AbortSignal | undefined,
 ): AsyncGenerator<ConverseStreamOutput> {
   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]);

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 {

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<string, unknown>,
+  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<unknown>): Promise<void> {
+  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<string, unknown> = {
+        [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<unknown>)).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();
+  });
+});