refactor(coding-agent): add runtime host for session switching closes #2024

packages/coding-agent/test/agent-session-branching.test.ts

@@ -10,126 +10,112 @@
 import { existsSync, mkdirSync, rmSync } from "node:fs";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
-import { Agent } from "@mariozechner/pi-agent-core";
 import { getModel } from "@mariozechner/pi-ai";
 import { afterEach, beforeEach, describe, expect, it } from "vitest";
-import { AgentSession } from "../src/core/agent-session.js";
+import type { AgentSession } from "../src/core/agent-session.js";
+import {
+	type AgentSessionRuntimeHost,
+	createAgentSessionRuntime,
+	AgentSessionRuntimeHost as RuntimeHost,
+} from "../src/core/agent-session-runtime.js";
 import { AuthStorage } from "../src/core/auth-storage.js";
-import { ModelRegistry } from "../src/core/model-registry.js";
 import { SessionManager } from "../src/core/session-manager.js";
-import { SettingsManager } from "../src/core/settings-manager.js";
 import { codingTools } from "../src/core/tools/index.js";
-import { API_KEY, createTestResourceLoader } from "./utilities.js";
+import { API_KEY } from "./utilities.js";
 
 describe.skipIf(!API_KEY)("AgentSession forking", () => {
 	let session: AgentSession;
+	let runtimeHost: AgentSessionRuntimeHost;
 	let tempDir: string;
 	let sessionManager: SessionManager;
 
 	beforeEach(() => {
-		// Create temp directory for session files
 		tempDir = join(tmpdir(), `pi-branching-test-${Date.now()}`);
 		mkdirSync(tempDir, { recursive: true });
 	});
 
 	afterEach(async () => {
-		if (session) {
-			session.dispose();
+		if (runtimeHost) {
+			await runtimeHost.dispose();
 		}
+		process.chdir(tmpdir());
 		if (tempDir && existsSync(tempDir)) {
 			rmSync(tempDir, { recursive: true });
 		}
 	});
 
-	function createSession(noSession: boolean = false) {
+	async function createSession(noSession: boolean = false) {
 		const model = getModel("anthropic", "claude-sonnet-4-5")!;
-		const agent = new Agent({
-			getApiKey: () => API_KEY,
-			initialState: {
-				model,
-				systemPrompt: "You are a helpful assistant. Be extremely concise, reply with just a few words.",
-				tools: codingTools,
-			},
-		});
-
-		sessionManager = noSession ? SessionManager.inMemory() : SessionManager.create(tempDir);
-		const settingsManager = SettingsManager.create(tempDir, tempDir);
+		sessionManager = noSession ? SessionManager.inMemory(tempDir) : SessionManager.create(tempDir);
 		const authStorage = AuthStorage.create(join(tempDir, "auth.json"));
-		const modelRegistry = ModelRegistry.create(authStorage, tempDir);
+		authStorage.setRuntimeApiKey("anthropic", API_KEY!);
 
-		session = new AgentSession({
-			agent,
-			sessionManager,
-			settingsManager,
+		const bootstrap = {
+			agentDir: tempDir,
+			authStorage,
+			model,
+			tools: codingTools,
+			resourceLoader: {
+				noExtensions: true,
+				noSkills: true,
+				noPromptTemplates: true,
+				noThemes: true,
+			},
+		};
+		const runtime = await createAgentSessionRuntime(bootstrap, {
 			cwd: tempDir,
-			modelRegistry,
-			resourceLoader: createTestResourceLoader(),
+			sessionManager,
 		});
-
-		// Must subscribe to enable session persistence
+		runtimeHost = new RuntimeHost(bootstrap, runtime);
+		session = runtimeHost.session;
 		session.subscribe(() => {});
-
 		return session;
 	}
 
 	it("should allow forking from single message", async () => {
-		createSession();
+		await createSession();
 
-		// Send one message
 		await session.prompt("Say hello");
 		await session.agent.waitForIdle();
 
-		// Should have exactly 1 user message available for forking
 		const userMessages = session.getUserMessagesForForking();
 		expect(userMessages.length).toBe(1);
 		expect(userMessages[0].text).toBe("Say hello");
 
-		// Fork from the first message
-		const result = await session.fork(userMessages[0].entryId);
-		expect(result.selectedText).toBe("Say hello");
+		const result = await runtimeHost.fork(userMessages[0].entryId);
 		expect(result.cancelled).toBe(false);
+		session = runtimeHost.session;
+		expect(result.selectedText).toBe("Say hello");
 
-		// After forking, conversation should be empty (forked before the first message)
 		expect(session.messages.length).toBe(0);
-
-		// Session file path should be set, but file is created lazily after first assistant message
 		expect(session.sessionFile).not.toBeNull();
 		expect(existsSync(session.sessionFile!)).toBe(false);
 	});
 
 	it("should support in-memory forking in --no-session mode", async () => {
-		createSession(true);
+		await createSession(true);
 
-		// Verify sessions are disabled
 		expect(session.sessionFile).toBeUndefined();
 
-		// Send one message
 		await session.prompt("Say hi");
 		await session.agent.waitForIdle();
 
-		// Should have 1 user message
 		const userMessages = session.getUserMessagesForForking();
 		expect(userMessages.length).toBe(1);
-
-		// Verify we have messages before forking
 		expect(session.messages.length).toBeGreaterThan(0);
 
-		// Fork from the first message
-		const result = await session.fork(userMessages[0].entryId);
-		expect(result.selectedText).toBe("Say hi");
+		const result = await runtimeHost.fork(userMessages[0].entryId);
 		expect(result.cancelled).toBe(false);
+		session = runtimeHost.session;
+		expect(result.selectedText).toBe("Say hi");
 
-		// After forking, conversation should be empty
 		expect(session.messages.length).toBe(0);
-
-		// Session file should still be undefined (no file created)
 		expect(session.sessionFile).toBeUndefined();
 	});
 
 	it("should fork from middle of conversation", async () => {
-		createSession();
+		await createSession();
 
-		// Send multiple messages
 		await session.prompt("Say one");
 		await session.agent.waitForIdle();
 
@@ -139,18 +125,17 @@ describe.skipIf(!API_KEY)("AgentSession forking", () => {
 		await session.prompt("Say three");
 		await session.agent.waitForIdle();
 
-		// Should have 3 user messages
 		const userMessages = session.getUserMessagesForForking();
 		expect(userMessages.length).toBe(3);
 
-		// Fork from second message (keeps first message + response)
 		const secondMessage = userMessages[1];
-		const result = await session.fork(secondMessage.entryId);
+		const result = await runtimeHost.fork(secondMessage.entryId);
+		expect(result.cancelled).toBe(false);
+		session = runtimeHost.session;
 		expect(result.selectedText).toBe("Say two");
 
-		// After forking, should have first user message + assistant response
 		expect(session.messages.length).toBe(2);
 		expect(session.messages[0].role).toBe("user");
 		expect(session.messages[1].role).toBe("assistant");
-	});
+	}, 60000);
 });

packages/coding-agent/test/agent-session-runtime-events.test.ts

@@ -0,0 +1,172 @@
+import { existsSync, mkdirSync, 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 session lifecycle events", () => {
+	const cleanups: Array<() => Promise<void> | void> = [];
+
+	afterEach(async () => {
+		while (cleanups.length > 0) {
+			await cleanups.pop()?.();
+		}
+	});
+
+	async function createRuntimeHost(extensionFactory: ExtensionFactory) {
+		const tempDir = join(tmpdir(), `pi-runtime-events-${Date.now()}-${Math.random().toString(36).slice(2)}`);
+		mkdirSync(tempDir, { recursive: true });
+
+		const faux = registerFauxProvider();
+		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: faux.getModel(),
+			resourceLoader: {
+				extensionFactories: [extensionFactory],
+				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();
+			process.chdir(tmpdir());
+			if (existsSync(tempDir)) {
+				rmSync(tempDir, { recursive: true, force: true });
+			}
+		});
+
+		return { runtimeHost, faux };
+	}
+
+	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;
+		expect(originalSessionFile).toBeTruthy();
+
+		const newSessionResult = await runtimeHost.newSession();
+		expect(newSessionResult.cancelled).toBe(false);
+		await runtimeHost.session.bindExtensions({});
+		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;
+		expect(secondSessionFile).toBeTruthy();
+
+		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", async () => {
+		const events: RecordedSessionEvent[] = [];
+		const { runtimeHost } = await createRuntimeHost((pi) => {
+			pi.on("session_before_switch", (event) => {
+				events.push(event);
+				return { cancel: true };
+			});
+			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 result = await runtimeHost.newSession();
+		expect(result.cancelled).toBe(true);
+		expect(runtimeHost.session.sessionFile).toBe(originalSessionFile);
+		expect(events).toEqual([{ type: "session_before_switch", reason: "new", targetSessionFile: undefined }]);
+	});
+
+	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);
+			});
+		});
+
+		expect(events).toEqual([{ type: "session_start", reason: "startup" }]);
+		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 }]);
+	});
+});

