Skip to content

Commit 0050c91

Browse files
christian-bromannchristian-bromann
authored
fix(langchain): reset shared currentSystemMessage on middleware handler retry (#10114)
1 parent 66df7fa commit 0050c91

File tree

3 files changed

+347
-0
lines changed

3 files changed

+347
-0
lines changed

‎.changeset/happy-coats-fry.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"langchain": patch
3+
---
4+
5+
fix(langchain): reset shared currentSystemMessage on middleware handler retry

‎libs/langchain/src/agents/nodes/AgentNode.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -451,6 +451,8 @@ export class AgentNode<
451451
unknown
452452
>
453453
): Promise<InternalModelResponse<StructuredResponseFormat>> => {
454+
const baselineSystemMessage = currentSystemMessage;
455+
454456
/**
455457
* Merge context with default context of middleware
456458
*/
@@ -501,6 +503,8 @@ export class AgentNode<
501503
unknown
502504
>
503505
): Promise<InternalModelResponse<StructuredResponseFormat>> => {
506+
currentSystemMessage = baselineSystemMessage;
507+
504508
/**
505509
* Verify that the user didn't add any new tools.
506510
* We can't allow this as the ToolNode is already initiated with given tools.

‎libs/langchain/src/agents/nodes/tests/AgentNode.test.ts‎

Lines changed: 338 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,3 +100,341 @@ describe("AgentNode concurrency", () => {
100100
expect(systemTexts).toEqual(["msg:X", "msg:Y"]);
101101
});
102102
});
103+
104+
describe("AgentNode system message reset on handler retry", () => {
105+
it("should not throw when outer middleware retries after inner middleware modified systemMessage", async () => {
106+
let handlerCallCount = 0;
107+
const model = new FakeToolCallingChatModel({
108+
responses: [new AIMessage("ok")],
109+
sleep: 0,
110+
});
111+
112+
const innerMiddleware = createMiddleware({
113+
name: "InnerMiddleware",
114+
wrapModelCall: async (request, handler) => {
115+
return handler({
116+
...request,
117+
systemMessage: request.systemMessage.concat(
118+
"\nExtra instructions from inner middleware"
119+
),
120+
});
121+
},
122+
});
123+
124+
const outerMiddleware = createMiddleware({
125+
name: "OuterMiddleware",
126+
wrapModelCall: async (request, handler) => {
127+
try {
128+
handlerCallCount += 1;
129+
return await handler(request);
130+
} catch {
131+
handlerCallCount += 1;
132+
return handler(request);
133+
}
134+
},
135+
});
136+
137+
const node = new AgentNode({
138+
model,
139+
systemMessage: new SystemMessage("base prompt"),
140+
toolClasses: [],
141+
shouldReturnDirect: new Set(),
142+
middleware: [outerMiddleware, innerMiddleware],
143+
wrapModelCallHookMiddleware: [
144+
[outerMiddleware, () => ({})],
145+
[innerMiddleware, () => ({})],
146+
],
147+
});
148+
149+
let generateCalls = 0;
150+
const originalGenerate = model._generate.bind(model);
151+
vi.spyOn(model, "_generate").mockImplementation(async (...args) => {
152+
generateCalls += 1;
153+
if (generateCalls === 1) {
154+
throw new Error("simulated context overflow");
155+
}
156+
return originalGenerate(...args);
157+
});
158+
159+
await expect(
160+
node.invoke(
161+
{ messages: [new HumanMessage("hello")], structuredResponse: {} },
162+
{ configurable: {} }
163+
)
164+
).resolves.toBeDefined();
165+
166+
expect(handlerCallCount).toBe(2);
167+
});
168+
169+
it("should not throw when outer middleware retries after inner middleware modified systemPrompt", async () => {
170+
let handlerCallCount = 0;
171+
const model = new FakeToolCallingChatModel({
172+
responses: [new AIMessage("ok")],
173+
sleep: 0,
174+
});
175+
176+
const innerMiddleware = createMiddleware({
177+
name: "InnerMiddleware",
178+
wrapModelCall: async (request, handler) => {
179+
return handler({
180+
...request,
181+
systemPrompt: `${request.systemPrompt}\nExtra prompt from inner`,
182+
});
183+
},
184+
});
185+
186+
const outerMiddleware = createMiddleware({
187+
name: "OuterMiddleware",
188+
wrapModelCall: async (request, handler) => {
189+
try {
190+
handlerCallCount += 1;
191+
return await handler(request);
192+
} catch {
193+
handlerCallCount += 1;
194+
return handler(request);
195+
}
196+
},
197+
});
198+
199+
const node = new AgentNode({
200+
model,
201+
systemMessage: new SystemMessage("base prompt"),
202+
toolClasses: [],
203+
shouldReturnDirect: new Set(),
204+
middleware: [outerMiddleware, innerMiddleware],
205+
wrapModelCallHookMiddleware: [
206+
[outerMiddleware, () => ({})],
207+
[innerMiddleware, () => ({})],
208+
],
209+
});
210+
211+
let generateCalls = 0;
212+
const originalGenerate = model._generate.bind(model);
213+
vi.spyOn(model, "_generate").mockImplementation(async (...args) => {
214+
generateCalls += 1;
215+
if (generateCalls === 1) {
216+
throw new Error("simulated error");
217+
}
218+
return originalGenerate(...args);
219+
});
220+
221+
await expect(
222+
node.invoke(
223+
{ messages: [new HumanMessage("hello")], structuredResponse: {} },
224+
{ configurable: {} }
225+
)
226+
).resolves.toBeDefined();
227+
228+
expect(handlerCallCount).toBe(2);
229+
});
230+
231+
it("should preserve inner middleware system message changes on each retry", async () => {
232+
const model = new FakeToolCallingChatModel({
233+
responses: [new AIMessage("ok")],
234+
sleep: 0,
235+
});
236+
const spy = vi.spyOn(model, "invoke");
237+
238+
const innerMiddleware = createMiddleware({
239+
name: "InnerMiddleware",
240+
wrapModelCall: async (request, handler) => {
241+
return handler({
242+
...request,
243+
systemMessage: request.systemMessage.concat("\n[inner-addition]"),
244+
});
245+
},
246+
});
247+
248+
let attempt = 0;
249+
const outerMiddleware = createMiddleware({
250+
name: "OuterMiddleware",
251+
wrapModelCall: async (request, handler) => {
252+
attempt += 1;
253+
if (attempt === 1) {
254+
try {
255+
return await handler(request);
256+
} catch {
257+
// fall through to retry
258+
}
259+
}
260+
return handler(request);
261+
},
262+
});
263+
264+
const node = new AgentNode({
265+
model,
266+
systemMessage: new SystemMessage("base"),
267+
toolClasses: [],
268+
shouldReturnDirect: new Set(),
269+
middleware: [outerMiddleware, innerMiddleware],
270+
wrapModelCallHookMiddleware: [
271+
[outerMiddleware, () => ({})],
272+
[innerMiddleware, () => ({})],
273+
],
274+
});
275+
276+
let generateCalls = 0;
277+
const originalGenerate = model._generate.bind(model);
278+
vi.spyOn(model, "_generate").mockImplementation(async (...args) => {
279+
generateCalls += 1;
280+
if (generateCalls === 1) {
281+
throw new Error("first attempt fails");
282+
}
283+
return originalGenerate(...args);
284+
});
285+
286+
await node.invoke(
287+
{ messages: [new HumanMessage("test")], structuredResponse: {} },
288+
{ configurable: {} }
289+
);
290+
291+
const lastCallMessages = spy.mock.calls.at(-1)?.[0] as BaseMessage[];
292+
const systemText = lastCallMessages[0].text;
293+
expect(systemText).toContain("base");
294+
expect(systemText).toContain("[inner-addition]");
295+
});
296+
297+
it("should handle three middleware layers with retry in outermost", async () => {
298+
const model = new FakeToolCallingChatModel({
299+
responses: [new AIMessage("ok")],
300+
sleep: 0,
301+
});
302+
const spy = vi.spyOn(model, "invoke");
303+
304+
const middlewareA = createMiddleware({
305+
name: "MiddlewareA",
306+
wrapModelCall: async (request, handler) => {
307+
return handler({
308+
...request,
309+
systemMessage: request.systemMessage.concat("\n[A]"),
310+
});
311+
},
312+
});
313+
314+
const middlewareB = createMiddleware({
315+
name: "MiddlewareB",
316+
wrapModelCall: async (request, handler) => {
317+
return handler({
318+
...request,
319+
systemMessage: request.systemMessage.concat("\n[B]"),
320+
});
321+
},
322+
});
323+
324+
let attempt = 0;
325+
const retryMiddleware = createMiddleware({
326+
name: "RetryMiddleware",
327+
wrapModelCall: async (request, handler) => {
328+
attempt += 1;
329+
if (attempt === 1) {
330+
try {
331+
return await handler(request);
332+
} catch {
333+
// retry
334+
}
335+
}
336+
return handler(request);
337+
},
338+
});
339+
340+
const node = new AgentNode({
341+
model,
342+
systemMessage: new SystemMessage("root"),
343+
toolClasses: [],
344+
shouldReturnDirect: new Set(),
345+
middleware: [retryMiddleware, middlewareA, middlewareB],
346+
wrapModelCallHookMiddleware: [
347+
[retryMiddleware, () => ({})],
348+
[middlewareA, () => ({})],
349+
[middlewareB, () => ({})],
350+
],
351+
});
352+
353+
let generateCalls = 0;
354+
const originalGenerate = model._generate.bind(model);
355+
vi.spyOn(model, "_generate").mockImplementation(async (...args) => {
356+
generateCalls += 1;
357+
if (generateCalls === 1) {
358+
throw new Error("fail first");
359+
}
360+
return originalGenerate(...args);
361+
});
362+
363+
await node.invoke(
364+
{ messages: [new HumanMessage("go")], structuredResponse: {} },
365+
{ configurable: {} }
366+
);
367+
368+
const lastCallMessages = spy.mock.calls.at(-1)?.[0] as BaseMessage[];
369+
const systemText = lastCallMessages[0].text;
370+
expect(systemText).toContain("root");
371+
expect(systemText).toContain("[A]");
372+
expect(systemText).toContain("[B]");
373+
});
374+
375+
it("should allow middleware to call handler multiple times for fallback logic", async () => {
376+
const model = new FakeToolCallingChatModel({
377+
responses: [new AIMessage("ok")],
378+
sleep: 0,
379+
});
380+
const spy = vi.spyOn(model, "invoke");
381+
382+
const innerMiddleware = createMiddleware({
383+
name: "InnerMiddleware",
384+
wrapModelCall: async (request, handler) => {
385+
return handler({
386+
...request,
387+
systemPrompt: `${request.systemPrompt}\n[inner]`,
388+
});
389+
},
390+
});
391+
392+
const fallbackMiddleware = createMiddleware({
393+
name: "FallbackMiddleware",
394+
wrapModelCall: async (request, handler) => {
395+
try {
396+
return await handler(request);
397+
} catch {
398+
return handler({
399+
...request,
400+
messages: [new HumanMessage("summarized context")],
401+
});
402+
}
403+
},
404+
});
405+
406+
const node = new AgentNode({
407+
model,
408+
systemMessage: new SystemMessage("original"),
409+
toolClasses: [],
410+
shouldReturnDirect: new Set(),
411+
middleware: [fallbackMiddleware, innerMiddleware],
412+
wrapModelCallHookMiddleware: [
413+
[fallbackMiddleware, () => ({})],
414+
[innerMiddleware, () => ({})],
415+
],
416+
});
417+
418+
let generateCalls = 0;
419+
const originalGenerate = model._generate.bind(model);
420+
vi.spyOn(model, "_generate").mockImplementation(async (...args) => {
421+
generateCalls += 1;
422+
if (generateCalls === 1) {
423+
throw new Error("context too large");
424+
}
425+
return originalGenerate(...args);
426+
});
427+
428+
await node.invoke(
429+
{ messages: [new HumanMessage("long message")], structuredResponse: {} },
430+
{ configurable: {} }
431+
);
432+
433+
expect(spy).toHaveBeenCalledTimes(2);
434+
const retryMessages = spy.mock.calls[1][0] as BaseMessage[];
435+
expect(retryMessages[0].text).toContain("original");
436+
expect(retryMessages[0].text).toContain("[inner]");
437+
expect(retryMessages[1]).toBeInstanceOf(HumanMessage);
438+
expect(retryMessages[1].text).toBe("summarized context");
439+
});
440+
});

0 commit comments

Comments
 (0)
⚡