test(coding-agent): add session lifecycle characterization suite

packages/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 "@sinclair/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");
+	});
+});

packages/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();
+	});
+});

packages/coding-agent/test/suite/agent-session-model-extension.test.ts

@@ -0,0 +1,325 @@
+import type { AgentTool, ThinkingLevel } from "@mariozechner/pi-agent-core";
+import { fauxAssistantMessage, fauxToolCall, type Model } from "@mariozechner/pi-ai";
+import { Type } from "@sinclair/typebox";
+import { afterEach, describe, expect, it } from "vitest";
+import type { ExtensionAPI } from "../../src/index.js";
+import { createHarness, getAssistantTexts, type Harness } from "./harness.js";
+
+describe("AgentSession model and extension characterization", () => {
+	const harnesses: Harness[] = [];
+
+	afterEach(() => {
+		while (harnesses.length > 0) {
+			harnesses.pop()?.cleanup();
+		}
+	});
+
+	it("setModel saves the model and emits model_select", async () => {
+		const modelEvents: string[] = [];
+		const harness = await createHarness({
+			models: [
+				{ id: "faux-1", name: "One", reasoning: true },
+				{ id: "faux-2", name: "Two", reasoning: true },
+			],
+			extensionFactories: [
+				(pi) => {
+					pi.on("model_select", async (event) => {
+						modelEvents.push(`${event.previousModel?.id ?? "none"}->${event.model.id}:${event.source}`);
+					});
+				},
+			],
+		});
+		harnesses.push(harness);
+		const nextModel = harness.getModel("faux-2")!;
+
+		await harness.session.setModel(nextModel);
+
+		expect(harness.session.model?.id).toBe("faux-2");
+		expect(modelEvents).toEqual(["faux-1->faux-2:set"]);
+		expect(
+			harness.sessionManager
+				.getEntries()
+				.filter((entry) => entry.type === "model_change")
+				.map((entry) => `${entry.provider}/${entry.modelId}`),
+		).toEqual([`${nextModel.provider}/${nextModel.id}`]);
+	});
+
+	it("cycles through scoped models and preserves the scoped thinking preference", async () => {
+		const harness = await createHarness({
+			models: [
+				{ id: "faux-1", name: "One", reasoning: true },
+				{ id: "faux-2", name: "Two", reasoning: false },
+			],
+		});
+		harnesses.push(harness);
+		const modelOne = harness.getModel("faux-1")!;
+		const modelTwo = harness.getModel("faux-2")!;
+		harness.session.setScopedModels([{ model: modelOne, thinkingLevel: "high" }, { model: modelTwo }] as Array<{
+			model: Model<string>;
+			thinkingLevel?: ThinkingLevel;
+		}>);
+		harness.session.setThinkingLevel("high");
+
+		await harness.session.cycleModel();
+		expect(harness.session.model?.id).toBe("faux-2");
+		expect(harness.session.thinkingLevel).toBe("off");
+
+		await harness.session.cycleModel();
+		expect(harness.session.model?.id).toBe("faux-1");
+		expect(harness.session.thinkingLevel).toBe("high");
+	});
+
+	it("clamps thinking levels to model capabilities and cycles available levels", async () => {
+		const harness = await createHarness({ models: [{ id: "faux-1", reasoning: false }] });
+		harnesses.push(harness);
+
+		harness.session.setThinkingLevel("high");
+		expect(harness.session.thinkingLevel).toBe("off");
+		expect(harness.session.cycleThinkingLevel()).toBeUndefined();
+	});
+
+	it("throws when setModel is called without configured auth", async () => {
+		const harness = await createHarness({
+			models: [
+				{ id: "faux-1", name: "One", reasoning: true },
+				{ id: "faux-2", name: "Two", reasoning: true },
+			],
+			withConfiguredAuth: false,
+		});
+		harnesses.push(harness);
+
+		await expect(harness.session.setModel(harness.getModel("faux-2")!)).rejects.toThrow(
+			`No API key for ${harness.getModel().provider}/faux-2`,
+		);
+	});
+
+	it("allows extension tool_call handlers to block tool execution", async () => {
+		const echoTool: AgentTool = {
+			name: "echo",
+			label: "Echo",
+			description: "Echo text back",
+			parameters: Type.Object({ text: Type.String() }),
+			execute: async () => {
+				throw new Error("tool should have been blocked");
+			},
+		};
+		const harness = await createHarness({
+			tools: [echoTool],
+			extensionFactories: [
+				(pi) => {
+					pi.on("tool_call", async () => ({ block: true, reason: "Blocked by test" }));
+				},
+			],
+		});
+		harnesses.push(harness);
+		harness.setResponses([
+			fauxAssistantMessage([fauxToolCall("echo", { text: "hello" })], { stopReason: "toolUse" }),
+			(context) => {
+				const toolResult = context.messages.find((message) => message.role === "toolResult");
+				const errorText =
+					toolResult?.role === "toolResult"
+						? toolResult.content
+								.filter((part): part is { type: "text"; text: string } => part.type === "text")
+								.map((part) => part.text)
+								.join("\n")
+						: "";
+				return fauxAssistantMessage(errorText);
+			},
+		]);
+
+		await harness.session.prompt("hi");
+
+		expect(getAssistantTexts(harness)).toContain("Blocked by test");
+		expect(
+			harness.session.messages.find((message) => message.role === "toolResult" && message.isError),
+		).toBeDefined();
+	});
+
+	it("allows extension tool_result handlers to modify tool results", 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 }], details: { text } };
+			},
+		};
+		const harness = await createHarness({
+			tools: [echoTool],
+			extensionFactories: [
+				(pi) => {
+					pi.on("tool_result", async () => ({
+						content: [{ type: "text", text: "patched result" }],
+						details: { patched: true },
+					}));
+				},
+			],
+		});
+		harnesses.push(harness);
+		harness.setResponses([
+			fauxAssistantMessage([fauxToolCall("echo", { text: "hello" })], { stopReason: "toolUse" }),
+			(context) => {
+				const toolResult = context.messages.find((message) => message.role === "toolResult");
+				const text =
+					toolResult?.role === "toolResult"
+						? toolResult.content
+								.filter((part): part is { type: "text"; text: string } => part.type === "text")
+								.map((part) => part.text)
+								.join("\n")
+						: "";
+				return fauxAssistantMessage(text);
+			},
+		]);
+
+		await harness.session.prompt("hi");
+
+		expect(getAssistantTexts(harness)).toContain("patched result");
+		expect(
+			harness.session.messages.find((message) => message.role === "toolResult" && message.details?.patched === true),
+		).toBeDefined();
+	});
+
+	it("allows extension context handlers to modify messages before the LLM call", async () => {
+		const harness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.on("context", async (event) => ({
+						messages: event.messages.map((message) =>
+							message.role === "user"
+								? { ...message, content: [{ type: "text", text: "rewritten" }], timestamp: message.timestamp }
+								: message,
+						),
+					}));
+				},
+			],
+		});
+		harnesses.push(harness);
+		let providerUserText = "";
+		harness.setResponses([
+			(context) => {
+				const user = context.messages.find((message) => message.role === "user");
+				providerUserText =
+					user && typeof user.content !== "string"
+						? user.content
+								.filter((part): part is { type: "text"; text: string } => part.type === "text")
+								.map((part) => part.text)
+								.join("\n")
+						: "";
+				return fauxAssistantMessage("done");
+			},
+		]);
+
+		await harness.session.prompt("original");
+
+		expect(providerUserText).toBe("rewritten");
+		const storedUserMessage = harness.session.messages.find((message) => message.role === "user");
+		expect(storedUserMessage?.role).toBe("user");
+		if (storedUserMessage?.role === "user") {
+			expect(storedUserMessage.content).toEqual([{ type: "text", text: "original" }]);
+		}
+	});
+
+	it("allows extension input handlers to transform or handle input", async () => {
+		let extensionApi: ExtensionAPI | undefined;
+		const transformedHarness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					extensionApi = pi;
+					pi.on("input", async (event) => {
+						if (event.text === "ping") {
+							return { action: "handled" };
+						}
+						return { action: "transform", text: `transformed:${event.text}` };
+					});
+				},
+			],
+		});
+		harnesses.push(transformedHarness);
+		let providerUserText = "";
+		transformedHarness.setResponses([
+			(context) => {
+				const user = context.messages.find((message) => message.role === "user");
+				providerUserText =
+					user && typeof user.content !== "string"
+						? user.content
+								.filter((part): part is { type: "text"; text: string } => part.type === "text")
+								.map((part) => part.text)
+								.join("\n")
+						: "";
+				return fauxAssistantMessage("done");
+			},
+		]);
+
+		await transformedHarness.session.prompt("hello");
+		await transformedHarness.session.prompt("ping");
+
+		expect(providerUserText).toBe("transformed:hello");
+		expect(transformedHarness.session.messages.filter((message) => message.role === "user")).toHaveLength(1);
+		expect(extensionApi).toBeDefined();
+	});
+
+	it("allows before_agent_start handlers to inject custom messages and modify the system prompt", async () => {
+		const harness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.on("before_agent_start", async (event) => ({
+						message: {
+							customType: "before-start",
+							content: "injected",
+							display: true,
+							details: { injected: true },
+						},
+						systemPrompt: `${event.systemPrompt}\n\nextra instructions`,
+					}));
+				},
+			],
+		});
+		harnesses.push(harness);
+		let providerSystemPrompt = "";
+		let sawInjectedUserMessage = false;
+		harness.setResponses([
+			(context) => {
+				providerSystemPrompt = context.systemPrompt ?? "";
+				sawInjectedUserMessage = context.messages.some(
+					(message) =>
+						message.role === "user" &&
+						typeof message.content !== "string" &&
+						message.content.some((part) => part.type === "text" && part.text === "injected"),
+				);
+				return fauxAssistantMessage("done");
+			},
+		]);
+
+		await harness.session.prompt("hello");
+
+		expect(providerSystemPrompt).toContain("extra instructions");
+		expect(sawInjectedUserMessage).toBe(true);
+		expect(
+			harness.session.messages.some((message) => message.role === "custom" && message.customType === "before-start"),
+		).toBe(true);
+	});
+
+	it("bindExtensions emits session_start and reload emits session_shutdown then session_start", async () => {
+		const lifecycleEvents: string[] = [];
+		const harness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.on("session_start", async (event) => {
+						lifecycleEvents.push(`start:${event.reason}`);
+					});
+					pi.on("session_shutdown", async () => {
+						lifecycleEvents.push("shutdown");
+					});
+				},
+			],
+		});
+		harnesses.push(harness);
+
+		await harness.session.bindExtensions({ shutdownHandler: () => {} });
+		await harness.session.reload();
+
+		expect(lifecycleEvents).toEqual(["start:startup", "shutdown", "start:reload"]);
+	});
+});

