Custom tools with session lifecycle, examples for hooks and tools

- Custom tools: TypeScript modules that extend pi with new tools - Custom TUI rendering via renderCall/renderResult - User interaction via pi.ui (select, confirm, input, notify) - Session lifecycle via onSession callback for state reconstruction - Examples: todo.ts, question.ts, hello.ts - Hook examples: permission-gate, git-checkpoint, protected-paths - Session lifecycle centralized in AgentSession - Works across all modes (interactive, print, RPC) - Unified session event for hooks (replaces session_start/session_switch) - Box component added to pi-tui - Examples bundled in npm and binary releases Fixes #190
2026-04-15 17:00:59 +00:00 · 2025-12-17 16:03:23 +01:00 · 2025-12-17 16:03:23 +01:00 · e7097d911a
commit e7097d911a
parent 295f51b53f
33 changed files with 1926 additions and 117 deletions
--- a/packages/coding-agent/docs/custom-tools.md
+++ b/packages/coding-agent/docs/custom-tools.md
@ -0,0 +1,341 @@
+# Custom Tools
+
+Custom tools extend pi with new capabilities beyond the built-in read/write/edit/bash tools. They are TypeScript modules that define one or more tools with optional custom rendering for the TUI.
+
+## Quick Start
+
+Create a file `~/.pi/agent/tools/hello.ts`:
+
+```typescript
+import { Type } from "@sinclair/typebox";
+import type { CustomToolFactory } from "@mariozechner/pi-coding-agent";
+
+const factory: CustomToolFactory = (pi) => ({
+  name: "hello",
+  label: "Hello",
+  description: "A simple greeting tool",
+  parameters: Type.Object({
+    name: Type.String({ description: "Name to greet" }),
+  }),
+
+  async execute(toolCallId, params) {
+    return {
+      content: [{ type: "text", text: `Hello, ${params.name}!` }],
+      details: { greeted: params.name },
+    };
+  },
+});
+
+export default factory;
+```
+
+The tool is automatically discovered and available in your next pi session.
+
+## Tool Locations
+
+| Location | Scope | Auto-discovered |
+|----------|-------|-----------------|
+| `~/.pi/agent/tools/*.ts` | Global (all projects) | Yes |
+| `.pi/tools/*.ts` | Project-local | Yes |
+| `settings.json` `customTools` array | Configured paths | Yes |
+| `--tool <path>` CLI flag | One-off/debugging | No |
+
+**Priority:** Later sources win on name conflicts. CLI `--tool` takes highest priority.
+
+**Reserved names:** Custom tools cannot use built-in tool names (`read`, `write`, `edit`, `bash`, `grep`, `find`, `ls`).
+
+## Tool Definition
+
+```typescript
+import { Type } from "@sinclair/typebox";
+import { StringEnum } from "@mariozechner/pi-ai";
+import { Text } from "@mariozechner/pi-tui";
+import type { CustomToolFactory, ToolSessionEvent } from "@mariozechner/pi-coding-agent";
+
+const factory: CustomToolFactory = (pi) => ({
+  name: "my_tool",
+  label: "My Tool",
+  description: "What this tool does (be specific for LLM)",
+  parameters: Type.Object({
+    // Use StringEnum for string enums (Google API compatible)
+    action: StringEnum(["list", "add", "remove"] as const),
+    text: Type.Optional(Type.String()),
+  }),
+
+  async execute(toolCallId, params, signal, onUpdate) {
+    // signal - AbortSignal for cancellation
+    // onUpdate - Callback for streaming partial results
+    return {
+      content: [{ type: "text", text: "Result for LLM" }],
+      details: { /* structured data for rendering */ },
+    };
+  },
+
+  // Optional: Session lifecycle callback
+  onSession(event) { /* reconstruct state from entries */ },
+
+  // 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;
+```
+
+**Important:** Use `StringEnum` from `@mariozechner/pi-ai` instead of `Type.Union`/`Type.Literal` for string enums. The latter doesn't work with Google's API.
+
+## ToolAPI Object
+
+The factory receives a `ToolAPI` object (named `pi` by convention):
+
+```typescript
+interface ToolAPI {
+  cwd: string;  // Current working directory
+  exec(command: string, args: string[]): Promise<ExecResult>;
+  ui: {
+    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
+}
+```
+
+Always check `pi.hasUI` before using UI methods.
+
+## Session Lifecycle
+
+Tools can implement `onSession` to react to session changes:
+
+```typescript
+interface ToolSessionEvent {
+  entries: SessionEntry[];      // All session entries
+  sessionFile: string | null;   // Current session file
+  previousSessionFile: string | null;  // Previous session file
+  reason: "start" | "switch" | "branch" | "clear";
+}
+```
+
+**Reasons:**
+- `start`: Initial session load on startup
+- `switch`: User switched to a different session (`/session`)
+- `branch`: User branched from a previous message (`/branch`)
+- `clear`: User cleared the session (`/clear`)
+
+### State Management Pattern
+
+Tools that maintain state should store it in `details` of their results, not external files. This allows branching to work correctly, as the state is reconstructed from the session history.
+
+```typescript
+interface MyToolDetails {
+  items: string[];
+}
+
+const factory: CustomToolFactory = (pi) => {
+  // In-memory state
+  let items: string[] = [];
+
+  // Reconstruct state from session entries
+  const reconstructState = (event: ToolSessionEvent) => {
+    items = [];
+    for (const entry of event.entries) {
+      if (entry.type !== "message") continue;
+      const msg = entry.message;
+      if (msg.role !== "toolResult") continue;
+      if (msg.toolName !== "my_tool") continue;
+      
+      const details = msg.details as MyToolDetails | undefined;
+      if (details) {
+        items = details.items;
+      }
+    }
+  };
+
+  return {
+    name: "my_tool",
+    label: "My Tool",
+    description: "...",
+    parameters: Type.Object({ ... }),
+    
+    onSession: reconstructState,
+    
+    async execute(toolCallId, params) {
+      // Modify items...
+      items.push("new item");
+      
+      return {
+        content: [{ type: "text", text: "Added item" }],
+        // Store current state in details for reconstruction
+        details: { items: [...items] },
+      };
+    },
+  };
+};
+```
+
+This pattern ensures:
+- When user branches, state is correct for that point in history
+- When user switches sessions, state matches that session
+- When user clears, state resets
+
+## Custom Rendering
+
+Custom tools can provide `renderCall` and `renderResult` methods to control how they appear in the TUI. Both are optional.
+
+### How It Works
+
+Tool output is wrapped in a `Box` component that handles:
+- Padding (1 character horizontal, 1 line vertical)
+- Background color based on state (pending/success/error)
+
+Your render methods return `Component` instances (typically `Text`) that go inside this box. Use `Text(content, 0, 0)` since the Box handles padding.
+
+### renderCall
+
+Renders the tool call (before/during execution):
+
+```typescript
+renderCall(args, theme) {
+  let text = theme.fg("toolTitle", theme.bold("my_tool "));
+  text += theme.fg("muted", args.action);
+  if (args.text) {
+    text += " " + theme.fg("dim", `"${args.text}"`);
+  }
+  return new Text(text, 0, 0);
+}
+```
+
+Called when:
+- Tool call starts (may have partial args during streaming)
+- Args are updated during streaming
+
+### renderResult
+
+Renders the tool result:
+
+```typescript
+renderResult(result, { expanded, isPartial }, theme) {
+  const { details } = result;
+
+  // Handle streaming/partial results
+  if (isPartial) {
+    return new Text(theme.fg("warning", "Processing..."), 0, 0);
+  }
+
+  // Handle errors
+  if (details?.error) {
+    return new Text(theme.fg("error", `Error: ${details.error}`), 0, 0);
+  }
+
+  // Normal result
+  let text = theme.fg("success", "✓ ") + theme.fg("muted", "Done");
+  
+  // Support expanded view (Ctrl+O)
+  if (expanded && details?.items) {
+    for (const item of details.items) {
+      text += "\n" + theme.fg("dim", `  ${item}`);
+    }
+  }
+
+  return new Text(text, 0, 0);
+}
+```
+
+**Options:**
+- `expanded`: User pressed Ctrl+O to expand
+- `isPartial`: Result is from `onUpdate` (streaming), not final
+
+### Best Practices
+
+1. **Use `Text` with padding `(0, 0)`** - The Box handles padding
+2. **Use `\n` for multi-line content** - Not multiple Text components
+3. **Handle `isPartial`** - Show progress during streaming
+4. **Support `expanded`** - Show more detail when user requests
+5. **Use theme colors** - For consistent appearance
+6. **Keep it compact** - Show summary by default, details when expanded
+
+### Theme Colors
+
+```typescript
+// Foreground
+theme.fg("toolTitle", text)   // Tool names
+theme.fg("accent", text)      // Highlights
+theme.fg("success", text)     // Success
+theme.fg("error", text)       // Errors
+theme.fg("warning", text)     // Warnings
+theme.fg("muted", text)       // Secondary text
+theme.fg("dim", text)         // Tertiary text
+theme.fg("toolOutput", text)  // Output content
+
+// Styles
+theme.bold(text)
+theme.italic(text)
+```
+
+### Fallback Behavior
+
+If `renderCall` or `renderResult` is not defined or throws an error:
+- `renderCall`: Shows tool name
+- `renderResult`: Shows raw text output from `content`
+
+## Execute Function
+
+```typescript
+async execute(toolCallId, args, signal, onUpdate) {
+  // Type assertion for params (TypeBox schema doesn't flow through)
+  const params = args as { action: "list" | "add"; text?: string };
+
+  // Check for abort
+  if (signal?.aborted) {
+    return { content: [...], details: { status: "aborted" } };
+  }
+
+  // Stream progress
+  onUpdate?.({
+    content: [{ type: "text", text: "Working..." }],
+    details: { progress: 50 },
+  });
+
+  // Return final result
+  return {
+    content: [{ type: "text", text: "Done" }],  // Sent to LLM
+    details: { data: result },  // For rendering only
+  };
+}
+```
+
+## Multiple Tools from One File
+
+Return an array to share state between related tools:
+
+```typescript
+const factory: CustomToolFactory = (pi) => {
+  // Shared state
+  let connection = null;
+
+  return [
+    { name: "db_connect", ... },
+    { name: "db_query", ... },
+    {
+      name: "db_close",
+      dispose() { connection?.close(); }
+    },
+  ];
+};
+```
+
+## Examples
+
+See [`examples/custom-tools/todo.ts`](../examples/custom-tools/todo.ts) for a complete example with:
+- `onSession` for state reconstruction
+- Custom `renderCall` and `renderResult`
+- Proper branching support via details storage
+
+Test with:
+```bash
+pi --tool packages/coding-agent/examples/custom-tools/todo.ts
+```
--- a/packages/coding-agent/docs/hooks.md
+++ b/packages/coding-agent/docs/hooks.md
@ -43,8 +43,8 @@ A hook is a TypeScript file that exports a default function. The function receiv
 import type { HookAPI } from "@mariozechner/pi-coding-agent/hooks";

 export default function (pi: HookAPI) {
-  pi.on("session_start", async (event, ctx) => {
-    ctx.ui.notify(`Session: ${ctx.sessionFile ?? "ephemeral"}`, "info");
+  pi.on("session", async (event, ctx) => {
+    ctx.ui.notify(`Session ${event.reason}: ${ctx.sessionFile ?? "ephemeral"}`, "info");
  });
 }
 ```
@ -76,7 +76,7 @@ Hooks are loaded using [jiti](https://github.com/unjs/jiti), so TypeScript works
 ```
 pi starts
  │
-  ├─► session_start
+  ├─► session (reason: "start")
  │
  ▼
 user sends prompt ─────────────────────────────────────────┐
@ -98,36 +98,54 @@ user sends prompt ────────────────────
                                                           │
 user sends another prompt ◄────────────────────────────────┘

-user branches or switches session
+user branches (/branch)
  │
-  └─► session_switch
+  ├─► branch (BEFORE branch, can control)
+  └─► session (reason: "switch", AFTER branch)
+
+user switches session (/session)
+  │
+  └─► session (reason: "switch")
+
+user clears session (/clear)
+  │
+  └─► session (reason: "clear")
 ```

 A **turn** is one LLM response plus any tool calls. Complex tasks loop through multiple turns until the LLM responds without calling tools.

-### session_start
+### session

-Fired once when pi starts.
+Fired on startup and when session changes.

 ```typescript
-pi.on("session_start", async (event, ctx) => {
-  // ctx.sessionFile: string | null
-  // ctx.hasUI: boolean
+pi.on("session", async (event, ctx) => {
+  // event.entries: SessionEntry[] - all session entries
+  // event.sessionFile: string | null - current session file
+  // event.previousSessionFile: string | null - previous session file
+  // event.reason: "start" | "switch" | "clear"
 });
 ```

-### session_switch
+**Reasons:**
+- `start`: Initial session load on startup
+- `switch`: User switched sessions (`/session`) or branched (`/branch`)
+- `clear`: User cleared the session (`/clear`)

-Fired when session changes (`/branch` or session switch).
+### branch
+
+Fired BEFORE a branch happens. Can control branch behavior.

 ```typescript
-pi.on("session_switch", async (event, ctx) => {
-  // event.newSessionFile: string | null (null in --no-session mode)
-  // event.previousSessionFile: string | null (null in --no-session mode)
-  // event.reason: "branch" | "switch"
+pi.on("branch", async (event, ctx) => {
+  // event.targetTurnIndex: number
+  // event.entries: SessionEntry[]
+  return { skipConversationRestore: true }; // or undefined
 });
 ```

+Note: After branch completes, a `session` event fires with `reason: "switch"`.
+
 ### agent_start / agent_end

 Fired once per user prompt.
@ -192,18 +210,6 @@ pi.on("tool_result", async (event, ctx) => {
 });
 ```

-### branch
-
-Fired when user branches via `/branch`.
-
-```typescript
-pi.on("branch", async (event, ctx) => {
-  // event.targetTurnIndex: number
-  // event.entries: SessionEntry[]
-  return { skipConversationRestore: true }; // or undefined
-});
-```
-
 ## Context API

 Every event handler receives a context object with these methods:
@ -309,7 +315,9 @@ import * as fs from "node:fs";
 import type { HookAPI } from "@mariozechner/pi-coding-agent/hooks";

 export default function (pi: HookAPI) {
-  pi.on("session_start", async (event, ctx) => {
+  pi.on("session", async (event, ctx) => {
+    if (event.reason !== "start") return;
+    
    // Watch a trigger file
    const triggerFile = "/tmp/agent-trigger.txt";
    
@ -339,7 +347,9 @@ import * as http from "node:http";
 import type { HookAPI } from "@mariozechner/pi-coding-agent/hooks";

 export default function (pi: HookAPI) {
-  pi.on("session_start", async (event, ctx) => {
+  pi.on("session", async (event, ctx) => {
+    if (event.reason !== "start") return;
+    
    const server = http.createServer((req, res) => {
      let body = "";
      req.on("data", chunk => body += chunk);
@ -563,7 +573,7 @@ Key behaviors:
 Mode initialization:
  -> hookRunner.setUIContext(ctx, hasUI)
  -> hookRunner.setSessionFile(path)
-  -> hookRunner.emit({ type: "session_start" })
+  -> hookRunner.emit({ type: "session", reason: "start", ... })

 User sends prompt:
  -> AgentSession.prompt()
@ -581,9 +591,19 @@ User sends prompt:
     -> [repeat if more tool calls]
  -> hookRunner.emit({ type: "agent_end", messages })

-Branch or session switch:
-  -> AgentSession.branch() or AgentSession.switchSession()
-     -> hookRunner.emit({ type: "session_switch", ... })
+Branch:
+  -> AgentSession.branch()
+     -> hookRunner.emit({ type: "branch", ... })  # BEFORE branch
+     -> [branch happens]
+     -> hookRunner.emit({ type: "session", reason: "switch", ... })  # AFTER
+
+Session switch:
+  -> AgentSession.switchSession()
+     -> hookRunner.emit({ type: "session", reason: "switch", ... })
+
+Clear:
+  -> AgentSession.reset()
+     -> hookRunner.emit({ type: "session", reason: "clear", ... })
 ```

 ## UI Context by Mode