Pi Agent 自定义工具开发
通过 Extension 注册自定义工具,让 AI 能够调用你编写的函数来完成特定任务。
工具注册基础
使用 pi.registerTool() 注册一个 AI 可调用的工具:
实例
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import { StringEnum } from "@earendil-works/pi-ai";
// 定义待办事项列表(内存存储)
let todos: string[] = [];
export default function (pi: ExtensionAPI) {
// 启动时从会话中恢复状态
pi.on("session_start", async (_event, ctx) => {
todos = [];
for (const entry of ctx.sessionManager.getBranch()) {
if (entry.type === "message"
&& entry.message.role === "toolResult"
&& entry.message.toolName === "todo") {
todos = entry.message.details?.todos ?? [];
}
}
});
pi.registerTool({
name: "todo",
label: "待办事项",
description: "管理项目的待办事项列表",
// 简短描述,显示在系统提示词的 Available tools 区域
promptSnippet: "管理项目待办事项:list 列表, add 添加, done 完成",
// 使用指南,追加到系统提示词的 Guidelines 区域
promptGuidelines: [
"使用 todo 工具进行任务规划,不要直接编辑待办文件",
],
parameters: Type.Object({
// StringEnum 确保与 Google API 兼容
action: StringEnum(["list", "add", "done"] as const),
text: Type.Optional(Type.String({
description: "添加时:事项描述。完成时:事项编号或描述"
})),
}),
// 兼容旧参数格式
prepareArguments(args) {
if (!args || typeof args !== "object") return args;
const input = args as { action?: string; oldAction?: string };
// 将旧字段映射到新字段
if (typeof input.oldAction === "string"
&& input.action === undefined) {
return { ...input, action: input.oldAction };
}
return args;
},
async execute(toolCallId, params, signal, onUpdate) {
if (signal?.aborted) {
return {
content: [{ type: "text", text: "操作已取消" }]
};
}
// 发送进度更新
onUpdate?.({
content: [{ type: "text", text: `正在处理: ${params.action}` }],
});
switch (params.action) {
case "list":
return {
content: [{
type: "text",
text: todos.length === 0
? "待办列表为空"
: todos.map((t, i) => `${i + 1}. ${t}`).join("\n"),
}],
details: { todos: [...todos] },
};
case "add":
if (!params.text) {
throw new Error("add 操作需要 text 参数");
}
todos.push(params.text);
return {
content: [{
type: "text",
text: `已添加: ${params.text}(共 ${todos.length} 项)`,
}],
details: { todos: [...todos] },
};
case "done":
if (!params.text) {
throw new Error("done 操作需要 text 参数");
}
const idx = todos.findIndex(
t => t.includes(params.text!)
);
if (idx === -1) {
return {
content: [{
type: "text",
text: `未找到匹配项: ${params.text}`
}],
details: { todos: [...todos] },
};
}
const removed = todos.splice(idx, 1)[0];
return {
content: [{
type: "text",
text: `已完成: ${removed}(剩余 ${todos.length} 项)`,
}],
details: { todos: [...todos] },
};
default:
throw new Error("未知操作");
}
},
});
}
工具定义详解
| 字段 | 类型 | 必填 | 说明 |
|---|---|---|---|
| name | string | 是 | 工具的唯一标识,AI 通过此名称调用 |
| label | string | 是 | 工具显示名 |
| description | string | 是 | 工具功能描述,AI 据此判断何时使用 |
| parameters | TypeBox Schema | 是 | 工具参数的类型定义 |
| promptSnippet | string | 否 | 在 Available tools 区域的单行简介 |
| promptGuidelines | string[] | 否 | 追加到系统提示 Guidelines 区域的指南 |
| prepareArguments | function | 否 | 在 Schema 验证前转换参数,用于兼容旧格式 |
| execute | async function | 是 | 工具的实际执行逻辑 |
promptGuidelines 中的每条指南都必须明确指出工具名称。不要写「使用此工具时…」,而要写「使用 my_tool 时…」,因为 AI 无法根据「此」来判断指的是哪个工具。
execute 函数详解
execute 函数接收以下参数:
| 参数 | 类型 | 说明 |
|---|---|---|
| toolCallId | string | 工具调用的唯一 ID |
| params | 泛型 T | 经过 Schema 验证的参数对象 |
| signal | AbortSignal | undefined | 取消信号,用户按 Escape 时触发 |
| onUpdate | function | 发送进度更新的回调 |
| ctx | ExtensionContext | 扩展上下文,提供 UI、会话等能力 |
返回值格式:
实例
// 发送给 LLM 的内容(必需)
content: [{ type: "text", text: "操作完成" }],
// 自定义详情数据,用于状态重建和渲染(可选)
details: { result: "..." },
// 嵌套 LLM 调用的用量统计(可选)
usage: nestedModelUsage,
// 终止标记:当同一批次所有工具都返回 terminate 时
// 跳过后续的 LLM 调用(可选)
terminate: true,
};
要标记工具执行失败,使用 throw new Error(),不要尝试在返回值中设置错误标记。只有抛出异常才能将 isError 设为 true 并通知 LLM 发生了错误。
StringEnum 与类型安全
使用 StringEnum 而不是 Type.Union/Type.Literal 来定义枚举参数:
实例
// 正确:Google API 兼容
parameters: Type.Object({
action: StringEnum(["list", "add", "done"] as const),
})
// 错误:Google API 不兼容
parameters: Type.Object({
action: Type.Union([
Type.Literal("list"),
Type.Literal("add"),
Type.Literal("done"),
]),
})
文件修改安全队列
如果你的工具会修改文件,使用 withFileMutationQueue() 参与文件修改队列:
实例
import { mkdir, readFile, writeFile } from "node:fs/promises";
import { dirname, resolve } from "node:path";
async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
// 将路径解析为绝对路径
const absolutePath = resolve(ctx.cwd, params.path);
return withFileMutationQueue(absolutePath, async () => {
// 确保目录存在
await mkdir(dirname(absolutePath), { recursive: true });
// 读取当前内容
const current = await readFile(absolutePath, "utf8");
// 应用修改
const next = current.replace(params.oldText, params.newText);
// 写回文件
await writeFile(absolutePath, next, "utf8");
return {
content: [{ type: "text", text: `已更新 ${params.path}` }],
details: {},
};
});
}
文件修改队列确保同一文件的并发修改不会互相覆盖——当内置 edit 工具和你的自定义工具同时修改同一个文件时,它们会排队执行。
输出截断
工具输出必须截断,避免撑满 LLM 上下文窗口。内置限制为 50KB(约 1 万 token)和 2000 行:
实例
truncateHead,
truncateTail,
formatSize,
DEFAULT_MAX_BYTES,
DEFAULT_MAX_LINES,
} from "@earendil-works/pi-coding-agent";
async execute(toolCallId, params) {
const output = await runSomeCommand();
// 保留开头部分(适合文件读取、搜索结果)
const truncation = truncateHead(output, {
maxLines: DEFAULT_MAX_LINES,
maxBytes: DEFAULT_MAX_BYTES,
});
let result = truncation.content;
if (truncation.truncated) {
result += `\n\n[输出已截断:${truncation.outputLines}/`;
result += `${truncation.totalLines} 行`;
result += `(${formatSize(truncation.outputBytes)}/`;
result += `${formatSize(truncation.totalBytes)})]`;
}
return {
content: [{ type: "text", text: result }],
};
}
截断不是可选的。过大的工具输出会导致上下文溢出、压缩失败和模型性能下降。始终对工具输出进行截断处理。
覆盖内置工具
你可以通过注册同名工具来覆盖内置工具(read、bash、edit、write、grep、find、ls):
$ pi -e ./tool-override.ts
覆盖时,渲染器(renderCall/renderResult)会按槽位继承——如果你省略了 renderCall,内置的 renderCall 仍会使用。这让你可以在不重写 UI 的情况下为内置工具添加日志或权限控制。
但对于 promptSnippet 和 promptGuidelines,它们不会从内置工具继承,需要显式定义。