packages/coding-agent/test/suite/agent-session-prompt.test.ts

@@ -0,0 +1,323 @@
+import { mkdirSync, rmSync, writeFileSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import type { AgentTool } from "@mariozechner/pi-agent-core";
+import { fauxAssistantMessage, fauxToolCall, type Model } from "@mariozechner/pi-ai";
+import { Type } from "@sinclair/typebox";
+import { afterEach, describe, expect, it } from "vitest";
+import type { PromptTemplate } from "../../src/core/prompt-templates.js";
+import { createSyntheticSourceInfo } from "../../src/core/source-info.js";
+import { createTestResourceLoader } from "../utilities.js";
+import { createHarness, getMessageText, type Harness } from "./harness.js";
+
+describe("AgentSession prompt characterization", () => {
+	const harnesses: Harness[] = [];
+	const tempDirs: string[] = [];
+
+	afterEach(() => {
+		while (harnesses.length > 0) {
+			harnesses.pop()?.cleanup();
+		}
+		while (tempDirs.length > 0) {
+			const tempDir = tempDirs.pop();
+			if (tempDir) {
+				rmSync(tempDir, { recursive: true, force: true });
+			}
+		}
+	});
+
+	it("prompts while idle and records a single text response", async () => {
+		const harness = await createHarness();
+		harnesses.push(harness);
+
+		harness.setResponses([fauxAssistantMessage("hello")]);
+
+		await harness.session.prompt("hi");
+
+		expect(harness.session.messages.map((message) => message.role)).toEqual(["user", "assistant"]);
+		expect(getMessageText(harness.session.messages[0]!)).toBe("hi");
+		expect(harness.getPendingResponseCount()).toBe(0);
+	});
+
+	it("handles a tool call turn and waits for the follow-up LLM response", async () => {
+		const toolRuns: string[] = [];
+		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) : "";
+				toolRuns.push(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.prompt("start");
+
+		expect(toolRuns).toEqual(["hello"]);
+		expect(harness.session.messages.map((message) => message.role)).toEqual([
+			"user",
+			"assistant",
+			"toolResult",
+			"assistant",
+		]);
+		expect(harness.session.messages[2]?.role).toBe("toolResult");
+		expect(harness.session.messages[3]?.role).toBe("assistant");
+	});
+
+	it("executes multiple tool calls from one response and continues with a single follow-up response", async () => {
+		const toolRuns: string[] = [];
+		const makeTool = (name: string, delayMs: number): AgentTool => ({
+			name,
+			label: name,
+			description: `${name} tool`,
+			parameters: Type.Object({ value: Type.String() }),
+			execute: async (_toolCallId, params) => {
+				const value =
+					typeof params === "object" && params !== null && "value" in params ? String(params.value) : "";
+				await new Promise((resolve) => setTimeout(resolve, delayMs));
+				toolRuns.push(`${name}:${value}`);
+				return {
+					content: [{ type: "text", text: `${name}:${value}` }],
+					details: { value },
+				};
+			},
+		});
+		const harness = await createHarness({ tools: [makeTool("slow", 25), makeTool("fast", 0)] });
+		harnesses.push(harness);
+
+		harness.setResponses([
+			fauxAssistantMessage([fauxToolCall("slow", { value: "a" }), fauxToolCall("fast", { value: "b" })], {
+				stopReason: "toolUse",
+			}),
+			(context) => {
+				const toolResults = context.messages.filter((message) => message.role === "toolResult");
+				return fauxAssistantMessage(`tool results: ${toolResults.length}`);
+			},
+		]);
+
+		await harness.session.prompt("run tools");
+
+		expect(toolRuns.sort()).toEqual(["fast:b", "slow:a"]);
+		expect(harness.session.messages.filter((message) => message.role === "toolResult")).toHaveLength(2);
+		expect(harness.session.messages[harness.session.messages.length - 1]?.role).toBe("assistant");
+	});
+
+	it("preserves image attachments in the provider context", async () => {
+		const harness = await createHarness();
+		harnesses.push(harness);
+		let sawImage = false;
+
+		harness.setResponses([
+			(context) => {
+				const user = context.messages.find((message) => message.role === "user");
+				sawImage =
+					user?.role === "user" &&
+					typeof user.content !== "string" &&
+					user.content.some((part) => part.type === "image");
+				return fauxAssistantMessage("ok");
+			},
+		]);
+
+		await harness.session.prompt("describe", {
+			images: [
+				{
+					type: "image",
+					mimeType: "image/png",
+					data: "ZmFrZQ==",
+				},
+			],
+		});
+
+		expect(sawImage).toBe(true);
+	});
+
+	it("expands skill commands before sending the prompt", async () => {
+		const tempDir = join(tmpdir(), `pi-skill-${Date.now()}-${Math.random().toString(36).slice(2)}`);
+		mkdirSync(tempDir, { recursive: true });
+		tempDirs.push(tempDir);
+		const skillPath = join(tempDir, "test-skill.md");
+		writeFileSync(skillPath, "# Test Skill\n\nUse the skill body.");
+
+		const resourceLoader = {
+			...createTestResourceLoader(),
+			getSkills: () => ({
+				skills: [
+					{
+						name: "test",
+						description: "Test skill",
+						filePath: skillPath,
+						disableModelInvocation: false,
+						baseDir: tempDir,
+						sourceInfo: createSyntheticSourceInfo(skillPath, {
+							source: "local",
+							scope: "project",
+							origin: "top-level",
+							baseDir: tempDir,
+						}),
+					},
+				],
+				diagnostics: [],
+			}),
+		};
+		const harness = await createHarness({ resourceLoader });
+		harnesses.push(harness);
+		let expandedPrompt = "";
+
+		harness.setResponses([
+			(context) => {
+				const user = context.messages.find((message) => message.role === "user");
+				expandedPrompt = user ? getMessageText(user) : "";
+				return fauxAssistantMessage("ok");
+			},
+		]);
+
+		await harness.session.prompt("/skill:test explain this");
+
+		expect(expandedPrompt).toContain('<skill name="test" location="');
+		expect(expandedPrompt).toContain("Use the skill body.");
+		expect(expandedPrompt).toContain("explain this");
+	});
+
+	it("expands prompt templates before sending the prompt", async () => {
+		const template: PromptTemplate = {
+			name: "review",
+			description: "Review template",
+			content: "Review this code: $1",
+			filePath: "/virtual/review.md",
+			sourceInfo: createSyntheticSourceInfo("/virtual/review.md", {
+				source: "local",
+				scope: "temporary",
+				origin: "top-level",
+			}),
+		};
+		const resourceLoader = {
+			...createTestResourceLoader(),
+			getPrompts: () => ({ prompts: [template], diagnostics: [] }),
+		};
+		const harness = await createHarness({ resourceLoader });
+		harnesses.push(harness);
+		let expandedPrompt = "";
+
+		harness.setResponses([
+			(context) => {
+				const user = context.messages.find((message) => message.role === "user");
+				expandedPrompt = user ? getMessageText(user) : "";
+				return fauxAssistantMessage("ok");
+			},
+		]);
+
+		await harness.session.prompt("/review src/index.ts");
+
+		expect(expandedPrompt).toBe("Review this code: src/index.ts");
+	});
+
+	it("dispatches extension commands without consuming a provider response", async () => {
+		const commandRuns: string[] = [];
+		const harness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.registerCommand("testcmd", {
+						description: "Test command",
+						handler: async (args) => {
+							commandRuns.push(args);
+						},
+					});
+				},
+			],
+		});
+		harnesses.push(harness);
+		harness.setResponses([fauxAssistantMessage("should stay queued")]);
+
+		await harness.session.prompt("/testcmd hello world");
+
+		expect(commandRuns).toEqual(["hello world"]);
+		expect(harness.session.messages).toEqual([]);
+		expect(harness.getPendingResponseCount()).toBe(1);
+	});
+
+	it("sendUserMessage while idle triggers a turn", async () => {
+		const harness = await createHarness();
+		harnesses.push(harness);
+
+		harness.setResponses([fauxAssistantMessage("response")]);
+
+		await harness.session.sendUserMessage("from extension");
+
+		expect(harness.session.messages.map((message) => message.role)).toEqual(["user", "assistant"]);
+		expect(getMessageText(harness.session.messages[0]!)).toBe("from extension");
+	});
+
+	it("throws when prompted during streaming without a streamingBehavior", 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"),
+		]);
+
+		const sawToolStart = new Promise<void>((resolve) => {
+			const unsubscribe = harness.session.subscribe((event) => {
+				if (event.type === "tool_execution_start") {
+					unsubscribe();
+					resolve();
+				}
+			});
+		});
+
+		const promptPromise = harness.session.prompt("start");
+		await sawToolStart;
+
+		await expect(harness.session.prompt("second")).rejects.toThrow(
+			"Agent is already processing. Specify streamingBehavior ('steer' or 'followUp') to queue the message.",
+		);
+
+		releaseToolExecution?.();
+		await promptPromise;
+	});
+
+	it("throws when prompting 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.prompt("hi")).rejects.toThrow("No model selected.");
+	});
+
+	it("throws when prompting without configured auth", async () => {
+		const harness = await createHarness({ withConfiguredAuth: false });
+		harnesses.push(harness);
+
+		await expect(harness.session.prompt("hi")).rejects.toThrow(
+			`No API key found for ${harness.getModel().provider}.`,
+		);
+	});
+});