packages/coding-agent/test/footer-width.test.ts

@@ -48,6 +48,7 @@ function createSession(options: {
 		sessionManager: {
 			getEntries: () => entries,
 			getSessionName: () => options.sessionName,
+			getCwd: () => "/tmp/project",
 		},
 		getContextUsage: () => ({ contextWindow: 200_000, percent: 12.3 }),
 		modelRegistry: {

packages/coding-agent/test/print-mode.test.ts

@@ -17,11 +17,15 @@ type FakeSession = {
 	bindExtensions: ReturnType<typeof vi.fn>;
 	subscribe: ReturnType<typeof vi.fn>;
 	prompt: ReturnType<typeof vi.fn>;
+	reload: ReturnType<typeof vi.fn>;
+};
+
+type FakeRuntimeHost = {
+	session: FakeSession;
 	newSession: ReturnType<typeof vi.fn>;
 	fork: ReturnType<typeof vi.fn>;
-	navigateTree: ReturnType<typeof vi.fn>;
 	switchSession: ReturnType<typeof vi.fn>;
-	reload: ReturnType<typeof vi.fn>;
+	dispose: ReturnType<typeof vi.fn>;
 };
 
 function createAssistantMessage(options?: {
@@ -49,7 +53,7 @@ function createAssistantMessage(options?: {
 	};
 }
 
-function createSession(assistantMessage: AssistantMessage): FakeSession {
+function createRuntimeHost(assistantMessage: AssistantMessage): FakeRuntimeHost {
 	const extensionRunner: FakeExtensionRunner = {
 		hasHandlers: (eventType: string) => eventType === "session_shutdown",
 		emit: vi.fn(async () => {}),
@@ -57,7 +61,7 @@ function createSession(assistantMessage: AssistantMessage): FakeSession {
 
 	const state = { messages: [assistantMessage] };
 
-	return {
+	const session: FakeSession = {
 		sessionManager: { getHeader: () => undefined },
 		agent: { waitForIdle: async () => {} },
 		state,
@@ -65,12 +69,18 @@ function createSession(assistantMessage: AssistantMessage): FakeSession {
 		bindExtensions: vi.fn(async () => {}),
 		subscribe: vi.fn(() => () => {}),
 		prompt: vi.fn(async () => {}),
-		newSession: vi.fn(async () => true),
-		fork: vi.fn(async () => ({ cancelled: false })),
-		navigateTree: vi.fn(async () => ({ cancelled: false })),
-		switchSession: vi.fn(async () => true),
 		reload: vi.fn(async () => {}),
 	};
+
+	return {
+		session,
+		newSession: vi.fn(async () => undefined),
+		fork: vi.fn(async () => ({ selectedText: "" })),
+		switchSession: vi.fn(async () => undefined),
+		dispose: vi.fn(async () => {
+			await session.extensionRunner.emit({ type: "session_shutdown" });
+		}),
+	};
 }
 
 afterEach(() => {
@@ -79,10 +89,11 @@ afterEach(() => {
 
 describe("runPrintMode", () => {
 	it("emits session_shutdown in text mode", async () => {
-		const session = createSession(createAssistantMessage({ text: "done" }));
+		const runtimeHost = createRuntimeHost(createAssistantMessage({ text: "done" }));
+		const { session } = runtimeHost;
 		const images: ImageContent[] = [{ type: "image", mimeType: "image/png", data: "abc" }];
 
-		const exitCode = await runPrintMode(session as unknown as Parameters<typeof runPrintMode>[0], {
+		const exitCode = await runPrintMode(runtimeHost as unknown as Parameters<typeof runPrintMode>[0], {
 			mode: "text",
 			initialMessage: "Say done",
 			initialImages: images,
@@ -95,9 +106,10 @@ describe("runPrintMode", () => {
 	});
 
 	it("emits session_shutdown in json mode", async () => {
-		const session = createSession(createAssistantMessage({ text: "done" }));
+		const runtimeHost = createRuntimeHost(createAssistantMessage({ text: "done" }));
+		const { session } = runtimeHost;
 
-		const exitCode = await runPrintMode(session as unknown as Parameters<typeof runPrintMode>[0], {
+		const exitCode = await runPrintMode(runtimeHost as unknown as Parameters<typeof runPrintMode>[0], {
 			mode: "json",
 			messages: ["hello"],
 		});
@@ -109,10 +121,13 @@ describe("runPrintMode", () => {
 	});
 
 	it("emits session_shutdown and returns non-zero on assistant error", async () => {
-		const session = createSession(createAssistantMessage({ stopReason: "error", errorMessage: "provider failure" }));
+		const runtimeHost = createRuntimeHost(
+			createAssistantMessage({ stopReason: "error", errorMessage: "provider failure" }),
+		);
+		const { session } = runtimeHost;
 		const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {});
 
-		const exitCode = await runPrintMode(session as unknown as Parameters<typeof runPrintMode>[0], {
+		const exitCode = await runPrintMode(runtimeHost as unknown as Parameters<typeof runPrintMode>[0], {
 			mode: "text",
 		});