From 41e1f6c54a1f361994527d16f65be0a34d426b71 Mon Sep 17 00:00:00 2001 From: RulaKhaled Date: Mon, 26 Oct 2026 20:16:41 +0100 Subject: [PATCH 1/4] fix(server-utils): Create an invoke_agent span for LangGraph stream() LangGraph runs the graph as soon as stream() is called, whether or not the stream is read, so the span ends on the run's own end or error callback. A break never reaches those callbacks, so it ends the span too. Co-Authored-By: Seongho.Bak Co-Authored-By: Claude Opus 5.5 --- .../suites/tracing/langgraph/index.ts | 6 +- .../suites/tracing/langgraph/test.ts | 21 +- .../langgraph/agent-stream-scenario.mjs | 53 +++ .../tracing/langgraph/scenario-stream.mjs | 70 ++++ .../suites/tracing/langgraph/test.ts | 80 +++++ .../tracing/langgraph/v1/scenario-stream.mjs | 102 ++++++ .../suites/tracing/langgraph/v1/test.ts | 63 ++++ .../server-utils/src/ai/langchain/utils.ts | 32 +- .../src/ai/langgraph/agent-span.ts | 130 +++++++ .../server-utils/src/ai/langgraph/index.ts | 255 +++++-------- .../server-utils/src/ai/langgraph/stream.ts | 155 ++++++++ .../server-utils/src/ai/langgraph/types.ts | 1 + .../src/integrations/langgraph.ts | 48 +-- .../ai/lib/tracing/langchain-utils.test.ts | 29 ++ .../ai/lib/tracing/langgraph-stream.test.ts | 335 ++++++++++++++++++ 15 files changed, 1169 insertions(+), 211 deletions(-) create mode 100644 dev-packages/node-integration-tests/suites/tracing/langgraph/agent-stream-scenario.mjs create mode 100644 dev-packages/node-integration-tests/suites/tracing/langgraph/scenario-stream.mjs create mode 100644 dev-packages/node-integration-tests/suites/tracing/langgraph/v1/scenario-stream.mjs create mode 100644 packages/server-utils/src/ai/langgraph/agent-span.ts create mode 100644 packages/server-utils/src/ai/langgraph/stream.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 8bf8f8acd720..625b77bc139e 100644 --- a/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/index.ts +++ b/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/index.ts @@ -60,7 +60,11 @@ export default Sentry.withSentry( messages: [{ role: 'user', content: 'What is the weather in SF?' }], }); - return new Response(JSON.stringify({ success: true })); + const stream = await compiled.stream( + { messages: [{ role: 'user', content: 'Stream the weather in SF' }] }, + { encoding: 'text/event-stream' }, + ); + return new Response(stream as unknown as ReadableStream, { headers: { 'content-type': 'text/event-stream' } }); }, }, ); diff --git a/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/test.ts b/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/test.ts index 87008d27d870..4aa2a3ca4623 100644 --- a/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/test.ts +++ b/dev-packages/cloudflare-integration-tests/suites/tracing/langgraph/test.ts @@ -5,6 +5,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 { getSpanOp, getSpansFromEnvelope } from '../../../spanUtils'; // want to test that the instrumentation does not break in our // cloudflare SDK. -it('traces langgraph compile and invoke operations', async ({ signal }) => { +it('traces langgraph invoke and stream operations', async ({ signal }) => { const runner = createRunner(__dirname) .ignore('event') .expect(envelope => { @@ -26,10 +27,22 @@ it('traces langgraph compile and invoke operations', async ({ signal }) => { expect(segmentSpan?.name).toBe('GET /'); const genAiSpans = spans.filter(span => getSpanOp(span)?.startsWith('gen_ai.')); - expect(genAiSpans).toHaveLength(1); - expect(genAiSpans.map(span => span.name).sort()).toEqual(['invoke_agent weather_assistant']); + expect(genAiSpans).toHaveLength(2); + expect(genAiSpans.map(span => span.name)).toEqual([ + 'invoke_agent weather_assistant', + 'invoke_agent weather_assistant', + ]); - const invokeAgentSpan = genAiSpans.find(span => span.name === 'invoke_agent weather_assistant'); + const streamAgentSpan = genAiSpans.find(span => span.attributes[GEN_AI_RESPONSE_STREAMING]); + expect(streamAgentSpan!.status).toBe('ok'); + expect(streamAgentSpan!.attributes[GEN_AI_RESPONSE_STREAMING]).toEqual({ type: 'boolean', value: true }); + expect(streamAgentSpan!.attributes[GEN_AI_INPUT_MESSAGES]).toEqual({ + type: 'string', + value: '[{"role":"user","content":"Stream the weather in SF"}]', + }); + expect(streamAgentSpan!.attributes[GEN_AI_RESPONSE_MODEL]).toEqual({ type: 'string', value: 'mock-model' }); + + const invokeAgentSpan = genAiSpans.find(span => !span.attributes[GEN_AI_RESPONSE_STREAMING]); expect(invokeAgentSpan).toBeDefined(); expect(invokeAgentSpan!.status).toBe('ok'); expect(invokeAgentSpan!.attributes[GEN_AI_OPERATION_NAME]).toEqual({ diff --git a/dev-packages/node-integration-tests/suites/tracing/langgraph/agent-stream-scenario.mjs b/dev-packages/node-integration-tests/suites/tracing/langgraph/agent-stream-scenario.mjs new file mode 100644 index 000000000000..cca20e95102a --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/langgraph/agent-stream-scenario.mjs @@ -0,0 +1,53 @@ +import { ChatAnthropic } from '@langchain/anthropic'; +import { HumanMessage, SystemMessage } from '@langchain/core/messages'; +import { createReactAgent } from '@langchain/langgraph/prebuilt'; +import * as Sentry from '@sentry/node'; +import express from 'express'; + +function startMockAnthropicServer() { + const app = express(); + app.use(express.json()); + + app.post('/v1/messages', (req, res) => { + res.json({ + id: 'msg_react_agent_stream_123', + type: 'message', + role: 'assistant', + content: [{ type: 'text', text: 'Paris is the capital of France.' }], + model: req.body.model, + stop_reason: 'end_turn', + stop_sequence: null, + usage: { input_tokens: 20, output_tokens: 10 }, + }); + }); + + return new Promise(resolve => { + const server = app.listen(0, () => resolve(server)); + }); +} + +async function run() { + const server = await startMockAnthropicServer(); + const baseUrl = `http://localhost:${server.address().port}`; + + await Sentry.startSpan({ op: 'function', name: 'main' }, async () => { + const llm = new ChatAnthropic({ + model: 'claude-3-5-sonnet-20241022', + apiKey: 'mock-api-key', + clientOptions: { baseURL: baseUrl }, + }); + const agent = createReactAgent({ llm, tools: [], name: 'helpful_assistant' }); + + const stream = await agent.stream({ + messages: [new SystemMessage('You are a helpful assistant.'), new HumanMessage('What is the capital of France?')], + }); + for await (const _chunk of stream) { + // Consume the stream so the agent and its child spans complete before the root span ends. + } + }); + + await Sentry.flush(2000); + server.close(); +} + +run(); diff --git a/dev-packages/node-integration-tests/suites/tracing/langgraph/scenario-stream.mjs b/dev-packages/node-integration-tests/suites/tracing/langgraph/scenario-stream.mjs new file mode 100644 index 000000000000..086e9caa4b76 --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/langgraph/scenario-stream.mjs @@ -0,0 +1,70 @@ +import { FakeListChatModel } from '@langchain/core/utils/testing'; +import { END, MessagesAnnotation, START, StateGraph } from '@langchain/langgraph'; +import * as Sentry from '@sentry/node'; + +const sleep = ms => new Promise(resolve => setTimeout(resolve, ms)); +const user = content => ({ messages: [{ role: 'user', content }] }); + +function weatherGraph(name, model) { + return new StateGraph(MessagesAnnotation) + .addNode('agent', async state => ({ messages: [await model.invoke(state.messages)] })) + .addEdge(START, 'agent') + .addEdge('agent', END) + .compile({ name }); +} + +async function run() { + const model = new FakeListChatModel({ responses: ['Sunny'] }); + + await Sentry.startSpan({ op: 'function', name: 'main' }, async () => { + const graph = weatherGraph('weather_graph', model); + + for await (const _chunk of await graph.stream(user('Read the whole stream'))) { + // drain + } + + await graph.stream(user('Never read the stream')); + + for await (const _chunk of await graph.stream(user('Break out of the stream'))) { + break; + } + + const canceled = await graph.stream(user('Cancel the stream')); + await canceled.cancel(); + + const failing = new StateGraph(MessagesAnnotation) + .addNode('agent', async () => { + throw new Error('Weather service down'); + }) + .addEdge(START, 'agent') + .addEdge('agent', END) + .compile({ name: 'failing_graph' }); + try { + for await (const _chunk of await failing.stream(user('Fail the stream'))) { + // drain + } + } catch { + // expected + } + + const outer = new StateGraph(MessagesAnnotation) + .addNode('inner', weatherGraph('inner_graph', model)) + .addEdge(START, 'inner') + .addEdge('inner', END) + .compile({ name: 'outer_graph' }); + for await (const _chunk of await outer.stream(user('Stream a subgraph'))) { + // drain + } + + for await (const _event of graph.streamEvents(user('Stream events'), { version: 'v2' })) { + // drain + } + + await graph.invoke(user('Invoke the graph')); + + // Lets the unread stream's run finish before the root span ends. + await sleep(200); + }); +} + +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 53dc3f6846af..befc8f05cfbd 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_OUTPUT_MESSAGES, 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, @@ -365,4 +366,83 @@ describe('LangGraph integration', () => { .completed(); }); }); + + createEsmAndCjsTests(__dirname, 'scenario-stream.mjs', 'instrument-with-pii.mjs', (createRunner, test) => { + test('creates one agent span per stream() run, however the stream is consumed', { timeout: 30000 }, async () => { + await createRunner() + .expect({ + span: container => { + const spans = container.items; + const agentSpans = spans.filter(span => span.attributes[SENTRY_OP]?.value === 'gen_ai.invoke_agent'); + const agentSpanFor = (prompt: string) => + agentSpans.find(span => + getStringAttributeValue(span.attributes[GEN_AI_INPUT_MESSAGES]?.value)?.includes(prompt), + ); + const chatSpanUnder = (parent: { span_id: string } | undefined) => + spans.find( + span => span.attributes[SENTRY_OP]?.value === 'gen_ai.chat' && span.parent_span_id === parent?.span_id, + ); + + expect(agentSpans).toHaveLength(9); + + const read = agentSpanFor('Read the whole stream'); + expect(read?.status).toBe('ok'); + expect(read?.attributes[GEN_AI_RESPONSE_STREAMING]?.value).toBe(true); + expect(read?.attributes[GEN_AI_RESPONSE_TEXT]?.value).toBe('[{"role":"assistant","content":"Sunny"}]'); + expect(chatSpanUnder(read)).toBeDefined(); + + const unread = agentSpanFor('Never read the stream'); + expect(unread?.status).toBe('ok'); + expect(unread?.attributes[GEN_AI_RESPONSE_TEXT]?.value).toBe('[{"role":"assistant","content":"Sunny"}]'); + expect(chatSpanUnder(unread)).toBeDefined(); + + expect(agentSpanFor('Break out of the stream')?.status).toBe('ok'); + expect(agentSpanFor('Cancel the stream')?.status).toBe('ok'); + + const failed = agentSpans.find(span => span.name === 'invoke_agent failing_graph'); + expect(failed?.status).toBe('error'); + + const outer = agentSpans.find(span => span.name === 'invoke_agent outer_graph'); + const inner = agentSpans.find(span => span.name === 'invoke_agent inner_graph'); + expect(outer?.status).toBe('ok'); + expect(inner?.parent_span_id).toBe(outer?.span_id); + + expect(agentSpanFor('Stream events')?.attributes[GEN_AI_RESPONSE_STREAMING]?.value).toBe(true); + + const invoked = agentSpanFor('Invoke the graph'); + expect(invoked?.status).toBe('ok'); + expect(invoked?.attributes[GEN_AI_RESPONSE_STREAMING]).toBeUndefined(); + }, + }) + .start() + .completed(); + }); + }); + + createEsmAndCjsTests(__dirname, 'agent-stream-scenario.mjs', 'instrument-agent.mjs', (createRunner, test) => { + test( + 'parents the chat span of a streamed createReactAgent run under its agent span', + { timeout: 30000 }, + async () => { + await createRunner() + .expect({ + span: container => { + const agentSpans = container.items.filter( + span => span.attributes[SENTRY_OP]?.value === 'gen_ai.invoke_agent', + ); + const chatSpans = container.items.filter(span => span.attributes[SENTRY_OP]?.value === 'gen_ai.chat'); + expect(agentSpans).toHaveLength(1); + expect(chatSpans).toHaveLength(1); + expect(agentSpans[0]!.name).toBe('invoke_agent helpful_assistant'); + expect(agentSpans[0]!.status).toBe('ok'); + expect(agentSpans[0]!.attributes[GEN_AI_RESPONSE_STREAMING]?.value).toBe(true); + expect(chatSpans[0]!.parent_span_id).toBe(agentSpans[0]!.span_id); + expect(chatSpans[0]!.attributes[GEN_AI_AGENT_NAME]?.value).toBe('helpful_assistant'); + }, + }) + .start() + .completed(); + }, + ); + }); }); 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..ae9d9a6e0bef --- /dev/null +++ b/dev-packages/node-integration-tests/suites/tracing/langgraph/v1/scenario-stream.mjs @@ -0,0 +1,102 @@ +import { RunnableLambda } from '@langchain/core/runnables'; +import { FakeListChatModel } from '@langchain/core/utils/testing'; +import { Command, END, interrupt, MemorySaver, MessagesAnnotation, START, StateGraph } from '@langchain/langgraph'; +import * as Sentry from '@sentry/node'; + +const sleep = ms => new Promise(resolve => setTimeout(resolve, ms)); +const user = content => ({ messages: [{ role: 'user', content }] }); +const drain = async stream => { + for await (const _chunk of stream) { + // drain + } +}; + +function weatherBuilder(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'] }); + + await Sentry.startSpan({ op: 'function', name: 'main' }, async () => { + const graph = weatherBuilder(model).compile({ name: 'weather_graph' }); + + await drain(await graph.stream(user('Read the whole stream'))); + await graph.stream(user('Never read the stream')); + for await (const _chunk of await graph.stream(user('Break out of the stream'))) { + break; + } + const canceled = await graph.stream(user('Cancel the stream')); + await canceled.cancel(); + + const slow = new StateGraph(MessagesAnnotation) + .addNode('agent', async () => { + await sleep(200); + return { messages: [] }; + }) + .addEdge(START, 'agent') + .addEdge('agent', END) + .compile({ name: 'slow_graph' }); + const controller = new AbortController(); + setTimeout(() => controller.abort(), 50); + try { + await drain(await slow.stream(user('Abort the stream'), { signal: controller.signal })); + } catch { + // expected + } + + const failing = new StateGraph(MessagesAnnotation) + .addNode('agent', async () => { + throw new Error('Weather service down'); + }) + .addEdge(START, 'agent') + .addEdge('agent', END) + .compile({ name: 'failing_graph' }); + try { + await drain(await failing.stream(user('Fail the stream'))); + } catch { + // expected + } + + const approval = new StateGraph(MessagesAnnotation) + .addNode('ask', async () => ({ messages: [{ role: 'user', content: String(interrupt('Approve?')) }] })) + .addEdge(START, 'ask') + .addEdge('ask', END) + .compile({ name: 'approval_graph', checkpointer: new MemorySaver() }); + const thread = { configurable: { thread_id: 'approval' } }; + await drain(await approval.stream(user('Interrupt the stream'), thread)); + await drain(await approval.stream(new Command({ resume: 'yes' }), thread)); + + await drain(await graph.streamEvents(user('Stream events v3'), { version: 'v3' })); + + const recursive = new StateGraph(MessagesAnnotation) + .addNode('agent', async state => { + if (state.messages.length === 1 && state.messages[0].content === 'Invoke with a nested stream') { + await drain(await recursive.stream(user('Nested stream of the same graph'))); + } + return { messages: [{ role: 'assistant', content: 'done' }] }; + }) + .addEdge(START, 'agent') + .addEdge('agent', END) + .compile({ name: 'recursive_graph' }); + await recursive.invoke(user('Invoke with a nested stream')); + + // Passing the parent's config hands stream() a CallbackManager rather than an array. + await RunnableLambda.from(async (input, config) => drain(await graph.stream(input, config))).invoke( + user('Stream inside a runnable'), + ); + + const manual = Sentry.instrumentStateGraph(weatherBuilder(model), { recordInputs: false }).compile({ + name: 'manual_graph', + }); + await drain(await manual.stream(user('Keep this prompt private'))); + + // Lets the unread stream's run finish before the root span ends. + await sleep(200); + }); +} + +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..3efcceee9c0d 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,65 @@ describe('LangGraph integration (v1)', () => { }, }, ); + + createEsmAndCjsTests( + __dirname, + 'scenario-stream.mjs', + 'instrument-with-pii.mjs', + (createRunner, test) => { + test('creates one agent span per stream() run, ended by the graph run', 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'); + const agentSpanFor = (prompt: string) => + agentSpans.find(span => String(span.attributes[GEN_AI_INPUT_MESSAGES]?.value).includes(prompt)); + const byName = (name: string) => agentSpans.filter(span => span.name === `invoke_agent ${name}`); + + expect(agentSpans).toHaveLength(13); + + for (const prompt of ['Read the whole stream', 'Never read the stream', 'Stream inside a runnable']) { + const span = agentSpanFor(prompt)!; + expect(span.status).toBe('ok'); + expect(span.attributes[GEN_AI_RESPONSE_STREAMING]?.value).toBe(true); + expect(span.attributes[GEN_AI_RESPONSE_TEXT]?.value).toBe('[{"role":"assistant","content":"Sunny"}]'); + expect( + spans.filter( + child => child.attributes[SENTRY_OP]?.value === 'gen_ai.chat' && child.parent_span_id === span.span_id, + ), + ).toHaveLength(1); + } + + expect(agentSpanFor('Break out of the stream')?.status).toBe('ok'); + expect(agentSpanFor('Cancel the stream')?.status).toBe('ok'); + expect(byName('slow_graph')[0]?.status).toBe('ok'); + expect(byName('failing_graph')[0]?.status).toBe('error'); + + const approvals = byName('approval_graph'); + expect(approvals).toHaveLength(2); + expect(approvals.map(span => span.status)).toEqual(['ok', 'ok']); + + expect(agentSpanFor('Stream events v3')?.attributes[GEN_AI_RESPONSE_STREAMING]?.value).toBe(true); + + const invoked = agentSpanFor('Invoke with a nested stream')!; + const nested = agentSpanFor('Nested stream of the same graph')!; + expect(invoked.attributes[GEN_AI_RESPONSE_STREAMING]).toBeUndefined(); + expect(nested.parent_span_id).toBe(invoked.span_id); + expect(byName('recursive_graph')).toHaveLength(2); + + const manual = byName('manual_graph'); + expect(manual).toHaveLength(1); + expect(manual[0]!.attributes[GEN_AI_INPUT_MESSAGES]).toBeUndefined(); + }); + }, + { + additionalDependencies: { + '@langchain/core': '^1.0.0', + '@langchain/langgraph': '^1.0.0', + }, + }, + ); }); diff --git a/packages/server-utils/src/ai/langchain/utils.ts b/packages/server-utils/src/ai/langchain/utils.ts index 8ae219e42437..5cae51d3ac53 100644 --- a/packages/server-utils/src/ai/langchain/utils.ts +++ b/packages/server-utils/src/ai/langchain/utils.ts @@ -606,24 +606,32 @@ function containsSentryHandler(handlers: unknown[]): boolean { * @internal Exported for cross-package instrumentation. */ export function _INTERNAL_mergeLangChainCallbackHandler(existing: unknown, sentryHandler: unknown): unknown { + const handlers = isCallbackManager(existing) + ? (existing.handlers ?? []) + : Array.isArray(existing) + ? existing + : existing + ? [existing] + : []; + + return containsSentryHandler(handlers) ? existing : _INTERNAL_addLangChainCallbackHandler(existing, sentryHandler); +} + +/** + * Add `handler` to a given set of LangChain callbacks or callback manager, without deduplicating it. + * A callback manager is copied rather than mutated. + * @internal Exported for cross-package instrumentation. + */ +export function _INTERNAL_addLangChainCallbackHandler(existing: unknown, handler: unknown): unknown { if (!existing) { - return [sentryHandler]; + return [handler]; } if (isCallbackManager(existing)) { - if (containsSentryHandler(existing.handlers ?? [])) { - return existing; - } - const copied = existing.copy() as { addHandler: (handler: unknown, inherit?: boolean) => void }; - copied.addHandler(sentryHandler, true); + copied.addHandler(handler, true); return copied; } - const handlers = Array.isArray(existing) ? existing : [existing]; - if (containsSentryHandler(handlers)) { - return existing; - } - - return [...handlers, sentryHandler]; + return [...(Array.isArray(existing) ? existing : [existing]), handler]; } diff --git a/packages/server-utils/src/ai/langgraph/agent-span.ts b/packages/server-utils/src/ai/langgraph/agent-span.ts new file mode 100644 index 000000000000..4bfbb075f06d --- /dev/null +++ b/packages/server-utils/src/ai/langgraph/agent-span.ts @@ -0,0 +1,130 @@ +import { stringify } from '@sentry/core'; +import type { Span, SpanAttributes } from '@sentry/core'; +import { + GEN_AI_AGENT_NAME, + GEN_AI_INPUT_MESSAGES, + GEN_AI_OPERATION_NAME, + GEN_AI_PIPELINE_NAME, + GEN_AI_REQUEST_MODEL, + GEN_AI_SYSTEM_INSTRUCTIONS, + GEN_AI_TOOL_DEFINITIONS, + SENTRY_OP, + SENTRY_ORIGIN, +} from '@sentry/conventions/attributes'; +import { GEN_AI_INVOKE_AGENT } from '@sentry/conventions/op'; +import { extractSystemInstructions } from '../core/utils'; +import type { BaseChatModel, LangChainMessage } from '../langchain/types'; +import { + _INTERNAL_mergeLangChainCallbackHandler, + getConversationIdFromMetadata, + getConversationIdMetadataFromConfig, + normalizeLangChainMessages, +} from '../langchain/utils'; +import { LANGGRAPH_ORIGIN } from './constants'; +import type { CompiledGraph, LangGraphOptions } from './types'; +import { extractToolsFromCompiledGraph } from './utils'; + +// invoke() drains this.stream() internally. The marker sits on invoke's own span scope, so only that call is +// skipped, and a stream() of the same graph started later in the run (e.g. from a tool) still gets a span. +export const LANGGRAPH_INVOKE_ACTIVE = 'sentry_langgraph_invoke_active'; +const INVOKE_MARKERS = new WeakMap(); + +/** The marker invoke() leaves for its own internal stream() call on `graph`. */ +export 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; +} + +/** Name and initial attributes of an agent span, from the call's config. */ +export function getAgentSpanOptions(args: unknown[]): { name: string; attributes: SpanAttributes } { + const config = args.length > 1 ? (args[1] as Record | undefined) : undefined; + return { + name: 'invoke_agent', + attributes: { + // Set before `spanStart`, so an id from `Sentry.setConversationId()` wins, as on the child spans + ...getConversationIdFromMetadata(config?.configurable as Record | undefined), + [SENTRY_ORIGIN]: LANGGRAPH_ORIGIN, + [SENTRY_OP]: GEN_AI_INVOKE_AGENT, + [GEN_AI_OPERATION_NAME]: 'invoke_agent', + }, + }; +} + +/** Sets the request attributes on an agent span and injects the callback handler into the call's config. */ +export function prepareAgentSpan( + span: Span, + args: unknown[], + graphInstance: CompiledGraph, + compileOptions: Record, + options: LangGraphOptions, + llm: BaseChatModel | null | undefined, + sentryCallbackHandler: unknown, +): LangChainMessage[] | null { + const graphName = compileOptions?.name; + + if (graphName && typeof graphName === 'string') { + span.setAttribute(GEN_AI_PIPELINE_NAME, graphName); + span.setAttribute(GEN_AI_AGENT_NAME, graphName); + span.updateName(`invoke_agent ${graphName}`); + } + + const modelName = llm?.modelName ?? llm?.model; + if (modelName) { + span.setAttribute(GEN_AI_REQUEST_MODEL, modelName); + } + + // Inject callback handler and agent name into invoke config + if (sentryCallbackHandler) { + const invokeConfig = (args[1] ?? {}) as Record; + args[1] = invokeConfig; + + const existingMetadata = (invokeConfig.metadata ?? {}) as Record; + invokeConfig.metadata = { + ...getConversationIdMetadataFromConfig(invokeConfig), + ...existingMetadata, + __sentry_langgraph__: true, + ...(typeof graphName === 'string' ? { lc_agent_name: graphName } : {}), + }; + + invokeConfig.callbacks = _INTERNAL_mergeLangChainCallbackHandler(invokeConfig.callbacks, sentryCallbackHandler); + } + + const recordInputs = options.recordInputs; + + // Extract available tools from the graph instance + const tools = recordInputs ? extractToolsFromCompiledGraph(graphInstance) : null; + if (tools) { + span.setAttribute(GEN_AI_TOOL_DEFINITIONS, JSON.stringify(tools)); + } + + // Custom state annotations have no `messages` array, the whole state is recorded instead. + const inputState = args[0] as { messages?: LangChainMessage[]; lg_name?: string } | null | undefined; + const inputMessages = Array.isArray(inputState?.messages) ? inputState.messages : null; + // `new Command({ resume })` resumes an interrupted run and carries no user turn (LangGraph + // tags it `lg_name: 'Command'`), so skip it like a `null` resume rather than recording the + // control object as a message nobody wrote. + const isResumeCommand = inputState?.lg_name === 'Command'; + + if (recordInputs) { + if (inputMessages) { + const normalizedMessages = normalizeLangChainMessages(inputMessages); + const { systemInstructions, filteredMessages } = extractSystemInstructions(normalizedMessages); + + if (systemInstructions) { + span.setAttribute(GEN_AI_SYSTEM_INSTRUCTIONS, systemInstructions); + } + + span.setAttributes({ + [GEN_AI_INPUT_MESSAGES]: stringify(filteredMessages), + }); + } else if (inputState && typeof inputState === 'object' && !isResumeCommand) { + span.setAttribute(GEN_AI_INPUT_MESSAGES, stringify([{ role: 'user', content: stringify(inputState) }])); + } + } + + return inputMessages; +} diff --git a/packages/server-utils/src/ai/langgraph/index.ts b/packages/server-utils/src/ai/langgraph/index.ts index 0b1dc9a566d6..8bc775830aef 100644 --- a/packages/server-utils/src/ai/langgraph/index.ts +++ b/packages/server-utils/src/ai/langgraph/index.ts @@ -1,42 +1,38 @@ /* eslint-disable typescript-eslint/no-deprecated */ -import { SPAN_STATUS_ERROR, startSpan, stringify } from '@sentry/core'; -import { - GEN_AI_AGENT_NAME, - GEN_AI_INPUT_MESSAGES, - GEN_AI_OPERATION_NAME, - GEN_AI_PIPELINE_NAME, - GEN_AI_REQUEST_MODEL, - GEN_AI_SYSTEM_INSTRUCTIONS, - GEN_AI_TOOL_DEFINITIONS, - SENTRY_OP, - SENTRY_ORIGIN, -} from '@sentry/conventions/attributes'; -import { GEN_AI_INVOKE_AGENT } from '@sentry/conventions/op'; -import { extractSystemInstructions, resolveAIRecordingOptions } from '../core/utils'; +import { getCurrentScope, SPAN_STATUS_ERROR, startSpan } from '@sentry/core'; +import { resolveAIRecordingOptions } from '../core/utils'; import { createLangChainCallbackHandler } from '../langchain'; -import type { BaseChatModel, LangChainMessage } from '../langchain/types'; -import { normalizeLangChainMessages } from '../langchain/utils'; -import { LANGGRAPH_ORIGIN } from './constants'; +import type { BaseChatModel } from '../langchain/types'; +import { getAgentSpanOptions, getInvokeMarker, LANGGRAPH_INVOKE_ACTIVE, prepareAgentSpan } from './agent-span'; +import { instrumentCompiledGraphStream } from './stream'; import type { CompiledGraph, LangGraphOptions } from './types'; -import { - extractAgentNameFromParams, - extractLLMFromParams, - extractToolsFromCompiledGraph, - setResponseAttributes, - wrapToolsWithSpans, -} from './utils'; -import { - getConversationIdMetadataFromConfig, - _INTERNAL_mergeLangChainCallbackHandler, - getConversationIdFromMetadata, -} from '../langchain/utils'; +import { extractAgentNameFromParams, extractLLMFromParams, setResponseAttributes, wrapToolsWithSpans } from './utils'; let _insideCreateReactAgent = false; +// Set while a manual wrapper runs the original `compile`/`createReactAgent`. The integration's channel hooks +// fire synchronously inside that call, and skip the graph so the options passed explicitly win. +let _insideManualInstrumentation = false; const SENTRY_PATCHED = '__sentry_patched__'; +const INSTRUMENTED_GRAPHS = new WeakSet(); + +/** Whether manual `instrumentStateGraph()`/`instrumentCreateReactAgent()` is creating a graph right now. */ +export function isInsideManualLangGraphInstrumentation(): boolean { + return _insideManualInstrumentation; +} + +function runManually(callback: () => T): T { + const previous = _insideManualInstrumentation; + _insideManualInstrumentation = true; + try { + return callback(); + } finally { + _insideManualInstrumentation = previous; + } +} /** - * Instruments StateGraph's compile method to wrap the returned compiled graph's invoke() with a + * Instruments StateGraph's compile method to wrap the returned compiled graph's invoke() and stream() with a * `gen_ai.invoke_agent` span. */ export function instrumentStateGraphCompile( @@ -59,21 +55,9 @@ export function instrumentStateGraphCompile( return Reflect.apply(target, thisArg, args); } - const compiledGraph = Reflect.apply(target, thisArg, args); + const compiledGraph = runManually(() => Reflect.apply(target, thisArg, args)); const compileOptions = args.length > 0 ? (args[0] as Record) : {}; - - // Instrument agent invoke method on the compiled graph - const originalInvoke = compiledGraph.invoke; - if (originalInvoke && typeof originalInvoke === 'function') { - compiledGraph.invoke = instrumentCompiledGraphInvoke( - originalInvoke.bind(compiledGraph) as (...args: unknown[]) => Promise, - compiledGraph, - compileOptions, - options, - undefined, - sentryHandler, - ); - } + instrumentCompiledGraph(compiledGraph, compileOptions, options, undefined, sentryHandler); return compiledGraph; }, @@ -83,6 +67,45 @@ 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' || INSTRUMENTED_GRAPHS.has(graph)) { + return; + } + INSTRUMENTED_GRAPHS.add(graph); + + const compiledGraph = graph as CompiledGraph; + const { invoke, stream } = compiledGraph; + if (typeof invoke === 'function') { + compiledGraph.invoke = instrumentCompiledGraphInvoke( + invoke.bind(compiledGraph), + compiledGraph, + compileOptions, + options, + llm, + sentryCallbackHandler, + ); + } + if (typeof stream === 'function') { + compiledGraph.stream = instrumentCompiledGraphStream( + stream.bind(compiledGraph), + compiledGraph, + compileOptions, + options, + llm, + sentryCallbackHandler, + ); + } +} + /** * Instruments CompiledGraph's invoke method to create spans for agent invocation * @@ -96,104 +119,37 @@ export function instrumentCompiledGraphInvoke( llm?: BaseChatModel | null, sentryCallbackHandler?: unknown, ): (...args: unknown[]) => Promise { + const invokeMarker = getInvokeMarker(graphInstance); + return new Proxy(originalInvoke, { apply(target, thisArg, args: unknown[]): Promise { - const modelName = llm?.modelName ?? llm?.model; - const config = args.length > 1 ? (args[1] as Record | undefined) : undefined; - return startSpan( - { - name: 'invoke_agent', - attributes: { - // Set before `spanStart`, so an id from `Sentry.setConversationId()` wins, as on the child spans - ...getConversationIdFromMetadata(config?.configurable as Record | undefined), - [SENTRY_ORIGIN]: LANGGRAPH_ORIGIN, - [SENTRY_OP]: GEN_AI_INVOKE_AGENT, - [GEN_AI_OPERATION_NAME]: 'invoke_agent', - }, - }, - 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); - } - - // Inject callback handler and agent name into invoke config - if (sentryCallbackHandler) { - const invokeConfig = (args[1] ?? {}) as Record; - args[1] = invokeConfig; - - const existingMetadata = (invokeConfig.metadata ?? {}) as Record; - invokeConfig.metadata = { - ...getConversationIdMetadataFromConfig(invokeConfig), - ...existingMetadata, - __sentry_langgraph__: true, - ...(typeof graphName === 'string' ? { lc_agent_name: graphName } : {}), - }; - - invokeConfig.callbacks = _INTERNAL_mergeLangChainCallbackHandler( - invokeConfig.callbacks, - sentryCallbackHandler, - ); - } - - const recordInputs = options.recordInputs; - const recordOutputs = options.recordOutputs; - - // Extract available tools from the graph instance - const tools = recordInputs ? extractToolsFromCompiledGraph(graphInstance) : null; - if (tools) { - span.setAttribute(GEN_AI_TOOL_DEFINITIONS, JSON.stringify(tools)); - } - - // Custom state annotations have no `messages` array, the whole state is recorded instead. - const inputState = args[0] as { messages?: LangChainMessage[]; lg_name?: string } | null | undefined; - const inputMessages = Array.isArray(inputState?.messages) ? inputState.messages : null; - // `new Command({ resume })` resumes an interrupted run and carries no user turn (LangGraph - // tags it `lg_name: 'Command'`), so skip it like a `null` resume rather than recording the - // control object as a message nobody wrote. - const isResumeCommand = inputState?.lg_name === 'Command'; - - if (recordInputs) { - if (inputMessages) { - const normalizedMessages = normalizeLangChainMessages(inputMessages); - const { systemInstructions, filteredMessages } = extractSystemInstructions(normalizedMessages); - - if (systemInstructions) { - span.setAttribute(GEN_AI_SYSTEM_INSTRUCTIONS, systemInstructions); - } - - span.setAttributes({ - [GEN_AI_INPUT_MESSAGES]: stringify(filteredMessages), - }); - } else if (inputState && typeof inputState === 'object' && !isResumeCommand) { - span.setAttribute(GEN_AI_INPUT_MESSAGES, stringify([{ role: 'user', content: stringify(inputState) }])); - } - } - - // Call original invoke - const result = await Reflect.apply(target, thisArg, args); - - if (recordOutputs) { - setResponseAttributes(span, inputMessages, result); - } - - 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; + return startSpan(getAgentSpanOptions(args), async span => { + try { + const inputMessages = prepareAgentSpan( + span, + args, + graphInstance, + compileOptions, + options, + llm, + sentryCallbackHandler, + ); + getCurrentScope().setSDKProcessingMetadata({ [LANGGRAPH_INVOKE_ACTIVE]: invokeMarker }); + + const result = await Reflect.apply(target, thisArg, args); + + if (options.recordOutputs) { + setResponseAttributes(span, inputMessages, result); } - }, - ); + + 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; + } + }); }, }); } @@ -227,28 +183,13 @@ export function instrumentCreateReactAgent( _insideCreateReactAgent = true; let compiledGraph: CompiledGraph; try { - compiledGraph = Reflect.apply(target, thisArg, args); + compiledGraph = runManually(() => Reflect.apply(target, thisArg, args)); } finally { _insideCreateReactAgent = false; } - // Wrap invoke() on the returned compiled graph - const originalInvoke = compiledGraph.invoke; - if (originalInvoke && typeof originalInvoke === 'function') { - const compileOptions: Record = {}; - if (agentName) { - compileOptions.name = agentName; - } - - compiledGraph.invoke = instrumentCompiledGraphInvoke( - originalInvoke.bind(compiledGraph) as (...args: unknown[]) => Promise, - compiledGraph, - compileOptions, - resolvedOptions, - llm, - sentryHandler, - ); - } + const compileOptions: Record = agentName ? { name: agentName } : {}; + instrumentCompiledGraph(compiledGraph, compileOptions, resolvedOptions, llm, sentryHandler); return compiledGraph; }, diff --git a/packages/server-utils/src/ai/langgraph/stream.ts b/packages/server-utils/src/ai/langgraph/stream.ts new file mode 100644 index 000000000000..007490dd0b49 --- /dev/null +++ b/packages/server-utils/src/ai/langgraph/stream.ts @@ -0,0 +1,155 @@ +import { getCurrentScope, SPAN_STATUS_ERROR, startSpanManual } from '@sentry/core'; +import { GEN_AI_RESPONSE_STREAMING } from '@sentry/conventions/attributes'; +import type { BaseChatModel, LangChainMessage } from '../langchain/types'; +import { _INTERNAL_addLangChainCallbackHandler } from '../langchain/utils'; +import { getAgentSpanOptions, getInvokeMarker, LANGGRAPH_INVOKE_ACTIVE, prepareAgentSpan } from './agent-span'; +import type { CompiledGraph, LangGraphOptions } from './types'; +import { setResponseAttributes } from './utils'; + +/** + * Instruments CompiledGraph's stream method with a `gen_ai.invoke_agent` span that lasts for the graph run. + * + * LangGraph runs the graph as soon as `stream()` is called, whether or not the stream is read, so the span + * ends on the run's own end or error callback rather than when the stream is consumed. + */ +export function instrumentCompiledGraphStream( + originalStream: (...args: unknown[]) => Promise, + graphInstance: CompiledGraph, + compileOptions: Record, + options: LangGraphOptions, + llm?: BaseChatModel | null, + sentryCallbackHandler?: unknown, +): (...args: unknown[]) => Promise { + const invokeMarker = getInvokeMarker(graphInstance); + + return new Proxy(originalStream, { + apply(target, thisArg, args: unknown[]): Promise { + const scope = getCurrentScope(); + if (scope.getScopeData().sdkProcessingMetadata[LANGGRAPH_INVOKE_ACTIVE] === invokeMarker) { + scope.setSDKProcessingMetadata({ [LANGGRAPH_INVOKE_ACTIVE]: undefined }); + return Reflect.apply(target, thisArg, args); + } + + // A copy, so the per-call listener never lands on a config object the caller reuses. + const config: Record = { ...(args[1] as Record | undefined) }; + const callArgs = [args[0], config, ...args.slice(2)]; + const spanOptions = getAgentSpanOptions(callArgs); + + return startSpanManual( + { ...spanOptions, attributes: { ...spanOptions.attributes, [GEN_AI_RESPONSE_STREAMING]: true } }, + async span => { + let inputMessages: LangChainMessage[] | null = null; + let stream: { signal?: AbortSignal } | undefined; + let ended = false; + const finish = (outcome?: { outputs: unknown } | { error: unknown }): void => { + if (ended) { + return; + } + ended = true; + try { + if (outcome && 'error' in outcome) { + // A cancel or abort by the caller is not a failed run. LangGraph reports both, and step + // timeouts too, as `Error('Abort')`, so only the signals tell them apart. + if (!stream?.signal?.aborted && !(config.signal as AbortSignal | undefined)?.aborted) { + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); + } + } else if (outcome && options.recordOutputs) { + setResponseAttributes(span, inputMessages, outcome.outputs); + } + } catch { + // Recording must never change the outcome of the run. + } + span.end(); + }; + + const listener = createRootRunListener(finish); + try { + inputMessages = prepareAgentSpan( + span, + callArgs, + graphInstance, + compileOptions, + options, + llm, + sentryCallbackHandler, + ); + config.callbacks = _INTERNAL_addLangChainCallbackHandler(config.callbacks, listener); + stream = (await Reflect.apply(target, thisArg, callArgs)) as typeof stream; + } catch (error) { + finish({ error }); + throw error; + } + + if (listener.rootRunId === undefined) { + // Not a LangGraph run we can follow, so don't leave the span open. + finish(); + } else { + endOnEarlyExit(stream, finish); + } + return stream; + }, + ); + }, + }); +} + +function createRootRunListener(finish: (outcome: { outputs: unknown } | { error: unknown }) => void): { + rootRunId: string | undefined; +} { + // The root run starts before any of its children, so it is the first run this per-call listener sees. + const listener = { + name: 'SentryLangGraphStreamListener', + awaitHandlers: true, + rootRunId: undefined as string | undefined, + handleChainStart(_chain: unknown, _inputs: unknown, runId: string): void { + listener.rootRunId ??= runId; + }, + handleChainEnd(outputs: unknown, runId: string): void { + if (runId === listener.rootRunId) { + finish({ outputs }); + } + }, + handleChainError(error: unknown, runId: string): void { + if (runId === listener.rootRunId) { + finish({ error }); + } + }, + copy() { + return listener; + }, + }; + return listener; +} + +// After a `break`, LangGraph keeps running the graph but never reports the run's end, so the span ends when +// the caller stops. Nodes still running at that point start spans after the agent span ended. +function endOnEarlyExit(stream: unknown, finish: () => void): void { + if (!stream || typeof stream !== 'object') { + return; + } + const target = stream as { return?: (...args: unknown[]) => unknown; cancel?: (...args: unknown[]) => unknown }; + const { return: originalReturn, cancel: originalCancel } = target; + try { + if (typeof originalReturn === 'function') { + target.return = function (this: unknown, ...args: unknown[]): unknown { + try { + return originalReturn.apply(this, args); + } finally { + finish(); + } + }; + } + if (typeof originalCancel === 'function') { + target.cancel = function (this: unknown, ...args: unknown[]): unknown { + const result = originalCancel.apply(this, args); + Promise.resolve(result).then( + () => finish(), + () => undefined, + ); + return result; + }; + } + } catch { + finish(); + } +} diff --git a/packages/server-utils/src/ai/langgraph/types.ts b/packages/server-utils/src/ai/langgraph/types.ts index 153fa9b4caba..f8ac291c5ee4 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 d9d5bd863144..6097149e4aac 100644 --- a/packages/server-utils/src/integrations/langgraph.ts +++ b/packages/server-utils/src/integrations/langgraph.ts @@ -3,9 +3,9 @@ import type { IntegrationFn } from '@sentry/core'; import { debug, defineIntegration, isObjectLike } from '@sentry/core'; import { resolveAIRecordingOptions } from '../ai/core/utils'; import { createLangChainCallbackHandler } from '../ai/langchain'; -import { instrumentCompiledGraphInvoke } from '../ai/langgraph'; +import { instrumentCompiledGraph, isInsideManualLangGraphInstrumentation } from '../ai/langgraph'; import { LANGGRAPH_INTEGRATION_NAME } from '../ai/langgraph/constants'; -import type { CompiledGraph, LangGraphOptions } from '../ai/langgraph/types'; +import type { LangGraphOptions } from '../ai/langgraph/types'; import { extractAgentNameFromParams, extractLLMFromParams, wrapToolsWithSpans } from '../ai/langgraph/utils'; import { DEBUG_BUILD } from '../debug-build'; import type { OrchestrionChannelContext } from '../orchestrion/types'; @@ -38,19 +38,19 @@ function instrumentLanggraph(options: LangGraphOptions): void { const resolvedOptions = resolveAIRecordingOptions(options); const sentryHandler = createLangChainCallbackHandler(resolvedOptions); - // StateGraph.compile returns synchronously; wrap the returned graph's `invoke` at `end`. + // StateGraph.compile returns synchronously; wrap the returned graph's `invoke` and `stream` at `end`. diagnosticsChannel .tracingChannel(CHANNELS.LANGGRAPH_STATE_GRAPH_COMPILE) .end.subscribe(message => { - if (insideCreateReactAgent) { + if (insideCreateReactAgent || isInsideManualLangGraphInstrumentation()) { return; } const { arguments: args, result } = message as OrchestrionChannelContext; - wrapCompiledGraphInvoke(result, getFirstArgObject(args) ?? {}, resolvedOptions, null, sentryHandler); + instrumentCompiledGraph(result, getFirstArgObject(args) ?? {}, resolvedOptions, null, sentryHandler); }); - // createReactAgent only wraps tools and the returned graph's `invoke`. Tools are wrapped at - // `start` (before the agent runs), invoke at `end`. + // createReactAgent only wraps tools and the returned graph's `invoke` and `stream`. Tools are wrapped at + // `start` (before the agent runs), the graph at `end`. const reactAgentChannel = diagnosticsChannel.tracingChannel( CHANNELS.LANGGRAPH_CREATE_REACT_AGENT, ); @@ -72,10 +72,13 @@ function instrumentLanggraph(options: LangGraphOptions): void { }); reactAgentChannel.end.subscribe(message => { insideCreateReactAgent = false; + if (isInsideManualLangGraphInstrumentation()) { + return; + } const { arguments: args, result } = message as OrchestrionChannelContext; const agentName = extractAgentNameFromParams(args) ?? undefined; const compileOptions = agentName ? { name: agentName } : {}; - wrapCompiledGraphInvoke(result, compileOptions, resolvedOptions, extractLLMFromParams(args), sentryHandler); + instrumentCompiledGraph(result, compileOptions, resolvedOptions, extractLLMFromParams(args), sentryHandler); }); // Make sure a thrown `createReactAgent` doesn't leave the suppression flag stuck on. reactAgentChannel.error.subscribe(() => { @@ -89,35 +92,6 @@ function getFirstArgObject(args: unknown[] | undefined): Record return isObjectLike(first) ? first : undefined; } -/** - * Wrap the compiled graph's `invoke` with the shared `invoke_agent` instrumentation, exactly as the - * OTel path does on the returned graph. - */ -function wrapCompiledGraphInvoke( - graph: unknown, - compileOptions: Record, - options: LangGraphOptions, - llm: ReturnType, - sentryHandler: unknown, -): void { - if (!graph || typeof graph !== 'object') { - return; - } - - const compiledGraph = graph as CompiledGraph; - const originalInvoke = compiledGraph.invoke; - if (typeof originalInvoke === 'function') { - compiledGraph.invoke = instrumentCompiledGraphInvoke( - originalInvoke.bind(compiledGraph), - compiledGraph, - compileOptions, - options, - llm, - sentryHandler, - ); - } -} - /** * Diagnostics-channel-based LangGraph integration. Subscribes to the diagnostics_channels * injected into `@langchain/langgraph`'s `StateGraph.compile` and `createReactAgent`, so it requires diff --git a/packages/server-utils/test/ai/lib/tracing/langchain-utils.test.ts b/packages/server-utils/test/ai/lib/tracing/langchain-utils.test.ts index a8038425e158..660ea39a3f3c 100644 --- a/packages/server-utils/test/ai/lib/tracing/langchain-utils.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/langchain-utils.test.ts @@ -2,6 +2,7 @@ import { GEN_AI_CONVERSATION_ID } from '@sentry/conventions/attributes'; import { describe, expect, it, vi } from 'vitest'; import type { LangChainMessage } from '../../../../src/ai/langchain/types'; import { + _INTERNAL_addLangChainCallbackHandler, _INTERNAL_mergeLangChainCallbackHandler, getConversationIdFromMetadata, getConversationIdMetadataFromConfig, @@ -68,6 +69,34 @@ describe('normalizeLangChainMessages', () => { }); }); +describe('_INTERNAL_addLangChainCallbackHandler', () => { + const sentryHandler = { name: 'SentryCallbackHandler' }; + const listener = { name: 'Listener' }; + + it('appends even when the sentry handler is already present', () => { + expect(_INTERNAL_addLangChainCallbackHandler([sentryHandler], listener)).toStrictEqual([sentryHandler, listener]); + expect(_INTERNAL_addLangChainCallbackHandler(sentryHandler, listener)).toStrictEqual([sentryHandler, listener]); + expect(_INTERNAL_addLangChainCallbackHandler(undefined, listener)).toStrictEqual([listener]); + }); + + it('adds to a copy of a callback manager that already holds the sentry handler', () => { + const manager = { + handlers: [sentryHandler] as unknown[], + addHandler: vi.fn(function (this: { handlers: unknown[] }, handler: unknown) { + this.handlers.push(handler); + }), + copy: vi.fn(() => ({ ...manager, handlers: [...manager.handlers] })), + }; + + const result = _INTERNAL_addLangChainCallbackHandler(manager, listener) as typeof manager; + + expect(result).not.toBe(manager); + expect(manager.handlers).toEqual([sentryHandler]); + expect(result.handlers).toEqual([sentryHandler, listener]); + expect(result.addHandler).toHaveBeenCalledWith(listener, true); + }); +}); + describe('_INTERNAL_mergeLangChainCallbackHandler', () => { const sentryHandler = { name: 'SentryCallbackHandler' }; 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..e80dc42fcda2 --- /dev/null +++ b/packages/server-utils/test/ai/lib/tracing/langgraph-stream.test.ts @@ -0,0 +1,335 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { getMainCarrier, setCurrentClient, spanToStaticSpanJSON, startSpan } from '@sentry/core'; +import type { Span } from '@sentry/core'; +import { GEN_AI_RESPONSE_STREAMING, GEN_AI_RESPONSE_TEXT } from '@sentry/conventions/attributes'; +import { instrumentCompiledGraph } from '../../../../src/ai/langgraph'; +import { getDefaultTestClientOptions, TestClient } from '../../../mocks/client'; + +type Handler = Record void) | undefined>; +type Emit = (event: string, ...args: unknown[]) => void; + +// The parts of LangGraph's `IterableReadableStreamWithAbortSignal` the instrumentation relies on: the stream is +// its own async iterator, `return()` is what a `break` calls, and `cancel()` aborts `signal`. +function createStream(chunks: unknown[] = ['chunk']) { + const controller = new AbortController(); + let index = 0; + const stream = { + signal: controller.signal, + [Symbol.asyncIterator]() { + return stream; + }, + async next(): Promise> { + return index < chunks.length ? { done: false, value: chunks[index++] } : { done: true, value: undefined }; + }, + return: vi.fn(async (): Promise> => ({ done: true, value: undefined })), + cancel: vi.fn(async (): Promise => controller.abort()), + }; + return stream; +} + +function getHandlers(callbacks: unknown): Handler[] { + if (Array.isArray(callbacks)) { + return callbacks as Handler[]; + } + return ((callbacks as { handlers?: Handler[] } | undefined)?.handlers ?? []) as Handler[]; +} + +// A compiled graph whose `stream()` reports its root run to `config.callbacks` the way LangGraph does: the root +// run starts before `stream()` resolves, and ends or fails later. +function createGraph(run: (emit: Emit, input: unknown) => unknown) { + return { + async stream(input: unknown, config?: { callbacks?: unknown }): Promise { + const emit: Emit = (event, ...args) => + getHandlers(config?.callbacks).forEach(handler => handler[event]?.(...args)); + emit('handleChainStart', {}, input, 'root'); + return run(emit, input); + }, + }; +} + +describe('LangGraph stream() agent span', () => { + let endedSpans: Span[]; + + beforeEach(() => { + getMainCarrier().__SENTRY__ = undefined; + const client = new TestClient( + getDefaultTestClientOptions({ dsn: 'https://public@dsn.ingest.sentry.io/1337', tracesSampleRate: 1 }), + ); + setCurrentClient(client); + client.init(); + endedSpans = []; + client.on('spanEnd', span => endedSpans.push(span)); + }); + + afterEach(() => { + getMainCarrier().__SENTRY__ = undefined; + }); + + function instrument(graph: T, options = {}): T { + instrumentCompiledGraph(graph, { name: 'weather_graph' }, options, undefined, undefined); + return graph; + } + + const agentSpans = (): Span[] => endedSpans.filter(span => spanToStaticSpanJSON(span).op === 'gen_ai.invoke_agent'); + + it('ends the span when the root run ends, not when a child run does', async () => { + let emit!: Emit; + const graph = instrument( + createGraph(e => { + emit = e; + return createStream(); + }), + { recordOutputs: true }, + ); + + await graph.stream({ messages: [{ role: 'user', content: 'Weather?' }] }); + emit('handleChainStart', {}, {}, 'child', 'root'); + emit('handleChainEnd', {}, 'child', 'root'); + expect(agentSpans()).toHaveLength(0); + + emit( + 'handleChainEnd', + { + messages: [ + { role: 'user', content: 'Weather?' }, + { role: 'assistant', content: 'Sunny' }, + ], + }, + 'root', + ); + + expect(agentSpans()).toHaveLength(1); + const span = spanToStaticSpanJSON(agentSpans()[0]!); + expect(span.description).toBe('invoke_agent weather_graph'); + expect(span.status).not.toBe('internal_error'); + expect(span.data[GEN_AI_RESPONSE_STREAMING]).toBe(true); + expect(span.data[GEN_AI_RESPONSE_TEXT]).toBe('[{"role":"assistant","content":"Sunny"}]'); + }); + + it('ends the span when the run finishes although the stream is never read', async () => { + const graph = instrument( + createGraph(emit => { + setTimeout(() => emit('handleChainEnd', {}, 'root')); + return createStream(); + }), + ); + + await graph.stream({}); + await vi.waitFor(() => expect(agentSpans()).toHaveLength(1)); + }); + + it('ends the span when the caller breaks out of the stream', async () => { + const stream = createStream(['first', 'second']); + const originalReturn = stream.return; + const graph = instrument(createGraph(() => stream)); + + for await (const _chunk of (await graph.stream({})) as AsyncIterable) { + break; + } + + expect(originalReturn).toHaveBeenCalledTimes(1); + expect(agentSpans()).toHaveLength(1); + expect(spanToStaticSpanJSON(agentSpans()[0]!).status).not.toBe('internal_error'); + }); + + it('passes a rejection from the original return() through to the caller', async () => { + const error = new Error('return failed'); + const stream = createStream(); + stream.return.mockRejectedValueOnce(error); + const graph = instrument(createGraph(() => stream)); + + const consume = async (): Promise => { + for await (const _chunk of (await graph.stream({})) as AsyncIterable) { + break; + } + }; + + await expect(consume()).rejects.toBe(error); + expect(agentSpans()).toHaveLength(1); + }); + + it('treats a cancel by the caller as a finished run', async () => { + let emit!: Emit; + const graph = instrument( + createGraph(e => { + emit = e; + return createStream(); + }), + ); + + const stream = (await graph.stream({})) as ReturnType; + await stream.cancel(); + emit('handleChainError', new Error('Abort'), 'root'); + + expect(agentSpans()).toHaveLength(1); + expect(spanToStaticSpanJSON(agentSpans()[0]!).status).not.toBe('internal_error'); + }); + + it('keeps the span open when cancel() is rejected, and still reports a later failure', async () => { + let emit!: Emit; + const graph = instrument( + createGraph(e => { + emit = e; + const stream = createStream(); + stream.cancel.mockRejectedValueOnce(new TypeError('locked')); + return stream; + }), + ); + + const stream = (await graph.stream({})) as ReturnType; + await expect(stream.cancel()).rejects.toThrow(TypeError); + expect(agentSpans()).toHaveLength(0); + + emit('handleChainError', new Error('node failed'), 'root'); + expect(agentSpans()).toHaveLength(1); + expect(spanToStaticSpanJSON(agentSpans()[0]!).status).toBe('internal_error'); + }); + + it('does not mark an aborted run as failed when the caller aborted its own signal', async () => { + const controller = new AbortController(); + const graph = instrument( + createGraph(emit => { + controller.abort(); + emit('handleChainError', new Error('Abort'), 'root'); + return createStream(); + }), + ); + + await graph.stream({}, { signal: controller.signal }); + + expect(agentSpans()).toHaveLength(1); + expect(spanToStaticSpanJSON(agentSpans()[0]!).status).not.toBe('internal_error'); + }); + + it('marks the span failed and rethrows when stream() rejects', async () => { + const error = new Error('bad input'); + const graph = instrument( + createGraph(() => { + throw error; + }), + ); + + await expect(graph.stream({})).rejects.toBe(error); + expect(agentSpans()).toHaveLength(1); + expect(spanToStaticSpanJSON(agentSpans()[0]!).status).toBe('internal_error'); + }); + + it('ends the span right away when stream() reports no run to follow', async () => { + const graph = instrument({ stream: async () => createStream() }); + + await graph.stream(); + + expect(agentSpans()).toHaveLength(1); + }); + + it('still ends the span when recording the output fails', async () => { + let emit!: Emit; + const graph = instrument( + createGraph(e => { + emit = e; + return createStream(); + }), + { recordOutputs: true }, + ); + const outputs = { + get messages(): never { + throw new Error('unreadable state'); + }, + }; + + await graph.stream({}); + expect(() => emit('handleChainEnd', outputs, 'root')).not.toThrow(); + expect(agentSpans()).toHaveLength(1); + }); + + it('does not add the per-call listener to the config object the caller passed', async () => { + const callerHandler = { name: 'callerHandler' }; + const config = { callbacks: [callerHandler], metadata: { user: 'u' } }; + const graph = instrument(createGraph(() => createStream())); + + await graph.stream({}, config); + + expect(config).toEqual({ callbacks: [callerHandler], metadata: { user: 'u' } }); + }); + + it('copies a callback manager instead of adding the listener to it', async () => { + let received: unknown; + const manager = { + handlers: [] as unknown[], + addHandler(handler: unknown) { + this.handlers.push(handler); + }, + copy() { + return { ...manager, handlers: [...manager.handlers], addHandler: manager.addHandler, copy: manager.copy }; + }, + }; + const graph = instrument({ + async stream(_input: unknown, config: { callbacks: unknown }) { + received = config.callbacks; + return createStream(); + }, + }); + + await graph.stream({}, { callbacks: manager }); + + expect(manager.handlers).toHaveLength(0); + expect(received).not.toBe(manager); + expect((received as typeof manager).handlers).toHaveLength(1); + }); + + it('creates one span for invoke(), which streams internally', async () => { + const graph = instrument({ + ...createGraph(emit => { + emit('handleChainEnd', {}, 'root'); + return createStream(); + }), + async invoke(this: { stream: (input: unknown) => Promise }, input: unknown) { + for await (const _chunk of (await this.stream(input)) as AsyncIterable) { + // drain, as LangGraph's invoke does + } + return {}; + }, + }); + + await graph.invoke({}); + + expect(agentSpans()).toHaveLength(1); + expect(spanToStaticSpanJSON(agentSpans()[0]!).data[GEN_AI_RESPONSE_STREAMING]).toBeUndefined(); + }); + + it('creates a child span for a stream() of the same graph started inside invoke()', async () => { + const graph = instrument({ + ...createGraph(emit => { + emit('handleChainEnd', {}, 'root'); + return createStream(); + }), + async invoke(this: { stream: (input: unknown) => Promise }, input: unknown) { + await this.stream(input); + await this.stream(input); + return {}; + }, + }); + + await graph.invoke({}); + + const [streamSpan, invokeSpan] = agentSpans(); + expect(agentSpans()).toHaveLength(2); + expect(spanToStaticSpanJSON(streamSpan!).parent_span_id).toBe(invokeSpan!.spanContext().spanId); + expect(spanToStaticSpanJSON(streamSpan!).data[GEN_AI_RESPONSE_STREAMING]).toBe(true); + }); + + it('wraps a graph only once when it is instrumented twice', async () => { + const graph = instrument( + instrument( + createGraph(emit => { + emit('handleChainEnd', {}, 'root'); + return createStream(); + }), + ), + ); + + await startSpan({ name: 'root' }, () => graph.stream({})); + + expect(agentSpans()).toHaveLength(1); + }); +}); From dd0ab81972272f2502c4fcc1376e507e0e8d164a Mon Sep 17 00:00:00 2001 From: RulaKhaled Date: Mon, 26 Oct 2026 20:20:00 +0100 Subject: [PATCH 2/4] ref(server-utils): Simplify LangGraph stream listener and callback helpers Co-Authored-By: Claude Opus 5.5 --- .../server-utils/src/ai/langchain/utils.ts | 17 +- .../src/ai/langgraph/agent-span.ts | 6 +- .../server-utils/src/ai/langgraph/index.ts | 1 - .../server-utils/src/ai/langgraph/stream.ts | 148 ++++++++---------- 4 files changed, 81 insertions(+), 91 deletions(-) diff --git a/packages/server-utils/src/ai/langchain/utils.ts b/packages/server-utils/src/ai/langchain/utils.ts index 5cae51d3ac53..405340e71c56 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()/stream() wrappers (`prepareAgentSpan`) const agentName = metadata?.lc_agent_name; if (typeof agentName === 'string') { attrs[GEN_AI_AGENT_NAME] = agentName; @@ -606,14 +606,7 @@ function containsSentryHandler(handlers: unknown[]): boolean { * @internal Exported for cross-package instrumentation. */ export function _INTERNAL_mergeLangChainCallbackHandler(existing: unknown, sentryHandler: unknown): unknown { - const handlers = isCallbackManager(existing) - ? (existing.handlers ?? []) - : Array.isArray(existing) - ? existing - : existing - ? [existing] - : []; - + const handlers = isCallbackManager(existing) ? (existing.handlers ?? []) : toHandlerArray(existing); return containsSentryHandler(handlers) ? existing : _INTERNAL_addLangChainCallbackHandler(existing, sentryHandler); } @@ -633,5 +626,9 @@ export function _INTERNAL_addLangChainCallbackHandler(existing: unknown, handler return copied; } - return [...(Array.isArray(existing) ? existing : [existing]), handler]; + return [...toHandlerArray(existing), handler]; +} + +function toHandlerArray(existing: unknown): unknown[] { + return Array.isArray(existing) ? existing : existing ? [existing] : []; } diff --git a/packages/server-utils/src/ai/langgraph/agent-span.ts b/packages/server-utils/src/ai/langgraph/agent-span.ts index 4bfbb075f06d..f2bce259eb8b 100644 --- a/packages/server-utils/src/ai/langgraph/agent-span.ts +++ b/packages/server-utils/src/ai/langgraph/agent-span.ts @@ -40,7 +40,10 @@ export function getInvokeMarker(graph: object): symbol { } /** Name and initial attributes of an agent span, from the call's config. */ -export function getAgentSpanOptions(args: unknown[]): { name: string; attributes: SpanAttributes } { +export function getAgentSpanOptions( + args: unknown[], + extraAttributes?: SpanAttributes, +): { name: string; attributes: SpanAttributes } { const config = args.length > 1 ? (args[1] as Record | undefined) : undefined; return { name: 'invoke_agent', @@ -50,6 +53,7 @@ export function getAgentSpanOptions(args: unknown[]): { name: string; attributes [SENTRY_ORIGIN]: LANGGRAPH_ORIGIN, [SENTRY_OP]: GEN_AI_INVOKE_AGENT, [GEN_AI_OPERATION_NAME]: 'invoke_agent', + ...extraAttributes, }, }; } diff --git a/packages/server-utils/src/ai/langgraph/index.ts b/packages/server-utils/src/ai/langgraph/index.ts index 8bc775830aef..6d7e3458342b 100644 --- a/packages/server-utils/src/ai/langgraph/index.ts +++ b/packages/server-utils/src/ai/langgraph/index.ts @@ -1,4 +1,3 @@ -/* eslint-disable typescript-eslint/no-deprecated */ import { getCurrentScope, SPAN_STATUS_ERROR, startSpan } from '@sentry/core'; import { resolveAIRecordingOptions } from '../core/utils'; import { createLangChainCallbackHandler } from '../langchain'; diff --git a/packages/server-utils/src/ai/langgraph/stream.ts b/packages/server-utils/src/ai/langgraph/stream.ts index 007490dd0b49..935120cd13e2 100644 --- a/packages/server-utils/src/ai/langgraph/stream.ts +++ b/packages/server-utils/src/ai/langgraph/stream.ts @@ -33,92 +33,82 @@ export function instrumentCompiledGraphStream( // A copy, so the per-call listener never lands on a config object the caller reuses. const config: Record = { ...(args[1] as Record | undefined) }; const callArgs = [args[0], config, ...args.slice(2)]; - const spanOptions = getAgentSpanOptions(callArgs); - return startSpanManual( - { ...spanOptions, attributes: { ...spanOptions.attributes, [GEN_AI_RESPONSE_STREAMING]: true } }, - async span => { - let inputMessages: LangChainMessage[] | null = null; - let stream: { signal?: AbortSignal } | undefined; - let ended = false; - const finish = (outcome?: { outputs: unknown } | { error: unknown }): void => { - if (ended) { - return; - } - ended = true; - try { - if (outcome && 'error' in outcome) { - // A cancel or abort by the caller is not a failed run. LangGraph reports both, and step - // timeouts too, as `Error('Abort')`, so only the signals tell them apart. - if (!stream?.signal?.aborted && !(config.signal as AbortSignal | undefined)?.aborted) { - span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); - } - } else if (outcome && options.recordOutputs) { - setResponseAttributes(span, inputMessages, outcome.outputs); + return startSpanManual(getAgentSpanOptions(callArgs, { [GEN_AI_RESPONSE_STREAMING]: true }), async span => { + let inputMessages: LangChainMessage[] | null = null; + let stream: { signal?: AbortSignal } | undefined; + let ended = false; + const finish = (outcome?: { outputs: unknown } | { error: unknown }): void => { + if (ended) { + return; + } + ended = true; + try { + if (outcome && 'error' in outcome) { + // A cancel or abort by the caller is not a failed run. LangGraph reports both, and step + // timeouts too, as `Error('Abort')`, so only the signals tell them apart. + if (!stream?.signal?.aborted && !(config.signal as AbortSignal | undefined)?.aborted) { + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); } - } catch { - // Recording must never change the outcome of the run. + } else if (outcome && options.recordOutputs) { + setResponseAttributes(span, inputMessages, outcome.outputs); } - span.end(); - }; - - const listener = createRootRunListener(finish); - try { - inputMessages = prepareAgentSpan( - span, - callArgs, - graphInstance, - compileOptions, - options, - llm, - sentryCallbackHandler, - ); - config.callbacks = _INTERNAL_addLangChainCallbackHandler(config.callbacks, listener); - stream = (await Reflect.apply(target, thisArg, callArgs)) as typeof stream; - } catch (error) { - finish({ error }); - throw error; + } catch { + // Recording must never change the outcome of the run. } + span.end(); + }; - if (listener.rootRunId === undefined) { - // Not a LangGraph run we can follow, so don't leave the span open. - finish(); - } else { - endOnEarlyExit(stream, finish); - } - return stream; - }, - ); - }, - }); -} + // The root run starts before any of its children, so it is the first run this per-call listener sees. + let rootRunId: string | undefined; + const listener = { + name: 'SentryLangGraphStreamListener', + awaitHandlers: true, + handleChainStart(_chain: unknown, _inputs: unknown, runId: string): void { + rootRunId ??= runId; + }, + handleChainEnd(outputs: unknown, runId: string): void { + if (runId === rootRunId) { + finish({ outputs }); + } + }, + handleChainError(error: unknown, runId: string): void { + if (runId === rootRunId) { + finish({ error }); + } + }, + copy() { + return listener; + }, + }; -function createRootRunListener(finish: (outcome: { outputs: unknown } | { error: unknown }) => void): { - rootRunId: string | undefined; -} { - // The root run starts before any of its children, so it is the first run this per-call listener sees. - const listener = { - name: 'SentryLangGraphStreamListener', - awaitHandlers: true, - rootRunId: undefined as string | undefined, - handleChainStart(_chain: unknown, _inputs: unknown, runId: string): void { - listener.rootRunId ??= runId; - }, - handleChainEnd(outputs: unknown, runId: string): void { - if (runId === listener.rootRunId) { - finish({ outputs }); - } - }, - handleChainError(error: unknown, runId: string): void { - if (runId === listener.rootRunId) { - finish({ error }); - } - }, - copy() { - return listener; + try { + inputMessages = prepareAgentSpan( + span, + callArgs, + graphInstance, + compileOptions, + options, + llm, + sentryCallbackHandler, + ); + config.callbacks = _INTERNAL_addLangChainCallbackHandler(config.callbacks, listener); + stream = (await Reflect.apply(target, thisArg, callArgs)) as typeof stream; + } catch (error) { + finish({ error }); + throw error; + } + + if (rootRunId === undefined) { + // Not a LangGraph run we can follow, so don't leave the span open. + finish(); + } else { + endOnEarlyExit(stream, finish); + } + return stream; + }); }, - }; - return listener; + }); } // After a `break`, LangGraph keeps running the graph but never reports the run's end, so the span ends when From 151cd766b5363b51abe47984112a5ac6da0869d2 Mon Sep 17 00:00:00 2001 From: RulaKhaled Date: Mon, 26 Oct 2026 20:33:59 +0100 Subject: [PATCH 3/4] fix(server-utils): Keep the caller's LangGraph config and inherited callbacks invoke() and stream() now copy the config instead of writing Sentry's metadata and callbacks into the caller's object, and start from the parent run's callbacks when none are passed, as LangChain would. Co-Authored-By: Claude Opus 5.5 --- .../tracing/langgraph/scenario-stream.mjs | 20 ++++++++++++++++ .../suites/tracing/langgraph/test.ts | 7 +++++- .../src/ai/langgraph/agent-span.ts | 12 ++++++---- .../server-utils/src/ai/langgraph/stream.ts | 10 ++++++-- .../ai/lib/tracing/langgraph-stream.test.ts | 24 +++++++++++++++++++ 5 files changed, 66 insertions(+), 7 deletions(-) diff --git a/dev-packages/node-integration-tests/suites/tracing/langgraph/scenario-stream.mjs b/dev-packages/node-integration-tests/suites/tracing/langgraph/scenario-stream.mjs index 086e9caa4b76..2f6f441301a6 100644 --- a/dev-packages/node-integration-tests/suites/tracing/langgraph/scenario-stream.mjs +++ b/dev-packages/node-integration-tests/suites/tracing/langgraph/scenario-stream.mjs @@ -1,3 +1,4 @@ +import { RunnableLambda } from '@langchain/core/runnables'; import { FakeListChatModel } from '@langchain/core/utils/testing'; import { END, MessagesAnnotation, START, StateGraph } from '@langchain/langgraph'; import * as Sentry from '@sentry/node'; @@ -62,6 +63,25 @@ async function run() { await graph.invoke(user('Invoke the graph')); + // stream() called without config inside a parent run must keep reporting to the parent's callbacks. + let graphRunsSeenByParent = 0; + const parentCallback = { + handleChainStart: (_chain, _inputs, _runId, _parentRunId, _tags, metadata) => { + if (metadata?.lc_agent_name === 'weather_graph') { + graphRunsSeenByParent++; + } + }, + }; + await RunnableLambda.from(async input => { + for await (const _chunk of await graph.stream(input)) { + // drain + } + }).invoke(user('Stream inside a parent run'), { callbacks: [parentCallback] }); + Sentry.startSpan( + { name: 'parent callbacks', attributes: { 'test.graph_runs_seen': graphRunsSeenByParent } }, + () => {}, + ); + // Lets the unread stream's run finish before the root span ends. await sleep(200); }); 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 befc8f05cfbd..c62337558f50 100644 --- a/dev-packages/node-integration-tests/suites/tracing/langgraph/test.ts +++ b/dev-packages/node-integration-tests/suites/tracing/langgraph/test.ts @@ -383,7 +383,10 @@ describe('LangGraph integration', () => { span => span.attributes[SENTRY_OP]?.value === 'gen_ai.chat' && span.parent_span_id === parent?.span_id, ); - expect(agentSpans).toHaveLength(9); + expect(agentSpans).toHaveLength(10); + expect(agentSpanFor('Stream inside a parent run')?.status).toBe('ok'); + const parentCallbacks = spans.find(span => span.name === 'parent callbacks'); + expect(parentCallbacks?.attributes['test.graph_runs_seen']?.value).toBeGreaterThan(0); const read = agentSpanFor('Read the whole stream'); expect(read?.status).toBe('ok'); @@ -414,6 +417,7 @@ describe('LangGraph integration', () => { expect(invoked?.attributes[GEN_AI_RESPONSE_STREAMING]).toBeUndefined(); }, }) + .unordered() .start() .completed(); }); @@ -440,6 +444,7 @@ describe('LangGraph integration', () => { expect(chatSpans[0]!.attributes[GEN_AI_AGENT_NAME]?.value).toBe('helpful_assistant'); }, }) + .unordered() .start() .completed(); }, diff --git a/packages/server-utils/src/ai/langgraph/agent-span.ts b/packages/server-utils/src/ai/langgraph/agent-span.ts index f2bce259eb8b..7a093798bac3 100644 --- a/packages/server-utils/src/ai/langgraph/agent-span.ts +++ b/packages/server-utils/src/ai/langgraph/agent-span.ts @@ -18,6 +18,7 @@ import { _INTERNAL_mergeLangChainCallbackHandler, getConversationIdFromMetadata, getConversationIdMetadataFromConfig, + getInheritedLangChainCallbacks, normalizeLangChainMessages, } from '../langchain/utils'; import { LANGGRAPH_ORIGIN } from './constants'; @@ -81,9 +82,9 @@ export function prepareAgentSpan( span.setAttribute(GEN_AI_REQUEST_MODEL, modelName); } - // Inject callback handler and agent name into invoke config if (sentryCallbackHandler) { - const invokeConfig = (args[1] ?? {}) as Record; + // A copy, so nothing is written back to a config object the caller reuses for other calls. + const invokeConfig: Record = { ...(args[1] as Record | undefined) }; args[1] = invokeConfig; const existingMetadata = (invokeConfig.metadata ?? {}) as Record; @@ -94,12 +95,15 @@ export function prepareAgentSpan( ...(typeof graphName === 'string' ? { lc_agent_name: graphName } : {}), }; - invokeConfig.callbacks = _INTERNAL_mergeLangChainCallbackHandler(invokeConfig.callbacks, sentryCallbackHandler); + // Explicit callbacks replace the parent run's, which LangChain would otherwise inherit. + invokeConfig.callbacks = _INTERNAL_mergeLangChainCallbackHandler( + invokeConfig.callbacks ?? getInheritedLangChainCallbacks(), + sentryCallbackHandler, + ); } const recordInputs = options.recordInputs; - // Extract available tools from the graph instance const tools = recordInputs ? extractToolsFromCompiledGraph(graphInstance) : null; if (tools) { span.setAttribute(GEN_AI_TOOL_DEFINITIONS, JSON.stringify(tools)); diff --git a/packages/server-utils/src/ai/langgraph/stream.ts b/packages/server-utils/src/ai/langgraph/stream.ts index 935120cd13e2..df38e4318bb0 100644 --- a/packages/server-utils/src/ai/langgraph/stream.ts +++ b/packages/server-utils/src/ai/langgraph/stream.ts @@ -1,7 +1,7 @@ import { getCurrentScope, SPAN_STATUS_ERROR, startSpanManual } from '@sentry/core'; import { GEN_AI_RESPONSE_STREAMING } from '@sentry/conventions/attributes'; import type { BaseChatModel, LangChainMessage } from '../langchain/types'; -import { _INTERNAL_addLangChainCallbackHandler } from '../langchain/utils'; +import { _INTERNAL_addLangChainCallbackHandler, getInheritedLangChainCallbacks } from '../langchain/utils'; import { getAgentSpanOptions, getInvokeMarker, LANGGRAPH_INVOKE_ACTIVE, prepareAgentSpan } from './agent-span'; import type { CompiledGraph, LangGraphOptions } from './types'; import { setResponseAttributes } from './utils'; @@ -57,6 +57,7 @@ export function instrumentCompiledGraphStream( // Recording must never change the outcome of the run. } span.end(); + inputMessages = null; }; // The root run starts before any of its children, so it is the first run this per-call listener sees. @@ -92,7 +93,12 @@ export function instrumentCompiledGraphStream( llm, sentryCallbackHandler, ); - config.callbacks = _INTERNAL_addLangChainCallbackHandler(config.callbacks, listener); + // `prepareAgentSpan` may have replaced the config with its own copy. + const callConfig = callArgs[1] as Record; + callConfig.callbacks = _INTERNAL_addLangChainCallbackHandler( + callConfig.callbacks ?? getInheritedLangChainCallbacks(), + listener, + ); stream = (await Reflect.apply(target, thisArg, callArgs)) as typeof stream; } catch (error) { finish({ error }); 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 e80dc42fcda2..7ca844b168f0 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 @@ -332,4 +332,28 @@ describe('LangGraph stream() agent span', () => { expect(agentSpans()).toHaveLength(1); }); + + it('leaves the caller config untouched in invoke() and stream() when the Sentry handler is injected', async () => { + const graph = createGraph(emit => { + emit('handleChainEnd', {}, 'root'); + return createStream(); + }); + const instrumented = Object.assign(graph, { + async invoke( + this: { stream: (input: unknown, config: unknown) => Promise }, + input: unknown, + config: unknown, + ) { + await this.stream(input, config); + return {}; + }, + }); + instrumentCompiledGraph(instrumented, { name: 'weather_graph' }, {}, undefined, { name: 'SentryCallbackHandler' }); + const config = { configurable: { thread_id: 't1' } }; + + await instrumented.invoke({}, config); + await instrumented.stream({}, config); + + expect(config).toEqual({ configurable: { thread_id: 't1' } }); + }); }); From 692321985f4596d214e22f0e459b06d1e8e9d4b0 Mon Sep 17 00:00:00 2001 From: RulaKhaled Date: Mon, 26 Oct 2026 21:10:46 +0100 Subject: [PATCH 4/4] ref(server-utils): Copy the LangGraph config in one place Co-Authored-By: Claude Opus 5.5 --- .../tracing/langgraph/scenario-stream.mjs | 6 +-- .../suites/tracing/langgraph/test.ts | 2 +- .../src/ai/langgraph/agent-span.ts | 25 ++++++------ .../server-utils/src/ai/langgraph/stream.ts | 23 ++++------- .../ai/lib/tracing/langgraph-stream.test.ts | 39 ++++++++++--------- 5 files changed, 46 insertions(+), 49 deletions(-) diff --git a/dev-packages/node-integration-tests/suites/tracing/langgraph/scenario-stream.mjs b/dev-packages/node-integration-tests/suites/tracing/langgraph/scenario-stream.mjs index 2f6f441301a6..747b7d4e11eb 100644 --- a/dev-packages/node-integration-tests/suites/tracing/langgraph/scenario-stream.mjs +++ b/dev-packages/node-integration-tests/suites/tracing/langgraph/scenario-stream.mjs @@ -64,11 +64,11 @@ async function run() { await graph.invoke(user('Invoke the graph')); // stream() called without config inside a parent run must keep reporting to the parent's callbacks. - let graphRunsSeenByParent = 0; + let parentSawGraphRun = false; const parentCallback = { handleChainStart: (_chain, _inputs, _runId, _parentRunId, _tags, metadata) => { if (metadata?.lc_agent_name === 'weather_graph') { - graphRunsSeenByParent++; + parentSawGraphRun = true; } }, }; @@ -78,7 +78,7 @@ async function run() { } }).invoke(user('Stream inside a parent run'), { callbacks: [parentCallback] }); Sentry.startSpan( - { name: 'parent callbacks', attributes: { 'test.graph_runs_seen': graphRunsSeenByParent } }, + { name: 'parent callbacks', attributes: { 'test.parent_saw_graph_run': parentSawGraphRun } }, () => {}, ); 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 c62337558f50..18910bfd7b29 100644 --- a/dev-packages/node-integration-tests/suites/tracing/langgraph/test.ts +++ b/dev-packages/node-integration-tests/suites/tracing/langgraph/test.ts @@ -386,7 +386,7 @@ describe('LangGraph integration', () => { expect(agentSpans).toHaveLength(10); expect(agentSpanFor('Stream inside a parent run')?.status).toBe('ok'); const parentCallbacks = spans.find(span => span.name === 'parent callbacks'); - expect(parentCallbacks?.attributes['test.graph_runs_seen']?.value).toBeGreaterThan(0); + expect(parentCallbacks?.attributes['test.parent_saw_graph_run']?.value).toBe(true); const read = agentSpanFor('Read the whole stream'); expect(read?.status).toBe('ok'); diff --git a/packages/server-utils/src/ai/langgraph/agent-span.ts b/packages/server-utils/src/ai/langgraph/agent-span.ts index 7a093798bac3..eee5cddae3e4 100644 --- a/packages/server-utils/src/ai/langgraph/agent-span.ts +++ b/packages/server-utils/src/ai/langgraph/agent-span.ts @@ -82,24 +82,25 @@ export function prepareAgentSpan( span.setAttribute(GEN_AI_REQUEST_MODEL, modelName); } - if (sentryCallbackHandler) { - // A copy, so nothing is written back to a config object the caller reuses for other calls. - const invokeConfig: Record = { ...(args[1] as Record | undefined) }; - args[1] = invokeConfig; + // A copy, so nothing is written back to a config object the caller reuses for other calls. + const config: Record = { ...(args[1] as Record | undefined) }; + args[1] = config; + // Explicit callbacks replace the parent run's, which LangChain would otherwise inherit. + const callbacks = config.callbacks ?? getInheritedLangChainCallbacks(); + if (callbacks) { + config.callbacks = callbacks; + } - const existingMetadata = (invokeConfig.metadata ?? {}) as Record; - invokeConfig.metadata = { - ...getConversationIdMetadataFromConfig(invokeConfig), + if (sentryCallbackHandler) { + const existingMetadata = (config.metadata ?? {}) as Record; + config.metadata = { + ...getConversationIdMetadataFromConfig(config), ...existingMetadata, __sentry_langgraph__: true, ...(typeof graphName === 'string' ? { lc_agent_name: graphName } : {}), }; - // Explicit callbacks replace the parent run's, which LangChain would otherwise inherit. - invokeConfig.callbacks = _INTERNAL_mergeLangChainCallbackHandler( - invokeConfig.callbacks ?? getInheritedLangChainCallbacks(), - sentryCallbackHandler, - ); + config.callbacks = _INTERNAL_mergeLangChainCallbackHandler(config.callbacks, sentryCallbackHandler); } const recordInputs = options.recordInputs; diff --git a/packages/server-utils/src/ai/langgraph/stream.ts b/packages/server-utils/src/ai/langgraph/stream.ts index df38e4318bb0..843d0bf573fa 100644 --- a/packages/server-utils/src/ai/langgraph/stream.ts +++ b/packages/server-utils/src/ai/langgraph/stream.ts @@ -1,7 +1,7 @@ import { getCurrentScope, SPAN_STATUS_ERROR, startSpanManual } from '@sentry/core'; import { GEN_AI_RESPONSE_STREAMING } from '@sentry/conventions/attributes'; import type { BaseChatModel, LangChainMessage } from '../langchain/types'; -import { _INTERNAL_addLangChainCallbackHandler, getInheritedLangChainCallbacks } from '../langchain/utils'; +import { _INTERNAL_addLangChainCallbackHandler } from '../langchain/utils'; import { getAgentSpanOptions, getInvokeMarker, LANGGRAPH_INVOKE_ACTIVE, prepareAgentSpan } from './agent-span'; import type { CompiledGraph, LangGraphOptions } from './types'; import { setResponseAttributes } from './utils'; @@ -30,11 +30,7 @@ export function instrumentCompiledGraphStream( return Reflect.apply(target, thisArg, args); } - // A copy, so the per-call listener never lands on a config object the caller reuses. - const config: Record = { ...(args[1] as Record | undefined) }; - const callArgs = [args[0], config, ...args.slice(2)]; - - return startSpanManual(getAgentSpanOptions(callArgs, { [GEN_AI_RESPONSE_STREAMING]: true }), async span => { + return startSpanManual(getAgentSpanOptions(args, { [GEN_AI_RESPONSE_STREAMING]: true }), async span => { let inputMessages: LangChainMessage[] | null = null; let stream: { signal?: AbortSignal } | undefined; let ended = false; @@ -47,7 +43,7 @@ export function instrumentCompiledGraphStream( if (outcome && 'error' in outcome) { // A cancel or abort by the caller is not a failed run. LangGraph reports both, and step // timeouts too, as `Error('Abort')`, so only the signals tell them apart. - if (!stream?.signal?.aborted && !(config.signal as AbortSignal | undefined)?.aborted) { + if (!stream?.signal?.aborted && !(args[1] as { signal?: AbortSignal } | undefined)?.signal?.aborted) { span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); } } else if (outcome && options.recordOutputs) { @@ -86,20 +82,17 @@ export function instrumentCompiledGraphStream( try { inputMessages = prepareAgentSpan( span, - callArgs, + args, graphInstance, compileOptions, options, llm, sentryCallbackHandler, ); - // `prepareAgentSpan` may have replaced the config with its own copy. - const callConfig = callArgs[1] as Record; - callConfig.callbacks = _INTERNAL_addLangChainCallbackHandler( - callConfig.callbacks ?? getInheritedLangChainCallbacks(), - listener, - ); - stream = (await Reflect.apply(target, thisArg, callArgs)) as typeof stream; + // `prepareAgentSpan` gave the call its own copy of the config, so the listener stays on this call. + const config = args[1] as Record; + config.callbacks = _INTERNAL_addLangChainCallbackHandler(config.callbacks, listener); + stream = (await Reflect.apply(target, thisArg, args)) as typeof stream; } catch (error) { finish({ error }); throw error; 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 7ca844b168f0..ed2630971206 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 @@ -65,8 +65,8 @@ describe('LangGraph stream() agent span', () => { getMainCarrier().__SENTRY__ = undefined; }); - function instrument(graph: T, options = {}): T { - instrumentCompiledGraph(graph, { name: 'weather_graph' }, options, undefined, undefined); + function instrument(graph: T, options = {}, sentryCallbackHandler?: unknown): T { + instrumentCompiledGraph(graph, { name: 'weather_graph' }, options, undefined, sentryCallbackHandler); return graph; } @@ -334,25 +334,28 @@ describe('LangGraph stream() agent span', () => { }); it('leaves the caller config untouched in invoke() and stream() when the Sentry handler is injected', async () => { - const graph = createGraph(emit => { - emit('handleChainEnd', {}, 'root'); - return createStream(); - }); - const instrumented = Object.assign(graph, { - async invoke( - this: { stream: (input: unknown, config: unknown) => Promise }, - input: unknown, - config: unknown, - ) { - await this.stream(input, config); - return {}; + const graph = instrument( + { + ...createGraph(emit => { + emit('handleChainEnd', {}, 'root'); + return createStream(); + }), + async invoke( + this: { stream: (input: unknown, config: unknown) => Promise }, + input: unknown, + config: unknown, + ) { + await this.stream(input, config); + return {}; + }, }, - }); - instrumentCompiledGraph(instrumented, { name: 'weather_graph' }, {}, undefined, { name: 'SentryCallbackHandler' }); + {}, + { name: 'SentryCallbackHandler' }, + ); const config = { configurable: { thread_id: 't1' } }; - await instrumented.invoke({}, config); - await instrumented.stream({}, config); + await graph.invoke({}, config); + await graph.stream({}, config); expect(config).toEqual({ configurable: { thread_id: 't1' } }); });