packages/coding-agent/test/suite/agent-session-queue.test.ts

@@ -3,32 +3,58 @@ import { fauxAssistantMessage, fauxToolCall } from "@mariozechner/pi-ai";
 import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
 import { Type } from "@sinclair/typebox";
 import { afterEach, describe, expect, it } from "vitest";
-import { createHarness, type Harness } from "./harness.js";
+import { createHarness, getAssistantTexts, getMessageText, getUserTexts, type Harness } from "./harness.js";
 
-function getUserTexts(harness: Harness): string[] {
+async function createWaitingHarness(
-	return harness.session.messages
+	options: {
-		.filter((message) => message.role === "user")
+		tools?: AgentTool[];
-		.map((message) => {
+		extensionFactories?: Harness["session"]["extensionRunner"] extends never
-			const content = message.content;
+			? never
-			if (typeof content === "string") {
+			: Array<(pi: ExtensionAPI) => void>;
-				return content;
+	} = {},
+): Promise<{
+	harness: Harness;
+	releaseToolExecution: () => void;
+	promptPromise: Promise<void>;
+	waitForToolStart: Promise<void>;
+}> {
+	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, ...(options.tools ?? [])],
+		extensionFactories: options.extensionFactories,
+	});
+
+	const waitForToolStart = new Promise<void>((resolve) => {
+		const unsubscribe = harness.session.subscribe((event) => {
+			if (event.type === "tool_execution_start" && event.toolName === "wait") {
+				unsubscribe();
+				resolve();
 			}
-			return content
-				.filter((part): part is { type: "text"; text: string } => part.type === "text")
-				.map((part) => part.text)
-				.join("\n");
 		});
-}
+	});
 
-function getAssistantTexts(harness: Harness): string[] {
+	return {
-	return harness.session.messages
+		harness,
-		.filter((message) => message.role === "assistant")
+		releaseToolExecution: () => releaseToolExecution?.(),
-		.map((message) =>
+		promptPromise: harness.session.prompt("start"),
-			message.content
+		waitForToolStart,
-				.filter((part): part is { type: "text"; text: string } => part.type === "text")
+	};
-				.map((part) => part.text)
-				.join("\n"),
-		);
 }
 
 describe("AgentSession queue characterization", () => {
@@ -65,128 +91,332 @@ describe("AgentSession queue characterization", () => {
 
 	it("delivers extension-origin steering messages before the next LLM call", async () => {
 		let extensionApi: ExtensionAPI | undefined;
-		let releaseToolExecution: (() => void) | undefined;
+		const waiting = await createWaitingHarness({
-		const toolStarted = new Promise<void>((resolve) => {
-			releaseToolExecution = resolve;
-		});
-		const waitTool: AgentTool = {
-			name: "wait",
-			label: "Wait",
-			description: "Wait for the test to release execution",
-			parameters: Type.Object({}),
-			execute: async () => {
-				await toolStarted;
-				return {
-					content: [{ type: "text", text: "released" }],
-					details: {},
-				};
-			},
-		};
-		const harness = await createHarness({
-			tools: [waitTool],
 			extensionFactories: [
 				(pi) => {
 					extensionApi = pi;
 				},
 			],
 		});
+		const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting;
 		harnesses.push(harness);
 
 		harness.setResponses([
 			fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }),
 			(context) => {
 				const sawSteer = context.messages.some(
-					(message) =>
+					(message) => message.role === "user" && getMessageText(message) === "steer now",
-						message.role === "user" &&
-						typeof message.content !== "string" &&
-						message.content.some((part) => part.type === "text" && part.text === "steer now"),
 				);
 				return fauxAssistantMessage(sawSteer ? "saw steer" : "missing steer");
 			},
 		]);
 
-		const sawToolStart = new Promise<void>((resolve) => {
+		await waitForToolStart;
-			const unsubscribe = harness.session.subscribe((event) => {
-				if (event.type === "tool_execution_start" && event.toolName === "wait") {
-					unsubscribe();
-					resolve();
-				}
-			});
-		});
-
-		const promptPromise = harness.session.prompt("start");
-		await sawToolStart;
 		await new Promise((resolve) => setTimeout(resolve, 0));
 
 		extensionApi?.sendUserMessage("steer now", { deliverAs: "steer" });
-		releaseToolExecution?.();
+		releaseToolExecution();
 		await promptPromise;
 
 		expect(getUserTexts(harness)).toEqual(["start", "steer now"]);
 		expect(getAssistantTexts(harness)).toContain("saw steer");
 	});
 
-	it("treats queued slash-command follow-ups as raw user text instead of dispatching the command", async () => {
+	it("delivers follow-up messages only after the current run finishes", async () => {
-		let extensionApi: ExtensionAPI | undefined;
+		const waiting = await createWaitingHarness();
-		const commandRuns: string[] = [];
+		const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting;
-		let releaseToolExecution: (() => void) | undefined;
+		harnesses.push(harness);
-		const toolRelease = new Promise<void>((resolve) => {
+		const assistantSeenBeforeFollowUp: string[] = [];
-			releaseToolExecution = resolve;
+
-		});
+		harness.setResponses([
-		const waitTool: AgentTool = {
+			fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }),
-			name: "wait",
+			(context) => {
-			label: "Wait",
+				assistantSeenBeforeFollowUp.push(
-			description: "Wait for the test to release execution",
+					...context.messages
-			parameters: Type.Object({}),
+						.filter((message) => message.role === "assistant")
-			execute: async () => {
+						.map((message) =>
-				await toolRelease;
+							message.content
-				return {
+								.filter((part): part is { type: "text"; text: string } => part.type === "text")
-					content: [{ type: "text", text: "released" }],
+								.map((part) => part.text)
-					details: {},
+								.join("\n"),
-				};
+						),
+				);
+				return fauxAssistantMessage("follow-up response");
 			},
-		};
+		]);
+
+		await waitForToolStart;
+		await harness.session.followUp("after current run");
+		releaseToolExecution();
+		await promptPromise;
+
+		expect(getUserTexts(harness)).toEqual(["start", "after current run"]);
+		expect(assistantSeenBeforeFollowUp).toContain("");
+		expect(getAssistantTexts(harness)).toContain("follow-up response");
+	});
+
+	it("delivers multiple steering messages in order in one-at-a-time mode", async () => {
+		const waiting = await createWaitingHarness();
+		const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting;
+		harnesses.push(harness);
+
+		harness.setResponses([
+			fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }),
+			fauxAssistantMessage("handled steer 1"),
+			fauxAssistantMessage("handled steer 2"),
+		]);
+
+		await waitForToolStart;
+		await harness.session.steer("steer 1");
+		await harness.session.steer("steer 2");
+		releaseToolExecution();
+		await promptPromise;
+
+		expect(getUserTexts(harness)).toEqual(["start", "steer 1", "steer 2"]);
+		expect(getAssistantTexts(harness)).toEqual(["", "handled steer 1", "handled steer 2"]);
+	});
+
+	it("delivers multiple follow-up messages in order in one-at-a-time mode", async () => {
+		const waiting = await createWaitingHarness();
+		const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting;
+		harnesses.push(harness);
+
+		harness.setResponses([
+			fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }),
+			fauxAssistantMessage("original turn complete"),
+			fauxAssistantMessage("handled follow-up 1"),
+			fauxAssistantMessage("handled follow-up 2"),
+		]);
+
+		await waitForToolStart;
+		await harness.session.followUp("follow-up 1");
+		await harness.session.followUp("follow-up 2");
+		releaseToolExecution();
+		await promptPromise;
+
+		expect(getUserTexts(harness)).toEqual(["start", "follow-up 1", "follow-up 2"]);
+		expect(getAssistantTexts(harness)).toEqual([
+			"",
+			"original turn complete",
+			"handled follow-up 1",
+			"handled follow-up 2",
+		]);
+	});
+
+	it("delivers all steering messages in one batch in all mode", async () => {
+		const waiting = await createWaitingHarness();
+		const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting;
+		harnesses.push(harness);
+		harness.session.setSteeringMode("all");
+		let batchedUserMessages: string[] = [];
+
+		harness.setResponses([
+			fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }),
+			(context) => {
+				batchedUserMessages = context.messages
+					.filter((message) => message.role === "user")
+					.map((message) => getMessageText(message));
+				return fauxAssistantMessage("batched steer response");
+			},
+		]);
+
+		await waitForToolStart;
+		await harness.session.steer("steer 1");
+		await harness.session.steer("steer 2");
+		releaseToolExecution();
+		await promptPromise;
+
+		expect(batchedUserMessages).toEqual(["start", "steer 1", "steer 2"]);
+		expect(getAssistantTexts(harness)).toEqual(["", "batched steer response"]);
+	});
+
+	it("delivers all follow-up messages in one batch in all mode", async () => {
+		const waiting = await createWaitingHarness();
+		const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting;
+		harnesses.push(harness);
+		harness.session.setFollowUpMode("all");
+		let batchedUserMessages: string[] = [];
+
+		harness.setResponses([
+			fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }),
+			fauxAssistantMessage("original turn complete"),
+			(context) => {
+				batchedUserMessages = context.messages
+					.filter((message) => message.role === "user")
+					.map((message) => getMessageText(message));
+				return fauxAssistantMessage("batched follow-up response");
+			},
+		]);
+
+		await waitForToolStart;
+		await harness.session.followUp("follow-up 1");
+		await harness.session.followUp("follow-up 2");
+		releaseToolExecution();
+		await promptPromise;
+
+		expect(batchedUserMessages).toEqual(["start", "follow-up 1", "follow-up 2"]);
+		expect(getAssistantTexts(harness)).toEqual(["", "original turn complete", "batched follow-up response"]);
+	});
+
+	it("queues custom messages with deliverAs steer while streaming", async () => {
+		const waiting = await createWaitingHarness();
+		const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting;
+		harnesses.push(harness);
+		let sawCustomMessage = false;
+
+		harness.setResponses([
+			fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }),
+			(context) => {
+				sawCustomMessage = context.messages.some(
+					(message) =>
+						message.role === "user" &&
+						typeof message.content !== "string" &&
+						message.content.some((part) => part.type === "text" && part.text === "steer custom"),
+				);
+				return fauxAssistantMessage("done");
+			},
+		]);
+
+		await waitForToolStart;
+		await harness.session.sendCustomMessage(
+			{ customType: "queue-test", content: "steer custom", display: true, details: { value: 1 } },
+			{ deliverAs: "steer" },
+		);
+		releaseToolExecution();
+		await promptPromise;
+
+		expect(sawCustomMessage).toBe(true);
+		expect(
+			harness.session.messages.some((message) => message.role === "custom" && message.customType === "queue-test"),
+		).toBe(true);
+	});
+
+	it("queues custom messages with deliverAs followUp while streaming", async () => {
+		const waiting = await createWaitingHarness();
+		const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting;
+		harnesses.push(harness);
+		let sawCustomMessage = false;
+
+		harness.setResponses([
+			fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }),
+			fauxAssistantMessage("original turn complete"),
+			(context) => {
+				sawCustomMessage = context.messages.some(
+					(message) =>
+						message.role === "user" &&
+						typeof message.content !== "string" &&
+						message.content.some((part) => part.type === "text" && part.text === "follow-up custom"),
+				);
+				return fauxAssistantMessage("done");
+			},
+		]);
+
+		await waitForToolStart;
+		await harness.session.sendCustomMessage(
+			{ customType: "queue-test", content: "follow-up custom", display: true, details: { value: 1 } },
+			{ deliverAs: "followUp" },
+		);
+		releaseToolExecution();
+		await promptPromise;
+
+		expect(sawCustomMessage).toBe(true);
+		expect(
+			harness.session.messages.some((message) => message.role === "custom" && message.customType === "queue-test"),
+		).toBe(true);
+	});
+
+	it("injects nextTurn custom messages into the next prompt", async () => {
+		const harness = await createHarness();
+		harnesses.push(harness);
+		let sawCustomMessage = false;
+
+		await harness.session.sendCustomMessage(
+			{ customType: "next-turn", content: "carry this", display: true, details: {} },
+			{ deliverAs: "nextTurn" },
+		);
+
+		harness.setResponses([
+			(context) => {
+				sawCustomMessage = context.messages.some(
+					(message) =>
+						message.role === "user" &&
+						typeof message.content !== "string" &&
+						message.content.some((part) => part.type === "text" && part.text === "carry this"),
+				);
+				return fauxAssistantMessage("done");
+			},
+		]);
+
+		await harness.session.prompt("normal prompt");
+
+		expect(sawCustomMessage).toBe(true);
+		expect(harness.session.messages.map((message) => message.role)).toEqual(["user", "custom", "assistant"]);
+	});
+
+	it("updates pendingMessageCount and removes queued text before message_start is emitted", async () => {
+		const waiting = await createWaitingHarness();
+		const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting;
+		harnesses.push(harness);
+		const countsAtQueuedMessageStart: number[] = [];
+
+		harness.setResponses([
+			fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }),
+			fauxAssistantMessage("done"),
+		]);
+
+		harness.session.subscribe((event) => {
+			if (
+				event.type === "message_start" &&
+				event.message.role === "user" &&
+				getMessageText(event.message) === "queued"
+			) {
+				countsAtQueuedMessageStart.push(harness.session.pendingMessageCount);
+			}
+		});
+
+		await waitForToolStart;
+		await harness.session.steer("queued");
+		expect(harness.session.pendingMessageCount).toBe(1);
+		releaseToolExecution();
+		await promptPromise;
+
+		expect(countsAtQueuedMessageStart).toEqual([0]);
+		expect(harness.session.pendingMessageCount).toBe(0);
+	});
+
+	it("throws when queueing an extension command with steer", async () => {
 		const harness = await createHarness({
-			tools: [waitTool],
 			extensionFactories: [
 				(pi) => {
-					extensionApi = pi;
 					pi.registerCommand("testcmd", {
 						description: "Test command",
-						handler: async (args) => {
+						handler: async () => {},
-							commandRuns.push(args);
-						},
 					});
 				},
 			],
 		});
 		harnesses.push(harness);
 
-		harness.setResponses([
+		await expect(harness.session.steer("/testcmd queued")).rejects.toThrow(
-			fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }),
+			'Extension command "/testcmd" cannot be queued. Use prompt() or execute the command when not streaming.',
-			fauxAssistantMessage("first turn complete"),
+		);
-			fauxAssistantMessage("queued follow-up handled by model"),
+	});
-		]);
 
