@esso0428/pi-subagents 0.17.1 → 0.17.3

package/src/agent-manager.ts

@@ -11,6 +11,14 @@ import { statSync } from "node:fs";
 import { isAbsolute } from "node:path";
 import type { Model } from "@earendil-works/pi-ai";
 import type { AgentSession, ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { readAgentHistory } from "./agent-history.js";
+import {
+  type AgentRecoveryCheckpoint,
+  type AgentRecoveryStatus,
+  readAgentRecoveryCheckpoints,
+  removeAgentRecoveryCheckpoint,
+  writeAgentRecoveryCheckpoint,
+} from "./agent-recovery.js";
 import { resumeAgent, runAgent, type ToolActivity } from "./agent-runner.js";
 import type { AgentInvocation, AgentRecord, IsolationMode, SubagentType, ThinkingLevel } from "./types.js";
 import { addUsage } from "./usage.js";
@@ -190,6 +198,8 @@ interface SpawnOptions {
   onAssistantUsage?: (usage: { input: number; output: number; cacheWrite: number }) => void;
   /** Called when the session successfully compacts. */
   onCompaction?: (info: CompactionInfo) => void;
+  /** Called synchronously after the record exists, before it is queued or started. */
+  onSpawned?: (id: string) => void;
 }
 
 export class AgentManager {
@@ -202,11 +212,15 @@ export class AgentManager {
   /** Base repos worktrees were created from — so dispose() can prune them all,
    *  not just the parent repo (caller-supplied cwd can target other repos). */
   private worktreeRepos = new Set<string>();
+  /** Project cwd for each record's durable checkpoint. */
+  private recoveryCwds = new Map<string, string>();
 
   /** Queue of background agents waiting to start. */
   private queue: { id: string; args: SpawnArgs }[] = [];
   /** Number of currently running background agents. */
   private runningBackground = 0;
+  /** Prevent late promise settlement from decrementing a replacement run. */
+  private runningBackgroundIds = new Set<string>();
 
   constructor(
     onComplete?: OnAgentComplete,
@@ -234,6 +248,99 @@ export class AgentManager {
     return this.maxConcurrent;
   }
 
+  private finishBackground(id: string): void {
+    if (!this.runningBackgroundIds.delete(id)) return;
+    this.runningBackground = Math.max(0, this.runningBackground - 1);
+  }
+
+  private checkpointStatus(record: AgentRecord): AgentRecoveryStatus {
+    return record.status;
+  }
+
+  private makeCheckpoint(record: AgentRecord): AgentRecoveryCheckpoint {
+    const checkpoint: AgentRecoveryCheckpoint = {
+      version: 1,
+      id: record.id,
+      type: record.type,
+      description: record.description,
+      status: this.checkpointStatus(record),
+      startedAt: record.startedAt,
+      toolUses: record.toolUses,
+      lifetimeUsage: { ...record.lifetimeUsage },
+      compactionCount: record.compactionCount,
+      ...(record.completedAt !== undefined && { completedAt: record.completedAt }),
+      // A durable transcript is the source of truth for partial/full output.
+      // Avoid duplicating potentially sensitive or very large result text.
+      ...(!record.transcriptPath && record.result !== undefined && { result: record.result }),
+      ...(record.error !== undefined && { error: record.error }),
+      ...(record.transcriptPath !== undefined && { transcriptPath: record.transcriptPath }),
+      ...(record.invocation !== undefined && { invocation: cloneInvocation(record.invocation) }),
+    };
+    return checkpoint;
+  }
+
+  private checkpoint(record: AgentRecord): void {
+    const cwd = this.recoveryCwds.get(record.id);
+    if (!cwd) return;
+    writeAgentRecoveryCheckpoint(cwd, this.makeCheckpoint(record));
+  }
+
+  private flushOutput(record: AgentRecord): void {
+    if (!record.outputCleanup) return;
+    try { record.outputCleanup(); } catch { /* recovery must remain best effort */ }
+    record.outputCleanup = undefined;
+  }
+
+  /** Set the durable transcript locator and checkpoint the current state. */
+  setTranscript(id: string, historyFile: string, transcriptPath: string, cwd?: string): void {
+    const record = this.agents.get(id);
+    if (!record) return;
+    record.historyFile = historyFile;
+    record.transcriptPath = transcriptPath;
+    if (cwd) this.recoveryCwds.set(id, cwd);
+    this.checkpoint(record);
+  }
+
+  /** Checkpoint one record explicitly (used after transcript setup). */
+  checkpointRecord(id: string): void {
+    const record = this.agents.get(id);
+    if (record) this.checkpoint(record);
+  }
+
+  /**
+   * Reload durable records from this project's checkpoint directory. A
+   * running/queued checkpoint means the process was killed before it could
+   * write its stopped state; treat it as stopped and retain its transcript.
+   * SIGKILL cannot run a final flush/checkpoint, so this active snapshot is
+   * necessarily the last recoverable state.
+   */
+  restoreRecovered(cwd: string): void {
+    for (const checkpoint of readAgentRecoveryCheckpoints(cwd)) {
+      if (!checkpoint.transcriptPath || !readAgentHistory(cwd, checkpoint.transcriptPath)) continue;
+      const status: RestorableAgentStatus = checkpoint.status === "running" || checkpoint.status === "queued"
+        ? "stopped"
+        : checkpoint.status;
+      const completedAt = checkpoint.completedAt ?? Date.now();
+      const existing = this.agents.get(checkpoint.id);
+      if (existing) {
+        // Parent-branch records can still carry an unread in-memory result.
+        // Never replace that richer record with the checkpoint's transcript
+        // stub during the same session. Merge only durable locator metadata.
+        if (!existing.transcriptPath && checkpoint.transcriptPath) {
+          existing.transcriptPath = checkpoint.transcriptPath;
+        }
+        this.recoveryCwds.set(checkpoint.id, cwd);
+        continue;
+      }
+      this.agents.set(checkpoint.id, this.createRestoredRecord({
+        ...checkpoint,
+        status,
+        completedAt,
+      }));
+      this.recoveryCwds.set(checkpoint.id, cwd);
+    }
+  }
+
   /**
    * Spawn an agent and return its ID immediately (for background use).
    * If the concurrency limit is reached, the agent is queued.
@@ -271,6 +378,19 @@ export class AgentManager {
       invocation: options.invocation,
     };
     this.agents.set(id, record);
+    this.recoveryCwds.set(id, ctx.cwd);
+    // Give callers a chance to create the durable transcript before the first
+    // checkpoint. This closes the small spawn→attach window in which a queued
+    // or running agent could be left recoverable only as metadata.
+    try {
+      options.onSpawned?.(id);
+      this.checkpoint(record);
+    } catch (err) {
+      this.agents.delete(id);
+      this.recoveryCwds.delete(id);
+      removeAgentRecoveryCheckpoint(ctx.cwd, id);
+      throw err;
+    }
 
     const args: SpawnArgs = { pi, ctx, type, prompt, options };
 
@@ -286,6 +406,8 @@ export class AgentManager {
       this.startAgent(id, record, args);
     } catch (err) {
       this.agents.delete(id);
+      this.recoveryCwds.delete(id);
+      removeAgentRecoveryCheckpoint(ctx.cwd, id);
       throw err;
     }
     return id;
@@ -327,7 +449,11 @@ export class AgentManager {
 
     record.status = "running";
     record.startedAt = Date.now();
-    if (options.isBackground) this.runningBackground++;
+    this.checkpoint(record);
+    if (options.isBackground) {
+      this.runningBackground++;
+      this.runningBackgroundIds.add(id);
+    }
     this.onStart?.(record);
 
     // Wire parent abort signal to stop the subagent when the parent is interrupted
@@ -386,9 +512,13 @@ export class AgentManager {
           record.pendingSteers = undefined;
         }
         options.onSessionCreated?.(session);
+        this.checkpoint(record);
       },
     })
       .then(({ responseText, session, aborted, steered, failure }) => {
+        // A disposed manager no longer owns this run. Avoid late callbacks
+        // mutating a dead session or emitting completion side effects.
+        if (this.agents.get(id) !== record) return responseText;
         // Don't overwrite status if externally stopped via abort()
         if (record.status !== "stopped") {
           // Precedence: a hard abort keeps "aborted"; then a failed final turn
@@ -427,6 +557,7 @@ export class AgentManager {
               `\n\n---\nChanges saved to branch \`${wtResult.branch}\`${repoNote}. Merge with: \`git merge ${wtResult.branch}\`${customCwd !== undefined ? ` (run in \`${baseCwd}\`)` : ""}`;
           }
         }
+        this.checkpoint(record);
 
         // Fire onComplete for foreground agents too — lifecycle symmetry.
         // Mark resultConsumed so the callback skips notifications (result returned inline).
@@ -434,13 +565,16 @@ export class AgentManager {
           record.resultConsumed = true;
           try { this.onComplete?.(record); } catch { /* ignore completion side-effect errors */ }
         } else {
-          this.runningBackground--;
+          this.finishBackground(id);
           try { this.onComplete?.(record); } catch { /* ignore completion side-effect errors */ }
           this.drainQueue();
         }
         return responseText;
       })
       .catch((err) => {
+        // A disposed manager no longer owns this run. Avoid late callbacks
+        // mutating a dead session or emitting completion side effects.
+        if (this.agents.get(id) !== record) return "";
         // Don't overwrite status if externally stopped via abort()
         if (record.status !== "stopped") {
           record.status = "error";
@@ -463,6 +597,7 @@ export class AgentManager {
             record.worktreeResult = wtResult;
           } catch { /* ignore cleanup errors */ }
         }
+        this.checkpoint(record);
 
         // Fire onComplete for foreground agents too — lifecycle symmetry.
         // Mark resultConsumed so the callback skips notifications (result returned inline).
@@ -470,7 +605,7 @@ export class AgentManager {
           record.resultConsumed = true;
           this.onComplete?.(record);
         } else {
-          this.runningBackground--;
+          this.finishBackground(id);
           this.onComplete?.(record);
           this.drainQueue();
         }
@@ -478,11 +613,6 @@ export class AgentManager {
       });
 
     record.promise = promise;
-
-    // Notify caller that spawn is complete (record is in the map, promise is set).
-    // Called synchronously — onSessionCreated fires asynchronously inside runAgent.
-    // Used by spawnAndWait to let the caller set up output files before streaming starts.
-    this.onSpawned?.(id);
   }
 
   /** Start queued agents up to the concurrency limit. */
