Skip to content

Commit 102b0f7

Browse files
authored
Fix streaming tool calls being overwritten at index 0 (#211)
## Problem When streaming multiple tool calls, new tool calls were corrupting existing ones. The bug was in `combineUIMessages`: ```typescript // BUGGY CODE const previousPartIndex = newParts.findIndex( (p) => getToolCallId(p) === toolCallId, ); const previousPart = newParts.splice(previousPartIndex, 1)[0]; if (!previousPart) { newParts.push(part); continue; } ``` When `findIndex` returns `-1` (tool call ID not found in previous parts): - `splice(-1, 1)` removes the **last element** of the array (JavaScript quirk!) - `previousPart` is that removed element (truthy, not undefined) - The new part incorrectly merges with the last element instead of being added - This corrupts the parts array, causing tool calls to disappear ## Solution Check for `-1` explicitly before calling splice: ```typescript // FIXED CODE if (previousPartIndex === -1) { // Tool call not found, add as new newParts.push(part); continue; } const previousPart = newParts.splice(previousPartIndex, 1)[0]; newParts.push(mergeParts(previousPart, part)); ``` Fixes #182 ## Test plan - [x] Test: preserves all tool calls when combining messages - [x] Test: accumulates tool calls progressively (A → A,B → A,B,C) - [x] Test: merges tool calls with same toolCallId correctly 🤖 Generated with [Claude Code](https://claude.ai/code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Fixed an issue in message combination that could unintentionally remove elements when specific tool calls were not found. * **Tests** * Added comprehensive test suite validating preservation and progressive accumulation of tool calls, state transitions, and deduplication during message combination. <sub>✏️ Tip: You can customize this high-level summary in your review settings.</sub> <!-- end of auto-generated comment: release notes by coderabbit.ai -->
2 parents 74cbf01 + 445dfff commit 102b0f7

File tree

2 files changed

+242
-2
lines changed

2 files changed

+242
-2
lines changed
Lines changed: 239 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,239 @@
1+
import { describe, it, expect } from "vitest";
2+
import { combineUIMessages, type UIMessage } from "./UIMessages.js";
3+
4+
describe("combineUIMessages", () => {
5+
it("should preserve all tool calls when combining messages", () => {
6+
const messages: UIMessage[] = [
7+
{
8+
id: "msg1",
9+
key: "thread-1-0",
10+
order: 1,
11+
stepOrder: 0,
12+
status: "streaming",
13+
role: "assistant",
14+
parts: [
15+
{
16+
type: "tool-toolA",
17+
toolCallId: "call_A",
18+
state: "input-available",
19+
input: {},
20+
},
21+
],
22+
text: "",
23+
_creationTime: Date.now(),
24+
},
25+
{
26+
id: "msg1",
27+
key: "thread-1-0",
28+
order: 1,
29+
stepOrder: 0,
30+
status: "streaming",
31+
role: "assistant",
32+
parts: [
33+
{
34+
type: "tool-toolB",
35+
toolCallId: "call_B",
36+
state: "input-available",
37+
input: {},
38+
},
39+
],
40+
text: "",
41+
_creationTime: Date.now(),
42+
},
43+
];
44+
45+
const result = combineUIMessages(messages);
46+
47+
expect(result).toHaveLength(1);
48+
expect(result[0].parts).toHaveLength(2);
49+
50+
const toolCallIds = result[0].parts
51+
.filter((p) => p.type.startsWith("tool-"))
52+
.map((p: any) => p.toolCallId);
53+
54+
expect(toolCallIds).toContain("call_A");
55+
expect(toolCallIds).toContain("call_B");
56+
});
57+
58+
it("should accumulate tool calls progressively (issue #182)", () => {
59+
// Simulating: A(started) → B → C → A(result)
60+
const messages: UIMessage[] = [
61+
{
62+
id: "msg1",
63+
key: "thread-1-0",
64+
order: 1,
65+
stepOrder: 0,
66+
status: "streaming",
67+
role: "assistant",
68+
parts: [
69+
{
70+
type: "tool-toolA",
71+
toolCallId: "call_A",
72+
state: "input-available",
73+
input: {},
74+
},
75+
],
76+
text: "",
77+
_creationTime: Date.now(),
78+
},
79+
{
80+
id: "msg1",
81+
key: "thread-1-0",
82+
order: 1,
83+
stepOrder: 0,
84+
status: "streaming",
85+
role: "assistant",
86+
parts: [
87+
{
88+
type: "tool-toolA",
89+
toolCallId: "call_A",
90+
state: "input-available",
91+
input: {},
92+
},
93+
{
94+
type: "tool-toolB",
95+
toolCallId: "call_B",
96+
state: "input-available",
97+
input: {},
98+
},
99+
],
100+
text: "",
101+
_creationTime: Date.now(),
102+
},
103+
{
104+
id: "msg1",
105+
key: "thread-1-0",
106+
order: 1,
107+
stepOrder: 0,
108+
status: "streaming",
109+
role: "assistant",
110+
parts: [
111+
{
112+
type: "tool-toolA",
113+
toolCallId: "call_A",
114+
state: "input-available",
115+
input: {},
116+
},
117+
{
118+
type: "tool-toolB",
119+
toolCallId: "call_B",
120+
state: "input-available",
121+
input: {},
122+
},
123+
{
124+
type: "tool-toolC",
125+
toolCallId: "call_C",
126+
state: "input-available",
127+
input: {},
128+
},
129+
],
130+
text: "",
131+
_creationTime: Date.now(),
132+
},
133+
{
134+
id: "msg1",
135+
key: "thread-1-0",
136+
order: 1,
137+
stepOrder: 0,
138+
status: "success",
139+
role: "assistant",
140+
parts: [
141+
{
142+
type: "tool-toolA",
143+
toolCallId: "call_A",
144+
state: "output-available",
145+
input: {},
146+
output: "success",
147+
},
148+
{
149+
type: "tool-toolB",
150+
toolCallId: "call_B",
151+
state: "input-available",
152+
input: {},
153+
},
154+
{
155+
type: "tool-toolC",
156+
toolCallId: "call_C",
157+
state: "input-available",
158+
input: {},
159+
},
160+
],
161+
text: "",
162+
_creationTime: Date.now(),
163+
},
164+
];
165+
166+
const result = combineUIMessages(messages);
167+
168+
expect(result).toHaveLength(1);
169+
expect(result[0].parts).toHaveLength(3);
170+
171+
const toolCallIds = result[0].parts
172+
.filter((p) => p.type.startsWith("tool-"))
173+
.map((p: any) => p.toolCallId);
174+
175+
// All tool calls should be present
176+
expect(toolCallIds).toContain("call_A");
177+
expect(toolCallIds).toContain("call_B");
178+
expect(toolCallIds).toContain("call_C");
179+
180+
// Tool A should have the final state (output-available)
181+
const toolA = result[0].parts.find(
182+
(p: any) => p.type === "tool-toolA" && p.toolCallId === "call_A",
183+
) as any;
184+
expect(toolA.state).toBe("output-available");
185+
expect(toolA.output).toBe("success");
186+
});
187+
188+
it("should merge tool calls with same toolCallId", () => {
189+
const messages: UIMessage[] = [
190+
{
191+
id: "msg1",
192+
key: "thread-1-0",
193+
order: 1,
194+
stepOrder: 0,
195+
status: "streaming",
196+
role: "assistant",
197+
parts: [
198+
{
199+
type: "tool-toolA",
200+
toolCallId: "call_A",
201+
state: "input-available",
202+
input: { test: "input" },
203+
},
204+
],
205+
text: "",
206+
_creationTime: Date.now(),
207+
},
208+
{
209+
id: "msg1",
210+
key: "thread-1-0",
211+
order: 1,
212+
stepOrder: 0,
213+
status: "success",
214+
role: "assistant",
215+
parts: [
216+
{
217+
type: "tool-toolA",
218+
toolCallId: "call_A",
219+
state: "output-available",
220+
input: { test: "input" },
221+
output: "completed",
222+
},
223+
],
224+
text: "",
225+
_creationTime: Date.now(),
226+
},
227+
];
228+
229+
const result = combineUIMessages(messages);
230+
231+
expect(result).toHaveLength(1);
232+
expect(result[0].parts).toHaveLength(1);
233+
234+
const toolCall = result[0].parts[0] as any;
235+
expect(toolCall.toolCallId).toBe("call_A");
236+
expect(toolCall.state).toBe("output-available");
237+
expect(toolCall.output).toBe("completed");
238+
});
239+
});

src/UIMessages.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -583,11 +583,12 @@ export function combineUIMessages(messages: UIMessage[]): UIMessage[] {
583583
const previousPartIndex = newParts.findIndex(
584584
(p) => getToolCallId(p) === toolCallId,
585585
);
586-
const previousPart = newParts.splice(previousPartIndex, 1)[0];
587-
if (!previousPart) {
586+
if (previousPartIndex === -1) {
587+
// Tool call not found in previous parts, add it as new
588588
newParts.push(part);
589589
continue;
590590
}
591+
const previousPart = newParts.splice(previousPartIndex, 1)[0];
591592
newParts.push(mergeParts(previousPart, part));
592593
}
593594
acc.push({

0 commit comments

Comments
 (0)