-		const sawToolStart = new Promise<void>((resolve) => {
+	it("throws when queueing an extension command with followUp", async () => {
-			const unsubscribe = harness.session.subscribe((event) => {
+		const harness = await createHarness({
-				if (event.type === "tool_execution_start" && event.toolName === "wait") {
+			extensionFactories: [
-					unsubscribe();
+				(pi) => {
-					resolve();
+					pi.registerCommand("testcmd", {
-				}
+						description: "Test command",
-			});
+						handler: async () => {},
+					});
+				},
+			],
 		});
+		harnesses.push(harness);
 
-		const promptPromise = harness.session.prompt("start");
+		await expect(harness.session.followUp("/testcmd queued")).rejects.toThrow(
-		await sawToolStart;
+			'Extension command "/testcmd" cannot be queued. Use prompt() or execute the command when not streaming.',
-		await new Promise((resolve) => setTimeout(resolve, 0));
+		);
-
-		extensionApi?.sendUserMessage("/testcmd queued", { deliverAs: "followUp" });
-		releaseToolExecution?.();
-		await promptPromise;
-
-		expect(commandRuns).toEqual([]);
-		expect(getUserTexts(harness)).toEqual(["start", "/testcmd queued"]);
-		expect(getAssistantTexts(harness)).toContain("queued follow-up handled by model");
 	});
 });

