fix(server-utils): Record LangGraph span I/O for non-MessagesAnnotation state

dev-packages/node-integration-tests/suites/tracing/langgraph/scenario-custom-state.mjs

@@ -0,0 +1,30 @@
+import { Annotation, END, START, StateGraph } from '@langchain/langgraph';
+import * as Sentry from '@sentry/node';
+
+// A graph built on a custom `Annotation.Root` state (no `messages` channel). The invoke_agent span
+// records the whole state object on both the input and the output side, unlike the MessagesAnnotation
+// graphs the other scenarios use.
+const CustomState = Annotation.Root({
+  topic: Annotation(),
+  summary: Annotation(),
+});
+
+async function run() {
+  await Sentry.startSpan({ op: 'function', name: 'main' }, async () => {
+    const summarize = state => {
+      return { summary: `Summary of ${state.topic}` };
+    };
+
+    const graph = new StateGraph(CustomState)
+      .addNode('summarize', summarize)
+      .addEdge(START, 'summarize')
+      .addEdge('summarize', END)
+      .compile({ name: 'custom_state_agent' });
+
+    await graph.invoke({ topic: 'weather' });
+  });
+
+  await Sentry.flush(2000);
+}
+
+run();

dev-packages/node-integration-tests/suites/tracing/langgraph/test.ts

@@ -4,6 +4,7 @@ import {
   GEN_AI_CONVERSATION_ID,
   GEN_AI_INPUT_MESSAGES,
   GEN_AI_OPERATION_NAME,
+  GEN_AI_OUTPUT_MESSAGES,
   GEN_AI_PIPELINE_NAME,
   GEN_AI_RESPONSE_MODEL,
   GEN_AI_RESPONSE_TEXT,
@@ -226,6 +227,33 @@ describe('LangGraph integration', () => {
     });
   });
 
+  // Custom `Annotation.Root` state has no `messages` channel, so the whole state object is recorded on
+  // both the input and the output side of the invoke_agent span.
+  createEsmAndCjsTests(__dirname, 'scenario-custom-state.mjs', 'instrument-with-pii.mjs', (createRunner, test) => {
+    test('records custom Annotation.Root state as input and output on the invoke_agent span', async () => {
+      await createRunner()
+        .expect({
+          span: container => {
+            const invokeAgentSpan = container.items.find(span => span.name === 'invoke_agent custom_state_agent');
+            expect(invokeAgentSpan).toBeDefined();
+            expect(invokeAgentSpan!.status).toBe('ok');
+            expect(invokeAgentSpan!.attributes[SENTRY_OP].value).toBe('gen_ai.invoke_agent');
+            expect(invokeAgentSpan!.attributes[SENTRY_ORIGIN].value).toBe('auto.ai.langgraph');
+
+            const inputMessages = getStringAttributeValue(invokeAgentSpan!.attributes[GEN_AI_INPUT_MESSAGES]?.value);
+            expect(inputMessages).toContain('"role":"user"');
+            expect(inputMessages).toContain('weather');
+
+            const outputMessages = getStringAttributeValue(invokeAgentSpan!.attributes[GEN_AI_OUTPUT_MESSAGES]?.value);
+            expect(outputMessages).toContain('"role":"assistant"');
+            expect(outputMessages).toContain('Summary of weather');
+          },
+        })
+        .start()
+        .completed();
+    });
+  });
+
   // createReactAgent tests.
   // Spans are asserted order-independently: the span-array order is not a protocol guarantee (Sentry
   // rebuilds the tree from `parent_span_id`), and the provider emits tree order while the OTel exporter

packages/server-utils/src/ai/langgraph/index.ts

@@ -153,28 +153,36 @@ export function instrumentCompiledGraphInvoke(
               span.setAttribute(GEN_AI_TOOL_DEFINITIONS, JSON.stringify(tools));
             }
 
