fix(coding-agent): emit session_shutdown in print mode closes #2576

packages/coding-agent/CHANGELOG.md

@@ -8,6 +8,7 @@
 - Fixed interactive bash execution collapsed previews to recompute visual line wrapping at render time, so previews respect the current terminal width after resizes and split-pane width changes ([#2569](https://github.com/badlogic/pi-mono/issues/2569))
 - Fixed RPC `get_session_stats` to expose `contextUsage`, so headless clients can read actual current context-window usage instead of deriving it from token totals ([#2550](https://github.com/badlogic/pi-mono/issues/2550))
 - Fixed `pi update` for git packages to fetch only the tracked target branch with `--no-tags`, reducing unrelated branch and tag noise while preserving force-push-safe updates ([#2548](https://github.com/badlogic/pi-mono/issues/2548))
+- Fixed print and JSON modes to emit `session_shutdown` before exit, so extensions can release long-lived resources and non-interactive runs terminate cleanly ([#2576](https://github.com/badlogic/pi-mono/issues/2576))
 
 ## [0.62.0] - 2026-03-23
 

packages/coding-agent/src/main.ts

@@ -882,7 +882,7 @@ export async function main(args: string[]) {
 		printTimings();
 		await interactiveMode.run();
 	} else {
-		await runPrintMode(session, {
+		const exitCode = await runPrintMode(session, {
 			mode,
 			messages: parsed.messages,
 			initialMessage,
@@ -890,6 +890,9 @@ export async function main(args: string[]) {
 		});
 		stopThemeWatcher();
 		restoreStdout();
+		if (exitCode !== 0) {
+			process.exitCode = exitCode;
+		}
 		return;
 	}
 }

packages/coding-agent/src/modes/print-mode.ts

@@ -28,93 +28,104 @@ export interface PrintModeOptions {
  * Run in print (single-shot) mode.
  * Sends prompts to the agent and outputs the result.
  */
-export async function runPrintMode(session: AgentSession, options: PrintModeOptions): Promise<void> {
+export async function runPrintMode(session: AgentSession, options: PrintModeOptions): Promise<number> {
 	const { mode, messages = [], initialMessage, initialImages } = options;
-	if (mode === "json") {
-		const header = session.sessionManager.getHeader();
-		if (header) {
-			writeRawStdout(`${JSON.stringify(header)}\n`);
-		}
-	}
-	// Set up extensions for print mode (no UI)
-	await session.bindExtensions({
-		commandContextActions: {
-			waitForIdle: () => session.agent.waitForIdle(),
-			newSession: async (options) => {
-				const success = await session.newSession({ parentSession: options?.parentSession });
-				if (success && options?.setup) {
-					await options.setup(session.sessionManager);
-				}
-				return { cancelled: !success };
-			},
-			fork: async (entryId) => {
-				const result = await session.fork(entryId);
-				return { cancelled: result.cancelled };
-			},
-			navigateTree: async (targetId, options) => {
-				const result = await session.navigateTree(targetId, {
-					summarize: options?.summarize,
-					customInstructions: options?.customInstructions,
-					replaceInstructions: options?.replaceInstructions,
-					label: options?.label,
-				});
-				return { cancelled: result.cancelled };
-			},
-			switchSession: async (sessionPath) => {
-				const success = await session.switchSession(sessionPath);
-				return { cancelled: !success };
-			},
-			reload: async () => {
-				await session.reload();
-			},
-		},
-		onError: (err) => {
-			console.error(`Extension error (${err.extensionPath}): ${err.error}`);
-		},
-	});
+	let exitCode = 0;
 
-	// Always subscribe to enable session persistence via _handleAgentEvent
-	session.subscribe((event) => {
-		// In JSON mode, output all events
+	try {
 		if (mode === "json") {
-			writeRawStdout(`${JSON.stringify(event)}\n`);
-		}
-	});
-
-	// Send initial message with attachments
-	if (initialMessage) {
-		await session.prompt(initialMessage, { images: initialImages });
-	}
-
-	// Send remaining messages
-	for (const message of messages) {
-		await session.prompt(message);
-	}
-
-	// In text mode, output final response
-	if (mode === "text") {
-		const state = session.state;
-		const lastMessage = state.messages[state.messages.length - 1];
-
-		if (lastMessage?.role === "assistant") {
-			const assistantMsg = lastMessage as AssistantMessage;
-
-			// Check for error/aborted
-			if (assistantMsg.stopReason === "error" || assistantMsg.stopReason === "aborted") {
-				console.error(assistantMsg.errorMessage || `Request ${assistantMsg.stopReason}`);
-				process.exit(1);
+			const header = session.sessionManager.getHeader();
+			if (header) {
+				writeRawStdout(`${JSON.stringify(header)}\n`);
 			}
+		}
+		// Set up extensions for print mode (no UI)
+		await session.bindExtensions({
+			commandContextActions: {
+				waitForIdle: () => session.agent.waitForIdle(),
+				newSession: async (options) => {
+					const success = await session.newSession({ parentSession: options?.parentSession });
+					if (success && options?.setup) {
+						await options.setup(session.sessionManager);
+					}
+					return { cancelled: !success };
+				},
+				fork: async (entryId) => {
+					const result = await session.fork(entryId);
+					return { cancelled: result.cancelled };
+				},
+				navigateTree: async (targetId, options) => {
+					const result = await session.navigateTree(targetId, {
+						summarize: options?.summarize,
+						customInstructions: options?.customInstructions,
+						replaceInstructions: options?.replaceInstructions,
+						label: options?.label,
+					});
+					return { cancelled: result.cancelled };
+				},
+				switchSession: async (sessionPath) => {
+					const success = await session.switchSession(sessionPath);
+					return { cancelled: !success };
+				},
+				reload: async () => {
+					await session.reload();
+				},
+			},
+			onError: (err) => {
+				console.error(`Extension error (${err.extensionPath}): ${err.error}`);
+			},
+		});
 
-			// Output text content
-			for (const content of assistantMsg.content) {
-				if (content.type === "text") {
-					writeRawStdout(`${content.text}\n`);
+		// Always subscribe to enable session persistence via _handleAgentEvent
+		session.subscribe((event) => {
+			// In JSON mode, output all events
+			if (mode === "json") {
+				writeRawStdout(`${JSON.stringify(event)}\n`);
+			}
+		});
+
+		// Send initial message with attachments
+		if (initialMessage) {
+			await session.prompt(initialMessage, { images: initialImages });
+		}
+
+		// Send remaining messages
+		for (const message of messages) {
+			await session.prompt(message);
+		}
+
+		// In text mode, output final response
+		if (mode === "text") {
+			const state = session.state;
+			const lastMessage = state.messages[state.messages.length - 1];
+
+			if (lastMessage?.role === "assistant") {
+				const assistantMsg = lastMessage as AssistantMessage;
+
+				// Check for error/aborted
+				if (assistantMsg.stopReason === "error" || assistantMsg.stopReason === "aborted") {
+					console.error(assistantMsg.errorMessage || `Request ${assistantMsg.stopReason}`);
+					exitCode = 1;
+				} else {
+					// Output text content
+					for (const content of assistantMsg.content) {
+						if (content.type === "text") {
+							writeRawStdout(`${content.text}\n`);
+						}
+					}
 				}
 			}
 		}
-	}
 
-	// Ensure stdout is fully flushed before returning
-	// This prevents race conditions where the process exits before all output is written
-	await flushRawStdout();
+		return exitCode;
+	} finally {
+		const extensionRunner = session.extensionRunner;
+		if (extensionRunner?.hasHandlers("session_shutdown")) {
+			await extensionRunner.emit({ type: "session_shutdown" });
+		}
+
+		// Ensure stdout is fully flushed before returning
+		// This prevents race conditions where the process exits before all output is written
+		await flushRawStdout();
+	}
 }

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

@@ -0,0 +1,124 @@
+import type { AssistantMessage, ImageContent } from "@mariozechner/pi-ai";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { runPrintMode } from "../src/modes/print-mode.js";
+
+type EmitEvent = { type: string };
+
+type FakeExtensionRunner = {
+	hasHandlers: (eventType: string) => boolean;
+	emit: ReturnType<typeof vi.fn<(event: EmitEvent) => Promise<void>>>;
+};
+
+type FakeSession = {
+	sessionManager: { getHeader: () => object | undefined };
+	agent: { waitForIdle: () => Promise<void> };
+	state: { messages: AssistantMessage[] };
+	extensionRunner: FakeExtensionRunner;
+	bindExtensions: ReturnType<typeof vi.fn>;
+	subscribe: ReturnType<typeof vi.fn>;
+	prompt: ReturnType<typeof vi.fn>;
+	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>;
+};
+
+function createAssistantMessage(options?: {
+	text?: string;
+	stopReason?: AssistantMessage["stopReason"];
+	errorMessage?: string;
+}): AssistantMessage {
+	return {
+		role: "assistant",
+		content: options?.text ? [{ type: "text", text: options.text }] : [],
+		api: "openai-responses",
+		provider: "openai",
+		model: "gpt-4o-mini",
+		usage: {
+			input: 0,
+			output: 0,
+			cacheRead: 0,
+			cacheWrite: 0,
+			totalTokens: 0,
+			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+		},
+		stopReason: options?.stopReason ?? "stop",
+		errorMessage: options?.errorMessage,
+		timestamp: Date.now(),
+	};
+}
+
+function createSession(assistantMessage: AssistantMessage): FakeSession {
+	const extensionRunner: FakeExtensionRunner = {
+		hasHandlers: (eventType: string) => eventType === "session_shutdown",
+		emit: vi.fn(async () => {}),
+	};
+
+	const state = { messages: [assistantMessage] };
+
+	return {
+		sessionManager: { getHeader: () => undefined },
+		agent: { waitForIdle: async () => {} },
+		state,
+		extensionRunner,
+		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 () => {}),
+	};
+}
+
+afterEach(() => {
+	vi.restoreAllMocks();
+});
+
+describe("runPrintMode", () => {
+	it("emits session_shutdown in text mode", async () => {
+		const session = createSession(createAssistantMessage({ text: "done" }));
+		const images: ImageContent[] = [{ type: "image", mimeType: "image/png", data: "abc" }];
+
+		const exitCode = await runPrintMode(session as unknown as Parameters<typeof runPrintMode>[0], {
+			mode: "text",
+			initialMessage: "Say done",
+			initialImages: images,
+		});
+
+		expect(exitCode).toBe(0);
+		expect(session.prompt).toHaveBeenCalledWith("Say done", { images });
+		expect(session.extensionRunner.emit).toHaveBeenCalledTimes(1);
+		expect(session.extensionRunner.emit).toHaveBeenCalledWith({ type: "session_shutdown" });
+	});
+
+	it("emits session_shutdown in json mode", async () => {
+		const session = createSession(createAssistantMessage({ text: "done" }));
+
+		const exitCode = await runPrintMode(session as unknown as Parameters<typeof runPrintMode>[0], {
+			mode: "json",
+			messages: ["hello"],
+		});
+
+		expect(exitCode).toBe(0);
+		expect(session.prompt).toHaveBeenCalledWith("hello");
+		expect(session.extensionRunner.emit).toHaveBeenCalledTimes(1);
+		expect(session.extensionRunner.emit).toHaveBeenCalledWith({ type: "session_shutdown" });
+	});
+
+	it("emits session_shutdown and returns non-zero on assistant error", async () => {
+		const session = createSession(createAssistantMessage({ stopReason: "error", errorMessage: "provider failure" }));
+		const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {});
+
+		const exitCode = await runPrintMode(session as unknown as Parameters<typeof runPrintMode>[0], {
+			mode: "text",
+		});
+
+		expect(exitCode).toBe(1);
+		expect(errorSpy).toHaveBeenCalledWith("provider failure");
+		expect(session.extensionRunner.emit).toHaveBeenCalledTimes(1);
+		expect(session.extensionRunner.emit).toHaveBeenCalledWith({ type: "session_shutdown" });
+	});
+});