Skip to content

Commit d34ea30

Browse files
logaretmlogaretm
committed
ref(server-utils): Tighten LangGraph state recording and its tests
1 parent c0df556 commit d34ea30

3 files changed

Lines changed: 45 additions & 94 deletions

File tree

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

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -154,13 +154,9 @@ export function instrumentCompiledGraphInvoke(
154154
span.setAttribute(GEN_AI_TOOL_DEFINITIONS, JSON.stringify(tools));
155155
}
156156

157-
// Parse input state. MessagesAnnotation graphs expose a `messages` array (possibly empty);
158-
// a custom state annotation exposes arbitrary keys instead. Route on whether `messages` is
159-
// an array, mirroring the output side in setResponseAttributes, so an empty chat history is
160-
// recorded as an empty chat array rather than misread as custom state and wrapped.
161-
const inputState = args.length > 0 ? args[0] : undefined;
162-
const stateMessages = (inputState as { messages?: LangChainMessage[] } | null)?.messages;
163-
const inputMessages = Array.isArray(stateMessages) ? stateMessages : null;
157+
// Custom state annotations have no `messages` array, the whole state is recorded instead.
158+
const inputState = args[0] as { messages?: LangChainMessage[] } | null | undefined;
159+
const inputMessages = Array.isArray(inputState?.messages) ? inputState.messages : null;
164160

