Squashed 'third_party/pi/' content from commit d981de122

git-subtree-dir: third_party/pi
git-subtree-split: d981de1229ef899957bbe968bc8dcda02a21f477

packages/coding-agent/test/branch-summarization.test.ts

@@ -0,0 +1,135 @@
+import type { StreamFn } from "@earendil-works/pi-agent-core";
+import {
+	type AssistantMessage,
+	createAssistantMessageEventStream,
+	fauxAssistantMessage,
+	type Model,
+	type SimpleStreamOptions,
+} from "@earendil-works/pi-ai";
+import { describe, expect, it } from "vitest";
+import { generateBranchSummary } from "../src/core/compaction/index.ts";
+import type { SessionEntry } from "../src/core/session-manager.ts";
+
+const model: Model<"anthropic-messages"> = {
+	id: "test-model",
+	name: "Test Model",
+	api: "anthropic-messages",
+	provider: "anthropic",
+	baseUrl: "https://api.anthropic.com",
+	reasoning: false,
+	input: ["text"],
+	cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+	contextWindow: 200000,
+	maxTokens: 8192,
+};
+
+const entries: SessionEntry[] = [
+	{
+		type: "message",
+		id: "branch-user",
+		parentId: null,
+		timestamp: new Date(1).toISOString(),
+		message: { role: "user", content: "Abandoned request", timestamp: 1 },
+	},
+];
+
+function response(content: AssistantMessage["content"]): AssistantMessage {
+	return {
+		...fauxAssistantMessage(""),
+		content,
+		api: model.api,
+		provider: model.provider,
+		model: model.id,
+	};
+}
+
+describe("branch summarization", () => {
+	it("does not override tool choice for branch summaries", async () => {
+		let requestOptions: SimpleStreamOptions | undefined;
+		const streamFn: StreamFn = (_model, _context, options) => {
+			requestOptions = options;
+			const stream = createAssistantMessageEventStream();
+			queueMicrotask(() =>
+				stream.push({ type: "done", reason: "stop", message: response([{ type: "text", text: "summary" }]) }),
+			);
+			return stream;
+		};
+
+		await generateBranchSummary(entries, {
+			model,
+			signal: new AbortController().signal,
+			streamFn,
+		});
+
+		expect(requestOptions?.maxTokens).toBe(4096);
+		expect(requestOptions?.toolChoice).toBeUndefined();
+	});
+
+	it("clamps the branch summary output cap to the model limit", async () => {
+		let requestOptions: SimpleStreamOptions | undefined;
+		const streamFn: StreamFn = (_model, _context, options) => {
+			requestOptions = options;
+			const stream = createAssistantMessageEventStream();
+			queueMicrotask(() =>
+				stream.push({ type: "done", reason: "stop", message: response([{ type: "text", text: "summary" }]) }),
+			);
+			return stream;
+		};
+
+		await generateBranchSummary(entries, {
+			model: { ...model, maxTokens: 1024 },
+			signal: new AbortController().signal,
+			streamFn,
+		});
+
+		expect(requestOptions?.maxTokens).toBe(1024);
+	});
+
+	it("rejects tool calls from branch summaries", async () => {
+		const streamFn: StreamFn = () => {
+			const stream = createAssistantMessageEventStream();
+			queueMicrotask(() =>
+				stream.push({
+					type: "done",
+					reason: "toolUse",
+					message: response([
+						{ type: "toolCall", id: "tool-call-1", name: "read", arguments: { path: "README.md" } },
+					]),
+				}),
+			);
+			return stream;
+		};
+
+		const result = await generateBranchSummary(entries, {
+			model,
+			signal: new AbortController().signal,
+			streamFn,
+		});
+
+		expect(result.error).toBe("Branch summarization attempted to call a tool");
+	});
+
+	it("rejects length-limited branch summaries", async () => {
+		const streamFn: StreamFn = () => {
+			const stream = createAssistantMessageEventStream();
+			queueMicrotask(() =>
+				stream.push({
+					type: "done",
+					reason: "length",
+					message: { ...response([{ type: "text", text: "partial" }]), stopReason: "length" },
+				}),
+			);
+			return stream;
+		};
+
+		const result = await generateBranchSummary(entries, {
+			model,
+			signal: new AbortController().signal,
+			streamFn,
+		});
+
+		expect(result.error).toBe(
+			"Branch summarization failed: generation hit the token cap and the summary is incomplete",
+		);
+	});
+});