diff --git a/dev-packages/node-integration-tests/package.json b/dev-packages/node-integration-tests/package.json index e801ce94c331..edb41d914cc7 100644 --- a/dev-packages/node-integration-tests/package.json +++ b/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", diff --git a/dev-packages/node-integration-tests/suites/aws-serverless/bedrock/scenario-stream.mjs b/dev-packages/node-integration-tests/suites/aws-serverless/bedrock/scenario-stream.mjs new file mode 100644 index 000000000000..aaa3ecb93ed0 --- /dev/null +++ b/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(); 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 de0d7878c826..24249bf2dcc8 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 @@ -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(); + }); + }); }); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 1e247feb36f4..7f73498a633c 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -720,6 +720,9 @@ importers: '@sentry/server-utils': specifier: 11.11.0 version: link:../../packages/server-utils + '@smithy/core': + specifier: 3.38.1 + version: 3.38.1 '@smithy/node-http-handler': specifier: 4.12.1 version: 4.12.1