packages/coding-agent/test/suite/agent-session-retry-events.test.ts

@@ -0,0 +1,358 @@
+import type { AgentTool } from "@mariozechner/pi-agent-core";
+import { fauxAssistantMessage, fauxThinking, fauxToolCall } from "@mariozechner/pi-ai";
+import { Type } from "@sinclair/typebox";
+import { afterEach, describe, expect, it } from "vitest";
+import { createHarness, type Harness } from "./harness.js";
+
+function normalizeEventOrder(events: Harness["events"]): string[] {
+	const normalized: string[] = [];
+	for (const event of events) {
+		const label =
+			event.type === "message_start" || event.type === "message_end"
+				? `${event.type}:${event.message.role}`
+				: event.type === "tool_execution_start" || event.type === "tool_execution_end"
+					? `${event.type}:${event.toolName}`
+					: event.type;
+		if (label === "message_update" && normalized[normalized.length - 1] === "message_update") {
+			continue;
+		}
+		normalized.push(label);
+	}
+	return normalized;
+}
+
+describe("AgentSession retry and event characterization", () => {
+	const harnesses: Harness[] = [];
+
+	afterEach(() => {
+		while (harnesses.length > 0) {
+			harnesses.pop()?.cleanup();
+		}
+	});
+
+	it("retries after a transient error and succeeds", async () => {
+		const harness = await createHarness({ settings: { retry: { enabled: true, maxRetries: 3, baseDelayMs: 1 } } });
+		harnesses.push(harness);
+		const retryEvents: string[] = [];
+		harness.session.subscribe((event) => {
+			if (event.type === "auto_retry_start") retryEvents.push(`start:${event.attempt}`);
+			if (event.type === "auto_retry_end") retryEvents.push(`end:${event.success}`);
+		});
+
+		harness.setResponses([
+			fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }),
+			fauxAssistantMessage("recovered"),
+		]);
+
+		await harness.session.prompt("test");
+
+		expect(retryEvents).toEqual(["start:1", "end:true"]);
+		expect(harness.faux.state.callCount).toBe(2);
+		expect(harness.session.isRetrying).toBe(false);
+	});
+
+	it("retries multiple transient failures and succeeds on the final attempt", async () => {
+		const harness = await createHarness({ settings: { retry: { enabled: true, maxRetries: 3, baseDelayMs: 1 } } });
+		harnesses.push(harness);
+		const retryEvents: string[] = [];
+		harness.session.subscribe((event) => {
+			if (event.type === "auto_retry_start") retryEvents.push(`start:${event.attempt}`);
+			if (event.type === "auto_retry_end") retryEvents.push(`end:${event.success}`);
+		});
+
+		harness.setResponses([
+			fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }),
+			fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }),
+			fauxAssistantMessage("success"),
+		]);
+
+		await harness.session.prompt("test");
+
+		expect(retryEvents).toEqual(["start:1", "start:2", "end:true"]);
+		expect(harness.faux.state.callCount).toBe(3);
+	});
+
+	it("exhausts max retries and emits a failure event", async () => {
+		const harness = await createHarness({ settings: { retry: { enabled: true, maxRetries: 2, baseDelayMs: 1 } } });
+		harnesses.push(harness);
+		const retryEvents: string[] = [];
+		harness.session.subscribe((event) => {
+			if (event.type === "auto_retry_start") retryEvents.push(`start:${event.attempt}`);
+			if (event.type === "auto_retry_end") retryEvents.push(`end:${event.success}`);
+		});
+
+		harness.setResponses([
+			fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }),
+			fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }),
+			fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }),
+		]);
+
+		await harness.session.prompt("test");
+
+		expect(retryEvents).toEqual(["start:1", "start:2", "end:false"]);
+		expect(harness.faux.state.callCount).toBe(3);
+		expect(harness.session.isRetrying).toBe(false);
+	});
+
+	it("prompt waits for retry completion even when assistant message_end handling is delayed", async () => {
+		const harness = await createHarness({
+			settings: { retry: { enabled: true, maxRetries: 3, baseDelayMs: 1 } },
+			extensionFactories: [
+				(pi) => {
+					pi.on("message_end", async (event) => {
+						if (event.message.role === "assistant") {
+							await new Promise((resolve) => setTimeout(resolve, 40));
+						}
+					});
+				},
+			],
+		});
+		harnesses.push(harness);
+		harness.setResponses([
+			fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }),
+			fauxAssistantMessage("recovered"),
+		]);
+
+		await harness.session.prompt("test");
+
+		expect(harness.faux.state.callCount).toBe(2);
+		expect(harness.session.isRetrying).toBe(false);
+	});
+
+	it("does not retry when retry is disabled", async () => {
+		const harness = await createHarness({ settings: { retry: { enabled: false } } });
+		harnesses.push(harness);
+		harness.setResponses([fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" })]);
+
+		await harness.session.prompt("test");
+
+		expect(harness.faux.state.callCount).toBe(1);
+		expect(harness.eventsOfType("auto_retry_start")).toEqual([]);
+	});
+
+	it("does not retry non-retryable errors", async () => {
+		const harness = await createHarness({ settings: { retry: { enabled: true, maxRetries: 3, baseDelayMs: 1 } } });
+		harnesses.push(harness);
+		harness.setResponses([fauxAssistantMessage("", { stopReason: "error", errorMessage: "invalid_api_key" })]);
+
+		await harness.session.prompt("test");
+
+		expect(harness.faux.state.callCount).toBe(1);
+		expect(harness.eventsOfType("auto_retry_start")).toEqual([]);
+	});
+
+	it("cancels retry sleep when abortRetry is called", async () => {
+		const harness = await createHarness({ settings: { retry: { enabled: true, maxRetries: 3, baseDelayMs: 100 } } });
+		harnesses.push(harness);
+		harness.setResponses([fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" })]);
+
+		const sawRetryStart = new Promise<void>((resolve) => {
+			const unsubscribe = harness.session.subscribe((event) => {
+				if (event.type === "auto_retry_start") {
+					unsubscribe();
+					resolve();
+				}
+			});
+		});
+
+		const promptPromise = harness.session.prompt("test");
+		await sawRetryStart;
+		harness.session.abortRetry();
+		await promptPromise;
+
+		expect(harness.session.isRetrying).toBe(false);
+		expect(harness.eventsOfType("auto_retry_end").map((event) => event.finalError)).toContain("Retry cancelled");
+		expect(harness.faux.state.callCount).toBe(1);
+	});
+
+	it("waits for the full loop when retry recovery produces tool calls", async () => {
+		const toolRuns: string[] = [];
+		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) : "";
+				toolRuns.push(text);
+				return { content: [{ type: "text", text: `echo:${text}` }], details: { text } };
+			},
+		};
+		const harness = await createHarness({
+			tools: [echoTool],
+			settings: { retry: { enabled: true, maxRetries: 3, baseDelayMs: 1 } },
+		});
+		harnesses.push(harness);
+		harness.setResponses([
+			fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }),
+			fauxAssistantMessage([fauxToolCall("echo", { text: "hello" })], { stopReason: "toolUse" }),
+			fauxAssistantMessage("final answer"),
+		]);
+
+		await harness.session.prompt("test");
+
+		expect(harness.faux.state.callCount).toBe(3);
+		expect(toolRuns).toEqual(["hello"]);
+		expect(harness.session.isStreaming).toBe(false);
+		await harness.session.prompt("follow-up");
+		expect(harness.faux.state.callCount).toBe(4);
+	});
+
+	it("emits extension events before public event subscribers", async () => {
+		const order: string[] = [];
+		const harness = await createHarness({
+			extensionFactories: [
+				(pi) => {
+					pi.on("message_start", async (event) => {
+						order.push(`extension:${event.type}:${event.message.role}`);
+					});
+					pi.on("message_end", async (event) => {
+						order.push(`extension:${event.type}:${event.message.role}`);
+					});
+				},
+			],
+		});
+		harnesses.push(harness);
+		harness.session.subscribe((event) => {
+			if (event.type === "message_start" || event.type === "message_end") {
+				order.push(`public:${event.type}:${event.message.role}`);
+			}
+		});
+		harness.setResponses([fauxAssistantMessage("done")]);
+
+		await harness.session.prompt("hi");
+
+		expect(order).toEqual([
+			"extension:message_start:user",
+			"public:message_start:user",
+			"extension:message_end:user",
+			"public:message_end:user",
+			"extension:message_start:assistant",
+			"public:message_start:assistant",
+			"extension:message_end:assistant",
+			"public:message_end:assistant",
+		]);
+	});
+
+	it("emits the expected event order for a single prompt", async () => {
+		const harness = await createHarness();
+		harnesses.push(harness);
+		harness.setResponses([fauxAssistantMessage("hello")]);
+
+		await harness.session.prompt("hi");
+
+		expect(normalizeEventOrder(harness.events)).toEqual([
+			"agent_start",
+			"turn_start",
+			"message_start:user",
+			"message_end:user",
+			"message_start:assistant",
+			"message_update",
+			"message_end:assistant",
+			"turn_end",
+			"agent_end",
+		]);
+	});
+
+	it("emits the expected event order for a tool call turn", async () => {
+		const toolRuns: string[] = [];
+		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) : "";
+				toolRuns.push(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.prompt("hi");
+
+		expect(toolRuns).toEqual(["hello"]);
+		expect(normalizeEventOrder(harness.events)).toEqual([
+			"agent_start",
+			"turn_start",
+			"message_start:user",
+			"message_end:user",
+			"message_start:assistant",
+			"message_update",
+			"message_end:assistant",
+			"tool_execution_start:echo",
+			"tool_execution_end:echo",
+			"message_start:toolResult",
+			"message_end:toolResult",
+			"turn_end",
+			"turn_start",
+			"message_start:assistant",
+			"message_update",
+			"message_end:assistant",
+			"turn_end",
+			"agent_end",
+		]);
+	});
+
+	it("emits streaming deltas for text, thinking, and tool calls in message_update events", async () => {
+		const harness = await createHarness();
+		harnesses.push(harness);
+		harness.setResponses([
+			fauxAssistantMessage(
+				[fauxThinking("plan"), { type: "text", text: "answer" }, fauxToolCall("echo", { text: "hello" })],
+				{
+					stopReason: "toolUse",
+				},
+			),
+		]);
+
+		await harness.session.prompt("hi").catch(() => {});
+
+		const updateTypes = harness.eventsOfType("message_update").map((event) => event.assistantMessageEvent.type);
+		expect(updateTypes).toContain("thinking_delta");
+		expect(updateTypes).toContain("text_delta");
+		expect(updateTypes).toContain("toolcall_delta");
+	});
+
+	it("emits agent_end for error responses", async () => {
+		const harness = await createHarness();
+		harnesses.push(harness);
+		harness.setResponses([fauxAssistantMessage("", { stopReason: "error", errorMessage: "broken" })]);
+
+		await harness.session.prompt("hi");
+
+		expect(harness.events[harness.events.length - 1]?.type).toBe("agent_end");
+	});
+
+	it("emits agent_end for aborted runs and persists the aborted assistant message", 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;
+
+		expect(harness.events[harness.events.length - 1]?.type).toBe("agent_end");
+		const lastMessage = harness.session.messages[harness.session.messages.length - 1];
+		expect(lastMessage?.role).toBe("assistant");
+		if (lastMessage?.role === "assistant") {
+			expect(lastMessage.stopReason).toBe("aborted");
+		}
+	});
+});