165161
if (recordInputs) {
166162
if (inputMessages) {
@@ -175,9 +171,7 @@ export function instrumentCompiledGraphInvoke(
175171
[GEN_AI_INPUT_MESSAGES]: stringify(filteredMessages),
176172
});
177173
} else if (inputState && typeof inputState === 'object') {
178-
span.setAttributes({
179-
[GEN_AI_INPUT_MESSAGES]: stringify([{ role: 'user', content: stringify(inputState) }]),
180-
});
174+
span.setAttribute(GEN_AI_INPUT_MESSAGES, stringify([{ role: 'user', content: stringify(inputState) }]));
181175
}
182176
}
183177

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -276,8 +276,7 @@ export function setResponseAttributes(span: Span, inputMessages: LangChainMessag
276276
const outputMessages = resultObj?.messages;
277277

278278
if (!outputMessages || !Array.isArray(outputMessages)) {
279-
// A custom state annotation has no `messages` array. Record the whole output state as a
280-
// fallback so it is not dropped silently.
279+
// Custom state annotations have no `messages` array, the whole state is recorded instead.
281280
if (result && typeof result === 'object') {
282281
span.setAttribute(GEN_AI_RESPONSE_TEXT, stringify([{ role: 'assistant', content: stringify(result) }]));
283282
}

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

Lines changed: 40 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -61,106 +61,64 @@ describe('invoke_agent input/output recording', () => {
6161
return endedSpans;
6262
}
6363

64-
// Regression test for a graph built on a custom state annotation (no `messages` key). The
65-
// instrumentation used to read `args[0].messages` only, so the whole state was dropped.
66-
it('records the full state for a graph that does not use MessagesAnnotation', async () => {
64+
async function getInvokeAttributes<T>(invoke: (input: T) => Promise<unknown>, input: T) {
6765
const endedSpans = setupClient();
68-
69-
const compiled = {
70-
invoke: async (input: Record<string, unknown>) => ({ ...input, expanded: 'expanded idea', validated: true }),
71-
};
72-
const stateGraph = { compile: () => compiled };
66+
const stateGraph = { compile: () => ({ invoke }) };
7367

7468
instrumentStateGraph(stateGraph, { recordInputs: true, recordOutputs: true });
75-
const graph = stateGraph.compile();
76-
const result = await graph.invoke({ idea: 'test idea' });
69+
await stateGraph.compile().invoke(input);
7770

78-
expect(result).toEqual({ idea: 'test idea', expanded: 'expanded idea', validated: true });
7971
expect(endedSpans).toHaveLength(1);
72+
return spanToJSON(endedSpans[0]!).attributes;
73+
}
8074

81-
const data = spanToJSON(endedSpans[0]!).attributes;
82-
83-
const inputMessages = data[GEN_AI_INPUT_MESSAGES] as string | undefined;
84-
expect(inputMessages).toBeDefined();
85-
const parsedInput = JSON.parse(inputMessages!) as Array<{ role: string; content: string }>;
86-
expect(parsedInput).toHaveLength(1);
87-
expect(parsedInput[0]!.role).toBe('user');
88-
expect(JSON.parse(parsedInput[0]!.content)).toEqual({ idea: 'test idea' });
89-
90-
const responseText = data[GEN_AI_RESPONSE_TEXT] as string | undefined;
91-
expect(responseText).toBeDefined();
92-
const parsedOutput = JSON.parse(responseText!) as Array<{ role: string; content: string }>;
93-
expect(parsedOutput[0]!.role).toBe('assistant');
94-
expect(JSON.parse(parsedOutput[0]!.content)).toEqual({
95-
idea: 'test idea',
96-
expanded: 'expanded idea',
97-
validated: true,
98-
});
75+
it('records the full state for a graph that does not use MessagesAnnotation', async () => {
76+
const attributes = await getInvokeAttributes(
77+
async (input: Record<string, unknown>) => ({ ...input, expanded: 'expanded idea', validated: true }),
78+
{ idea: 'test idea' },
79+
);
80+
81+
expect(JSON.parse(attributes[GEN_AI_INPUT_MESSAGES] as string)).toEqual([
82+
{ role: 'user', content: JSON.stringify({ idea: 'test idea' }) },
83+
]);
84+
expect(JSON.parse(attributes[GEN_AI_RESPONSE_TEXT] as string)).toEqual([
85+
{
86+
role: 'assistant',
87+
content: JSON.stringify({ idea: 'test idea', expanded: 'expanded idea', validated: true }),
88+
},
89+
]);
9990
});
10091

10192
it('still records chat messages for a MessagesAnnotation graph', async () => {
102-
const endedSpans = setupClient();
103-
104-
const compiled = {
105-
invoke: async (input: { messages: Array<{ role: string; content: string }> }) => ({
93+
const attributes = await getInvokeAttributes(
94+
async (input: { messages: Array<{ role: string; content: string }> }) => ({
10695
messages: [...input.messages, { role: 'assistant', content: 'The weather is sunny' }],
10796
}),
108-
};
109-
const stateGraph = { compile: () => compiled };
110-
111-
instrumentStateGraph(stateGraph, { recordInputs: true, recordOutputs: true });
112-
const graph = stateGraph.compile();
113-
await graph.invoke({ messages: [{ role: 'user', content: 'What is the weather today?' }] });
114-
115-
const data = spanToJSON(endedSpans[0]!).attributes;
116-
117-
const inputMessages = data[GEN_AI_INPUT_MESSAGES] as string | undefined;
118-
expect(inputMessages).toBeDefined();
119-
expect(JSON.parse(inputMessages!)).toEqual([{ role: 'user', content: 'What is the weather today?' }]);
97+
{ messages: [{ role: 'user', content: 'What is the weather today?' }] },
98+
);
12099

121-
const responseText = data[GEN_AI_RESPONSE_TEXT] as string | undefined;
122-
expect(responseText).toBeDefined();
123-
expect(responseText).toContain('The weather is sunny');
100+
expect(JSON.parse(attributes[GEN_AI_INPUT_MESSAGES] as string)).toEqual([
101+
{ role: 'user', content: 'What is the weather today?' },
102+
]);
103+
expect(attributes[GEN_AI_RESPONSE_TEXT]).toContain('The weather is sunny');
124104
});
125105

126-
it('records an empty messages array on the chat path rather than wrapping it as custom state', async () => {
127-
const endedSpans = setupClient();
128-
129-
const compiled = {
130-
invoke: async (_input: { messages: Array<{ role: string; content: string }> }) => ({
131-
messages: [{ role: 'assistant', content: 'Hello' }],
132-
}),
133-
};
134-
const stateGraph = { compile: () => compiled };
135-
136-
instrumentStateGraph(stateGraph, { recordInputs: true, recordOutputs: true });
137-
const graph = stateGraph.compile();
138-
await graph.invoke({ messages: [] });
139-
140-
const data = spanToJSON(endedSpans[0]!).attributes;
141-
142-
const inputMessages = data[GEN_AI_INPUT_MESSAGES] as string | undefined;
143-
expect(inputMessages).toBeDefined();
144-
expect(JSON.parse(inputMessages!)).toEqual([]);
106+
it('records an empty messages array as an empty chat array', async () => {
107+
const attributes = await getInvokeAttributes(
108+
async (_input: { messages: unknown[] }) => ({ messages: [{ role: 'assistant', content: 'Hello' }] }),
109+
{ messages: [] },
110+
);
145111

146-
const responseText = data[GEN_AI_RESPONSE_TEXT] as string | undefined;
147-
expect(responseText).toBeDefined();
148-
expect(responseText).toContain('Hello');
112+
expect(attributes[GEN_AI_INPUT_MESSAGES]).toBe('[]');
113+
expect(attributes[GEN_AI_RESPONSE_TEXT]).toContain('Hello');
149114
});
150115

151116
it('does not record input messages when invoked with null input', async () => {
152-
const endedSpans = setupClient();
153-
154-
const compiled = {
155-
invoke: async (_input?: unknown) => ({ messages: [{ role: 'assistant', content: 'resumed' }] }),
156-
};
157-
const stateGraph = { compile: () => compiled };
158-
159-
instrumentStateGraph(stateGraph, { recordInputs: true, recordOutputs: true });
160-
const graph = stateGraph.compile();
161-
await expect(graph.invoke(null)).resolves.toBeDefined();
117+
const attributes = await getInvokeAttributes(
118+
async (_input: null) => ({ messages: [{ role: 'assistant', content: 'resumed' }] }),
119+
null,
120+
);
162121

163-
const data = spanToJSON(endedSpans[0]!).attributes;
164-
expect(data[GEN_AI_INPUT_MESSAGES]).toBeUndefined();
122+
expect(attributes[GEN_AI_INPUT_MESSAGES]).toBeUndefined();
165123
});
166124
});

0 commit comments

Comments
 (0)