1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
|
import type { Role, StepId } from "@dispatch/wire";
import { describe, expect, it } from "vitest";
import { groupRenderedChunks } from "./groups";
import type { RenderedChunk } from "./types";
const text = (seq: number, role: Role, t: string, provisional = false): RenderedChunk => ({
seq,
role,
chunk: { type: "text", text: t },
provisional,
});
const call = (seq: number, id: string, stepId?: string, provisional = false): RenderedChunk => ({
seq,
role: "assistant",
chunk: {
type: "tool-call",
toolCallId: id,
toolName: `tool-${id}`,
input: { id },
...(stepId !== undefined ? { stepId: stepId as StepId } : {}),
},
provisional,
});
const result = (seq: number, id: string, stepId?: string, provisional = false): RenderedChunk => ({
seq,
role: "tool",
chunk: {
type: "tool-result",
toolCallId: id,
toolName: `tool-${id}`,
content: `result-${id}`,
isError: false,
...(stepId !== undefined ? { stepId: stepId as StepId } : {}),
},
provisional,
});
describe("groupRenderedChunks", () => {
it("returns no groups for an empty stream", () => {
expect(groupRenderedChunks([])).toEqual([]);
});
it("passes non-tool chunks through as single groups, in order", () => {
const groups = groupRenderedChunks([text(1, "user", "hi"), text(2, "assistant", "hello")]);
expect(groups).toHaveLength(2);
expect(groups.every((g) => g.kind === "single")).toBe(true);
});
it("does NOT batch a single tool call (one per step) — call+result stay separate singles", () => {
const groups = groupRenderedChunks([call(1, "a", "s1"), result(2, "a", "s1")]);
expect(groups).toHaveLength(2);
expect(groups.map((g) => g.kind)).toEqual(["single", "single"]);
});
it("does NOT batch tool calls that have no stepId (pre-0.2.0 replay)", () => {
const groups = groupRenderedChunks([
call(1, "a"),
call(2, "b"),
result(3, "a"),
result(4, "b"),
]);
expect(groups).toHaveLength(4);
expect(groups.every((g) => g.kind === "single")).toBe(true);
});
it("batches 2+ calls sharing a stepId into one group, pairing each with its result", () => {
const groups = groupRenderedChunks([
call(1, "a", "s1"),
call(2, "b", "s1"),
result(3, "a", "s1"),
result(4, "b", "s1"),
]);
expect(groups).toHaveLength(1);
const g = groups[0];
if (g?.kind !== "tool-batch") throw new Error("expected a tool-batch group");
expect(g.stepId).toBe("s1");
expect(g.entries).toHaveLength(2);
expect(g.entries[0]?.call.toolCallId).toBe("a");
expect(g.entries[0]?.result?.content).toBe("result-a");
expect(g.entries[1]?.call.toolCallId).toBe("b");
expect(g.entries[1]?.result?.content).toBe("result-b");
});
it("positions the batch at the first call and keeps surrounding chunks in order", () => {
const groups = groupRenderedChunks([
text(1, "assistant", "before"),
call(2, "a", "s1"),
call(3, "b", "s1"),
result(4, "a", "s1"),
result(5, "b", "s1"),
text(6, "assistant", "after"),
]);
expect(groups.map((g) => g.kind)).toEqual(["single", "tool-batch", "single"]);
});
it("marks the batch provisional when any of its calls/results is provisional", () => {
const groups = groupRenderedChunks([call(1, "a", "s1"), call(2, "b", "s1", true)]);
const g = groups[0];
if (g?.kind !== "tool-batch") throw new Error("expected a tool-batch group");
expect(g.provisional).toBe(true);
expect(g.entries).toHaveLength(2);
expect(g.entries[1]?.result).toBeNull(); // dangling call (no result yet)
});
it("batches one step while leaving a different single-call step ungrouped", () => {
const groups = groupRenderedChunks([
call(1, "a", "s1"),
call(2, "b", "s1"),
call(3, "c", "s2"),
result(4, "a", "s1"),
result(5, "b", "s1"),
result(6, "c", "s2"),
]);
expect(groups.map((g) => g.kind)).toEqual(["tool-batch", "single", "single"]);
const batch = groups[0];
if (batch?.kind !== "tool-batch") throw new Error("expected a tool-batch group");
expect(batch.entries).toHaveLength(2);
// the s2 single call + its result remain as separate single groups
const singles = groups.slice(1);
expect(singles[0]?.kind === "single" && singles[0].chunk.chunk.type).toBe("tool-call");
expect(singles[1]?.kind === "single" && singles[1].chunk.chunk.type).toBe("tool-result");
});
});
|
