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
This commit is contained in:
Mario Zechner 2025-12-31 02:55:45 +01:00
parent 450d77fb79
commit ff78ac2f84
4 changed files with 46 additions and 34 deletions

View file

@ -82,7 +82,7 @@ Custom tools can import from these packages (automatically resolved by pi):
| Package | Purpose | | Package | Purpose |
|---------|---------| |---------|---------|
| `@sinclair/typebox` | Schema definitions (`Type.Object`, `Type.String`, etc.) | | `@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-ai` | AI utilities (`StringEnum` for Google-compatible enums) |
| `@mariozechner/pi-tui` | TUI components (`Text`, `Box`, etc. for custom rendering) | | `@mariozechner/pi-tui` | TUI components (`Text`, `Box`, etc. for custom rendering) |
@ -116,14 +116,17 @@ const factory: CustomToolFactory = (pi) => ({
}, },
// Optional: Session lifecycle callback // 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 // Optional: Custom rendering
renderCall(args, theme) { /* return Component */ }, renderCall(args, theme) { /* return Component */ },
renderResult(result, options, theme) { /* return Component */ }, renderResult(result, options, theme) { /* return Component */ },
// Optional: Cleanup on session end
dispose() { /* save state, close connections */ },
}); });
export default factory; export default factory;
@ -139,15 +142,18 @@ The factory receives a `ToolAPI` object (named `pi` by convention):
interface ToolAPI { interface ToolAPI {
cwd: string; // Current working directory cwd: string; // Current working directory
exec(command: string, args: string[], options?: ExecOptions): Promise<ExecResult>; exec(command: string, args: string[], options?: ExecOptions): Promise<ExecResult>;
ui: { ui: ToolUIContext;
select(title: string, options: string[]): Promise<string | null>;
confirm(title: string, message: string): Promise<boolean>;
input(title: string, placeholder?: string): Promise<string | null>;
notify(message: string, type?: "info" | "warning" | "error"): void;
};
hasUI: boolean; // false in --print or --mode rpc hasUI: boolean; // false in --print or --mode rpc
} }
interface ToolUIContext {
select(title: string, options: string[]): Promise<string | undefined>;
confirm(title: string, message: string): Promise<boolean>;
input(title: string, placeholder?: string): Promise<string | undefined>;
notify(message: string, type?: "info" | "warning" | "error"): void;
custom(component: Component & { dispose?(): void }): { close: () => void; requestRender: () => void };
}
interface ExecOptions { interface ExecOptions {
signal?: AbortSignal; // Cancel the process signal?: AbortSignal; // Cancel the process
timeout?: number; // Timeout in milliseconds timeout?: number; // Timeout in milliseconds
@ -182,11 +188,11 @@ async execute(toolCallId, params, signal) {
Tools can implement `onSession` to react to session changes: Tools can implement `onSession` to react to session changes:
```typescript ```typescript
interface ToolSessionEvent { interface SessionEvent {
entries: SessionEntry[]; // All session entries entries: SessionEntry[]; // All session entries
sessionFile: string | null; // Current session file sessionFile: string | undefined; // Current session file (undefined with --no-session)
previousSessionFile: string | null; // Previous session file previousSessionFile: string | undefined; // Previous session file
reason: "start" | "switch" | "branch" | "new"; reason: "start" | "switch" | "branch" | "new" | "tree";
} }
``` ```
@ -195,6 +201,8 @@ interface ToolSessionEvent {
- `switch`: User switched to a different session (`/resume`) - `switch`: User switched to a different session (`/resume`)
- `branch`: User branched from a previous message (`/branch`) - `branch`: User branched from a previous message (`/branch`)
- `new`: User started a new session (`/new`) - `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 ### State Management Pattern
@ -387,13 +395,16 @@ const factory: CustomToolFactory = (pi) => {
// Shared state // Shared state
let connection = null; let connection = null;
const handleSession = (event: ToolSessionEvent) => {
if (event.reason === "shutdown") {
connection?.close();
}
};
return [ return [
{ name: "db_connect", ... }, { name: "db_connect", onSession: handleSession, ... },
{ name: "db_query", ... }, { name: "db_query", onSession: handleSession, ... },
{ { name: "db_close", onSession: handleSession, ... },
name: "db_close",
dispose() { connection?.close(); }
},
]; ];
}; };
``` ```

View file

@ -698,7 +698,7 @@ export class AgentSession {
} }
// Emit session event to custom tools // Emit session event to custom tools
await this._emitToolSessionEvent("new", previousSessionFile); await this.emitToolSessionEvent("new", previousSessionFile);
return true; return true;
} }
@ -1473,7 +1473,7 @@ export class AgentSession {
} }
// Emit session event to custom tools // Emit session event to custom tools
await this._emitToolSessionEvent("switch", previousSessionFile); await this.emitToolSessionEvent("switch", previousSessionFile);
this.agent.replaceMessages(sessionContext.messages); this.agent.replaceMessages(sessionContext.messages);
@ -1550,7 +1550,7 @@ export class AgentSession {
} }
// Emit session event to custom tools (with reason "branch") // Emit session event to custom tools (with reason "branch")
await this._emitToolSessionEvent("branch", previousSessionFile); await this.emitToolSessionEvent("branch", previousSessionFile);
if (!skipConversationRestore) { if (!skipConversationRestore) {
this.agent.replaceMessages(sessionContext.messages); this.agent.replaceMessages(sessionContext.messages);
@ -1720,7 +1720,7 @@ export class AgentSession {
} }
// Emit to custom tools // Emit to custom tools
await this._emitToolSessionEvent("tree", this.sessionFile); await this.emitToolSessionEvent("tree", this.sessionFile);
this._branchSummaryAbortController = undefined; this._branchSummaryAbortController = undefined;
return { editorText, cancelled: false, summaryEntry }; return { editorText, cancelled: false, summaryEntry };
@ -1875,11 +1875,11 @@ export class AgentSession {
/** /**
* Emit session event to all custom tools. * 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"], reason: ToolSessionEvent["reason"],
previousSessionFile: string | undefined, previousSessionFile?: string | undefined,
): Promise<void> { ): Promise<void> {
const event: ToolSessionEvent = { const event: ToolSessionEvent = {
entries: this.sessionManager.getEntries(), entries: this.sessionManager.getEntries(),

View file

@ -40,10 +40,10 @@ export interface SessionEvent {
entries: SessionEntry[]; entries: SessionEntry[];
/** Current session file path, or undefined in --no-session mode */ /** Current session file path, or undefined in --no-session mode */
sessionFile: string | undefined; 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; previousSessionFile: string | undefined;
/** Reason for the session event */ /** Reason for the session event */
reason: "start" | "switch" | "branch" | "new" | "tree"; reason: "start" | "switch" | "branch" | "new" | "tree" | "shutdown";
} }
/** Rendering options passed to renderResult */ /** Rendering options passed to renderResult */
@ -85,14 +85,12 @@ export interface RenderResultOptions {
*/ */
export interface CustomAgentTool<TParams extends TSchema = TSchema, TDetails = any> export interface CustomAgentTool<TParams extends TSchema = TSchema, TDetails = any>
extends AgentTool<TParams, TDetails> { extends AgentTool<TParams, TDetails> {
/** 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<void>; onSession?: (event: SessionEvent) => void | Promise<void>;
/** Custom rendering for tool call display - return a Component */ /** Custom rendering for tool call display - return a Component */
renderCall?: (args: Static<TParams>, theme: Theme) => Component; renderCall?: (args: Static<TParams>, theme: Theme) => Component;
/** Custom rendering for tool result display - return a Component */ /** Custom rendering for tool result display - return a Component */
renderResult?: (result: AgentToolResult<TDetails>, options: RenderResultOptions, theme: Theme) => Component; renderResult?: (result: AgentToolResult<TDetails>, options: RenderResultOptions, theme: Theme) => Component;
/** Called when session ends - cleanup resources */
dispose?: () => Promise<void> | void;
} }
/** Factory function that creates a custom tool or array of tools */ /** Factory function that creates a custom tool or array of tools */

View file

@ -1239,7 +1239,7 @@ export class InteractiveMode {
/** /**
* Gracefully shutdown the agent. * Gracefully shutdown the agent.
* Emits shutdown event to hooks, then exits. * Emits shutdown event to hooks and tools, then exits.
*/ */
private async shutdown(): Promise<void> { private async shutdown(): Promise<void> {
// Emit shutdown event to hooks // 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(); this.stop();
process.exit(0); process.exit(0);
} }