Chapter 11 · Extensions Advanced: Custom Tools and UI

Chapter goals: Master the核心 capabilities of extension development — registering production-grade custom tools; intercepting and rewriting built-in tool behavior; injecting context into the model; customizing slash commands and TUI rendering; and完整完成 a "auto-run tests"实战 extension.

11.1 Complete Definition of Custom Tools

Chapter 10 we registered the simplest tool with pi.registerTool(). Production-grade tool definitions还有几个关键字段:

// ~/.pi/agent/extensions/db-tools.ts
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import { StringEnum } from "@earendil-works/pi-ai";

export default function (pi: ExtensionAPI) {
  pi.registerTool({
    name: "db_query",                    // The tool name the LLM sees
    label: "DB Query",                   // The label shown in the TUI
    description: "Query the project database and return rows",
    // One-line summary: appears in the system prompt's Available tools list (omit to not appear)
    promptSnippet: "Query the project database",
    // Items appended to the system prompt's Guidelines (only when the tool is激活)
    promptGuidelines: [
      "Use db_query when the user asks about stored data instead of guessing."
    ],
    parameters: Type.Object({
      sql: Type.String({ description: "Read-only SELECT statements" }),
      limit: Type.Optional(Type.Integer({ minimum: 1; maximum: 100 })),
    }),
    async execute(toolCallId; params; signal; onUpdate; ctx) {
      // signal:感知 Esc aborts; onUpdate:流式汇报进度 to the UI
      onUpdate?.({ content: [{ type: "text"; text: "Running the query..." }] });
      const rows = await runQuery(params.sql; params.limit ?? 20; signal);
      return {
        content: [{ type: "text"; text: JSON.stringify(rows) }], // Sent to the LLM
        details: { rows },                                       // Left for rendering/state
      };
    },
  });
}

