From 82e75ebc5f0439a632142720c804a8e5d326a799 Mon Sep 17 00:00:00 2001 From: "Seongho.Bak" Date: Wed, 26 Aug 2026 12:16:32 +0900 Subject: [PATCH 01/13] fix(server-utils): Instrument LangGraph stream executions Fixes #19626 --- .../suites/tracing/langgraph/index.ts | 7 + .../suites/tracing/langgraph/test.ts | 18 +- .../suites/tracing/langgraph/scenario.mjs | 7 + .../suites/tracing/langgraph/test.ts | 21 +- .../server-utils/src/ai/langgraph/index.ts | 274 ++++++++++++------ .../src/ai/langgraph/streaming.ts | 58 ++++ .../server-utils/src/ai/langgraph/types.ts | 1 + .../src/integrations/langgraph.ts | 30 +- .../ai/lib/tracing/langgraph-stream.test.ts | 116 ++++++++ 9 files changed, 433 insertions(+), 99 deletions(-) create mode 100644 packages/server-utils/src/ai/langgraph/streaming.ts create mode 100644 packages/server-utils/test/ai/lib/tracing/langgraph-stream.test.ts 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 6b22dcd4388a..5433280eb09b 100644 --- a/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/index.ts +++ b/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/index.ts @@ -61,6 +61,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 6615c79430b4..722db0d3cf1a 100644 --- a/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/test.ts +++ b/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/test.ts @@ -6,6 +6,7 @@ import { GEN_AI_OPERATION_NAME, GEN_AI_PIPELINE_NAME, GEN_AI_RESPONSE_MODEL, + GEN_AI_RESPONSE_STREAMING, GEN_AI_USAGE_INPUT_TOKENS, GEN_AI_USAGE_OUTPUT_TOKENS, GEN_AI_USAGE_TOTAL_TOKENS, @@ -17,7 +18,7 @@ import { createRunner } from '../../../runner'; // 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 => { @@ -29,13 +30,14 @@ it('traces langgraph compile and invoke operations', async ({ signal }) => { const container = envelope[1]?.[1]?.[1] as any; expect(container).toBeDefined(); - expect(container.items).toHaveLength(1); + expect(container.items).toHaveLength(2); expect(container.items.map((span: SerializedStreamedSpan) => span.name).sort()).toEqual([ 'invoke_agent weather_assistant', + 'invoke_agent weather_assistant', ]); const invokeAgentSpan = container.items.find( - (span: SerializedStreamedSpan) => span.name === 'invoke_agent weather_assistant', + (span: SerializedStreamedSpan) => span.attributes[GEN_AI_RESPONSE_STREAMING] === undefined, ); expect(invokeAgentSpan).toBeDefined(); expect(invokeAgentSpan!.status).toBe('ok'); @@ -73,6 +75,16 @@ it('traces langgraph compile and invoke operations', async ({ signal }) => { type: 'integer', value: 30, }); + + const streamSpan = container.items.find( + (span: SerializedStreamedSpan) => 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"}]', + }); }) .start(signal); await runner.makeRequest('get', '/'); 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..062cddaba699 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,13 @@ async function run() { { role: 'user', content: 'Tell me about the weather' }, ], }); + + const stream = await graph.stream({ + messages: [{ role: 'user', content: 'Stream the weather forecast' }], + }); + for await (const _chunk of stream) { + // Consuming the iterator is what runs the graph and completes the agent span. + } }); 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 3f85267b154e..2cde2a5d527c 100644 --- a/dev-packages/node-integration-tests/suites/tracing/langgraph/test.ts +++ b/dev-packages/node-integration-tests/suites/tracing/langgraph/test.ts @@ -7,6 +7,7 @@ import { GEN_AI_OPERATION_NAME, GEN_AI_PIPELINE_NAME, GEN_AI_RESPONSE_MODEL, + GEN_AI_RESPONSE_STREAMING, GEN_AI_RESPONSE_TEXT, GEN_AI_RESPONSE_TOOL_CALLS, GEN_AI_SYSTEM_INSTRUCTIONS, @@ -30,14 +31,15 @@ describe('LangGraph integration', () => { .expect({ transaction: { transaction: 'langgraph-test' } }) .expect({ span: container => { - expect(container.items).toHaveLength(2); + expect(container.items).toHaveLength(3); expect(container.items.map(span => span.name).sort()).toEqual([ 'invoke_agent weather_assistant', 'invoke_agent weather_assistant', + 'invoke_agent weather_assistant', ]); const invokeAgentSpans = container.items.filter(span => span.name === 'invoke_agent weather_assistant'); - expect(invokeAgentSpans).toHaveLength(2); + expect(invokeAgentSpans).toHaveLength(3); for (const span of invokeAgentSpans) { expect(span.status).toBe('ok'); expect(span.attributes['sentry.op'].value).toBe('gen_ai.invoke_agent'); @@ -46,6 +48,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 streamSpan = invokeAgentSpans.find( + span => span.attributes[GEN_AI_RESPONSE_STREAMING]?.value === true, + ); + expect(streamSpan).toBeDefined(); }, }) .start() @@ -59,7 +66,7 @@ describe('LangGraph integration', () => { .expect({ transaction: { transaction: 'langgraph-test' } }) .expect({ span: container => { - expect(container.items).toHaveLength(2); + expect(container.items).toHaveLength(3); const weatherTodaySpan = container.items.find(span => getStringAttributeValue(span.attributes[GEN_AI_INPUT_MESSAGES]?.value)?.includes( @@ -81,6 +88,14 @@ 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 = container.items.find(span => + getStringAttributeValue(span.attributes[GEN_AI_INPUT_MESSAGES]?.value)?.includes( + 'Stream the weather forecast', + ), + ); + expect(weatherStreamSpan).toBeDefined(); + expect(weatherStreamSpan!.attributes[GEN_AI_RESPONSE_STREAMING].value).toBe(true); }, }) .start() diff --git a/packages/server-utils/src/ai/langgraph/index.ts b/packages/server-utils/src/ai/langgraph/index.ts index de8c0ec962c1..a599107852e9 100644 --- a/packages/server-utils/src/ai/langgraph/index.ts +++ b/packages/server-utils/src/ai/langgraph/index.ts @@ -1,5 +1,13 @@ /* eslint-disable typescript-eslint/no-deprecated */ -import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, SPAN_STATUS_ERROR, startSpan, stringify } from '@sentry/core'; +import { + getCurrentScope, + SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, + SPAN_STATUS_ERROR, + startSpan, + startSpanManual, + stringify, +} from '@sentry/core'; +import type { Span } from '@sentry/core'; import { GEN_AI_AGENT_NAME, GEN_AI_CONVERSATION_ID, @@ -7,6 +15,7 @@ import { 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, @@ -15,8 +24,9 @@ import { GEN_AI_INVOKE_AGENT } from '@sentry/conventions/op'; import { extractSystemInstructions, resolveAIRecordingOptions } from '../core/utils'; import { createLangChainCallbackHandler } from '../langchain'; import type { BaseChatModel, LangChainMessage } from '../langchain/types'; -import { normalizeLangChainMessages } from '../langchain/utils'; +import { _INTERNAL_mergeLangChainCallbackHandler, normalizeLangChainMessages } from '../langchain/utils'; import { LANGGRAPH_ORIGIN } from './constants'; +import { getGraphInstrumentationId, instrumentStreamResult, isAsyncIterable } from './streaming'; import type { CompiledGraph, LangGraphOptions } from './types'; import { extractAgentNameFromParams, @@ -25,15 +35,15 @@ import { setResponseAttributes, wrapToolsWithSpans, } from './utils'; -import { _INTERNAL_mergeLangChainCallbackHandler } from '../langchain/utils'; let _insideCreateReactAgent = false; const SENTRY_PATCHED = '__sentry_patched__'; +const LANGGRAPH_INVOKE_ACTIVE = 'sentry_langgraph_invoke_active'; /** - * 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, @@ -55,7 +65,7 @@ export function instrumentStateGraphCompile( const compiledGraph = Reflect.apply(target, thisArg, args); const compileOptions = args.length > 0 ? (args[0] as Record) : {}; - // Instrument agent invoke method on the compiled graph + // Instrument agent methods on the compiled graph const originalInvoke = compiledGraph.invoke; if (originalInvoke && typeof originalInvoke === 'function') { compiledGraph.invoke = instrumentCompiledGraphInvoke( @@ -68,6 +78,18 @@ export function instrumentStateGraphCompile( ); } + const originalStream = compiledGraph.stream; + if (originalStream && typeof originalStream === 'function') { + compiledGraph.stream = instrumentCompiledGraphStream( + originalStream.bind(compiledGraph), + compiledGraph, + compileOptions, + options, + undefined, + sentryHandler, + ); + } + return compiledGraph; }, }); @@ -89,100 +111,167 @@ export function instrumentCompiledGraphInvoke( llm?: BaseChatModel | null, sentryCallbackHandler?: unknown, ): (...args: unknown[]) => Promise { - return new Proxy(originalInvoke, { + return instrumentCompiledGraphOperation( + originalInvoke, + graphInstance, + compileOptions, + options, + llm, + sentryCallbackHandler, + false, + ); +} + +export function instrumentCompiledGraphStream( + originalStream: (...args: unknown[]) => Promise>, + graphInstance: CompiledGraph, + compileOptions: Record, + options: LangGraphOptions, + llm?: BaseChatModel | null, + sentryCallbackHandler?: unknown, +): (...args: unknown[]) => Promise> { + return instrumentCompiledGraphOperation( + originalStream, + graphInstance, + compileOptions, + options, + llm, + sentryCallbackHandler, + true, + ) as (...args: unknown[]) => Promise>; +} + +function instrumentCompiledGraphOperation( + originalOperation: (...args: unknown[]) => Promise, + graphInstance: CompiledGraph, + compileOptions: Record, + options: LangGraphOptions, + llm: BaseChatModel | null | undefined, + sentryCallbackHandler: unknown, + streaming: boolean, +): (...args: unknown[]) => Promise { + const graphInstrumentationId = getGraphInstrumentationId(graphInstance); + + return new Proxy(originalOperation, { apply(target, thisArg, args: unknown[]): Promise { + if ( + streaming && + getCurrentScope().getScopeData().sdkProcessingMetadata[LANGGRAPH_INVOKE_ACTIVE] === graphInstrumentationId + ) { + return Reflect.apply(target, thisArg, args); + } + const modelName = llm?.modelName ?? llm?.model; - return startSpan( - { - name: 'invoke_agent', - attributes: { - [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: LANGGRAPH_ORIGIN, - [SENTRY_OP]: GEN_AI_INVOKE_AGENT, - [GEN_AI_OPERATION_NAME]: 'invoke_agent', - }, + const spanOptions = { + name: 'invoke_agent', + attributes: { + [SEMANTIC_ATTRIBUTE_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}`); - } - - if (modelName) { - span.setAttribute(GEN_AI_REQUEST_MODEL, modelName); - } + }; + 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}`); + } - // Extract thread_id from the config (second argument) - // LangGraph uses config.configurable.thread_id for conversation/session linking - const config = args.length > 1 ? (args[1] as Record | undefined) : undefined; - const configurable = config?.configurable as Record | undefined; - const threadId = configurable?.thread_id; - if (threadId && typeof threadId === 'string') { - span.setAttribute(GEN_AI_CONVERSATION_ID, threadId); - } + 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 = { - ...existingMetadata, - __sentry_langgraph__: true, - ...(typeof graphName === 'string' ? { lc_agent_name: graphName } : {}), - }; - - invokeConfig.callbacks = _INTERNAL_mergeLangChainCallbackHandler( - invokeConfig.callbacks, - sentryCallbackHandler, - ); - } + // Extract thread_id from the config (second argument) + // LangGraph uses config.configurable.thread_id for conversation/session linking + const config = args.length > 1 ? (args[1] as Record | undefined) : undefined; + const configurable = config?.configurable as Record | undefined; + const threadId = configurable?.thread_id; + if (threadId && typeof threadId === 'string') { + span.setAttribute(GEN_AI_CONVERSATION_ID, threadId); + } - // Extract available tools from the graph instance - const tools = extractToolsFromCompiledGraph(graphInstance); - if (tools) { - span.setAttribute(GEN_AI_TOOL_DEFINITIONS, JSON.stringify(tools)); - } + // 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 = { + ...existingMetadata, + __sentry_langgraph__: true, + ...(typeof graphName === 'string' ? { lc_agent_name: graphName } : {}), + }; + + invokeConfig.callbacks = _INTERNAL_mergeLangChainCallbackHandler( + invokeConfig.callbacks, + sentryCallbackHandler, + ); + } - // Parse input messages - const recordInputs = options.recordInputs; - const recordOutputs = options.recordOutputs; - const inputMessages = - args.length > 0 ? ((args[0] as { messages?: LangChainMessage[] } | null)?.messages ?? []) : []; + // Extract available tools from the graph instance + const tools = extractToolsFromCompiledGraph(graphInstance); + if (tools) { + span.setAttribute(GEN_AI_TOOL_DEFINITIONS, JSON.stringify(tools)); + } - if (inputMessages && recordInputs) { - const normalizedMessages = normalizeLangChainMessages(inputMessages); - const { systemInstructions, filteredMessages } = extractSystemInstructions(normalizedMessages); + // Parse input messages + const recordInputs = options.recordInputs; + const recordOutputs = options.recordOutputs; + const inputMessages = + args.length > 0 ? ((args[0] as { messages?: LangChainMessage[] } | null)?.messages ?? []) : []; - if (systemInstructions) { - span.setAttribute(GEN_AI_SYSTEM_INSTRUCTIONS, systemInstructions); - } + if (inputMessages && recordInputs) { + const normalizedMessages = normalizeLangChainMessages(inputMessages); + const { systemInstructions, filteredMessages } = extractSystemInstructions(normalizedMessages); - span.setAttributes({ - [GEN_AI_INPUT_MESSAGES]: stringify(filteredMessages), - }); + if (systemInstructions) { + span.setAttribute(GEN_AI_SYSTEM_INSTRUCTIONS, systemInstructions); } - // Call original invoke - const result = await Reflect.apply(target, thisArg, args); + span.setAttributes({ + [GEN_AI_INPUT_MESSAGES]: stringify(filteredMessages), + }); + } - if (recordOutputs) { - setResponseAttributes(span, inputMessages ?? null, result); + if (!streaming) { + // LangGraph implements invoke() by consuming stream(); the shared async scope keeps that internal + // call from producing a nested duplicate while leaving direct concurrent stream() calls unaffected. + getCurrentScope().setSDKProcessingMetadata({ + [LANGGRAPH_INVOKE_ACTIVE]: graphInstrumentationId, + }); + } + + const result = await Reflect.apply(target, thisArg, args); + + if (streaming) { + if (isAsyncIterable(result)) { + span.setAttribute(GEN_AI_RESPONSE_STREAMING, true); + return instrumentStreamResult(result, span); } + 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 ?? null, 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); }, }); } @@ -221,7 +310,7 @@ export function instrumentCreateReactAgent( _insideCreateReactAgent = false; } - // Wrap invoke() on the returned compiled graph + // Wrap agent methods on the returned compiled graph const originalInvoke = compiledGraph.invoke; if (originalInvoke && typeof originalInvoke === 'function') { const compileOptions: Record = {}; @@ -239,6 +328,23 @@ export function instrumentCreateReactAgent( ); } + const originalStream = compiledGraph.stream; + if (originalStream && typeof originalStream === 'function') { + const compileOptions: Record = {}; + if (agentName) { + compileOptions.name = agentName; + } + + compiledGraph.stream = instrumentCompiledGraphStream( + originalStream.bind(compiledGraph), + 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..27d088341b47 --- /dev/null +++ b/packages/server-utils/src/ai/langgraph/streaming.ts @@ -0,0 +1,58 @@ +import { SPAN_STATUS_ERROR, withActiveSpan } from '@sentry/core'; +import type { Span } from '@sentry/core'; +import type { CompiledGraph } from './types'; + +const graphInstrumentationIds = new WeakMap(); +let nextGraphInstrumentationId = 0; + +export function getGraphInstrumentationId(graph: CompiledGraph): number { + const existingId = graphInstrumentationIds.get(graph); + if (existingId !== undefined) { + return existingId; + } + + const id = nextGraphInstrumentationId++; + graphInstrumentationIds.set(graph, id); + return id; +} + +export function isAsyncIterable(value: unknown): value is AsyncIterable { + return !!value && typeof (value as AsyncIterable)[Symbol.asyncIterator] === 'function'; +} + +export function instrumentStreamResult>(stream: T, span: Span): T { + const iterate = stream[Symbol.asyncIterator].bind(stream); + const instrumented = instrumentStreamIterator({ [Symbol.asyncIterator]: iterate }, span); + stream[Symbol.asyncIterator] = () => instrumented; + return stream; +} + +async function* instrumentStreamIterator( + stream: AsyncIterable, + span: Span, +): AsyncGenerator { + const iterator = stream[Symbol.asyncIterator](); + let completed = false; + + try { + while (true) { + const result = await withActiveSpan(span, () => iterator.next()); + if (result.done) { + completed = true; + return; + } + yield result.value; + } + } catch (error) { + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + throw error; + } finally { + try { + if (!completed) { + await withActiveSpan(span, () => iterator.return?.()); + } + } finally { + span.end(); + } + } +} 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/integrations/langgraph.ts b/packages/server-utils/src/integrations/langgraph.ts index 124260fe4029..e7d642e37f9f 100644 --- a/packages/server-utils/src/integrations/langgraph.ts +++ b/packages/server-utils/src/integrations/langgraph.ts @@ -3,7 +3,7 @@ import type { IntegrationFn } from '@sentry/core'; import { debug, defineIntegration } from '@sentry/core'; import { resolveAIRecordingOptions } from '../ai/core/utils'; import { createLangChainCallbackHandler } from '../ai/langchain'; -import { instrumentCompiledGraphInvoke } from '../ai/langgraph'; +import { instrumentCompiledGraphInvoke, instrumentCompiledGraphStream } from '../ai/langgraph'; import { LANGGRAPH_INTEGRATION_NAME } from '../ai/langgraph/constants'; import type { CompiledGraph, LangGraphOptions } from '../ai/langgraph/types'; import { extractAgentNameFromParams, extractLLMFromParams, wrapToolsWithSpans } from '../ai/langgraph/utils'; @@ -47,7 +47,7 @@ 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 => { @@ -55,11 +55,11 @@ function instrumentLanggraph(options: LangGraphOptions): void { return; } const { arguments: args, result } = message as CompileChannelContext; - wrapCompiledGraphInvoke(result, getFirstArgObject(args) ?? {}, resolvedOptions, null, sentryHandler); + wrapCompiledGraphMethods(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, ); @@ -84,7 +84,7 @@ function instrumentLanggraph(options: LangGraphOptions): void { const { arguments: args, result } = message as CreateReactAgentChannelContext; const agentName = extractAgentNameFromParams(args) ?? undefined; const compileOptions = agentName ? { name: agentName } : {}; - wrapCompiledGraphInvoke(result, compileOptions, resolvedOptions, extractLLMFromParams(args), sentryHandler); + wrapCompiledGraphMethods(result, compileOptions, resolvedOptions, extractLLMFromParams(args), sentryHandler); }); // Make sure a thrown `createReactAgent` doesn't leave the suppression flag stuck on. reactAgentChannel.error.subscribe(() => { @@ -99,10 +99,10 @@ function getFirstArgObject(args: unknown[] | undefined): Record } /** - * Wrap the compiled graph's `invoke` with the shared `invoke_agent` instrumentation, exactly as the - * OTel path does on the returned graph. + * Wrap the compiled graph's agent methods with the shared `invoke_agent` instrumentation, exactly as + * the OTel path does on the returned graph. */ -function wrapCompiledGraphInvoke( +function wrapCompiledGraphMethods( graph: unknown, compileOptions: Record, options: LangGraphOptions, @@ -125,6 +125,18 @@ function wrapCompiledGraphInvoke( sentryHandler, ); } + + const originalStream = compiledGraph.stream; + if (typeof originalStream === 'function') { + compiledGraph.stream = instrumentCompiledGraphStream( + originalStream.bind(compiledGraph), + compiledGraph, + compileOptions, + options, + llm, + sentryHandler, + ); + } } /** 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..2e81bdd5783a --- /dev/null +++ b/packages/server-utils/test/ai/lib/tracing/langgraph-stream.test.ts @@ -0,0 +1,116 @@ +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_STREAMING } 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(); + +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); + }); +}); From 591bb2bc21081b86f4a4ec2fe9460d9492e3abf1 Mon Sep 17 00:00:00 2001 From: "Seongho.Bak" Date: Wed, 26 Aug 2026 13:18:45 +0900 Subject: [PATCH 02/13] fix(server-utils): Complete LangGraph stream spans Handle ReadableStream consumption and record accumulated stream responses. --- .../suites/tracing/langgraph/test.ts | 26 ++ .../suites/tracing/langgraph/scenario.mjs | 9 +- .../suites/tracing/langgraph/test.ts | 28 ++- .../server-utils/src/ai/langgraph/index.ts | 2 +- .../src/ai/langgraph/streaming.ts | 233 +++++++++++++++++- .../ai/lib/tracing/langgraph-stream.test.ts | 161 +++++++++++- 6 files changed, 445 insertions(+), 14 deletions(-) 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 722db0d3cf1a..28ebc0ba18b5 100644 --- a/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/test.ts +++ b/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/test.ts @@ -5,8 +5,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, @@ -85,6 +87,30 @@ it('traces langgraph invoke and stream operations', async ({ signal }) => { 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/scenario.mjs b/dev-packages/node-integration-tests/suites/tracing/langgraph/scenario.mjs index 062cddaba699..86d3f34a73a4 100644 --- a/dev-packages/node-integration-tests/suites/tracing/langgraph/scenario.mjs +++ b/dev-packages/node-integration-tests/suites/tracing/langgraph/scenario.mjs @@ -48,9 +48,12 @@ async function run() { const stream = await graph.stream({ messages: [{ role: 'user', content: 'Stream the weather forecast' }], }); - for await (const _chunk of stream) { - // Consuming the iterator is what runs the graph and completes the agent span. - } + 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 2cde2a5d527c..e3bce89ef868 100644 --- a/dev-packages/node-integration-tests/suites/tracing/langgraph/test.ts +++ b/dev-packages/node-integration-tests/suites/tracing/langgraph/test.ts @@ -6,6 +6,7 @@ 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, @@ -31,15 +32,16 @@ describe('LangGraph integration', () => { .expect({ transaction: { transaction: 'langgraph-test' } }) .expect({ span: container => { - expect(container.items).toHaveLength(3); + expect(container.items).toHaveLength(4); expect(container.items.map(span => span.name).sort()).toEqual([ 'invoke_agent weather_assistant', 'invoke_agent weather_assistant', 'invoke_agent weather_assistant', + 'invoke_agent weather_assistant', ]); const invokeAgentSpans = container.items.filter(span => span.name === 'invoke_agent weather_assistant'); - expect(invokeAgentSpans).toHaveLength(3); + 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'); @@ -49,10 +51,10 @@ describe('LangGraph integration', () => { expect(span.attributes[GEN_AI_PIPELINE_NAME].value).toBe('weather_assistant'); } - const streamSpan = invokeAgentSpans.find( + const streamSpans = invokeAgentSpans.filter( span => span.attributes[GEN_AI_RESPONSE_STREAMING]?.value === true, ); - expect(streamSpan).toBeDefined(); + expect(streamSpans).toHaveLength(2); }, }) .start() @@ -66,7 +68,7 @@ describe('LangGraph integration', () => { .expect({ transaction: { transaction: 'langgraph-test' } }) .expect({ span: container => { - expect(container.items).toHaveLength(3); + expect(container.items).toHaveLength(4); const weatherTodaySpan = container.items.find(span => getStringAttributeValue(span.attributes[GEN_AI_INPUT_MESSAGES]?.value)?.includes( @@ -96,6 +98,22 @@ describe('LangGraph integration', () => { ); expect(weatherStreamSpan).toBeDefined(); 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 = container.items.find(span => + getStringAttributeValue(span.attributes[GEN_AI_INPUT_MESSAGES]?.value)?.includes( + 'Cancel the weather forecast', + ), + ); + expect(canceledStreamSpan).toBeDefined(); + expect(canceledStreamSpan!.attributes[GEN_AI_RESPONSE_STREAMING].value).toBe(true); }, }) .start() diff --git a/packages/server-utils/src/ai/langgraph/index.ts b/packages/server-utils/src/ai/langgraph/index.ts index a599107852e9..37549e39b0c2 100644 --- a/packages/server-utils/src/ai/langgraph/index.ts +++ b/packages/server-utils/src/ai/langgraph/index.ts @@ -249,7 +249,7 @@ function instrumentCompiledGraphOperation( if (streaming) { if (isAsyncIterable(result)) { span.setAttribute(GEN_AI_RESPONSE_STREAMING, true); - return instrumentStreamResult(result, span); + return instrumentStreamResult(result, span, inputMessages ?? null, recordOutputs); } span.end(); diff --git a/packages/server-utils/src/ai/langgraph/streaming.ts b/packages/server-utils/src/ai/langgraph/streaming.ts index 27d088341b47..a7bb3fd36165 100644 --- a/packages/server-utils/src/ai/langgraph/streaming.ts +++ b/packages/server-utils/src/ai/langgraph/streaming.ts @@ -1,6 +1,8 @@ import { SPAN_STATUS_ERROR, withActiveSpan } from '@sentry/core'; import type { Span } from '@sentry/core'; +import type { LangChainMessage } from '../langchain/types'; import type { CompiledGraph } from './types'; +import { setResponseAttributes } from './utils'; const graphInstrumentationIds = new WeakMap(); let nextGraphInstrumentationId = 0; @@ -20,16 +22,237 @@ export function isAsyncIterable(value: unknown): value is AsyncIterable return !!value && typeof (value as AsyncIterable)[Symbol.asyncIterator] === 'function'; } -export function instrumentStreamResult>(stream: T, span: Span): T { +export function instrumentStreamResult>( + stream: T, + span: Span, + inputMessages: LangChainMessage[] | null, + recordOutputs: boolean | undefined, +): T { + const responseState: StreamResponseState = { updateMessages: [] }; + const lifecycle = createStreamLifecycle( + span, + recordOutputs ? chunk => accumulateStreamResponse(responseState, chunk) : undefined, + recordOutputs + ? () => setResponseAttributes(span, inputMessages, getStreamResponseResult(responseState, inputMessages)) + : undefined, + ); + + if (isReadableStream(stream)) { + instrumentReadableStream(stream, span, lifecycle); + return stream; + } + const iterate = stream[Symbol.asyncIterator].bind(stream); - const instrumented = instrumentStreamIterator({ [Symbol.asyncIterator]: iterate }, span); + const instrumented = instrumentStreamIterator({ [Symbol.asyncIterator]: iterate }, span, lifecycle); stream[Symbol.asyncIterator] = () => instrumented; return stream; } +interface StreamLifecycle { + recordChunk: (chunk: unknown) => void; + complete: () => void; + fail: () => void; +} + +interface StreamResponseState { + finalState?: { messages: LangChainMessage[] }; + updateMessages: LangChainMessage[]; +} + +interface ReadableStreamReaderLike { + read: (...args: unknown[]) => Promise>; + cancel?: (reason?: unknown) => Promise; +} + +interface InstrumentableReadableStream extends AsyncIterable { + getReader: (...args: unknown[]) => ReadableStreamReaderLike; + cancel?: (reason?: unknown) => Promise; + pipeThrough?: ( + transform: ReadableWritablePair, + options?: StreamPipeOptions, + ) => ReadableStream; + pipeTo?: (destination: WritableStream, options?: StreamPipeOptions) => Promise; +} + +function createStreamLifecycle( + span: Span, + recordChunk?: (chunk: unknown) => void, + completeResponse?: () => void, +): StreamLifecycle { + let completed = false; + + const complete = (): void => { + if (completed) { + return; + } + + completed = true; + try { + completeResponse?.(); + } finally { + span.end(); + } + }; + + return { + recordChunk(chunk: unknown): void { + recordChunk?.(chunk); + }, + complete, + fail(): void { + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + complete(); + }, + }; +} + +function accumulateStreamResponse(state: StreamResponseState, chunk: unknown): void { + const payload = + Array.isArray(chunk) && chunk.length === 2 && typeof chunk[0] === 'string' ? (chunk[1] as unknown) : chunk; + const directState = getMessageState(payload); + if (directState) { + state.finalState = directState; + return; + } + + if (!payload || typeof payload !== 'object') { + return; + } + + for (const update of Object.values(payload)) { + const updateState = getMessageState(update); + if (updateState) { + state.updateMessages.push(...updateState.messages); + } + } +} + +function getMessageState(value: unknown): { messages: LangChainMessage[] } | undefined { + if (!value || typeof value !== 'object' || !('messages' in value) || !Array.isArray(value.messages)) { + return undefined; + } + + return { messages: value.messages as LangChainMessage[] }; +} + +function getStreamResponseResult( + state: StreamResponseState, + inputMessages: LangChainMessage[] | null, +): { messages: LangChainMessage[] } | undefined { + if (state.finalState) { + return state.finalState; + } + + if (state.updateMessages.length === 0) { + return undefined; + } + + return { messages: [...(inputMessages ?? []), ...state.updateMessages] }; +} + +function isReadableStream(stream: AsyncIterable): stream is InstrumentableReadableStream { + return typeof (stream as Partial).getReader === 'function'; +} + +function instrumentReadableStream(stream: InstrumentableReadableStream, span: Span, lifecycle: StreamLifecycle): void { + const originalGetReader = stream.getReader.bind(stream); + stream.getReader = (...args: unknown[]): ReadableStreamReaderLike => { + const reader = withActiveSpan(span, () => originalGetReader(...args)); + return instrumentReader(reader, span, lifecycle); + }; + + if (stream.cancel) { + const originalCancel = stream.cancel.bind(stream); + stream.cancel = (reason?: unknown): Promise => + completeWithLifecycle( + withActiveSpan(span, () => originalCancel(reason)), + lifecycle, + ); + } + + if (stream.pipeTo) { + const originalPipeTo = stream.pipeTo.bind(stream); + const originalPipeThrough = stream.pipeThrough?.bind(stream); + stream.pipeTo = (destination: WritableStream, options?: StreamPipeOptions): Promise => { + let pipePromise: Promise; + try { + pipePromise = withActiveSpan(span, () => { + if (!originalPipeThrough) { + return originalPipeTo(destination, options); + } + + const passthrough = new TransformStream({ + transform(chunk, controller) { + lifecycle.recordChunk(chunk); + controller.enqueue(chunk); + }, + }); + const outputStream: ReadableStream = originalPipeThrough(passthrough); + return outputStream.pipeTo(destination, options); + }); + } catch (error) { + lifecycle.fail(); + throw error; + } + + return completeWithLifecycle(pipePromise, lifecycle); + }; + } +} + +function instrumentReader( + reader: ReadableStreamReaderLike, + span: Span, + lifecycle: StreamLifecycle, +): ReadableStreamReaderLike { + const originalRead = reader.read.bind(reader); + reader.read = (...args: unknown[]): Promise> => { + const readPromise: Promise> = withActiveSpan(span, () => originalRead(...args)); + return readPromise.then( + result => { + if (result.done) { + lifecycle.complete(); + } else { + lifecycle.recordChunk(result.value); + } + return result; + }, + error => { + lifecycle.fail(); + throw error; + }, + ); + }; + + if (reader.cancel) { + const originalCancel = reader.cancel.bind(reader); + reader.cancel = (reason?: unknown): Promise => + completeWithLifecycle( + withActiveSpan(span, () => originalCancel(reason)), + lifecycle, + ); + } + + return reader; +} + +function completeWithLifecycle(promise: Promise, lifecycle: StreamLifecycle): Promise { + return promise.then( + result => { + lifecycle.complete(); + return result; + }, + error => { + lifecycle.fail(); + throw error; + }, + ); +} + async function* instrumentStreamIterator( stream: AsyncIterable, span: Span, + lifecycle: StreamLifecycle, ): AsyncGenerator { const iterator = stream[Symbol.asyncIterator](); let completed = false; @@ -39,12 +262,14 @@ async function* instrumentStreamIterator( const result = await withActiveSpan(span, () => iterator.next()); if (result.done) { completed = true; + lifecycle.complete(); return; } + lifecycle.recordChunk(result.value); yield result.value; } } catch (error) { - span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + lifecycle.fail(); throw error; } finally { try { @@ -52,7 +277,7 @@ async function* instrumentStreamIterator( await withActiveSpan(span, () => iterator.return?.()); } } finally { - span.end(); + lifecycle.complete(); } } } 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 index 2e81bdd5783a..73c788dfb837 100644 --- a/packages/server-utils/test/ai/lib/tracing/langgraph-stream.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/langgraph-stream.test.ts @@ -2,7 +2,15 @@ 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_STREAMING } from '@sentry/conventions/attributes'; +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(); @@ -10,6 +18,45 @@ 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 { @@ -113,4 +160,116 @@ describe('instrumentStateGraphCompile stream instrumentation', () => { 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('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); + }); }); From 7b86b353922e7e6e7243d53532b77e627509e1df Mon Sep 17 00:00:00 2001 From: "Seongho.Bak" Date: Wed, 26 Aug 2026 13:55:21 +0900 Subject: [PATCH 03/13] fix(server-utils): Preserve LangGraph pipe lifecycle Record source chunks and complete spans across pipeTo and pipeThrough consumption. --- .../src/ai/langgraph/streaming.ts | 59 +++++++++---- .../ai/lib/tracing/langgraph-stream.test.ts | 88 +++++++++++++++++++ 2 files changed, 129 insertions(+), 18 deletions(-) diff --git a/packages/server-utils/src/ai/langgraph/streaming.ts b/packages/server-utils/src/ai/langgraph/streaming.ts index a7bb3fd36165..5e0ed86fff29 100644 --- a/packages/server-utils/src/ai/langgraph/streaming.ts +++ b/packages/server-utils/src/ai/langgraph/streaming.ts @@ -172,31 +172,54 @@ function instrumentReadableStream(stream: InstrumentableReadableStream, span: Sp if (stream.pipeTo) { const originalPipeTo = stream.pipeTo.bind(stream); - const originalPipeThrough = stream.pipeThrough?.bind(stream); - stream.pipeTo = (destination: WritableStream, options?: StreamPipeOptions): Promise => { + const instrumentedPipeTo = (destination: WritableStream, options?: StreamPipeOptions): Promise => { + let destinationWriter: WritableStreamDefaultWriter; + try { + destinationWriter = destination.getWriter(); + } catch (error) { + lifecycle.fail(); + return Promise.reject(error); + } + + const recordingDestination = new WritableStream({ + write(chunk) { + lifecycle.recordChunk(chunk); + return destinationWriter.write(chunk); + }, + close() { + return destinationWriter.close(); + }, + abort(reason) { + return destinationWriter.abort(reason); + }, + }); + let pipePromise: Promise; try { - pipePromise = withActiveSpan(span, () => { - if (!originalPipeThrough) { - return originalPipeTo(destination, options); - } - - const passthrough = new TransformStream({ - transform(chunk, controller) { - lifecycle.recordChunk(chunk); - controller.enqueue(chunk); - }, - }); - const outputStream: ReadableStream = originalPipeThrough(passthrough); - return outputStream.pipeTo(destination, options); - }); + pipePromise = withActiveSpan(span, () => originalPipeTo(recordingDestination, options)); } catch (error) { + destinationWriter.releaseLock(); lifecycle.fail(); - throw error; + return Promise.reject(error); } - return completeWithLifecycle(pipePromise, lifecycle); + return completeWithLifecycle(pipePromise, lifecycle).finally(() => { + destinationWriter.releaseLock(); + }); }; + + stream.pipeTo = instrumentedPipeTo; + + if (stream.pipeThrough) { + stream.pipeThrough = ( + transform: ReadableWritablePair, + options?: StreamPipeOptions, + ): ReadableStream => { + // pipeThrough exposes pipeline failures through the returned readable instead of its internal promise. + void instrumentedPipeTo(transform.writable, options).catch(() => {}); + return transform.readable; + }; + } } } 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 index 73c788dfb837..fbffcc7689b7 100644 --- a/packages/server-utils/test/ai/lib/tracing/langgraph-stream.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/langgraph-stream.test.ts @@ -197,6 +197,94 @@ describe('instrumentStateGraphCompile stream instrumentation', () => { 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('forwards pipe options to the source pipeTo operation', async () => { + const stream = new TestLangGraphStream(['first update']); + const originalPipeTo = stream.pipeTo.bind(stream); + const sourcePipeTo = vi.fn((destination: WritableStream, options?: StreamPipeOptions) => + originalPipeTo(destination, options), + ); + Object.defineProperty(stream, 'pipeTo', { configurable: true, value: sourcePipeTo, writable: true }); + const compiledGraph = { stream: vi.fn().mockResolvedValue(stream) }; + const compile = instrumentStateGraphCompile(() => compiledGraph, {}); + const graph = compile() as { stream: () => Promise> }; + const destination = new WritableStream(); + const options = { preventClose: true, signal: new AbortController().signal }; + + await (await graph.stream()).pipeTo(destination, options); + + expect(sourcePipeTo).toHaveBeenCalledWith(expect.any(WritableStream), options); + expect(destination.locked).toBe(false); + 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 Promise.resolve(); + + expect(spanSetAttribute).toHaveBeenCalledWith( + GEN_AI_RESPONSE_TEXT, + '[{"role":"assistant","content":"Clear skies"}]', + ); + 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) }; From ca153693ef164da4eb7d53ffecf112e174bcee6d Mon Sep 17 00:00:00 2001 From: "Seongho.Bak" Date: Wed, 26 Aug 2026 14:20:44 +0900 Subject: [PATCH 04/13] fix(server-utils): Preserve pipeThrough errors Throw for locked pipeThrough inputs and forward asynchronous pipeline failures through the returned readable. --- .../src/ai/langgraph/streaming.ts | 13 +++++- .../ai/lib/tracing/langgraph-stream.test.ts | 46 +++++++++++++++++++ 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/packages/server-utils/src/ai/langgraph/streaming.ts b/packages/server-utils/src/ai/langgraph/streaming.ts index 5e0ed86fff29..57715ecda602 100644 --- a/packages/server-utils/src/ai/langgraph/streaming.ts +++ b/packages/server-utils/src/ai/langgraph/streaming.ts @@ -65,6 +65,7 @@ interface ReadableStreamReaderLike { } interface InstrumentableReadableStream extends AsyncIterable { + locked: boolean; getReader: (...args: unknown[]) => ReadableStreamReaderLike; cancel?: (reason?: unknown) => Promise; pipeThrough?: ( @@ -215,8 +216,16 @@ function instrumentReadableStream(stream: InstrumentableReadableStream, span: Sp transform: ReadableWritablePair, options?: StreamPipeOptions, ): ReadableStream => { - // pipeThrough exposes pipeline failures through the returned readable instead of its internal promise. - void instrumentedPipeTo(transform.writable, options).catch(() => {}); + if (stream.locked || transform.writable.locked) { + lifecycle.fail(); + throw new TypeError('Cannot pipe through a locked stream.'); + } + + void instrumentedPipeTo(transform.writable, options).catch(error => { + void transform.writable.abort(error).catch(() => { + // The pipe may already have propagated the failure through the transform. + }); + }); return transform.readable; }; } 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 index fbffcc7689b7..700bbec0be9d 100644 --- a/packages/server-utils/test/ai/lib/tracing/langgraph-stream.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/langgraph-stream.test.ts @@ -285,6 +285,52 @@ describe('instrumentStateGraphCompile stream instrumentation', () => { 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); + + 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); + + writer.releaseLock(); + }); + + it('errors the pipeThrough readable when the source pipe rejects', async () => { + const error = new Error('pipe failed'); + const stream = new TestLangGraphStream(['first update']); + Object.defineProperty(stream, 'pipeTo', { + configurable: true, + value: vi.fn().mockRejectedValue(error), + writable: true, + }); + 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) }; From cb32555f0063f196f3c1f435d691cd70cffce914 Mon Sep 17 00:00:00 2001 From: "Seongho.Bak" Date: Wed, 26 Aug 2026 14:31:37 +0900 Subject: [PATCH 05/13] fix(server-utils): Keep span open on pipeThrough preconditions Preserve the native TypeError for locked pipeThrough inputs without reporting an internal stream failure or ending the active span. --- packages/server-utils/src/ai/langgraph/streaming.ts | 1 - .../server-utils/test/ai/lib/tracing/langgraph-stream.test.ts | 4 ++++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/packages/server-utils/src/ai/langgraph/streaming.ts b/packages/server-utils/src/ai/langgraph/streaming.ts index 57715ecda602..4a83958af33e 100644 --- a/packages/server-utils/src/ai/langgraph/streaming.ts +++ b/packages/server-utils/src/ai/langgraph/streaming.ts @@ -217,7 +217,6 @@ function instrumentReadableStream(stream: InstrumentableReadableStream, span: Sp options?: StreamPipeOptions, ): ReadableStream => { if (stream.locked || transform.writable.locked) { - lifecycle.fail(); throw new TypeError('Cannot pipe through a locked stream.'); } 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 index 700bbec0be9d..2b4f9ab29a63 100644 --- a/packages/server-utils/test/ai/lib/tracing/langgraph-stream.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/langgraph-stream.test.ts @@ -294,6 +294,8 @@ describe('instrumentStateGraphCompile stream instrumentation', () => { const reader = instrumentedStream.getReader(); expect(() => instrumentedStream.pipeThrough(new TransformStream())).toThrow(TypeError); + expect(spanSetStatus).not.toHaveBeenCalled(); + expect(spanEnd).not.toHaveBeenCalled(); reader.releaseLock(); }); @@ -308,6 +310,8 @@ describe('instrumentStateGraphCompile stream instrumentation', () => { const instrumentedStream = await graph.stream(); expect(() => instrumentedStream.pipeThrough(transform)).toThrow(TypeError); + expect(spanSetStatus).not.toHaveBeenCalled(); + expect(spanEnd).not.toHaveBeenCalled(); writer.releaseLock(); }); From bf66118f51482e90732f224693b525bdc13b9b96 Mon Sep 17 00:00:00 2001 From: "Seongho.Bak" Date: Tue, 6 Oct 2026 10:14:56 +0900 Subject: [PATCH 06/13] fix(server-utils): Address LangGraph stream review feedback --- .../langgraph/agent-stream-scenario.mjs | 53 ++++++++ .../suites/tracing/langgraph/test.ts | 32 +++++ .../server-utils/src/ai/langgraph/index.ts | 2 + .../src/ai/langgraph/streaming.ts | 33 ++++- .../tracing/langgraph-invoke-agent.test.ts | 52 ++++++- .../ai/lib/tracing/langgraph-stream.test.ts | 14 +- .../lib/tracing/langgraph-streaming.test.ts | 127 ++++++++++++++++++ 7 files changed, 302 insertions(+), 11 deletions(-) create mode 100644 dev-packages/node-integration-tests/suites/tracing/langgraph/agent-stream-scenario.mjs create mode 100644 packages/server-utils/test/ai/lib/tracing/langgraph-streaming.test.ts 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/test.ts b/dev-packages/node-integration-tests/suites/tracing/langgraph/test.ts index e3bce89ef868..08b82fd48cd8 100644 --- a/dev-packages/node-integration-tests/suites/tracing/langgraph/test.ts +++ b/dev-packages/node-integration-tests/suites/tracing/langgraph/test.ts @@ -329,6 +329,38 @@ 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({ transaction: { transaction: 'main' } }) + .expect({ + span: container => { + const agentSpans = container.items.filter( + span => span.attributes[SEMANTIC_ATTRIBUTE_SENTRY_OP]?.value === 'gen_ai.invoke_agent', + ); + const chatSpans = container.items.filter( + span => span.attributes[SEMANTIC_ATTRIBUTE_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(agentSpan.span_id).toMatch(/^[\da-f]{16}$/); + 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/packages/server-utils/src/ai/langgraph/index.ts b/packages/server-utils/src/ai/langgraph/index.ts index 37549e39b0c2..f4dfa36843aa 100644 --- a/packages/server-utils/src/ai/langgraph/index.ts +++ b/packages/server-utils/src/ai/langgraph/index.ts @@ -158,6 +158,8 @@ function instrumentCompiledGraphOperation( streaming && getCurrentScope().getScopeData().sdkProcessingMetadata[LANGGRAPH_INVOKE_ACTIVE] === graphInstrumentationId ) { + // 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); } diff --git a/packages/server-utils/src/ai/langgraph/streaming.ts b/packages/server-utils/src/ai/langgraph/streaming.ts index 4a83958af33e..276a8e9f0f31 100644 --- a/packages/server-utils/src/ai/langgraph/streaming.ts +++ b/packages/server-utils/src/ai/langgraph/streaming.ts @@ -28,7 +28,7 @@ export function instrumentStreamResult>( inputMessages: LangChainMessage[] | null, recordOutputs: boolean | undefined, ): T { - const responseState: StreamResponseState = { updateMessages: [] }; + const responseState: StreamResponseState = { updateMessages: [], messageChunks: [] }; const lifecycle = createStreamLifecycle( span, recordOutputs ? chunk => accumulateStreamResponse(responseState, chunk) : undefined, @@ -57,6 +57,7 @@ interface StreamLifecycle { interface StreamResponseState { finalState?: { messages: LangChainMessage[] }; updateMessages: LangChainMessage[]; + messageChunks: LangChainMessage[]; } interface ReadableStreamReaderLike { @@ -110,6 +111,26 @@ function createStreamLifecycle( function accumulateStreamResponse(state: StreamResponseState, chunk: unknown): void { const payload = Array.isArray(chunk) && chunk.length === 2 && typeof chunk[0] === 'string' ? (chunk[1] as unknown) : chunk; + 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 directState = getMessageState(payload); if (directState) { state.finalState = directState; @@ -144,11 +165,13 @@ function getStreamResponseResult( return state.finalState; } - if (state.updateMessages.length === 0) { + // 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 ?? []), ...state.updateMessages] }; + return { messages: [...(inputMessages ?? []), ...messages] }; } function isReadableStream(stream: AsyncIterable): stream is InstrumentableReadableStream { @@ -174,6 +197,10 @@ function instrumentReadableStream(stream: InstrumentableReadableStream, span: Sp if (stream.pipeTo) { const originalPipeTo = stream.pipeTo.bind(stream); const instrumentedPipeTo = (destination: WritableStream, options?: StreamPipeOptions): Promise => { + if (stream.locked || destination.locked) { + return Promise.reject(new TypeError('Cannot pipe to or from a locked stream.')); + } + let destinationWriter: WritableStreamDefaultWriter; try { destinationWriter = destination.getWriter(); 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 1e8a204bd85e..bd78cfd89b10 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 @@ -1,7 +1,7 @@ import { afterEach, beforeEach, describe, expect, it } from 'vitest'; import { getMainCarrier, setCurrentClient, spanToStaticSpanJSON } from '@sentry/core'; import type { Span } from '@sentry/core'; -import { instrumentCompiledGraphInvoke } from '../../../../src/ai/langgraph'; +import { instrumentCompiledGraphInvoke, instrumentStateGraphCompile } from '../../../../src/ai/langgraph'; import type { CompiledGraph } from '../../../../src/ai/langgraph/types'; import { getDefaultTestClientOptions, TestClient } from '../../../mocks/client'; @@ -54,3 +54,53 @@ describe('LangGraph invoke_agent span names', () => { expect(spanToStaticSpanJSON(endedSpans[0]!).description).toBe('invoke_agent'); }); }); + +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); + }); +}); 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 index 2b4f9ab29a63..8aff9f7deca5 100644 --- a/packages/server-utils/test/ai/lib/tracing/langgraph-stream.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/langgraph-stream.test.ts @@ -276,13 +276,13 @@ describe('instrumentStateGraphCompile stream instrumentation', () => { const transformed = (await graph.stream()).pipeThrough(new TransformStream()); await transformed.pipeTo(new WritableStream()); - await Promise.resolve(); - - expect(spanSetAttribute).toHaveBeenCalledWith( - GEN_AI_RESPONSE_TEXT, - '[{"role":"assistant","content":"Clear skies"}]', - ); - expect(spanEnd).toHaveBeenCalledTimes(1); + 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 () => { 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..44c63eecb4bc --- /dev/null +++ b/packages/server-utils/test/ai/lib/tracing/langgraph-streaming.test.ts @@ -0,0 +1,127 @@ +import { describe, expect, it, vi } from 'vitest'; +import type { Span } 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 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(); + }); +}); From a78a23e76b666b3e7d6c28f901dcea1a11896059 Mon Sep 17 00:00:00 2001 From: "Seongho.Bak" Date: Tue, 6 Oct 2026 11:19:12 +0900 Subject: [PATCH 07/13] fix(server-utils): Preserve spans on stream caller errors --- .../src/ai/langgraph/streaming.ts | 18 +++--- .../lib/tracing/langgraph-streaming.test.ts | 64 +++++++++++++++++++ 2 files changed, 73 insertions(+), 9 deletions(-) diff --git a/packages/server-utils/src/ai/langgraph/streaming.ts b/packages/server-utils/src/ai/langgraph/streaming.ts index 276a8e9f0f31..a4859e4635aa 100644 --- a/packages/server-utils/src/ai/langgraph/streaming.ts +++ b/packages/server-utils/src/ai/langgraph/streaming.ts @@ -188,24 +188,24 @@ function instrumentReadableStream(stream: InstrumentableReadableStream, span: Sp if (stream.cancel) { const originalCancel = stream.cancel.bind(stream); stream.cancel = (reason?: unknown): Promise => - completeWithLifecycle( - withActiveSpan(span, () => originalCancel(reason)), - lifecycle, - ); + stream.locked + ? originalCancel(reason) + : completeWithLifecycle( + withActiveSpan(span, () => originalCancel(reason)), + lifecycle, + ); } if (stream.pipeTo) { const originalPipeTo = stream.pipeTo.bind(stream); const instrumentedPipeTo = (destination: WritableStream, options?: StreamPipeOptions): Promise => { - if (stream.locked || destination.locked) { - return Promise.reject(new TypeError('Cannot pipe to or from a locked stream.')); - } - let destinationWriter: WritableStreamDefaultWriter; try { + if (stream.locked || destination.locked) { + return Promise.reject(new TypeError('Cannot pipe to or from a locked stream.')); + } destinationWriter = destination.getWriter(); } catch (error) { - lifecycle.fail(); return Promise.reject(error); } 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 index 44c63eecb4bc..c107ff043d73 100644 --- a/packages/server-utils/test/ai/lib/tracing/langgraph-streaming.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/langgraph-streaming.test.ts @@ -125,3 +125,67 @@ describe('LangGraph pipeTo preconditions', () => { 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(); + }); +}); From 539657443ba4180362386beac0e26b175c5df92b Mon Sep 17 00:00:00 2001 From: "Seongho.Bak" Date: Tue, 6 Oct 2026 14:53:28 +0900 Subject: [PATCH 08/13] fix(server-utils): Preserve stream outcomes on recording failures --- .../src/ai/langgraph/streaming.ts | 2 + .../langgraph-stream-lifecycle.test.ts | 45 +++++++++++++++++++ 2 files changed, 47 insertions(+) create mode 100644 packages/server-utils/test/ai/lib/tracing/langgraph-stream-lifecycle.test.ts diff --git a/packages/server-utils/src/ai/langgraph/streaming.ts b/packages/server-utils/src/ai/langgraph/streaming.ts index 55b1cb4c68e6..f2e63ee30ab4 100644 --- a/packages/server-utils/src/ai/langgraph/streaming.ts +++ b/packages/server-utils/src/ai/langgraph/streaming.ts @@ -90,6 +90,8 @@ function createStreamLifecycle( completed = true; try { completeResponse?.(); + } catch { + // Recording failures must not replace the application's stream outcome. } finally { span.end(); } 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..37fb3f00c783 --- /dev/null +++ b/packages/server-utils/test/ai/lib/tracing/langgraph-stream-lifecycle.test.ts @@ -0,0 +1,45 @@ +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'; + +describe('LangGraph response recording failures', () => { + it.each([false, true])('preserves the stream outcome when serialization fails (stream error: %s)', async fails => { + const span = { + end: vi.fn(), + setAttribute: vi.fn(), + setStatus: vi.fn(), + spanContext: () => ({ traceId: 'a'.repeat(32), spanId: 'b'.repeat(16), traceFlags: 1 }), + }; + 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); + }); +}); From 29855c7df6b7eed67e1e39e51a5764f023d5373a Mon Sep 17 00:00:00 2001 From: RulaKhaled Date: Wed, 21 Oct 2026 15:33:14 +0200 Subject: [PATCH 09/13] fix(server-utils): Address LangGraph stream review feedback Wrap each compiled graph once when manual and auto instrumentation both see it, stop closing a source iterator that already threw, and keep streaming when recording a chunk fails. Add a LangGraph v1 stream suite. Co-Authored-By: Claude Opus 5.5 --- .../tracing/langgraph/v1/scenario-stream.mjs | 30 ++++ .../suites/tracing/langgraph/v1/test.ts | 48 ++++++ packages/server-utils/src/ai/core/utils.ts | 5 + .../server-utils/src/ai/langchain/utils.ts | 2 +- .../server-utils/src/ai/langgraph/index.ts | 142 +++++++----------- .../src/ai/langgraph/streaming.ts | 124 +++++++-------- .../server-utils/src/ai/mistral/streaming.ts | 8 +- .../src/integrations/langgraph.ts | 49 +----- .../tracing/langgraph-invoke-agent.test.ts | 41 ++++- .../langgraph-stream-lifecycle.test.ts | 102 ++++++++++++- 10 files changed, 338 insertions(+), 213 deletions(-) create mode 100644 dev-packages/node-integration-tests/suites/tracing/langgraph/v1/scenario-stream.mjs 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..e9b1c5050389 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,50 @@ 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', + }, + }, + ); }); diff --git a/packages/server-utils/src/ai/core/utils.ts b/packages/server-utils/src/ai/core/utils.ts index 93949ec40eca..ee934b68d38c 100644 --- a/packages/server-utils/src/ai/core/utils.ts +++ b/packages/server-utils/src/ai/core/utils.ts @@ -88,6 +88,11 @@ export function buildMethodPath(currentPath: string, prop: string): string { return currentPath ? `${currentPath}.${prop}` : prop; } +/** 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 cea5e0df9892..ca4f9631e86c 100644 --- a/packages/server-utils/src/ai/langgraph/index.ts +++ b/packages/server-utils/src/ai/langgraph/index.ts @@ -14,7 +14,7 @@ import { 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 { @@ -24,7 +24,7 @@ import { normalizeLangChainMessages, } from '../langchain/utils'; import { LANGGRAPH_ORIGIN } from './constants'; -import { getGraphInstrumentationId, instrumentStreamResult, isAsyncIterable } from './streaming'; +import { instrumentStreamResult } from './streaming'; import type { CompiledGraph, LangGraphOptions } from './types'; import { extractAgentNameFromParams, @@ -37,8 +37,23 @@ import { let _insideCreateReactAgent = false; 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() and stream() * methods with a `gen_ai.invoke_agent` span. @@ -65,31 +80,7 @@ export function instrumentStateGraphCompile( const compiledGraph = Reflect.apply(target, thisArg, args); const compileOptions = args.length > 0 ? (args[0] as Record) : {}; - - // Instrument agent methods 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, - ); - } - - const originalStream = compiledGraph.stream; - if (originalStream && typeof originalStream === 'function') { - compiledGraph.stream = instrumentCompiledGraphStream( - originalStream.bind(compiledGraph), - compiledGraph, - compileOptions, - options, - undefined, - sentryHandler, - ); - } + instrumentCompiledGraph(compiledGraph, compileOptions, options, undefined, sentryHandler); return compiledGraph; }, @@ -99,6 +90,38 @@ export function instrumentStateGraphCompile( return wrapped; } +/** + * Wraps a compiled graph's `invoke()` and `stream()` with `gen_ai.invoke_agent` spans, once per graph. + */ +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', + ); + } + } +} + /** * Instruments CompiledGraph's invoke method to create spans for agent invocation * @@ -123,25 +146,6 @@ export function instrumentCompiledGraphInvoke( ); } -export function instrumentCompiledGraphStream( - originalStream: (...args: unknown[]) => Promise>, - graphInstance: CompiledGraph, - compileOptions: Record, - options: LangGraphOptions, - llm?: BaseChatModel | null, - sentryCallbackHandler?: unknown, -): (...args: unknown[]) => Promise> { - return instrumentCompiledGraphOperation( - originalStream, - graphInstance, - compileOptions, - options, - llm, - sentryCallbackHandler, - true, - ) as (...args: unknown[]) => Promise>; -} - function instrumentCompiledGraphOperation( originalOperation: (...args: unknown[]) => Promise, graphInstance: CompiledGraph, @@ -151,13 +155,13 @@ function instrumentCompiledGraphOperation( sentryCallbackHandler: unknown, streaming: boolean, ): (...args: unknown[]) => Promise { - const graphInstrumentationId = getGraphInstrumentationId(graphInstance); + const invokeMarker = getInvokeMarker(graphInstance); return new Proxy(originalOperation, { apply(target, thisArg, args: unknown[]): Promise { if ( streaming && - getCurrentScope().getScopeData().sdkProcessingMetadata[LANGGRAPH_INVOKE_ACTIVE] === graphInstrumentationId + 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 }); @@ -244,11 +248,7 @@ function instrumentCompiledGraphOperation( } if (!streaming) { - // LangGraph implements invoke() by consuming stream(); the shared async scope keeps that internal - // call from producing a nested duplicate while leaving direct concurrent stream() calls unaffected. - getCurrentScope().setSDKProcessingMetadata({ - [LANGGRAPH_INVOKE_ACTIVE]: graphInstrumentationId, - }); + getCurrentScope().setSDKProcessingMetadata({ [LANGGRAPH_INVOKE_ACTIVE]: invokeMarker }); } const result = await Reflect.apply(target, thisArg, args); @@ -317,40 +317,8 @@ export function instrumentCreateReactAgent( _insideCreateReactAgent = false; } - // Wrap agent methods 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 originalStream = compiledGraph.stream; - if (originalStream && typeof originalStream === 'function') { - const compileOptions: Record = {}; - if (agentName) { - compileOptions.name = agentName; - } - - compiledGraph.stream = instrumentCompiledGraphStream( - originalStream.bind(compiledGraph), - 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 index f2e63ee30ab4..e713763b06f7 100644 --- a/packages/server-utils/src/ai/langgraph/streaming.ts +++ b/packages/server-utils/src/ai/langgraph/streaming.ts @@ -1,49 +1,25 @@ -import { SPAN_STATUS_ERROR, withActiveSpan } from '@sentry/core'; -import type { Span } from '@sentry/core'; +import { SPAN_STATUS_ERROR, withActiveSpan, withScope } from '@sentry/core'; +import type { Scope, Span } from '@sentry/core'; import type { LangChainMessage } from '../langchain/types'; -import type { CompiledGraph } from './types'; import { setResponseAttributes } from './utils'; -const graphInstrumentationIds = new WeakMap(); -let nextGraphInstrumentationId = 0; - -export function getGraphInstrumentationId(graph: CompiledGraph): number { - const existingId = graphInstrumentationIds.get(graph); - if (existingId !== undefined) { - return existingId; - } - - const id = nextGraphInstrumentationId++; - graphInstrumentationIds.set(graph, id); - return id; -} - -export function isAsyncIterable(value: unknown): value is AsyncIterable { - return !!value && typeof (value as AsyncIterable)[Symbol.asyncIterator] === 'function'; -} - export function instrumentStreamResult>( stream: T, span: Span, inputMessages: LangChainMessage[] | null, recordOutputs: boolean | undefined, ): T { - const responseState: StreamResponseState = { updateMessages: [], messageChunks: [] }; - const lifecycle = createStreamLifecycle( - span, - recordOutputs ? chunk => accumulateStreamResponse(responseState, chunk) : undefined, - recordOutputs - ? () => setResponseAttributes(span, inputMessages, getStreamResponseResult(responseState, inputMessages)) - : undefined, - ); + // 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); if (isReadableStream(stream)) { - instrumentReadableStream(stream, span, lifecycle); + instrumentReadableStream(stream, spanScope, lifecycle); return stream; } const iterate = stream[Symbol.asyncIterator].bind(stream); - stream[Symbol.asyncIterator] = () => instrumentStreamIterator({ [Symbol.asyncIterator]: iterate }, span, lifecycle); + stream[Symbol.asyncIterator] = () => instrumentStreamIterator(iterate(), spanScope, lifecycle); return stream; } @@ -54,7 +30,7 @@ interface StreamLifecycle { } interface StreamResponseState { - finalState?: { messages: LangChainMessage[] }; + finalMessages?: LangChainMessage[]; updateMessages: LangChainMessage[]; messageChunks: LangChainMessage[]; } @@ -77,9 +53,10 @@ interface InstrumentableReadableStream extends AsyncIterable { function createStreamLifecycle( span: Span, - recordChunk?: (chunk: unknown) => void, - completeResponse?: () => void, + inputMessages: LangChainMessage[] | null, + recordOutputs: boolean | undefined, ): StreamLifecycle { + let state: StreamResponseState | undefined = recordOutputs ? { updateMessages: [], messageChunks: [] } : undefined; let completed = false; const complete = (): void => { @@ -89,17 +66,27 @@ function createStreamLifecycle( completed = true; try { - completeResponse?.(); + if (state) { + setResponseAttributes(span, inputMessages, getStreamResponseResult(state, inputMessages)); + } } 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; span.end(); } }; return { recordChunk(chunk: unknown): void { - recordChunk?.(chunk); + try { + if (state) { + accumulateStreamResponse(state, chunk); + } + } catch { + // Same as `complete`: recording must never break the application's stream. + } }, complete, fail(): void { @@ -132,9 +119,9 @@ function accumulateStreamResponse(state: StreamResponseState, chunk: unknown): v } return; } - const directState = getMessageState(payload); - if (directState) { - state.finalState = directState; + const directMessages = getMessages(payload); + if (directMessages) { + state.finalMessages = directMessages; return; } @@ -143,27 +130,24 @@ function accumulateStreamResponse(state: StreamResponseState, chunk: unknown): v } for (const update of Object.values(payload)) { - const updateState = getMessageState(update); - if (updateState) { - state.updateMessages.push(...updateState.messages); - } + state.updateMessages.push(...(getMessages(update) ?? [])); } } -function getMessageState(value: unknown): { messages: LangChainMessage[] } | undefined { +function getMessages(value: unknown): LangChainMessage[] | undefined { if (!value || typeof value !== 'object' || !('messages' in value) || !Array.isArray(value.messages)) { return undefined; } - return { messages: value.messages as LangChainMessage[] }; + return value.messages as LangChainMessage[]; } function getStreamResponseResult( state: StreamResponseState, inputMessages: LangChainMessage[] | null, ): { messages: LangChainMessage[] } | undefined { - if (state.finalState) { - return state.finalState; + if (state.finalMessages) { + return { messages: state.finalMessages }; } // Updates contain complete messages, including the same responses emitted as token chunks. @@ -179,11 +163,15 @@ function isReadableStream(stream: AsyncIterable): stream is Instrumenta return typeof (stream as Partial).getReader === 'function'; } -function instrumentReadableStream(stream: InstrumentableReadableStream, span: Span, lifecycle: StreamLifecycle): void { +function instrumentReadableStream( + stream: InstrumentableReadableStream, + scope: Scope, + lifecycle: StreamLifecycle, +): void { const originalGetReader = stream.getReader.bind(stream); stream.getReader = (...args: unknown[]): ReadableStreamReaderLike => { - const reader = withActiveSpan(span, () => originalGetReader(...args)); - return instrumentReader(reader, span, lifecycle); + const reader = withScope(scope, () => originalGetReader(...args)); + return instrumentReader(reader, scope, lifecycle); }; if (stream.cancel) { @@ -192,7 +180,7 @@ function instrumentReadableStream(stream: InstrumentableReadableStream, span: Sp stream.locked ? originalCancel(reason) : completeWithLifecycle( - withActiveSpan(span, () => originalCancel(reason)), + withScope(scope, () => originalCancel(reason)), lifecycle, ); } @@ -225,7 +213,7 @@ function instrumentReadableStream(stream: InstrumentableReadableStream, span: Sp let pipePromise: Promise; try { - pipePromise = withActiveSpan(span, () => originalPipeTo(recordingDestination, options)); + pipePromise = withScope(scope, () => originalPipeTo(recordingDestination, options)); } catch (error) { destinationWriter.releaseLock(); lifecycle.fail(); @@ -261,12 +249,12 @@ function instrumentReadableStream(stream: InstrumentableReadableStream, span: Sp function instrumentReader( reader: ReadableStreamReaderLike, - span: Span, + scope: Scope, lifecycle: StreamLifecycle, ): ReadableStreamReaderLike { const originalRead = reader.read.bind(reader); reader.read = (...args: unknown[]): Promise> => { - const readPromise: Promise> = withActiveSpan(span, () => originalRead(...args)); + const readPromise: Promise> = withScope(scope, () => originalRead(...args)); return readPromise.then( result => { if (result.done) { @@ -287,7 +275,7 @@ function instrumentReader( const originalCancel = reader.cancel.bind(reader); reader.cancel = (reason?: unknown): Promise => completeWithLifecycle( - withActiveSpan(span, () => originalCancel(reason)), + withScope(scope, () => originalCancel(reason)), lifecycle, ); } @@ -309,31 +297,35 @@ function completeWithLifecycle(promise: Promise, lifecycle: StreamLifecycl } async function* instrumentStreamIterator( - stream: AsyncIterable, - span: Span, + iterator: AsyncIterator, + scope: Scope, lifecycle: StreamLifecycle, ): AsyncGenerator { - const iterator = stream[Symbol.asyncIterator](); - let completed = false; + // Like `for await`, only close the source when the consumer stops early, not after it threw or finished. + let sourceDone = false; try { while (true) { - const result = await withActiveSpan(span, () => iterator.next()); + let result: IteratorResult; + try { + result = await withScope(scope, () => iterator.next()); + } catch (error) { + sourceDone = true; + lifecycle.fail(); + throw error; + } if (result.done) { - completed = true; + sourceDone = true; lifecycle.complete(); return; } lifecycle.recordChunk(result.value); yield result.value; } - } catch (error) { - lifecycle.fail(); - throw error; } finally { try { - if (!completed) { - await withActiveSpan(span, () => iterator.return?.()); + if (!sourceDone) { + await withScope(scope, () => iterator.return?.()); } } finally { lifecycle.complete(); diff --git a/packages/server-utils/src/ai/mistral/streaming.ts b/packages/server-utils/src/ai/mistral/streaming.ts index 74acc6e43b70..bdd3f6abfe27 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, setOutputMessagesAttribute } from '../core/utils'; import type { MistralCompletionChunk, MistralToolCall } from './types'; import { contentToString } from './utils'; @@ -22,7 +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 }>; @@ -37,11 +36,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: [], diff --git a/packages/server-utils/src/integrations/langgraph.ts b/packages/server-utils/src/integrations/langgraph.ts index 3a375d2a8bd9..662621442dfc 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, instrumentCompiledGraphStream } from '../ai/langgraph'; +import { instrumentCompiledGraph } 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'; @@ -46,7 +46,7 @@ function instrumentLanggraph(options: LangGraphOptions): void { return; } const { arguments: args, result } = message as OrchestrionChannelContext; - wrapCompiledGraphMethods(result, getFirstArgObject(args) ?? {}, resolvedOptions, null, sentryHandler); + instrumentCompiledGraph(result, getFirstArgObject(args) ?? {}, resolvedOptions, null, sentryHandler); }); // createReactAgent only wraps tools and the returned graph's agent methods. Tools are wrapped at @@ -75,7 +75,7 @@ function instrumentLanggraph(options: LangGraphOptions): void { const { arguments: args, result } = message as OrchestrionChannelContext; const agentName = extractAgentNameFromParams(args) ?? undefined; const compileOptions = agentName ? { name: agentName } : {}; - wrapCompiledGraphMethods(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,47 +89,6 @@ function getFirstArgObject(args: unknown[] | undefined): Record return isObjectLike(first) ? first : undefined; } -/** - * Wrap the compiled graph's agent methods with the shared `invoke_agent` instrumentation, exactly as - * the OTel path does on the returned graph. - */ -function wrapCompiledGraphMethods( - 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, - ); - } - - const originalStream = compiledGraph.stream; - if (typeof originalStream === 'function') { - compiledGraph.stream = instrumentCompiledGraphStream( - originalStream.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 a6bd5ed14554..5cdf7e370ab9 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 @@ -9,7 +9,11 @@ import { } from '@sentry/core'; import type { Span } from '@sentry/core'; import { GEN_AI_CONVERSATION_ID } from '@sentry/conventions/attributes'; -import { instrumentCompiledGraphInvoke, instrumentStateGraphCompile } from '../../../../src/ai/langgraph'; +import { + instrumentCompiledGraph, + instrumentCompiledGraphInvoke, + instrumentStateGraphCompile, +} from '../../../../src/ai/langgraph'; import type { CompiledGraph } from '../../../../src/ai/langgraph/types'; import { wrapToolsWithSpans } from '../../../../src/ai/langgraph/utils'; import { getDefaultTestClientOptions, TestClient } from '../../../mocks/client'; @@ -190,4 +194,39 @@ describe('LangGraph invoke stream suppression', () => { 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); + }); }); 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 index 37fb3f00c783..47c09dc08c53 100644 --- 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 @@ -3,14 +3,18 @@ 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 = { - end: vi.fn(), - setAttribute: vi.fn(), - setStatus: vi.fn(), - spanContext: () => ({ traceId: 'a'.repeat(32), spanId: 'b'.repeat(16), traceFlags: 1 }), - }; + const span = createSpan(); const args: Record = {}; args.self = args; const chunk = { @@ -42,4 +46,90 @@ describe('LangGraph response recording failures', () => { } 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); + }); }); From 284729767d2527e5a1a242c6d8d207e47dd5ba4a Mon Sep 17 00:00:00 2001 From: RulaKhaled Date: Wed, 21 Oct 2026 15:51:20 +0200 Subject: [PATCH 10/13] fix(server-utils): Harden LangGraph stream span lifecycle Record only the root graph's output from subgraph and multi-mode chunks, treat a caller's cancel as a finished run, and route tee/pipeTo/ pipeThrough through the instrumented reader (shared with Mistral) so those consumers end the span. Manual instrumentation now wins over the integration, and an unpatchable stream is returned to the caller. Co-Authored-By: Claude Opus 5.5 --- .../langgraph/v1/scenario-stream-edge.mjs | 52 +++++++ .../suites/tracing/langgraph/v1/test.ts | 38 +++++ packages/server-utils/src/ai/core/utils.ts | 32 +++++ .../server-utils/src/ai/langgraph/index.ts | 54 ++++---- .../src/ai/langgraph/streaming.ts | 130 +++++++++--------- .../server-utils/src/ai/mistral/streaming.ts | 46 ++----- .../src/integrations/langgraph.ts | 10 +- .../tracing/langgraph-invoke-agent.test.ts | 56 ++++++-- .../ai/lib/tracing/langgraph-stream.test.ts | 33 ++--- .../lib/tracing/langgraph-streaming.test.ts | 120 ++++++++++++---- 10 files changed, 385 insertions(+), 186 deletions(-) create mode 100644 dev-packages/node-integration-tests/suites/tracing/langgraph/v1/scenario-stream-edge.mjs 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/test.ts b/dev-packages/node-integration-tests/suites/tracing/langgraph/v1/test.ts index e9b1c5050389..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 @@ -133,4 +133,42 @@ describe('LangGraph integration (v1)', () => { }, }, ); + + 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 ee934b68d38c..a781d60252e0 100644 --- a/packages/server-utils/src/ai/core/utils.ts +++ b/packages/server-utils/src/ai/core/utils.ts @@ -88,6 +88,38 @@ export function buildMethodPath(currentPath: string, prop: string): string { return currentPath ? `${currentPath}.${prop}` : prop; } +/** + * A stream that pulls through `reader`, typically an instrumented one. + * + * `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. + */ +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'; diff --git a/packages/server-utils/src/ai/langgraph/index.ts b/packages/server-utils/src/ai/langgraph/index.ts index ca4f9631e86c..4aec034f6634 100644 --- a/packages/server-utils/src/ai/langgraph/index.ts +++ b/packages/server-utils/src/ai/langgraph/index.ts @@ -35,6 +35,24 @@ import { } from './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 @@ -78,7 +96,7 @@ 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) : {}; instrumentCompiledGraph(compiledGraph, compileOptions, options, undefined, sentryHandler); @@ -122,30 +140,6 @@ export function instrumentCompiledGraph( } } -/** - * Instruments CompiledGraph's invoke method to create spans for agent invocation - * - * Creates a `gen_ai.invoke_agent` span when invoke() is called - */ -export function instrumentCompiledGraphInvoke( - originalInvoke: (...args: unknown[]) => Promise, - graphInstance: CompiledGraph, - compileOptions: Record, - options: LangGraphOptions, - llm?: BaseChatModel | null, - sentryCallbackHandler?: unknown, -): (...args: unknown[]) => Promise { - return instrumentCompiledGraphOperation( - originalInvoke, - graphInstance, - compileOptions, - options, - llm, - sentryCallbackHandler, - false, - ); -} - function instrumentCompiledGraphOperation( originalOperation: (...args: unknown[]) => Promise, graphInstance: CompiledGraph, @@ -256,7 +250,13 @@ function instrumentCompiledGraphOperation( if (streaming) { if (isAsyncIterable(result)) { span.setAttribute(GEN_AI_RESPONSE_STREAMING, true); - return instrumentStreamResult(result, span, inputMessages, recordOutputs); + try { + return instrumentStreamResult(result, span, inputMessages, recordOutputs); + } catch { + // An unpatchable stream (e.g. frozen) still belongs to the caller; only the span is lost. + span.end(); + return result; + } } span.end(); @@ -312,7 +312,7 @@ 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; } diff --git a/packages/server-utils/src/ai/langgraph/streaming.ts b/packages/server-utils/src/ai/langgraph/streaming.ts index e713763b06f7..2afc970e37ca 100644 --- a/packages/server-utils/src/ai/langgraph/streaming.ts +++ b/packages/server-utils/src/ai/langgraph/streaming.ts @@ -1,5 +1,6 @@ import { 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'; @@ -27,8 +28,12 @@ interface StreamLifecycle { recordChunk: (chunk: unknown) => void; complete: () => void; fail: () => void; + markCancelled: () => 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[]; @@ -44,6 +49,7 @@ interface InstrumentableReadableStream extends AsyncIterable { locked: boolean; getReader: (...args: unknown[]) => ReadableStreamReaderLike; cancel?: (reason?: unknown) => Promise; + tee?: () => [ReadableStream, ReadableStream]; pipeThrough?: ( transform: ReadableWritablePair, options?: StreamPipeOptions, @@ -58,6 +64,8 @@ function createStreamLifecycle( ): StreamLifecycle { let state: StreamResponseState | undefined = recordOutputs ? { updateMessages: [], messageChunks: [] } : undefined; 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) { @@ -90,15 +98,40 @@ function createStreamLifecycle( }, complete, fail(): void { - span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + if (!cancelled && !completed) { + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + } complete(); }, + markCancelled(): void { + cancelled = true; + }, }; } +// 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 || (typeof mode === 'string' && !OUTPUT_STREAM_MODES.has(mode))) { + return undefined; + } + + return { payload: chunk[chunk.length - 1] }; +} + function accumulateStreamResponse(state: StreamResponseState, chunk: unknown): void { - const payload = - Array.isArray(chunk) && chunk.length === 2 && typeof chunk[0] === 'string' ? (chunk[1] as unknown) : chunk; + 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 && @@ -125,10 +158,6 @@ function accumulateStreamResponse(state: StreamResponseState, chunk: unknown): v return; } - if (!payload || typeof payload !== 'object') { - return; - } - for (const update of Object.values(payload)) { state.updateMessages.push(...(getMessages(update) ?? [])); } @@ -176,74 +205,47 @@ function instrumentReadableStream( if (stream.cancel) { const originalCancel = stream.cancel.bind(stream); - stream.cancel = (reason?: unknown): Promise => - stream.locked + stream.cancel = (reason?: unknown): Promise => { + lifecycle.markCancelled(); + return stream.locked ? originalCancel(reason) : completeWithLifecycle( withScope(scope, () => originalCancel(reason)), lifecycle, ); + }; } - if (stream.pipeTo) { - const originalPipeTo = stream.pipeTo.bind(stream); - const instrumentedPipeTo = (destination: WritableStream, options?: StreamPipeOptions): Promise => { - let destinationWriter: WritableStreamDefaultWriter; - try { - if (stream.locked || destination.locked) { - return Promise.reject(new TypeError('Cannot pipe to or from a locked stream.')); - } - destinationWriter = destination.getWriter(); - } catch (error) { - return Promise.reject(error); - } + const readerSource = (): ReadableStream => + readableFromReader(stream.getReader()); + // 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): ReadableStream => { + if (stream.locked || typeof destination?.getWriter !== 'function' || destination.locked) { + throw new TypeError('Cannot pipe a locked stream, or to a missing or locked destination.'); + } + return readerSource(); + }; + + if (stream.tee) { + stream.tee = () => readerSource().tee(); + } - const recordingDestination = new WritableStream({ - write(chunk) { - lifecycle.recordChunk(chunk); - return destinationWriter.write(chunk); - }, - close() { - return destinationWriter.close(); - }, - abort(reason) { - return destinationWriter.abort(reason); - }, - }); - - let pipePromise: Promise; + if (stream.pipeTo) { + stream.pipeTo = (destination: WritableStream, options?: StreamPipeOptions): Promise => { try { - pipePromise = withScope(scope, () => originalPipeTo(recordingDestination, options)); + return instrumentedSource(destination).pipeTo(destination, options); } catch (error) { - destinationWriter.releaseLock(); - lifecycle.fail(); return Promise.reject(error); } - - return completeWithLifecycle(pipePromise, lifecycle).finally(() => { - destinationWriter.releaseLock(); - }); }; + } - stream.pipeTo = instrumentedPipeTo; - - if (stream.pipeThrough) { - stream.pipeThrough = ( - transform: ReadableWritablePair, - options?: StreamPipeOptions, - ): ReadableStream => { - if (stream.locked || transform.writable.locked) { - throw new TypeError('Cannot pipe through a locked stream.'); - } - - void instrumentedPipeTo(transform.writable, options).catch(error => { - void transform.writable.abort(error).catch(() => { - // The pipe may already have propagated the failure through the transform. - }); - }); - return transform.readable; - }; - } + if (stream.pipeThrough) { + stream.pipeThrough = ( + transform: ReadableWritablePair, + options?: StreamPipeOptions, + ): ReadableStream => instrumentedSource(transform?.writable).pipeThrough(transform, options); } } @@ -273,11 +275,13 @@ function instrumentReader( if (reader.cancel) { const originalCancel = reader.cancel.bind(reader); - reader.cancel = (reason?: unknown): Promise => - completeWithLifecycle( + reader.cancel = (reason?: unknown): Promise => { + lifecycle.markCancelled(); + return completeWithLifecycle( withScope(scope, () => originalCancel(reason)), lifecycle, ); + }; } return reader; diff --git a/packages/server-utils/src/ai/mistral/streaming.ts b/packages/server-utils/src/ai/mistral/streaming.ts index bdd3f6abfe27..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, isAsyncIterable, setOutputMessagesAttribute } from '../core/utils'; +import { endStreamSpan, isAsyncIterable, readableFromReader, setOutputMessagesAttribute } from '../core/utils'; import type { MistralCompletionChunk, MistralToolCall } from './types'; import { contentToString } from './utils'; @@ -22,7 +22,6 @@ interface StreamingState { /** Which drain path owns accumulation. See {@link instrumentEventStream}. */ type StreamConsumer = 'iterator' | 'reader'; - interface StreamReaderLike { read: () => Promise<{ done: boolean; value?: unknown }>; cancel?: (reason?: unknown) => Promise; @@ -181,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. @@ -299,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 662621442dfc..cdd804065f58 100644 --- a/packages/server-utils/src/integrations/langgraph.ts +++ b/packages/server-utils/src/integrations/langgraph.ts @@ -3,7 +3,7 @@ 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 { instrumentCompiledGraph } from '../ai/langgraph'; +import { instrumentCompiledGraph, isInsideManualLangGraphInstrumentation } from '../ai/langgraph'; import { LANGGRAPH_INTEGRATION_NAME } from '../ai/langgraph/constants'; import type { LangGraphOptions } from '../ai/langgraph/types'; import { extractAgentNameFromParams, extractLLMFromParams, wrapToolsWithSpans } from '../ai/langgraph/utils'; @@ -42,7 +42,7 @@ function instrumentLanggraph(options: LangGraphOptions): void { diagnosticsChannel .tracingChannel(CHANNELS.LANGGRAPH_STATE_GRAPH_COMPILE) .end.subscribe(message => { - if (insideCreateReactAgent) { + if (insideCreateReactAgent || isInsideManualLangGraphInstrumentation()) { return; } const { arguments: args, result } = message as OrchestrionChannelContext; @@ -60,6 +60,9 @@ function instrumentLanggraph(options: LangGraphOptions): void { // throw) and cleared in both `end` and `error`, so it can neither stick on across calls nor // stay off during this call's nested compile. Tool wrapping is guarded for the same reason. insideCreateReactAgent = true; + if (isInsideManualLangGraphInstrumentation()) { + return; + } try { const { arguments: args } = message as OrchestrionChannelContext; const params = getFirstArgObject(args); @@ -72,6 +75,9 @@ 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 } : {}; 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 5cdf7e370ab9..28c760846bd5 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 @@ -11,13 +11,22 @@ import type { Span } from '@sentry/core'; import { GEN_AI_CONVERSATION_ID } from '@sentry/conventions/attributes'; import { instrumentCompiledGraph, - instrumentCompiledGraphInvoke, instrumentStateGraphCompile, + isInsideManualLangGraphInstrumentation, } from '../../../../src/ai/langgraph'; -import type { CompiledGraph } from '../../../../src/ai/langgraph/types'; 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; @@ -45,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); } @@ -96,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]!; } @@ -117,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' } }); @@ -229,4 +230,33 @@ describe('LangGraph invoke stream suppression', () => { 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); + }); }); 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 index 8aff9f7deca5..a5983ebdc973 100644 --- a/packages/server-utils/test/ai/lib/tracing/langgraph-stream.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/langgraph-stream.test.ts @@ -247,23 +247,19 @@ describe('instrumentStateGraphCompile stream instrumentation', () => { expect(spanEnd).toHaveBeenCalledTimes(1); }); - it('forwards pipe options to the source pipeTo operation', async () => { + it('applies pipe options to the destination', async () => { const stream = new TestLangGraphStream(['first update']); - const originalPipeTo = stream.pipeTo.bind(stream); - const sourcePipeTo = vi.fn((destination: WritableStream, options?: StreamPipeOptions) => - originalPipeTo(destination, options), - ); - Object.defineProperty(stream, 'pipeTo', { configurable: true, value: sourcePipeTo, writable: true }); const compiledGraph = { stream: vi.fn().mockResolvedValue(stream) }; const compile = instrumentStateGraphCompile(() => compiledGraph, {}); const graph = compile() as { stream: () => Promise> }; - const destination = new WritableStream(); - const options = { preventClose: true, signal: new AbortController().signal }; + const written: string[] = []; + const destination = new WritableStream({ write: chunk => void written.push(chunk) }); - await (await graph.stream()).pipeTo(destination, options); + await (await graph.stream()).pipeTo(destination, { preventClose: true }); - expect(sourcePipeTo).toHaveBeenCalledWith(expect.any(WritableStream), options); - expect(destination.locked).toBe(false); + const writer = destination.getWriter(); + await writer.write('after pipe'); + expect(written).toEqual(['first update', 'after pipe']); expect(spanEnd).toHaveBeenCalledTimes(1); }); @@ -316,17 +312,16 @@ describe('instrumentStateGraphCompile stream instrumentation', () => { writer.releaseLock(); }); - it('errors the pipeThrough readable when the source pipe rejects', async () => { - const error = new Error('pipe failed'); - const stream = new TestLangGraphStream(['first update']); - Object.defineProperty(stream, 'pipeTo', { - configurable: true, - value: vi.fn().mockRejectedValue(error), - writable: true, + 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 graph = compile() as { stream: () => Promise> }; const reader = (await graph.stream()).pipeThrough(new TransformStream()).getReader(); 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 index 1b405386ac0d..2c033e1b2e0c 100644 --- a/packages/server-utils/test/ai/lib/tracing/langgraph-streaming.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/langgraph-streaming.test.ts @@ -1,6 +1,5 @@ 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'; @@ -43,18 +42,12 @@ describe('LangGraph iterable protocol', () => { }); describe('LangGraph pipeThrough execution lifecycle', () => { - it('awaits async transform and flush while allowing buffered output to outlive execution', async () => { + it('ends the span when the graph finishes, while buffered output outlives it', async () => { const span = createSpan(); let finishTransform!: () => void; - let finishFlush!: () => void; const transformPending = new Promise(resolve => { finishTransform = resolve; }); - const flushPending = new Promise(resolve => { - finishFlush = resolve; - }); - const transformStarted = vi.fn(); - const flushStarted = vi.fn(); const stream = instrumentStreamResult( new ReadableStream({ start(controller) { @@ -70,26 +63,17 @@ describe('LangGraph pipeThrough execution lifecycle', () => { new TransformStream( { async transform(chunk, controller) { - transformStarted(); await transformPending; controller.enqueue(chunk); }, - async flush() { - flushStarted(); - await flushPending; - }, }, undefined, { highWaterMark: 2 }, ), ); - await vi.waitFor(() => expect(transformStarted).toHaveBeenCalledTimes(1)); - expect(span.end).not.toHaveBeenCalled(); - finishTransform(); - await vi.waitFor(() => expect(flushStarted).toHaveBeenCalledTimes(1)); - expect(span.end).not.toHaveBeenCalled(); - finishFlush(); + 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 }); @@ -97,7 +81,7 @@ describe('LangGraph pipeThrough execution lifecycle', () => { reader.releaseLock(); }); - it('marks async flush errors as failures of the source execution', async () => { + 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( @@ -119,10 +103,8 @@ describe('LangGraph pipeThrough execution lifecycle', () => { ); const reader = output.getReader(); await expect(reader.read()).rejects.toThrow(error); - await vi.waitFor(() => { - expect(span.setStatus).toHaveBeenCalledWith({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); - expect(span.end).toHaveBeenCalledTimes(1); - }); + expect(span.setStatus).not.toHaveBeenCalled(); + expect(span.end).toHaveBeenCalledTimes(1); reader.releaseLock(); }); }); @@ -303,3 +285,93 @@ describe('LangGraph stream caller errors', () => { 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); + }); +}); From a6d40c8e89505f370f4c7223f652f52e9ac0ceee Mon Sep 17 00:00:00 2001 From: RulaKhaled Date: Thu, 22 Oct 2026 12:35:59 +0200 Subject: [PATCH 11/13] ref(server-utils): Simplify LangGraph stream lifecycle Reuse core handleCallbackErrors for reader and cancel settling, move the unpatchable-stream fallback into instrumentStreamResult, end the span once a locked cancel takes effect, and skip output recording for stream modes that don't carry the graph output. Co-Authored-By: Claude Opus 5.5 --- packages/server-utils/src/ai/core/utils.ts | 13 +-- .../server-utils/src/ai/langgraph/index.ts | 11 +- .../src/ai/langgraph/streaming.ts | 101 ++++++++---------- .../src/integrations/langgraph.ts | 3 - .../tracing/langgraph-invoke-agent.test.ts | 25 ++++- .../lib/tracing/langgraph-streaming.test.ts | 13 +++ 6 files changed, 85 insertions(+), 81 deletions(-) diff --git a/packages/server-utils/src/ai/core/utils.ts b/packages/server-utils/src/ai/core/utils.ts index a781d60252e0..46516dbb3d6a 100644 --- a/packages/server-utils/src/ai/core/utils.ts +++ b/packages/server-utils/src/ai/core/utils.ts @@ -89,17 +89,8 @@ export function buildMethodPath(currentPath: string, prop: string): string { } /** - * A stream that pulls through `reader`, typically an instrumented one. - * - * `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. + * 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 }>; diff --git a/packages/server-utils/src/ai/langgraph/index.ts b/packages/server-utils/src/ai/langgraph/index.ts index 4aec034f6634..c8f2aa228175 100644 --- a/packages/server-utils/src/ai/langgraph/index.ts +++ b/packages/server-utils/src/ai/langgraph/index.ts @@ -24,7 +24,7 @@ import { normalizeLangChainMessages, } from '../langchain/utils'; import { LANGGRAPH_ORIGIN } from './constants'; -import { instrumentStreamResult } from './streaming'; +import { instrumentStreamResult, isOutputStreamMode } from './streaming'; import type { CompiledGraph, LangGraphOptions } from './types'; import { extractAgentNameFromParams, @@ -250,13 +250,8 @@ function instrumentCompiledGraphOperation( if (streaming) { if (isAsyncIterable(result)) { span.setAttribute(GEN_AI_RESPONSE_STREAMING, true); - try { - return instrumentStreamResult(result, span, inputMessages, recordOutputs); - } catch { - // An unpatchable stream (e.g. frozen) still belongs to the caller; only the span is lost. - span.end(); - return result; - } + const recordStreamOutputs = recordOutputs && isOutputStreamMode(config?.streamMode); + return instrumentStreamResult(result, span, inputMessages, recordStreamOutputs); } span.end(); diff --git a/packages/server-utils/src/ai/langgraph/streaming.ts b/packages/server-utils/src/ai/langgraph/streaming.ts index 2afc970e37ca..74e23e58e9d9 100644 --- a/packages/server-utils/src/ai/langgraph/streaming.ts +++ b/packages/server-utils/src/ai/langgraph/streaming.ts @@ -1,4 +1,4 @@ -import { SPAN_STATUS_ERROR, withActiveSpan, withScope } from '@sentry/core'; +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'; @@ -14,13 +14,17 @@ export function instrumentStreamResult>( const spanScope = withActiveSpan(span, scope => scope); const lifecycle = createStreamLifecycle(span, inputMessages, recordOutputs); - if (isReadableStream(stream)) { - instrumentReadableStream(stream, spanScope, lifecycle); - return stream; + 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(); } - - const iterate = stream[Symbol.asyncIterator].bind(stream); - stream[Symbol.asyncIterator] = () => instrumentStreamIterator(iterate(), spanScope, lifecycle); return stream; } @@ -63,6 +67,7 @@ function createStreamLifecycle( 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; @@ -75,13 +80,14 @@ function createStreamLifecycle( completed = true; try { if (state) { - setResponseAttributes(span, inputMessages, getStreamResponseResult(state, inputMessages)); + 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(); } }; @@ -118,13 +124,18 @@ function getOutputPayload(chunk: unknown): { payload: unknown } | undefined { const namespace = Array.isArray(chunk[0]) ? (chunk[0] as unknown[]) : []; const mode = chunk.length === 3 ? chunk[1] : chunk[0]; - if (namespace.length > 0 || (typeof mode === 'string' && !OUTPUT_STREAM_MODES.has(mode))) { + 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. @@ -207,38 +218,33 @@ function instrumentReadableStream( const originalCancel = stream.cancel.bind(stream); stream.cancel = (reason?: unknown): Promise => { lifecycle.markCancelled(); - return stream.locked - ? originalCancel(reason) - : completeWithLifecycle( - withScope(scope, () => originalCancel(reason)), - lifecycle, - ); + // A plain locked stream rejects the cancel and keeps streaming, so only a cancel that took effect ends it. + const onError = stream.locked ? () => undefined : () => lifecycle.fail(); + return handleCallbackErrors( + () => withScope(scope, () => originalCancel(reason)), + onError, + undefined, + () => lifecycle.complete(), + ); }; } - const readerSource = (): ReadableStream => - readableFromReader(stream.getReader()); // 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): ReadableStream => { if (stream.locked || typeof destination?.getWriter !== 'function' || destination.locked) { throw new TypeError('Cannot pipe a locked stream, or to a missing or locked destination.'); } - return readerSource(); + return readableFromReader(stream.getReader()); }; if (stream.tee) { - stream.tee = () => readerSource().tee(); + stream.tee = () => readableFromReader(stream.getReader()).tee(); } if (stream.pipeTo) { - stream.pipeTo = (destination: WritableStream, options?: StreamPipeOptions): Promise => { - try { - return instrumentedSource(destination).pipeTo(destination, options); - } catch (error) { - return Promise.reject(error); - } - }; + stream.pipeTo = async (destination: WritableStream, options?: StreamPipeOptions): Promise => + instrumentedSource(destination).pipeTo(destination, options); } if (stream.pipeThrough) { @@ -255,31 +261,23 @@ function instrumentReader( lifecycle: StreamLifecycle, ): ReadableStreamReaderLike { const originalRead = reader.read.bind(reader); - reader.read = (...args: unknown[]): Promise> => { - const readPromise: Promise> = withScope(scope, () => originalRead(...args)); - return readPromise.then( - result => { - if (result.done) { - lifecycle.complete(); - } else { - lifecycle.recordChunk(result.value); - } - return result; - }, - error => { - lifecycle.fail(); - throw error; - }, + 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(); - return completeWithLifecycle( - withScope(scope, () => originalCancel(reason)), - lifecycle, + return handleCallbackErrors( + () => withScope(scope, () => originalCancel(reason)), + () => lifecycle.fail(), + () => lifecycle.complete(), ); }; } @@ -287,19 +285,6 @@ function instrumentReader( return reader; } -function completeWithLifecycle(promise: Promise, lifecycle: StreamLifecycle): Promise { - return promise.then( - result => { - lifecycle.complete(); - return result; - }, - error => { - lifecycle.fail(); - throw error; - }, - ); -} - async function* instrumentStreamIterator( iterator: AsyncIterator, scope: Scope, diff --git a/packages/server-utils/src/integrations/langgraph.ts b/packages/server-utils/src/integrations/langgraph.ts index cdd804065f58..3b7dafd7d3cb 100644 --- a/packages/server-utils/src/integrations/langgraph.ts +++ b/packages/server-utils/src/integrations/langgraph.ts @@ -60,9 +60,6 @@ function instrumentLanggraph(options: LangGraphOptions): void { // throw) and cleared in both `end` and `error`, so it can neither stick on across calls nor // stay off during this call's nested compile. Tool wrapping is guarded for the same reason. insideCreateReactAgent = true; - if (isInsideManualLangGraphInstrumentation()) { - return; - } try { const { arguments: args } = message as OrchestrionChannelContext; const params = getFirstArgObject(args); 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 28c760846bd5..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,7 +8,7 @@ import { startSpan, } from '@sentry/core'; import type { Span } from '@sentry/core'; -import { GEN_AI_CONVERSATION_ID } from '@sentry/conventions/attributes'; +import { GEN_AI_CONVERSATION_ID, GEN_AI_RESPONSE_TEXT } from '@sentry/conventions/attributes'; import { instrumentCompiledGraph, instrumentStateGraphCompile, @@ -259,4 +259,27 @@ describe('LangGraph invoke stream suppression', () => { 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-streaming.test.ts b/packages/server-utils/test/ai/lib/tracing/langgraph-streaming.test.ts index 2c033e1b2e0c..dbf41bdd7cb3 100644 --- a/packages/server-utils/test/ai/lib/tracing/langgraph-streaming.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/langgraph-streaming.test.ts @@ -374,4 +374,17 @@ describe('LangGraph stream cancellation and tee', () => { 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); + }); }); From aac1c7a1840b9e159f0c6281b8c4ee78a24483ac Mon Sep 17 00:00:00 2001 From: RulaKhaled Date: Thu, 22 Oct 2026 15:32:59 +0200 Subject: [PATCH 12/13] fix(server-utils): Keep stream errors after a rejected LangGraph cancel A cancel that rejects no longer leaves the stream marked as cancelled, so a later read failure still ends the agent span as an error. Co-Authored-By: Claude Opus 5.5 --- .../src/ai/langgraph/streaming.ts | 25 +++++++++++----- .../lib/tracing/langgraph-streaming.test.ts | 29 +++++++++++++++++++ 2 files changed, 46 insertions(+), 8 deletions(-) diff --git a/packages/server-utils/src/ai/langgraph/streaming.ts b/packages/server-utils/src/ai/langgraph/streaming.ts index 74e23e58e9d9..ca5b7e78da7e 100644 --- a/packages/server-utils/src/ai/langgraph/streaming.ts +++ b/packages/server-utils/src/ai/langgraph/streaming.ts @@ -32,7 +32,7 @@ interface StreamLifecycle { recordChunk: (chunk: unknown) => void; complete: () => void; fail: () => void; - markCancelled: () => void; + markCancelled: (cancelled: boolean) => void; } // Modes whose payloads carry the graph's output; `debug`, `tasks`, `checkpoints` and `custom` do not. @@ -109,8 +109,8 @@ function createStreamLifecycle( } complete(); }, - markCancelled(): void { - cancelled = true; + markCancelled(value: boolean): void { + cancelled = value; }, }; } @@ -217,12 +217,18 @@ function instrumentReadableStream( if (stream.cancel) { const originalCancel = stream.cancel.bind(stream); stream.cancel = (reason?: unknown): Promise => { - lifecycle.markCancelled(); + // 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 onError = stream.locked ? () => undefined : () => lifecycle.fail(); + const locked = stream.locked; return handleCallbackErrors( () => withScope(scope, () => originalCancel(reason)), - onError, + () => { + lifecycle.markCancelled(false); + if (!locked) { + lifecycle.fail(); + } + }, undefined, () => lifecycle.complete(), ); @@ -273,10 +279,13 @@ function instrumentReader( if (reader.cancel) { const originalCancel = reader.cancel.bind(reader); reader.cancel = (reason?: unknown): Promise => { - lifecycle.markCancelled(); + lifecycle.markCancelled(true); return handleCallbackErrors( () => withScope(scope, () => originalCancel(reason)), - () => lifecycle.fail(), + () => { + lifecycle.markCancelled(false); + lifecycle.fail(); + }, () => lifecycle.complete(), ); }; 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 index dbf41bdd7cb3..723582497c6f 100644 --- a/packages/server-utils/test/ai/lib/tracing/langgraph-streaming.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/langgraph-streaming.test.ts @@ -1,5 +1,6 @@ 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'; @@ -387,4 +388,32 @@ describe('LangGraph stream cancellation and tee', () => { 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); + }); }); From e5b4875bdc5aa1b5016b973f91f38332d36c4013 Mon Sep 17 00:00:00 2001 From: RulaKhaled Date: Thu, 22 Oct 2026 16:13:05 +0200 Subject: [PATCH 13/13] fix(server-utils): Check pipe options before locking a LangGraph stream pipeThrough without a readable side, or pipeTo/pipeThrough with an invalid signal, took the source reader before the native call threw, leaving the stream locked and its span open. Co-Authored-By: Claude Opus 5.5 --- .../src/ai/langgraph/streaming.ts | 22 +++++++++++---- .../lib/tracing/langgraph-streaming.test.ts | 27 +++++++++++++++++++ 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/packages/server-utils/src/ai/langgraph/streaming.ts b/packages/server-utils/src/ai/langgraph/streaming.ts index ca5b7e78da7e..899e791bddb5 100644 --- a/packages/server-utils/src/ai/langgraph/streaming.ts +++ b/packages/server-utils/src/ai/langgraph/streaming.ts @@ -237,9 +237,20 @@ function instrumentReadableStream( // 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): ReadableStream => { - if (stream.locked || typeof destination?.getWriter !== 'function' || destination.locked) { - throw new TypeError('Cannot pipe a locked stream, or to a missing or locked destination.'); + 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()); }; @@ -250,14 +261,15 @@ function instrumentReadableStream( if (stream.pipeTo) { stream.pipeTo = async (destination: WritableStream, options?: StreamPipeOptions): Promise => - instrumentedSource(destination).pipeTo(destination, options); + instrumentedSource(destination, options).pipeTo(destination, options); } if (stream.pipeThrough) { stream.pipeThrough = ( transform: ReadableWritablePair, options?: StreamPipeOptions, - ): ReadableStream => instrumentedSource(transform?.writable).pipeThrough(transform, options); + ): ReadableStream => + instrumentedSource(transform?.writable, options, !!transform?.readable).pipeThrough(transform, options); } } 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 index 723582497c6f..5918c993ef92 100644 --- a/packages/server-utils/test/ai/lib/tracing/langgraph-streaming.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/langgraph-streaming.test.ts @@ -416,4 +416,31 @@ describe('LangGraph stream cancellation and tee', () => { 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); + }); });