@@ -499,18 +629,12 @@ export class AgentManager {
         record.status = "error";
         record.error = err instanceof Error ? err.message : String(err);
         record.completedAt = Date.now();
+        this.checkpoint(record);
         this.onComplete?.(record);
       }
     }
   }
 
-  /**
-   * Called synchronously right after spawn, before onSessionCreated fires.
-   * Lets the caller set up the output file path on the record.
-   * The record is guaranteed to be in this.agents at this point.
-   */
-  private onSpawned?: (id: string) => void;
-
   /**
    * Spawn an agent and wait for completion (foreground use).
    * Foreground agents bypass the concurrency queue.
@@ -527,17 +651,14 @@ export class AgentManager {
     options: Omit<SpawnOptions, "isBackground">,
     onSpawned?: (id: string) => void,
   ): Promise<{ id: string; record: AgentRecord }> {
-    // Temporarily register the onSpawned hook so startAgent can call it.
-    const prevOnSpawned = this.onSpawned;
-    this.onSpawned = onSpawned;
-    try {
-      const id = this.spawn(pi, ctx, type, prompt, { ...options, isBackground: false });
-      const record = this.agents.get(id)!;
-      await record.promise;
-      return { id, record };
-    } finally {
-      this.onSpawned = prevOnSpawned;
-    }
+    const id = this.spawn(pi, ctx, type, prompt, {
+      ...options,
+      isBackground: false,
+      onSpawned,
+    });
+    const record = this.agents.get(id)!;
+    await record.promise;
+    return { id, record };
   }
 
   /**
@@ -562,6 +683,7 @@ export class AgentManager {
       ...(resumedModel && { effectiveModelName: resumedModel.name ?? resumedModel.id }),
       effectiveThinking: record.session.thinkingLevel,
     };
+    this.checkpoint(record);
 
     try {
       const { text, failure } = await resumeAgent(record.session, prompt, {
@@ -583,10 +705,12 @@ export class AgentManager {
       if (failure) record.error = failure;
       record.result = text;
       record.completedAt = Date.now();
+      this.checkpoint(record);
     } catch (err) {
       record.status = "error";
       record.error = err instanceof Error ? err.message : String(err);
       record.completedAt = Date.now();
+      this.checkpoint(record);
     }
 
     return record;
@@ -652,6 +776,7 @@ export class AgentManager {
     lifetimeUsage?: { input: number; output: number; cacheWrite: number };
     transcriptPath?: string;
     invocation?: AgentInvocation;
+    compactionCount?: number;
   }): AgentRecord {
     return {
       id: record.id,
@@ -668,7 +793,7 @@ export class AgentManager {
       lifetimeUsage: record.lifetimeUsage
         ? { ...record.lifetimeUsage }
         : { input: 0, output: 0, cacheWrite: 0 },
-      compactionCount: 0,
+      compactionCount: record.compactionCount ?? 0,
     };
   }
 
@@ -681,6 +806,7 @@ export class AgentManager {
       this.queue = this.queue.filter(q => q.id !== id);
       record.status = "stopped";
       record.completedAt = Date.now();
+      this.checkpoint(record);
       return true;
     }
 
@@ -688,6 +814,10 @@ export class AgentManager {
     record.abortController?.abort();
     record.status = "stopped";
     record.completedAt = Date.now();
+    this.flushOutput(record);
+    this.checkpoint(record);
+    this.finishBackground(id);
+    this.drainQueue();
     return true;
   }
 
@@ -714,6 +844,10 @@ export class AgentManager {
         record.outputCleanup = undefined;
         record.outputFile = undefined;
         record.historyFile = undefined;
+        // The durable transcript is the source of truth after the TTL. Keep
+        // only the small identity/status record in memory; get_subagent_result
+        // reloads the final answer from transcriptPath on demand.
+        record.result = undefined;
         continue;
       }
       this.removeRecord(id, record);
@@ -750,16 +884,21 @@ export class AgentManager {
       if (record) {
         record.status = "stopped";
         record.completedAt = Date.now();
+        this.checkpoint(record);
         count++;
       }
     }
     this.queue = [];
-    // Abort running agents
+    // Abort running agents. Flush before checkpointing so a catchable
+    // shutdown/session switch leaves the latest assistant message available.
     for (const record of this.agents.values()) {
       if (record.status === "running") {
         record.abortController?.abort();
         record.status = "stopped";
         record.completedAt = Date.now();
+        this.flushOutput(record);
+        this.finishBackground(record.id);
+        this.checkpoint(record);
         count++;
       }
     }
@@ -789,6 +928,9 @@ export class AgentManager {
       record.session?.dispose();
     }
     this.agents.clear();
+    this.recoveryCwds.clear();
+    this.runningBackgroundIds.clear();
+    this.runningBackground = 0;
     // Prune any orphaned git worktrees (crash recovery)
     try { pruneWorktrees(process.cwd()); } catch { /* ignore */ }
     // Also prune repos that caller-supplied cwds created worktrees in — a clean

