AgentSession 源码分析 — Pi Coding Agent 核心生命周期管理

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

export interface ParsedSkillBlock {
  name: string;
  location: string;
  content: string;
  userMessage: string | undefined;
}

export function parseSkillBlock(text: string): ParsedSkillBlock | null {
  const match = text.match(
    /^<skill name="([^"]+)" location="([^"]+)">\n([\s\S]*?)\n<\/skill>(?:\n\n([\s\S]+))?$/
  );
  if (!match) return null;
  return { name: match[1], location: match[2], content: match[3], userMessage: match[4]?.trim() || undefined };
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

export type AgentSessionEvent =
  | Exclude<AgentEvent, { type: "agent_end" }>
  | { type: "agent_end"; messages: AgentMessage[]; willRetry: boolean }
  | { type: "queue_update"; steering: readonly string[]; followUp: readonly string[] }
  | { type: "compaction_start"; reason: "manual" | "threshold" | "overflow" }
  | { type: "compaction_end"; reason: ...; result: ...; aborted: boolean; willRetry: boolean; errorMessage?: string }
  | { type: "auto_retry_start"; attempt: number; maxAttempts: number; delayMs: number; errorMessage: string }
  | { type: "auto_retry_end"; success: boolean; attempt: number; finalError?: string }
  | { type: "session_info_changed"; name: string | undefined }
  | { type: "thinking_level_changed"; level: ThinkingLevel };

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

export interface AgentSessionConfig {
  agent: Agent;                              // agent-core 的 Agent 实例
  sessionManager: SessionManager;            // 会话持久化管理器
  settingsManager: SettingsManager;          // 设置管理器
  cwd: string;
  scopedModels?: Array<{ model: Model<any>; thinkingLevel?: ThinkingLevel }>; // --models 标志
  resourceLoader: ResourceLoader;
  customTools?: ToolDefinition[];
  modelRegistry: ModelRegistry;
  initialActiveToolNames?: string[];
  allowedToolNames?: string[];               // 工具 allowlist
  excludedToolNames?: string[];              // 工具 denylist
  baseToolsOverride?: Record<string, AgentTool>;
  extensionRunnerRef?: { current?: ExtensionRunner };
  sessionStartEvent?: SessionStartEvent;
}

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

export class AgentSession {
  // ~30 个私有状态字段
  private _steeringMessages: string[] = [];
  private _followUpMessages: string[] = [];
  private _pendingNextTurnMessages: CustomMessage[] = [];
  private _compactionAbortController: AbortController | undefined;
  private _retryAbortController: AbortController | undefined;
  private _bashAbortController: AbortController | undefined;
  // ...

  constructor(config: AgentSessionConfig) {
    // 1. 保存配置
    this.agent = config.agent;
    this.sessionManager = config.sessionManager;
    // ...
    // 2. 订阅 agent 事件
    this._unsubscribeAgent = this.agent.subscribe(this._handleAgentEvent);
    // 3. 安装工具钩子（beforeToolCall / afterToolCall）
    this._installAgentToolHooks();
    // 4. 构建运行时（工具注册表、ExtensionRunner、系统提示）
    this._buildRuntime({ activeToolNames: this._initialActiveToolNames, includeAllExtensionTools: true });
  }

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

private _installAgentToolHooks(): void {
  this.agent.beforeToolCall = async ({ toolCall, args }) => {
    const runner = this._extensionRunner;
    if (!runner.hasHandlers("tool_call")) return undefined;
    return await runner.emitToolCall({ type: "tool_call", toolName: toolCall.name, toolCallId: toolCall.id, input: args as Record<string, unknown> });
  };

  this.agent.afterToolCall = async ({ toolCall, args, result, isError }) => {
    const runner = this._extensionRunner;
    if (!runner.hasHandlers("tool_result")) return undefined;
    const hookResult = await runner.emitToolResult({ type: "tool_result", toolName: toolCall.name, ... });
    if (!hookResult) return undefined;
    return { content: hookResult.content, details: hookResult.details, isError: hookResult.isError ?? isError };
  };
}

1
2
3
4

subscribe(listener: AgentSessionEventListener): () => void {
  this._eventListeners.push(listener);
  return () => { const i = this._eventListeners.indexOf(listener); if (i !== -1) this._eventListeners.splice(i, 1); };
}

1
2
3
4
5
6

private _replaceMessageInPlace(target: AgentMessage, replacement: AgentMessage): void {
  if (target === replacement) return;
  const targetRecord = target as unknown as Record<string, unknown>;
  for (const key of Object.keys(targetRecord)) delete targetRecord[key];
  Object.assign(targetRecord, replacement);
}

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

private async _runAgentPrompt(messages): Promise<void> {
  try {
    await this.agent.prompt(messages);
    while (await this._handlePostAgentRun()) {
      await this.agent.continue();
    }
  } finally {
    this._flushPendingBashMessages();
  }
}

private async _handlePostAgentRun(): Promise<boolean> {
  // 1. 重试检查 → 是则准备重试，返回 true
  if (this._isRetryableError(msg) && await this._prepareRetry(msg)) return true;
  // 2. 压缩检查 → 是则执行压缩，返回 true（overflow 会继续重试）
  if (await this._checkCompaction(msg)) return true;
  // 3. 队列检查 → extension 在 agent_end 中排队了消息
  return this.agent.hasQueuedMessages();
}

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

async steer(text: string, images?: ImageContent[]): Promise<void> {
  if (text.startsWith("/")) this._throwIfExtensionCommand(text); // 命令不能排队
  let expandedText = this._expandSkillCommand(text);
  expandedText = expandPromptTemplate(expandedText, [...this.promptTemplates]);
  await this._queueSteer(expandedText, images);
}

private async _queueSteer(text: string, images?: ImageContent[]): Promise<void> {
  this._steeringMessages.push(text);
  this._emitQueueUpdate();
  this.agent.steer({ role: "user", content: [{ type: "text", text }, ...(images ?? [])], timestamp: Date.now() });
}

1
2
3
4
5
6
7
8
9

async setModel(model: Model<any>): Promise<void> {
  if (!this._modelRegistry.hasConfiguredAuth(model)) throw new Error(...);
  const previousModel = this.model;       // 保存旧模型
  this.agent.state.model = model;          // 设置新模型
  this.sessionManager.appendModelChange(model.provider, model.id);  // 持久化
  this.settingsManager.setDefaultModelAndProvider(model.provider, model.id);  // 保存为默认
  this.setThinkingLevel(thinkingLevel);    // 重新钳位 thinking level
  await this._emitModelSelect(model, previousModel, "set");  // 通知扩展
}

1
2
3
4

async cycleModel(direction: "forward" | "backward"): Promise<ModelCycleResult | undefined> {
  if (this._scopedModels.length > 0) return this._cycleScopedModel(direction);     // --models 优先
  return this._cycleAvailableModel(direction);  // 回退到 ModelRegistry
}

1
2
3
4
5
6
7

setThinkingLevel(level: ThinkingLevel): void {
  const effectiveLevel = this.getAvailableThinkingLevels().includes(level)
    ? level
    : this._clampThinkingLevel(level, this.getAvailableThinkingLevels());
  // 只在实际改变时持久化
  if (effectiveLevel !== this.agent.state.thinkingLevel) { ... }
}

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

private async _checkCompaction(assistantMessage, skipAbortedCheck = true): Promise<boolean> {
  // 1. 跳过中止消息
  if (skipAbortedCheck && stopReason === "aborted") return false;
  // 2. 跳过来自不同模型的消息（用户已切换模型）
  if (!sameModel) return false;
  // 3. 跳过最近一次压缩之前的消息（防重复触发）
  if (assistantIsFromBeforeCompaction) return false;
  // 4. Overflow → 压缩 + 重试
  if (isContextOverflow(assistantMessage, contextWindow)) { ... }
  // 5. Threshold → 压缩
  if (shouldCompact(contextTokens, contextWindow, settings)) { ... }
}

1
2
3
4
5
6
7
8

private _isRetryableError(message: AssistantMessage): boolean {
  const err = message.errorMessage;
  if (this._isNonRetryableProviderLimitError(err)) return false;
  // 覆盖数百种错误：
  return /overloaded|provider.?returned.?error|rate.?limit|too many requests|
         429|500|502|503|504|service.?unavailable|connection.?error|connection.?refused|
         websocket.?closed|fetch failed|socket hang up|timed? out|terminated|retry delay/i.test(err);
}

1
2
3
4
5
6
7
8
9

private async _prepareRetry(message: AssistantMessage): Promise<boolean> {
  this._retryAttempt++;
  if (this._retryAttempt > settings.maxRetries) { this._retryAttempt--; return false; }

  const delayMs = settings.baseDelayMs * 2 ** (this._retryAttempt - 1);
  // auto_retry_start 事件 → 移除错误消息 → sleep(delayMs, abortable)
  // 若 sleep 被中止 → auto_retry_end(success: false)
  return true;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

async executeBash(command, onChunk?, options?): Promise<BashResult> {
  this._bashAbortController = new AbortController();
  const result = await executeBashWithOperations(
    resolvedCommand,
    this.sessionManager.getCwd(),
    options?.operations ?? createLocalBashOperations({ shellPath }),
    { onChunk, signal: this._bashAbortController.signal },
  );
  this.recordBashResult(command, result, options);
  return result;
}

1
2
3
4

// ExtensionCommandContextActions 中定义，由各运行模式实现
newSession(options?: { parentSession?, setup?, withSession? })
fork(entryId, options?: { position?, withSession? })
switchSession(sessionPath, options?: { withSession? })

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

runner.bindCore(
  {
    // ExtensionActions — 扩展的 pi.sendMessage / pi.sendUserMessage 等
    sendMessage: (message, options) => this.sendCustomMessage(message, options),
    sendUserMessage: (content, options) => this.sendUserMessage(content, options),
    appendEntry: (customType, data) => this.sessionManager.appendCustomEntry(customType, data),
    setSessionName: (name) => this.setSessionName(name),
    getSessionName: () => this.sessionManager.getSessionName(),
    getActiveTools: () => this.getActiveToolNames(),
    getAllTools: () => this.getAllTools(),
    setActiveTools: (names) => this.setActiveToolsByName(names),
    refreshTools: () => this._refreshToolRegistry(),
    getCommands: () => [...extensionCommands, ...templates, ...skills],
    setModel: async (model) => { ... await this.setModel(model); return true; },
    getThinkingLevel: () => this.thinkingLevel,
    setThinkingLevel: (level) => this.setThinkingLevel(level),
  },
  { /* ExtensionContextActions — ctx.* 方法 */ },
  { /* Provider registration */ },
);

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

setActiveToolsByName(toolNames: string[]): void {
  const tools: AgentTool[] = [];
  for (const name of toolNames) {
    const tool = this._toolRegistry.get(name);
    if (tool) { tools.push(tool); }
  }
  this.agent.state.tools = tools;                   // 更新 agent
  this._baseSystemPrompt = this._rebuildSystemPrompt(validToolNames);  // 重建提示
  this.agent.state.systemPrompt = this._baseSystemPrompt;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

getContextUsage(): ContextUsage | undefined {
  // 压缩后：只信任压缩边界之后的 assistant usage
  // （压缩前的 usage 反映的是老的大上下文）
  const latestCompaction = getLatestCompactionEntry(branchEntries);
  if (latestCompaction && !hasPostCompactionUsage) {
    return { tokens: null, contextWindow, percent: null };  // 未知
  }
  // 正常情况：估算 token 数
  const estimate = estimateContextTokens(this.messages);
  return { tokens: estimate.tokens, contextWindow, percent: (estimate.tokens / contextWindow) * 100 };
}

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

async exportToHtml(outputPath?: string): Promise<string> {
  // 创建工具渲染器（支持自定义工具 HTML 渲染）
  const toolRenderer = createToolHtmlRenderer({ getToolDefinition: (name) => this.getToolDefinition(name), ... });
  return await exportSessionToHtml(this.sessionManager, this.state, { outputPath, themeName, toolRenderer });
}

exportToJsonl(outputPath?: string): string {
  // 线性化分支：重新建立 parentId 链，写入 JSONL 文件
  let prevId: string | null = null;
  for (const entry of branchEntries) {
    const linear = { ...entry, parentId: prevId };
    writeFileSync(filePath, `${lines.join("\n")}\n`);
    prevId = entry.id;
  }
}

1
2
3
4
5
6
7

createReplacedSessionContext(): ReplacedSessionContext {
  // 使用 Object.defineProperties 复制惰性 getter
  const context = Object.defineProperties({}, Object.getOwnPropertyDescriptors(this._extensionRunner.createCommandContext())) as ReplacedSessionContext;
  context.sendMessage = (message, options) => this.sendCustomMessage(message, options);
  context.sendUserMessage = (content, options) => this.sendUserMessage(content, options);
  return context;
}