pi-usereq 0.7.0 → 0.11.0

package/pi.dev-src/pi-mono/coding-agent/test/suite/agent-session-bash-persistence.test.ts

@@ -0,0 +1,242 @@
+import { Buffer } from "node:buffer";
+import type { AgentTool } from "@mariozechner/pi-agent-core";
+import { fauxAssistantMessage, fauxToolCall } from "@mariozechner/pi-ai";
+import { Type } from "typebox";
+import { afterEach, describe, expect, it } from "vitest";
+import type { BashOperations } from "../../src/core/tools/bash.js";
+import { createHarness, type Harness } from "./harness.js";
+
+function getEntryTypes(harness: Harness): string[] {
+	return harness.sessionManager.getEntries().map((entry) => entry.type);
+}
+
+describe("AgentSession bash and persistence characterization", () => {
+	const harnesses: Harness[] = [];
+
+	afterEach(() => {
+		while (harnesses.length > 0) {
+			harnesses.pop()?.cleanup();
+		}
+	});
+
+	it("records bash results immediately while idle", async () => {
+		const harness = await createHarness();
+		harnesses.push(harness);
+
+		harness.session.recordBashResult("echo hi", {
+			output: "hi",
+			exitCode: 0,
+			cancelled: false,
+			truncated: false,
+		});
+
+		expect(harness.session.hasPendingBashMessages).toBe(false);
+		expect(harness.session.messages[harness.session.messages.length - 1]?.role).toBe("bashExecution");
+		expect(getEntryTypes(harness)).toContain("message");
+	});
+
+	it("defers bash results while streaming and flushes them before the next prompt", async () => {
+		let releaseToolExecution: (() => void) | undefined;
+		const toolRelease = new Promise<void>((resolve) => {
+			releaseToolExecution = resolve;
+		});
+		const waitTool: AgentTool = {
+			name: "wait",
+			label: "Wait",
+			description: "Wait for release",
+			parameters: Type.Object({}),
+			execute: async () => {
+				await toolRelease;
+				return {
+					content: [{ type: "text", text: "released" }],
+					details: {},
+				};
+			},
+		};
+		const harness = await createHarness({ tools: [waitTool] });
+		harnesses.push(harness);
+		harness.setResponses([
+			fauxAssistantMessage([fauxToolCall("wait", {})], { stopReason: "toolUse" }),
+			fauxAssistantMessage("done"),
+			fauxAssistantMessage("after flush"),
+		]);
+
+		const sawToolStart = new Promise<void>((resolve) => {
+			const unsubscribe = harness.session.subscribe((event) => {
+				if (event.type === "tool_execution_start") {
+					unsubscribe();
+					resolve();
+				}
+			});
+		});
+
+		const firstPrompt = harness.session.prompt("start");
+		await sawToolStart;
+		harness.session.recordBashResult("echo hi", {
+			output: "hi",
+			exitCode: 0,
+			cancelled: false,
+			truncated: false,
+		});
+
+		expect(harness.session.hasPendingBashMessages).toBe(true);
+		expect(harness.session.messages.some((message) => message.role === "bashExecution")).toBe(false);
+
+		releaseToolExecution?.();
+		await firstPrompt;
+
+		expect(harness.session.hasPendingBashMessages).toBe(true);
+		expect(harness.session.messages.some((message) => message.role === "bashExecution")).toBe(false);
+
+		await harness.session.prompt("next turn");
+
+		expect(harness.session.hasPendingBashMessages).toBe(false);
+		expect(harness.session.messages.some((message) => message.role === "bashExecution")).toBe(true);
+		expect(getEntryTypes(harness).filter((type) => type === "message").length).toBeGreaterThan(0);
+	});
+
+	it("executes bash commands and records the result", async () => {
+		const harness = await createHarness();
+		harnesses.push(harness);
+
+		const result = await harness.session.executeBash("printf 'hello'");
+
+		expect(result.output).toContain("hello");
+		expect(harness.session.messages[harness.session.messages.length - 1]?.role).toBe("bashExecution");
+	});
+
+	it("cancels running bash commands with abortBash", async () => {
+		const harness = await createHarness();
+		harnesses.push(harness);
+		const operations: BashOperations = {
+			exec: async (_command, _cwd, options) => {
+				return await new Promise<{ exitCode: number | null }>((_resolve, reject) => {
+					options.signal?.addEventListener(
+						"abort",
+						() => {
+							reject(new Error("aborted"));
+						},
+						{ once: true },
+					);
+				});
+			},
+		};
+
+		const bashPromise = harness.session.executeBash("sleep", undefined, { operations });
+		await new Promise((resolve) => setTimeout(resolve, 0));
+		expect(harness.session.isBashRunning).toBe(true);
+		harness.session.abortBash();
+
+		const result = await bashPromise;
+		expect(result.cancelled).toBe(true);
+		expect(harness.session.isBashRunning).toBe(false);
+	});
+
+	it("persists user, assistant, toolResult, and custom messages in order", async () => {
+		const echoTool: AgentTool = {
+			name: "echo",
+			label: "Echo",
+			description: "Echo text back",
+			parameters: Type.Object({ text: Type.String() }),
+			execute: async (_toolCallId, params) => {
+				const text = typeof params === "object" && params !== null && "text" in params ? String(params.text) : "";
+				return { content: [{ type: "text", text: `echo:${text}` }], details: { text } };
+			},
+		};
+		const harness = await createHarness({ tools: [echoTool] });
+		harnesses.push(harness);
+		harness.setResponses([
+			fauxAssistantMessage([fauxToolCall("echo", { text: "hello" })], { stopReason: "toolUse" }),
+			fauxAssistantMessage("done"),
+		]);
+
+		await harness.session.sendCustomMessage({
+			customType: "note",
+			content: "hello",
+			display: true,
+			details: { a: 1 },
+		});
+		await harness.session.prompt("start");
+
+		const entries = harness.sessionManager.getEntries();
+		expect(entries.map((entry) => entry.type)).toEqual([
+			"custom_message",
+			"message",
+			"message",
+			"message",
+			"message",
+		]);
+		expect(harness.session.messages.map((message) => message.role)).toEqual([
+			"custom",
+			"user",
+			"assistant",
+			"toolResult",
+			"assistant",
+		]);
+	});
+
+	it("does not emit message_end for bash execution messages", async () => {
+		const harness = await createHarness();
+		harnesses.push(harness);
+		const messageEndRoles: string[] = [];
+		harness.session.subscribe((event) => {
+			if (event.type === "message_end") {
+				messageEndRoles.push(event.message.role);
+			}
+		});
+
+		harness.session.recordBashResult("echo hi", {
+			output: "hi",
+			exitCode: 0,
+			cancelled: false,
+			truncated: false,
+		});
+
+		expect(messageEndRoles).toEqual([]);
+	});
+
+	it("persists aborted assistant messages", async () => {
+		const harness = await createHarness();
+		harnesses.push(harness);
+		harness.setResponses([fauxAssistantMessage("x".repeat(20_000))]);
+
+		const sawMessageUpdate = new Promise<void>((resolve) => {
+			const unsubscribe = harness.session.subscribe((event) => {
+				if (event.type === "message_update") {
+					unsubscribe();
+					resolve();
+				}
+			});
+		});
+
+		const promptPromise = harness.session.prompt("hi");
+		await sawMessageUpdate;
+		await harness.session.abort();
+		await promptPromise;
+
+		const lastEntry = harness.sessionManager.getEntries()[harness.sessionManager.getEntries().length - 1];
+		expect(lastEntry?.type).toBe("message");
+		if (lastEntry?.type === "message") {
+			expect(lastEntry.message.role).toBe("assistant");
+			if (lastEntry.message.role === "assistant") {
+				expect(lastEntry.message.stopReason).toBe("aborted");
+			}
+		}
+	});
+
+	it("records bash output through custom operations", async () => {
+		const harness = await createHarness();
+		harnesses.push(harness);
+		const operations: BashOperations = {
+			exec: async (_command, _cwd, options) => {
+				options.onData(Buffer.from("hello from custom ops"));
+				return { exitCode: 0 };
+			},
+		};
+
+		const result = await harness.session.executeBash("custom", undefined, { operations });
+
+		expect(result.output).toContain("hello from custom ops");
+		expect(harness.session.messages[harness.session.messages.length - 1]?.role).toBe("bashExecution");
+	});
+});

