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

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

@@ -154,28 +154,32 @@ 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[] } | null | undefined;
+            const inputMessages = Array.isArray(inputState?.messages) ? inputState.messages : null;
+
+            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') {
+                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 { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, SPAN_STATUS_ERROR, startSpan } from '@sentry/core';
+import { SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN, SPAN_STATUS_ERROR, startSpan, stringify } from '@sentry/core';
 import type { Span, SpanAttributes } from '@sentry/core';
 import {
   GEN_AI_AGENT_NAME,
@@ -276,6 +276,10 @@ 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') {
+      span.setAttribute(GEN_AI_RESPONSE_TEXT, stringify([{ role: 'assistant', content: stringify(result) }]));
+    }
     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_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,89 @@ 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();
+  });
+});