Pi Coding Agent 扩展系统源码分析

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

import { Type, type Static } from "typebox";

const MyParams = Type.Object({
  name: Type.String(),
  count: Type.Optional(Type.Integer({ minimum: 1 })),
});

pi.registerTool({
  name: "my-tool",
  parameters: MyParams,    // TSchema 约束
  execute: (id, params: Static<typeof MyParams>, signal, onUpdate, ctx) => {
    // params.name 类型为 string
    // params.count 类型为 number | undefined
  },
});

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

import type { AgentMessage, AgentToolResult, AgentToolUpdateCallback, ThinkingLevel, ToolExecutionMode } from "@earendil-works/pi-agent-core";
import type { Api, AssistantMessageEvent, AssistantMessageEventStream, Context, ImageContent, Model, OAuthCredentials, OAuthLoginCallbacks, SimpleStreamOptions, TextContent, ToolResultMessage } from "@earendil-works/pi-ai";
import type { AutocompleteItem, AutocompleteProvider, Component, EditorComponent, EditorTheme, KeyId, OverlayHandle, OverlayOptions, TUI } from "@earendil-works/pi-tui";
import type { Static, TSchema } from "typebox";
import type { Theme } from "../../modes/interactive/theme/theme.ts";
import type { BashResult } from "../bash-executor.ts";
import type { CompactionPreparation, CompactionResult } from "../compaction/index.ts";
import type { EventBus } from "../event-bus.ts";
import type { ExecOptions, ExecResult } from "../exec.ts";
import type { ReadonlyFooterDataProvider } from "../footer-data-provider.ts";
import type { KeybindingsManager } from "../keybindings.ts";
import type { CustomMessage } from "../messages.ts";
import type { ModelRegistry } from "../model-registry.ts";
import type { SessionManager, ReadonlySessionManager, SessionEntry, BranchSummaryEntry, CompactionEntry } from "../session-manager.ts";
import type { SlashCommandInfo } from "../slash-commands.ts";
import type { SourceInfo } from "../source-info.ts";
import type { BuildSystemPromptOptions } from "../system-prompt.ts";
import type { BashOperations } from "../tools/bash.ts";
import type { EditToolDetails } from "../tools/edit.ts";
import type { BashToolDetails, BashToolInput, EditToolInput, FindToolDetails, FindToolInput, GrepToolDetails, GrepToolInput, LsToolDetails, LsToolInput, ReadToolDetails, ReadToolInput, WriteToolInput } from "../tools/index.ts";

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

import { Type, type Static } from "typebox";

const MyParams = Type.Object({
  name: Type.String(),
  count: Type.Optional(Type.Integer({ minimum: 1 })),
});

pi.registerTool({
  name: "my-tool",
  label: "My Tool",
  description: "A custom tool with validated parameters",
  parameters: MyParams,      // TSchema 约束确保合法性
  execute: (id, params: Static<typeof MyParams>, signal, onUpdate, ctx) => {
    // TypeScript 自动推断：params.name 是 string
    // TypeScript 自动推断：params.count 是 number | undefined
    // 运行时 TypeBox 校验确保传入参数符合 schema
  },
});

1
2
3
4
5
6
7

