From e500300e205e8b252faa6f2cc0473fd3ea6a6e45 Mon Sep 17 00:00:00 2001 From: "Seongho.Bak" Date: Wed, 26 Aug 2026 12:16:32 +0900 Subject: [PATCH 1/5] 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 | 278 ++++++++++++------ .../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(+), 103 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 7b04ee85f65d..d913c86bbe39 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, @@ -31,14 +32,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'); @@ -47,6 +49,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() @@ -61,7 +68,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( @@ -83,6 +90,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 5a12a972dc9d..ad121f1b3aed 100644 --- a/packages/server-utils/src/ai/langgraph/index.ts +++ b/packages/server-utils/src/ai/langgraph/index.ts @@ -1,12 +1,15 @@ /* eslint-disable typescript-eslint/no-deprecated */ import { captureException, + getCurrentScope, SEMANTIC_ATTRIBUTE_SENTRY_OP, 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, @@ -14,6 +17,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, } from '@sentry/conventions/attributes'; @@ -21,8 +25,9 @@ import { GEN_AI_INVOKE_AGENT_OPERATION_ATTRIBUTE } from '../core/gen-ai-attribut 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, @@ -31,15 +36,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, @@ -61,7 +66,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( @@ -74,6 +79,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; }, }); @@ -95,105 +112,173 @@ 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( - { - op: 'gen_ai.invoke_agent', - name: 'invoke_agent', - attributes: { - [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: LANGGRAPH_ORIGIN, - [SEMANTIC_ATTRIBUTE_SENTRY_OP]: GEN_AI_INVOKE_AGENT_OPERATION_ATTRIBUTE, - [GEN_AI_OPERATION_NAME]: 'invoke_agent', - }, + const spanOptions = { + op: 'gen_ai.invoke_agent', + name: 'invoke_agent', + attributes: { + [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: LANGGRAPH_ORIGIN, + [SEMANTIC_ATTRIBUTE_SENTRY_OP]: GEN_AI_INVOKE_AGENT_OPERATION_ATTRIBUTE, + [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 (!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 (recordOutputs) { - setResponseAttributes(span, inputMessages ?? null, result); + if (streaming) { + if (isAsyncIterable(result)) { + span.setAttribute(GEN_AI_RESPONSE_STREAMING, true); + return instrumentStreamResult(result, span); } + span.end(); return result; - } catch (error) { - span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); - captureException(error, { - mechanism: { - handled: false, - type: 'auto.ai.langgraph.error', - }, - }); - throw error; } - }, - ); + + if (recordOutputs) { + setResponseAttributes(span, inputMessages ?? null, result); + } + + return result; + } catch (error) { + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + captureException(error, { + mechanism: { + handled: false, + type: 'auto.ai.langgraph.error', + }, + }); + if (streaming) { + span.end(); + } + throw error; + } + }; + + return streaming ? startSpanManual(spanOptions, run) : startSpan(spanOptions, run); }, }); } @@ -232,7 +317,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 = {}; @@ -250,6 +335,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 bbd119d85e193efb383faca922e54b5138768b73 Mon Sep 17 00:00:00 2001 From: "Seongho.Bak" Date: Wed, 26 Aug 2026 13:18:45 +0900 Subject: [PATCH 2/5] 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 d913c86bbe39..aa04d631140d 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, @@ -32,15 +33,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'); @@ -50,10 +52,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() @@ -68,7 +70,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( @@ -98,6 +100,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 ad121f1b3aed..aefccf04bc9c 100644 --- a/packages/server-utils/src/ai/langgraph/index.ts +++ b/packages/server-utils/src/ai/langgraph/index.ts @@ -251,7 +251,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 794a3362fc5ef343fd0bbc900d9f310c76e4817a Mon Sep 17 00:00:00 2001 From: "Seongho.Bak" Date: Wed, 26 Aug 2026 13:55:21 +0900 Subject: [PATCH 3/5] fix(server-utils): Preserve LangGraph pipe lifecycle Record source chunks and complete spans across pipeTo and pipeThrough consumption. Co-Authored-By: OpenAI Codex --- .../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 23aa45fe69199a1f25a96b2a43827b0383984d7c Mon Sep 17 00:00:00 2001 From: "Seongho.Bak" Date: Wed, 26 Aug 2026 14:20:44 +0900 Subject: [PATCH 4/5] fix(server-utils): Preserve pipeThrough errors Throw for locked pipeThrough inputs and forward asynchronous pipeline failures through the returned readable. Co-Authored-By: OpenAI Codex --- .../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 10f506d1571666a46eb3163c1e9325507420e0e0 Mon Sep 17 00:00:00 2001 From: "Seongho.Bak" Date: Wed, 26 Aug 2026 14:31:37 +0900 Subject: [PATCH 5/5] 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. Co-Authored-By: OpenAI Codex --- 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(); });