ultimate-pi 0.11.0 → 0.13.0

package/.pi/extensions/lib/harness-subagents/vendored/agent-manager.ts

@@ -1,558 +0,0 @@
-/**
- * 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 type { Model } from "@earendil-works/pi-ai";
-import type {
-	AgentSession,
-	ExtensionAPI,
-	ExtensionContext,
-} from "@earendil-works/pi-coding-agent";
-import type { ParentHarnessUiHooks } from "../parent-harness-ui-bridge.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";
-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;
-
-interface SpawnArgs {
-	pi: ExtensionAPI;
-	ctx: ExtensionContext;
-	type: SubagentType;
-	prompt: string;
-	options: SpawnOptions;
-}
-
-interface SpawnOptions {
-	description: string;
-	model?: Model<any>;
-	maxTurns?: number;
-	isolated?: boolean;
-	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;
-	/** 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. */
-	onTextDelta?: (delta: string, fullText: string) => void;
-	/** Called when the agent session is created (for accessing session stats). */
-	onSessionCreated?: (session: AgentSession) => void;
-	/** Called at the end of each agentic turn with the cumulative count. */
-	onTurnEnd?: (turnCount: number) => void;
-	/** 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;
-	/** Spawn context (blackboard injection) for subagent system prompt. */
-	systemPromptAppendix?: string;
-	/** Parent UI hooks for plan approval sync and draft transcript. */
-	parentHarnessUiHooks?: ParentHarnessUiHooks;
-}
-
-export class AgentManager {
-	private agents = new Map<string, AgentRecord>();
-	private cleanupInterval: ReturnType<typeof setInterval>;
-	private onComplete?: OnAgentComplete;
-	private onStart?: OnAgentStart;
-	private onCompact?: OnAgentCompact;
-	private maxConcurrent: number;
-
-	/** 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,
-		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. */
-	setMaxConcurrent(n: number) {
-		this.maxConcurrent = Math.max(1, n);
-		// Start queued agents if the new limit allows
-		this.drainQueue();
-	}
-
-	getMaxConcurrent(): number {
-		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: ExtensionAPI,
-		ctx: ExtensionContext,
-		type: SubagentType,
-		prompt: string,
-		options: SpawnOptions,
-	): string {
-		const id = randomUUID().slice(0, 17);
-		const abortController = new AbortController();
-		const record: AgentRecord = {
-			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,
-			invocation: options.invocation,
-		};
-		this.agents.set(id, record);
-
-		const args: SpawnArgs = { 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). */
-	private startAgent(
-		id: string,
-		record: AgentRecord,
-		{ pi, ctx, type, prompt, options }: SpawnArgs,
-	) {
-		// 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(ctx.cwd, 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;
-			worktreeCwd = wt.path;
-		}
-
-		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: (() => 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;
-		};
-
-		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,
-			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);
-			},
-			systemPromptAppendix: options.systemPromptAppendix,
-			parentExtensionContext: ctx,
-			parentHarnessUiHooks: options.parentHarnessUiHooks,
-		})
-			.then(({ responseText, session, aborted, steered }) => {
-				// Don't overwrite status if externally stopped via abort()
-				if (record.status !== "stopped") {
-					record.status = aborted
-						? "aborted"
-						: 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(
-						ctx.cwd,
-						record.worktree,
-						options.description,
-					);
-					record.worktreeResult = wtResult;
-					if (wtResult.hasChanges && wtResult.branch) {
-						record.result =
-							(record.result ?? "") +
-							`\n\n---\nChanges saved to branch \`${wtResult.branch}\`. Merge with: \`git merge ${wtResult.branch}\``;
-					}
-				}
-
-				if (options.isBackground) {
-					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(
-							ctx.cwd,
-							record.worktree,
-							options.description,
-						);
-						record.worktreeResult = wtResult;
-					} catch {
-						/* ignore cleanup errors */
-					}
-				}
-
-				if (options.isBackground) {
-					this.runningBackground--;
-					this.onComplete?.(record);
-					this.drainQueue();
-				}
-				return "";
-			});
-
-		record.promise = promise;
-	}
-
-	/** Start queued agents up to the concurrency limit. */
-	private 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);
-			}
-		}
-	}
-
-	/**
-	 * Spawn an agent and wait for completion (foreground use).
-	 * Foreground agents bypass the concurrency queue.
-	 */
-	async spawnAndWait(
-		pi: ExtensionAPI,
-		ctx: ExtensionContext,
-		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;
-	}
-
-	/**
-	 * Resume an existing agent session with a new prompt.
-	 */
-	async resume(
-		id: string,
-		prompt: string,
-		signal?: AbortSignal,
-	): Promise<AgentRecord | undefined> {
-		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 responseText = 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,
-			});
-			record.status = "completed";
-			record.result = responseText;
-			record.completedAt = Date.now();
-		} catch (err) {
-			record.status = "error";
-			record.error = err instanceof Error ? err.message : String(err);
-			record.completedAt = Date.now();
-		}
-
-		return record;
-	}
-
-	getRecord(id: string): AgentRecord | undefined {
-		return this.agents.get(id);
-	}
-
-	listAgents(): AgentRecord[] {
-		return [...this.agents.values()].sort((a, b) => b.startedAt - a.startedAt);
-	}
-
-	abort(id: string): boolean {
-		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. */
-	private removeRecord(id: string, record: AgentRecord): void {
-		record.session?.dispose?.();
-		record.session = undefined;
-		this.agents.delete(id);
-	}
-
-	private 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.
-	 */
-	clearCompleted(): void {
-		for (const [id, record] of this.agents) {
-			if (record.status === "running" || record.status === "queued") continue;
-			this.removeRecord(id, record);
-		}
-	}
-
-	/** Whether any agents are still running or queued. */
-	hasRunning(): boolean {
-		return [...this.agents.values()].some(
-			(r) => r.status === "running" || r.status === "queued",
-		);
-	}
-
-	/** Abort all running and queued agents immediately. */
-	abortAll(): number {
-		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(): Promise<void> {
-		// 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 */
-		}
-	}
-}