package/pi.dev-src/pi-mono/coding-agent/test/suite/agent-session-compaction.test.ts

@@ -0,0 +1,343 @@
+import { type AssistantMessage, fauxAssistantMessage, type Model } from "@mariozechner/pi-ai";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { createHarness, type Harness } from "./harness.js";
+
+type SessionWithCompactionInternals = {
+	_checkCompaction: (assistantMessage: AssistantMessage, skipAbortedCheck?: boolean) => Promise<void>;
+	_runAutoCompaction: (reason: "overflow" | "threshold", willRetry: boolean) => Promise<void>;
+};
+
+function createUsage(totalTokens: number) {
+	return {
+		input: totalTokens,
+		output: 0,
+		cacheRead: 0,
+		cacheWrite: 0,
+		totalTokens,
+		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+	};
+}
+
+function createAssistant(
+	harness: Harness,
+	options: {
+		stopReason?: AssistantMessage["stopReason"];
+		errorMessage?: string;
+		totalTokens?: number;
+		timestamp?: number;
+	},
+): AssistantMessage {
+	const model = harness.getModel();
+	return {
+		...fauxAssistantMessage("", {
+			stopReason: options.stopReason,
+			errorMessage: options.errorMessage,
+			timestamp: options.timestamp,
+		}),
+		api: model.api,
+		provider: model.provider,
+		model: model.id,
+		usage: createUsage(options.totalTokens ?? 0),
+	};
+}
+
+describe("AgentSession compaction characterization", () => {
+	const harnesses: Harness[] = [];
+
+	afterEach(() => {
+		vi.useRealTimers();
+		vi.restoreAllMocks();
+		while (harnesses.length > 0) {
+			harnesses.pop()?.cleanup();
+		}
+	});
+
+	it("manually compacts using an extension-provided summary", async () => {
+		const harness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.on("session_before_compact", async (event) => ({
+						compaction: {
+							summary: "summary from extension",
+							firstKeptEntryId: event.preparation.firstKeptEntryId,
+							tokensBefore: event.preparation.tokensBefore,
+							details: { source: "extension" },
+						},
+					}));
+				},
+			],
+		});
+		harnesses.push(harness);
+
+		await harness.session.prompt("one");
+		await harness.session.prompt("two");
+
+		const result = await harness.session.compact();
+		const compactionEntries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction");
+
+		expect(result.summary).toBe("summary from extension");
+		expect(compactionEntries).toHaveLength(1);
+		expect(harness.session.messages[0]?.role).toBe("compactionSummary");
+	});
+
+	it("throws when compacting without a model", async () => {
+		const harness = await createHarness();
+		harnesses.push(harness);
+		harness.session.agent.state.model = undefined as unknown as Model<any>;
+
+		await expect(harness.session.compact()).rejects.toThrow("No model selected");
+	});
+
+	it("throws when compacting without configured auth", async () => {
+		const harness = await createHarness({ withConfiguredAuth: false });
+		harnesses.push(harness);
+
+		await expect(harness.session.compact()).rejects.toThrow(`No API key found for ${harness.getModel().provider}.`);
+	});
+
+	it("cancels in-progress manual compaction when abortCompaction is called", async () => {
+		const harness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.on("session_before_compact", async (event) => {
+						return await new Promise<{ cancel: true }>((resolve) => {
+							event.signal.addEventListener("abort", () => resolve({ cancel: true }), { once: true });
+						});
+					});
+				},
+			],
+		});
+		harnesses.push(harness);
+
+		await harness.session.prompt("one");
+		await harness.session.prompt("two");
+
+		const compactPromise = harness.session.compact();
+		await new Promise((resolve) => setTimeout(resolve, 0));
+		harness.session.abortCompaction();
+
+		await expect(compactPromise).rejects.toThrow("Compaction cancelled");
+	});
+
+	it("resumes after threshold compaction when only agent-level queued messages exist", async () => {
+		vi.useFakeTimers();
+		const harness = await createHarness({
+			settings: { compaction: { keepRecentTokens: 1 } },
+			extensionFactories: [
+				(pi) => {
+					pi.on("session_before_compact", async (event) => ({
+						compaction: {
+							summary: "auto compacted",
+							firstKeptEntryId: event.preparation.firstKeptEntryId,
+							tokensBefore: event.preparation.tokensBefore,
+							details: {},
+						},
+					}));
+				},
+			],
+		});
+		harnesses.push(harness);
+		harness.setResponses([fauxAssistantMessage("one"), fauxAssistantMessage("two")]);
+		await harness.session.prompt("first");
+		await harness.session.prompt("second");
+
+		harness.session.agent.followUp({
+			role: "custom",
+			customType: "test",
+			content: [{ type: "text", text: "queued custom" }],
+			display: false,
+			timestamp: Date.now(),
+		});
+
+		const continueSpy = vi.spyOn(harness.session.agent, "continue").mockResolvedValue();
+		const sessionInternals = harness.session as unknown as SessionWithCompactionInternals;
+
+		await sessionInternals._runAutoCompaction("threshold", false);
+		await vi.advanceTimersByTimeAsync(100);
+
+		expect(continueSpy).toHaveBeenCalledTimes(1);
+	});
+
+	it("does not retry overflow recovery more than once", async () => {
+		const harness = await createHarness();
+		harnesses.push(harness);
+		const sessionInternals = harness.session as unknown as SessionWithCompactionInternals;
+		const overflowMessage = createAssistant(harness, {
+			stopReason: "error",
+			errorMessage: "prompt is too long",
+			timestamp: Date.now(),
+		});
+		const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue();
+		const compactionErrors: string[] = [];
+		harness.session.subscribe((event) => {
+			if (event.type === "compaction_end" && event.errorMessage) {
+				compactionErrors.push(event.errorMessage);
+			}
+		});
+
+		await sessionInternals._checkCompaction(overflowMessage);
+		await sessionInternals._checkCompaction({ ...overflowMessage, timestamp: Date.now() + 1 });
+
+		expect(runAutoCompactionSpy).toHaveBeenCalledTimes(1);
+		expect(compactionErrors).toContain(
+			"Context overflow recovery failed after one compact-and-retry attempt. Try reducing context or switching to a larger-context model.",
+		);
+	});
+
+	it("ignores stale pre-compaction assistant usage on pre-prompt checks", async () => {
+		const harness = await createHarness();
+		harnesses.push(harness);
+		const sessionInternals = harness.session as unknown as SessionWithCompactionInternals;
+		const staleTimestamp = Date.now() - 10_000;
+		const staleAssistant = createAssistant(harness, {
+			stopReason: "stop",
+			totalTokens: 610_000,
+			timestamp: staleTimestamp,
+		});
+
+		harness.sessionManager.appendMessage({
+			role: "user",
+			content: [{ type: "text", text: "before compaction" }],
+			timestamp: staleTimestamp - 1000,
+		});
+		harness.sessionManager.appendMessage(staleAssistant);
+		const firstKeptEntryId = harness.sessionManager.getEntries()[0]!.id;
+		harness.sessionManager.appendCompaction(
+			"summary",
+			firstKeptEntryId,
+			staleAssistant.usage.totalTokens,
+			undefined,
+			false,
+		);
+		harness.sessionManager.appendMessage({
+			role: "user",
+			content: [{ type: "text", text: "after compaction" }],
+			timestamp: Date.now(),
+		});
+
+		const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue();
+
+		await sessionInternals._checkCompaction(staleAssistant, false);
+
+		expect(runAutoCompactionSpy).not.toHaveBeenCalled();
+	});
+
+	it("triggers threshold compaction for error messages using the last successful usage", async () => {
+		const harness = await createHarness();
+		harnesses.push(harness);
+		const sessionInternals = harness.session as unknown as SessionWithCompactionInternals;
+		const successfulAssistant = createAssistant(harness, {
+			stopReason: "stop",
+			totalTokens: 190_000,
+			timestamp: Date.now(),
+		});
+		const errorAssistant = createAssistant(harness, {
+			stopReason: "error",
+			errorMessage: "529 overloaded",
+			timestamp: Date.now() + 1000,
+		});
+		harness.session.agent.state.messages = [
+			{ role: "user", content: [{ type: "text", text: "hello" }], timestamp: Date.now() - 1000 },
+			successfulAssistant,
+			{ role: "user", content: [{ type: "text", text: "retry" }], timestamp: Date.now() + 500 },
+			errorAssistant,
+		];
+
+		const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue();
+
+		await sessionInternals._checkCompaction(errorAssistant);
+
+		expect(runAutoCompactionSpy).toHaveBeenCalledWith("threshold", false);
+	});
+
+	it("does not trigger threshold compaction for error messages when no prior usage exists", async () => {
+		const harness = await createHarness();
+		harnesses.push(harness);
+		const sessionInternals = harness.session as unknown as SessionWithCompactionInternals;
+		const errorAssistant = createAssistant(harness, {
+			stopReason: "error",
+			errorMessage: "529 overloaded",
+			timestamp: Date.now(),
+		});
+		harness.session.agent.state.messages = [
+			{ role: "user", content: [{ type: "text", text: "hello" }], timestamp: Date.now() - 1000 },
+			errorAssistant,
+		];
+
+		const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue();
+
+		await sessionInternals._checkCompaction(errorAssistant);
+
+		expect(runAutoCompactionSpy).not.toHaveBeenCalled();
+	});
+
+	it("does not trigger threshold compaction when only kept pre-compaction usage exists", async () => {
+		const harness = await createHarness();
+		harnesses.push(harness);
+		const sessionInternals = harness.session as unknown as SessionWithCompactionInternals;
+		const preCompactionTimestamp = Date.now() - 10_000;
+		const keptAssistant = createAssistant(harness, {
+			stopReason: "stop",
+			totalTokens: 190_000,
+			timestamp: preCompactionTimestamp,
+		});
+
+		harness.sessionManager.appendMessage({
+			role: "user",
+			content: [{ type: "text", text: "before compaction" }],
+			timestamp: preCompactionTimestamp - 1000,
+		});
+		harness.sessionManager.appendMessage(keptAssistant);
+		const firstKeptEntryId = harness.sessionManager.getEntries()[0]!.id;
+		harness.sessionManager.appendCompaction(
+			"summary",
+			firstKeptEntryId,
+			keptAssistant.usage.totalTokens,
+			undefined,
+			false,
+		);
+
+		const errorAssistant = createAssistant(harness, {
+			stopReason: "error",
+			errorMessage: "529 overloaded",
+			timestamp: Date.now(),
+		});
+		harness.session.agent.state.messages = [
+			{ role: "user", content: [{ type: "text", text: "kept user" }], timestamp: preCompactionTimestamp - 1000 },
+			keptAssistant,
+			{ role: "user", content: [{ type: "text", text: "new prompt" }], timestamp: Date.now() - 500 },
+			errorAssistant,
+		];
+
+		const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue();
+
+		await sessionInternals._checkCompaction(errorAssistant);
+
+		expect(runAutoCompactionSpy).not.toHaveBeenCalled();
+	});
+
+	it("does not trigger threshold compaction below the threshold or when disabled", async () => {
+		const belowThresholdHarness = await createHarness({
+			settings: { compaction: { enabled: true, reserveTokens: 1000 } },
+			models: [{ id: "faux-1", contextWindow: 200_000 }],
+		});
+		harnesses.push(belowThresholdHarness);
+		const disabledHarness = await createHarness({ settings: { compaction: { enabled: false } } });
+		harnesses.push(disabledHarness);
+
+		const belowThresholdInternals = belowThresholdHarness.session as unknown as SessionWithCompactionInternals;
+		const disabledInternals = disabledHarness.session as unknown as SessionWithCompactionInternals;
+		const belowThresholdSpy = vi.spyOn(belowThresholdInternals, "_runAutoCompaction").mockResolvedValue();
+		const disabledSpy = vi.spyOn(disabledInternals, "_runAutoCompaction").mockResolvedValue();
+
+		await belowThresholdInternals._checkCompaction(
+			createAssistant(belowThresholdHarness, { stopReason: "stop", totalTokens: 1_000, timestamp: Date.now() }),
+		);
+		await disabledInternals._checkCompaction(
+			createAssistant(disabledHarness, { stopReason: "stop", totalTokens: 1_000_000, timestamp: Date.now() }),
+		);
+
+		expect(belowThresholdSpy).not.toHaveBeenCalled();
+		expect(disabledSpy).not.toHaveBeenCalled();
+	});
+});