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

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}`));
 		}