From ff78ac2f84d0353a3ad4206f8b57082f810d8822 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Wed, 31 Dec 2025 02:55:45 +0100 Subject: [PATCH] Replace custom tool dispose() with shutdown session event Breaking change: CustomAgentTool.dispose() removed. Use onSession with reason 'shutdown' instead for cleanup. - Add 'shutdown' to SessionEvent.reason for custom tools - Remove dispose() method from CustomAgentTool interface - Make emitToolSessionEvent() public on AgentSession - Emit shutdown event to tools in InteractiveMode.shutdown() - Update custom-tools.md with new API and examples --- packages/coding-agent/docs/custom-tools.md | 53 +++++++++++-------- .../coding-agent/src/core/agent-session.ts | 14 ++--- .../src/core/custom-tools/types.ts | 8 ++- .../src/modes/interactive/interactive-mode.ts | 5 +- 4 files changed, 46 insertions(+), 34 deletions(-) diff --git a/packages/coding-agent/docs/custom-tools.md b/packages/coding-agent/docs/custom-tools.md index 854a9619..d10861b4 100644 --- a/packages/coding-agent/docs/custom-tools.md +++ b/packages/coding-agent/docs/custom-tools.md @@ -82,7 +82,7 @@ Custom tools can import from these packages (automatically resolved by pi): | Package | Purpose | |---------|---------| | `@sinclair/typebox` | Schema definitions (`Type.Object`, `Type.String`, etc.) | -| `@mariozechner/pi-coding-agent` | Types (`CustomToolFactory`, `ToolSessionEvent`, etc.) | +| `@mariozechner/pi-coding-agent` | Types (`CustomToolFactory`, `ToolSessionEvent` (alias for `SessionEvent`), etc.) | | `@mariozechner/pi-ai` | AI utilities (`StringEnum` for Google-compatible enums) | | `@mariozechner/pi-tui` | TUI components (`Text`, `Box`, etc. for custom rendering) | @@ -116,14 +116,17 @@ const factory: CustomToolFactory = (pi) => ({ }, // Optional: Session lifecycle callback - onSession(event) { /* reconstruct state from entries */ }, + onSession(event) { + if (event.reason === "shutdown") { + // Cleanup resources (close connections, save state, etc.) + return; + } + // Reconstruct state from entries for other events + }, // Optional: Custom rendering renderCall(args, theme) { /* return Component */ }, renderResult(result, options, theme) { /* return Component */ }, - - // Optional: Cleanup on session end - dispose() { /* save state, close connections */ }, }); export default factory; @@ -139,15 +142,18 @@ The factory receives a `ToolAPI` object (named `pi` by convention): interface ToolAPI { cwd: string; // Current working directory exec(command: string, args: string[], options?: ExecOptions): Promise; - ui: { - select(title: string, options: string[]): Promise; - confirm(title: string, message: string): Promise; - input(title: string, placeholder?: string): Promise; - notify(message: string, type?: "info" | "warning" | "error"): void; - }; + ui: ToolUIContext; hasUI: boolean; // false in --print or --mode rpc } +interface ToolUIContext { + select(title: string, options: string[]): Promise; + confirm(title: string, message: string): Promise; + input(title: string, placeholder?: string): Promise; + notify(message: string, type?: "info" | "warning" | "error"): void; + custom(component: Component & { dispose?(): void }): { close: () => void; requestRender: () => void }; +} + interface ExecOptions { signal?: AbortSignal; // Cancel the process timeout?: number; // Timeout in milliseconds @@ -182,11 +188,11 @@ async execute(toolCallId, params, signal) { Tools can implement `onSession` to react to session changes: ```typescript -interface ToolSessionEvent { +interface SessionEvent { entries: SessionEntry[]; // All session entries - sessionFile: string | null; // Current session file - previousSessionFile: string | null; // Previous session file - reason: "start" | "switch" | "branch" | "new"; + sessionFile: string | undefined; // Current session file (undefined with --no-session) + previousSessionFile: string | undefined; // Previous session file + reason: "start" | "switch" | "branch" | "new" | "tree"; } ``` @@ -195,6 +201,8 @@ interface ToolSessionEvent { - `switch`: User switched to a different session (`/resume`) - `branch`: User branched from a previous message (`/branch`) - `new`: User started a new session (`/new`) +- `tree`: User navigated to a different point in the session tree (`/tree`) +- `shutdown`: Process is exiting (Ctrl+C, Ctrl+D, or SIGTERM) - use to cleanup resources ### State Management Pattern @@ -387,13 +395,16 @@ const factory: CustomToolFactory = (pi) => { // Shared state let connection = null; + const handleSession = (event: ToolSessionEvent) => { + if (event.reason === "shutdown") { + connection?.close(); + } + }; + return [ - { name: "db_connect", ... }, - { name: "db_query", ... }, - { - name: "db_close", - dispose() { connection?.close(); } - }, + { name: "db_connect", onSession: handleSession, ... }, + { name: "db_query", onSession: handleSession, ... }, + { name: "db_close", onSession: handleSession, ... }, ]; }; ``` diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 58fe2def..6916d8e0 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -698,7 +698,7 @@ export class AgentSession { } // Emit session event to custom tools - await this._emitToolSessionEvent("new", previousSessionFile); + await this.emitToolSessionEvent("new", previousSessionFile); return true; } @@ -1473,7 +1473,7 @@ export class AgentSession { } // Emit session event to custom tools - await this._emitToolSessionEvent("switch", previousSessionFile); + await this.emitToolSessionEvent("switch", previousSessionFile); this.agent.replaceMessages(sessionContext.messages); @@ -1550,7 +1550,7 @@ export class AgentSession { } // Emit session event to custom tools (with reason "branch") - await this._emitToolSessionEvent("branch", previousSessionFile); + await this.emitToolSessionEvent("branch", previousSessionFile); if (!skipConversationRestore) { this.agent.replaceMessages(sessionContext.messages); @@ -1720,7 +1720,7 @@ export class AgentSession { } // Emit to custom tools - await this._emitToolSessionEvent("tree", this.sessionFile); + await this.emitToolSessionEvent("tree", this.sessionFile); this._branchSummaryAbortController = undefined; return { editorText, cancelled: false, summaryEntry }; @@ -1875,11 +1875,11 @@ export class AgentSession { /** * Emit session event to all custom tools. - * Called on session switch, branch, and clear. + * Called on session switch, branch, tree navigation, and shutdown. */ - private async _emitToolSessionEvent( + async emitToolSessionEvent( reason: ToolSessionEvent["reason"], - previousSessionFile: string | undefined, + previousSessionFile?: string | undefined, ): Promise { const event: ToolSessionEvent = { entries: this.sessionManager.getEntries(), diff --git a/packages/coding-agent/src/core/custom-tools/types.ts b/packages/coding-agent/src/core/custom-tools/types.ts index 28192b05..7bf99407 100644 --- a/packages/coding-agent/src/core/custom-tools/types.ts +++ b/packages/coding-agent/src/core/custom-tools/types.ts @@ -40,10 +40,10 @@ export interface SessionEvent { entries: SessionEntry[]; /** Current session file path, or undefined in --no-session mode */ sessionFile: string | undefined; - /** Previous session file path, or undefined for "start" and "new" */ + /** Previous session file path, or undefined for "start", "new", and "shutdown" */ previousSessionFile: string | undefined; /** Reason for the session event */ - reason: "start" | "switch" | "branch" | "new" | "tree"; + reason: "start" | "switch" | "branch" | "new" | "tree" | "shutdown"; } /** Rendering options passed to renderResult */ @@ -85,14 +85,12 @@ export interface RenderResultOptions { */ export interface CustomAgentTool extends AgentTool { - /** Called on session start/switch/branch/clear - use to reconstruct state from entries */ + /** Called on session lifecycle events - use to reconstruct state or cleanup resources */ onSession?: (event: SessionEvent) => void | Promise; /** Custom rendering for tool call display - return a Component */ renderCall?: (args: Static, theme: Theme) => Component; /** Custom rendering for tool result display - return a Component */ renderResult?: (result: AgentToolResult, options: RenderResultOptions, theme: Theme) => Component; - /** Called when session ends - cleanup resources */ - dispose?: () => Promise | void; } /** Factory function that creates a custom tool or array of tools */ diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 17fef8c0..b6ea952d 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -1239,7 +1239,7 @@ export class InteractiveMode { /** * Gracefully shutdown the agent. - * Emits shutdown event to hooks, then exits. + * Emits shutdown event to hooks and tools, then exits. */ private async shutdown(): Promise { // Emit shutdown event to hooks @@ -1250,6 +1250,9 @@ export class InteractiveMode { }); } + // Emit shutdown event to custom tools + await this.session.emitToolSessionEvent("shutdown"); + this.stop(); process.exit(0); }