Skip to content

Commit e6d895b

Browse files
christian-bromannchristian-bromann
authored
fix(middleware): avoid unnecessary REMOVE_ALL_MESSAGES in PatchToolCallsMiddleware (#169)
* fix(middleware): avoid unnecessary REMOVE_ALL_MESSAGES in PatchToolCallsMiddleware The middleware was unconditionally triggering REMOVE_ALL_MESSAGES on every request, even when there were no dangling tool calls to patch. This caused unnecessary frontend re-renders and UI flickering when using useStream with fetchStateHistory: true. Added a needsPatch flag to track whether any dangling tool calls are found, and only trigger the state rebuild when patching is actually needed. Fixes #127 * retrigger * fix test * add changeset
1 parent 0fe09a5 commit e6d895b

File tree

4 files changed

+145
-102
lines changed

4 files changed

+145
-102
lines changed

‎.changeset/easy-oranges-watch.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"deepagents": patch
3+
---
4+
5+
fix(middleware): avoid unnecessary REMOVE_ALL_MESSAGES in PatchToolCallsMiddleware

‎libs/deepagents/src/middleware/index.test.ts‎

Lines changed: 2 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -248,13 +248,7 @@ describe("PatchToolCallsMiddleware", () => {
248248
const stateUpdate = await beforeAgentHook({
249249
messages: inputMessages,
250250
});
251-
expect(stateUpdate).toBeDefined();
252-
expect(stateUpdate.messages).toHaveLength(3);
253-
expect(stateUpdate.messages[0]._getType()).toBe("remove");
254-
expect(stateUpdate.messages[1].content).toBe(
255-
"You are a helpful assistant.",
256-
);
257-
expect(stateUpdate.messages[2].content).toBe("Hello, how are you?");
251+
expect(stateUpdate).toBeUndefined();
258252
});
259253

260254
it("should patch a single missing tool call", async () => {
@@ -331,14 +325,7 @@ describe("PatchToolCallsMiddleware", () => {
331325
messages: inputMessages,
332326
});
333327

334-
expect(stateUpdate).toBeDefined();
335-
expect(stateUpdate.messages).toHaveLength(6);
336-
expect(stateUpdate.messages[0]._getType()).toBe("remove");
337-
expect(stateUpdate.messages.slice(1)).toEqual(inputMessages);
338-
339-
const updatedMessages = addMessages(inputMessages, stateUpdate.messages);
340-
expect(updatedMessages).toHaveLength(5);
341-
expect(updatedMessages).toEqual(inputMessages);
328+
expect(stateUpdate).toBeUndefined();
342329
});
343330

344331
it("should patch multiple missing tool calls", async () => {

‎libs/deepagents/src/middleware/patch_tool_calls.test.ts‎

Lines changed: 129 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -1,35 +1,95 @@
11
import { describe, it, expect } from "vitest";
2-
import { HumanMessage, AIMessage, ToolMessage } from "@langchain/core/messages";
32
import { createPatchToolCallsMiddleware } from "./patch_tool_calls.js";
3+
import { AIMessage, HumanMessage, ToolMessage } from "@langchain/core/messages";
4+
import { RemoveMessage } from "@langchain/core/messages";
5+
import { REMOVE_ALL_MESSAGES } from "@langchain/langgraph";
46

57
describe("createPatchToolCallsMiddleware", () => {
6-
describe("basic functionality", () => {
7-
it("should return undefined when no messages", async () => {
8+
describe("no patching needed (should return undefined)", () => {
9+
it("should return undefined when messages is empty", async () => {
810
const middleware = createPatchToolCallsMiddleware();
9-
10-
// @ts-expect-error - typing issue
11+
// @ts-expect-error - typing issue in LangChain
1112
const result = await middleware.beforeAgent?.({ messages: [] });
1213
expect(result).toBeUndefined();
1314
});
1415

15-
it("should not modify messages without tool calls", async () => {
16+
it("should return undefined when messages is undefined", async () => {
17+
const middleware = createPatchToolCallsMiddleware();
18+
// @ts-expect-error - typing issue in LangChain
19+
const result = await middleware.beforeAgent?.({ messages: undefined });
20+
expect(result).toBeUndefined();
21+
});
22+
23+
it("should return undefined when there are no AI messages with tool calls", async () => {
1624
const middleware = createPatchToolCallsMiddleware();
1725
const messages = [
1826
new HumanMessage({ content: "Hello" }),
1927
new AIMessage({ content: "Hi there!" }),
2028
new HumanMessage({ content: "How are you?" }),
2129
];
2230

23-
// @ts-expect-error - typing issue
31+
// @ts-expect-error - typing issue in LangChain
2432
const result = await middleware.beforeAgent?.({ messages });
33+
expect(result).toBeUndefined();
34+
});
2535

26-
expect(result).toBeDefined();
27-
// Should have RemoveMessage + original messages
28-
expect(result?.messages.length).toBe(messages.length + 1);
36+
it("should return undefined when all tool calls have corresponding ToolMessages", async () => {
37+
const middleware = createPatchToolCallsMiddleware();
38+
const messages = [
39+
new HumanMessage({ content: "Read a file" }),
40+
new AIMessage({
41+
content: "",
42+
tool_calls: [
43+
{
44+
id: "call_123",
45+
name: "read_file",
46+
args: { path: "/test.txt" },
47+
},
48+
],
49+
}),
50+
new ToolMessage({
51+
content: "File contents here",
52+
name: "read_file",
53+
tool_call_id: "call_123",
54+
}),
55+
new AIMessage({ content: "Here's the file content" }),
56+
];
57+
58+
// @ts-expect-error - typing issue in LangChain
59+
const result = await middleware.beforeAgent?.({ messages });
60+
expect(result).toBeUndefined();
61+
});
62+
63+
it("should return undefined when AI message has empty tool_calls array", async () => {
64+
const middleware = createPatchToolCallsMiddleware();
65+
const messages = [
66+
new AIMessage({
67+
content: "No tools",
68+
tool_calls: [],
69+
}),
70+
];
71+
72+
// @ts-expect-error - typing issue in LangChain
73+
const result = await middleware.beforeAgent?.({ messages });
74+
expect(result).toBeUndefined();
75+
});
76+
77+
it("should return undefined when AI message has null tool_calls", async () => {
78+
const middleware = createPatchToolCallsMiddleware();
79+
const messages = [
80+
new AIMessage({
81+
content: "Also no tools",
82+
tool_calls: null as any,
83+
}),
84+
];
85+
86+
// @ts-expect-error - typing issue in LangChain
87+
const result = await middleware.beforeAgent?.({ messages });
88+
expect(result).toBeUndefined();
2989
});
3090
});
3191

32-
describe("dangling tool calls", () => {
92+
describe("dangling tool calls (should patch)", () => {
3393
it("should add synthetic ToolMessage for dangling tool call", async () => {
3494
const middleware = createPatchToolCallsMiddleware();
3595
const messages = [
@@ -47,13 +107,18 @@ describe("createPatchToolCallsMiddleware", () => {
47107
new HumanMessage({ content: "Never mind" }),
48108
];
49109

50-
// @ts-expect-error - typing issue
110+
// @ts-expect-error - typing issue in LangChain
51111
const result = await middleware.beforeAgent?.({ messages });
52112

53113
expect(result).toBeDefined();
54114
// Should have RemoveMessage + 3 original + 1 synthetic ToolMessage
55115
expect(result?.messages.length).toBe(5);
56116

117+
// First message should be RemoveMessage
118+
const firstMsg = result?.messages[0];
119+
expect(firstMsg).toBeInstanceOf(RemoveMessage);
120+
expect((firstMsg as RemoveMessage).id).toBe(REMOVE_ALL_MESSAGES);
121+
57122
// Find the synthetic ToolMessage and verify its content
58123
const toolMessage = result?.messages.find(
59124
(m: any) => ToolMessage.isInstance(m) && m.tool_call_id === "call_123",
@@ -63,37 +128,71 @@ describe("createPatchToolCallsMiddleware", () => {
63128
expect(toolMessage?.name).toBe("read_file");
64129
});
65130

66-
it("should not add ToolMessage when corresponding ToolMessage already exists", async () => {
131+
it("should patch multiple dangling tool calls in a single AI message", async () => {
67132
const middleware = createPatchToolCallsMiddleware();
68133
const messages = [
69-
new HumanMessage({ content: "Read a file" }),
134+
new HumanMessage({ content: "Do multiple things" }),
70135
new AIMessage({
71136
content: "",
72137
tool_calls: [
73-
{
74-
id: "call_123",
75-
name: "read_file",
76-
args: { path: "/test.txt" },
77-
},
138+
{ id: "call_1", name: "tool_a", args: {} },
139+
{ id: "call_2", name: "tool_b", args: {} },
78140
],
79141
}),
80-
new ToolMessage({
81-
content: "File contents here",
82-
name: "read_file",
83-
tool_call_id: "call_123",
84-
}),
85-
new AIMessage({ content: "Here's the file content" }),
142+
// Both tool calls are dangling
86143
];
87144

88-
// @ts-expect-error - typing issue
145+
// @ts-expect-error - typing issue in LangChain
89146
const result = await middleware.beforeAgent?.({ messages });
90147

91148
expect(result).toBeDefined();
92-
// Should have RemoveMessage + 4 original messages, no synthetic ones
149+
// RemoveMessage + 2 original + 2 synthetic ToolMessages
93150
expect(result?.messages.length).toBe(5);
151+
152+
// Should have synthetic ToolMessages for both dangling calls
153+
const syntheticMsgs = result?.messages.filter(
154+
(m: any) =>
155+
ToolMessage.isInstance(m) &&
156+
(m.tool_call_id === "call_1" || m.tool_call_id === "call_2"),
157+
);
158+
expect(syntheticMsgs?.length).toBe(2);
159+
});
160+
161+
it("should handle multiple AI messages with dangling tool calls", async () => {
162+
const middleware = createPatchToolCallsMiddleware();
163+
const messages = [
164+
new AIMessage({
165+
content: "",
166+
tool_calls: [{ id: "call_1", name: "tool_a", args: {} }],
167+
}),
168+
new HumanMessage({ content: "msg1" }),
169+
new AIMessage({
170+
content: "",
171+
tool_calls: [{ id: "call_2", name: "tool_b", args: {} }],
172+
}),
173+
new HumanMessage({ content: "msg2" }),
174+
];
175+
176+
// @ts-expect-error - typing issue in LangChain
177+
const result = await middleware.beforeAgent?.({ messages });
178+
179+
expect(result).toBeDefined();
180+
// RemoveMessage + 4 original + 2 synthetic ToolMessages
181+
expect(result?.messages.length).toBe(7);
182+
183+
// Both tool calls should have synthetic responses
184+
const toolMessage1 = result?.messages.find(
185+
(m: any) => ToolMessage.isInstance(m) && m.tool_call_id === "call_1",
186+
);
187+
const toolMessage2 = result?.messages.find(
188+
(m: any) => ToolMessage.isInstance(m) && m.tool_call_id === "call_2",
189+
);
190+
191+
expect(toolMessage1).toBeDefined();
192+
expect(toolMessage2).toBeDefined();
94193
});
95194

96-
it("should handle mixed scenario: some tool calls have responses, some don't", async () => {
195+
it("should only patch dangling tool calls, not ones with responses", async () => {
97196
const middleware = createPatchToolCallsMiddleware();
98197
const messages = [
99198
new HumanMessage({ content: "Do two things" }),
@@ -120,11 +219,11 @@ describe("createPatchToolCallsMiddleware", () => {
120219
new HumanMessage({ content: "Thanks" }),
121220
];
122221

123-
// @ts-expect-error - typing issue
222+
// @ts-expect-error - typing issue in LangChain
124223
const result = await middleware.beforeAgent?.({ messages });
125224

126225
expect(result).toBeDefined();
127-
// Should have RemoveMessage + 4 original + 1 synthetic ToolMessage for call_1
226+
// RemoveMessage + 4 original + 1 synthetic ToolMessage for call_1
128227
expect(result?.messages.length).toBe(6);
129228

130229
// Check synthetic ToolMessage for call_1 exists (dangling)
@@ -147,61 +246,4 @@ describe("createPatchToolCallsMiddleware", () => {
147246
expect(originalToolMessage).toBeDefined();
148247
});
149248
});
150-
151-
describe("edge cases", () => {
152-
it("should handle AI message with empty or null tool_calls", async () => {
153-
const middleware = createPatchToolCallsMiddleware();
154-
const messages = [
155-
new AIMessage({
156-
content: "No tools",
157-
tool_calls: [],
158-
}),
159-
new AIMessage({
160-
content: "Also no tools",
161-
tool_calls: null as any,
162-
}),
163-
];
164-
165-
// @ts-expect-error - typing issue
166-
const result = await middleware.beforeAgent?.({ messages });
167-
168-
expect(result).toBeDefined();
169-
// Should have RemoveMessage + 2 original messages
170-
expect(result?.messages.length).toBe(3);
171-
});
172-
173-
it("should handle multiple AI messages with dangling tool calls", async () => {
174-
const middleware = createPatchToolCallsMiddleware();
175-
const messages = [
176-
new AIMessage({
177-
content: "",
178-
tool_calls: [{ id: "call_1", name: "tool_a", args: {} }],
179-
}),
180-
new HumanMessage({ content: "msg1" }),
181-
new AIMessage({
182-
content: "",
183-
tool_calls: [{ id: "call_2", name: "tool_b", args: {} }],
184-
}),
185-
new HumanMessage({ content: "msg2" }),
186-
];
187-
188-
// @ts-expect-error - typing issue
189-
const result = await middleware.beforeAgent?.({ messages });
190-
191-
expect(result).toBeDefined();
192-
// RemoveMessage + 4 original + 2 synthetic ToolMessages
193-
expect(result?.messages.length).toBe(7);
194-
195-
// Both tool calls should have synthetic responses
196-
const toolMessage1 = result?.messages.find(
197-
(m: any) => ToolMessage.isInstance(m) && m.tool_call_id === "call_1",
198-
);
199-
const toolMessage2 = result?.messages.find(
200-
(m: any) => ToolMessage.isInstance(m) && m.tool_call_id === "call_2",
201-
);
202-
203-
expect(toolMessage1).toBeDefined();
204-
expect(toolMessage2).toBeDefined();
205-
});
206-
});
207249
});

‎libs/deepagents/src/middleware/patch_tool_calls.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ export function createPatchToolCallsMiddleware() {
4141
}
4242

4343
const patchedMessages: any[] = [];
44+
let needsPatch = false;
4445

4546
// Iterate over the messages and add any dangling tool calls
4647
for (let i = 0; i < messages.length; i++) {
@@ -60,6 +61,7 @@ export function createPatchToolCallsMiddleware() {
6061

6162
if (!correspondingToolMsg) {
6263
// We have a dangling tool call which needs a ToolMessage
64+
needsPatch = true;
6365
const toolMsg = `Tool call ${toolCall.name} with id ${toolCall.id} was cancelled - another message came in before it could be completed.`;
6466
patchedMessages.push(
6567
new ToolMessage({
@@ -73,6 +75,13 @@ export function createPatchToolCallsMiddleware() {
7375
}
7476
}
7577

78+
/**
79+
* Only trigger REMOVE_ALL_MESSAGES if patching is actually needed
80+
*/
81+
if (!needsPatch) {
82+
return;
83+
}
84+
7685
// Return state update with RemoveMessage followed by patched messages
7786
return {
7887
messages: [

0 commit comments

Comments
 (0)