From 7c18481c742e943aabfcf7e6aad8ff4901928960 Mon Sep 17 00:00:00 2001 From: RulaKhaled Date: Fri, 18 Sep 2026 09:22:29 +0300 Subject: [PATCH 1/5] fix(server-utils): End Anthropic stream spans drained through the raw Response Hook `Stream.fromSSEResponse` and wrap the `Response` body it is built from, so the `gen_ai.chat` span ends however the caller drains the stream. Every consumption path bottoms out in `response.body`, but only the SDK `Stream`'s async iterator was instrumented. A caller who takes `.asResponse()` or `.withResponse()` and reads the body themselves never touches that iterator, so nothing ever ended the span and it was dropped. `tee()` was lost the same way, since it calls `this.iterator()` directly. Fixes #24258 Co-Authored-By: Claude Opus 5 --- .../tracing/anthropic/instrument-raw-body.mjs | 10 ++ .../anthropic/scenario-stream-raw-body.mjs | 92 +++++++++++ .../suites/tracing/anthropic/test.ts | 29 ++++ .../src/ai/anthropic-ai/streaming.ts | 146 +++++++++++++++--- .../src/integrations/anthropic.ts | 58 ++++++- .../src/orchestrion/config/anthropic-ai.ts | 14 ++ 6 files changed, 322 insertions(+), 27 deletions(-) create mode 100644 dev-packages/node-integration-tests/suites/tracing/anthropic/instrument-raw-body.mjs create mode 100644 dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-stream-raw-body.mjs diff --git a/dev-packages/node-integration-tests/suites/tracing/anthropic/instrument-raw-body.mjs b/dev-packages/node-integration-tests/suites/tracing/anthropic/instrument-raw-body.mjs new file mode 100644 index 000000000000..1c507ba84b9c --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/anthropic/instrument-raw-body.mjs @@ -0,0 +1,10 @@ +import * as Sentry from '@sentry/node'; +import { loggingTransport } from '@sentry-internal/node-integration-tests'; + +Sentry.init({ + dsn: 'https://public@dsn.ingest.sentry.io/1337', + release: '1.0', + tracesSampleRate: 1.0, + dataCollection: { genAI: { inputs: true, outputs: true } }, + transport: loggingTransport, +}); diff --git a/dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-stream-raw-body.mjs b/dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-stream-raw-body.mjs new file mode 100644 index 000000000000..bac244368d87 --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-stream-raw-body.mjs @@ -0,0 +1,92 @@ +import Anthropic from '@anthropic-ai/sdk'; +import * as Sentry from '@sentry/node'; +import express from 'express'; + +function startMockAnthropicServer() { + const app = express(); + app.use(express.json()); + + app.post('/anthropic/v1/messages', (req, res) => { + res.writeHead(200, { + 'Content-Type': 'text/event-stream', + 'Cache-Control': 'no-cache', + Connection: 'keep-alive', + }); + + const model = req.body.model; + const events = [ + { + type: 'message_start', + message: { + id: 'msg_raw_body', + type: 'message', + role: 'assistant', + model, + content: [], + usage: { input_tokens: 10 }, + }, + }, + { type: 'content_block_start', index: 0, content_block: { type: 'text', text: '' } }, + { type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: 'Raw ' } }, + { type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: 'body!' } }, + { type: 'content_block_stop', index: 0 }, + { + type: 'message_delta', + delta: { stop_reason: 'end_turn', stop_sequence: null }, + usage: { output_tokens: 15 }, + }, + { type: 'message_stop' }, + ]; + + events.forEach((event, index) => { + setTimeout(() => { + res.write(`event: ${event.type}\n`); + res.write(`data: ${JSON.stringify(event)}\n\n`); + if (index === events.length - 1) { + res.end(); + } + }, index * 10); + }); + }); + + return new Promise(resolve => { + const server = app.listen(0, () => { + resolve(server); + }); + }); +} + +async function run() { + const server = await startMockAnthropicServer(); + + await Sentry.startSpan({ op: 'function', name: 'main' }, async () => { + const client = new Anthropic({ + apiKey: 'mock-api-key', + baseURL: `http://localhost:${server.address().port}/anthropic`, + }); + + const params = { + model: 'claude-3-haiku-20240307', + messages: [{ role: 'user', content: 'Stream this please' }], + stream: true, + }; + + // 1) Drain the raw `Response` body, never touching the SDK `Stream` + const response = await client.messages.create({ ...params }).asResponse(); + for await (const _ of response.body) { + void _; + } + + // 2) Drain the SDK `Stream`, so both consumption styles are covered in one run + const stream = await client.messages.create({ ...params }); + for await (const _ of stream) { + void _; + } + }); + + await Sentry.flush(2000); + + server.close(); +} + +run(); diff --git a/dev-packages/node-integration-tests/suites/tracing/anthropic/test.ts b/dev-packages/node-integration-tests/suites/tracing/anthropic/test.ts index ee2bb16b37dc..7971efaaa731 100644 --- a/dev-packages/node-integration-tests/suites/tracing/anthropic/test.ts +++ b/dev-packages/node-integration-tests/suites/tracing/anthropic/test.ts @@ -265,6 +265,35 @@ describe('Anthropic integration', () => { }); }); + createEsmAndCjsTests(__dirname, 'scenario-stream-raw-body.mjs', 'instrument-raw-body.mjs', (createRunner, test) => { + test('ends the span when a stream is drained through the raw Response body', async () => { + await createRunner() + .unordered() + .expect({ + span: container => { + const genAiSpans = container.items.filter(span => span.attributes['sentry.op']?.value === 'gen_ai.chat'); + // Two calls, one drained via `.asResponse()` and one via the SDK `Stream`. Both must end, + // and both must carry the response attributes accumulated off the SSE frames. + expect(genAiSpans).toHaveLength(2); + for (const span of genAiSpans) { + expect(span.name).toBe('chat claude-3-haiku-20240307'); + expect(span.status).toBe('ok'); + expect(span.attributes[GEN_AI_RESPONSE_STREAMING].value).toBe(true); + expect(span.attributes[GEN_AI_RESPONSE_ID].value).toBe('msg_raw_body'); + expect(span.attributes[GEN_AI_RESPONSE_MODEL].value).toBe('claude-3-haiku-20240307'); + expect(span.attributes[GEN_AI_RESPONSE_FINISH_REASONS].value).toBe('["end_turn"]'); + expect(span.attributes[GEN_AI_RESPONSE_TEXT].value).toBe('Raw body!'); + expect(span.attributes[GEN_AI_USAGE_INPUT_TOKENS].value).toBe(10); + expect(span.attributes[GEN_AI_USAGE_OUTPUT_TOKENS].value).toBe(15); + expect(span.attributes[GEN_AI_USAGE_TOTAL_TOKENS].value).toBe(25); + } + }, + }) + .start() + .completed(); + }); + }); + createEsmAndCjsTests(__dirname, 'scenario-stream.mjs', 'instrument-with-pii.mjs', (createRunner, test) => { test('streams record response text when PII true', async () => { await createRunner() diff --git a/packages/server-utils/src/ai/anthropic-ai/streaming.ts b/packages/server-utils/src/ai/anthropic-ai/streaming.ts index 3daf61c2b57e..971b21f7fb02 100644 --- a/packages/server-utils/src/ai/anthropic-ai/streaming.ts +++ b/packages/server-utils/src/ai/anthropic-ai/streaming.ts @@ -37,6 +37,21 @@ interface StreamingState { >; } +function createStreamingState(): StreamingState { + return { + responseTexts: [], + finishReasons: [], + responseId: '', + responseModel: '', + promptTokens: undefined, + completionTokens: undefined, + cacheCreationInputTokens: undefined, + cacheReadInputTokens: undefined, + toolCalls: [], + activeToolBlocks: {}, + }; +} + /** * Checks if an event is an error event * @param event - The event to process @@ -205,18 +220,7 @@ export async function* instrumentAsyncIterableStream( span: Span, recordOutputs: boolean, ): AsyncGenerator { - const state: StreamingState = { - responseTexts: [], - finishReasons: [], - responseId: '', - responseModel: '', - promptTokens: undefined, - completionTokens: undefined, - cacheCreationInputTokens: undefined, - cacheReadInputTokens: undefined, - toolCalls: [], - activeToolBlocks: {}, - }; + const state = createStreamingState(); try { for await (const event of stream) { @@ -236,18 +240,7 @@ export function instrumentMessageStream span: Span, recordOutputs: boolean, ): R { - const state: StreamingState = { - responseTexts: [], - finishReasons: [], - responseId: '', - responseModel: '', - promptTokens: undefined, - completionTokens: undefined, - cacheCreationInputTokens: undefined, - cacheReadInputTokens: undefined, - toolCalls: [], - activeToolBlocks: {}, - }; + const state = createStreamingState(); stream.on('streamEvent', (event: unknown) => { processEvent(event as AnthropicAiStreamingEvent, state, recordOutputs, span); @@ -277,3 +270,108 @@ export function instrumentMessageStream return stream; } + +/** Handle returned by {@link instrumentRawSseBody} for the `Stream`-iterator path to claim the span. */ +export interface RawSseBodyHandle { + /** Called when the SDK `Stream`'s async iterator takes over, so the body wrapper stays a pass-through. */ + claim: () => void; +} + +/** + * Replace `response.body` with a pass-through that accumulates the SSE frames flowing through it and + * ends `span` when the body is exhausted, cancelled or errors. + * + * Every way of draining an Anthropic stream bottoms out in `response.body`: the SDK `Stream`'s async + * iterator, `tee()`, and a caller reading `.asResponse()`/`.withResponse()`'s raw `Response`. Only the + * first of those is visible to {@link instrumentAsyncIterableStream}, so without this the other two end + * no span at all. When the iterator path does run it claims the span and this wrapper goes quiet, so a + * chunk is never accounted for twice. + * + * Returns `undefined` — leaving the response untouched — for a body we can't wrap. + */ +export function instrumentRawSseBody( + response: { body?: unknown }, + span: Span, + recordOutputs: boolean, +): RawSseBodyHandle | undefined { + const body = response.body as ReadableStream | null | undefined; + if (!body || typeof body.getReader !== 'function') { + return undefined; + } + + const state = createStreamingState(); + const decoder = new TextDecoder(); + let buffered = ''; + let claimed = false; + let settled = false; + + // Never lets an accumulation failure reach the caller: their stream matters more than our attributes. + const consume = (chunk: Uint8Array): void => { + try { + buffered += decoder.decode(chunk, { stream: true }); + + let newline = buffered.indexOf('\n'); + while (newline !== -1) { + const line = buffered.slice(0, newline).trim(); + buffered = buffered.slice(newline + 1); + // An SSE frame's `event:` line only repeats the `type` already carried by the JSON payload. + if (line.startsWith('data:')) { + processEvent(JSON.parse(line.slice(5)) as AnthropicAiStreamingEvent, state, recordOutputs, span); + } + newline = buffered.indexOf('\n'); + } + } catch { + // A frame we can't decode or parse is not worth breaking the caller's stream over. + } + }; + + const settle = (error?: unknown): void => { + if (settled || claimed) { + return; + } + settled = true; + if (error !== undefined) { + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + } + endStreamSpan(span, state, recordOutputs); + }; + + const reader = body.getReader(); + const instrumented = new ReadableStream({ + async pull(controller) { + try { + const { done, value } = await reader.read(); + if (done) { + settle(); + controller.close(); + return; + } + if (!claimed) { + consume(value); + } + controller.enqueue(value); + } catch (error) { + settle(error); + controller.error(error); + } + }, + async cancel(reason) { + settle(); + await reader.cancel(reason); + }, + }); + + try { + // `body` is a prototype getter, so an own data property shadows it for every later read. + Object.defineProperty(response, 'body', { value: instrumented, configurable: true }); + } catch { + reader.releaseLock(); + return undefined; + } + + return { + claim: () => { + claimed = true; + }, + }; +} diff --git a/packages/server-utils/src/integrations/anthropic.ts b/packages/server-utils/src/integrations/anthropic.ts index 87bce3b9c199..dfa072ae9db6 100644 --- a/packages/server-utils/src/integrations/anthropic.ts +++ b/packages/server-utils/src/integrations/anthropic.ts @@ -4,14 +4,21 @@ import type { IntegrationFn, Span, SpanAttributeValue } from '@sentry/core'; import { _INTERNAL_shouldSkipAiProviderWrapping, defineIntegration, + getActiveSpan, getClient, hasSpanStreamingEnabled, + isObjectLike, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, startInactiveSpan, } from '@sentry/core'; import { getGenAiSpanOp, resolveAIRecordingOptions } from '../ai/core/utils'; import { addPrivateRequestAttributes, addResponseAttributes, extractRequestAttributes } from '../ai/anthropic-ai'; -import { instrumentAsyncIterableStream, instrumentMessageStream } from '../ai/anthropic-ai/streaming'; +import type { RawSseBodyHandle } from '../ai/anthropic-ai/streaming'; +import { + instrumentAsyncIterableStream, + instrumentMessageStream, + instrumentRawSseBody, +} from '../ai/anthropic-ai/streaming'; import type { AnthropicAiOptions, AnthropicAiResponse } from '../ai/anthropic-ai/types'; import { CHANNELS } from '../orchestrion/channels'; import { bindTracingChannelToSpan } from '../tracing-channel'; @@ -41,6 +48,14 @@ interface AnthropicChannelContext { result?: unknown; } +// Spans opened for `messages.create({ stream: true })`, i.e. the ones whose stream is drained through +// the SDK's `Stream`. `messages.stream()` spans are excluded: `instrumentMessageStream` already owns +// when those end, so the raw-body wrapper must keep its hands off them. +const asyncIterableStreamSpans = new WeakSet(); + +// The `Stream` a raw-body wrapper was installed for, so the iterator path can claim the span. +const rawSseBodyHandles = new WeakMap(); + const _anthropicAIIntegration = ((options: AnthropicAiOptions = {}) => { return { name: INTEGRATION_NAME, @@ -54,7 +69,7 @@ function instrumentAnthropic(options: AnthropicAiOptions): void { for (const { channel, operation, stream } of INSTRUMENTED_CHANNELS) { bindTracingChannelToSpan( diagnosticsChannel.tracingChannel(channel), - data => createGenAiSpan(data, operation, options), + data => createGenAiSpan(data, operation, stream, options), { beforeSpanEnd: (span, data) => { addResponseAttributes( @@ -67,6 +82,34 @@ function instrumentAnthropic(options: AnthropicAiOptions): void { }, ); } + + subscribeToSseStream(options); +} + +/** + * Wrap the raw `Response` behind every SSE stream so the span ends no matter how the caller drains it. + * + * `Stream.fromSSEResponse` runs inside the traced `messages.create` call, so the active span here is + * the `gen_ai` span that call opened — that is what links a response to its span, since nothing on the + * `Stream` the SDK hands back points at the `Response` it was built from. + */ +function subscribeToSseStream(options: AnthropicAiOptions): void { + const { recordOutputs } = resolveAIRecordingOptions(options); + + diagnosticsChannel.tracingChannel(CHANNELS.ANTHROPIC_SSE_STREAM).end.subscribe(message => { + const data = message as AnthropicChannelContext; + const span = getActiveSpan(); + const stream = data.result; + const response = data.arguments?.[0]; + if (!span || !asyncIterableStreamSpans.has(span) || !isObjectLike(stream) || !isObjectLike(response)) { + return; + } + + const handle = instrumentRawSseBody(response, span, recordOutputs); + if (handle) { + rawSseBodyHandles.set(stream, handle); + } + }); } /** @@ -76,6 +119,7 @@ function instrumentAnthropic(options: AnthropicAiOptions): void { function createGenAiSpan( data: AnthropicChannelContext, operation: string, + stream: StreamMode, options: AnthropicAiOptions, ): Span | undefined { const args = data.arguments ?? []; @@ -114,6 +158,10 @@ function createGenAiSpan( addPrivateRequestAttributes(span, params); } + if (stream === 'async-iterable') { + asyncIterableStreamSpans.add(span); + } + return span; } @@ -146,9 +194,13 @@ function wrapStreamResult( const result = data.result; if (stream === 'async-iterable' && isAsyncIterable(result)) { + const handle = rawSseBodyHandles.get(result); const iterate = result[Symbol.asyncIterator].bind(result); const instrumented = instrumentAsyncIterableStream({ [Symbol.asyncIterator]: iterate }, span, recordOutputs); - result[Symbol.asyncIterator] = () => instrumented; + result[Symbol.asyncIterator] = () => { + handle?.claim(); + return instrumented; + }; return true; } diff --git a/packages/server-utils/src/orchestrion/config/anthropic-ai.ts b/packages/server-utils/src/orchestrion/config/anthropic-ai.ts index 40cbdc515072..0e696b64ab45 100644 --- a/packages/server-utils/src/orchestrion/config/anthropic-ai.ts +++ b/packages/server-utils/src/orchestrion/config/anthropic-ai.ts @@ -26,6 +26,19 @@ export const anthropicAiConfig = [ module: { name: '@anthropic-ai/sdk', versionRange: '>=0.19.2 <1', filePath }, functionQuery: { className: 'Messages', methodName: 'stream', kind: 'Sync' as const }, })), + // `Stream.fromSSEResponse` is the one place the SDK's `Stream` and the raw `Response` meet. Hooking + // it lets the span survive the consumption paths the `Stream`'s async iterator never sees — chiefly + // `.asResponse()`/`.withResponse()`, where the caller drains `response.body` itself. + { + channelName: 'sse-stream', + module: { + name: '@anthropic-ai/sdk', + versionRange: '>=0.19.2 <1', + // `streaming.js` moved under `core/` in 0.60. + filePath: /^(?:core\/)?streaming\.(?:js|mjs)$/, + }, + functionQuery: { className: 'Stream', methodName: 'fromSSEResponse', kind: 'Sync' as const }, + }, ] satisfies InstrumentationConfig[]; export const anthropicAiModuleNames = getModuleNames(anthropicAiConfig); @@ -33,4 +46,5 @@ export const anthropicAiModuleNames = getModuleNames(anthropicAiConfig); export const anthropicAiChannels = { ANTHROPIC_CHAT: 'orchestrion:@anthropic-ai/sdk:chat', ANTHROPIC_MESSAGES_STREAM: 'orchestrion:@anthropic-ai/sdk:messages-stream', + ANTHROPIC_SSE_STREAM: 'orchestrion:@anthropic-ai/sdk:sse-stream', } as const; From 8127927f11126f4c8651af76616cf9c2811ec7be Mon Sep 17 00:00:00 2001 From: RulaKhaled Date: Fri, 18 Sep 2026 09:34:43 +0300 Subject: [PATCH 2/5] fix(server-utils): Don't disturb the Anthropic response body when wrapping it The wrapper acquired a reader as soon as it was installed, and its stream defaulted to a high-water mark of 1, so it read a chunk ahead before anyone asked for one. Both disturb the body, which made `text()`, `arrayBuffer()` and `clone()` throw on a `.asResponse()` result the caller had not read yet. Co-Authored-By: Claude Opus 5 --- .../anthropic/scenario-stream-raw-body.mjs | 7 +++ .../src/ai/anthropic-ai/streaming.ts | 54 +++++++++++-------- 2 files changed, 38 insertions(+), 23 deletions(-) diff --git a/dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-stream-raw-body.mjs b/dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-stream-raw-body.mjs index bac244368d87..8085803b91a9 100644 --- a/dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-stream-raw-body.mjs +++ b/dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-stream-raw-body.mjs @@ -73,6 +73,13 @@ async function run() { // 1) Drain the raw `Response` body, never touching the SDK `Stream` const response = await client.messages.create({ ...params }).asResponse(); + + // Wrapping the body must not disturb it, or `text()`, `arrayBuffer()` and `clone()` would throw + // on a response the caller has not read yet. + if (response.bodyUsed) { + throw new Error('raw Response body was consumed before the caller read it'); + } + for await (const _ of response.body) { void _; } diff --git a/packages/server-utils/src/ai/anthropic-ai/streaming.ts b/packages/server-utils/src/ai/anthropic-ai/streaming.ts index 971b21f7fb02..ded21326aa95 100644 --- a/packages/server-utils/src/ai/anthropic-ai/streaming.ts +++ b/packages/server-utils/src/ai/anthropic-ai/streaming.ts @@ -336,36 +336,44 @@ export function instrumentRawSseBody( endStreamSpan(span, state, recordOutputs); }; - const reader = body.getReader(); - const instrumented = new ReadableStream({ - async pull(controller) { - try { - const { done, value } = await reader.read(); - if (done) { - settle(); - controller.close(); - return; - } - if (!claimed) { - consume(value); + // Acquired on the first read, never at wrap time: taking a reader disturbs the body, which would + // make `response.text()`, `arrayBuffer()` and `clone()` throw on a response nobody has read yet. + let reader: ReadableStreamDefaultReader | undefined; + + const instrumented = new ReadableStream( + { + async pull(controller) { + try { + reader ??= body.getReader(); + const { done, value } = await reader.read(); + if (done) { + settle(); + controller.close(); + return; + } + if (!claimed) { + consume(value); + } + controller.enqueue(value); + } catch (error) { + settle(error); + controller.error(error); } - controller.enqueue(value); - } catch (error) { - settle(error); - controller.error(error); - } - }, - async cancel(reason) { - settle(); - await reader.cancel(reason); + }, + async cancel(reason) { + settle(); + await (reader ? reader.cancel(reason) : body.cancel(reason)); + }, }, - }); + // A high-water mark of 0 keeps the stream from pulling a chunk before anyone asks for one. The + // default of 1 would read ahead the moment we wrap, disturbing a body the caller may never read. + { highWaterMark: 0 }, + ); try { // `body` is a prototype getter, so an own data property shadows it for every later read. Object.defineProperty(response, 'body', { value: instrumented, configurable: true }); } catch { - reader.releaseLock(); return undefined; } From b288e3d319d6e77250e4850bd28912942a42b32e Mon Sep 17 00:00:00 2001 From: RulaKhaled Date: Fri, 18 Sep 2026 09:35:47 +0300 Subject: [PATCH 3/5] chore(server-utils): Correct the Anthropic streaming.js path comment Co-Authored-By: Claude Opus 5 --- packages/server-utils/src/orchestrion/config/anthropic-ai.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/server-utils/src/orchestrion/config/anthropic-ai.ts b/packages/server-utils/src/orchestrion/config/anthropic-ai.ts index 0e696b64ab45..e6e1d073c6a1 100644 --- a/packages/server-utils/src/orchestrion/config/anthropic-ai.ts +++ b/packages/server-utils/src/orchestrion/config/anthropic-ai.ts @@ -34,7 +34,8 @@ export const anthropicAiConfig = [ module: { name: '@anthropic-ai/sdk', versionRange: '>=0.19.2 <1', - // `streaming.js` moved under `core/` in 0.60. + // `class Stream` sits at the package root up to 0.5x and under `core/` from 0.59 on, where the + // root file is left behind as a re-export shim that matches nothing. filePath: /^(?:core\/)?streaming\.(?:js|mjs)$/, }, functionQuery: { className: 'Stream', methodName: 'fromSSEResponse', kind: 'Sync' as const }, From 50af097badb2cd4c8ce5eed06f4f751c7667939c Mon Sep 17 00:00:00 2001 From: RulaKhaled Date: Fri, 18 Sep 2026 09:57:34 +0300 Subject: [PATCH 4/5] ref(server-utils): Let the SSE body wrapper own Anthropic stream spans The body wrapper sees every chunk whichever way the caller drains the stream, so keeping the async-iterator instrumentation alongside it bought nothing and cost a claim handshake to stop the two from double-counting. Drop it from the channel path: one accumulator, one owner, and `wrapStreamResult` becomes a membership test. The manual `instrumentAnthropicAiClient` path still uses the iterator. A streaming call whose body we couldn't wrap now ends at `asyncEnd` with request attributes rather than hanging, so the degraded case stays graceful. Co-Authored-By: Claude Opus 5 --- .../src/ai/anthropic-ai/streaming.ts | 48 +++++----------- .../src/integrations/anthropic.ts | 57 ++++++------------- 2 files changed, 31 insertions(+), 74 deletions(-) diff --git a/packages/server-utils/src/ai/anthropic-ai/streaming.ts b/packages/server-utils/src/ai/anthropic-ai/streaming.ts index ded21326aa95..2f14f24e4af6 100644 --- a/packages/server-utils/src/ai/anthropic-ai/streaming.ts +++ b/packages/server-utils/src/ai/anthropic-ai/streaming.ts @@ -271,38 +271,25 @@ export function instrumentMessageStream return stream; } -/** Handle returned by {@link instrumentRawSseBody} for the `Stream`-iterator path to claim the span. */ -export interface RawSseBodyHandle { - /** Called when the SDK `Stream`'s async iterator takes over, so the body wrapper stays a pass-through. */ - claim: () => void; -} - /** * Replace `response.body` with a pass-through that accumulates the SSE frames flowing through it and - * ends `span` when the body is exhausted, cancelled or errors. + * ends `span` once the body is exhausted, cancelled or errors. * * Every way of draining an Anthropic stream bottoms out in `response.body`: the SDK `Stream`'s async - * iterator, `tee()`, and a caller reading `.asResponse()`/`.withResponse()`'s raw `Response`. Only the - * first of those is visible to {@link instrumentAsyncIterableStream}, so without this the other two end - * no span at all. When the iterator path does run it claims the span and this wrapper goes quiet, so a - * chunk is never accounted for twice. + * iterator, `tee()`, and a caller reading the raw `Response` from `.asResponse()`/`.withResponse()`. + * Instrumenting the body instead of the `Stream` covers all of them with one accumulator. * - * Returns `undefined` — leaving the response untouched — for a body we can't wrap. + * Returns `false`, leaving the response untouched, for a body we can't wrap. */ -export function instrumentRawSseBody( - response: { body?: unknown }, - span: Span, - recordOutputs: boolean, -): RawSseBodyHandle | undefined { +export function instrumentRawSseBody(response: { body?: unknown }, span: Span, recordOutputs: boolean): boolean { const body = response.body as ReadableStream | null | undefined; if (!body || typeof body.getReader !== 'function') { - return undefined; + return false; } const state = createStreamingState(); const decoder = new TextDecoder(); let buffered = ''; - let claimed = false; let settled = false; // Never lets an accumulation failure reach the caller: their stream matters more than our attributes. @@ -325,14 +312,13 @@ export function instrumentRawSseBody( } }; - const settle = (error?: unknown): void => { - if (settled || claimed) { + // No error status on a torn-down body, matching `instrumentAsyncIterableStream`: the SDK surfaces the + // failure to the caller, and an `error` SSE frame already marks the span through `isErrorEvent`. + const settle = (): void => { + if (settled) { return; } settled = true; - if (error !== undefined) { - span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); - } endStreamSpan(span, state, recordOutputs); }; @@ -351,12 +337,10 @@ export function instrumentRawSseBody( controller.close(); return; } - if (!claimed) { - consume(value); - } + consume(value); controller.enqueue(value); } catch (error) { - settle(error); + settle(); controller.error(error); } }, @@ -374,12 +358,8 @@ export function instrumentRawSseBody( // `body` is a prototype getter, so an own data property shadows it for every later read. Object.defineProperty(response, 'body', { value: instrumented, configurable: true }); } catch { - return undefined; + return false; } - return { - claim: () => { - claimed = true; - }, - }; + return true; } diff --git a/packages/server-utils/src/integrations/anthropic.ts b/packages/server-utils/src/integrations/anthropic.ts index dfa072ae9db6..5898ba9bbd57 100644 --- a/packages/server-utils/src/integrations/anthropic.ts +++ b/packages/server-utils/src/integrations/anthropic.ts @@ -13,12 +13,7 @@ import { } from '@sentry/core'; import { getGenAiSpanOp, resolveAIRecordingOptions } from '../ai/core/utils'; import { addPrivateRequestAttributes, addResponseAttributes, extractRequestAttributes } from '../ai/anthropic-ai'; -import type { RawSseBodyHandle } from '../ai/anthropic-ai/streaming'; -import { - instrumentAsyncIterableStream, - instrumentMessageStream, - instrumentRawSseBody, -} from '../ai/anthropic-ai/streaming'; +import { instrumentMessageStream, instrumentRawSseBody } from '../ai/anthropic-ai/streaming'; import type { AnthropicAiOptions, AnthropicAiResponse } from '../ai/anthropic-ai/types'; import { CHANNELS } from '../orchestrion/channels'; import { bindTracingChannelToSpan } from '../tracing-channel'; @@ -48,13 +43,12 @@ interface AnthropicChannelContext { result?: unknown; } -// Spans opened for `messages.create({ stream: true })`, i.e. the ones whose stream is drained through -// the SDK's `Stream`. `messages.stream()` spans are excluded: `instrumentMessageStream` already owns -// when those end, so the raw-body wrapper must keep its hands off them. -const asyncIterableStreamSpans = new WeakSet(); +// Spans from `messages.create()`. `messages.stream()` spans are excluded because +// `instrumentMessageStream` already owns when those end. +const createSpans = new WeakSet(); -// The `Stream` a raw-body wrapper was installed for, so the iterator path can claim the span. -const rawSseBodyHandles = new WeakMap(); +// Spans whose SSE response body we wrapped, so `wrapStreamResult` knows the wrapper will end them. +const bodyOwnedSpans = new WeakSet(); const _anthropicAIIntegration = ((options: AnthropicAiOptions = {}) => { return { @@ -99,15 +93,13 @@ function subscribeToSseStream(options: AnthropicAiOptions): void { diagnosticsChannel.tracingChannel(CHANNELS.ANTHROPIC_SSE_STREAM).end.subscribe(message => { const data = message as AnthropicChannelContext; const span = getActiveSpan(); - const stream = data.result; const response = data.arguments?.[0]; - if (!span || !asyncIterableStreamSpans.has(span) || !isObjectLike(stream) || !isObjectLike(response)) { + if (!span || !createSpans.has(span) || !isObjectLike(response)) { return; } - const handle = instrumentRawSseBody(response, span, recordOutputs); - if (handle) { - rawSseBodyHandles.set(stream, handle); + if (instrumentRawSseBody(response, span, recordOutputs)) { + bodyOwnedSpans.add(span); } }); } @@ -159,29 +151,24 @@ function createGenAiSpan( } if (stream === 'async-iterable') { - asyncIterableStreamSpans.add(span); + createSpans.add(span); } return span; } -type AsyncIterableStream = { [Symbol.asyncIterator]: () => AsyncIterator }; type MessageStreamEmitter = { on: (...args: unknown[]) => void }; -function isAsyncIterable(value: unknown): value is AsyncIterableStream { - return !!value && typeof (value as AsyncIterableStream)[Symbol.asyncIterator] === 'function'; -} - function isMessageStream(value: unknown): value is MessageStreamEmitter { return !!value && typeof (value as MessageStreamEmitter).on === 'function'; } /** * Hand span-ending ownership to a streamed result: returns `true` to skip the normal `beforeSpanEnd`, - * `false` for non-streaming results (which end via `beforeSpanEnd`). + * `false` for results that end via `beforeSpanEnd`. * - * - `async-iterable`: patch the `Stream`'s async iterator in place so `instrumentAsyncIterableStream` ends - * the span when iteration finishes. + * - `async-iterable`: the SSE body wrapper ends the span once the body is drained. A streaming call we + * couldn't wrap ends here instead, carrying request attributes only. * - `message-stream`: `instrumentMessageStream` attaches `'message'`/`'error'` listeners that end the span. */ function wrapStreamResult( @@ -190,22 +177,12 @@ function wrapStreamResult( stream: StreamMode, options: AnthropicAiOptions, ): boolean { - const { recordOutputs } = resolveAIRecordingOptions(options); - const result = data.result; - - if (stream === 'async-iterable' && isAsyncIterable(result)) { - const handle = rawSseBodyHandles.get(result); - const iterate = result[Symbol.asyncIterator].bind(result); - const instrumented = instrumentAsyncIterableStream({ [Symbol.asyncIterator]: iterate }, span, recordOutputs); - result[Symbol.asyncIterator] = () => { - handle?.claim(); - return instrumented; - }; - return true; + if (stream === 'async-iterable') { + return bodyOwnedSpans.has(span); } - if (stream === 'message-stream' && isMessageStream(result)) { - instrumentMessageStream(result, span, recordOutputs); + if (stream === 'message-stream' && isMessageStream(data.result)) { + instrumentMessageStream(data.result, span, resolveAIRecordingOptions(options).recordOutputs); return true; } From f26ca68dedaa06ecfde5ca04762a78dc5fbe1407 Mon Sep 17 00:00:00 2001 From: RulaKhaled Date: Fri, 18 Sep 2026 11:08:06 +0300 Subject: [PATCH 5/5] fix(server-utils): Cover every way an Anthropic SSE body gets drained Handing the span to the SSE body wrapper left several paths uncovered. A body that is not a web ReadableStream, which is what an injected node-fetch or undici shim hands back, could not be wrapped at all, so the span ended the moment create() resolved with request attributes only. Patching the SDK Stream's iterator now serves as the fallback for those. clone() tees the response's internal body and swaps in one branch, which left the stream the wrapper had captured locked, so the next read threw. The wrapper now resolves the source through the prototype getter on every read instead of holding on to the stream it was handed. text(), json() and arrayBuffer() read the internal body and never touch the property we shadow, so the span never ended. Those are wrapped too, and text() and arrayBuffer() feed their result through the accumulator so the response attributes survive. Three smaller ones: recordOutputs is resolved per call rather than once at subscribe time, since subscribeToSseStream runs once per process and a later client can carry different options. The frame parser's try now covers a single frame, so one unparsable line no longer costs us the rest of its chunk, where message_delta and message_stop ride. An unsampled span skips the wrap entirely, and settle() flushes a trailing frame left without a newline. The pass-through is a byte stream now, so getReader({ mode: 'byob' }) keeps working on a response that supported it before. instrumentRawSseBody moved to its own file to stay under the line cap. Co-Authored-By: Claude Opus 5 --- .../anthropic/scenario-stream-node-body.mjs | 97 +++++++++ .../anthropic/scenario-stream-raw-body.mjs | 28 +++ .../suites/tracing/anthropic/test.ts | 32 ++- .../src/ai/anthropic-ai/sse-body.ts | 202 ++++++++++++++++++ .../src/ai/anthropic-ai/streaming.ts | 108 ++-------- .../src/integrations/anthropic.ts | 35 ++- .../ai/lib/tracing/anthropic-sse-body.test.ts | 158 ++++++++++++++ 7 files changed, 554 insertions(+), 106 deletions(-) create mode 100644 dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-stream-node-body.mjs create mode 100644 packages/server-utils/src/ai/anthropic-ai/sse-body.ts create mode 100644 packages/server-utils/test/ai/lib/tracing/anthropic-sse-body.test.ts diff --git a/dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-stream-node-body.mjs b/dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-stream-node-body.mjs new file mode 100644 index 000000000000..ac6dabdf5b42 --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-stream-node-body.mjs @@ -0,0 +1,97 @@ +import { Readable } from 'node:stream'; +import Anthropic from '@anthropic-ai/sdk'; +import * as Sentry from '@sentry/node'; +import express from 'express'; + +function startMockAnthropicServer() { + const app = express(); + app.use(express.json()); + + app.post('/anthropic/v1/messages', (req, res) => { + res.writeHead(200, { + 'Content-Type': 'text/event-stream', + 'Cache-Control': 'no-cache', + Connection: 'keep-alive', + }); + + const model = req.body.model; + const events = [ + { + type: 'message_start', + message: { + id: 'msg_node_body', + type: 'message', + role: 'assistant', + model, + content: [], + usage: { input_tokens: 10 }, + }, + }, + { type: 'content_block_start', index: 0, content_block: { type: 'text', text: '' } }, + { type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: 'Node ' } }, + { type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: 'body!' } }, + { type: 'content_block_stop', index: 0 }, + { + type: 'message_delta', + delta: { stop_reason: 'end_turn', stop_sequence: null }, + usage: { output_tokens: 15 }, + }, + { type: 'message_stop' }, + ]; + + events.forEach((event, index) => { + setTimeout(() => { + res.write(`event: ${event.type}\n`); + res.write(`data: ${JSON.stringify(event)}\n\n`); + if (index === events.length - 1) { + res.end(); + } + }, index * 10); + }); + }); + + return new Promise(resolve => { + const server = app.listen(0, () => { + resolve(server); + }); + }); +} + +// Stands in for the node-fetch/undici-compat shims callers pass as `fetch`, whose responses carry a +// Node `Readable` body the SSE body wrapper can't wrap. +async function fetchWithNodeStreamBody(url, init) { + const response = await fetch(url, init); + Object.defineProperty(response, 'body', { + value: Readable.fromWeb(response.body), + configurable: true, + }); + return response; +} + +async function run() { + const server = await startMockAnthropicServer(); + + await Sentry.startSpan({ op: 'function', name: 'main' }, async () => { + const client = new Anthropic({ + apiKey: 'mock-api-key', + baseURL: `http://localhost:${server.address().port}/anthropic`, + fetch: fetchWithNodeStreamBody, + }); + + const stream = await client.messages.create({ + model: 'claude-3-haiku-20240307', + messages: [{ role: 'user', content: 'Stream this please' }], + stream: true, + }); + + for await (const _ of stream) { + void _; + } + }); + + await Sentry.flush(2000); + + server.close(); +} + +run(); diff --git a/dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-stream-raw-body.mjs b/dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-stream-raw-body.mjs index 8085803b91a9..e64dab46e7f8 100644 --- a/dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-stream-raw-body.mjs +++ b/dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-stream-raw-body.mjs @@ -89,6 +89,34 @@ async function run() { for await (const _ of stream) { void _; } + + // 3) Clone first, then drain. `clone()` tees the response's internal body and swaps in one branch, + // so the wrapper has to re-read the body rather than hold on to the stream it was handed. + const cloned = await client.messages.create({ ...params }).asResponse(); + const copy = cloned.clone(); + for await (const _ of cloned.body) { + void _; + } + void copy; + + // 4) Read the body as text, which never touches the `body` property at all + const asText = await client.messages.create({ ...params }).asResponse(); + const text = await asText.text(); + if (!text.includes('message_stop')) { + throw new Error('raw Response text did not contain the streamed frames'); + } + + // 5) Read the body through a BYOB reader, which only a byte stream supports + const byob = await client.messages.create({ ...params }).asResponse(); + const reader = byob.body.getReader({ mode: 'byob' }); + let buffer = new ArrayBuffer(1024); + for (;;) { + const { done, value } = await reader.read(new Uint8Array(buffer, 0, 1024)); + if (done) { + break; + } + buffer = value.buffer; + } }); await Sentry.flush(2000); diff --git a/dev-packages/node-integration-tests/suites/tracing/anthropic/test.ts b/dev-packages/node-integration-tests/suites/tracing/anthropic/test.ts index 7971efaaa731..238d94cbfa87 100644 --- a/dev-packages/node-integration-tests/suites/tracing/anthropic/test.ts +++ b/dev-packages/node-integration-tests/suites/tracing/anthropic/test.ts @@ -272,9 +272,10 @@ describe('Anthropic integration', () => { .expect({ span: container => { const genAiSpans = container.items.filter(span => span.attributes['sentry.op']?.value === 'gen_ai.chat'); - // Two calls, one drained via `.asResponse()` and one via the SDK `Stream`. Both must end, - // and both must carry the response attributes accumulated off the SSE frames. - expect(genAiSpans).toHaveLength(2); + // One call per way of draining the stream: the raw `.asResponse()` body, the SDK `Stream`, + // a cloned response, `text()`, and a BYOB reader. Every one of them must end its span with + // the response attributes accumulated off the SSE frames. + expect(genAiSpans).toHaveLength(5); for (const span of genAiSpans) { expect(span.name).toBe('chat claude-3-haiku-20240307'); expect(span.status).toBe('ok'); @@ -294,6 +295,31 @@ describe('Anthropic integration', () => { }); }); + createEsmAndCjsTests(__dirname, 'scenario-stream-node-body.mjs', 'instrument-raw-body.mjs', (createRunner, test) => { + test('ends the span when the response body is not a web ReadableStream', async () => { + await createRunner() + .unordered() + .expect({ + span: container => { + const genAiSpan = container.items.find(span => span.attributes['sentry.op']?.value === 'gen_ai.chat'); + // The body wrapper has nothing to hold on to here, so the SDK `Stream`'s iterator has to + // carry the span instead — otherwise it would end with request attributes only. + expect(genAiSpan).toBeDefined(); + expect(genAiSpan!.status).toBe('ok'); + expect(genAiSpan!.attributes[GEN_AI_RESPONSE_STREAMING].value).toBe(true); + expect(genAiSpan!.attributes[GEN_AI_RESPONSE_ID].value).toBe('msg_node_body'); + expect(genAiSpan!.attributes[GEN_AI_RESPONSE_FINISH_REASONS].value).toBe('["end_turn"]'); + expect(genAiSpan!.attributes[GEN_AI_RESPONSE_TEXT].value).toBe('Node body!'); + expect(genAiSpan!.attributes[GEN_AI_USAGE_INPUT_TOKENS].value).toBe(10); + expect(genAiSpan!.attributes[GEN_AI_USAGE_OUTPUT_TOKENS].value).toBe(15); + expect(genAiSpan!.attributes[GEN_AI_USAGE_TOTAL_TOKENS].value).toBe(25); + }, + }) + .start() + .completed(); + }); + }); + createEsmAndCjsTests(__dirname, 'scenario-stream.mjs', 'instrument-with-pii.mjs', (createRunner, test) => { test('streams record response text when PII true', async () => { await createRunner() diff --git a/packages/server-utils/src/ai/anthropic-ai/sse-body.ts b/packages/server-utils/src/ai/anthropic-ai/sse-body.ts new file mode 100644 index 000000000000..a99f007fa292 --- /dev/null +++ b/packages/server-utils/src/ai/anthropic-ai/sse-body.ts @@ -0,0 +1,202 @@ +import type { Span } from '@sentry/core'; +import { endStreamSpan } from '../core/utils'; +import { createStreamingState, processEvent } from './streaming'; +import type { AnthropicAiStreamingEvent } from './types'; + +/** The slice of a stream controller the SSE pass-through uses, shared by the byte and default variants. */ +interface SseStreamController { + enqueue: (chunk: Uint8Array) => void; + close: () => void; + error: (reason?: unknown) => void; +} + +/** Walks the prototype chain for a getter, so a shadowing own property doesn't hide the original. */ +function findBoundGetter(target: object, key: string): (() => unknown) | undefined { + for (let proto = Object.getPrototypeOf(target); proto; proto = Object.getPrototypeOf(proto)) { + const descriptor = Object.getOwnPropertyDescriptor(proto, key); + if (descriptor?.get) { + return descriptor.get.bind(target); + } + } + return undefined; +} + +/** + * Replace `response.body` with a pass-through that accumulates the SSE frames flowing through it and + * ends `span` once the body is exhausted, cancelled or errors. + * + * Every way of draining an Anthropic stream bottoms out in `response.body`: the SDK `Stream`'s async + * iterator, `tee()`, and a caller reading the raw `Response` from `.asResponse()`/`.withResponse()`. + * Instrumenting the body instead of the `Stream` covers all of them with one accumulator. + * + * Returns `false`, leaving the response untouched, for a body we can't wrap. + * + * @internal Exported for the diagnostics-channel integration. + */ +export function instrumentRawSseBody(response: { body?: unknown }, span: Span, recordOutputs: boolean): boolean { + const body = response.body as ReadableStream | null | undefined; + // An unsampled span is discarded by `endStreamSpan`, so decoding and parsing every frame for it buys + // nothing — leave the body alone and let the caller's response stay exactly as the SDK built it. + if (!body || typeof body.getReader !== 'function' || !span.isRecording()) { + return false; + } + + const state = createStreamingState(); + const decoder = new TextDecoder(); + let buffered = ''; + let settled = false; + + // Never lets an accumulation failure reach the caller: their stream matters more than our attributes. + // Scoped to a single frame so one line we can't parse doesn't cost us its neighbours — `message_delta` + // (token usage) and `message_stop` ride in the last chunk, where a bad frame would hurt most. + const consumeFrame = (line: string): void => { + // An SSE frame's `event:` line only repeats the `type` already carried by the JSON payload. + if (!line.startsWith('data:')) { + return; + } + try { + processEvent(JSON.parse(line.slice(5)) as AnthropicAiStreamingEvent, state, recordOutputs, span); + } catch { + // A frame we can't parse is not worth breaking the caller's stream over. + } + }; + + const consumeText = (text: string): void => { + buffered += text; + + let newline = buffered.indexOf('\n'); + while (newline !== -1) { + consumeFrame(buffered.slice(0, newline).trim()); + buffered = buffered.slice(newline + 1); + newline = buffered.indexOf('\n'); + } + }; + + const consumeBytes = (chunk: Uint8Array): void => { + try { + consumeText(decoder.decode(chunk, { stream: true })); + } catch { + // As above: a chunk we can't decode is not worth breaking the caller's stream over. + } + }; + + // No error status on a torn-down body, matching `instrumentAsyncIterableStream`: the SDK surfaces the + // failure to the caller, and an `error` SSE frame already marks the span through `isErrorEvent`. + const settle = (): void => { + if (settled) { + return; + } + settled = true; + // A body that ends without a trailing newline leaves its last frame sitting in `buffered`. + const trailing = buffered.trim(); + buffered = ''; + if (trailing) { + consumeFrame(trailing); + } + endStreamSpan(span, state, recordOutputs); + }; + + // Resolved on every read rather than captured at wrap time: `clone()` tees the response's *internal* + // body and swaps in one branch, which leaves the stream we were handed permanently locked. + const readBody = findBoundGetter(response, 'body'); + const source = (): ReadableStream => (readBody?.() as ReadableStream | null) ?? body; + + // Acquired on the first read, never at wrap time: taking a reader disturbs the body, which would + // make `response.text()`, `arrayBuffer()` and `clone()` throw on a response nobody has read yet. + let reader: ReadableStreamDefaultReader | undefined; + + const underlyingSource = { + async pull(controller: SseStreamController): Promise { + try { + reader ??= source().getReader(); + const { done, value } = await reader.read(); + if (done) { + settle(); + controller.close(); + return; + } + // Must precede the enqueue: a byte stream transfers the chunk's buffer, detaching it. + consumeBytes(value); + controller.enqueue(value); + } catch (error) { + settle(); + controller.error(error); + } + }, + async cancel(reason: unknown): Promise { + settle(); + await (reader ? reader.cancel(reason) : source().cancel(reason)); + }, + }; + + // A high-water mark of 0 keeps the stream from pulling a chunk before anyone asks for one. The + // default of 1 would read ahead the moment we wrap, disturbing a body the caller may never read. + const strategy = { highWaterMark: 0 }; + let instrumented: ReadableStream; + try { + // `response.body` is a byte stream, so a plain one here would break `getReader({ mode: 'byob' })` + // on a response that supported it before we touched it. + instrumented = new ReadableStream( + { ...underlyingSource, type: 'bytes' } as unknown as UnderlyingSource, + strategy, + ); + } catch { + instrumented = new ReadableStream(underlyingSource as UnderlyingSource, strategy); + } + + try { + // `body` is a prototype getter, so an own data property shadows it for every later read. + Object.defineProperty(response, 'body', { value: instrumented, configurable: true }); + } catch { + return false; + } + + // `text()`, `json()` and `arrayBuffer()` read the response's internal body and never touch the + // property we just shadowed. Without wrapping them too, a caller draining the stream that way would + // leave the span unended forever, since nothing else ends it once we take ownership. + instrumentBodyConsumers(response, consumeText, settle); + + return true; +} + +const BODY_CONSUMERS = ['text', 'json', 'arrayBuffer'] as const; + +function instrumentBodyConsumers( + response: Record, + consumeText: (text: string) => void, + settle: () => void, +): void { + for (const name of BODY_CONSUMERS) { + const original = response[name]; + if (typeof original !== 'function') { + continue; + } + + const wrapped = function (this: unknown, ...args: unknown[]): Promise { + return Promise.resolve((original as (...a: unknown[]) => unknown).apply(this ?? response, args)).then( + result => { + // Whatever came back is the stream we would have accumulated off the body, so read the frames + // out of it rather than settling for a span with request attributes only. `json()` returns + // neither shape, since parsing an SSE body as JSON is a caller error to begin with. + if (typeof result === 'string') { + consumeText(result); + } else if (result instanceof ArrayBuffer) { + consumeText(new TextDecoder().decode(result)); + } + settle(); + return result; + }, + error => { + settle(); + throw error; + }, + ); + }; + + try { + Object.defineProperty(response, name, { value: wrapped, configurable: true, writable: true }); + } catch { + // A consumer we can't wrap just means the span leans on the body wrapper to end. + } + } +} diff --git a/packages/server-utils/src/ai/anthropic-ai/streaming.ts b/packages/server-utils/src/ai/anthropic-ai/streaming.ts index 2f14f24e4af6..4b28caf137f3 100644 --- a/packages/server-utils/src/ai/anthropic-ai/streaming.ts +++ b/packages/server-utils/src/ai/anthropic-ai/streaming.ts @@ -6,8 +6,10 @@ import { mapAnthropicErrorToStatusMessage } from './utils'; /** * State object used to accumulate information from a stream of Anthropic AI events. + * + * @internal Exported for the SSE body wrapper. */ -interface StreamingState { +export interface StreamingState { /** Collected response text fragments (for output recording). */ responseTexts: string[]; /** Reasons for finishing the response, as reported by the API. */ @@ -37,7 +39,8 @@ interface StreamingState { >; } -function createStreamingState(): StreamingState { +/** @internal Exported for the SSE body wrapper. */ +export function createStreamingState(): StreamingState { return { responseTexts: [], finishReasons: [], @@ -180,12 +183,14 @@ function handleContentBlockStop(event: AnthropicAiStreamingEvent, state: Streami /** * Processes an event + * + * @internal Exported for the SSE body wrapper. * @param event - The event to process * @param state - The state of the streaming process * @param recordOutputs - Whether to record outputs * @param span - The span to update */ -function processEvent( +export function processEvent( event: AnthropicAiStreamingEvent, state: StreamingState, recordOutputs: boolean, @@ -214,6 +219,8 @@ function processEvent( * Instruments an async iterable stream of Anthropic events, updates the span with * streaming attributes and (optionally) the aggregated output text, and yields * each event from the input stream unchanged. + * + * @internal Exported for the Anthropic instrumentation. */ export async function* instrumentAsyncIterableStream( stream: AsyncIterable, @@ -234,6 +241,8 @@ export async function* instrumentAsyncIterableStream( /** * Instruments a MessageStream by registering event handlers and preserving the original stream API. + * + * @internal Exported for the Anthropic instrumentation. */ export function instrumentMessageStream void }>( stream: R, @@ -270,96 +279,3 @@ export function instrumentMessageStream return stream; } - -/** - * Replace `response.body` with a pass-through that accumulates the SSE frames flowing through it and - * ends `span` once the body is exhausted, cancelled or errors. - * - * Every way of draining an Anthropic stream bottoms out in `response.body`: the SDK `Stream`'s async - * iterator, `tee()`, and a caller reading the raw `Response` from `.asResponse()`/`.withResponse()`. - * Instrumenting the body instead of the `Stream` covers all of them with one accumulator. - * - * Returns `false`, leaving the response untouched, for a body we can't wrap. - */ -export function instrumentRawSseBody(response: { body?: unknown }, span: Span, recordOutputs: boolean): boolean { - const body = response.body as ReadableStream | null | undefined; - if (!body || typeof body.getReader !== 'function') { - return false; - } - - const state = createStreamingState(); - const decoder = new TextDecoder(); - let buffered = ''; - let settled = false; - - // Never lets an accumulation failure reach the caller: their stream matters more than our attributes. - const consume = (chunk: Uint8Array): void => { - try { - buffered += decoder.decode(chunk, { stream: true }); - - let newline = buffered.indexOf('\n'); - while (newline !== -1) { - const line = buffered.slice(0, newline).trim(); - buffered = buffered.slice(newline + 1); - // An SSE frame's `event:` line only repeats the `type` already carried by the JSON payload. - if (line.startsWith('data:')) { - processEvent(JSON.parse(line.slice(5)) as AnthropicAiStreamingEvent, state, recordOutputs, span); - } - newline = buffered.indexOf('\n'); - } - } catch { - // A frame we can't decode or parse is not worth breaking the caller's stream over. - } - }; - - // No error status on a torn-down body, matching `instrumentAsyncIterableStream`: the SDK surfaces the - // failure to the caller, and an `error` SSE frame already marks the span through `isErrorEvent`. - const settle = (): void => { - if (settled) { - return; - } - settled = true; - endStreamSpan(span, state, recordOutputs); - }; - - // Acquired on the first read, never at wrap time: taking a reader disturbs the body, which would - // make `response.text()`, `arrayBuffer()` and `clone()` throw on a response nobody has read yet. - let reader: ReadableStreamDefaultReader | undefined; - - const instrumented = new ReadableStream( - { - async pull(controller) { - try { - reader ??= body.getReader(); - const { done, value } = await reader.read(); - if (done) { - settle(); - controller.close(); - return; - } - consume(value); - controller.enqueue(value); - } catch (error) { - settle(); - controller.error(error); - } - }, - async cancel(reason) { - settle(); - await (reader ? reader.cancel(reason) : body.cancel(reason)); - }, - }, - // A high-water mark of 0 keeps the stream from pulling a chunk before anyone asks for one. The - // default of 1 would read ahead the moment we wrap, disturbing a body the caller may never read. - { highWaterMark: 0 }, - ); - - try { - // `body` is a prototype getter, so an own data property shadows it for every later read. - Object.defineProperty(response, 'body', { value: instrumented, configurable: true }); - } catch { - return false; - } - - return true; -} diff --git a/packages/server-utils/src/integrations/anthropic.ts b/packages/server-utils/src/integrations/anthropic.ts index 5898ba9bbd57..41a682365394 100644 --- a/packages/server-utils/src/integrations/anthropic.ts +++ b/packages/server-utils/src/integrations/anthropic.ts @@ -13,7 +13,8 @@ import { } from '@sentry/core'; import { getGenAiSpanOp, resolveAIRecordingOptions } from '../ai/core/utils'; import { addPrivateRequestAttributes, addResponseAttributes, extractRequestAttributes } from '../ai/anthropic-ai'; -import { instrumentMessageStream, instrumentRawSseBody } from '../ai/anthropic-ai/streaming'; +import { instrumentAsyncIterableStream, instrumentMessageStream } from '../ai/anthropic-ai/streaming'; +import { instrumentRawSseBody } from '../ai/anthropic-ai/sse-body'; import type { AnthropicAiOptions, AnthropicAiResponse } from '../ai/anthropic-ai/types'; import { CHANNELS } from '../orchestrion/channels'; import { bindTracingChannelToSpan } from '../tracing-channel'; @@ -88,8 +89,6 @@ function instrumentAnthropic(options: AnthropicAiOptions): void { * `Stream` the SDK hands back points at the `Response` it was built from. */ function subscribeToSseStream(options: AnthropicAiOptions): void { - const { recordOutputs } = resolveAIRecordingOptions(options); - diagnosticsChannel.tracingChannel(CHANNELS.ANTHROPIC_SSE_STREAM).end.subscribe(message => { const data = message as AnthropicChannelContext; const span = getActiveSpan(); @@ -98,7 +97,7 @@ function subscribeToSseStream(options: AnthropicAiOptions): void { return; } - if (instrumentRawSseBody(response, span, recordOutputs)) { + if (instrumentRawSseBody(response, span, resolveAIRecordingOptions(options).recordOutputs)) { bodyOwnedSpans.add(span); } }); @@ -157,8 +156,13 @@ function createGenAiSpan( return span; } +type AsyncIterableStream = { [Symbol.asyncIterator]: () => AsyncIterator }; type MessageStreamEmitter = { on: (...args: unknown[]) => void }; +function isAsyncIterable(value: unknown): value is AsyncIterableStream { + return !!value && typeof (value as AsyncIterableStream)[Symbol.asyncIterator] === 'function'; +} + function isMessageStream(value: unknown): value is MessageStreamEmitter { return !!value && typeof (value as MessageStreamEmitter).on === 'function'; } @@ -167,8 +171,9 @@ function isMessageStream(value: unknown): value is MessageStreamEmitter { * Hand span-ending ownership to a streamed result: returns `true` to skip the normal `beforeSpanEnd`, * `false` for results that end via `beforeSpanEnd`. * - * - `async-iterable`: the SSE body wrapper ends the span once the body is drained. A streaming call we - * couldn't wrap ends here instead, carrying request attributes only. + * - `async-iterable`: the SSE body wrapper ends the span once the body is drained. When there was no body + * to wrap — an injected `fetch` handing back a Node `Readable`, say — fall back to patching the SDK + * `Stream`'s iterator, which covers every body shape but only the callers who iterate the `Stream`. * - `message-stream`: `instrumentMessageStream` attaches `'message'`/`'error'` listeners that end the span. */ function wrapStreamResult( @@ -178,7 +183,23 @@ function wrapStreamResult( options: AnthropicAiOptions, ): boolean { if (stream === 'async-iterable') { - return bodyOwnedSpans.has(span); + if (bodyOwnedSpans.has(span)) { + return true; + } + + if (span.isRecording() && isAsyncIterable(data.result)) { + const result = data.result; + const iterate = result[Symbol.asyncIterator].bind(result); + const instrumented = instrumentAsyncIterableStream( + { [Symbol.asyncIterator]: iterate }, + span, + resolveAIRecordingOptions(options).recordOutputs, + ); + result[Symbol.asyncIterator] = () => instrumented; + return true; + } + + return false; } if (stream === 'message-stream' && isMessageStream(data.result)) { diff --git a/packages/server-utils/test/ai/lib/tracing/anthropic-sse-body.test.ts b/packages/server-utils/test/ai/lib/tracing/anthropic-sse-body.test.ts new file mode 100644 index 000000000000..13a1c2d88a82 --- /dev/null +++ b/packages/server-utils/test/ai/lib/tracing/anthropic-sse-body.test.ts @@ -0,0 +1,158 @@ +import { afterEach, beforeEach, describe, expect, it } from 'vitest'; +import { + GEN_AI_RESPONSE_FINISH_REASONS, + GEN_AI_RESPONSE_TEXT, + GEN_AI_USAGE_OUTPUT_TOKENS, +} from '@sentry/conventions/attributes'; +import { getMainCarrier, setCurrentClient, spanToStaticSpanJSON, startInactiveSpan } from '@sentry/core'; +import type { Span } from '@sentry/core'; +import { instrumentRawSseBody } from '../../../../src/ai/anthropic-ai/sse-body'; +import { getDefaultTestClientOptions, TestClient } from '../../../mocks/client'; + +const FRAMES = { + start: { type: 'message_start', message: { id: 'msg_1', model: 'claude-3-haiku-20240307', usage: {} } }, + blockStart: { type: 'content_block_start', index: 0, content_block: { type: 'text', text: '' } }, + delta: { type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: 'Hi' } }, + blockStop: { type: 'content_block_stop', index: 0 }, + messageDelta: { type: 'message_delta', delta: { stop_reason: 'end_turn' }, usage: { output_tokens: 7 } }, + stop: { type: 'message_stop' }, +}; + +function sse(frame: unknown): string { + return `event: x\ndata: ${JSON.stringify(frame)}\n\n`; +} + +/** Mirrors the shape the wrapper relies on: `body` as a prototype getter over a byte stream. */ +class FakeResponse { + private _stream: ReadableStream; + + public constructor(chunks: string[]) { + const encoder = new TextEncoder(); + const queue = [...chunks]; + this._stream = new ReadableStream({ + type: 'bytes', + pull(controller) { + const next = queue.shift(); + if (next === undefined) { + controller.close(); + return; + } + controller.enqueue(encoder.encode(next)); + }, + }); + } + + public get body(): ReadableStream { + return this._stream; + } + + public async text(): Promise { + const decoder = new TextDecoder(); + const reader = this._stream.getReader(); + let out = ''; + for (;;) { + const { done, value } = await reader.read(); + if (done) { + return out; + } + out += decoder.decode(value, { stream: true }); + } + } +} + +async function drain(stream: ReadableStream): Promise { + const reader = stream.getReader(); + for (;;) { + const { done } = await reader.read(); + if (done) { + return; + } + } +} + +describe('instrumentRawSseBody', () => { + function setupClient(tracesSampleRate = 1): Span[] { + const client = new TestClient( + getDefaultTestClientOptions({ + dsn: 'https://public@dsn.ingest.sentry.io/1337', + tracesSampleRate, + }), + ); + setCurrentClient(client); + client.init(); + + const endedSpans: Span[] = []; + client.on('spanEnd', span => endedSpans.push(span)); + return endedSpans; + } + + beforeEach(() => { + getMainCarrier().__SENTRY__ = undefined; + }); + + afterEach(() => { + getMainCarrier().__SENTRY__ = undefined; + }); + + it('keeps the frames that follow one it cannot parse', async () => { + setupClient(); + const span = startInactiveSpan({ name: 'chat' }); + // The unparsable frame shares a chunk with the two that carry usage and the finish reason. + const response = new FakeResponse([ + sse(FRAMES.start) + sse(FRAMES.blockStart) + sse(FRAMES.delta) + sse(FRAMES.blockStop), + `event: x\ndata: {not json\n\n${sse(FRAMES.messageDelta)}${sse(FRAMES.stop)}`, + ]); + + expect(instrumentRawSseBody(response, span, true)).toBe(true); + await drain(response.body); + + const data = spanToStaticSpanJSON(span).data; + expect(data[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(7); + expect(data[GEN_AI_RESPONSE_FINISH_REASONS]).toBe('["end_turn"]'); + expect(data[GEN_AI_RESPONSE_TEXT]).toBe('Hi'); + }); + + it('records the final frame of a body that ends without a trailing newline', async () => { + setupClient(); + const span = startInactiveSpan({ name: 'chat' }); + const response = new FakeResponse([ + sse(FRAMES.start) + sse(FRAMES.blockStart) + sse(FRAMES.delta) + sse(FRAMES.blockStop), + `data: ${JSON.stringify(FRAMES.messageDelta)}`, + ]); + + expect(instrumentRawSseBody(response, span, true)).toBe(true); + await drain(response.body); + + expect(spanToStaticSpanJSON(span).data[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(7); + }); + + it('ends the span when the body is drained through text()', async () => { + const endedSpans = setupClient(); + const span = startInactiveSpan({ name: 'chat' }); + const response = new FakeResponse([sse(FRAMES.start), sse(FRAMES.messageDelta), sse(FRAMES.stop)]); + + expect(instrumentRawSseBody(response, span, true)).toBe(true); + await response.text(); + + expect(endedSpans).toHaveLength(1); + expect(spanToStaticSpanJSON(endedSpans[0]!).data[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(7); + }); + + it('leaves the response untouched for an unsampled span', () => { + setupClient(0); + const span = startInactiveSpan({ name: 'chat' }); + const response = new FakeResponse([sse(FRAMES.stop)]); + const body = response.body; + + expect(instrumentRawSseBody(response, span, true)).toBe(false); + expect(response.body).toBe(body); + }); + + it('leaves a response whose body is not a web stream untouched', () => { + setupClient(); + const span = startInactiveSpan({ name: 'chat' }); + const response = { body: { [Symbol.asyncIterator]: () => ({}) } }; + + expect(instrumentRawSseBody(response, span, true)).toBe(false); + }); +});