@esso0428/pi-subagents 0.15.0

package/dist/agent-manager.js

@@ -0,0 +1,542 @@
+/**
+ * agent-manager.ts — Tracks agents, background execution, resume support.
+ *
+ * Background agents are subject to a configurable concurrency limit (default: 4).
+ * Excess agents are queued and auto-started as running agents complete.
+ * Foreground agents bypass the queue (they block the parent anyway).
+ */
+import { randomUUID } from "node:crypto";
+import { statSync } from "node:fs";
+import { isAbsolute } from "node:path";
+import { resumeAgent, runAgent } from "./agent-runner.js";
+import { addUsage } from "./usage.js";
+import { cleanupWorktree, createWorktree, pruneWorktrees, } from "./worktree.js";
+/** 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) {
+    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}"`);
+    }
+}
+export class AgentManager {
+    agents = new Map();
+    cleanupInterval;
+    onComplete;
+    onStart;
+    onCompact;
+    maxConcurrent;
+    /** Base repos worktrees were created from — so dispose() can prune them all,
+     *  not just the parent repo (caller-supplied cwd can target other repos). */
+    worktreeRepos = new Set();
+    /** Queue of background agents waiting to start. */
+    queue = [];
+    /** Number of currently running background agents. */
+    runningBackground = 0;
+    constructor(onComplete, maxConcurrent = DEFAULT_MAX_CONCURRENT, onStart, onCompact) {
+        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. */
+    setMaxConcurrent(n) {
+        this.maxConcurrent = Math.max(1, n);
+        // Start queued agents if the new limit allows
+        this.drainQueue();
+    }
+    getMaxConcurrent() {
+        return this.maxConcurrent;
+    }
+    /**
+     * Spawn an agent and return its ID immediately (for background use).
+     * If the concurrency limit is reached, the agent is queued.
+     */
+    spawn(pi, ctx, type, prompt, options) {
+        // 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 = {
+            id,
+            type,
+            description: options.description,
+            status: options.isBackground ? "queued" : "running",
+            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 = { pi, ctx, type, prompt, options };
+        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;
+        }
+        // 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). */
+    startAgent(id, record, { pi, ctx, type, prompt, options }) {
+        // 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;
+        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.isBackground)
+            this.runningBackground++;
+        this.onStart?.(record);
+        // Wire parent abort signal to stop the subagent when the parent is interrupted
+        let detachParentSignal;
+        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; };
+        const promise = runAgent(ctx, type, prompt, {
+            pi,
+            agentId: id,
+            model: options.model,
+            maxTurns: options.maxTurns,
+            isolated: options.isolated,
+            inheritContext: options.inheritContext,
+            thinkingLevel: options.thinkingLevel,
+            // 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++;
+                options.onToolActivity?.(activity);
+            },
+            onTurnEnd: options.onTurnEnd,
+            onTextDelta: options.onTextDelta,
+            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
+                if (record.pendingSteers?.length) {
+                    for (const msg of record.pendingSteers) {
+                        session.steer(msg).catch(() => { });
+                    }
+                    record.pendingSteers = undefined;
+                }
+                options.onSessionCreated?.(session);
+            },
+        })
+            .then(({ responseText, session, aborted, steered, failure }) => {
+            // Don't overwrite status if externally stopped via abort()
+            if (record.status !== "stopped") {
+                // Precedence: a hard abort keeps "aborted"; then a failed final turn
+                // (provider error that pi resolved instead of rejecting, #144) is an
+                // honest "error" — not a completion with an empty or stale result.
+                if (aborted) {
+                    record.status = "aborted";
+                }
+                else if (failure) {
+                    record.status = "error";
+                    record.error = failure;
+                }
+                else {
+                    record.status = steered ? "steered" : "completed";
+                }
+            }
+            record.result = responseText;
+            record.session = session;
+            record.completedAt ??= Date.now();
+            detach();
+            // Final flush of streaming output file
+            if (record.outputCleanup) {
+                try {
+                    record.outputCleanup();
+                }
+                catch { /* ignore */ }
+                record.outputCleanup = undefined;
+            }
+            // Clean up worktree if used
+            if (record.worktree) {
+                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}\`${repoNote}. Merge with: \`git merge ${wtResult.branch}\`${customCwd !== undefined ? ` (run in \`${baseCwd}\`)` : ""}`;
+                }
+            }
+            // 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--;
+                try {
+                    this.onComplete?.(record);
+                }
+                catch { /* ignore completion side-effect errors */ }
+                this.drainQueue();
+            }
+            return responseText;
+        })
+            .catch((err) => {
+            // Don't overwrite status if externally stopped via abort()
+            if (record.status !== "stopped") {
+                record.status = "error";
+            }
+            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 */ }
+                record.outputCleanup = undefined;
+            }
+            // Best-effort worktree cleanup on error
+            if (record.worktree) {
+                try {
+                    const wtResult = cleanupWorktree(baseCwd, record.worktree, options.description);
+                    record.worktreeResult = wtResult;
+                }
+                catch { /* ignore cleanup errors */ }
+            }
+            // 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();
+            }
+            return "";
+        });
+        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. */
+    drainQueue() {
+        while (this.queue.length > 0 && this.runningBackground < this.maxConcurrent) {
+            const next = this.queue.shift();
+            const record = this.agents.get(next.id);
+            if (!record || record.status !== "queued")
+                continue;
+            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.
+     */
+    onSpawned;
+    /**
+     * 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, ctx, type, prompt, options, onSpawned) {
+        // 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;
+        }
+    }
+    /**
+     * Resume an existing agent session with a new prompt.
+     */
+    async resume(id, prompt, signal) {
+        const record = this.agents.get(id);
+        if (!record?.session)
+            return undefined;
+        record.status = "running";
+        record.startedAt = Date.now();
+        record.completedAt = undefined;
+        record.result = undefined;
+        record.error = undefined;
+        try {
+            const { text, failure } = await resumeAgent(record.session, prompt, {
+                onToolActivity: (activity) => {
+                    if (activity.type === "end")
+                        record.toolUses++;
+                },
+                onAssistantUsage: (usage) => {
+                    addUsage(record.lifetimeUsage, usage);
+                },
+                onCompaction: (info) => {
+                    record.compactionCount++;
+                    this.onCompact?.(record, info);
+                },
+                signal,
+            });
+            // Same contract as the spawn path (#144): a failed final turn is an
+            // error, not a completion — but the resumed text stays available.
+            record.status = failure ? "error" : "completed";
+            if (failure)
+                record.error = failure;
+            record.result = text;
+            record.completedAt = Date.now();
+        }
+        catch (err) {
+            record.status = "error";
+            record.error = err instanceof Error ? err.message : String(err);
+            record.completedAt = Date.now();
+        }
+        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, message) {
+        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) {
+        return this.agents.get(id);
+    }
+    listAgents() {
+        return [...this.agents.values()].sort((a, b) => b.startedAt - a.startedAt);
+    }
+    abort(id) {
+        const record = this.agents.get(id);
+        if (!record)
+            return false;
+        // Remove from queue if queued
+        if (record.status === "queued") {
+            this.queue = this.queue.filter(q => q.id !== id);
+            record.status = "stopped";
+            record.completedAt = Date.now();
+            return true;
+        }
+        if (record.status !== "running")
+            return false;
+        record.abortController?.abort();
+        record.status = "stopped";
+        record.completedAt = Date.now();
+        return true;
+    }
+    /** Dispose a record's session and remove it from the map. */
+    removeRecord(id, record) {
+        record.session?.dispose?.();
+        record.session = undefined;
+        this.agents.delete(id);
+    }
+    cleanup() {
+        const cutoff = Date.now() - 10 * 60_000;
+        for (const [id, record] of this.agents) {
+            if (record.status === "running" || record.status === "queued")
+                continue;
+            if ((record.completedAt ?? 0) >= cutoff)
+                continue;
+            this.removeRecord(id, record);
+        }
+    }
+    /**
+     * 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(skipUnconsumed = false) {
+        for (const [id, record] of this.agents) {
+            if (record.status === "running" || record.status === "queued")
+                continue;
+            if (skipUnconsumed && !record.resultConsumed)
+                continue;
+            this.removeRecord(id, record);
+        }
+    }
+    /** Whether any agents are still running or queued. */
+    hasRunning() {
+        return [...this.agents.values()].some(r => r.status === "running" || r.status === "queued");
+    }
+    /** Abort all running and queued agents immediately. */
+    abortAll() {
+        let count = 0;
+        // Clear queued agents first
+        for (const queued of this.queue) {
+            const record = this.agents.get(queued.id);
+            if (record) {
+                record.status = "stopped";
+                record.completedAt = Date.now();
+                count++;
+            }
+        }
+        this.queue = [];
+        // Abort running agents
+        for (const record of this.agents.values()) {
+            if (record.status === "running") {
+                record.abortController?.abort();
+                record.status = "stopped";
+                record.completedAt = Date.now();
+                count++;
+            }
+        }
+        return count;
+    }
+    /** Wait for all running and queued agents to complete (including queued ones). */
+    async waitForAll() {
+        // Loop because drainQueue respects the concurrency limit — as running
+        // agents finish they start queued ones, which need awaiting too.
+        while (true) {
+            this.drainQueue();
+            const pending = [...this.agents.values()]
+                .filter(r => r.status === "running" || r.status === "queued")
+                .map(r => r.promise)
+                .filter(Boolean);
+            if (pending.length === 0)
+                break;
+            await Promise.allSettled(pending);
+        }
+    }
+    dispose() {
+        clearInterval(this.cleanupInterval);
+        // Clear queue
+        this.queue = [];
+        for (const record of this.agents.values()) {
+            record.session?.dispose();
+        }
+        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 */ }
+        }
+    }
+}