@ilya-lesikov/pi-pi 0.8.0 → 0.10.0

package/3p/pi-subagents/src/agent-manager.ts

@@ -7,18 +7,45 @@
  */
 
 import { randomUUID } from "node:crypto";
+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 { resumeAgent, runAgent, type ToolActivity } from "./agent-runner.js";
-import type { AgentRecord, IsolationMode, SubagentType, ThinkingLevel } from "./types.js";
+import type { AgentInvocation, AgentRecord, IsolationMode, SubagentType, ThinkingLevel } from "./types.js";
+import { addUsage } from "./usage.js";
 import { cleanupWorktree, createWorktree, pruneWorktrees, } from "./worktree.js";
 
 export type OnAgentComplete = (record: AgentRecord) => void;
 export type OnAgentStart = (record: AgentRecord) => void;
+export type OnAgentCompact = (record: AgentRecord, info: CompactionInfo) => void;
+export type CompactionInfo = { reason: "manual" | "threshold" | "overflow"; tokensBefore: number };
 
 /** Default max concurrent background agents. */
 const DEFAULT_MAX_CONCURRENT = 4;
 
+/**
+ * Validate a caller-supplied SpawnOptions.cwd. `undefined`/`null` mean "unset"
+ * (parent cwd). Anything else must be an absolute path to an existing
+ * directory — curated errors instead of TypeErrors from path/fs internals
+ * (RPC callers send arbitrary JSON: null, numbers, file paths).
+ */
+function assertValidSpawnCwd(cwd: unknown): asserts cwd is string | undefined | null {
+  if (cwd == null) return;
+  if (typeof cwd !== "string" || !isAbsolute(cwd)) {
+    throw new Error(`SpawnOptions.cwd must be an absolute path: "${String(cwd)}"`);
+  }
+  let isDirectory = false;
+  try {
+    isDirectory = statSync(cwd).isDirectory();
+  } catch {
+    throw new Error(`SpawnOptions.cwd does not exist: "${cwd}"`);
+  }
+  if (!isDirectory) {
+    throw new Error(`SpawnOptions.cwd is not a directory: "${cwd}"`);
+  }
+}
+
 interface SpawnArgs {
   pi: ExtensionAPI;
   ctx: ExtensionContext;
@@ -35,8 +62,27 @@ interface SpawnOptions {
   inheritContext?: boolean;
   thinkingLevel?: ThinkingLevel;
   isBackground?: boolean;
+  /**
+   * Skip the maxConcurrent queue check for this spawn — start immediately even
+   * if the configured concurrency limit would otherwise queue it. Used by the
+   * scheduler so a fired job can't be deferred past its trigger window.
+   */
+  bypassQueue?: boolean;
   /** Isolation mode — "worktree" creates a temp git worktree for the agent. */
   isolation?: IsolationMode;
+  /**
+   * Working directory for the agent (absolute path). Default: parent session
+   * cwd. The agent's tools operate here, but .pi config (extensions, skills,
+   * settings, memory) still loads from the parent session's project — the
+   * target directory's `.pi` extensions never execute. With isolation:
+   * "worktree", the worktree is created FROM this directory and the result
+   * branch lands in that repo.
+   */
+  cwd?: string;
+  /** Resolved invocation snapshot captured for UI display. */
+  invocation?: AgentInvocation;
+  /** Parent abort signal — when aborted, the subagent is also stopped. */
+  signal?: AbortSignal;
   /** Called on tool start/end with activity info (for streaming progress to UI). */
   onToolActivity?: (activity: ToolActivity) => void;
   /** Called on streaming text deltas from the assistant response. */
@@ -45,10 +91,10 @@ interface SpawnOptions {
   onSessionCreated?: (session: AgentSession) => void;
   /** Called at the end of each agentic turn with the cumulative count. */
   onTurnEnd?: (turnCount: number) => void;
-  validateCompletion?: () => string | undefined;
-  maxValidationRetries?: number;
-  /** toolCallId of the parent tool call that spawned this subagent (for trace correlation). */
-  toolCallId?: string;
+  /** Called once per assistant message_end with that message's usage delta. */
+  onAssistantUsage?: (usage: { input: number; output: number; cacheWrite: number }) => void;
+  /** Called when the session successfully compacts. */
+  onCompaction?: (info: CompactionInfo) => void;
 }
 
 export class AgentManager {
@@ -56,19 +102,30 @@ export class AgentManager {
   private cleanupInterval: ReturnType<typeof setInterval>;
   private onComplete?: OnAgentComplete;
   private onStart?: OnAgentStart;
+  private onCompact?: OnAgentCompact;
   private maxConcurrent: number;
+  /** 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>();
 
   /** Queue of background agents waiting to start. */
   private queue: { id: string; args: SpawnArgs }[] = [];
   /** Number of currently running background agents. */
   private runningBackground = 0;
 
-  constructor(onComplete?: OnAgentComplete, maxConcurrent = DEFAULT_MAX_CONCURRENT, onStart?: OnAgentStart) {
+  constructor(
+    onComplete?: OnAgentComplete,
+    maxConcurrent = DEFAULT_MAX_CONCURRENT,
+    onStart?: OnAgentStart,
+    onCompact?: OnAgentCompact,
+  ) {
     this.onComplete = onComplete;
     this.onStart = onStart;
+    this.onCompact = onCompact;
     this.maxConcurrent = maxConcurrent;
     // Cleanup completed agents after 10 minutes (but keep sessions for resume)
     this.cleanupInterval = setInterval(() => this.cleanup(), 60_000);
+    this.cleanupInterval.unref();
   }
 
   /** Update the max concurrent background agents limit. */
@@ -93,6 +150,11 @@ export class AgentManager {
     prompt: string,
     options: SpawnOptions,
   ): string {
+    // Validate before the queue branch — a queued spawn should fail at the
+    // call, not minutes later at drain. Throw (not warn): programmatic callers
+    // can fix and retry; the RPC layer converts throws into error envelopes.
+    assertValidSpawnCwd(options.cwd);
+
     const id = randomUUID().slice(0, 17);
     const abortController = new AbortController();
     const record: AgentRecord = {
@@ -103,99 +165,116 @@ export class AgentManager {
       toolUses: 0,
       startedAt: Date.now(),
       abortController,
+      lifetimeUsage: { input: 0, output: 0, cacheWrite: 0 },
+      compactionCount: 0,
+      // Raw tri-state (not coerced to a boolean): true = background, false =
+      // foreground (has an inline tool-result surface), undefined = caller never
+      // declared it (e.g. a cross-extension RPC spawn). The widget's background-
+      // only filter excludes only explicit `false`, so undefined agents — which
+      // have no inline surface — stay visible instead of vanishing.
+      isBackground: options.isBackground,
+      invocation: options.invocation,
     };
     this.agents.set(id, record);
 
     const args: SpawnArgs = { pi, ctx, type, prompt, options };
 
-    if (options.isBackground && this.runningBackground >= this.maxConcurrent) {
+    if (options.isBackground && !options.bypassQueue && this.runningBackground >= this.maxConcurrent) {
       // Queue it — will be started when a running agent completes
       this.queue.push({ id, args });
       return id;
     }
 
-    this.startAgent(id, record, args);
+    // startAgent can throw (e.g. strict worktree-isolation failure) — clean
+    // up the record so callers don't see an orphan in `listAgents()`.
+    try {
+      this.startAgent(id, record, args);
+    } catch (err) {
+      this.agents.delete(id);
+      throw err;
+    }
     return id;
   }
 
   /** Actually start an agent (called immediately or from queue drain). */
   private startAgent(id: string, record: AgentRecord, { pi, ctx, type, prompt, options }: SpawnArgs) {
+    // Re-validate a caller-supplied cwd: queued spawns can start minutes after
+    // spawn()'s check, and the directory may be gone by then (TOCTOU). Same
+    // curated errors; drainQueue parks a throw on the record as an error.
+    assertValidSpawnCwd(options.cwd);
+    // Single resolution point for the caller-supplied cwd — the worktree base
+    // repo and both cleanup calls below MUST agree on this value forever.
+    const customCwd = options.cwd ?? undefined; // null (RPC "unset") → undefined
+    const baseCwd = customCwd ?? ctx.cwd;
+
+    // Worktree isolation: try to create a temporary git worktree. Strict —
+    // fail loud if not possible (no silent fallback to main tree). Done
+    // BEFORE state mutation so a throw doesn't leave the record half-running.
+    let worktreeCwd: string | undefined;
+    if (options.isolation === "worktree") {
+      const wt = createWorktree(baseCwd, id);
+      if (!wt) {
+        throw new Error(
+          'Cannot run with isolation: "worktree" — not a git repo, no commits yet, or `git worktree add` failed. ' +
+          'Initialize git and commit at least once, or omit `isolation`.',
+        );
+      }
+      record.worktree = wt;
+      // workPath preserves subdirectory scoping for caller-supplied cwds: a
+      // cwd deep in a monorepo maps to the same subdir inside the copy, not
+      // the copied repo's root. Plain worktree spawns keep the historical
+      // behavior (agent at the copy's root) — moving them to workPath would
+      // also move .pi config discovery when the parent session sits in a repo
+      // subdirectory, silently dropping extensions/skills.
+      worktreeCwd = customCwd !== undefined ? wt.workPath : wt.path;
+      this.worktreeRepos.add(baseCwd);
+    }
+
     record.status = "running";
     record.startedAt = Date.now();
-    if (options.toolCallId && !record.toolCallId) record.toolCallId = options.toolCallId;
-    // Reset first-progress flags: a re-started (e.g. dequeued) record must be able
-    // to emit first_tool/first_turn again for its fresh run.
-    record.firstToolEmitted = false;
-    record.firstTurnEmitted = false;
     if (options.isBackground) this.runningBackground++;
     this.onStart?.(record);
 
-    // Worktree isolation: create a temporary git worktree if requested
-    let worktreeCwd: string | undefined;
-    let worktreeWarning = "";
-    if (options.isolation === "worktree") {
-      const wt = createWorktree(ctx.cwd, id);
-      if (wt) {
-        record.worktree = wt;
-        worktreeCwd = wt.path;
-      } else {
-        worktreeWarning = "\n\n[WARNING: Worktree isolation was requested but failed (not a git repo, or no commits yet). Running in the main working directory instead.]";
-      }
+    // Wire parent abort signal to stop the subagent when the parent is interrupted
+    let detachParentSignal: (() => void) | undefined;
+    if (options.signal) {
+      const onParentAbort = () => this.abort(id);
+      options.signal.addEventListener("abort", onParentAbort, { once: true });
+      detachParentSignal = () => options.signal!.removeEventListener("abort", onParentAbort);
     }
+    const detach = () => { detachParentSignal?.(); detachParentSignal = undefined; };
 
-    // Prepend worktree warning to prompt if isolation failed
-    const effectivePrompt = worktreeWarning ? worktreeWarning + "\n\n" + prompt : prompt;
-
-    const promise = runAgent(ctx, type, effectivePrompt, {
+    const promise = runAgent(ctx, type, prompt, {
       pi,
+      agentId: id,
       model: options.model,
       maxTurns: options.maxTurns,
       isolated: options.isolated,
       inheritContext: options.inheritContext,
       thinkingLevel: options.thinkingLevel,
-      cwd: worktreeCwd,
+      // Worktree wins for the working dir (the agent must run in the copy —
+      // which, with a custom cwd, was created from that target). Config stays
+      // with the parent project when a caller-supplied cwd is in play; it must
+      // stay undefined otherwise so plain worktree runs keep resolving config
+      // (incl. relative extension paths and memory) inside the worktree copy.
+      cwd: worktreeCwd ?? customCwd,
+      configCwd: customCwd !== undefined ? ctx.cwd : undefined,
       signal: record.abortController!.signal,
       onToolActivity: (activity) => {
         if (activity.type === "end") record.toolUses++;
-        // Emit first_tool once per run from this single choke point, so ALL spawn
-        // paths (RPC panels, Agent-tool background, foreground) report it — not
-        // just the Agent tool's background branch. Panels spawn via RPC and were
-        // previously silent, making it impossible to tell if a reviewer actually
-        // investigated (ran tools) or wrote from context alone.
-        if (activity.type === "start" && !record.firstToolEmitted) {
-          record.firstToolEmitted = true;
-          try {
-            pi.events.emit("subagents:first_tool", {
-              id,
-              type,
-              description: options.description,
-              toolName: activity.toolName,
-            });
-          } catch { /* ignore */ }
-        }
         options.onToolActivity?.(activity);
       },
-      onTurnEnd: (turnCount) => {
-        if (!record.firstTurnEmitted) {
-          record.firstTurnEmitted = true;
-          try {
-            pi.events.emit("subagents:first_turn", {
-              id,
-              type,
-              description: options.description,
-              turnCount,
-            });
-          } catch { /* ignore */ }
-        }
-        options.onTurnEnd?.(turnCount);
-      },
+      onTurnEnd: options.onTurnEnd,
       onTextDelta: options.onTextDelta,
-      validateCompletion: options.validateCompletion,
-      maxValidationRetries: options.maxValidationRetries,
-      subagentId: id,
-      subagentType: type,
-      subagentDescription: options.description,
-      parentToolCallId: options.toolCallId ?? record.toolCallId,
+      onAssistantUsage: (usage) => {
+        addUsage(record.lifetimeUsage, usage);
+        options.onAssistantUsage?.(usage);
+      },
+      onCompaction: (info) => {
+        record.compactionCount++;
+        this.onCompact?.(record, info);
+        options.onCompaction?.(info);
+      },
       onSessionCreated: (session) => {
         record.session = session;
         // Flush any steers that arrived before the session was ready
@@ -217,6 +296,8 @@ export class AgentManager {
         record.session = session;
         record.completedAt ??= Date.now();
 
+        detach();
+
         // Final flush of streaming output file
         if (record.outputCleanup) {
           try { record.outputCleanup(); } catch { /* ignore */ }
@@ -225,17 +306,25 @@ export class AgentManager {
 
         // Clean up worktree if used
         if (record.worktree) {
-          const wtResult = cleanupWorktree(ctx.cwd, record.worktree, options.description);
+          const wtResult = cleanupWorktree(baseCwd, record.worktree, options.description);
           record.worktreeResult = wtResult;
           if (wtResult.hasChanges && wtResult.branch) {
+            // With a caller-supplied cwd the branch lives in THAT repo, not the
+            // parent session's — say so, or the orchestrator merges in the wrong repo.
+            const repoNote = customCwd !== undefined ? ` in \`${baseCwd}\`` : "";
             record.result = (record.result ?? "") +
-              `\n\n---\nChanges saved to branch \`${wtResult.branch}\`. Merge with: \`git merge ${wtResult.branch}\``;
+              `\n\n---\nChanges saved to branch \`${wtResult.branch}\`${repoNote}. Merge with: \`git merge ${wtResult.branch}\`${customCwd !== undefined ? ` (run in \`${baseCwd}\`)` : ""}`;
           }
         }
 
-        if (options.isBackground) {
+        // Fire onComplete for foreground agents too — lifecycle symmetry.
+        // Mark resultConsumed so the callback skips notifications (result returned inline).
+        if (!options.isBackground) {
+          record.resultConsumed = true;
+          try { this.onComplete?.(record); } catch { /* ignore completion side-effect errors */ }
+        } else {
           this.runningBackground--;
-          this.onComplete?.(record);
+          try { this.onComplete?.(record); } catch { /* ignore completion side-effect errors */ }
           this.drainQueue();
         }
         return responseText;
@@ -248,6 +337,8 @@ export class AgentManager {
         record.error = err instanceof Error ? err.message : String(err);
         record.completedAt ??= Date.now();
 
+        detach();
+
         // Final flush of streaming output file on error
         if (record.outputCleanup) {
           try { record.outputCleanup(); } catch { /* ignore */ }
@@ -257,12 +348,17 @@ export class AgentManager {
         // Best-effort worktree cleanup on error
         if (record.worktree) {
           try {
-            const wtResult = cleanupWorktree(ctx.cwd, record.worktree, options.description);
+            const wtResult = cleanupWorktree(baseCwd, record.worktree, options.description);
             record.worktreeResult = wtResult;
           } catch { /* ignore cleanup errors */ }
         }
 
-        if (options.isBackground) {
+        // Fire onComplete for foreground agents too — lifecycle symmetry.
+        // Mark resultConsumed so the callback skips notifications (result returned inline).
+        if (!options.isBackground) {
+          record.resultConsumed = true;
+          this.onComplete?.(record);
+        } else {
           this.runningBackground--;
           this.onComplete?.(record);
           this.drainQueue();
@@ -271,6 +367,11 @@ 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. */
@@ -279,13 +380,33 @@ export class AgentManager {
       const next = this.queue.shift()!;
       const record = this.agents.get(next.id);
       if (!record || record.status !== "queued") continue;
-      this.startAgent(next.id, record, next.args);
+      try {
+        this.startAgent(next.id, record, next.args);
+      } catch (err) {
+        // Late failure (e.g. strict worktree-isolation) — surface on the record
+        // so the user/agent can see it via /agents, then keep draining.
+        record.status = "error";
+        record.error = err instanceof Error ? err.message : String(err);
+        record.completedAt = Date.now();
+        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.
+   * Returns { id, record } so callers can access the agent ID.
+   *
+   * @param onSpawned - Called synchronously after spawn(), before onSessionCreated fires.
+   *   Use this to set record.outputFile so streamToOutputFile can pick it up.
    */
   async spawnAndWait(
     pi: ExtensionAPI,
@@ -293,11 +414,19 @@ export class AgentManager {
     type: SubagentType,
     prompt: string,
     options: Omit<SpawnOptions, "isBackground">,
-  ): Promise<AgentRecord> {
-    const id = this.spawn(pi, ctx, type, prompt, { ...options, isBackground: false });
-    const record = this.agents.get(id)!;
-    await record.promise;
-    return record;
+    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;
+    }
   }
 
   /**
@@ -322,8 +451,14 @@ export class AgentManager {
         onToolActivity: (activity) => {
           if (activity.type === "end") record.toolUses++;
         },
+        onAssistantUsage: (usage) => {
+          addUsage(record.lifetimeUsage, usage);
+        },
+        onCompaction: (info) => {
+          record.compactionCount++;
+          this.onCompact?.(record, info);
+        },
         signal,
-        subagentId: id,
       });
       record.status = "completed";
       record.result = responseText;
@@ -337,6 +472,27 @@ export class AgentManager {
     return record;
   }
 
+  /**
+   * Send a steering message to an agent from the UI (mirrors the steer_subagent
+   * tool). A live session delivers it now — it interrupts the agent after its
+   * current tool execution and appears as a user message. If the session isn't
+   * ready yet, the message is queued on `pendingSteers` and flushed when the
+   * session is created. Returns false if the agent can't accept steering
+   * (unknown id, or no longer running/queued).
+   */
+  steer(id: string, message: string): boolean {
+    const record = this.agents.get(id);
+    if (!record) return false;
+    if (record.status !== "running" && record.status !== "queued") return false;
+    if (record.session) {
+      record.session.steer(message).catch(() => {});
+    } else {
+      if (!record.pendingSteers) record.pendingSteers = [];
+      record.pendingSteers.push(message);
+    }
+    return true;
+  }
+
   getRecord(id: string): AgentRecord | undefined {
     return this.agents.get(id);
   }
@@ -385,10 +541,13 @@ export class AgentManager {
   /**
    * Remove all completed/stopped/errored records immediately.
    * Called on session start/switch so tasks from a prior session don't persist.
+   * Pass skipUnconsumed=true to preserve records the LLM hasn't read yet
+   * (resultConsumed=false) — they will be evicted by the 10-minute cleanup timer instead.
    */
-  clearCompleted(): void {
+  clearCompleted(skipUnconsumed = false): void {
     for (const [id, record] of this.agents) {
       if (record.status === "running" || record.status === "queued") continue;
+      if (skipUnconsumed && !record.resultConsumed) continue;
       this.removeRecord(id, record);
     }
   }
@@ -450,5 +609,10 @@ export class AgentManager {
     this.agents.clear();
     // 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
+    // exit with in-flight agents would otherwise leave stale registrations there.
+    for (const repo of this.worktreeRepos) {
+      try { pruneWorktrees(repo); } catch { /* ignore */ }
+    }
   }
 }