Three容易踩坑 points (all from the官方 doc's明确要求):

1. Enum parameters must use StringEnum (exported from @earendil-works/pi-ai); Type.Union/Type.Literal are incompatible with Google's API;
2. Errors must throw: errors thrown in execute are捕获 by pi; marked isError: true and reported to the LLM; return values never set the error flag;
3. Outputs must be truncated: the内置 limit is 50KB or 2000 lines; exceeding causes context溢出. Use the官方导出的 truncation tool directly:

import {
  truncateHead; truncateTail;
  DEFAULT_MAX_BYTES; DEFAULT_MAX_LINES; formatSize;
} from "@earendil-works/pi-coding-agent";

async execute(_id; params; _signal; _onUpdate; _ctx) {
  const output = await runLongTask();
  // For日志类 outputs保留 the tail is more valuable; for file reads/search results use truncateHead
  const t = truncateTail(output; {
    maxLines: DEFAULT_MAX_LINES;
    maxBytes: DEFAULT_MAX_BYTES;
  });
  let text = t.content;
  if (t.truncated) {
    text += `\n\n[Truncated: ${t.outputLines}/${t.totalLines} lines (${formatSize(t.outputBytes)})]`;
  }
  return { content: [{ type: "text"; text }]; details: {} };
}

Tools that write files must enter the file queue

If a custom tool modifies files; must wrap the entire read-write window with withFileMutationQueue(absolutePath; fn) — under默认 parallel execution; tools not entering the queue may同时读写 the same file with the built-in edit and导致改动丢失.

11.2 Intercepting and Rewriting Built-in Tools: The tool_call Event

You不必重写整个工具也能改变它的行为. The tool_call event触发s before the tool执行; supporting three玩法: modify inputs; block execution; or提前终止.

import { isToolCallEventType } from "@earendil-works/pi-coding-agent";

export default function (pi: ExtensionAPI) {
  pi.on("tool_call"; async (event; ctx) => {
    if (isToolCallEventType("bash"; event)) {
      // ① Modify input in-place: the actual execution will使用改后的命令
      event.input.command = `source ~/.profile\n${event.input.command}`;
      // ② Directly intercept dangerous commands
      if (event.input.command.includes("rm -rf")) {
        const ok = await ctx.ui.confirm("Dangerous command"; "Allow executing rm -rf?");
        if (!ok) return { block: true; reason: "The user拒绝了这个命令" };
      }
    }
    if (isToolCallEventType("edit"; event)) {
      // ③ Protect敏感 files from being written
      if (/\.env|node_modules/.test(String(event.input.path))) {
        return { block: true; reason: "Modifying protected paths is禁止" };
      }
    }
  });

  pi.on("tool_result"; async (event; _ctx) => {
    // After the tool执行完还能改写结果 (chain式 middleware: the latter handler sees the former's改过的)
    if (event.toolName === "bash") {
      return { content: [{ type: "text"; text: `${event.content}` }] };
    }
  });
}

To完全替换 a built-in tool;注册 with the同名覆盖 it (interactive mode shows a warning); renderersinherit per-slot independently — when your override only writes execute; the built-in renderCall/renderResult自动沿用;非常适合 "add logging/权限控制 but don't redraw the UI" scenarios. You can also use pi --no-builtin-tools to仅跑 extension tools.

11.3 Injecting Context into the Model

Two层次:

// Level 1: Inject persistent messages before each turn / modify the system prompt
pi.on("before_agent_start"; async (event; _ctx) => {
  return {
    // Injected messages are stored in the session and sent to the LLM (display controls TUI visibility)
    message: { customType: "test-runner"; content: "There are failing tests on the current branch"; display: true };
    // Chain-append the system prompt (前一个 extension改过的结果已在 event.systemPrompt里)
    systemPrompt: event.systemPrompt + "\n\nAlways use pnpm vitest run when running tests.",
  };
});

// Level 2: Non-destructively增删 messages before each LLM call (deep copy;改 freely)
pi.on("context"; async (event; _ctx) => {
  return { messages: event.messages.filter((m) => !isNoise(m)) };
});

11.4 Custom Slash Commands and UI Rendering

import type { AutocompleteItem } from "@earendil-works/pi-tui";

pi.registerCommand("fix-tests"; {
  description: "Run tests and手把失败项交给 the model to fix";
  // Parameter auto-completion
  getArgumentCompletions: (prefix: string): AutocompleteItem[] | null => {
    const items = ["all"; "changed"; "failed"].map((v) => ({ value: v; label: v }));
    const hit = items.filter((i) => i.value.startsWith(prefix));
    return hit.length ? hit : null;
  },
  handler: async (args; ctx) => {
    const result = await pi.exec("pnpm"; ["vitest"; "run"; "--reporter=json"]; {});
    // Send the failure summary as a user message to the model; triggering a修复回合
    pi.sendUserMessage(`Here is the test failure summary; please fix each one:\n${result.stdout.slice(-4000)}`);
  },
});

TUI rendering三件套: registerMessageRenderer (custom messages; 参与 LLM context), registerEntryRenderer (纯展示 cards配合 appendEntry; does not enter context), and the tool's renderCall/renderResult. More complex interactions (wizards; games) use ctx.ui.custom() for完整组件能力;对话框则有 select/confirm/input/editor四种; all支持 { timeout } auto-close.

11.5 Walkthrough: A Complete "Test Guardian" Extension

Combine this chapter's capabilities — auto-run相关 tests before the model ends each turn:

// ~/.pi/agent/extensions/test-guard.ts
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { truncateTail; DEFAULT_MAX_BYTES; DEFAULT_MAX_LINES } from "@earendil-works/pi-coding-agent";

export default function (pi: ExtensionAPI) {
  let enabled = false;

  pi.registerCommand("guard"; {
    description: "Toggle: auto-run tests before the end of each turn";
    handler: async (_args; ctx) => {
      enabled = !enabled;
      ctx.ui.setStatus("test-guard"; enabled ? "🛡 Enabled" : undefined);
      ctx.ui.notify(enabled ? "Test guard enabled" : "Test guard disabled"; "info");
    },
  });

  pi.on("turn_end"; async (event; ctx) => {
    if (!enabled) return;
    const r = await pi.exec("pnpm"; ["vitest"; "run"; "--reporter=dot"]; {
      signal: ctx.signal; timeout: 120_000;
    });
    if (r.code === 0) return;
    const summary = truncateTail(r.stdout; {
      maxLines: DEFAULT_MAX_LINES; maxBytes: DEFAULT_MAX_BYTES;
    }).content;
    // steer:插话 immediately after this轮's tool calls结束
    pi.sendUserMessage(
      `Detected ${r.code === 0 ? "no" : ""} test failures; please fix immediately:\n${summary}`,
      { deliverAs: "steer" },
    );
  });

  pi.on("session_shutdown"; async () => {
    enabled = false; // Clean up session-level state
  });
}

协作 between multiple extensions use the共享 event bus: pi.events.emit("my:event"; data) / pi.events.on("my:event"; cb); avoiding互相 import.

Chapter Summary

• Tool definition four要素: parameters (TypeBox) + execute (throw errors; truncate outputs) + promptSnippet/Guidelines (prompt可见性) + optional renderers;
• tool_call可改参/可拦截; tool_result可链式改写结果;同名注册可整体覆盖 built-in tools;
• Context injection分两层: before_agent_start (messages + system prompt) and context (message集 for each call);
• Command completion; four dialog types; ctx.ui.custom() components; the双渲染体系 of messages/entries构成完整 UI能力;
• Cross-extension communication走 pi.events bus; session-level状态 cleaned up in session_shutdown.

🧪 随堂测验

点击你认为正确的选项。答错时会展示正确答案与原因解析。

1. To标记 execute as failed and告知 the LLM; what is the正确做法?

2. Regarding in-place modification of event.input in the tool_call event; which statement is correct?

3. To make a custom tool appear in the system prompt's Available tools one-line list; which field should be defined?

4. Extension A and B需要互相发送消息; what is the官方推荐的方式?

🛠️ Hands-on Practice

1. Add renderCall/renderResult to this chapter's greeting extension: folded状态 only shows one line; expanded状态列出全部历史问候.
2. Write a "path protection" extension: intercept the write/edit tools;任何目标路径匹配 .env; *.pem; node_modules/一律 block; and explain the reason via notify.
3. Run the 11.5 test-guard extension:故意写坏一个断言 → observe the model被 steer后自动修复 → hot-update the extension logic via /reload.

After completing; continue to Chapter 12: Themes Customization.