test(aws-serverless): Add Bedrock streaming integration scenario

dev-packages/node-integration-tests/package.json

@@ -68,6 +68,7 @@
     "@sentry/node-native": "11.11.0",
     "@sentry/server-runtime-injection": "11.11.0",
     "@sentry/server-utils": "11.11.0",
+    "@smithy/core": "3.38.1",
     "@smithy/node-http-handler": "4.12.1",
     "@types/express": "^4.17.21",
     "@types/node-schedule": "^2.1.8",

dev-packages/node-integration-tests/suites/aws-serverless/bedrock/scenario-stream.mjs

@@ -0,0 +1,200 @@
+import * as Sentry from '@sentry/aws-serverless';
+import {
+  BedrockRuntimeClient,
+  ConverseStreamCommand,
+  InvokeModelWithResponseStreamCommand,
+} 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 { EventStreamCodec } from '@smithy/core/event-streams';
+import nock from 'nock';
+import { Buffer } from 'node:buffer';
+
+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`;
+const modelId = 'anthropic.claude-3-5-sonnet-20240620-v1:0';
+
+// Bedrock streams responses in the binary AWS event stream format, not JSON. The SDK decodes and
+// checks each message (headers, payload, CRC32), so the mocked response must use the same format.
+// The SDK's own codec encodes the messages. It needs UTF-8 helpers to encode the header strings.
+// https://docs.aws.amazon.com/AmazonS3/latest/API/RESTSelectObjectAppendix.html
+const codec = new EventStreamCodec(
+  bytes => Buffer.from(bytes).toString('utf8'),
+  text => Buffer.from(text),
+);
+
+// The codec needs a type for each header value. All Bedrock headers are strings.
+function encodeMessage(headers, payload) {
+  return codec.encode({
+    headers: Object.fromEntries(Object.entries(headers).map(([name, value]) => [name, { type: 'string', value }])),
+    body: Buffer.from(JSON.stringify(payload)),
+  });
+}
+
+// The SDK reads `:event-type` to find the stream member (for example `messageStart` or `chunk`).
+function event(eventType, payload) {
+  return encodeMessage(
+    { ':message-type': 'event', ':event-type': eventType, ':content-type': 'application/json' },
+    payload,
+  );
+}
+
+// An exception message makes the SDK throw while the caller iterates the stream. This is how a
+// stream fails after the response starts.
+function exception(exceptionType, payload) {
+  return encodeMessage(
+    { ':message-type': 'exception', ':exception-type': exceptionType, ':content-type': 'application/json' },
+    payload,
+  );
+}
+
+// InvokeModelWithResponseStream wraps each model chunk as base64 JSON in a `chunk` event.
+function invokeModelChunk(chunk) {
+  return event('chunk', { bytes: Buffer.from(JSON.stringify(chunk)).toString('base64') });
+}
+
+function mockStream(path, messages) {
+  nock(host).post(path).reply(200, Buffer.concat(messages), { 'content-type': 'application/vnd.amazon.eventstream' });
+}
+
+function createClient() {
+  return new BedrockRuntimeClient({ region, credentials, maxAttempts: 1, requestHandler: new NodeHttpHandler() });
+}
+
+async function converseStream() {
+  mockStream(/\/model\/.*\/converse-stream$/, [
+    event('messageStart', { role: 'assistant' }),
+    event('contentBlockDelta', { contentBlockIndex: 0, delta: { text: 'Hello ' } }),
+    event('contentBlockDelta', { contentBlockIndex: 0, delta: { text: 'from Bedrock' } }),
+    event('contentBlockStop', { contentBlockIndex: 0 }),
+    event('contentBlockStart', {
+      contentBlockIndex: 1,
+      start: { toolUse: { toolUseId: 'tooluse_1', name: 'get_weather' } },
+    }),
+    event('contentBlockDelta', { contentBlockIndex: 1, delta: { toolUse: { input: '{"city":' } } }),
+    event('contentBlockDelta', { contentBlockIndex: 1, delta: { toolUse: { input: '"Vienna"}' } } }),
+    event('contentBlockStop', { contentBlockIndex: 1 }),
+    event('messageStop', { stopReason: 'tool_use' }),
+    event('metadata', { usage: { inputTokens: 12, outputTokens: 8, totalTokens: 20 }, metrics: { latencyMs: 100 } }),
+  ]);
+
+  const response = await createClient().send(
+    new ConverseStreamCommand({
+      modelId,
+      system: [{ text: 'You are a helpful assistant.' }],
+      messages: [{ role: 'user', content: [{ text: 'Hello' }] }],
+    }),
+  );
+
+  let text = '';
+  for await (const item of response.stream) {
+    text += item.contentBlockDelta?.delta?.text ?? '';
+  }
+  if (text !== 'Hello from Bedrock') {
+    throw new Error(`Unexpected streamed text: ${text}`);
+  }
+}
+
+async function invokeModelWithResponseStream() {
+  mockStream(/\/model\/.*\/invoke-with-response-stream$/, [
+    invokeModelChunk({
+      type: 'message_start',
+      message: { id: 'msg_bdrk_2', model: 'claude-3-5-sonnet-20240620', usage: { input_tokens: 15, output_tokens: 1 } },
+    }),
+    invokeModelChunk({ type: 'content_block_start', index: 0, content_block: { type: 'text', text: '' } }),
+    invokeModelChunk({ type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: 'Streamed ' } }),
+    invokeModelChunk({ type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: 'reply' } }),
+    invokeModelChunk({ type: 'content_block_stop', index: 0 }),
+    invokeModelChunk({ type: 'message_delta', delta: { stop_reason: 'end_turn' }, usage: { output_tokens: 9 } }),
+    invokeModelChunk({ type: 'message_stop' }),
+  ]);
+
+  const response = await createClient().send(
+    new InvokeModelWithResponseStreamCommand({
+      modelId,
+      contentType: 'application/json',
+      body: JSON.stringify({
+        anthropic_version: 'bedrock-2023-05-31',
+        max_tokens: 100,
+        messages: [{ role: 'user', content: 'Stream please' }],
+      }),
+    }),
+  );
+
+  let chunks = 0;
+  for await (const item of response.body) {
+    if (item.chunk?.bytes) {
+      chunks++;
+    }
+  }
+  if (chunks !== 7) {
+    throw new Error(`Unexpected chunk count: ${chunks}`);
+  }
+}
+
+async function failingConverseStream() {
+  mockStream(/\/model\/.*\/converse-stream$/, [
+    event('messageStart', { role: 'assistant' }),
+    event('contentBlockDelta', { contentBlockIndex: 0, delta: { text: 'Partial' } }),
+    exception('modelStreamErrorException', { message: 'The model stream failed', originalStatusCode: 500 }),
+  ]);
+
+  const response = await createClient().send(
+    new ConverseStreamCommand({
+      modelId: 'amazon.nova-pro-v1:0',
+      messages: [{ role: 'user', content: [{ text: 'Fail please' }] }],
+    }),
+  );
+
+  try {
+    for await (const _ of response.stream) {
+      // drain
+    }
+  } catch (error) {
+    if (error.name !== 'ModelStreamErrorException') {
+      throw error;
+    }
+    return;
+  }
+  throw new Error('Expected the stream to fail');
+}
+
+async function abortedConverseStream() {
+  mockStream(/\/model\/.*\/converse-stream$/, [
+    event('messageStart', { role: 'assistant' }),
+    event('contentBlockDelta', { contentBlockIndex: 0, delta: { text: 'Partial' } }),
+    exception('modelStreamErrorException', { message: 'The model stream failed', originalStatusCode: 500 }),
+  ]);
+
+  const controller = new AbortController();
+  const response = await createClient().send(
+    new ConverseStreamCommand({
+      modelId: 'amazon.nova-lite-v1:0',
+      messages: [{ role: 'user', content: [{ text: 'Stop please' }] }],
+    }),
+    { abortSignal: controller.signal },
+  );
+
+  try {
+    for await (const _ of response.stream) {
+      controller.abort();
+    }
+  } catch {
+    return;
+  }
+  throw new Error('Expected the stream to fail');
+}
+
+async function run() {
+  await Sentry.startSpan({ name: 'Test Transaction' }, async () => {
+    await converseStream();
+    await invokeModelWithResponseStream();
+    await failingConverseStream();
+    await abortedConverseStream();
+  });
+}
+
+run();

dev-packages/node-integration-tests/suites/aws-serverless/bedrock/test.ts

@@ -11,6 +11,7 @@ import {
   SENTRY_OP,
   SENTRY_ORIGIN,
 } from '@sentry/conventions/attributes';
+import { GEN_AI_CHAT, GEN_AI_GENERATE_CONTENT } from '@sentry/conventions/op';
 import type { SerializedStreamedSpanContainer } from '@sentry/core';
 import { afterAll, describe, expect } from 'vitest';
 import { cleanupChildProcesses, createEsmAndCjsTests } from '../../../utils/runner';
@@ -52,6 +53,29 @@ function assertBedrockSpans(container: SerializedStreamedSpanContainer): void {
   expect(invokeModelSpan!.attributes[GEN_AI_RESPONSE_FINISH_REASONS].value).toBe('["end_turn"]');
 }
 
+function assertStreamSpans(container: SerializedStreamedSpanContainer): void {
+  const converseStreamSpan = container.items.find(span => span.name === `chat ${MODEL_ID}`);
+  expect(converseStreamSpan!.status).toBe('ok');
+  expect(converseStreamSpan!.attributes[SENTRY_OP].value).toBe(GEN_AI_CHAT);
+  expect(converseStreamSpan!.attributes[GEN_AI_USAGE_INPUT_TOKENS].value).toBe(12);
+  expect(converseStreamSpan!.attributes[GEN_AI_USAGE_OUTPUT_TOKENS].value).toBe(8);
+  expect(converseStreamSpan!.attributes[GEN_AI_RESPONSE_FINISH_REASONS].value).toBe('["tool_use"]');
+
+  const invokeStreamSpan = container.items.find(span => span.name === `generate_content ${MODEL_ID}`);
+  expect(invokeStreamSpan!.status).toBe('ok');
+  expect(invokeStreamSpan!.attributes[SENTRY_OP].value).toBe(GEN_AI_GENERATE_CONTENT);
+  expect(invokeStreamSpan!.attributes[GEN_AI_USAGE_INPUT_TOKENS].value).toBe(15);
+  expect(invokeStreamSpan!.attributes[GEN_AI_USAGE_OUTPUT_TOKENS].value).toBe(9);
+  expect(invokeStreamSpan!.attributes[GEN_AI_RESPONSE_FINISH_REASONS].value).toBe('["end_turn"]');
+
+  const failedStreamSpan = container.items.find(span => span.name === 'chat amazon.nova-pro-v1:0');
+  expect(failedStreamSpan!.status).toBe('error');
+
+  // `cancelled` is sent as `ok` under span streaming.
+  const abortedStreamSpan = container.items.find(span => span.name === 'chat amazon.nova-lite-v1:0');
+  expect(abortedStreamSpan!.status).toBe('ok');
+}
+
 describe('awsIntegration - Bedrock', () => {
   afterAll(() => {
     cleanupChildProcesses();
@@ -62,4 +86,10 @@ describe('awsIntegration - Bedrock', () => {
       await createTestRunner().ignore('event').expect({ span: assertBedrockSpans }).start().completed();
     });
   });
+
+  createEsmAndCjsTests(__dirname, 'scenario-stream.mjs', 'instrument.mjs', (createTestRunner, test) => {
+    test('ends Bedrock stream spans when the streams finish, fail, or are aborted', { timeout: 90_000 }, async () => {
+      await createTestRunner().ignore('event').expect({ span: assertStreamSpans }).start().completed();
+    });
+  });
 });