diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index bfa89814..b3179a1c 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/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 diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index 97a163e3..e47da49b 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/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. diff --git a/packages/coding-agent/examples/extensions/shutdown-command.ts b/packages/coding-agent/examples/extensions/shutdown-command.ts new file mode 100644 index 00000000..7cfdf56e --- /dev/null +++ b/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 }, + }; + }, + }); +} diff --git a/packages/coding-agent/src/core/extensions/index.ts b/packages/coding-agent/src/core/extensions/index.ts index 755f4ee6..ec2aa1e1 100644 --- a/packages/coding-agent/src/core/extensions/index.ts +++ b/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, diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index c81f7624..b417164a 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/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 { + 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(), }; } diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index e333ca04..168f6899 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/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; } /** diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index e2a04582..e04288bd 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/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) diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 3e2e8a19..02d7ca3a 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/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 { + 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 { + if (!this.shutdownRequested) return; + await this.shutdown(); + } + private handleCtrlZ(): void { // Set up handler to restore TUI when resumed process.once("SIGCONT", () => { diff --git a/packages/coding-agent/src/modes/print-mode.ts b/packages/coding-agent/src/modes/print-mode.ts index d49f4753..a7061b9f 100644 --- a/packages/coding-agent/src/modes/print-mode.ts +++ b/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") { diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index 4808cd51..1d044a74 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -63,6 +63,9 @@ export async function runRpcMode(session: AgentSession): Promise { { resolve: (value: any) => void; reject: (error: Error) => void } >(); + // Shutdown request flag + let shutdownRequested = false; + /** Helper for dialog methods with signal/timeout support */ function createDialogPromise( opts: ExtensionUIDialogOptions | undefined, @@ -265,6 +268,9 @@ export async function runRpcMode(session: AgentSession): Promise { 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 { } }; + /** + * 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 { + 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 { 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}`)); } diff --git a/packages/coding-agent/test/compaction-extensions.test.ts b/packages/coding-agent/test/compaction-extensions.test.ts index 344227a4..f997937e 100644 --- a/packages/coding-agent/test/compaction-extensions.test.ts +++ b/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: () => {}, }, );