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-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 new file mode 100644 index 000000000000..e64dab46e7f8 --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/anthropic/scenario-stream-raw-body.mjs @@ -0,0 +1,127 @@ +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(); + + // 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 _; + } + + // 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 _; + } + + // 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); + + 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..238d94cbfa87 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,61 @@ 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'); + // 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'); + 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-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 3daf61c2b57e..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,6 +39,22 @@ interface StreamingState { >; } +/** @internal Exported for the SSE body wrapper. */ +export 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 @@ -165,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, @@ -199,24 +219,15 @@ 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, 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) { @@ -230,24 +241,15 @@ 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, 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); diff --git a/packages/server-utils/src/integrations/anthropic.ts b/packages/server-utils/src/integrations/anthropic.ts index 87bce3b9c199..41a682365394 100644 --- a/packages/server-utils/src/integrations/anthropic.ts +++ b/packages/server-utils/src/integrations/anthropic.ts @@ -4,14 +4,17 @@ 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 { 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'; @@ -41,6 +44,13 @@ interface AnthropicChannelContext { result?: unknown; } +// Spans from `messages.create()`. `messages.stream()` spans are excluded because +// `instrumentMessageStream` already owns when those end. +const createSpans = new WeakSet(); + +// Spans whose SSE response body we wrapped, so `wrapStreamResult` knows the wrapper will end them. +const bodyOwnedSpans = new WeakSet(); + const _anthropicAIIntegration = ((options: AnthropicAiOptions = {}) => { return { name: INTEGRATION_NAME, @@ -54,7 +64,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 +77,30 @@ 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 { + diagnosticsChannel.tracingChannel(CHANNELS.ANTHROPIC_SSE_STREAM).end.subscribe(message => { + const data = message as AnthropicChannelContext; + const span = getActiveSpan(); + const response = data.arguments?.[0]; + if (!span || !createSpans.has(span) || !isObjectLike(response)) { + return; + } + + if (instrumentRawSseBody(response, span, resolveAIRecordingOptions(options).recordOutputs)) { + bodyOwnedSpans.add(span); + } + }); } /** @@ -76,6 +110,7 @@ function instrumentAnthropic(options: AnthropicAiOptions): void { function createGenAiSpan( data: AnthropicChannelContext, operation: string, + stream: StreamMode, options: AnthropicAiOptions, ): Span | undefined { const args = data.arguments ?? []; @@ -114,6 +149,10 @@ function createGenAiSpan( addPrivateRequestAttributes(span, params); } + if (stream === 'async-iterable') { + createSpans.add(span); + } + return span; } @@ -130,10 +169,11 @@ function isMessageStream(value: unknown): value is MessageStreamEmitter { /** * 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. 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( @@ -142,18 +182,28 @@ function wrapStreamResult( stream: StreamMode, options: AnthropicAiOptions, ): boolean { - const { recordOutputs } = resolveAIRecordingOptions(options); - const result = data.result; - - if (stream === 'async-iterable' && isAsyncIterable(result)) { - const iterate = result[Symbol.asyncIterator].bind(result); - const instrumented = instrumentAsyncIterableStream({ [Symbol.asyncIterator]: iterate }, span, recordOutputs); - result[Symbol.asyncIterator] = () => instrumented; - return true; + if (stream === 'async-iterable') { + 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(result)) { - instrumentMessageStream(result, span, recordOutputs); + if (stream === 'message-stream' && isMessageStream(data.result)) { + instrumentMessageStream(data.result, span, resolveAIRecordingOptions(options).recordOutputs); 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..e6e1d073c6a1 100644 --- a/packages/server-utils/src/orchestrion/config/anthropic-ai.ts +++ b/packages/server-utils/src/orchestrion/config/anthropic-ai.ts @@ -26,6 +26,20 @@ 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', + // `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 }, + }, ] satisfies InstrumentationConfig[]; export const anthropicAiModuleNames = getModuleNames(anthropicAiConfig); @@ -33,4 +47,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; 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); + }); +});