package/src/agent-recovery.ts

@@ -0,0 +1,185 @@
+/** Durable checkpoints for agents that may be interrupted by a catchable lifecycle event. */
+
+import { randomUUID } from "node:crypto";
+import { existsSync, mkdirSync, readdirSync, readFileSync, renameSync, unlinkSync, writeFileSync } from "node:fs";
+import { join } from "node:path";
+import { ensureSubagentsGitignore } from "./agent-history.js";
+import type { AgentInvocation } from "./types.js";
+import type { LifetimeUsage } from "./usage.js";
+
+const SUBAGENTS_DIR = ".pi-subagents";
+const CHECKPOINTS_DIR = "agent-checkpoints";
+const CHECKPOINT_VERSION = 1;
+const TERMINAL_STATUSES = new Set(["completed", "steered", "stopped", "aborted", "error"] as const);
+const THINKING_LEVELS = new Set(["minimal", "low", "medium", "high", "xhigh", "max", "off"]);
+
+type ActiveStatus = "running" | "queued";
+type TerminalStatus = "completed" | "steered" | "stopped" | "aborted" | "error";
+export type AgentRecoveryStatus = ActiveStatus | TerminalStatus;
+
+export interface AgentRecoveryCheckpoint {
+  version: 1;
+  id: string;
+  type: string;
+  description: string;
+  status: AgentRecoveryStatus;
+  startedAt: number;
+  completedAt?: number;
+  result?: string;
+  error?: string;
+  toolUses: number;
+  lifetimeUsage: LifetimeUsage;
+  compactionCount: number;
+  transcriptPath?: string;
+  invocation?: AgentInvocation;
+}
+
+function isSafeString(value: unknown, maxLength: number): value is string {
+  return typeof value === "string" && value.length > 0 && value.length <= maxLength && !/[\0\r\n]/.test(value);
+}
+
+function isFiniteTimestamp(value: unknown): value is number {
+  return typeof value === "number" && Number.isFinite(value) && Number.isInteger(value) && value >= 0;
+}
+
+function isUsage(value: unknown): value is LifetimeUsage {
+  if (!value || typeof value !== "object") return false;
+  const usage = value as Record<string, unknown>;
+  return ["input", "output", "cacheWrite"].every((key) => {
+    const n = usage[key];
+    return typeof n === "number" && Number.isFinite(n) && n >= 0;
+  });
+}
+
+function isInvocation(value: unknown): value is AgentInvocation {
+  if (!value || typeof value !== "object") return false;
+  const invocation = value as Record<string, unknown>;
+  for (const key of ["modelName", "effectiveModelName"]) {
+    if (invocation[key] !== undefined && !isSafeString(invocation[key], 512)) return false;
+  }
+  for (const key of ["thinking", "effectiveThinking"]) {
+    if (invocation[key] !== undefined && (typeof invocation[key] !== "string" || !THINKING_LEVELS.has(invocation[key]))) return false;
+  }
+  if (invocation.maxTurns !== undefined && (!Number.isInteger(invocation.maxTurns) || (invocation.maxTurns as number) < 0)) return false;
+  for (const key of ["isolated", "inheritContext", "runInBackground"]) {
+    if (invocation[key] !== undefined && typeof invocation[key] !== "boolean") return false;
+  }
+  return invocation.isolation === undefined || invocation.isolation === "worktree";
+}
+
+function isActiveStatus(value: AgentRecoveryStatus): value is ActiveStatus {
+  return value === "running" || value === "queued";
+}
+
+function isTerminalStatus(value: AgentRecoveryStatus): value is TerminalStatus {
+  return TERMINAL_STATUSES.has(value as TerminalStatus);
+}
+
+function isSafeTranscriptPath(value: unknown): value is string {
+  return typeof value === "string"
+    && /^\.pi-subagents\/agent-transcripts\/[^/]+\.jsonl$/.test(value)
+    && !value.includes("..")
+    && !value.includes("\\")
+    && !value.includes("\0");
+}
+
+/** Validate untrusted JSON before it can enter the manager or UI. */
+export function isAgentRecoveryCheckpoint(value: unknown): value is AgentRecoveryCheckpoint {
+  if (!value || typeof value !== "object") return false;
+  const checkpoint = value as Record<string, unknown>;
+  if (checkpoint.version !== CHECKPOINT_VERSION
+    || !isSafeString(checkpoint.id, 256)
+    || !isSafeString(checkpoint.type, 256)
+    || !isSafeString(checkpoint.description, 4096)
+    || typeof checkpoint.status !== "string"
+    || !isFiniteTimestamp(checkpoint.startedAt)
+    || !Number.isInteger(checkpoint.toolUses)
+    || (checkpoint.toolUses as number) < 0
+    || !isUsage(checkpoint.lifetimeUsage)
+    || !Number.isInteger(checkpoint.compactionCount)
+    || (checkpoint.compactionCount as number) < 0) return false;
+
+  const status = checkpoint.status as AgentRecoveryStatus;
+  if (!isActiveStatus(status) && !isTerminalStatus(status)) return false;
+  if (checkpoint.completedAt !== undefined && !isFiniteTimestamp(checkpoint.completedAt)) return false;
+  if (isTerminalStatus(status)
+    && (checkpoint.completedAt === undefined || checkpoint.completedAt < checkpoint.startedAt)) return false;
+  if (isActiveStatus(status) && checkpoint.completedAt !== undefined) return false;
+  if (checkpoint.result !== undefined && !isSafeString(checkpoint.result, 2_000_000)) return false;
+  if (checkpoint.error !== undefined && !isSafeString(checkpoint.error, 64_000)) return false;
+  if (checkpoint.transcriptPath !== undefined && !isSafeTranscriptPath(checkpoint.transcriptPath)) return false;
+  if (checkpoint.invocation !== undefined && !isInvocation(checkpoint.invocation)) return false;
+  return true;
+}
+
+function checkpointDirectory(cwd: string): string {
+  return join(cwd, SUBAGENTS_DIR, CHECKPOINTS_DIR);
+}
+
+/** Return the on-disk path for an agent's single deduplicated checkpoint. */
+export function agentRecoveryCheckpointPath(cwd: string, agentId: string): string {
+  const safeId = agentId.replace(/[^A-Za-z0-9._-]+/g, "-") || "agent";
+  return join(checkpointDirectory(cwd), `${safeId}.json`);
+}
+
+/**
+ * Atomically write one checkpoint. Repeated writes for the same agent replace
+ * the same file; identical payloads are skipped, so shutdown + abort callbacks
+ * cannot create duplicate recovery records.
+ */
+export function removeAgentRecoveryCheckpoint(cwd: string, agentId: string): boolean {
+  try {
+    unlinkSync(agentRecoveryCheckpointPath(cwd, agentId));
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+export function writeAgentRecoveryCheckpoint(cwd: string, checkpoint: AgentRecoveryCheckpoint): boolean {
+  if (!isAgentRecoveryCheckpoint(checkpoint)) return false;
+  try {
+    ensureSubagentsGitignore(cwd);
+    const directory = checkpointDirectory(cwd);
+    mkdirSync(directory, { recursive: true });
+    const path = agentRecoveryCheckpointPath(cwd, checkpoint.id);
+    const contents = `${JSON.stringify(checkpoint)}\n`;
+    try {
+      if (readFileSync(path, "utf8") === contents) return false;
+    } catch {
+      // The file is new, missing, or unreadable; replace it below.
+    }
+    const temporary = `${path}.${process.pid}.${randomUUID()}.tmp`;
+    writeFileSync(temporary, contents, { encoding: "utf8", mode: 0o600 });
+    renameSync(temporary, path);
+    return true;
+  } catch {
+    // Recovery must never make a spawn or shutdown fail. A later checkpoint
+    // gets another chance to persist if the filesystem becomes available.
+    return false;
+  }
+}
+
+/** Load valid checkpoints, ignoring orphan, malformed, and corrupt files. */
+export function readAgentRecoveryCheckpoints(cwd: string): AgentRecoveryCheckpoint[] {
+  const directory = checkpointDirectory(cwd);
+  if (!existsSync(directory)) return [];
+  let names: string[];
+  try {
+    names = readdirSync(directory).filter((name) => /^[A-Za-z0-9._-]+\.json$/.test(name));
+  } catch {
+    return [];
+  }
+
+  const latest = new Map<string, AgentRecoveryCheckpoint>();
+  for (const name of names) {
+    try {
+      const value: unknown = JSON.parse(readFileSync(join(directory, name), "utf8"));
+      if (!isAgentRecoveryCheckpoint(value)) continue;
+      latest.set(value.id, value);
+    } catch {
+      // Ignore a partially written/corrupt checkpoint and continue indexing.
+    }
+  }
+  return [...latest.values()];
+}