refactor: ChatMessage.chunks[] union — interleaved thinking, tool batching, error/system chunks

1 files changed, 437 insertions, 0 deletions

diff --git a/packages/core/tests/chunks/append.test.ts b/packages/core/tests/chunks/append.test.ts
new file mode 100644
index 0000000..c8917c9
--- /dev/null
+++ b/packages/core/tests/chunks/append.test.ts
@@ -0,0 +1,437 @@
+import { describe, expect, it } from "vitest";
+import { appendEventToChunks, applySystemEvent } from "../../src/chunks/append.js";
+import type { AgentEvent, ChatMessage, Chunk } from "../../src/types/index.js";
+
+// ─── helpers ─────────────────────────────────────────────────────
+
+const td = (delta: string): AgentEvent => ({ type: "text-delta", delta });
+const rd = (delta: string): AgentEvent => ({ type: "reasoning-delta", delta });
+const tc = (id: string, name = "fake_tool", args: Record<string, unknown> = {}): AgentEvent => ({
+	type: "tool-call",
+	toolCall: { id, name, arguments: args },
+});
+const tr = (toolCallId: string, result: string, isError = false): AgentEvent => ({
+	type: "tool-result",
+	toolResult: { toolCallId, toolName: "fake_tool", result, isError },
+});
+const so = (data: string, stream: "stdout" | "stderr" = "stdout"): AgentEvent => ({
+	type: "shell-output",
+	data,
+	stream,
+});
+const err = (error: string, statusCode?: number): AgentEvent => ({
+	type: "error",
+	error,
+	...(statusCode !== undefined ? { statusCode } : {}),
+});
+const notice = (message: string): AgentEvent => ({ type: "notice", message });
+const modelChanged = (keyId: string, modelId: string): AgentEvent => ({
+	type: "model-changed",
+	keyId,
+	modelId,
+});
+const configReload: AgentEvent = { type: "config-reload" };
+
+function run(events: AgentEvent[]): Chunk[] {
+	const chunks: Chunk[] = [];
+	for (const e of events) appendEventToChunks(chunks, e);
+	return chunks;
+}
+
+// ─── Required cases from the plan ────────────────────────────────
+
+describe("appendEventToChunks — required cases from plan", () => {
+	it("empty chunks + text-delta → one text chunk with the delta", () => {
+		const chunks = run([td("Hello")]);
+		expect(chunks).toEqual([{ type: "text", text: "Hello" }]);
+	});
+
+	it("two consecutive text-deltas → one text chunk with concatenated text", () => {
+		const chunks = run([td("Hello, "), td("world!")]);
+		expect(chunks).toEqual([{ type: "text", text: "Hello, world!" }]);
+	});
+
+	it("text-delta then reasoning-delta → two chunks (text, thinking)", () => {
+		const chunks = run([td("ans: 42"), rd("I should explain")]);
+		expect(chunks).toEqual([
+			{ type: "text", text: "ans: 42" },
+			{ type: "thinking", text: "I should explain" },
+		]);
+	});
+
+	it("text-delta then tool-call → two chunks (text, tool-batch with one entry)", () => {
+		const chunks = run([td("Looking..."), tc("t1", "read_file", { path: "x" })]);
+		expect(chunks).toEqual([
+			{ type: "text", text: "Looking..." },
+			{
+				type: "tool-batch",
+				calls: [{ id: "t1", name: "read_file", arguments: { path: "x" } }],
+			},
+		]);
+	});
+
+	it("two consecutive tool-calls → one tool-batch with two entries", () => {
+		const chunks = run([tc("t1", "read_file"), tc("t2", "list_files")]);
+		expect(chunks).toHaveLength(1);
+		expect(chunks[0]).toMatchObject({
+			type: "tool-batch",
+			calls: [
+				{ id: "t1", name: "read_file" },
+				{ id: "t2", name: "list_files" },
+			],
+		});
+	});
+
+	it("tool-call then tool-call then text → two chunks (tool-batch with 2 entries, text)", () => {
+		const chunks = run([tc("t1"), tc("t2"), td("done")]);
+		expect(chunks).toHaveLength(2);
+		expect(chunks[0]).toMatchObject({
+			type: "tool-batch",
+			calls: [{ id: "t1" }, { id: "t2" }],
+		});
+		expect(chunks[1]).toEqual({ type: "text", text: "done" });
+	});
+
+	it("tool-result arrives → updates matching tool-call entry in the latest tool-batch chunk by id", () => {
+		const chunks = run([
+			tc("t1"),
+			tc("t2"),
+			tr("t1", "first-result"),
+			tr("t2", "second-result", true),
+		]);
+		expect(chunks).toHaveLength(1);
+		const batch = chunks[0];
+		expect(batch.type).toBe("tool-batch");
+		if (batch.type !== "tool-batch") throw new Error("type guard");
+		expect(batch.calls[0]).toMatchObject({ id: "t1", result: "first-result", isError: false });
+		expect(batch.calls[1]).toMatchObject({ id: "t2", result: "second-result", isError: true });
+	});
+
+	it("shell-output arrives → appends to the most recent tool-call's shellOutput", () => {
+		const chunks = run([
+			tc("t1", "run_shell"),
+			so("hello\n", "stdout"),
+			so("world\n", "stdout"),
+			so("err!\n", "stderr"),
+		]);
+		expect(chunks).toHaveLength(1);
+		const batch = chunks[0];
+		if (batch.type !== "tool-batch") throw new Error("type guard");
+		expect(batch.calls[0]?.shellOutput).toEqual({
+			stdout: "hello\nworld\n",
+			stderr: "err!\n",
+		});
+	});
+
+	it("error event → opens an error chunk; subsequent events go to new chunks", () => {
+		const chunks = run([td("partial..."), err("network failed", 503), td("recovery")]);
+		expect(chunks).toEqual([
+			{ type: "text", text: "partial..." },
+			{ type: "error", message: "network failed", statusCode: 503 },
+			{ type: "text", text: "recovery" },
+		]);
+	});
+
+	it("system event during text run → closes text, opens system, would re-open text on next text-delta", () => {
+		const chunks = run([td("first "), notice("model swap"), td("second")]);
+		expect(chunks).toEqual([
+			{ type: "text", text: "first " },
+			{ type: "system", kind: "notice", text: "model swap" },
+			{ type: "text", text: "second" },
+		]);
+	});
+
+	it("two consecutive system events → two separate system chunks (no coalescing)", () => {
+		const chunks = run([notice("a"), notice("b")]);
+		expect(chunks).toEqual([
+			{ type: "system", kind: "notice", text: "a" },
+			{ type: "system", kind: "notice", text: "b" },
+		]);
+	});
+
+	it("interleaved think → text → think → tool → think → text → 6 chunks in order", () => {
+		const chunks = run([
+			rd("planning..."),
+			td("here goes:"),
+			rd("hmm, actually"),
+			tc("t1", "read_file"),
+			rd("ok now"),
+			td("and so..."),
+		]);
+		expect(chunks.map((c) => c.type)).toEqual([
+			"thinking",
+			"text",
+			"thinking",
+			"tool-batch",
+			"thinking",
+			"text",
+		]);
+		expect(chunks[0]).toEqual({ type: "thinking", text: "planning..." });
+		expect(chunks[1]).toEqual({ type: "text", text: "here goes:" });
+		expect(chunks[2]).toEqual({ type: "thinking", text: "hmm, actually" });
+		expect(chunks[3]).toMatchObject({ type: "tool-batch", calls: [{ id: "t1" }] });
+		expect(chunks[4]).toEqual({ type: "thinking", text: "ok now" });
+		expect(chunks[5]).toEqual({ type: "text", text: "and so..." });
+	});
+});
+
+// ─── Additional transition coverage ──────────────────────────────
+
+describe("appendEventToChunks — transition matrix", () => {
+	it("thinking → thinking coalesces", () => {
+		const chunks = run([rd("a"), rd("b")]);
+		expect(chunks).toEqual([{ type: "thinking", text: "ab" }]);
+	});
+
+	it("thinking → text opens a new text chunk", () => {
+		const chunks = run([rd("think"), td("speak")]);
+		expect(chunks).toEqual([
+			{ type: "thinking", text: "think" },
+			{ type: "text", text: "speak" },
+		]);
+	});
+
+	it("tool-batch → text opens a new text chunk", () => {
+		const chunks = run([tc("t1"), td("after tool")]);
+		expect(chunks).toHaveLength(2);
+		expect(chunks[1]).toEqual({ type: "text", text: "after tool" });
+	});
+
+	it("text → reasoning-delta after a multi-delta text run still splits cleanly", () => {
+		const chunks = run([td("a"), td("b"), rd("x"), rd("y"), td("c")]);
+		expect(chunks).toEqual([
+			{ type: "text", text: "ab" },
+			{ type: "thinking", text: "xy" },
+			{ type: "text", text: "c" },
+		]);
+	});
+
+	it("error → text opens a fresh text chunk after the error", () => {
+		const chunks = run([err("boom"), td("recovered")]);
+		expect(chunks).toEqual([
+			{ type: "error", message: "boom" },
+			{ type: "text", text: "recovered" },
+		]);
+	});
+
+	it("two consecutive errors stay as two error chunks (no coalescing)", () => {
+		const chunks = run([err("first"), err("second", 429)]);
+		expect(chunks).toEqual([
+			{ type: "error", message: "first" },
+			{ type: "error", message: "second", statusCode: 429 },
+		]);
+	});
+
+	it("system → tool-call opens a new tool-batch (does not extend the system chunk)", () => {
+		const chunks = run([notice("info"), tc("t1")]);
+		expect(chunks).toHaveLength(2);
+		expect(chunks[1]).toMatchObject({ type: "tool-batch", calls: [{ id: "t1" }] });
+	});
+
+	it("tool-result with no matching call is silently dropped", () => {
+		const chunks = run([td("hi"), tr("no-such-id", "ignored")]);
+		expect(chunks).toEqual([{ type: "text", text: "hi" }]);
+	});
+
+	it("shell-output with no tool-batch in scope is silently dropped", () => {
+		const chunks = run([td("hi"), so("orphan")]);
+		expect(chunks).toEqual([{ type: "text", text: "hi" }]);
+	});
+
+	it("tool-result for an earlier batch still updates the right call (results can arrive late)", () => {
+		// Order: tc -> td -> tc(new batch) -> tr(for first batch's id)
+		const chunks = run([
+			tc("t1", "read_file"),
+			td("midstream text"),
+			tc("t2", "list_files"),
+			tr("t1", "late result for first"),
+		]);
+		// Two tool-batches, separated by the text chunk. The result must land
+		// inside the FIRST batch (the one containing t1), not the most-recent.
+		expect(chunks.map((c) => c.type)).toEqual(["tool-batch", "text", "tool-batch"]);
+		const first = chunks[0];
+		if (first?.type !== "tool-batch") throw new Error("type guard");
+		expect(first.calls[0]).toMatchObject({ id: "t1", result: "late result for first" });
+		const second = chunks[2];
+		if (second?.type !== "tool-batch") throw new Error("type guard");
+		// t2 in the second batch has no result yet.
+		expect(second.calls[0]?.result).toBeUndefined();
+	});
+
+	it("shell-output goes to the most recent tool-batch's most recent entry, even with intervening chunks", () => {
+		// First batch's tool runs, emits output, then later a second batch starts and emits output.
+		const chunks = run([
+			tc("t1", "run_shell"),
+			so("first-stdout\n"),
+			td("interlude"),
+			tc("t2", "run_shell"),
+			so("second-stdout\n"),
+		]);
+		expect(chunks.map((c) => c.type)).toEqual(["tool-batch", "text", "tool-batch"]);
+		const first = chunks[0];
+		const second = chunks[2];
+		if (first?.type !== "tool-batch" || second?.type !== "tool-batch") {
+			throw new Error("type guard");
+		}
+		expect(first.calls[0]?.shellOutput).toEqual({ stdout: "first-stdout\n", stderr: "" });
+		expect(second.calls[0]?.shellOutput).toEqual({ stdout: "second-stdout\n", stderr: "" });
+	});
+
+	it("model-changed event opens a system chunk with kind=model-changed", () => {
+		const chunks = run([modelChanged("anthropic-1", "claude-sonnet-4")]);
+		expect(chunks).toEqual([
+			{
+				type: "system",
+				kind: "model-changed",
+				text: "Switched to claude-sonnet-4 (anthropic-1)",
+			},
+		]);
+	});
+
+	it("config-reload event opens a system chunk with kind=config-reload", () => {
+		const chunks = run([configReload]);
+		expect(chunks).toEqual([
+			{ type: "system", kind: "config-reload", text: "Configuration reloaded" },
+		]);
+	});
+
+	it("non-content events (status / done / task-list-update / message-queued etc.) are no-ops", () => {
+		const chunks = run([
+			td("hello"),
+			{ type: "status", status: "running" },
+			{ type: "task-list-update", tasks: [] },
+			{
+				type: "tab-created",
+				id: "tab1",
+				title: "x",
+				keyId: null,
+				modelId: null,
+				parentTabId: null,
+				workingDirectory: null,
+			},
+			{ type: "message-queued", tabId: "t", messageId: "m", message: "queued" },
+			{ type: "message-consumed", tabId: "t", messageIds: ["m"] },
+			{ type: "message-cancelled", tabId: "t", messageId: "m" },
+			{
+				type: "done",
+				message: { role: "assistant", chunks: [] },
+			},
+			td(" world"),
+		]);
+		expect(chunks).toEqual([{ type: "text", text: "hello world" }]);
+	});
+
+	it("error chunk omits statusCode when not provided", () => {
+		const chunks = run([err("boom")]);
+		expect(chunks).toEqual([{ type: "error", message: "boom" }]);
+		// And no stray statusCode key:
+		expect(Object.hasOwn(chunks[0], "statusCode")).toBe(false);
+	});
+
+	it("tool-result updates isError=false correctly (default success path)", () => {
+		const chunks = run([tc("t1"), tr("t1", "ok", false)]);
+		const batch = chunks[0];
+		if (batch?.type !== "tool-batch") throw new Error("type guard");
+		expect(batch.calls[0]).toMatchObject({ result: "ok", isError: false });
+	});
+});
+
+// ─── applySystemEvent routing ────────────────────────────────────
+
+describe("applySystemEvent", () => {
+	type Msg = { id: string; role: "user" | "assistant" | "system"; chunks: Chunk[] };
+
+	let counter = 0;
+	const idFactory = () => `gen-${++counter}`;
+
+	it("creates a new role:system message when message list is empty", () => {
+		counter = 0;
+		const messages: Msg[] = [];
+		const result = applySystemEvent(messages, { kind: "notice", text: "hello" }, idFactory);
+		expect(result.messageId).toBe("gen-1");
+		expect(messages).toEqual([
+			{
+				id: "gen-1",
+				role: "system",
+				chunks: [{ type: "system", kind: "notice", text: "hello" }],
+			},
+		]);
+	});
+
+	it("creates a new role:system message when last message is user", () => {
+		counter = 0;
+		const messages: Msg[] = [{ id: "u1", role: "user", chunks: [{ type: "text", text: "hi" }] }];
+		const result = applySystemEvent(messages, { kind: "model-changed", text: "swap" }, idFactory);
+		expect(result.messageId).toBe("gen-1");
+		expect(messages).toHaveLength(2);
+		expect(messages[1]).toMatchObject({
+			id: "gen-1",
+			role: "system",
+			chunks: [{ type: "system", kind: "model-changed", text: "swap" }],
+		});
+	});
+
+	it("creates a new role:system message when last message is assistant", () => {
+		counter = 0;
+		const messages: Msg[] = [
+			{ id: "a1", role: "assistant", chunks: [{ type: "text", text: "done" }] },
+		];
+		applySystemEvent(messages, { kind: "config-reload", text: "reloaded" }, idFactory);
+		expect(messages).toHaveLength(2);
+		expect(messages[1]?.role).toBe("system");
+	});
+
+	it("appends a chunk to the existing system message when last message is role:system", () => {
+		counter = 0;
+		const messages: Msg[] = [
+			{
+				id: "s1",
+				role: "system",
+				chunks: [{ type: "system", kind: "notice", text: "first" }],
+			},
+		];
+		const result = applySystemEvent(messages, { kind: "notice", text: "second" }, idFactory);
+		expect(result.messageId).toBe("s1");
+		expect(messages).toHaveLength(1);
+		expect(messages[0]?.chunks).toEqual([
+			{ type: "system", kind: "notice", text: "first" },
+			{ type: "system", kind: "notice", text: "second" },
+		]);
+	});
+
+	it("multiple consecutive calls accumulate in the same system message", () => {
+		counter = 0;
+		const messages: Msg[] = [{ id: "u1", role: "user", chunks: [{ type: "text", text: "hi" }] }];
+		applySystemEvent(messages, { kind: "notice", text: "a" }, idFactory);
+		applySystemEvent(messages, { kind: "notice", text: "b" }, idFactory);
+		applySystemEvent(messages, { kind: "model-changed", text: "c" }, idFactory);
+		expect(messages).toHaveLength(2);
+		const sys = messages[1];
+		expect(sys?.role).toBe("system");
+		expect(sys?.chunks).toEqual([
+			{ type: "system", kind: "notice", text: "a" },
+			{ type: "system", kind: "notice", text: "b" },
+			{ type: "system", kind: "model-changed", text: "c" },
+		]);
+	});
+
+	it("returns the same messageId across appends to the same system message", () => {
+		counter = 0;
+		const messages: Msg[] = [];
+		const first = applySystemEvent(messages, { kind: "notice", text: "a" }, idFactory);
+		const second = applySystemEvent(messages, { kind: "notice", text: "b" }, idFactory);
+		expect(first.messageId).toBe(second.messageId);
+	});
+
+	it("works against the core ChatMessage shape (with id added by caller)", () => {
+		// Sanity: ChatMessage has {role, chunks}; the caller layers id on top.
+		// This test exists to prove the generic constraint doesn't reject the
+		// real persistence/in-memory shape we'll see in Phase 5.
+		counter = 0;
+		const messages: Array<ChatMessage & { id: string }> = [];
+		const result = applySystemEvent(messages, { kind: "cancelled", text: "user stop" }, idFactory);
+		expect(result.messageId).toBe("gen-1");
+		expect(messages[0]?.role).toBe("system");
+		expect(messages[0]?.chunks[0]).toMatchObject({ kind: "cancelled", text: "user stop" });
+	});
+});