-            // Parse input messages
-            const inputMessages =
-              args.length > 0 ? ((args[0] as { messages?: LangChainMessage[] } | null)?.messages ?? []) : [];
-
-            if (inputMessages && recordInputs) {
-              const normalizedMessages = normalizeLangChainMessages(inputMessages);
-              const { systemInstructions, filteredMessages } = extractSystemInstructions(normalizedMessages);
-
-              if (systemInstructions) {
-                span.setAttribute(GEN_AI_SYSTEM_INSTRUCTIONS, systemInstructions);
+            // 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) }]));
               }
-
-              span.setAttributes({
-                [GEN_AI_INPUT_MESSAGES]: stringify(filteredMessages),
-              });
             }
 
             // Call original invoke
             const result = await Reflect.apply(target, thisArg, args);
 
             if (recordOutputs) {
-              setResponseAttributes(span, inputMessages ?? null, result);
+              setResponseAttributes(span, inputMessages, result);
             }
 
             return result;

packages/server-utils/src/ai/langgraph/utils.ts

@@ -1,5 +1,5 @@
 /* eslint-disable typescript-eslint/no-deprecated */
-import { SPAN_STATUS_ERROR, startSpan } from '@sentry/core';
+import { SPAN_STATUS_ERROR, startSpan, stringify } from '@sentry/core';
 import type { Span, SpanAttributes } from '@sentry/core';
 import {
   GEN_AI_AGENT_NAME,
@@ -22,6 +22,7 @@ import {
 } from '@sentry/conventions/attributes';
 import { GEN_AI_EXECUTE_TOOL } from '@sentry/conventions/op';
 import { GEN_AI_TOOL_CALL_ID_ATTRIBUTE } from '../core/gen-ai-attributes';
+import { setOutputMessagesAttribute } from '../core/utils';
 import type { BaseChatModel, LangChainMessage } from '../langchain/types';
 import {
   extractMessageTokenUsageAttributes,
@@ -256,6 +257,14 @@ export function setResponseAttributes(span: Span, inputMessages: LangChainMessag
   const outputMessages = resultObj?.messages;
 
   if (!outputMessages || !Array.isArray(outputMessages)) {
+    // Custom state annotations have no `messages` array, the whole state is recorded instead.
+    if (result && typeof result === 'object') {
+      const serializedState = stringify(result);
+      // `gen_ai.output.messages` is what the product reads first; `gen_ai.response.text` is kept for
+      // back-compat (Relay still migrates it).
+      setOutputMessagesAttribute(span, { responseText: serializedState });
+      span.setAttribute(GEN_AI_RESPONSE_TEXT, stringify([{ role: 'assistant', content: serializedState }]));
+    }
     return;
   }
 

packages/server-utils/test/ai/lib/tracing/langgraph.test.ts

@@ -1,9 +1,13 @@
-import { describe, expect, it } from 'vitest';
+import { GEN_AI_INPUT_MESSAGES, GEN_AI_OUTPUT_MESSAGES, GEN_AI_RESPONSE_TEXT } from '@sentry/conventions/attributes';
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+import { getMainCarrier, setCurrentClient, spanToJSON } from '@sentry/core';
+import type { Span } from '@sentry/core';
 import {
   instrumentCreateReactAgent,
   instrumentStateGraph,
   instrumentStateGraphCompile,
 } from '../../../../src/ai/langgraph';
+import { getDefaultTestClientOptions, TestClient } from '../../../mocks/client';
 
 describe('langgraph double-patch guard', () => {
   it('instrumentStateGraphCompile returns the same wrapper when applied twice', () => {
@@ -32,3 +36,118 @@ describe('instrumentStateGraph', () => {
     expect(stateGraph.compile).not.toBe(originalCompile);
   });
 });
+
+describe('invoke_agent input/output recording', () => {
+  beforeEach(() => {
+    getMainCarrier().__SENTRY__ = undefined;
+  });
+
+  afterEach(() => {
+    getMainCarrier().__SENTRY__ = undefined;
+  });
+
+  function setupClient(): Span[] {
+    const client = new TestClient(
+      getDefaultTestClientOptions({
+        dsn: 'https://public@dsn.ingest.sentry.io/1337',
+        tracesSampleRate: 1,
+      }),
+    );
+    setCurrentClient(client);
+    client.init();
+
+    const endedSpans: Span[] = [];
+    client.on('spanEnd', span => endedSpans.push(span));
+    return endedSpans;
+  }
+
+  async function getInvokeAttributes<T>(invoke: (input: T) => Promise<unknown>, input: T) {
+    const endedSpans = setupClient();
+    const stateGraph = { compile: () => ({ invoke }) };
+
+    instrumentStateGraph(stateGraph, { recordInputs: true, recordOutputs: true });
+    await stateGraph.compile().invoke(input);
+
+    expect(endedSpans).toHaveLength(1);
+    return spanToJSON(endedSpans[0]!).attributes;
+  }
+
+  it('records the full state for a graph that does not use MessagesAnnotation', async () => {
+    const attributes = await getInvokeAttributes(
+      async (input: Record<string, unknown>) => ({ ...input, expanded: 'expanded idea', validated: true }),
+      { idea: 'test idea' },
+    );
+
+    expect(JSON.parse(attributes[GEN_AI_INPUT_MESSAGES] as string)).toEqual([
+      { role: 'user', content: JSON.stringify({ idea: 'test idea' }) },
+    ]);
+    expect(JSON.parse(attributes[GEN_AI_RESPONSE_TEXT] as string)).toEqual([
+      {
+        role: 'assistant',
+        content: JSON.stringify({ idea: 'test idea', expanded: 'expanded idea', validated: true }),
+      },
+    ]);
+  });
+
+  it('still records chat messages for a MessagesAnnotation graph', async () => {
+    const attributes = await getInvokeAttributes(
+      async (input: { messages: Array<{ role: string; content: string }> }) => ({
+        messages: [...input.messages, { role: 'assistant', content: 'The weather is sunny' }],
+      }),
+      { messages: [{ role: 'user', content: 'What is the weather today?' }] },
+    );
+
+    expect(JSON.parse(attributes[GEN_AI_INPUT_MESSAGES] as string)).toEqual([
+      { role: 'user', content: 'What is the weather today?' },
+    ]);
+    expect(attributes[GEN_AI_RESPONSE_TEXT]).toContain('The weather is sunny');
+  });
+
+  it('records an empty messages array as an empty chat array', async () => {
+    const attributes = await getInvokeAttributes(
+      async (_input: { messages: unknown[] }) => ({ messages: [{ role: 'assistant', content: 'Hello' }] }),
+      { messages: [] },
+    );
+
+    expect(attributes[GEN_AI_INPUT_MESSAGES]).toBe('[]');
+    expect(attributes[GEN_AI_RESPONSE_TEXT]).toContain('Hello');
+  });
+
+  it('does not record input messages when invoked with null input', async () => {
+    const attributes = await getInvokeAttributes(
+      async (_input: null) => ({ messages: [{ role: 'assistant', content: 'resumed' }] }),
+      null,
+    );
+
+    expect(attributes[GEN_AI_INPUT_MESSAGES]).toBeUndefined();
+  });
+
+  it('does not record a Command resume as input (skips it like a null resume)', async () => {
+    const attributes = await getInvokeAttributes(
+      async (_input: { lg_name: string; resume: string; goto: unknown[] }) => ({
+        messages: [{ role: 'assistant', content: 'resumed' }],
+      }),
+      { lg_name: 'Command', resume: 'approved', goto: [] },
+    );
+
+    expect(attributes[GEN_AI_INPUT_MESSAGES]).toBeUndefined();
+    expect(attributes[GEN_AI_RESPONSE_TEXT]).toContain('resumed');
+  });
+
+  it('records custom state output as gen_ai.output.messages alongside the deprecated response.text', async () => {
+    const attributes = await getInvokeAttributes(
+      async (input: Record<string, unknown>) => ({ ...input, summary: 'done' }),
+      { topic: 'weather' },
+    );
+
+    expect(JSON.parse(attributes[GEN_AI_OUTPUT_MESSAGES] as string)).toEqual([
+      {
+        role: 'assistant',
+        parts: [{ type: 'text', content: JSON.stringify({ topic: 'weather', summary: 'done' }) }],
+      },
+    ]);
+    expect(JSON.parse(attributes[GEN_AI_RESPONSE_TEXT] as string)).toEqual([
+      { role: 'assistant', content: JSON.stringify({ topic: 'weather', summary: 'done' }) },
+    ]);
+  });
+});