summaryrefslogtreecommitdiffhomepage
path: root/src/core/chunks
diff options
context:
space:
mode:
Diffstat (limited to 'src/core/chunks')
-rw-r--r--src/core/chunks/groups.test.ts125
-rw-r--r--src/core/chunks/groups.ts95
-rw-r--r--src/core/chunks/index.ts2
-rw-r--r--src/core/chunks/reducer.test.ts12
-rw-r--r--src/core/chunks/reducer.ts2
5 files changed, 234 insertions, 2 deletions
diff --git a/src/core/chunks/groups.test.ts b/src/core/chunks/groups.test.ts
new file mode 100644
index 0000000..fbfda83
--- /dev/null
+++ b/src/core/chunks/groups.test.ts
@@ -0,0 +1,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");
+ });
+});
diff --git a/src/core/chunks/groups.ts b/src/core/chunks/groups.ts
new file mode 100644
index 0000000..6dc7e10
--- /dev/null
+++ b/src/core/chunks/groups.ts
@@ -0,0 +1,95 @@
+import type { ToolCallChunk, ToolResultChunk } from "@dispatch/wire";
+import type { RenderedChunk } from "./types";
+
+/**
+ * One tool call within a batch, paired with its result (matched by `toolCallId`).
+ * `result` is null while the call is still pending (no result chunk yet).
+ */
+export interface ToolBatchEntry {
+ readonly call: ToolCallChunk;
+ readonly result: ToolResultChunk | null;
+}
+
+/**
+ * A render group: either a single rendered chunk (rendered as today) or a batch
+ * of tool calls the model emitted together in one step (shared `stepId`), to be
+ * rendered as one grouped unit.
+ */
+export type RenderGroup =
+ | { readonly kind: "single"; readonly chunk: RenderedChunk }
+ | {
+ readonly kind: "tool-batch";
+ readonly stepId: string;
+ readonly entries: readonly ToolBatchEntry[];
+ readonly provisional: boolean;
+ };
+
+/**
+ * Group a flat rendered-chunk stream for display. Tool calls sharing a `stepId`
+ * (the backend's authoritative batch key) where the step has 2+ calls become one
+ * `tool-batch` group, positioned at the first call and pairing each call with its
+ * `tool-result` (by `toolCallId`); the absorbed result chunks are not emitted on
+ * their own. Single tool calls (one per step, or no `stepId` — e.g. pre-0.2.0
+ * replay rows) and every non-tool chunk render as `single` groups, in order.
+ *
+ * Pure: input → output, no DOM, no Svelte.
+ */
+export function groupRenderedChunks(rendered: readonly RenderedChunk[]): readonly RenderGroup[] {
+ // 1. Steps that batched 2+ tool calls.
+ const callsPerStep = new Map<string, number>();
+ for (const rc of rendered) {
+ if (rc.chunk.type === "tool-call" && rc.chunk.stepId !== undefined) {
+ callsPerStep.set(rc.chunk.stepId, (callsPerStep.get(rc.chunk.stepId) ?? 0) + 1);
+ }
+ }
+ const batchSteps = new Set<string>();
+ for (const [stepId, count] of callsPerStep) {
+ if (count >= 2) batchSteps.add(stepId);
+ }
+
+ // 2. toolCallIds belonging to a batch (so their results are absorbed), and a
+ // lookup of result chunks by toolCallId for pairing.
+ const batchCallIds = new Set<string>();
+ const resultByCallId = new Map<string, ToolResultChunk>();
+ for (const rc of rendered) {
+ const chunk = rc.chunk;
+ if (chunk.type === "tool-call" && chunk.stepId !== undefined && batchSteps.has(chunk.stepId)) {
+ batchCallIds.add(chunk.toolCallId);
+ } else if (chunk.type === "tool-result" && !resultByCallId.has(chunk.toolCallId)) {
+ resultByCallId.set(chunk.toolCallId, chunk);
+ }
+ }
+
+ // 3. Emit groups in stream order; each batch lands at its first call.
+ const groups: RenderGroup[] = [];
+ const emittedSteps = new Set<string>();
+ for (const rc of rendered) {
+ const chunk = rc.chunk;
+
+ if (chunk.type === "tool-call" && chunk.stepId !== undefined && batchSteps.has(chunk.stepId)) {
+ const stepId = chunk.stepId;
+ if (emittedSteps.has(stepId)) continue;
+ emittedSteps.add(stepId);
+
+ const entries: ToolBatchEntry[] = [];
+ let provisional = false;
+ for (const inner of rendered) {
+ if (inner.chunk.type === "tool-call" && inner.chunk.stepId === stepId) {
+ const result = resultByCallId.get(inner.chunk.toolCallId) ?? null;
+ entries.push({ call: inner.chunk, result });
+ if (inner.provisional) provisional = true;
+ }
+ }
+ groups.push({ kind: "tool-batch", stepId, entries, provisional });
+ continue;
+ }
+
+ if (chunk.type === "tool-result" && batchCallIds.has(chunk.toolCallId)) {
+ continue; // absorbed into its batch
+ }
+
+ groups.push({ kind: "single", chunk: rc });
+ }
+
+ return groups;
+}
diff --git a/src/core/chunks/index.ts b/src/core/chunks/index.ts
index 67739bc..0718c0d 100644
--- a/src/core/chunks/index.ts
+++ b/src/core/chunks/index.ts
@@ -1,3 +1,5 @@
+export type { RenderGroup, ToolBatchEntry } from "./groups";
+export { groupRenderedChunks } from "./groups";
export { appendUserMessage, applyHistory, foldEvent, initialState } from "./reducer";
export { selectChunks, selectMessages } from "./selectors";
export type {
diff --git a/src/core/chunks/reducer.test.ts b/src/core/chunks/reducer.test.ts
index b7165e4..7ecc349 100644
--- a/src/core/chunks/reducer.test.ts
+++ b/src/core/chunks/reducer.test.ts
@@ -1,4 +1,5 @@
import type {
+ StepId,
StoredChunk,
TurnDoneEvent,
TurnErrorEvent,
@@ -39,6 +40,7 @@ const toolCall = (
toolCallId: string,
toolName: string,
input: unknown,
+ stepId = "s0",
): TurnToolCallEvent => ({
type: "tool-call",
conversationId: "c1",
@@ -46,6 +48,7 @@ const toolCall = (
toolCallId,
toolName,
input,
+ stepId: stepId as StepId,
});
const toolResult = (
@@ -53,6 +56,7 @@ const toolResult = (
toolCallId: string,
toolName: string,
content: string,
+ stepId = "s0",
): TurnToolResultEvent => ({
type: "tool-result",
conversationId: "c1",
@@ -61,6 +65,7 @@ const toolResult = (
toolName,
content,
isError: false,
+ stepId: stepId as StepId,
});
const usageEvent = (turnId: string, inputTokens: number, outputTokens: number): TurnUsageEvent => ({
@@ -161,15 +166,17 @@ describe("foldEvent — tool-call then tool-result", () => {
it("tool-call then tool-result render in order", () => {
let s = initialState();
s = foldEvent(s, turnStart("t1"));
- s = foldEvent(s, toolCall("t1", "tc1", "bash", { cmd: "ls" }));
- s = foldEvent(s, toolResult("t1", "tc1", "bash", "file.txt"));
+ s = foldEvent(s, toolCall("t1", "tc1", "bash", { cmd: "ls" }, "t1#0"));
+ s = foldEvent(s, toolResult("t1", "tc1", "bash", "file.txt", "t1#0"));
expect(s.provisional).toHaveLength(2);
expect(s.provisional[0]?.role).toBe("assistant");
+ // foldEvent copies the event's stepId onto the chunk (grouping key).
expect(s.provisional[0]?.chunk).toEqual({
type: "tool-call",
toolCallId: "tc1",
toolName: "bash",
input: { cmd: "ls" },
+ stepId: "t1#0",
});
expect(s.provisional[1]?.role).toBe("tool");
expect(s.provisional[1]?.chunk).toEqual({
@@ -178,6 +185,7 @@ describe("foldEvent — tool-call then tool-result", () => {
toolName: "bash",
content: "file.txt",
isError: false,
+ stepId: "t1#0",
});
});
diff --git a/src/core/chunks/reducer.ts b/src/core/chunks/reducer.ts
index d3b999d..1dcfa39 100644
--- a/src/core/chunks/reducer.ts
+++ b/src/core/chunks/reducer.ts
@@ -106,6 +106,7 @@ export function foldEvent(state: TranscriptState, event: AgentEvent): Transcript
toolCallId: event.toolCallId,
toolName: event.toolName,
input: event.input,
+ stepId: event.stepId,
};
return {
...state,
@@ -122,6 +123,7 @@ export function foldEvent(state: TranscriptState, event: AgentEvent): Transcript
toolName: event.toolName,
content: event.content,
isError: event.isError,
+ stepId: event.stepId,
};
return {
...state,
generated by cgit v1.2.3 (git 2.39.1) at 2026-06-29 02:35:12 +0000