feat(coding-agent): make ctx.shutdown() available for extensions (#542)

packages/coding-agent/CHANGELOG.md

@@ -21,6 +21,7 @@
 - Extension UI dialogs (`ctx.ui.select()`, `ctx.ui.confirm()`, `ctx.ui.input()`) now support a `timeout` option with live countdown display ([#522](https://github.com/badlogic/pi-mono/pull/522) by [@nicobailon](https://github.com/nicobailon))
 - Extensions can now provide custom editor components via `ctx.ui.setEditorComponent()`. See `examples/extensions/modal-editor.ts` and `docs/tui.md` Pattern 7.
 - Extension factories can now be async, enabling dynamic imports and lazy-loaded dependencies ([#513](https://github.com/badlogic/pi-mono/pull/513) by [@austinm911](https://github.com/austinm911))
+- `ctx.shutdown()` is now available in extension contexts for requesting a graceful shutdown. In interactive mode, shutdown is deferred until the agent becomes idle (after processing all queued steering and follow-up messages). In RPC mode, shutdown is deferred until after completing the current command response. In print mode, shutdown is a no-op as the process exits automatically when prompts complete.
 
 ### Fixed
 

packages/coding-agent/docs/extensions.md

@@ -556,6 +556,24 @@ Access to models and API keys.
 
 Control flow helpers.
 
+### ctx.shutdown()
+
+Request a graceful shutdown of pi.
+
+- **Interactive mode:** Deferred until the agent becomes idle (after processing all queued steering and follow-up messages).
+- **RPC mode:** Deferred until the next idle state (after completing the current command response, when waiting for the next command).
+- **Print mode:** No-op. The process exits automatically when all prompts are processed.
+
+Emits `session_shutdown` event to all extensions before exiting. Available in all contexts (event handlers, tools, commands, shortcuts).
+
+```typescript
+pi.on("tool_call", (event, ctx) => {
+  if (isFatal(event.input)) {
+    ctx.shutdown();
+  }
+});
+```
+
 ## ExtensionCommandContext
 
 Command handlers receive `ExtensionCommandContext`, which extends `ExtensionContext` with session control methods. These are only available in commands because they can deadlock if called from event handlers.

packages/coding-agent/examples/extensions/shutdown-command.ts

@@ -0,0 +1,63 @@
+/**
+ * Shutdown Command Extension
+ *
+ * Adds a /quit command that allows extensions to trigger clean shutdown.
+ * Demonstrates how extensions can use ctx.shutdown() to exit pi cleanly.
+ */
+
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import { Type } from "@sinclair/typebox";
+
+export default function (pi: ExtensionAPI) {
+	// Register a /quit command that cleanly exits pi
+	pi.registerCommand("quit", {
+		description: "Exit pi cleanly",
+		handler: async (_args, ctx) => {
+			await ctx.shutdown();
+		},
+	});
+
+	// You can also create a tool that shuts down after completing work
+	pi.registerTool({
+		name: "finish_and_exit",
+		label: "Finish and Exit",
+		description: "Complete a task and exit pi",
+		parameters: Type.Object({}),
+		async execute(_toolCallId, _params, _onUpdate, ctx, _signal) {
+			// Do any final work here...
+			// Then shutdown
+			await ctx.shutdown();
+
+			// This return won't be reached, but required by type
+			return {
+				content: [{ type: "text", text: "Shutting down..." }],
+				details: {},
+			};
+		},
+	});
+
+	// You could also create a more complex tool with parameters
+	pi.registerTool({
+		name: "deploy_and_exit",
+		label: "Deploy and Exit",
+		description: "Deploy the application and exit pi",
+		parameters: Type.Object({
+			environment: Type.String({ description: "Target environment (e.g., production, staging)" }),
+		}),
+		async execute(_toolCallId, params, onUpdate, ctx, _signal) {
+			onUpdate?.({ content: [{ type: "text", text: `Deploying to ${params.environment}...` }], details: {} });
+
+			// Example deployment logic
+			// const result = await pi.exec("npm", ["run", "deploy", params.environment], { signal });
+
+			// On success, shutdown
+			onUpdate?.({ content: [{ type: "text", text: "Deployment complete, exiting..." }], details: {} });
+			await ctx.shutdown();
+
+			return {
+				content: [{ type: "text", text: "Done!" }],
+				details: { environment: params.environment },
+			};
+		},
+	});
+}

packages/coding-agent/src/core/extensions/index.ts

@@ -8,7 +8,13 @@ export {
 	loadExtensionFromFactory,
 	loadExtensions,
 } from "./loader.js";
-export type { BranchHandler, ExtensionErrorListener, NavigateTreeHandler, NewSessionHandler } from "./runner.js";
+export type {
+	BranchHandler,
+	ExtensionErrorListener,
+	NavigateTreeHandler,
+	NewSessionHandler,
+	ShutdownHandler,
+} from "./runner.js";
 export { ExtensionRunner } from "./runner.js";
 export type {
 	AgentEndEvent,

packages/coding-agent/src/core/extensions/runner.ts

@@ -55,6 +55,22 @@ export type NavigateTreeHandler = (
 	options?: { summarize?: boolean },
 ) => Promise<{ cancelled: boolean }>;
 
+export type ShutdownHandler = () => void;
+
+/**
+ * Helper function to emit session_shutdown event to extensions.
+ * Returns true if the event was emitted, false if there were no handlers.
+ */
+export async function emitSessionShutdownEvent(extensionRunner: ExtensionRunner | undefined): Promise<boolean> {
+	if (extensionRunner?.hasHandlers("session_shutdown")) {
+		await extensionRunner.emit({
+			type: "session_shutdown",
+		});
+		return true;
+	}
+	return false;
+}
+
 const noOpUIContext: ExtensionUIContext = {
 	select: async () => undefined,
 	confirm: async () => false,
@@ -91,6 +107,7 @@ export class ExtensionRunner {
 	private newSessionHandler: NewSessionHandler = async () => ({ cancelled: false });
 	private branchHandler: BranchHandler = async () => ({ cancelled: false });
 	private navigateTreeHandler: NavigateTreeHandler = async () => ({ cancelled: false });
+	private shutdownHandler: ShutdownHandler = () => {};
 
 	constructor(
 		extensions: Extension[],
@@ -129,6 +146,7 @@ export class ExtensionRunner {
 		this.isIdleFn = contextActions.isIdle;
 		this.abortFn = contextActions.abort;
 		this.hasPendingMessagesFn = contextActions.hasPendingMessages;
+		this.shutdownHandler = contextActions.shutdown;
 
 		// Command context actions (optional, only for interactive mode)
 		if (commandContextActions) {
@@ -137,7 +155,6 @@ export class ExtensionRunner {
 			this.branchHandler = commandContextActions.branch;
 			this.navigateTreeHandler = commandContextActions.navigateTree;
 		}
-
 		this.uiContext = uiContext ?? noOpUIContext;
 	}
 
@@ -282,6 +299,7 @@ export class ExtensionRunner {
 			isIdle: () => this.isIdleFn(),
 			abort: () => this.abortFn(),
 			hasPendingMessages: () => this.hasPendingMessagesFn(),
+			shutdown: () => this.shutdownHandler(),
 		};
 	}
 

packages/coding-agent/src/core/extensions/types.ts

@@ -175,6 +175,8 @@ export interface ExtensionContext {
 	abort(): void;
 	/** Whether there are queued messages waiting */
 	hasPendingMessages(): boolean;
+	/** Gracefully shutdown pi and exit. Available in all contexts. */
+	shutdown(): void;
 }
 
 /**
@@ -775,6 +777,7 @@ export interface ExtensionContextActions {
 	isIdle: () => boolean;
 	abort: () => void;
 	hasPendingMessages: () => boolean;
+	shutdown: () => void;
 }
 
 /**

packages/coding-agent/src/core/sdk.ts

@@ -527,6 +527,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
 		abort: () => {
 			session.abort();
 		},
+		shutdown: () => {},
 	}));
 
 	// Create tool registry mapping name -> tool (for extension getTools/setTools)

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

@@ -174,6 +174,9 @@ export class InteractiveMode {
 	// Messages queued while compaction is running
 	private compactionQueuedMessages: CompactionQueuedMessage[] = [];
 
+	// Shutdown state
+	private shutdownRequested = false;
+
 	// Extension UI state
 	private extensionSelector: ExtensionSelectorComponent | undefined = undefined;
 	private extensionInput: ExtensionInputComponent | undefined = undefined;
@@ -628,6 +631,9 @@ export class InteractiveMode {
 				isIdle: () => !this.session.isStreaming,
 				abort: () => this.session.abort(),
 				hasPendingMessages: () => this.session.pendingMessageCount > 0,
+				shutdown: () => {
+					this.shutdownRequested = true;
+				},
 			},
 			// ExtensionCommandContextActions - for ctx.* in command handlers
 			{
@@ -760,6 +766,9 @@ export class InteractiveMode {
 			isIdle: () => !this.session.isStreaming,
 			abort: () => this.session.abort(),
 			hasPendingMessages: () => this.session.pendingMessageCount > 0,
+			shutdown: () => {
+				this.shutdownRequested = true;
+			},
 		});
 
 		// Set up the extension shortcut handler on the default editor
@@ -1617,6 +1626,9 @@ export class InteractiveMode {
 					this.streamingMessage = undefined;
 				}
 				this.pendingTools.clear();
+
+				await this.checkShutdownRequested();
+
 				this.ui.requestRender();
 				break;
 
@@ -1930,7 +1942,12 @@ export class InteractiveMode {
 	 * Gracefully shutdown the agent.
 	 * Emits shutdown event to extensions, then exits.
 	 */
+	private isShuttingDown = false;
+
 	private async shutdown(): Promise<void> {
+		if (this.isShuttingDown) return;
+		this.isShuttingDown = true;
+
 		// Emit shutdown event to extensions
 		const extensionRunner = this.session.extensionRunner;
 		if (extensionRunner?.hasHandlers("session_shutdown")) {
@@ -1943,6 +1960,14 @@ export class InteractiveMode {
 		process.exit(0);
 	}
 
+	/**
+	 * Check if shutdown was requested and perform shutdown if so.
+	 */
+	private async checkShutdownRequested(): Promise<void> {
+		if (!this.shutdownRequested) return;
+		await this.shutdown();
+	}
+
 	private handleCtrlZ(): void {
 		// Set up handler to restore TUI when resumed
 		process.once("SIGCONT", () => {

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

@@ -66,6 +66,7 @@ export async function runPrintMode(session: AgentSession, options: PrintModeOpti
 				isIdle: () => !session.isStreaming,
 				abort: () => session.abort(),
 				hasPendingMessages: () => session.pendingMessageCount > 0,
+				shutdown: () => {},
 			},
 			// ExtensionCommandContextActions - commands invokable via prompt("/command")
 			{

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

@@ -63,6 +63,9 @@ export async function runRpcMode(session: AgentSession): Promise<never> {
 		{ resolve: (value: any) => void; reject: (error: Error) => void }
 	>();
 
+	// Shutdown request flag
+	let shutdownRequested = false;
+
 	/** Helper for dialog methods with signal/timeout support */
 	function createDialogPromise<T>(
 		opts: ExtensionUIDialogOptions | undefined,
@@ -265,6 +268,9 @@ export async function runRpcMode(session: AgentSession): Promise<never> {
 				isIdle: () => !session.isStreaming,
 				abort: () => session.abort(),
 				hasPendingMessages: () => session.pendingMessageCount > 0,
+				shutdown: () => {
+					shutdownRequested = true;
+				},
 			},
 			// ExtensionCommandContextActions - commands invokable via prompt("/command")
 			{
@@ -515,6 +521,22 @@ export async function runRpcMode(session: AgentSession): Promise<never> {
 		}
 	};
 
+	/**
+	 * Check if shutdown was requested and perform shutdown if so.
+	 * Called after handling each command when waiting for the next command.
+	 */
+	async function checkShutdownRequested(): Promise<void> {
+		if (!shutdownRequested) return;
+
+		if (extensionRunner?.hasHandlers("session_shutdown")) {
+			await extensionRunner.emit({ type: "session_shutdown" });
+		}
+
+		// Close readline interface to stop waiting for input
+		rl.close();
+		process.exit(0);
+	}
+
 	// Listen for JSON input
 	const rl = readline.createInterface({
 		input: process.stdin,
@@ -541,6 +563,9 @@ export async function runRpcMode(session: AgentSession): Promise<never> {
 			const command = parsed as RpcCommand;
 			const response = await handleCommand(command);
 			output(response);
+
+			// Check for deferred shutdown request (idle between commands)
+			await checkShutdownRequested();
 		} catch (e: any) {
 			output(error(undefined, "parse", `Failed to parse command: ${e.message}`));
 		}

packages/coding-agent/test/compaction-extensions.test.ts

@@ -121,6 +121,7 @@ describe.skipIf(!API_KEY)("Compaction extensions", () => {
 				isIdle: () => !session.isStreaming,
 				abort: () => session.abort(),
 				hasPendingMessages: () => session.pendingMessageCount > 0,
+				shutdown: () => {},
 			},
 		);