diff --git a/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/index.ts b/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/index.ts index 8bf8f8acd720..991300071124 100644 --- a/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/index.ts +++ b/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/index.ts @@ -60,6 +60,13 @@ export default Sentry.withSentry( messages: [{ role: 'user', content: 'What is the weather in SF?' }], }); + const stream = await compiled.stream({ + messages: [{ role: 'user', content: 'Stream the weather in SF' }], + }); + for await (const _chunk of stream) { + // Consuming the iterator is what runs the graph and completes the agent span. + } + return new Response(JSON.stringify({ success: true })); }, }, diff --git a/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/test.ts b/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/test.ts index 87008d27d870..abfff2d13ba6 100644 --- a/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/test.ts +++ b/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/test.ts @@ -4,7 +4,10 @@ import { GEN_AI_INPUT_MESSAGES, GEN_AI_OPERATION_NAME, GEN_AI_PIPELINE_NAME, + GEN_AI_RESPONSE_FINISH_REASONS, GEN_AI_RESPONSE_MODEL, + GEN_AI_RESPONSE_STREAMING, + GEN_AI_RESPONSE_TEXT, GEN_AI_USAGE_INPUT_TOKENS, GEN_AI_USAGE_OUTPUT_TOKENS, GEN_AI_USAGE_TOTAL_TOKENS, @@ -17,7 +20,7 @@ import { getSpanOp, getSpansFromEnvelope } from '../../../spanUtils'; // want to test that the instrumentation does not break in our // cloudflare SDK. -it('traces langgraph compile and invoke operations', async ({ signal }) => { +it('traces langgraph invoke and stream operations', async ({ signal }) => { const runner = createRunner(__dirname) .ignore('event') .expect(envelope => { @@ -26,10 +29,13 @@ it('traces langgraph compile and invoke operations', async ({ signal }) => { expect(segmentSpan?.name).toBe('GET /'); const genAiSpans = spans.filter(span => getSpanOp(span)?.startsWith('gen_ai.')); - expect(genAiSpans).toHaveLength(1); - expect(genAiSpans.map(span => span.name).sort()).toEqual(['invoke_agent weather_assistant']); + expect(genAiSpans).toHaveLength(2); + expect(genAiSpans.map(span => span.name).sort()).toEqual([ + 'invoke_agent weather_assistant', + 'invoke_agent weather_assistant', + ]); - const invokeAgentSpan = genAiSpans.find(span => span.name === 'invoke_agent weather_assistant'); + const invokeAgentSpan = genAiSpans.find(span => span.attributes[GEN_AI_RESPONSE_STREAMING] === undefined); expect(invokeAgentSpan).toBeDefined(); expect(invokeAgentSpan!.status).toBe('ok'); expect(invokeAgentSpan!.attributes[GEN_AI_OPERATION_NAME]).toEqual({ @@ -66,6 +72,38 @@ it('traces langgraph compile and invoke operations', async ({ signal }) => { type: 'integer', value: 30, }); + + const streamSpan = genAiSpans.find(span => span.attributes[GEN_AI_RESPONSE_STREAMING]?.value === true); + expect(streamSpan).toBeDefined(); + expect(streamSpan!.status).toBe('ok'); + expect(streamSpan!.attributes[GEN_AI_INPUT_MESSAGES]).toEqual({ + type: 'string', + value: '[{"role":"user","content":"Stream the weather in SF"}]', + }); + expect(streamSpan!.attributes[GEN_AI_RESPONSE_TEXT]).toEqual({ + type: 'string', + value: '[{"role":"assistant","content":"Mock response from LangGraph agent"}]', + }); + expect(streamSpan!.attributes[GEN_AI_RESPONSE_MODEL]).toEqual({ + type: 'string', + value: 'mock-model', + }); + expect(streamSpan!.attributes[GEN_AI_RESPONSE_FINISH_REASONS]).toEqual({ + type: 'array', + value: ['stop'], + }); + expect(streamSpan!.attributes[GEN_AI_USAGE_INPUT_TOKENS]).toEqual({ + type: 'integer', + value: 20, + }); + expect(streamSpan!.attributes[GEN_AI_USAGE_OUTPUT_TOKENS]).toEqual({ + type: 'integer', + value: 10, + }); + expect(streamSpan!.attributes[GEN_AI_USAGE_TOTAL_TOKENS]).toEqual({ + type: 'integer', + value: 30, + }); }) .start(signal); await runner.makeRequest('get', '/'); diff --git a/dev-packages/node-integration-tests/suites/tracing/langgraph/agent-stream-scenario.mjs b/dev-packages/node-integration-tests/suites/tracing/langgraph/agent-stream-scenario.mjs new file mode 100644 index 000000000000..cca20e95102a --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/langgraph/agent-stream-scenario.mjs @@ -0,0 +1,53 @@ +import { ChatAnthropic } from '@langchain/anthropic'; +import { HumanMessage, SystemMessage } from '@langchain/core/messages'; +import { createReactAgent } from '@langchain/langgraph/prebuilt'; +import * as Sentry from '@sentry/node'; +import express from 'express'; + +function startMockAnthropicServer() { + const app = express(); + app.use(express.json()); + + app.post('/v1/messages', (req, res) => { + res.json({ + id: 'msg_react_agent_stream_123', + type: 'message', + role: 'assistant', + content: [{ type: 'text', text: 'Paris is the capital of France.' }], + model: req.body.model, + stop_reason: 'end_turn', + stop_sequence: null, + usage: { input_tokens: 20, output_tokens: 10 }, + }); + }); + + return new Promise(resolve => { + const server = app.listen(0, () => resolve(server)); + }); +} + +async function run() { + const server = await startMockAnthropicServer(); + const baseUrl = `http://localhost:${server.address().port}`; + + await Sentry.startSpan({ op: 'function', name: 'main' }, async () => { + const llm = new ChatAnthropic({ + model: 'claude-3-5-sonnet-20241022', + apiKey: 'mock-api-key', + clientOptions: { baseURL: baseUrl }, + }); + const agent = createReactAgent({ llm, tools: [], name: 'helpful_assistant' }); + + const stream = await agent.stream({ + messages: [new SystemMessage('You are a helpful assistant.'), new HumanMessage('What is the capital of France?')], + }); + for await (const _chunk of stream) { + // Consume the stream so the agent and its child spans complete before the root span ends. + } + }); + + await Sentry.flush(2000); + server.close(); +} + +run(); diff --git a/dev-packages/node-integration-tests/suites/tracing/langgraph/scenario.mjs b/dev-packages/node-integration-tests/suites/tracing/langgraph/scenario.mjs index d93c4b5491c7..86d3f34a73a4 100644 --- a/dev-packages/node-integration-tests/suites/tracing/langgraph/scenario.mjs +++ b/dev-packages/node-integration-tests/suites/tracing/langgraph/scenario.mjs @@ -44,6 +44,16 @@ async function run() { { role: 'user', content: 'Tell me about the weather' }, ], }); + + const stream = await graph.stream({ + messages: [{ role: 'user', content: 'Stream the weather forecast' }], + }); + await stream.pipeTo(new WritableStream()); + + const canceledStream = await graph.stream({ + messages: [{ role: 'user', content: 'Cancel the weather forecast' }], + }); + await canceledStream.cancel('no longer needed'); }); await Sentry.flush(2000); diff --git a/dev-packages/node-integration-tests/suites/tracing/langgraph/test.ts b/dev-packages/node-integration-tests/suites/tracing/langgraph/test.ts index 53dc3f6846af..975c7ab9521d 100644 --- a/dev-packages/node-integration-tests/suites/tracing/langgraph/test.ts +++ b/dev-packages/node-integration-tests/suites/tracing/langgraph/test.ts @@ -6,7 +6,9 @@ import { GEN_AI_OPERATION_NAME, GEN_AI_OUTPUT_MESSAGES, GEN_AI_PIPELINE_NAME, + GEN_AI_RESPONSE_FINISH_REASONS, GEN_AI_RESPONSE_MODEL, + GEN_AI_RESPONSE_STREAMING, GEN_AI_RESPONSE_TEXT, GEN_AI_RESPONSE_TOOL_CALLS, GEN_AI_SYSTEM_INSTRUCTIONS, @@ -34,14 +36,16 @@ describe('LangGraph integration', () => { const segment = container.items.find(span => span.is_segment && span.name === 'langgraph-test'); expect(segment).toBeDefined(); const spans = container.items.filter(span => span.attributes[SENTRY_ORIGIN]?.value === 'auto.ai.langgraph'); - expect(spans).toHaveLength(2); + expect(spans).toHaveLength(4); expect(spans.map(span => span.name).sort()).toEqual([ 'invoke_agent weather_assistant', 'invoke_agent weather_assistant', + 'invoke_agent weather_assistant', + 'invoke_agent weather_assistant', ]); const invokeAgentSpans = spans.filter(span => span.name === 'invoke_agent weather_assistant'); - expect(invokeAgentSpans).toHaveLength(2); + expect(invokeAgentSpans).toHaveLength(4); for (const span of invokeAgentSpans) { expect(span.status).toBe('ok'); expect(span.attributes[SENTRY_OP].value).toBe('gen_ai.invoke_agent'); @@ -50,6 +54,11 @@ describe('LangGraph integration', () => { expect(span.attributes[GEN_AI_AGENT_NAME].value).toBe('weather_assistant'); expect(span.attributes[GEN_AI_PIPELINE_NAME].value).toBe('weather_assistant'); } + + const streamSpans = invokeAgentSpans.filter( + span => span.attributes[GEN_AI_RESPONSE_STREAMING]?.value === true, + ); + expect(streamSpans).toHaveLength(2); }, }) .start() @@ -65,7 +74,7 @@ describe('LangGraph integration', () => { const segment = container.items.find(span => span.is_segment && span.name === 'langgraph-test'); expect(segment).toBeDefined(); const spans = container.items.filter(span => span.attributes[SENTRY_ORIGIN]?.value === 'auto.ai.langgraph'); - expect(spans).toHaveLength(2); + expect(spans).toHaveLength(4); const weatherTodaySpan = spans.find(span => getStringAttributeValue(span.attributes[GEN_AI_INPUT_MESSAGES]?.value)?.includes( @@ -85,6 +94,32 @@ describe('LangGraph integration', () => { expect(weatherDetailsSpan!.name).toBe('invoke_agent weather_assistant'); expect(weatherDetailsSpan!.status).toBe('ok'); expect(weatherDetailsSpan!.attributes[SENTRY_OP].value).toBe('gen_ai.invoke_agent'); + + const weatherStreamSpan = spans.find(span => + getStringAttributeValue(span.attributes[GEN_AI_INPUT_MESSAGES]?.value)?.includes( + 'Stream the weather forecast', + ), + ); + expect(weatherStreamSpan).toBeDefined(); + expect(weatherStreamSpan!.status).toBe('ok'); + expect(weatherStreamSpan!.attributes[GEN_AI_RESPONSE_STREAMING].value).toBe(true); + expect(weatherStreamSpan!.attributes[GEN_AI_RESPONSE_TEXT].value).toBe( + '[{"role":"assistant","content":"Mock LLM response"}]', + ); + expect(weatherStreamSpan!.attributes[GEN_AI_RESPONSE_MODEL].value).toBe('mock-model'); + expect(weatherStreamSpan!.attributes[GEN_AI_RESPONSE_FINISH_REASONS].value).toEqual(['stop']); + expect(weatherStreamSpan!.attributes[GEN_AI_USAGE_INPUT_TOKENS].value).toBe(20); + expect(weatherStreamSpan!.attributes[GEN_AI_USAGE_OUTPUT_TOKENS].value).toBe(10); + expect(weatherStreamSpan!.attributes[GEN_AI_USAGE_TOTAL_TOKENS].value).toBe(30); + + const canceledStreamSpan = spans.find(span => + getStringAttributeValue(span.attributes[GEN_AI_INPUT_MESSAGES]?.value)?.includes( + 'Cancel the weather forecast', + ), + ); + expect(canceledStreamSpan).toBeDefined(); + expect(canceledStreamSpan!.status).toBe('ok'); + expect(canceledStreamSpan!.attributes[GEN_AI_RESPONSE_STREAMING].value).toBe(true); }, }) .start() @@ -294,6 +329,34 @@ describe('LangGraph integration', () => { }); }); + createEsmAndCjsTests(__dirname, 'agent-stream-scenario.mjs', 'instrument-agent.mjs', (createRunner, test) => { + test('parents the streaming agent chat span under its agent span', { timeout: 30000 }, async () => { + await createRunner() + .expect({ + span: container => { + const agentSpans = container.items.filter( + span => span.attributes[SENTRY_OP]?.value === 'gen_ai.invoke_agent', + ); + const chatSpans = container.items.filter(span => span.attributes[SENTRY_OP]?.value === 'gen_ai.chat'); + expect(agentSpans).toHaveLength(1); + expect(chatSpans).toHaveLength(1); + + const agentSpan = agentSpans[0]!; + const chatSpan = chatSpans[0]!; + expect(agentSpan.name).toBe('invoke_agent helpful_assistant'); + expect(agentSpan.status).toBe('ok'); + expect(agentSpan.attributes[GEN_AI_RESPONSE_STREAMING]?.value).toBe(true); + expect(chatSpan.parent_span_id).toBe(agentSpan.span_id); + expect(chatSpan.trace_id).toBe(agentSpan.trace_id); + expect(chatSpan.status).toBe('ok'); + expect(chatSpan.attributes[GEN_AI_AGENT_NAME]?.value).toBe('helpful_assistant'); + }, + }) + .start() + .completed(); + }); + }); + // createReactAgent with tools - verifies tool execution spans (asserted order-independently, see above). createEsmAndCjsTests(__dirname, 'agent-tools-scenario.mjs', 'instrument-agent.mjs', (createRunner, test) => { test('should create tool execution spans for createReactAgent with tools', { timeout: 30000 }, async () => { diff --git a/dev-packages/node-integration-tests/suites/tracing/langgraph/v1/scenario-stream-edge.mjs b/dev-packages/node-integration-tests/suites/tracing/langgraph/v1/scenario-stream-edge.mjs new file mode 100644 index 000000000000..6200fbecb2dd --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/langgraph/v1/scenario-stream-edge.mjs @@ -0,0 +1,52 @@ +import { FakeListChatModel } from '@langchain/core/utils/testing'; +import { END, MessagesAnnotation, START, StateGraph } from '@langchain/langgraph'; +import * as Sentry from '@sentry/node'; + +function weatherGraph(model) { + return new StateGraph(MessagesAnnotation) + .addNode('agent', async state => ({ messages: [await model.invoke(state.messages)] })) + .addEdge(START, 'agent') + .addEdge('agent', END); +} + +async function run() { + const model = new FakeListChatModel({ responses: ['Sunny', 'Cloudy', 'Rainy', 'Windy'] }); + const user = content => ({ messages: [{ role: 'user', content }] }); + + await Sentry.startSpan({ op: 'function', name: 'main' }, async () => { + const graph = weatherGraph(model).compile({ name: 'weather_graph' }); + + const canceled = await graph.stream(user('Cancel while iterating')); + try { + for await (const _chunk of canceled) { + await canceled.cancel('done'); + } + } catch { + // LangGraph rejects the next read after a cancel; the caller asked for it. + } + + for await (const _chunk of await graph.stream(user('Stream with subgraphs'), { + streamMode: 'values', + subgraphs: true, + })) { + // drain + } + + const [first, second] = (await graph.stream(user('Tee the forecast'))).tee(); + for await (const _chunk of first) { + // drain + } + for await (const _chunk of second) { + // drain + } + + const manual = Sentry.instrumentStateGraph(weatherGraph(model), { recordInputs: false }).compile({ + name: 'manual_graph', + }); + for await (const _chunk of await manual.stream(user('Keep this prompt private'))) { + // drain + } + }); +} + +run(); diff --git a/dev-packages/node-integration-tests/suites/tracing/langgraph/v1/scenario-stream.mjs b/dev-packages/node-integration-tests/suites/tracing/langgraph/v1/scenario-stream.mjs new file mode 100644 index 000000000000..765bc3349d28 --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/langgraph/v1/scenario-stream.mjs @@ -0,0 +1,30 @@ +import { FakeListChatModel } from '@langchain/core/utils/testing'; +import { END, MessagesAnnotation, START, StateGraph } from '@langchain/langgraph'; +import * as Sentry from '@sentry/node'; + +async function run() { + const model = new FakeListChatModel({ responses: ['Sunny', 'Cloudy', 'Rainy'] }); + + await Sentry.startSpan({ op: 'function', name: 'main' }, async () => { + const graph = new StateGraph(MessagesAnnotation) + .addNode('agent', async state => ({ messages: [await model.invoke(state.messages)] })) + .addEdge(START, 'agent') + .addEdge('agent', END) + .compile({ name: 'weather_graph' }); + + await graph.invoke({ messages: [{ role: 'user', content: 'Invoke the forecast' }] }); + + for await (const _chunk of await graph.stream({ messages: [{ role: 'user', content: 'Stream the forecast' }] })) { + // Consume the whole stream. + } + + for await (const _chunk of await graph.stream( + { messages: [{ role: 'user', content: 'Abandon the forecast' }] }, + { streamMode: 'values' }, + )) { + break; + } + }); +} + +run(); diff --git a/dev-packages/node-integration-tests/suites/tracing/langgraph/v1/test.ts b/dev-packages/node-integration-tests/suites/tracing/langgraph/v1/test.ts index b09261efc29e..c595478f01a6 100644 --- a/dev-packages/node-integration-tests/suites/tracing/langgraph/v1/test.ts +++ b/dev-packages/node-integration-tests/suites/tracing/langgraph/v1/test.ts @@ -3,6 +3,8 @@ import { GEN_AI_AGENT_NAME, GEN_AI_INPUT_MESSAGES, GEN_AI_OUTPUT_MESSAGES, + GEN_AI_RESPONSE_STREAMING, + GEN_AI_RESPONSE_TEXT, SENTRY_OP, SENTRY_ORIGIN, } from '@sentry/conventions/attributes'; @@ -85,4 +87,88 @@ describe('LangGraph integration (v1)', () => { }, }, ); + + createEsmAndCjsTests( + __dirname, + 'scenario-stream.mjs', + 'instrument-with-pii.mjs', + (createRunner, test) => { + test('creates one agent span per invoke() and stream() run, with chat spans under it', async () => { + const runner = createRunner(); + const spansPromise = runner.collectStreamedSpansUntilSegment('main'); + + await runner.start().completed(); + + const spans = await spansPromise; + const rootSpan = spans.find(span => span.is_segment && span.name === 'main')!; + const agentSpans = spans.filter(span => span.attributes[SENTRY_OP]?.value === 'gen_ai.invoke_agent'); + const chatSpans = spans.filter(span => span.attributes[SENTRY_OP]?.value === 'gen_ai.chat'); + expect(agentSpans).toHaveLength(3); + expect(chatSpans).toHaveLength(3); + + const agentSpanFor = (prompt: string) => + agentSpans.find(span => String(span.attributes[GEN_AI_INPUT_MESSAGES]?.value).includes(prompt))!; + const invokeSpan = agentSpanFor('Invoke the forecast'); + const streamSpan = agentSpanFor('Stream the forecast'); + const abandonedSpan = agentSpanFor('Abandon the forecast'); + + expect(invokeSpan.attributes[GEN_AI_RESPONSE_STREAMING]).toBeUndefined(); + expect(invokeSpan.attributes[GEN_AI_RESPONSE_TEXT].value).toBe('[{"role":"assistant","content":"Sunny"}]'); + expect(streamSpan.attributes[GEN_AI_RESPONSE_STREAMING].value).toBe(true); + expect(streamSpan.attributes[GEN_AI_RESPONSE_TEXT].value).toBe('[{"role":"assistant","content":"Cloudy"}]'); + expect(abandonedSpan.attributes[GEN_AI_RESPONSE_STREAMING].value).toBe(true); + + for (const agentSpan of agentSpans) { + expect(agentSpan.name).toBe('invoke_agent weather_graph'); + expect(agentSpan.status).toBe('ok'); + expect(agentSpan.parent_span_id).toBe(rootSpan.span_id); + expect(chatSpans.filter(span => span.parent_span_id === agentSpan.span_id)).toHaveLength(1); + } + }); + }, + { + additionalDependencies: { + '@langchain/core': '^1.0.0', + '@langchain/langgraph': '^1.0.0', + }, + }, + ); + + createEsmAndCjsTests( + __dirname, + 'scenario-stream-edge.mjs', + 'instrument-with-pii.mjs', + (createRunner, test) => { + test('handles cancel, subgraphs, tee() and manual options on streamed runs', async () => { + const runner = createRunner(); + const spansPromise = runner.collectStreamedSpansUntilSegment('main'); + + await runner.start().completed(); + + const spans = await spansPromise; + const agentSpans = spans.filter(span => span.attributes[SENTRY_OP]?.value === 'gen_ai.invoke_agent'); + expect(agentSpans).toHaveLength(4); + const agentSpanFor = (prompt: string) => + agentSpans.find(span => String(span.attributes[GEN_AI_INPUT_MESSAGES]?.value).includes(prompt))!; + + expect(agentSpanFor('Cancel while iterating').status).toBe('ok'); + expect(agentSpanFor('Stream with subgraphs').attributes[GEN_AI_RESPONSE_TEXT].value).toBe( + '[{"role":"assistant","content":"Cloudy"}]', + ); + expect(agentSpanFor('Tee the forecast').attributes[GEN_AI_RESPONSE_TEXT].value).toBe( + '[{"role":"assistant","content":"Rainy"}]', + ); + + const manualSpan = agentSpans.find(span => span.name === 'invoke_agent manual_graph')!; + expect(manualSpan.attributes[GEN_AI_INPUT_MESSAGES]).toBeUndefined(); + expect(manualSpan.attributes[GEN_AI_RESPONSE_TEXT].value).toBe('[{"role":"assistant","content":"Windy"}]'); + }); + }, + { + additionalDependencies: { + '@langchain/core': '^1.0.0', + '@langchain/langgraph': '^1.0.0', + }, + }, + ); }); diff --git a/packages/server-utils/src/ai/core/utils.ts b/packages/server-utils/src/ai/core/utils.ts index 93949ec40eca..46516dbb3d6a 100644 --- a/packages/server-utils/src/ai/core/utils.ts +++ b/packages/server-utils/src/ai/core/utils.ts @@ -88,6 +88,34 @@ export function buildMethodPath(currentPath: string, prop: string): string { return currentPath ? `${currentPath}.${prop}` : prop; } +/** + * A stream that pulls through `reader`. `tee` and `pipe*` read through internal slots and bypass a + * patched `getReader`; this routes them back through it, pulling on demand and forwarding `cancel` (#24054). + */ +export function readableFromReader(reader: { + read: () => Promise<{ done: boolean; value?: unknown }>; + cancel?: (reason?: unknown) => Promise; +}): ReadableStream { + return new ReadableStream({ + async pull(controller) { + const { done, value } = await reader.read(); + if (done) { + controller.close(); + return; + } + controller.enqueue(value); + }, + async cancel(reason) { + await reader.cancel?.(reason); + }, + }); +} + +/** Whether a value can be drained with `for await`. */ +export function isAsyncIterable(value: unknown): value is AsyncIterable { + return !!value && typeof (value as AsyncIterable)[Symbol.asyncIterator] === 'function'; +} + /** * A minimal structural check for a web `ReadableStream`. * diff --git a/packages/server-utils/src/ai/langchain/utils.ts b/packages/server-utils/src/ai/langchain/utils.ts index 8ae219e42437..fcfc6c8fca25 100644 --- a/packages/server-utils/src/ai/langchain/utils.ts +++ b/packages/server-utils/src/ai/langchain/utils.ts @@ -505,7 +505,7 @@ export function extractLlmResponseAttributes( export function getAgentNameFromMetadata(metadata?: Record): Record { const attrs: Record = {}; - // lc_agent_name is injected by instrumentCompiledGraphInvoke (langgraph integration) + // lc_agent_name is injected by the langgraph invoke_agent wrapper (`instrumentCompiledGraph`) const agentName = metadata?.lc_agent_name; if (typeof agentName === 'string') { attrs[GEN_AI_AGENT_NAME] = agentName; diff --git a/packages/server-utils/src/ai/langgraph/index.ts b/packages/server-utils/src/ai/langgraph/index.ts index 0b1dc9a566d6..c8f2aa228175 100644 --- a/packages/server-utils/src/ai/langgraph/index.ts +++ b/packages/server-utils/src/ai/langgraph/index.ts @@ -1,22 +1,30 @@ /* eslint-disable typescript-eslint/no-deprecated */ -import { SPAN_STATUS_ERROR, startSpan, stringify } from '@sentry/core'; +import { getCurrentScope, SPAN_STATUS_ERROR, startSpan, startSpanManual, stringify } from '@sentry/core'; +import type { Span } from '@sentry/core'; import { GEN_AI_AGENT_NAME, GEN_AI_INPUT_MESSAGES, GEN_AI_OPERATION_NAME, GEN_AI_PIPELINE_NAME, GEN_AI_REQUEST_MODEL, + GEN_AI_RESPONSE_STREAMING, GEN_AI_SYSTEM_INSTRUCTIONS, GEN_AI_TOOL_DEFINITIONS, SENTRY_OP, SENTRY_ORIGIN, } from '@sentry/conventions/attributes'; import { GEN_AI_INVOKE_AGENT } from '@sentry/conventions/op'; -import { extractSystemInstructions, resolveAIRecordingOptions } from '../core/utils'; +import { extractSystemInstructions, isAsyncIterable, resolveAIRecordingOptions } from '../core/utils'; import { createLangChainCallbackHandler } from '../langchain'; import type { BaseChatModel, LangChainMessage } from '../langchain/types'; -import { normalizeLangChainMessages } from '../langchain/utils'; +import { + _INTERNAL_mergeLangChainCallbackHandler, + getConversationIdFromMetadata, + getConversationIdMetadataFromConfig, + normalizeLangChainMessages, +} from '../langchain/utils'; import { LANGGRAPH_ORIGIN } from './constants'; +import { instrumentStreamResult, isOutputStreamMode } from './streaming'; import type { CompiledGraph, LangGraphOptions } from './types'; import { extractAgentNameFromParams, @@ -25,19 +33,48 @@ import { setResponseAttributes, wrapToolsWithSpans, } from './utils'; -import { - getConversationIdMetadataFromConfig, - _INTERNAL_mergeLangChainCallbackHandler, - getConversationIdFromMetadata, -} from '../langchain/utils'; let _insideCreateReactAgent = false; +// Set while a manual wrapper runs the original sync `compile`/`createReactAgent`, so the integration's +// channel hooks (which fire inside that call) leave the graph to the options the user passed explicitly. +let _insideManualInstrumentation = false; + +/** Whether a manual `instrumentStateGraph()` / `instrumentCreateReactAgent()` call is in progress. */ +export function isInsideManualLangGraphInstrumentation(): boolean { + return _insideManualInstrumentation; +} + +function runManually(callback: () => T): T { + const previous = _insideManualInstrumentation; + _insideManualInstrumentation = true; + try { + return callback(); + } finally { + _insideManualInstrumentation = previous; + } +} const SENTRY_PATCHED = '__sentry_patched__'; +// invoke() drains this.stream() internally. Unlike Anthropic's sync flag, the marker lives on the span's +// scope, so it still holds if LangGraph ever awaits before that call. +const LANGGRAPH_INVOKE_ACTIVE = 'sentry_langgraph_invoke_active'; + +// One marker per graph, so a stream() of another graph inside invoke() still gets its own span. A graph +// already in here is instrumented: manual `instrumentStateGraph()` and the integration can both see it. +const INVOKE_MARKERS = new WeakMap(); + +function getInvokeMarker(graph: object): symbol { + let marker = INVOKE_MARKERS.get(graph); + if (!marker) { + marker = Symbol('sentry.langgraph.invoke'); + INVOKE_MARKERS.set(graph, marker); + } + return marker; +} /** - * Instruments StateGraph's compile method to wrap the returned compiled graph's invoke() with a - * `gen_ai.invoke_agent` span. + * Instruments StateGraph's compile method to wrap the returned compiled graph's invoke() and stream() + * methods with a `gen_ai.invoke_agent` span. */ export function instrumentStateGraphCompile( originalCompile: (...args: unknown[]) => CompiledGraph, @@ -59,21 +96,9 @@ export function instrumentStateGraphCompile( return Reflect.apply(target, thisArg, args); } - const compiledGraph = Reflect.apply(target, thisArg, args); + const compiledGraph = runManually(() => Reflect.apply(target, thisArg, args)); const compileOptions = args.length > 0 ? (args[0] as Record) : {}; - - // Instrument agent invoke method on the compiled graph - const originalInvoke = compiledGraph.invoke; - if (originalInvoke && typeof originalInvoke === 'function') { - compiledGraph.invoke = instrumentCompiledGraphInvoke( - originalInvoke.bind(compiledGraph) as (...args: unknown[]) => Promise, - compiledGraph, - compileOptions, - options, - undefined, - sentryHandler, - ); - } + instrumentCompiledGraph(compiledGraph, compileOptions, options, undefined, sentryHandler); return compiledGraph; }, @@ -84,116 +109,171 @@ export function instrumentStateGraphCompile( } /** - * Instruments CompiledGraph's invoke method to create spans for agent invocation - * - * Creates a `gen_ai.invoke_agent` span when invoke() is called + * Wraps a compiled graph's `invoke()` and `stream()` with `gen_ai.invoke_agent` spans, once per graph. */ -export function instrumentCompiledGraphInvoke( - originalInvoke: (...args: unknown[]) => Promise, +export function instrumentCompiledGraph( + graph: unknown, + compileOptions: Record, + options: LangGraphOptions, + llm: BaseChatModel | null | undefined, + sentryCallbackHandler: unknown, +): void { + if (!graph || typeof graph !== 'object' || INVOKE_MARKERS.has(graph)) { + return; + } + getInvokeMarker(graph); + + const compiledGraph = graph as CompiledGraph; + for (const method of ['invoke', 'stream'] as const) { + const original = compiledGraph[method]; + if (typeof original === 'function') { + (compiledGraph as Record)[method] = instrumentCompiledGraphOperation( + original.bind(compiledGraph) as (...args: unknown[]) => Promise, + compiledGraph, + compileOptions, + options, + llm, + sentryCallbackHandler, + method === 'stream', + ); + } + } +} + +function instrumentCompiledGraphOperation( + originalOperation: (...args: unknown[]) => Promise, graphInstance: CompiledGraph, compileOptions: Record, options: LangGraphOptions, - llm?: BaseChatModel | null, - sentryCallbackHandler?: unknown, + llm: BaseChatModel | null | undefined, + sentryCallbackHandler: unknown, + streaming: boolean, ): (...args: unknown[]) => Promise { - return new Proxy(originalInvoke, { + const invokeMarker = getInvokeMarker(graphInstance); + + return new Proxy(originalOperation, { apply(target, thisArg, args: unknown[]): Promise { + if ( + streaming && + getCurrentScope().getScopeData().sdkProcessingMetadata[LANGGRAPH_INVOKE_ACTIVE] === invokeMarker + ) { + // Only the first stream() is invoke's implementation detail; recursive calls need their own spans. + getCurrentScope().setSDKProcessingMetadata({ [LANGGRAPH_INVOKE_ACTIVE]: undefined }); + return Reflect.apply(target, thisArg, args); + } + const modelName = llm?.modelName ?? llm?.model; const config = args.length > 1 ? (args[1] as Record | undefined) : undefined; - return startSpan( - { - name: 'invoke_agent', - attributes: { - // Set before `spanStart`, so an id from `Sentry.setConversationId()` wins, as on the child spans - ...getConversationIdFromMetadata(config?.configurable as Record | undefined), - [SENTRY_ORIGIN]: LANGGRAPH_ORIGIN, - [SENTRY_OP]: GEN_AI_INVOKE_AGENT, - [GEN_AI_OPERATION_NAME]: 'invoke_agent', - }, + const spanOptions = { + name: 'invoke_agent', + attributes: { + // Set before `spanStart`, so an id from `Sentry.setConversationId()` wins, as on the child spans + ...getConversationIdFromMetadata(config?.configurable as Record | undefined), + [SENTRY_ORIGIN]: LANGGRAPH_ORIGIN, + [SENTRY_OP]: GEN_AI_INVOKE_AGENT, + [GEN_AI_OPERATION_NAME]: 'invoke_agent', }, - async span => { - try { - const graphName = compileOptions?.name; - - if (graphName && typeof graphName === 'string') { - span.setAttribute(GEN_AI_PIPELINE_NAME, graphName); - span.setAttribute(GEN_AI_AGENT_NAME, graphName); - span.updateName(`invoke_agent ${graphName}`); - } + }; + const run = async (span: Span): Promise => { + try { + const graphName = compileOptions?.name; + + if (graphName && typeof graphName === 'string') { + span.setAttribute(GEN_AI_PIPELINE_NAME, graphName); + span.setAttribute(GEN_AI_AGENT_NAME, graphName); + span.updateName(`invoke_agent ${graphName}`); + } - if (modelName) { - span.setAttribute(GEN_AI_REQUEST_MODEL, modelName); - } + if (modelName) { + span.setAttribute(GEN_AI_REQUEST_MODEL, modelName); + } - // Inject callback handler and agent name into invoke config - if (sentryCallbackHandler) { - const invokeConfig = (args[1] ?? {}) as Record; - args[1] = invokeConfig; - - const existingMetadata = (invokeConfig.metadata ?? {}) as Record; - invokeConfig.metadata = { - ...getConversationIdMetadataFromConfig(invokeConfig), - ...existingMetadata, - __sentry_langgraph__: true, - ...(typeof graphName === 'string' ? { lc_agent_name: graphName } : {}), - }; - - invokeConfig.callbacks = _INTERNAL_mergeLangChainCallbackHandler( - invokeConfig.callbacks, - sentryCallbackHandler, - ); - } + // Inject callback handler and agent name into invoke config + if (sentryCallbackHandler) { + const invokeConfig = (args[1] ?? {}) as Record; + args[1] = invokeConfig; + + const existingMetadata = (invokeConfig.metadata ?? {}) as Record; + invokeConfig.metadata = { + ...getConversationIdMetadataFromConfig(invokeConfig), + ...existingMetadata, + __sentry_langgraph__: true, + ...(typeof graphName === 'string' ? { lc_agent_name: graphName } : {}), + }; + + invokeConfig.callbacks = _INTERNAL_mergeLangChainCallbackHandler( + invokeConfig.callbacks, + sentryCallbackHandler, + ); + } - const recordInputs = options.recordInputs; - const recordOutputs = options.recordOutputs; + const recordInputs = options.recordInputs; + const recordOutputs = options.recordOutputs; - // Extract available tools from the graph instance - const tools = recordInputs ? extractToolsFromCompiledGraph(graphInstance) : null; - if (tools) { - span.setAttribute(GEN_AI_TOOL_DEFINITIONS, JSON.stringify(tools)); - } + // Extract available tools from the graph instance + const tools = recordInputs ? extractToolsFromCompiledGraph(graphInstance) : null; + if (tools) { + span.setAttribute(GEN_AI_TOOL_DEFINITIONS, JSON.stringify(tools)); + } - // Custom state annotations have no `messages` array, the whole state is recorded instead. - const inputState = args[0] as { messages?: LangChainMessage[]; lg_name?: string } | null | undefined; - const inputMessages = Array.isArray(inputState?.messages) ? inputState.messages : null; - // `new Command({ resume })` resumes an interrupted run and carries no user turn (LangGraph - // tags it `lg_name: 'Command'`), so skip it like a `null` resume rather than recording the - // control object as a message nobody wrote. - const isResumeCommand = inputState?.lg_name === 'Command'; - - if (recordInputs) { - if (inputMessages) { - const normalizedMessages = normalizeLangChainMessages(inputMessages); - const { systemInstructions, filteredMessages } = extractSystemInstructions(normalizedMessages); - - if (systemInstructions) { - span.setAttribute(GEN_AI_SYSTEM_INSTRUCTIONS, systemInstructions); - } - - span.setAttributes({ - [GEN_AI_INPUT_MESSAGES]: stringify(filteredMessages), - }); - } else if (inputState && typeof inputState === 'object' && !isResumeCommand) { - span.setAttribute(GEN_AI_INPUT_MESSAGES, stringify([{ role: 'user', content: stringify(inputState) }])); + // Custom state annotations have no `messages` array, the whole state is recorded instead. + const inputState = args[0] as { messages?: LangChainMessage[]; lg_name?: string } | null | undefined; + const inputMessages = Array.isArray(inputState?.messages) ? inputState.messages : null; + // `new Command({ resume })` resumes an interrupted run and carries no user turn (LangGraph + // tags it `lg_name: 'Command'`), so skip it like a `null` resume rather than recording the + // control object as a message nobody wrote. + const isResumeCommand = inputState?.lg_name === 'Command'; + + if (recordInputs) { + if (inputMessages) { + const normalizedMessages = normalizeLangChainMessages(inputMessages); + const { systemInstructions, filteredMessages } = extractSystemInstructions(normalizedMessages); + + if (systemInstructions) { + span.setAttribute(GEN_AI_SYSTEM_INSTRUCTIONS, systemInstructions); } + + span.setAttributes({ + [GEN_AI_INPUT_MESSAGES]: stringify(filteredMessages), + }); + } else if (inputState && typeof inputState === 'object' && !isResumeCommand) { + span.setAttribute(GEN_AI_INPUT_MESSAGES, stringify([{ role: 'user', content: stringify(inputState) }])); } + } - // Call original invoke - const result = await Reflect.apply(target, thisArg, args); + if (!streaming) { + getCurrentScope().setSDKProcessingMetadata({ [LANGGRAPH_INVOKE_ACTIVE]: invokeMarker }); + } + + const result = await Reflect.apply(target, thisArg, args); - if (recordOutputs) { - setResponseAttributes(span, inputMessages, result); + if (streaming) { + if (isAsyncIterable(result)) { + span.setAttribute(GEN_AI_RESPONSE_STREAMING, true); + const recordStreamOutputs = recordOutputs && isOutputStreamMode(config?.streamMode); + return instrumentStreamResult(result, span, inputMessages, recordStreamOutputs); } + span.end(); return result; - } catch (error) { - // The error is rethrown to the caller (invoke() rejects), so we only mark the span failed - // and do not record it. - span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); - throw error; } - }, - ); + + if (recordOutputs) { + setResponseAttributes(span, inputMessages, result); + } + + return result; + } catch (error) { + // The error is rethrown to the caller, so we only mark the span failed and do not record it. + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + if (streaming) { + span.end(); + } + throw error; + } + }; + + return streaming ? startSpanManual(spanOptions, run) : startSpan(spanOptions, run); }, }); } @@ -227,28 +307,13 @@ export function instrumentCreateReactAgent( _insideCreateReactAgent = true; let compiledGraph: CompiledGraph; try { - compiledGraph = Reflect.apply(target, thisArg, args); + compiledGraph = runManually(() => Reflect.apply(target, thisArg, args)); } finally { _insideCreateReactAgent = false; } - // Wrap invoke() on the returned compiled graph - const originalInvoke = compiledGraph.invoke; - if (originalInvoke && typeof originalInvoke === 'function') { - const compileOptions: Record = {}; - if (agentName) { - compileOptions.name = agentName; - } - - compiledGraph.invoke = instrumentCompiledGraphInvoke( - originalInvoke.bind(compiledGraph) as (...args: unknown[]) => Promise, - compiledGraph, - compileOptions, - resolvedOptions, - llm, - sentryHandler, - ); - } + const compileOptions: Record = agentName ? { name: agentName } : {}; + instrumentCompiledGraph(compiledGraph, compileOptions, resolvedOptions, llm, sentryHandler); return compiledGraph; }, diff --git a/packages/server-utils/src/ai/langgraph/streaming.ts b/packages/server-utils/src/ai/langgraph/streaming.ts new file mode 100644 index 000000000000..899e791bddb5 --- /dev/null +++ b/packages/server-utils/src/ai/langgraph/streaming.ts @@ -0,0 +1,344 @@ +import { handleCallbackErrors, SPAN_STATUS_ERROR, withActiveSpan, withScope } from '@sentry/core'; +import type { Scope, Span } from '@sentry/core'; +import { readableFromReader } from '../core/utils'; +import type { LangChainMessage } from '../langchain/types'; +import { setResponseAttributes } from './utils'; + +export function instrumentStreamResult>( + stream: T, + span: Span, + inputMessages: LangChainMessage[] | null, + recordOutputs: boolean | undefined, +): T { + // Captured once: entering it per chunk is cheaper than `withActiveSpan`, which forks a scope each call. + const spanScope = withActiveSpan(span, scope => scope); + const lifecycle = createStreamLifecycle(span, inputMessages, recordOutputs); + + try { + if (isReadableStream(stream)) { + instrumentReadableStream(stream, spanScope, lifecycle); + } else { + const iterate = stream[Symbol.asyncIterator].bind(stream); + stream[Symbol.asyncIterator] = () => instrumentStreamIterator(iterate(), spanScope, lifecycle); + } + } catch { + // An unpatchable stream (e.g. frozen) still belongs to the caller; only the span is lost. + lifecycle.complete(); + } + return stream; +} + +interface StreamLifecycle { + recordChunk: (chunk: unknown) => void; + complete: () => void; + fail: () => void; + markCancelled: (cancelled: boolean) => void; +} + +// Modes whose payloads carry the graph's output; `debug`, `tasks`, `checkpoints` and `custom` do not. +const OUTPUT_STREAM_MODES = new Set(['values', 'updates', 'messages']); + +interface StreamResponseState { + finalMessages?: LangChainMessage[]; + updateMessages: LangChainMessage[]; + messageChunks: LangChainMessage[]; +} + +interface ReadableStreamReaderLike { + read: (...args: unknown[]) => Promise>; + cancel?: (reason?: unknown) => Promise; +} + +interface InstrumentableReadableStream extends AsyncIterable { + locked: boolean; + getReader: (...args: unknown[]) => ReadableStreamReaderLike; + cancel?: (reason?: unknown) => Promise; + tee?: () => [ReadableStream, ReadableStream]; + pipeThrough?: ( + transform: ReadableWritablePair, + options?: StreamPipeOptions, + ) => ReadableStream; + pipeTo?: (destination: WritableStream, options?: StreamPipeOptions) => Promise; +} + +function createStreamLifecycle( + span: Span, + inputMessages: LangChainMessage[] | null, + recordOutputs: boolean | undefined, +): StreamLifecycle { + let state: StreamResponseState | undefined = recordOutputs ? { updateMessages: [], messageChunks: [] } : undefined; + let input = inputMessages; + let completed = false; + // LangGraph's cancel releases the reader under a pending read, which rejects; that is not a failed run. + let cancelled = false; + + const complete = (): void => { + if (completed) { + return; + } + + completed = true; + try { + if (state) { + setResponseAttributes(span, input, getStreamResponseResult(state, input)); + } + } catch { + // Recording failures must not replace the application's stream outcome. + } finally { + // The caller may hold the stream long after it ends; don't keep every chunk alive with it. + state = undefined; + input = null; + span.end(); + } + }; + + return { + recordChunk(chunk: unknown): void { + try { + if (state) { + accumulateStreamResponse(state, chunk); + } + } catch { + // Same as `complete`: recording must never break the application's stream. + } + }, + complete, + fail(): void { + if (!cancelled && !completed) { + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + } + complete(); + }, + markCancelled(value: boolean): void { + cancelled = value; + }, + }; +} + +// Multi-mode and `subgraphs` streams wrap payloads as `[mode, payload]`, `[namespace, payload]` or +// `[namespace, mode, payload]`. Only the root graph's output modes describe this span's response. +function getOutputPayload(chunk: unknown): { payload: unknown } | undefined { + if (!Array.isArray(chunk) || (typeof chunk[0] !== 'string' && !Array.isArray(chunk[0]))) { + return { payload: chunk }; + } + + const namespace = Array.isArray(chunk[0]) ? (chunk[0] as unknown[]) : []; + const mode = chunk.length === 3 ? chunk[1] : chunk[0]; + if (namespace.length > 0 || !isOutputStreamMode(mode)) { + return undefined; + } + + return { payload: chunk[chunk.length - 1] }; +} + +/** False for a `streamMode` whose payloads are not the graph's output; non-strings are chunk shapes. */ +export function isOutputStreamMode(mode: unknown): boolean { + return typeof mode !== 'string' || OUTPUT_STREAM_MODES.has(mode); +} + +function accumulateStreamResponse(state: StreamResponseState, chunk: unknown): void { + const payload = getOutputPayload(chunk)?.payload; + // `encoding: 'text/event-stream'` yields bytes; walking them would cost a lookup per byte for nothing. + if (!payload || typeof payload !== 'object' || ArrayBuffer.isView(payload)) { + return; + } + + if ( + Array.isArray(payload) && + payload.length === 2 && + payload[0] && + typeof payload[0] === 'object' && + 'content' in payload[0] + ) { + const message = payload[0] as LangChainMessage; + const previousIndex = + typeof message.id === 'string' ? state.messageChunks.findIndex(previous => previous.id === message.id) : -1; + const previous = state.messageChunks[previousIndex]; + if (previous && typeof previous.concat === 'function') { + state.messageChunks[previousIndex] = previous.concat(message) as LangChainMessage; + } else if (previous && typeof previous.content === 'string' && typeof message.content === 'string') { + state.messageChunks[previousIndex] = { ...previous, ...message, content: previous.content + message.content }; + } else { + state.messageChunks.push(message); + } + return; + } + const directMessages = getMessages(payload); + if (directMessages) { + state.finalMessages = directMessages; + return; + } + + for (const update of Object.values(payload)) { + state.updateMessages.push(...(getMessages(update) ?? [])); + } +} + +function getMessages(value: unknown): LangChainMessage[] | undefined { + if (!value || typeof value !== 'object' || !('messages' in value) || !Array.isArray(value.messages)) { + return undefined; + } + + return value.messages as LangChainMessage[]; +} + +function getStreamResponseResult( + state: StreamResponseState, + inputMessages: LangChainMessage[] | null, +): { messages: LangChainMessage[] } | undefined { + if (state.finalMessages) { + return { messages: state.finalMessages }; + } + + // Updates contain complete messages, including the same responses emitted as token chunks. + const messages = state.updateMessages.length > 0 ? state.updateMessages : state.messageChunks; + if (messages.length === 0) { + return undefined; + } + + return { messages: [...(inputMessages ?? []), ...messages] }; +} + +function isReadableStream(stream: AsyncIterable): stream is InstrumentableReadableStream { + return typeof (stream as Partial).getReader === 'function'; +} + +function instrumentReadableStream( + stream: InstrumentableReadableStream, + scope: Scope, + lifecycle: StreamLifecycle, +): void { + const originalGetReader = stream.getReader.bind(stream); + stream.getReader = (...args: unknown[]): ReadableStreamReaderLike => { + const reader = withScope(scope, () => originalGetReader(...args)); + return instrumentReader(reader, scope, lifecycle); + }; + + if (stream.cancel) { + const originalCancel = stream.cancel.bind(stream); + stream.cancel = (reason?: unknown): Promise => { + // Marked up front because LangGraph's cancel can reject an in-flight read before the cancel settles. + lifecycle.markCancelled(true); + // A plain locked stream rejects the cancel and keeps streaming, so only a cancel that took effect ends it. + const locked = stream.locked; + return handleCallbackErrors( + () => withScope(scope, () => originalCancel(reason)), + () => { + lifecycle.markCancelled(false); + if (!locked) { + lifecycle.fail(); + } + }, + undefined, + () => lifecycle.complete(), + ); + }; + } + + // Native piping reads through internal slots, bypassing the patched `getReader`. The checks mirror + // native piping's own, made before the reader is taken so a rejected pipe leaves the source unlocked. + const instrumentedSource = ( + destination: WritableStream | undefined, + options: StreamPipeOptions | undefined, + hasReadable = true, + ): ReadableStream => { + const signal = options?.signal; + if ( + stream.locked || + typeof destination?.getWriter !== 'function' || + destination.locked || + !hasReadable || + (signal !== undefined && !(signal instanceof AbortSignal)) + ) { + throw new TypeError('Cannot pipe a locked stream, or with a missing or locked destination or an invalid signal.'); + } + return readableFromReader(stream.getReader()); + }; + + if (stream.tee) { + stream.tee = () => readableFromReader(stream.getReader()).tee(); + } + + if (stream.pipeTo) { + stream.pipeTo = async (destination: WritableStream, options?: StreamPipeOptions): Promise => + instrumentedSource(destination, options).pipeTo(destination, options); + } + + if (stream.pipeThrough) { + stream.pipeThrough = ( + transform: ReadableWritablePair, + options?: StreamPipeOptions, + ): ReadableStream => + instrumentedSource(transform?.writable, options, !!transform?.readable).pipeThrough(transform, options); + } +} + +function instrumentReader( + reader: ReadableStreamReaderLike, + scope: Scope, + lifecycle: StreamLifecycle, +): ReadableStreamReaderLike { + const originalRead = reader.read.bind(reader); + reader.read = (...args: unknown[]): Promise> => + handleCallbackErrors( + () => withScope(scope, () => originalRead(...args)), + () => lifecycle.fail(), + undefined, + (result: ReadableStreamReadResult) => + result.done ? lifecycle.complete() : lifecycle.recordChunk(result.value), + ); + + if (reader.cancel) { + const originalCancel = reader.cancel.bind(reader); + reader.cancel = (reason?: unknown): Promise => { + lifecycle.markCancelled(true); + return handleCallbackErrors( + () => withScope(scope, () => originalCancel(reason)), + () => { + lifecycle.markCancelled(false); + lifecycle.fail(); + }, + () => lifecycle.complete(), + ); + }; + } + + return reader; +} + +async function* instrumentStreamIterator( + iterator: AsyncIterator, + scope: Scope, + lifecycle: StreamLifecycle, +): AsyncGenerator { + // Like `for await`, only close the source when the consumer stops early, not after it threw or finished. + let sourceDone = false; + + try { + while (true) { + let result: IteratorResult; + try { + result = await withScope(scope, () => iterator.next()); + } catch (error) { + sourceDone = true; + lifecycle.fail(); + throw error; + } + if (result.done) { + sourceDone = true; + lifecycle.complete(); + return; + } + lifecycle.recordChunk(result.value); + yield result.value; + } + } finally { + try { + if (!sourceDone) { + await withScope(scope, () => iterator.return?.()); + } + } finally { + lifecycle.complete(); + } + } +} diff --git a/packages/server-utils/src/ai/langgraph/types.ts b/packages/server-utils/src/ai/langgraph/types.ts index 153fa9b4caba..ff2d62cf38c2 100644 --- a/packages/server-utils/src/ai/langgraph/types.ts +++ b/packages/server-utils/src/ai/langgraph/types.ts @@ -61,6 +61,7 @@ export interface StateGraphBuilder { export interface CompiledGraph { [key: string]: unknown; invoke?: (...args: unknown[]) => Promise; + stream?: (...args: unknown[]) => Promise>; name?: string; graph_name?: string; lc_kwargs?: { diff --git a/packages/server-utils/src/ai/mistral/streaming.ts b/packages/server-utils/src/ai/mistral/streaming.ts index 74acc6e43b70..5907b8445265 100644 --- a/packages/server-utils/src/ai/mistral/streaming.ts +++ b/packages/server-utils/src/ai/mistral/streaming.ts @@ -1,6 +1,6 @@ import type { Span } from '@sentry/core'; import { SPAN_STATUS_ERROR } from '@sentry/core'; -import { endStreamSpan, setOutputMessagesAttribute } from '../core/utils'; +import { endStreamSpan, isAsyncIterable, readableFromReader, setOutputMessagesAttribute } from '../core/utils'; import type { MistralCompletionChunk, MistralToolCall } from './types'; import { contentToString } from './utils'; @@ -22,8 +22,6 @@ interface StreamingState { /** Which drain path owns accumulation. See {@link instrumentEventStream}. */ type StreamConsumer = 'iterator' | 'reader'; -type AsyncIterableStream = { [Symbol.asyncIterator]: () => AsyncIterator }; - interface StreamReaderLike { read: () => Promise<{ done: boolean; value?: unknown }>; cancel?: (reason?: unknown) => Promise; @@ -37,11 +35,6 @@ interface ReadableStreamLike { pipeThrough?: (transform: unknown, options?: unknown) => unknown; } -/** Whether a value can be drained with `for await`. */ -export function isAsyncIterable(value: unknown): value is AsyncIterableStream { - return !!value && typeof (value as AsyncIterableStream)[Symbol.asyncIterator] === 'function'; -} - function createStreamingState(): StreamingState { return { responseTexts: [], @@ -187,38 +180,6 @@ function wrapReader( }); } -/** - * A stream that pulls through `readable`'s instrumented reader. - * - * `tee`, `pipeTo` and `pipeThrough` acquire their reader through internal slots rather than by - * calling the public `getReader`, so patching that method alone leaves them uninstrumented: the - * chunks bypass accumulation and the span is never ended. Handing them this stream instead routes - * them back through the wrapped reader, so there is still exactly one accumulating consumer. - * - * Same idea as `monitorStream` in `@sentry/deno`, with two deliberate differences: chunks are pulled - * on demand rather than drained in a `start` loop, so backpressure still reaches the source, and - * `cancel` is forwarded (the hole #24054 fixed there) so the source stops producing when the - * consumer disconnects. - */ -function instrumentedSource(readable: ReadableStreamLike): ReadableStream { - // The patched `getReader`, so reads are accumulated and the span is ended by the shared logic. - const reader = readable.getReader!(); - - return new ReadableStream({ - async pull(controller) { - const { done, value } = await reader.read(); - if (done) { - controller.close(); - return; - } - controller.enqueue(value); - }, - async cancel(reason) { - await reader.cancel?.(reason); - }, - }); -} - /** * Instrument a Mistral event stream in place: accumulate response attributes as it is drained and * end `span` when it finishes. @@ -305,20 +266,23 @@ export function instrumentEventStream(stream: unknown, span: Span, recordOutputs }; } - // Only patched when `getReader` is present, since that is what `instrumentedSource` pulls through. + // Only patched when `getReader` is present, since that is what `readableFromReader` pulls through. if (typeof readable.getReader === 'function') { if (typeof readable.tee === 'function') { - readable.tee = () => instrumentedSource(readable).tee(); + readable.tee = () => readableFromReader(readable.getReader!()).tee(); } if (typeof readable.pipeTo === 'function') { readable.pipeTo = (destination: unknown, options?: unknown) => - instrumentedSource(readable).pipeTo(destination as WritableStream, options as StreamPipeOptions); + readableFromReader(readable.getReader!()).pipeTo( + destination as WritableStream, + options as StreamPipeOptions, + ); } if (typeof readable.pipeThrough === 'function') { readable.pipeThrough = (transform: unknown, options?: unknown) => - instrumentedSource(readable).pipeThrough( + readableFromReader(readable.getReader!()).pipeThrough( transform as ReadableWritablePair, options as StreamPipeOptions, ); diff --git a/packages/server-utils/src/integrations/langgraph.ts b/packages/server-utils/src/integrations/langgraph.ts index d9d5bd863144..3b7dafd7d3cb 100644 --- a/packages/server-utils/src/integrations/langgraph.ts +++ b/packages/server-utils/src/integrations/langgraph.ts @@ -3,9 +3,9 @@ import type { IntegrationFn } from '@sentry/core'; import { debug, defineIntegration, isObjectLike } from '@sentry/core'; import { resolveAIRecordingOptions } from '../ai/core/utils'; import { createLangChainCallbackHandler } from '../ai/langchain'; -import { instrumentCompiledGraphInvoke } from '../ai/langgraph'; +import { instrumentCompiledGraph, isInsideManualLangGraphInstrumentation } from '../ai/langgraph'; import { LANGGRAPH_INTEGRATION_NAME } from '../ai/langgraph/constants'; -import type { CompiledGraph, LangGraphOptions } from '../ai/langgraph/types'; +import type { LangGraphOptions } from '../ai/langgraph/types'; import { extractAgentNameFromParams, extractLLMFromParams, wrapToolsWithSpans } from '../ai/langgraph/utils'; import { DEBUG_BUILD } from '../debug-build'; import type { OrchestrionChannelContext } from '../orchestrion/types'; @@ -38,19 +38,19 @@ function instrumentLanggraph(options: LangGraphOptions): void { const resolvedOptions = resolveAIRecordingOptions(options); const sentryHandler = createLangChainCallbackHandler(resolvedOptions); - // StateGraph.compile returns synchronously; wrap the returned graph's `invoke` at `end`. + // StateGraph.compile returns synchronously; wrap the returned graph's agent methods at `end`. diagnosticsChannel .tracingChannel(CHANNELS.LANGGRAPH_STATE_GRAPH_COMPILE) .end.subscribe(message => { - if (insideCreateReactAgent) { + if (insideCreateReactAgent || isInsideManualLangGraphInstrumentation()) { return; } const { arguments: args, result } = message as OrchestrionChannelContext; - wrapCompiledGraphInvoke(result, getFirstArgObject(args) ?? {}, resolvedOptions, null, sentryHandler); + instrumentCompiledGraph(result, getFirstArgObject(args) ?? {}, resolvedOptions, null, sentryHandler); }); - // createReactAgent only wraps tools and the returned graph's `invoke`. Tools are wrapped at - // `start` (before the agent runs), invoke at `end`. + // createReactAgent only wraps tools and the returned graph's agent methods. Tools are wrapped at + // `start` (before the agent runs), agent methods at `end`. const reactAgentChannel = diagnosticsChannel.tracingChannel( CHANNELS.LANGGRAPH_CREATE_REACT_AGENT, ); @@ -72,10 +72,13 @@ function instrumentLanggraph(options: LangGraphOptions): void { }); reactAgentChannel.end.subscribe(message => { insideCreateReactAgent = false; + if (isInsideManualLangGraphInstrumentation()) { + return; + } const { arguments: args, result } = message as OrchestrionChannelContext; const agentName = extractAgentNameFromParams(args) ?? undefined; const compileOptions = agentName ? { name: agentName } : {}; - wrapCompiledGraphInvoke(result, compileOptions, resolvedOptions, extractLLMFromParams(args), sentryHandler); + instrumentCompiledGraph(result, compileOptions, resolvedOptions, extractLLMFromParams(args), sentryHandler); }); // Make sure a thrown `createReactAgent` doesn't leave the suppression flag stuck on. reactAgentChannel.error.subscribe(() => { @@ -89,35 +92,6 @@ function getFirstArgObject(args: unknown[] | undefined): Record return isObjectLike(first) ? first : undefined; } -/** - * Wrap the compiled graph's `invoke` with the shared `invoke_agent` instrumentation, exactly as the - * OTel path does on the returned graph. - */ -function wrapCompiledGraphInvoke( - graph: unknown, - compileOptions: Record, - options: LangGraphOptions, - llm: ReturnType, - sentryHandler: unknown, -): void { - if (!graph || typeof graph !== 'object') { - return; - } - - const compiledGraph = graph as CompiledGraph; - const originalInvoke = compiledGraph.invoke; - if (typeof originalInvoke === 'function') { - compiledGraph.invoke = instrumentCompiledGraphInvoke( - originalInvoke.bind(compiledGraph), - compiledGraph, - compileOptions, - options, - llm, - sentryHandler, - ); - } -} - /** * Diagnostics-channel-based LangGraph integration. Subscribes to the diagnostics_channels * injected into `@langchain/langgraph`'s `StateGraph.compile` and `createReactAgent`, so it requires diff --git a/packages/server-utils/test/ai/lib/tracing/langgraph-invoke-agent.test.ts b/packages/server-utils/test/ai/lib/tracing/langgraph-invoke-agent.test.ts index 3f894efae3ea..94b1cfe9d876 100644 --- a/packages/server-utils/test/ai/lib/tracing/langgraph-invoke-agent.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/langgraph-invoke-agent.test.ts @@ -8,12 +8,25 @@ import { startSpan, } from '@sentry/core'; import type { Span } from '@sentry/core'; -import { GEN_AI_CONVERSATION_ID } from '@sentry/conventions/attributes'; -import { instrumentCompiledGraphInvoke } from '../../../../src/ai/langgraph'; -import type { CompiledGraph } from '../../../../src/ai/langgraph/types'; +import { GEN_AI_CONVERSATION_ID, GEN_AI_RESPONSE_TEXT } from '@sentry/conventions/attributes'; +import { + instrumentCompiledGraph, + instrumentStateGraphCompile, + isInsideManualLangGraphInstrumentation, +} from '../../../../src/ai/langgraph'; import { wrapToolsWithSpans } from '../../../../src/ai/langgraph/utils'; import { getDefaultTestClientOptions, TestClient } from '../../../mocks/client'; +function instrumentInvoke( + invoke: (...args: unknown[]) => Promise, + compileOptions: Record = {}, + sentryCallbackHandler?: unknown, +): (...args: unknown[]) => Promise { + const graph = { invoke }; + instrumentCompiledGraph(graph, compileOptions, {}, undefined, sentryCallbackHandler); + return graph.invoke; +} + describe('LangGraph invoke_agent span names', () => { beforeEach(() => { getMainCarrier().__SENTRY__ = undefined; @@ -41,12 +54,7 @@ describe('LangGraph invoke_agent span names', () => { } async function invokeGraph(compileOptions: Record, config?: Record): Promise { - const invoke = instrumentCompiledGraphInvoke( - async () => ({ messages: [] }), - {} as CompiledGraph, - compileOptions, - {}, - ); + const invoke = instrumentInvoke(async () => ({ messages: [] }), compileOptions); await invoke({}, config); } @@ -92,7 +100,7 @@ describe('LangGraph conversation id', () => { async function invokeGraph(config: Record): Promise { const endedSpans = setupClient(); - const invoke = instrumentCompiledGraphInvoke(async () => ({ messages: [] }), {} as CompiledGraph, {}, {}); + const invoke = instrumentInvoke(async () => ({ messages: [] })); await invoke({}, config); return endedSpans[0]!; } @@ -113,16 +121,13 @@ describe('LangGraph conversation id', () => { it('copies the conversation id from configurable into the metadata the graph runs with', async () => { setupClient(); let receivedConfig: Record | undefined; - const invoke = instrumentCompiledGraphInvoke( + const invoke = instrumentInvoke( async (_input: unknown, config: unknown) => { receivedConfig = config as Record; return { messages: [] }; }, - {} as CompiledGraph, {}, {}, - undefined, - {}, ); await invoke({}, { configurable: { sessionId: 'session-1' }, metadata: { user: 'u' } }); @@ -141,3 +146,140 @@ describe('LangGraph conversation id', () => { expect(spanToStaticSpanJSON(toolSpan).data[GEN_AI_CONVERSATION_ID]).toBe('thread-1'); }); }); + +describe('LangGraph invoke stream suppression', () => { + beforeEach(() => { + getMainCarrier().__SENTRY__ = undefined; + }); + + afterEach(() => { + getMainCarrier().__SENTRY__ = undefined; + }); + + it('instruments a recursive stream after suppressing the internal invoke stream', async () => { + const client = new TestClient(getDefaultTestClientOptions({ tracesSampleRate: 1 })); + setCurrentClient(client); + client.init(); + const startedSpans: Span[] = []; + client.on('spanStart', span => startedSpans.push(span)); + let nestedStream: Promise> | undefined; + const graph = { + async invoke() { + return this.stream('internal'); + }, + async stream(input: string): Promise> { + if (input === 'internal') { + nestedStream = this.stream('recursive'); + } + return { + async *[Symbol.asyncIterator]() { + yield input; + }, + }; + }, + }; + instrumentStateGraphCompile(() => graph, {})(); + + const invokedStream = await graph.invoke(); + const recursiveStream = await nestedStream!; + const chunks = []; + for await (const chunk of invokedStream) { + chunks.push(chunk); + } + for await (const chunk of recursiveStream) { + chunks.push(chunk); + } + + expect(chunks).toEqual(['internal', 'recursive']); + expect(startedSpans).toHaveLength(2); + expect(startedSpans.map(span => spanToStaticSpanJSON(span).description)).toEqual(['invoke_agent', 'invoke_agent']); + expect(spanToStaticSpanJSON(startedSpans[1]!).parent_span_id).toBe(startedSpans[0]!.spanContext().spanId); + }); + + it('creates one agent span per run when manual and auto instrumentation both wrap the graph', async () => { + const client = new TestClient(getDefaultTestClientOptions({ tracesSampleRate: 1 })); + setCurrentClient(client); + client.init(); + const startedSpans: Span[] = []; + client.on('spanStart', span => startedSpans.push(span)); + const graph = { + async invoke() { + const stream = await this.stream(); + for await (const _chunk of stream) { + // drain, as LangGraph's invoke does + } + }, + async stream(): Promise> { + return { + async *[Symbol.asyncIterator]() { + yield 'chunk'; + }, + }; + }, + }; + instrumentStateGraphCompile( + instrumentStateGraphCompile(() => graph, {}), + {}, + )(); + instrumentCompiledGraph(graph, {}, {}, undefined, undefined); + + await graph.invoke(); + for await (const _chunk of await graph.stream()) { + // drain + } + + expect(startedSpans).toHaveLength(2); + }); + + it('returns a stream it cannot patch to the caller instead of throwing', async () => { + const client = new TestClient(getDefaultTestClientOptions({ tracesSampleRate: 1 })); + setCurrentClient(client); + client.init(); + const endedSpans: Span[] = []; + client.on('spanEnd', span => endedSpans.push(span)); + const frozenStream = Object.freeze({ + async *[Symbol.asyncIterator]() { + yield 'chunk'; + }, + }); + const graph = { stream: async () => frozenStream }; + instrumentCompiledGraph(graph, {}, {}, undefined, undefined); + + await expect(graph.stream()).resolves.toBe(frozenStream); + expect(endedSpans).toHaveLength(1); + }); + + it('leaves a graph compiled by manual instrumentation to the manual options', () => { + let sawManualFlag = false; + instrumentStateGraphCompile(() => { + sawManualFlag = isInsideManualLangGraphInstrumentation(); + return {}; + }, {})(); + + expect(sawManualFlag).toBe(true); + expect(isInsideManualLangGraphInstrumentation()).toBe(false); + }); + + it('records no output for a stream mode whose chunks are not the graph output', async () => { + const client = new TestClient(getDefaultTestClientOptions({ tracesSampleRate: 1 })); + setCurrentClient(client); + client.init(); + const endedSpans: Span[] = []; + client.on('spanEnd', span => endedSpans.push(span)); + const graph = { + stream: async () => ({ + async *[Symbol.asyncIterator]() { + yield { values: { messages: [{ role: 'assistant', content: 'checkpointed' }] } }; + }, + }), + }; + instrumentCompiledGraph(graph, {}, { recordOutputs: true }, undefined, undefined); + + for await (const _chunk of await graph.stream({}, { streamMode: 'checkpoints' })) { + // drain + } + + expect(endedSpans).toHaveLength(1); + expect(spanToStaticSpanJSON(endedSpans[0]!).data[GEN_AI_RESPONSE_TEXT]).toBeUndefined(); + }); +}); diff --git a/packages/server-utils/test/ai/lib/tracing/langgraph-stream-lifecycle.test.ts b/packages/server-utils/test/ai/lib/tracing/langgraph-stream-lifecycle.test.ts new file mode 100644 index 000000000000..47c09dc08c53 --- /dev/null +++ b/packages/server-utils/test/ai/lib/tracing/langgraph-stream-lifecycle.test.ts @@ -0,0 +1,135 @@ +import { describe, expect, it, vi } from 'vitest'; +import { SPAN_STATUS_ERROR } from '@sentry/core'; +import type { Span } from '@sentry/core'; +import { instrumentStreamResult } from '../../../../src/ai/langgraph/streaming'; + +function createSpan() { + return { + end: vi.fn(), + setAttribute: vi.fn(), + setStatus: vi.fn(), + spanContext: () => ({ traceId: 'a'.repeat(32), spanId: 'b'.repeat(16), traceFlags: 1 }), + }; +} + +describe('LangGraph response recording failures', () => { + it.each([false, true])('preserves the stream outcome when serialization fails (stream error: %s)', async fails => { + const span = createSpan(); + const args: Record = {}; + args.self = args; + const chunk = { + agent: { messages: [{ role: 'assistant', content: '', tool_calls: [{ name: 'weather', args }] }] }, + }; + const error = new Error('Graph execution failed'); + const stream = instrumentStreamResult( + { + async *[Symbol.asyncIterator]() { + yield chunk; + if (fails) { + throw error; + } + }, + }, + span as unknown as Span, + null, + true, + ); + const iterator = stream[Symbol.asyncIterator](); + expect(await iterator.next()).toEqual({ done: false, value: chunk }); + + if (fails) { + await expect(iterator.next()).rejects.toBe(error); + expect(span.setStatus).toHaveBeenCalledExactlyOnceWith({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + } else { + await expect(iterator.next()).resolves.toEqual({ done: true, value: undefined }); + expect(span.setStatus).not.toHaveBeenCalled(); + } + expect(span.end).toHaveBeenCalledTimes(1); + }); + + it('keeps streaming when recording a chunk throws', async () => { + const span = createSpan(); + const message = { + id: 'msg-1', + content: 'Sun', + concat: () => { + throw new Error('concat failed'); + }, + }; + const chunks = [ + ['messages', [message, {}]], + ['messages', [{ ...message, content: 'ny' }, {}]], + ]; + const stream = instrumentStreamResult( + { + async *[Symbol.asyncIterator]() { + yield* chunks; + }, + }, + span as unknown as Span, + null, + true, + ); + + const received: unknown[] = []; + for await (const chunk of stream) { + received.push(chunk); + } + + expect(received).toEqual(chunks); + expect(span.setStatus).not.toHaveBeenCalled(); + expect(span.end).toHaveBeenCalledTimes(1); + }); +}); + +describe('LangGraph source iterator closing', () => { + it('surfaces the iteration error without closing a source that already threw', async () => { + const span = createSpan(); + const error = new Error('Graph execution failed'); + const returnSpy = vi.fn(async () => { + throw new Error('return failed'); + }); + const stream = instrumentStreamResult( + { + [Symbol.asyncIterator]: () => ({ + next: async () => { + throw error; + }, + return: returnSpy, + }), + }, + span as unknown as Span, + null, + false, + ); + + await expect(stream[Symbol.asyncIterator]().next()).rejects.toBe(error); + expect(returnSpy).not.toHaveBeenCalled(); + expect(span.setStatus).toHaveBeenCalledExactlyOnceWith({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + expect(span.end).toHaveBeenCalledTimes(1); + }); + + it('closes the source when the consumer stops early', async () => { + const span = createSpan(); + const returnSpy = vi.fn(async () => ({ done: true as const, value: undefined })); + const stream = instrumentStreamResult( + { + [Symbol.asyncIterator]: () => ({ + next: async () => ({ done: false, value: 'chunk' }), + return: returnSpy, + }), + }, + span as unknown as Span, + null, + false, + ); + + for await (const _chunk of stream) { + break; + } + + expect(returnSpy).toHaveBeenCalledTimes(1); + expect(span.setStatus).not.toHaveBeenCalled(); + expect(span.end).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/server-utils/test/ai/lib/tracing/langgraph-stream.test.ts b/packages/server-utils/test/ai/lib/tracing/langgraph-stream.test.ts new file mode 100644 index 000000000000..a5983ebdc973 --- /dev/null +++ b/packages/server-utils/test/ai/lib/tracing/langgraph-stream.test.ts @@ -0,0 +1,408 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type * as SentryCore from '@sentry/core'; +import { SPAN_STATUS_ERROR } from '@sentry/core'; +import type { Span } from '@sentry/core'; +import { + GEN_AI_RESPONSE_FINISH_REASONS, + GEN_AI_RESPONSE_MODEL, + GEN_AI_RESPONSE_STREAMING, + GEN_AI_RESPONSE_TEXT, + GEN_AI_USAGE_INPUT_TOKENS, + GEN_AI_USAGE_OUTPUT_TOKENS, + GEN_AI_USAGE_TOTAL_TOKENS, +} from '@sentry/conventions/attributes'; +import { instrumentStateGraphCompile } from '../../../../src/ai/langgraph'; + +const spanEnd = vi.fn(); +const spanSetAttribute = vi.fn(); +const spanSetStatus = vi.fn(); +const startSpanManual = vi.fn(); + +class TestLangGraphStream extends ReadableStream implements AsyncIterableIterator { + private iteratorReader: ReadableStreamDefaultReader | undefined; + + public constructor(chunks: T[]) { + const queuedChunks = [...chunks]; + super({ + pull(controller) { + const chunk = queuedChunks.shift(); + if (chunk === undefined) { + controller.close(); + } else { + controller.enqueue(chunk); + } + }, + }); + } + + public async next(): Promise> { + this.iteratorReader ??= this.getReader(); + const result = await this.iteratorReader.read(); + if (result.done) { + this.iteratorReader.releaseLock(); + } + return result; + } + + public async return(): Promise> { + if (this.iteratorReader) { + await this.iteratorReader.cancel(); + this.iteratorReader.releaseLock(); + } + return { done: true, value: undefined }; + } + + public [Symbol.asyncIterator](): this { + return this; + } +} + +vi.mock('@sentry/core', async importOriginal => { + const actual = (await importOriginal()) as typeof SentryCore; + return { + ...actual, + startSpanManual: ( + options: SentryCore.StartSpanOptions, + callback: (span: Span, finish: () => void) => unknown, + ): unknown => { + const span = { + end: spanEnd, + isRecording: () => true, + setAttribute: spanSetAttribute, + setAttributes: vi.fn(), + setStatus: spanSetStatus, + updateName: vi.fn(), + } as unknown as Span; + + startSpanManual(options, callback); + return callback(span, spanEnd); + }, + }; +}); + +describe('instrumentStateGraphCompile stream instrumentation', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('keeps the span open until the stream is fully consumed', async () => { + const stream = { + async *[Symbol.asyncIterator]() { + yield { agent: { messages: ['first update'] } }; + yield { agent: { messages: ['final update'] } }; + }, + }; + const compiledGraph = { stream: vi.fn().mockResolvedValue(stream) }; + const compile = instrumentStateGraphCompile(() => compiledGraph, {}); + + const graph = compile({ name: 'weather_assistant' }) as { + stream: (input: unknown) => Promise>; + }; + const instrumentedStream = await graph.stream({ messages: ['What is the weather?'] }); + + expect(instrumentedStream).toBe(stream); + expect(spanEnd).not.toHaveBeenCalled(); + + const chunks = []; + for await (const chunk of instrumentedStream) { + chunks.push(chunk); + } + + expect(chunks).toEqual([{ agent: { messages: ['first update'] } }, { agent: { messages: ['final update'] } }]); + expect(startSpanManual).toHaveBeenCalledTimes(1); + expect(spanSetAttribute).toHaveBeenCalledWith(GEN_AI_RESPONSE_STREAMING, true); + expect(spanEnd).toHaveBeenCalledTimes(1); + }); + + it('ends the span and the underlying iterator when consumption stops early', async () => { + const streamCleanup = vi.fn(); + const stream = { + async *[Symbol.asyncIterator]() { + try { + yield 'first update'; + yield 'final update'; + } finally { + streamCleanup(); + } + }, + }; + const compiledGraph = { stream: vi.fn().mockResolvedValue(stream) }; + const compile = instrumentStateGraphCompile(() => compiledGraph, {}); + const graph = compile() as { stream: () => Promise> }; + + for await (const _chunk of await graph.stream()) { + break; + } + + expect(streamCleanup).toHaveBeenCalledTimes(1); + expect(spanEnd).toHaveBeenCalledTimes(1); + }); + + it('marks the span as failed when stream iteration throws', async () => { + const error = new Error('stream failed'); + const stream = { + [Symbol.asyncIterator]() { + return { + next: vi.fn().mockRejectedValue(error), + }; + }, + }; + const compiledGraph = { stream: vi.fn().mockResolvedValue(stream) }; + const compile = instrumentStateGraphCompile(() => compiledGraph, {}); + const graph = compile() as { stream: () => Promise> }; + + await expect(async () => { + for await (const _chunk of await graph.stream()) { + // The iterator throws before yielding. + } + }).rejects.toThrow(error); + + expect(spanSetStatus).toHaveBeenCalledWith({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + expect(spanEnd).toHaveBeenCalledTimes(1); + }); + + it('ends the span when a reader consumes the stream', async () => { + const stream = new TestLangGraphStream(['first update', 'final update']); + const compiledGraph = { stream: vi.fn().mockResolvedValue(stream) }; + const compile = instrumentStateGraphCompile(() => compiledGraph, {}); + const graph = compile() as { stream: () => Promise> }; + + const reader = (await graph.stream()).getReader(); + expect(await reader.read()).toEqual({ done: false, value: 'first update' }); + expect(await reader.read()).toEqual({ done: false, value: 'final update' }); + expect(spanEnd).not.toHaveBeenCalled(); + expect(await reader.read()).toEqual({ done: true, value: undefined }); + + expect(spanEnd).toHaveBeenCalledTimes(1); + }); + + it('ends the span when pipeTo consumes the stream', async () => { + const chunks: string[] = []; + const stream = new TestLangGraphStream(['first update', 'final update']); + const compiledGraph = { stream: vi.fn().mockResolvedValue(stream) }; + const compile = instrumentStateGraphCompile(() => compiledGraph, {}); + const graph = compile() as { stream: () => Promise> }; + + await ( + await graph.stream() + ).pipeTo( + new WritableStream({ + write(chunk) { + chunks.push(chunk); + }, + }), + ); + + expect(chunks).toEqual(['first update', 'final update']); + expect(spanEnd).toHaveBeenCalledTimes(1); + }); + + it('records response attributes when pipeThrough is unavailable', async () => { + const responseMessage = { role: 'assistant', content: 'Clear skies' }; + const stream = new TestLangGraphStream([{ agent: { messages: [responseMessage] } }]); + Object.defineProperty(stream, 'pipeThrough', { value: undefined }); + const compiledGraph = { stream: vi.fn().mockResolvedValue(stream) }; + const compile = instrumentStateGraphCompile(() => compiledGraph, { recordOutputs: true }); + const graph = compile() as { stream: () => Promise> }; + + await (await graph.stream()).pipeTo(new WritableStream()); + + expect(spanSetAttribute).toHaveBeenCalledWith( + GEN_AI_RESPONSE_TEXT, + '[{"role":"assistant","content":"Clear skies"}]', + ); + expect(spanEnd).toHaveBeenCalledTimes(1); + }); + + it('does not re-enter custom pipeThrough implementations from pipeTo', async () => { + const chunks: string[] = []; + const stream = new TestLangGraphStream(['first update', 'final update']); + Object.defineProperty(stream, 'pipeThrough', { + configurable: true, + value( + this: TestLangGraphStream, + transform: ReadableWritablePair, + options?: StreamPipeOptions, + ): ReadableStream { + void this.pipeTo(transform.writable, options).catch(() => {}); + return transform.readable; + }, + writable: true, + }); + const compiledGraph = { stream: vi.fn().mockResolvedValue(stream) }; + const compile = instrumentStateGraphCompile(() => compiledGraph, {}); + const graph = compile() as { stream: () => Promise> }; + + await ( + await graph.stream() + ).pipeTo( + new WritableStream({ + write(chunk) { + chunks.push(chunk); + }, + }), + ); + + expect(chunks).toEqual(['first update', 'final update']); + expect(spanEnd).toHaveBeenCalledTimes(1); + }); + + it('applies pipe options to the destination', async () => { + const stream = new TestLangGraphStream(['first update']); + const compiledGraph = { stream: vi.fn().mockResolvedValue(stream) }; + const compile = instrumentStateGraphCompile(() => compiledGraph, {}); + const graph = compile() as { stream: () => Promise> }; + const written: string[] = []; + const destination = new WritableStream({ write: chunk => void written.push(chunk) }); + + await (await graph.stream()).pipeTo(destination, { preventClose: true }); + + const writer = destination.getWriter(); + await writer.write('after pipe'); + expect(written).toEqual(['first update', 'after pipe']); + expect(spanEnd).toHaveBeenCalledTimes(1); + }); + + it('ends the span when pipeThrough output is consumed', async () => { + const responseMessage = { role: 'assistant', content: 'Clear skies' }; + const stream = new TestLangGraphStream([{ agent: { messages: [responseMessage] } }]); + const compiledGraph = { stream: vi.fn().mockResolvedValue(stream) }; + const compile = instrumentStateGraphCompile(() => compiledGraph, { recordOutputs: true }); + const graph = compile() as { stream: () => Promise> }; + + const transformed = (await graph.stream()).pipeThrough(new TransformStream()); + await transformed.pipeTo(new WritableStream()); + await vi.waitFor(() => { + expect(spanSetAttribute).toHaveBeenCalledWith( + GEN_AI_RESPONSE_TEXT, + '[{"role":"assistant","content":"Clear skies"}]', + ); + expect(spanEnd).toHaveBeenCalledTimes(1); + }); + }); + + it('throws when pipeThrough is called with a locked source', async () => { + const stream = new TestLangGraphStream(['first update']); + const compiledGraph = { stream: vi.fn().mockResolvedValue(stream) }; + const compile = instrumentStateGraphCompile(() => compiledGraph, {}); + const graph = compile() as { stream: () => Promise> }; + const instrumentedStream = await graph.stream(); + const reader = instrumentedStream.getReader(); + + expect(() => instrumentedStream.pipeThrough(new TransformStream())).toThrow(TypeError); + expect(spanSetStatus).not.toHaveBeenCalled(); + expect(spanEnd).not.toHaveBeenCalled(); + + reader.releaseLock(); + }); + + it('throws when pipeThrough is called with a locked writable', async () => { + const stream = new TestLangGraphStream(['first update']); + const compiledGraph = { stream: vi.fn().mockResolvedValue(stream) }; + const compile = instrumentStateGraphCompile(() => compiledGraph, {}); + const graph = compile() as { stream: () => Promise> }; + const transform = new TransformStream(); + const writer = transform.writable.getWriter(); + const instrumentedStream = await graph.stream(); + + expect(() => instrumentedStream.pipeThrough(transform)).toThrow(TypeError); + expect(spanSetStatus).not.toHaveBeenCalled(); + expect(spanEnd).not.toHaveBeenCalled(); + + writer.releaseLock(); + }); + + it('errors the pipeThrough readable when the source stream fails', async () => { + const error = new Error('graph failed'); + const stream = new ReadableStream({ + pull(controller) { + controller.error(error); + }, + }); + const compiledGraph = { stream: vi.fn().mockResolvedValue(stream) }; + const compile = instrumentStateGraphCompile(() => compiledGraph, {}); + const graph = compile() as { stream: () => Promise> }; + + const reader = (await graph.stream()).pipeThrough(new TransformStream()).getReader(); + + await expect(reader.read()).rejects.toThrow(error); + expect(spanSetStatus).toHaveBeenCalledWith({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + expect(spanEnd).toHaveBeenCalledTimes(1); + }); + + it('ends the span when direct next calls consume the stream', async () => { + const stream = new TestLangGraphStream(['first update']); + const compiledGraph = { stream: vi.fn().mockResolvedValue(stream) }; + const compile = instrumentStateGraphCompile(() => compiledGraph, {}); + const graph = compile() as { stream: () => Promise> }; + + const instrumentedStream = await graph.stream(); + expect(await instrumentedStream.next()).toEqual({ done: false, value: 'first update' }); + expect(spanEnd).not.toHaveBeenCalled(); + expect(await instrumentedStream.next()).toEqual({ done: true, value: undefined }); + + expect(spanEnd).toHaveBeenCalledTimes(1); + }); + + it('ends the span when the stream is canceled directly', async () => { + const stream = new TestLangGraphStream(['first update']); + const compiledGraph = { stream: vi.fn().mockResolvedValue(stream) }; + const compile = instrumentStateGraphCompile(() => compiledGraph, {}); + const graph = compile() as { stream: () => Promise> }; + + await (await graph.stream()).cancel('no longer needed'); + + expect(spanEnd).toHaveBeenCalledTimes(1); + }); + + it('records response attributes from the final streamed state', async () => { + const inputMessages = [{ role: 'user', content: 'What is the weather in Paris?' }]; + const intermediateMessage = { + role: 'assistant', + content: 'Checking the forecast', + response_metadata: { + model_name: 'weather-model-v2', + tokenUsage: { + promptTokens: 8, + completionTokens: 2, + totalTokens: 10, + }, + }, + }; + const outputMessage = { + role: 'assistant', + content: 'Clear skies', + response_metadata: { + model_name: 'weather-model-v2', + finish_reason: 'stop', + tokenUsage: { + promptTokens: 12, + completionTokens: 3, + totalTokens: 15, + }, + }, + }; + const stream = new TestLangGraphStream([ + { planner: { messages: [intermediateMessage] } }, + { agent: { messages: [outputMessage] } }, + ]); + const compiledGraph = { stream: vi.fn().mockResolvedValue(stream) }; + const compile = instrumentStateGraphCompile(() => compiledGraph, { recordInputs: true, recordOutputs: true }); + const graph = compile() as { stream: (input: unknown) => Promise> }; + + for await (const _chunk of await graph.stream({ messages: inputMessages })) { + // The final state is applied when iteration completes. + } + + expect(spanSetAttribute).toHaveBeenCalledWith( + GEN_AI_RESPONSE_TEXT, + '[{"role":"assistant","content":"Checking the forecast"},{"role":"assistant","content":"Clear skies"}]', + ); + expect(spanSetAttribute).toHaveBeenCalledWith(GEN_AI_RESPONSE_MODEL, 'weather-model-v2'); + expect(spanSetAttribute).toHaveBeenCalledWith(GEN_AI_RESPONSE_FINISH_REASONS, ['stop']); + expect(spanSetAttribute).toHaveBeenCalledWith(GEN_AI_USAGE_INPUT_TOKENS, 20); + expect(spanSetAttribute).toHaveBeenCalledWith(GEN_AI_USAGE_OUTPUT_TOKENS, 5); + expect(spanSetAttribute).toHaveBeenCalledWith(GEN_AI_USAGE_TOTAL_TOKENS, 25); + expect(spanEnd).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/server-utils/test/ai/lib/tracing/langgraph-streaming.test.ts b/packages/server-utils/test/ai/lib/tracing/langgraph-streaming.test.ts new file mode 100644 index 000000000000..5918c993ef92 --- /dev/null +++ b/packages/server-utils/test/ai/lib/tracing/langgraph-streaming.test.ts @@ -0,0 +1,446 @@ +import { describe, expect, it, vi } from 'vitest'; +import type { Span } from '@sentry/core'; +import { SPAN_STATUS_ERROR } from '@sentry/core'; +import { AIMessageChunk } from '@langchain/core/messages'; +import { GEN_AI_RESPONSE_TEXT, GEN_AI_USAGE_OUTPUT_TOKENS } from '@sentry/conventions/attributes'; +import { instrumentStreamResult } from '../../../../src/ai/langgraph/streaming'; + +function createSpan() { + return { + end: vi.fn(), + setAttribute: vi.fn(), + setStatus: vi.fn(), + spanContext: () => ({ traceId: 'a'.repeat(32), spanId: 'b'.repeat(16), traceFlags: 1 }), + }; +} + +describe('LangGraph iterable protocol', () => { + it('preserves repeated iteration over a reusable async iterable', async () => { + const span = createSpan(); + const stream = instrumentStreamResult( + { + async *[Symbol.asyncIterator]() { + yield 'Sunny in Seoul'; + yield 'Clear skies tomorrow'; + }, + }, + span as unknown as Span, + null, + false, + ); + const first: string[] = []; + const second: string[] = []; + for await (const chunk of stream) { + first.push(chunk); + } + for await (const chunk of stream) { + second.push(chunk); + } + expect(first).toEqual(['Sunny in Seoul', 'Clear skies tomorrow']); + expect(second).toEqual(first); + expect(span.end).toHaveBeenCalledTimes(1); + }); +}); + +describe('LangGraph pipeThrough execution lifecycle', () => { + it('ends the span when the graph finishes, while buffered output outlives it', async () => { + const span = createSpan(); + let finishTransform!: () => void; + const transformPending = new Promise(resolve => { + finishTransform = resolve; + }); + const stream = instrumentStreamResult( + new ReadableStream({ + start(controller) { + controller.enqueue('Sunny in Seoul'); + controller.close(); + }, + }), + span as unknown as Span, + null, + false, + ); + const output = stream.pipeThrough( + new TransformStream( + { + async transform(chunk, controller) { + await transformPending; + controller.enqueue(chunk); + }, + }, + undefined, + { highWaterMark: 2 }, + ), + ); + + await vi.waitFor(() => expect(span.end).toHaveBeenCalledTimes(1)); + finishTransform(); + const reader = output.getReader(); + expect(await reader.read()).toEqual({ done: false, value: 'Sunny in Seoul' }); + expect(await reader.read()).toEqual({ done: true, value: undefined }); + expect(span.setStatus).not.toHaveBeenCalled(); + reader.releaseLock(); + }); + + it('does not mark the agent span failed when a downstream transform fails', async () => { + const span = createSpan(); + const error = new Error('Weather transform flush failed'); + const stream = instrumentStreamResult( + new ReadableStream({ + start(controller) { + controller.close(); + }, + }), + span as unknown as Span, + null, + false, + ); + const output = stream.pipeThrough( + new TransformStream({ + async flush() { + throw error; + }, + }), + ); + const reader = output.getReader(); + await expect(reader.read()).rejects.toThrow(error); + expect(span.setStatus).not.toHaveBeenCalled(); + expect(span.end).toHaveBeenCalledTimes(1); + reader.releaseLock(); + }); +}); + +describe('LangGraph stream response recording', () => { + it.each([ + { updates: false, nativeConcat: false }, + { updates: false, nativeConcat: true }, + { updates: true, nativeConcat: false }, + { updates: true, nativeConcat: true }, + ])('combines chunks with the same id without duplicating updates: %o', async ({ updates, nativeConcat }) => { + const span = createSpan(); + const firstMessage = { + id: 'seoul-weather', + type: 'ai', + content: 'Sunny ', + }; + const lastMessage = { + id: firstMessage.id, + type: 'ai', + content: 'in Seoul', + usage_metadata: { input_tokens: 5, output_tokens: 4, total_tokens: 9 }, + }; + const first = nativeConcat ? new AIMessageChunk(firstMessage) : firstMessage; + const last = nativeConcat ? new AIMessageChunk(lastMessage) : lastMessage; + const stream = instrumentStreamResult( + { + async *[Symbol.asyncIterator]() { + yield ['messages', [first, { langgraph_node: 'agent' }]]; + yield ['messages', [last, { langgraph_node: 'agent' }]]; + if (updates) { + yield ['updates', { agent: { messages: [{ ...lastMessage, content: 'Sunny in Seoul' }] } }]; + } + }, + }, + span as unknown as Span, + null, + true, + ); + const chunks = []; + for await (const chunk of stream) { + chunks.push(chunk); + } + expect(chunks).toHaveLength(updates ? 3 : 2); + expect(span.setAttribute).toHaveBeenCalledWith( + GEN_AI_RESPONSE_TEXT, + JSON.stringify([{ role: 'assistant', content: 'Sunny in Seoul' }]), + ); + expect(span.setAttribute).toHaveBeenCalledWith(GEN_AI_USAGE_OUTPUT_TOKENS, 4); + }); + + it.each([false, true])('records messages tuples with multi-mode wrapping %s', async multiMode => { + const span = createSpan(); + const message = { type: 'ai', content: 'Sunny in Seoul', usage_metadata: { output_tokens: 4 } }; + const tuple = [message, { langgraph_node: 'agent' }]; + const chunk = multiMode ? ['messages', tuple] : tuple; + const stream = instrumentStreamResult( + { + async *[Symbol.asyncIterator]() { + yield chunk; + }, + }, + span as unknown as Span, + [{ type: 'human', content: 'Weather in Seoul?' }], + true, + ); + const chunks = []; + for await (const value of stream) { + chunks.push(value); + } + expect(chunks).toEqual([chunk]); + expect(span.setAttribute).toHaveBeenCalledWith( + GEN_AI_RESPONSE_TEXT, + JSON.stringify([{ role: 'assistant', content: 'Sunny in Seoul' }]), + ); + expect(span.setAttribute).toHaveBeenCalledWith(GEN_AI_USAGE_OUTPUT_TOKENS, 4); + expect(span.end).toHaveBeenCalledTimes(1); + }); +}); + +describe('LangGraph pipeTo preconditions', () => { + it.each(['source', 'destination'])('keeps consumption active after piping with a locked %s', async lockedTarget => { + const span = createSpan(); + const stream = instrumentStreamResult( + new ReadableStream({ + start(controller) { + controller.enqueue('Sunny in Seoul'); + controller.enqueue('Clear skies tomorrow'); + controller.close(); + }, + }), + span as unknown as Span, + null, + false, + ); + const destination = new WritableStream(); + const reader = lockedTarget === 'source' ? stream.getReader() : undefined; + const writer = lockedTarget === 'destination' ? destination.getWriter() : undefined; + if (reader) { + expect(await reader.read()).toEqual({ done: false, value: 'Sunny in Seoul' }); + } + await expect(stream.pipeTo(destination)).rejects.toThrow(TypeError); + expect(span.end).not.toHaveBeenCalled(); + expect(span.setStatus).not.toHaveBeenCalled(); + writer?.releaseLock(); + if (reader) { + expect(await reader.read()).toEqual({ done: false, value: 'Clear skies tomorrow' }); + expect(await reader.read()).toEqual({ done: true, value: undefined }); + reader.releaseLock(); + } else { + await stream.pipeTo(destination); + } + expect(span.end).toHaveBeenCalledTimes(1); + expect(span.setStatus).not.toHaveBeenCalled(); + }); +}); + +describe('LangGraph stream caller errors', () => { + it.each([undefined, null, {}])( + 'keeps the stream usable after an invalid pipeTo destination: %s', + async destination => { + const span = createSpan(); + const stream = instrumentStreamResult( + new ReadableStream({ + start(controller) { + controller.enqueue('Sunny in Seoul'); + controller.close(); + }, + }), + span as unknown as Span, + null, + false, + ); + + await expect(stream.pipeTo(destination as WritableStream)).rejects.toThrow(TypeError); + + expect(span.end).not.toHaveBeenCalled(); + expect(span.setStatus).not.toHaveBeenCalled(); + const chunks: string[] = []; + await stream.pipeTo( + new WritableStream({ + write: chunk => { + chunks.push(chunk); + }, + }), + ); + expect(chunks).toEqual(['Sunny in Seoul']); + expect(span.end).toHaveBeenCalledTimes(1); + expect(span.setStatus).not.toHaveBeenCalled(); + }, + ); + + it('keeps reader consumption active when cancel is called on a locked stream', async () => { + const span = createSpan(); + const stream = instrumentStreamResult( + new ReadableStream({ + start(controller) { + controller.enqueue('Sunny in Seoul'); + controller.enqueue('Clear skies tomorrow'); + controller.close(); + }, + }), + span as unknown as Span, + null, + false, + ); + const reader = stream.getReader(); + expect(await reader.read()).toEqual({ done: false, value: 'Sunny in Seoul' }); + + await expect(stream.cancel()).rejects.toThrow(TypeError); + + expect(span.end).not.toHaveBeenCalled(); + expect(span.setStatus).not.toHaveBeenCalled(); + expect(await reader.read()).toEqual({ done: false, value: 'Clear skies tomorrow' }); + expect(await reader.read()).toEqual({ done: true, value: undefined }); + reader.releaseLock(); + expect(span.end).toHaveBeenCalledTimes(1); + expect(span.setStatus).not.toHaveBeenCalled(); + }); +}); + +describe('LangGraph wrapped stream chunks', () => { + const user = { role: 'user', content: 'Weather?' }; + const assistant = { role: 'assistant', content: 'Sunny' }; + + async function recordedResponse(chunks: unknown[]): Promise { + const span = createSpan(); + const stream = instrumentStreamResult( + { + async *[Symbol.asyncIterator]() { + yield* chunks; + }, + }, + span as unknown as Span, + [user], + true, + ); + for await (const _chunk of stream) { + // drain + } + return span.setAttribute.mock.calls.find(([key]) => key === GEN_AI_RESPONSE_TEXT)?.[1]; + } + + it('records only the root graph output from `subgraphs` chunks', async () => { + const response = await recordedResponse([ + [[], { messages: [user] }], + [['inner:1'], { messages: [user, { role: 'assistant', content: 'inner thought' }] }], + [[], { messages: [user, assistant] }], + ]); + + expect(response).toBe('[{"role":"assistant","content":"Sunny"}]'); + }); + + it('ignores modes that do not carry the graph output', async () => { + const response = await recordedResponse([ + [[], 'checkpoints', { values: { messages: [user, { role: 'assistant', content: 'stale' }] } }], + [[], 'values', { messages: [user, assistant] }], + ]); + + expect(response).toBe('[{"role":"assistant","content":"Sunny"}]'); + }); +}); + +describe('LangGraph stream cancellation and tee', () => { + it('does not mark a caller cancel as a failure when it rejects the pending read', async () => { + const span = createSpan(); + const cancelError = new TypeError('Invalid state: Releasing reader'); + let rejectRead!: (error: unknown) => void; + const source = new ReadableStream(); + const stream = Object.assign(source, { + getReader: () => ({ read: () => new Promise((_resolve, reject) => (rejectRead = reject)) }), + cancel: async () => rejectRead(cancelError), + }); + instrumentStreamResult(stream, span as unknown as Span, null, false); + + const pendingRead = stream.getReader().read(); + Object.defineProperty(stream, 'locked', { value: true }); + await stream.cancel(); + + await expect(pendingRead).rejects.toBe(cancelError); + expect(span.setStatus).not.toHaveBeenCalled(); + expect(span.end).toHaveBeenCalledTimes(1); + }); + + it('records and ends the span when the stream is consumed through tee()', async () => { + const span = createSpan(); + const stream = instrumentStreamResult( + new ReadableStream({ + start(controller) { + controller.enqueue({ agent: { messages: [{ role: 'assistant', content: 'Sunny' }] } }); + controller.close(); + }, + }), + span as unknown as Span, + null, + true, + ); + + const [first, second] = stream.tee(); + for await (const _chunk of first) { + // drain + } + for await (const _chunk of second) { + // drain + } + + expect(span.setAttribute).toHaveBeenCalledWith(GEN_AI_RESPONSE_TEXT, '[{"role":"assistant","content":"Sunny"}]'); + expect(span.end).toHaveBeenCalledTimes(1); + }); + + it('ends the span when a cancel on a locked stream takes effect, even if nothing reads again', async () => { + const span = createSpan(); + const source = new ReadableStream(); + const stream = Object.assign(source, { cancel: async () => undefined }); + instrumentStreamResult(stream, span as unknown as Span, null, false); + Object.defineProperty(stream, 'locked', { value: true }); + + await stream.cancel('no longer needed'); + + expect(span.setStatus).not.toHaveBeenCalled(); + expect(span.end).toHaveBeenCalledTimes(1); + }); + + it('still marks a later stream failure as an error after a cancel that was rejected', async () => { + const span = createSpan(); + const error = new Error('graph failed'); + let pulls = 0; + const stream = instrumentStreamResult( + new ReadableStream({ + pull(controller) { + if (pulls++ === 0) { + controller.enqueue('Sunny'); + } else { + controller.error(error); + } + }, + }), + span as unknown as Span, + null, + false, + ); + const reader = stream.getReader(); + expect(await reader.read()).toEqual({ done: false, value: 'Sunny' }); + + await expect(stream.cancel()).rejects.toThrow(TypeError); + await expect(reader.read()).rejects.toBe(error); + + expect(span.setStatus).toHaveBeenCalledExactlyOnceWith({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + expect(span.end).toHaveBeenCalledTimes(1); + }); + + it.each([ + ['pipeThrough without a readable side', (stream: ReadableStream) => stream.pipeThrough({ writable: new WritableStream() } as ReadableWritablePair)], + ['pipeThrough with an invalid signal', (stream: ReadableStream) => stream.pipeThrough(new TransformStream(), { signal: {} as AbortSignal })], + ['pipeTo with an invalid signal', (stream: ReadableStream) => stream.pipeTo(new WritableStream(), { signal: {} as AbortSignal })], + ])('rejects %s without locking the source', async (_label, pipe) => { + const span = createSpan(); + const stream = instrumentStreamResult( + new ReadableStream({ + start(controller) { + controller.enqueue('Sunny'); + controller.close(); + }, + }), + span as unknown as Span, + null, + false, + ); + + await expect(async () => pipe(stream)).rejects.toThrow(TypeError); + + expect(stream.locked).toBe(false); + const reader = stream.getReader(); + expect(await reader.read()).toEqual({ done: false, value: 'Sunny' }); + expect(await reader.read()).toEqual({ done: true, value: undefined }); + expect(span.end).toHaveBeenCalledTimes(1); + }); +});