第 11 讲扩展系统：自定义工具、事件拦截与 UI 集成

import { type ExtensionAPI, isToolCallEventType } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";

export default function (pi: ExtensionAPI) {
  // 1. 订阅事件
  pi.on("session_start", async (_event, ctx) => {
    ctx.ui.notify("Extension loaded!", "info");
  });

  // 2. 拦截 tool_call，可以 block
  pi.on("tool_call", async (event, ctx) => {
    if (isToolCallEventType("bash", event) && event.input.command.includes("rm -rf")) {
      const ok = await ctx.ui.confirm("Dangerous!", "Allow rm -rf?");
      if (!ok) return { block: true, reason: "Blocked by user" };
    }
  });

  // 3. 注册 LLM 可调用的自定义工具
  pi.registerTool({
    name: "greet", label: "Greet", description: "Greet someone by name",
    parameters: Type.Object({ name: Type.String({ description: "Name to greet" }) }),
    async execute(toolCallId, params, signal, onUpdate, ctx) {
      return { content: [{ type: "text", text: `Hello, ${params.name}!` }], details: {} };
    },
  });

  // 4. 注册斜杠命令
  pi.registerCommand("hello", {
    description: "Say hello",
    handler: async (args, ctx) => { ctx.ui.notify(`Hello ${args || "world"}!`, "info"); },
  });
}

/**
 * Discovery rules:
 * 1. Direct files: `extensions/*.ts` or `*.js` → load
 * 2. Subdirectory with index: `extensions/*/index.ts` or `index.js` → load
 * 3. Subdirectory with package.json: `extensions/*/package.json`
 *    with "pi" field → load what it declares
 *
 * No recursion beyond one level. Complex packages must use package.json manifest.
 */

{
  "name": "my-extension",
  "dependencies": { "zod": "^3.0.0", "chalk": "^5.0.0" },
  "pi": {
    "extensions": ["./src/index.ts"]
  }
}

if (beforeResult?.block) {
  const result = createErrorToolResult(beforeResult.reason || "Tool execution was blocked");
  if (beforeResult.terminate === true) {
    result.terminate = true;
  }
  return { kind: "immediate", result, isError: true };
}

/**
 * Returning `{ block: true }` prevents the tool from executing. The loop emits
 * an error tool result instead. `reason` becomes the text shown in that error
 * result. If omitted, a default blocked message is used.
 */
export interface BeforeToolCallResult {
  block?: boolean;
  reason?: string;
  /** Early termination only happens when every finalized tool result in the batch sets this to true. */
  terminate?: boolean;
}

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

export default function (pi: ExtensionAPI) {
  const dangerousPatterns = [/\brm\s+(-rf?|--recursive)/i, /\bsudo\b/i, /\b(chmod|chown)\b.*777/i];

  pi.on("tool_call", async (event, ctx) => {
    if (event.toolName !== "bash") return undefined;

    const command = event.input.command as string;
    const isDangerous = dangerousPatterns.some((p) => p.test(command));

    if (isDangerous) {
      if (!ctx.hasUI) {
        // In non-interactive mode, block by default
        return { block: true, reason: "Dangerous command blocked (no UI for confirmation)" };
      }
      const choice = await ctx.ui.select(`⚠️ Dangerous command:\n\n  ${command}\n\nAllow?`, ["Yes", "No"]);
      if (choice !== "Yes") {
        return { block: true, reason: "Blocked by user" };
      }
    }
    return undefined;
  });
}

const showSummaryUi = async (summary: string, ctx: ExtensionCommandContext) => {
  if (ctx.mode !== "tui") {          // 模式守卫：非 TUI 直接放弃渲染
    return;
  }
  await ctx.ui.custom((_tui, theme, _kb, done) => {
    const container = new Container();
    // …DynamicBorder + Markdown(summary) 组装滚动视图…
    return {
      render: (width: number) => container.render(width),
      invalidate: () => container.invalidate(),
      handleInput: (data: string) => {
        if (matchesKey(data, "enter") || matchesKey(data, "escape")) {
          done(undefined);      // 关闭全屏组件
        }
      },
    };
  });
};

pi.appendEntry("my-state", { count: 42 });
pi.appendEntry("status-card", { title: "Indexed files", count: 17 });

// Restore on reload / resume / fork
pi.on("session_start", async (_event, ctx) => {
  for (const entry of ctx.sessionManager.getEntries()) {
    if (entry.type === "custom" && entry.customType === "my-state") {
      // Reconstruct from entry.data
    }
  }
});

import { Box, Text } from "@earendil-works/pi-tui";

pi.registerEntryRenderer("status-card", (entry, { expanded }, theme) => {
  const data = entry.data as { title: string; count: number };
  const box = new Box(1, 1, (text) => theme.bg("customMessageBg", text));
  box.addChild(new Text(`${theme.bold(data.title)}: ${data.count}`));
  if (expanded) {
    box.addChild(new Text(theme.fg("dim", JSON.stringify(data, null, 2))));
  }
  return box;
});
pi.appendEntry("status-card", { title: "Indexed files", count: 17 });

const checkpoints = new Map<string, string>();
let currentEntryId: string | undefined;

pi.on("tool_result", async (_event, ctx) => {
  const leaf = ctx.sessionManager.getLeafEntry();
  if (leaf) currentEntryId = leaf.id;
});

pi.on("turn_start", async () => {
  // Create a git stash entry before LLM makes changes
  const { stdout } = await pi.exec("git", ["stash", "create"]);
  const ref = stdout.trim();
  if (ref && currentEntryId) {
    checkpoints.set(currentEntryId, ref);
  }
});

pi.on("session_before_fork", async (event, ctx) => {
  const ref = checkpoints.get(event.entryId);
  if (!ref) return;
  if (!ctx.hasUI) return;  // 非 TUI 不自动恢复
  const choice = await ctx.ui.select("Restore code state?", [
    "Yes, restore code to that point", "No, keep current code",
  ]);
  if (choice?.startsWith("Yes")) {
    await pi.exec("git", ["stash", "apply", ref]);
    ctx.ui.notify("Code restored to checkpoint", "info");
  }
});