export interface ExtensionUIContext {
  select(title: string, options: string[], opts?: ExtensionUIDialogOptions): Promise<string | undefined>;
  confirm(title: string, message: string, opts?: ExtensionUIDialogOptions): Promise<boolean>;
  input(title: string, placeholder?: string, opts?: ExtensionUIDialogOptions): Promise<string | undefined>;
  notify(message: string, type?: "info" | "warning" | "error"): void;
  // ... ~30 个方法
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

export interface ExtensionContext {
  ui: ExtensionUIContext;
  mode: ExtensionMode;         // "tui" | "rpc" | "json" | "print"
  hasUI: boolean;
  cwd: string;
  sessionManager: ReadonlySessionManager;
  modelRegistry: ModelRegistry;
  model: Model<any> | undefined;
  isIdle(): boolean;
  isProjectTrusted(): boolean;
  signal: AbortSignal | undefined;
  abort(): void;
  hasPendingMessages(): boolean;
  shutdown(): void;
  getContextUsage(): ContextUsage | undefined;
  compact(options?: CompactOptions): void;
  getSystemPrompt(): string;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

export interface ToolDefinition<TParams extends TSchema = TSchema, TDetails = unknown, TState = any> {
  name: string;
  label: string;
  description: string;
  promptSnippet?: string;         // 可选：系统提示中的简短描述
  promptGuidelines?: string[];     // 可选：系统提示指南中的要点
  parameters: TParams;             // TypeBox 模式（运行时校验）
  renderShell?: "default" | "self";
  prepareArguments?: (args: unknown) => Static<TParams>;
  executionMode?: ToolExecutionMode;  // "sequential" | "parallel"
  execute(toolCallId, params, signal, onUpdate, ctx): Promise<AgentToolResult<TDetails>>;
  renderCall?: (args, theme, context) => Component;
  renderResult?: (result, options, theme, context) => Component;
}

1
2
3

export function defineTool<TParams extends TSchema, TDetails = unknown, TState = any>(
  tool: ToolDefinition<TParams, TDetails, TState>,
): ToolDefinition<TParams, TDetails, TState> & AnyToolDefinition { ... }

1
2
3
4
5
6

export interface ToolCallEventResult { block?: boolean; reason?: string; }
export interface ContextEventResult { messages?: AgentMessage[]; }
export interface InputEventResult =
  | { action: "continue" }
  | { action: "transform"; text: string; images?: ImageContent[] }
  | { action: "handled" };

1
2

export function isBashToolResult(e: ToolResultEvent): e is BashToolResultEvent { ... }
export function isToolCallEventType(toolName: string, event: ToolCallEvent): boolean { ... }

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

export interface ExtensionAPI {
  on(event, handler): void;                                  // 订阅事件
  registerTool(tool): void;                                  // 注册工具
  registerCommand(name, options): void;                      // 注册命令
  registerShortcut(shortcut, options): void;                  // 注册快捷键
  registerFlag(name, options): void;                          // 注册 CLI 标志
  registerMessageRenderer(customType, renderer): void;        // 注册消息渲染器
  sendMessage(message, options?): void;                        // 发送自定义消息
  sendUserMessage(content, options?): void;                    // 发送用户消息
  appendEntry(customType, data?): void;                        // 追加会话条目
  registerProvider(name, config): void;                        // 注册/覆盖 Provider
  unregisterProvider(name): void;                              // 取消注册 Provider
  events: EventBus;                                            // 扩展间事件总线
  // ... 更多方法
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

export interface ProviderConfig {
  name?: string;
  baseUrl?: string;
  apiKey?: string;     // "$PROXY_API_KEY" — 环境变量插值
  api?: Api;
  streamSimple?: (model, context, options?) => AssistantMessageEventStream;  // 自定义 API 流处理器
  headers?: Record<string, string>;
  authHeader?: boolean;
  models?: ProviderModelConfig[];
  oauth?: { login, refreshToken, getApiKey, modifyModels? };
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

export interface ExtensionRuntime extends ExtensionRuntimeState, ExtensionActions {}
// 继承关系：runtime = 状态 + 动作

export interface Extension {
  path: string;                    // 原始路径
  resolvedPath: string;            // 解析后路径
  sourceInfo: SourceInfo;          // 来源元数据
  handlers: Map<string, HandlerFn[]>;  // 事件 -> 处理器列表
  tools: Map<string, RegisteredTool>;
  messageRenderers: Map<string, MessageRenderer>;
  commands: Map<string, RegisteredCommand>;
  flags: Map<string, ExtensionFlag>;
  shortcuts: Map<KeyId, ExtensionShortcut>;
}

1
2
3
4
5
6

import { createJiti } from "jiti/static";
import * as _bundledTypebox from "typebox";
import { CONFIG_DIR_NAME, getAgentDir, isBunBinary } from "../../config.ts";
import * as _bundledPiCodingAgent from "../../index.ts";
import { createEventBus, type EventBus } from "../event-bus.ts";
import { createSyntheticSourceInfo } from "../source-info.ts";

1
2
3
4
5
6
7
8
9

const VIRTUAL_MODULES: Record<string, unknown> = {
  typebox: _bundledTypebox,
  "typebox/compile": _bundledTypeboxCompile,
  "@earendil-works/pi-agent-core": _bundledPiAgentCore,
  "@earendil-works/pi-tui": _bundledPiTui,
  "@earendil-works/pi-ai": _bundledPiAiCompat,    // compat 是核心的超集
  "@earendil-works/pi-coding-agent": _bundledPiCodingAgent,
  "@mariozechner/pi-*": ...,                       // 旧命名空间别名
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

function getAliases(): Record<string, string> {
  const __dirname = path.dirname(fileURLToPath(import.meta.url));
  const packageIndex = path.resolve(__dirname, "../..", "index.js");
  const packagesRoot = path.resolve(__dirname, "../../../../");
  
  const resolveWorkspaceOrImport = (workspaceRelativePath: string, specifier: string): string => {
    const workspacePath = path.join(packagesRoot, workspaceRelativePath);
    if (fs.existsSync(workspacePath)) { return workspacePath; }  // monorepo workspace
    return fileURLToPath(import.meta.resolve(specifier));        // fallback to node_modules
  };
  // ...
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

export function createExtensionRuntime(): ExtensionRuntime {
  const notInitialized = () => {
    throw new Error("Extension runtime not initialized");
  };
  const state: { staleMessage?: string } = {};
  const assertActive = () => { if (state.staleMessage) throw new Error(state.staleMessage); };

  const runtime: ExtensionRuntime = {
    sendMessage: notInitialized,    // 抛出错误的桩函数
    sendUserMessage: notInitialized,
    // ... 所有动作都是桩函数
    registerProvider: (name, config, extensionPath = "<unknown>") => {
      runtime.pendingProviderRegistrations.push({ name, config, extensionPath });
    },
    invalidate: (message) => { state.staleMessage = message; },
    // ...
  };
  return runtime;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

function createExtensionAPI(extension: Extension, runtime: ExtensionRuntime, cwd: string, eventBus: EventBus): ExtensionAPI {
  const api = {
    on(event: string, handler: HandlerFn): void {
      runtime.assertActive();
      const list = extension.handlers.get(event) ?? [];
      list.push(handler);
      extension.handlers.set(event, list);
    },
    registerTool(tool: ToolDefinition): void {
      runtime.assertActive();
      extension.tools.set(tool.name, { definition: tool, sourceInfo: extension.sourceInfo });
      runtime.refreshTools();
    },
    getFlag(name: string): boolean | string | undefined {
      runtime.assertActive();
      if (!extension.flags.has(name)) return undefined;  // 安全检查：只能读自己注册的
      return runtime.flagValues.get(name);
    },
    // ...
  };
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

async function loadExtensionModule(extensionPath: string, cacheToken?: ExtensionCacheToken) {
  if (isCurrentCacheToken(cacheToken)) {
    const cachedFactory = extensionCache.get(extensionPath);
    if (cachedFactory) return cachedFactory;  // 缓存命中
  }

  const jiti = createJiti(import.meta.url, {
    moduleCache: false,
    ...(isBunBinary ? { virtualModules: VIRTUAL_MODULES, tryNative: false } : { alias: getAliases() }),
  });

  const module = await jiti.import(extensionPath, { default: true });
  const factory = module as ExtensionFactory;
  if (typeof factory !== "function") return undefined;  // 不是有效的扩展
  // ...
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

function readPiManifest(packageJsonPath: string): PiManifest | null {
  const pkg = JSON.parse(fs.readFileSync(packageJsonPath, "utf-8"));
  return pkg.pi && typeof pkg.pi === "object" ? pkg.pi as PiManifest : null;
}

function discoverExtensionsInDir(dir: string): string[] {
  // 发现规则：
  // 1. *.ts / *.js 文件 → 直接加载
  // 2. 子目录有 index.ts/index.js → 加载
  // 3. 子目录有 package.json 含 "pi.extensions" 字段 → 按清单加载
  // 不递归超过一层，复杂包必须使用 package.json 清单
}

export async function discoverAndLoadExtensions(
  configuredPaths: string[], cwd: string, agentDir: string
): Promise<LoadExtensionsResult> {
  // 优先级：
  // 1. 项目本地: cwd/.pi/extensions/
  // 2. 全局: agentDir/extensions/
  // 3. 配置路径（逐条解析）
  // 去重：Set<resolvedPath>
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

let extensionCacheCwd: string | undefined;
let extensionCacheGeneration = 0;
const extensionCache = new Map<string, ExtensionFactory>();

export function clearExtensionCache(): void { extensionCacheGeneration++; }

function useExtensionCacheCwd(cwd: string): ExtensionCacheToken {
  if (extensionCacheCwd !== undefined && extensionCacheCwd !== resolvedCwd) clearExtensionCache();
  // 切换工作目录 → 清空缓存
}

1
2
3
4
5
6

import type { AgentMessage } from "@earendil-works/pi-agent-core";
import type { KeyId } from "@earendil-works/pi-tui";
import { type Theme, theme } from "../../modes/interactive/theme/theme.ts";
import type { ModelRegistry } from "../model-registry.ts";
import type { SessionManager } from "../session-manager.ts";
// ... 从同一目录的 types.ts 导入 ~30 个类型

1
2
3
4
5

const RESERVED_KEYBINDINGS_FOR_EXTENSION_CONFLICTS = [
  "app.interrupt", "app.clear", "app.exit", "app.suspend", "tui.input.submit", ...
] as const;

type BuiltInKeyBindings = Partial<Record<KeyId, { keybinding: string; restrictOverride: boolean }>>;

1
2
3
4
5

for (const key of keyList) {
  const existing = builtinKeybindings[normalizedKey];
  if (existing?.restrictOverride && !restrictOverride) continue;  // 保留的获胜
  builtinKeybindings[normalizedKey] = { keybinding, restrictOverride };
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

export class ExtensionRunner {
  private extensions: Extension[];
  private runtime: ExtensionRuntime;
  private uiContext: ExtensionUIContext;        // 默认为 noOpUIContext
  private mode: ExtensionMode = "print";
  private cwd: string;
  private sessionManager: SessionManager;
  private modelRegistry: ModelRegistry;
  
  // 所有内部函数都可被外部替换（通过 bindCore/bindCommandContext）
  private getModel: () => Model<any> | undefined = () => undefined;
  private isIdleFn: () => boolean = () => true;
  // ... 约 15 个可注入函数
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

bindCore(actions: ExtensionActions, contextActions: ExtensionContextActions, providerActions?): void {
  // 1. 将 actions 复制到共享 runtime（所有 ExtensionAPI 引用同一个 runtime 对象）
  this.runtime.sendMessage = actions.sendMessage;
  this.runtime.sendUserMessage = actions.sendUserMessage;
  // ...

  // 2. 设置 context actions（getter 函数）
  this.getModel = contextActions.getModel;
  this.isIdleFn = contextActions.isIdle;
  // ...

  // 3. 刷新 provider 注册队列
  for (const { name, config } of this.runtime.pendingProviderRegistrations) {
    this.modelRegistry.registerProvider(name, config);
  }
  this.runtime.pendingProviderRegistrations = [];

  // 4. 替换 registerProvider 为立即生效的版本
  this.runtime.registerProvider = (name, config) => {
    this.modelRegistry.registerProvider(name, config);
  };
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

createContext(): ExtensionContext {
  const runner = this;
  return {
    get ui() {
      runner.assertActive();
      return runner.uiContext;
    },
    get mode() {
      runner.assertActive();
      return runner.mode;
    },
    isIdle: () => {
      runner.assertActive();
      return runner.isIdleFn();
    },
    // ...
  };
}

1
2
3
4
5
6
7
8
9

createCommandContext(): ExtensionCommandContext {
  const context = Object.defineProperties(
    {},
    Object.getOwnPropertyDescriptors(this.createContext()),
  ) as ExtensionCommandContext;
  context.waitForIdle = () => { ... };
  context.newSession = (options) => { ... };
  // ...
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

async emit<TEvent extends RunnerEmitEvent>(event: TEvent): Promise<RunnerEmitResult<TEvent>> {
  const ctx = this.createContext();
  for (const ext of this.extensions) {
    const handlers = ext.handlers.get(event.type);
    if (!handlers) continue;
    for (const handler of handlers) {
      const handlerResult = await handler(event, ctx);
      if (this.isSessionBeforeEvent(event) && handlerResult?.cancel) return;
    }
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

async emitMessageEnd(event: MessageEndEvent): Promise<AgentMessage | undefined> {
  const ctx = this.createContext();
  let currentMessage = event.message;
  let modified = false;

  for (const ext of this.extensions) {
    for (const handler of ext.handlers.get("message_end") ?? []) {
      const currentEvent: MessageEndEvent = { ...event, message: currentMessage };
      const handlerResult = await handler(currentEvent, ctx) as MessageEndEventResult;
      if (handlerResult?.message) {
        if (handlerResult.message.role !== currentMessage.role) {
          this.emitError({...});  // 角色不变检查
          continue;
        }
        currentMessage = handlerResult.message;
        modified = true;
      }
    }
  }
  return modified ? currentMessage : undefined;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

async emitToolResult(event: ToolResultEvent): Promise<ToolResultEventResult | undefined> {
  // 对原事件对象展开（不会污染原始对象）
  const currentEvent: ToolResultEvent = { ...event };
  for (const handler of handlers) {
    const handlerResult = await handler(currentEvent, ctx) as ToolResultEventResult;
    if (handlerResult.content !== undefined) currentEvent.content = handlerResult.content;
    if (handlerResult.details !== undefined) currentEvent.details = handlerResult.details;
    if (handlerResult.isError !== undefined) currentEvent.isError = handlerResult.isError;
  }
  // 返回修改后的完整结果
}

1
2
3
4
5
6
7
8
9

async emitToolCall(event: ToolCallEvent): Promise<ToolCallEventResult | undefined> {
  for (const handler of handlers) {
    const handlerResult = await handler(event, ctx);
    if (handlerResult) {
      result = handlerResult as ToolCallEventResult;
      if (result.block) return result;  // block 立即终止
    }
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

async emitInput(text, images, source, streamingBehavior?): Promise<InputEventResult> {
  let currentText = text;
  let currentImages = images;
  for (const handler of ext.handlers.get("input") ?? []) {
    const event: InputEvent = { type: "input", text: currentText, images: currentImages, source, streamingBehavior };
    const result = await handler(event, ctx) as InputEventResult;
    if (result?.action === "handled") return result;      // 短路：扩展完全处理了输入
    if (result?.action === "transform") {                  // 转换：修改输入文本
      currentText = result.text;
      currentImages = result.images ?? currentImages;
    }
  }
  return currentText !== text
    ? { action: "transform", text: currentText, images: currentImages }
    : { action: "continue" };
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

getAllRegisteredTools(): RegisteredTool[] {
  const toolsByName = new Map<string, RegisteredTool>();
  for (const ext of this.extensions) {
    for (const tool of ext.tools.values()) {
      if (!toolsByName.has(tool.definition.name)) {
        toolsByName.set(tool.definition.name, tool);  // 先注册者获胜
      }
    }
  }
  return Array.from(toolsByName.values());
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

private resolveRegisteredCommands(): ResolvedCommand[] {
  const counts = new Map<string, number>();
  for (const ext of this.extensions) {
    for (const command of ext.commands.values()) {
      counts.set(command.name, (counts.get(command.name) ?? 0) + 1);
    }
  }
  // 重名命令自动添加数字后缀：command:1, command:2
  return commands.map((command) => {
    const invocationName = counts.get(command.name) > 1
      ? `${command.name}:${occurrence}`
      : command.name;
    // ...
  });
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

invalidate(message?: string): void {
  if (!this.staleMessage) {
    this.staleMessage = message;
    this.runtime.invalidate(message);  // 传播到共享 runtime
  }
}

private assertActive(): void {
  if (this.staleMessage) {
    throw new Error(this.staleMessage);
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

import type { AgentTool } from "@earendil-works/pi-agent-core";
import { wrapToolDefinition, wrapToolDefinitions } from "../tools/tool-definition-wrapper.ts";
import type { ExtensionRunner } from "./runner.ts";
import type { RegisteredTool } from "./types.ts";

export function wrapRegisteredTool(registeredTool: RegisteredTool, runner: ExtensionRunner): AgentTool {
  return wrapToolDefinition(registeredTool.definition, () => runner.createContext());
}

export function wrapRegisteredTools(registeredTools: RegisteredTool[], runner: ExtensionRunner): AgentTool[] {
  return wrapToolDefinitions(
    registeredTools.map((registeredTool) => registeredTool.definition),
    () => runner.createContext(),
  );
}

1
2
3
4
5
6

export {
  createExtensionRuntime,
  discoverAndLoadExtensions,
  loadExtensionFromFactory,
  loadExtensions,
} from "./loader.ts";

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

export {
  ExtensionRunner,
} from "./runner.ts";
export type {
  ExtensionErrorListener,
  ForkHandler,
  NavigateTreeHandler,
  NewSessionHandler,
  SwitchSessionHandler,
} from "./runner.ts";

1

export { wrapRegisteredTool, wrapRegisteredTools } from "./wrapper.ts";

1
2

export type { SlashCommandInfo, SlashCommandSource } from "../slash-commands.ts";
export type { SourceInfo } from "../source-info.ts";