packages/coding-agent/test/suite/agent-session-runtime-host.test.ts

@@ -0,0 +1,321 @@
+import { existsSync, mkdirSync, realpathSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { fauxAssistantMessage, registerFauxProvider } from "@mariozechner/pi-ai";
+import { afterEach, describe, expect, it } from "vitest";
+import {
+	type AgentSessionRuntimeBootstrap,
+	AgentSessionRuntimeHost,
+	createAgentSessionRuntime,
+} from "../../src/core/agent-session-runtime.js";
+import { AuthStorage } from "../../src/core/auth-storage.js";
+import { SessionManager } from "../../src/core/session-manager.js";
+import type {
+	ExtensionFactory,
+	SessionBeforeForkEvent,
+	SessionBeforeSwitchEvent,
+	SessionStartEvent,
+} from "../../src/index.js";
+
+type RecordedSessionEvent = SessionBeforeSwitchEvent | SessionBeforeForkEvent | SessionStartEvent;
+
+describe("AgentSessionRuntimeHost characterization", () => {
+	const cleanups: Array<() => Promise<void> | void> = [];
+
+	afterEach(async () => {
+		while (cleanups.length > 0) {
+			await cleanups.pop()?.();
+		}
+		process.chdir(tmpdir());
+	});
+
+	async function createRuntimeHost(
+		extensionFactory: ExtensionFactory,
+		options?: { cwd?: string; bootstrapModel?: boolean; bootstrapThinkingLevel?: boolean },
+	) {
+		const tempDir =
+			options?.cwd ?? join(tmpdir(), `pi-runtime-suite-${Date.now()}-${Math.random().toString(36).slice(2)}`);
+		mkdirSync(tempDir, { recursive: true });
+
+		const faux = registerFauxProvider({
+			models: [
+				{ id: "faux-1", reasoning: true },
+				{ id: "faux-2", reasoning: false },
+			],
+		});
+		faux.setResponses([fauxAssistantMessage("one"), fauxAssistantMessage("two"), fauxAssistantMessage("three")]);
+
+		const authStorage = AuthStorage.inMemory();
+		authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key");
+
+		const bootstrap: AgentSessionRuntimeBootstrap = {
+			agentDir: tempDir,
+			authStorage,
+			model: options?.bootstrapModel === false ? undefined : faux.getModel(),
+			thinkingLevel: options?.bootstrapThinkingLevel === false ? undefined : undefined,
+			resourceLoader: {
+				extensionFactories: [
+					(pi) => {
+						pi.registerProvider(faux.getModel().provider, {
+							baseUrl: faux.getModel().baseUrl,
+							apiKey: "faux-key",
+							api: faux.api,
+							models: faux.models.map((registeredModel) => ({
+								id: registeredModel.id,
+								name: registeredModel.name,
+								api: registeredModel.api,
+								reasoning: registeredModel.reasoning,
+								input: registeredModel.input,
+								cost: registeredModel.cost,
+								contextWindow: registeredModel.contextWindow,
+								maxTokens: registeredModel.maxTokens,
+							})),
+						});
+						extensionFactory(pi);
+					},
+				],
+				noSkills: true,
+				noPromptTemplates: true,
+				noThemes: true,
+			},
+		};
+		const runtime = await createAgentSessionRuntime(bootstrap, {
+			cwd: tempDir,
+			sessionManager: SessionManager.create(tempDir),
+		});
+		const runtimeHost = new AgentSessionRuntimeHost(bootstrap, runtime);
+		await runtimeHost.session.bindExtensions({});
+
+		cleanups.push(async () => {
+			await runtimeHost.dispose();
+			faux.unregister();
+			if (existsSync(tempDir)) {
+				rmSync(tempDir, { recursive: true, force: true });
+			}
+		});
+
+		return { runtimeHost, faux, tempDir };
+	}
+
+	it("emits session_before_switch and session_start for new and resume flows", async () => {
+		const events: RecordedSessionEvent[] = [];
+		const { runtimeHost } = await createRuntimeHost((pi) => {
+			pi.on("session_before_switch", (event) => {
+				events.push(event);
+			});
+			pi.on("session_start", (event) => {
+				events.push(event);
+			});
+		});
+
+		expect(events).toEqual([{ type: "session_start", reason: "startup" }]);
+		events.length = 0;
+
+		await runtimeHost.session.prompt("hello");
+		const originalSessionFile = runtimeHost.session.sessionFile;
+		const originalSession = runtimeHost.session;
+
+		const newSessionResult = await runtimeHost.newSession();
+		expect(newSessionResult.cancelled).toBe(false);
+		await runtimeHost.session.bindExtensions({});
+		expect(runtimeHost.session).not.toBe(originalSession);
+		expect(runtimeHost.session.messages).toEqual([]);
+		expect(events).toEqual([
+			{ type: "session_before_switch", reason: "new", targetSessionFile: undefined },
+			{ type: "session_start", reason: "new", previousSessionFile: originalSessionFile },
+		]);
+
+		events.length = 0;
+		const secondSessionFile = runtimeHost.session.sessionFile;
+
+		const switchResult = await runtimeHost.switchSession(originalSessionFile!);
+		expect(switchResult.cancelled).toBe(false);
+		await runtimeHost.session.bindExtensions({});
+		expect(events).toEqual([
+			{ type: "session_before_switch", reason: "resume", targetSessionFile: originalSessionFile },
+			{ type: "session_start", reason: "resume", previousSessionFile: secondSessionFile },
+		]);
+	});
+
+	it("honors session_before_switch cancellation for new and resume", async () => {
+		const events: RecordedSessionEvent[] = [];
+		let cancelReason: "new" | "resume" | undefined;
+		const { runtimeHost } = await createRuntimeHost((pi) => {
+			pi.on("session_before_switch", (event) => {
+				events.push(event);
+				if (event.reason === cancelReason) {
+					return { cancel: true };
+				}
+			});
+			pi.on("session_start", (event) => {
+				events.push(event);
+			});
+		});
+
+		await runtimeHost.session.prompt("hello");
+		const originalSessionFile = runtimeHost.session.sessionFile;
+
+		cancelReason = "new";
+		const newResult = await runtimeHost.newSession();
+		expect(newResult.cancelled).toBe(true);
+		expect(runtimeHost.session.sessionFile).toBe(originalSessionFile);
+
+		events.length = 0;
+		const otherDir = join(tmpdir(), `pi-runtime-other-${Date.now()}-${Math.random().toString(36).slice(2)}`);
+		mkdirSync(otherDir, { recursive: true });
+		const otherSession = SessionManager.create(otherDir);
+		otherSession.appendMessage({ role: "user", content: [{ type: "text", text: "other" }], timestamp: Date.now() });
+		const otherSessionFile = otherSession.getSessionFile();
+		cancelReason = "resume";
+		const resumeResult = await runtimeHost.switchSession(otherSessionFile!);
+		expect(resumeResult.cancelled).toBe(true);
+		expect(runtimeHost.session.sessionFile).toBe(originalSessionFile);
+	});
+
+	it("emits session_before_fork and session_start and honors cancellation", async () => {
+		const events: RecordedSessionEvent[] = [];
+		let cancelNextFork = false;
+		const { runtimeHost } = await createRuntimeHost((pi) => {
+			pi.on("session_before_fork", (event) => {
+				events.push(event);
+				if (cancelNextFork) {
+					cancelNextFork = false;
+					return { cancel: true };
+				}
+			});
+			pi.on("session_start", (event) => {
+				events.push(event);
+			});
+		});
+
+		events.length = 0;
+		await runtimeHost.session.prompt("hello");
+		const userMessage = runtimeHost.session.getUserMessagesForForking()[0]!;
+		const previousSessionFile = runtimeHost.session.sessionFile;
+
+		const successResult = await runtimeHost.fork(userMessage.entryId);
+		expect(successResult.cancelled).toBe(false);
+		expect(successResult.selectedText).toBe("hello");
+		await runtimeHost.session.bindExtensions({});
+		expect(events).toEqual([
+			{ type: "session_before_fork", entryId: userMessage.entryId },
+			{ type: "session_start", reason: "fork", previousSessionFile },
+		]);
+
+		events.length = 0;
+		cancelNextFork = true;
+		const cancelResult = await runtimeHost.fork(userMessage.entryId);
+		expect(cancelResult).toEqual({ cancelled: true });
+		expect(events).toEqual([{ type: "session_before_fork", entryId: userMessage.entryId }]);
+	});
+
+	it("throws when forking with an invalid entry id", async () => {
+		const { runtimeHost } = await createRuntimeHost(() => {});
+		await expect(runtimeHost.fork("missing-entry")).rejects.toThrow("Invalid entry ID for forking");
+	});
+
+	it("updates process.cwd() on cross-cwd session replacement", async () => {
+		const firstDir = join(tmpdir(), `pi-runtime-cwd-a-${Date.now()}-${Math.random().toString(36).slice(2)}`);
+		const secondDir = join(tmpdir(), `pi-runtime-cwd-b-${Date.now()}-${Math.random().toString(36).slice(2)}`);
+		mkdirSync(firstDir, { recursive: true });
+		mkdirSync(secondDir, { recursive: true });
+		const { runtimeHost, faux, tempDir } = await createRuntimeHost(() => {}, { cwd: firstDir });
+		const otherAuthStorage = AuthStorage.inMemory();
+		otherAuthStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key");
+		const otherRuntime = await createAgentSessionRuntime(
+			{
+				agentDir: tempDir,
+				authStorage: otherAuthStorage,
+				resourceLoader: {
+					extensionFactories: [
+						(pi) => {
+							pi.registerProvider(faux.getModel().provider, {
+								baseUrl: faux.getModel().baseUrl,
+								apiKey: "faux-key",
+								api: faux.api,
+								models: faux.models.map((registeredModel) => ({
+									id: registeredModel.id,
+									name: registeredModel.name,
+									api: registeredModel.api,
+									reasoning: registeredModel.reasoning,
+									input: registeredModel.input,
+									cost: registeredModel.cost,
+									contextWindow: registeredModel.contextWindow,
+									maxTokens: registeredModel.maxTokens,
+								})),
+							});
+						},
+					],
+					noSkills: true,
+					noPromptTemplates: true,
+					noThemes: true,
+				},
+			},
+			{ cwd: secondDir, sessionManager: SessionManager.create(secondDir) },
+		);
+		cleanups.push(async () => {
+			otherRuntime.session.dispose();
+		});
+		await otherRuntime.session.prompt("other");
+		const otherSessionFile = otherRuntime.session.sessionFile!;
+
+		await runtimeHost.switchSession(otherSessionFile);
+
+		expect(realpathSync(process.cwd())).toBe(realpathSync(secondDir));
+		expect(realpathSync(runtimeHost.session.sessionManager.getCwd())).toBe(realpathSync(secondDir));
+	});
+
+	it("restores model and thinking state from the destination session", async () => {
+		const { runtimeHost, faux, tempDir } = await createRuntimeHost(() => {}, {
+			bootstrapModel: false,
+			bootstrapThinkingLevel: false,
+		});
+		const otherDir = join(tempDir, "other");
+		mkdirSync(otherDir, { recursive: true });
+		const otherAuthStorage = AuthStorage.inMemory();
+		otherAuthStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key");
+		const otherRuntime = await createAgentSessionRuntime(
+			{
+				agentDir: tempDir,
+				authStorage: otherAuthStorage,
+				resourceLoader: {
+					extensionFactories: [
+						(pi) => {
+							pi.registerProvider(faux.getModel().provider, {
+								baseUrl: faux.getModel().baseUrl,
+								apiKey: "faux-key",
+								api: faux.api,
+								models: faux.models.map((registeredModel) => ({
+									id: registeredModel.id,
+									name: registeredModel.name,
+									api: registeredModel.api,
+									reasoning: registeredModel.reasoning,
+									input: registeredModel.input,
+									cost: registeredModel.cost,
+									contextWindow: registeredModel.contextWindow,
+									maxTokens: registeredModel.maxTokens,
+								})),
+							});
+						},
+					],
+					noSkills: true,
+					noPromptTemplates: true,
+					noThemes: true,
+				},
+			},
+			{ cwd: otherDir, sessionManager: SessionManager.create(otherDir) },
+		);
+		cleanups.push(async () => {
+			otherRuntime.session.dispose();
+		});
+		await otherRuntime.session.setModel(faux.getModel("faux-2")!);
+		otherRuntime.session.setThinkingLevel("off");
+		await otherRuntime.session.prompt("hello");
+		const targetSessionFile = otherRuntime.session.sessionFile!;
+
+		await runtimeHost.switchSession(targetSessionFile);
+
+		expect(runtimeHost.session.model?.id).toBe("faux-2");
+		expect(runtimeHost.session.thinkingLevel).toBe("off");
+	});
+});

packages/coding-agent/test/suite/harness.ts

@@ -5,12 +5,14 @@
 import { existsSync, mkdirSync, rmSync } from "node:fs";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
-import type { AgentTool } from "@mariozechner/pi-agent-core";
+import type { AgentMessage, AgentTool } from "@mariozechner/pi-agent-core";
 import { Agent } from "@mariozechner/pi-agent-core";
 import type { FauxModelDefinition, FauxProviderRegistration, FauxResponseStep, Model } from "@mariozechner/pi-ai";
 import { registerFauxProvider } from "@mariozechner/pi-ai";
 import { AgentSession, type AgentSessionEvent } from "../../src/core/agent-session.js";
 import { AuthStorage } from "../../src/core/auth-storage.js";
+import type { ExtensionRunner } from "../../src/core/extensions/index.js";
+import { convertToLlm } from "../../src/core/messages.js";
 import { ModelRegistry } from "../../src/core/model-registry.js";
 import { SessionManager } from "../../src/core/session-manager.js";
 import type { Settings } from "../../src/core/settings-manager.js";
@@ -22,6 +24,37 @@ import {
 	createTestResourceLoader,
 } from "../utilities.js";
 
+type MessageTextPart = { type: "text"; text: string };
+
+export function getMessageText(message: unknown): string {
+	if (!message || typeof message !== "object" || !("content" in message)) {
+		return "";
+	}
+	const content = (message as { content?: string | Array<{ type: string; text?: string }> }).content;
+	if (content === undefined) {
+		return "";
+	}
+	if (typeof content === "string") {
+		return content;
+	}
+	return content
+		.filter((part): part is MessageTextPart => part.type === "text")
+		.map((part) => part.text)
+		.join("\n");
+}
+
+export function getUserTexts(harness: Harness): string[] {
+	return harness.session.messages
+		.filter((message) => message.role === "user")
+		.map((message) => getMessageText(message));
+}
+
+export function getAssistantTexts(harness: Harness): string[] {
+	return harness.session.messages
+		.filter((message) => message.role === "assistant")
+		.map((message) => getMessageText(message));
+}
+
 export interface HarnessOptions {
 	models?: FauxModelDefinition[];
 	settings?: Partial<Settings>;
@@ -29,12 +62,14 @@ export interface HarnessOptions {
 	tools?: AgentTool[];
 	resourceLoader?: ResourceLoader;
 	extensionFactories?: Array<ExtensionFactory | CreateTestExtensionsResultInput>;
+	withConfiguredAuth?: boolean;
 }
 
 export interface Harness {
 	session: AgentSession;
 	sessionManager: SessionManager;
 	settingsManager: SettingsManager;
+	authStorage: AuthStorage;
 	faux: FauxProviderRegistration;
 	models: [Model<string>, ...Model<string>[]];
 	getModel(): Model<string>;
@@ -62,37 +97,56 @@ export async function createHarness(options: HarnessOptions = {}): Promise<Harne
 	fauxProvider.setResponses([]);
 	const model = fauxProvider.getModel();
 	const toolMap = options.tools ? Object.fromEntries(options.tools.map((tool) => [tool.name, tool])) : undefined;
-
+	const withConfiguredAuth = options.withConfiguredAuth ?? true;
-	const agent = new Agent({
+	const extensionRunnerRef: { current?: ExtensionRunner } = {};
-		getApiKey: () => "faux-key",
-		initialState: {
-			model,
-			systemPrompt: options.systemPrompt ?? "You are a test assistant.",
-			tools: [],
-		},
-	});
 
 	const sessionManager = SessionManager.inMemory();
 	const settingsManager = SettingsManager.inMemory(options.settings);
 
 	const authStorage = AuthStorage.inMemory();
-	authStorage.setRuntimeApiKey(model.provider, "faux-key");
+	if (withConfiguredAuth) {
+		authStorage.setRuntimeApiKey(model.provider, "faux-key");
+	}
 	const modelRegistry = ModelRegistry.inMemory(authStorage);
-	modelRegistry.registerProvider(model.provider, {
+	if (withConfiguredAuth) {
-		baseUrl: model.baseUrl,
+		modelRegistry.registerProvider(model.provider, {
-		apiKey: "faux-key",
+			baseUrl: model.baseUrl,
-		api: fauxProvider.api,
+			apiKey: "faux-key",
-		models: fauxProvider.models.map((registeredModel) => ({
+			api: fauxProvider.api,
-			id: registeredModel.id,
+			models: fauxProvider.models.map((registeredModel) => ({
-			name: registeredModel.name,
+				id: registeredModel.id,
-			api: registeredModel.api,
+				name: registeredModel.name,
-			reasoning: registeredModel.reasoning,
+				api: registeredModel.api,
-			input: registeredModel.input,
+				reasoning: registeredModel.reasoning,
-			cost: registeredModel.cost,
+				input: registeredModel.input,
-			contextWindow: registeredModel.contextWindow,
+				cost: registeredModel.cost,
-			maxTokens: registeredModel.maxTokens,
+				contextWindow: registeredModel.contextWindow,
-			baseUrl: registeredModel.baseUrl,
+				maxTokens: registeredModel.maxTokens,
-		})),
+				baseUrl: registeredModel.baseUrl,
+			})),
+		});
+	}
+
+	const agent = new Agent({
+		getApiKey: () => (withConfiguredAuth ? "faux-key" : undefined),
+		initialState: {
+			model,
+			systemPrompt: options.systemPrompt ?? "You are a test assistant.",
+			tools: [],
+		},
+		convertToLlm,
+		onPayload: async (payload) => {
+			const runner = extensionRunnerRef.current;
+			if (!runner?.hasHandlers("before_provider_request")) {
+				return payload;
+			}
+			return runner.emitBeforeProviderRequest(payload);
+		},
+		transformContext: async (messages: AgentMessage[]) => {
+			const runner = extensionRunnerRef.current;
+			if (!runner) return messages;
+			return runner.emitContext(messages);
+		},
 	});
 	const extensionsResult = options.extensionFactories
 		? await createTestExtensionsResult(options.extensionFactories, tempDir)
@@ -108,6 +162,7 @@ export async function createHarness(options: HarnessOptions = {}): Promise<Harne
 		modelRegistry,
 		resourceLoader,
 		baseToolsOverride: toolMap,
+		extensionRunnerRef,
 	});
 
 	const events: AgentSessionEvent[] = [];
@@ -119,6 +174,7 @@ export async function createHarness(options: HarnessOptions = {}): Promise<Harne
 		session,
 		sessionManager,
 		settingsManager,
+		authStorage,
 		faux: fauxProvider,
 		models: fauxProvider.models,
 		getModel: fauxProvider.getModel,

packages/coding-agent/test/suite/regressions/2023-queued-slash-command-followup.test.ts

@@ -0,0 +1,80 @@
+import type { AgentTool } from "@mariozechner/pi-agent-core";
+import { fauxAssistantMessage, fauxToolCall } from "@mariozechner/pi-ai";
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import { Type } from "@sinclair/typebox";
+import { afterEach, describe, expect, it } from "vitest";
+import { createHarness, getAssistantTexts, getUserTexts, type Harness } from "../harness.js";
+
+describe("issue #2023 queued slash-command follow-up", () => {
+	const harnesses: Harness[] = [];
+
+	afterEach(() => {
+		while (harnesses.length > 0) {
+			harnesses.pop()?.cleanup();
+		}
+	});
+
+	it("treats extension-origin queued slash-command follow-ups as raw user text instead of dispatching the command", async () => {
+		let extensionApi: ExtensionAPI | undefined;
+		const commandRuns: string[] = [];
+		let releaseToolExecution: (() => void) | undefined;
+		const toolRelease = new Promise<void>((resolve) => {
+			releaseToolExecution = resolve;
+		});
+		const waitTool: AgentTool = {
+			name: "wait",
+			label: "Wait",
+			description: "Wait for the test to release execution",
+			parameters: Type.Object({}),
+			execute: async () => {
+				await toolRelease;
+				return {
+					content: [{ type: "text", text: "released" }],
+					details: {},
+				};
+			},
+		};
+		const harness = await createHarness({
+			tools: [waitTool],
+			extensionFactories: [
+				(pi) => {
+					extensionApi = pi;
+					pi.registerCommand("testcmd", {
+						description: "Test command",
+						handler: async (args) => {
+							commandRuns.push(args);
+						},
+					});
+				},
+			],
+		});
+		harnesses.push(harness);
+
+		harness.setResponses([
+			fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }),
+			fauxAssistantMessage("first turn complete"),
+			fauxAssistantMessage("queued follow-up handled by model"),
+		]);
+
+		const sawToolStart = new Promise<void>((resolve) => {
+			const unsubscribe = harness.session.subscribe((event) => {
+				if (event.type === "tool_execution_start" && event.toolName === "wait") {
+					unsubscribe();
+					resolve();
+				}
+			});
+		});
+
+		const promptPromise = harness.session.prompt("start");
+		await sawToolStart;
+		await new Promise((resolve) => setTimeout(resolve, 0));
+
+		extensionApi?.sendUserMessage("/testcmd queued", { deliverAs: "followUp" });
+		releaseToolExecution?.();
+		await promptPromise;
+
+		expect(commandRuns).toEqual([]);
+		expect(getUserTexts(harness)).toEqual(["start", "/testcmd queued"]);
+		expect(getAssistantTexts(harness)).toContain("queued follow-up handled by model");
+	});
+});