/**
 * 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 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";
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;
const RESTORABLE_STATUSES = new Set(["completed", "steered", "stopped", "aborted", "error"] as const);
const THINKING_LEVELS = new Set(["minimal", "low", "medium", "high", "xhigh", "max", "off"]);

type RestorableAgentStatus = "completed" | "steered" | "stopped" | "aborted" | "error";

/** Narrow persisted strings before putting them into runtime/UI state. */
function isSafePersistedString(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 isValidUsage(value: unknown): value is { input: number; output: number; cacheWrite: number } {
  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 isSafeTranscriptLocator(value: unknown): value is string {
  return typeof value === "string"
    && /^\.pi-subagents\/agent-transcripts\/[^/]+\.jsonl$/.test(value)
    && !value.includes("..")
    && !value.includes("\\")
    && !value.includes("\0");
}

function isSafeInvocation(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 && !isSafePersistedString(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;
  }
  if (invocation.isolation !== undefined && invocation.isolation !== "worktree") return false;
  return true;
}

function cloneInvocation(value: AgentInvocation | undefined): AgentInvocation | undefined {
  return value ? {
    modelName: value.modelName,
    effectiveModelName: value.effectiveModelName,
    thinking: value.thinking,
    effectiveThinking: value.effectiveThinking,
    maxTurns: value.maxTurns,
    isolated: value.isolated,
    inheritContext: value.inheritContext,
    runInBackground: value.runInBackground,
    isolation: value.isolation,
  } : undefined;
}

/** Validate one persisted terminal record without constructing runtime handles. */
export function isRestorableAgentRecord(value: unknown): value is {
  id: string;
  type: string;
  description: string;
  status: RestorableAgentStatus;
  startedAt: number;
  completedAt: number;
  result?: string;
  error?: string;
  toolUses?: number;
  lifetimeUsage?: { input: number; output: number; cacheWrite: number };
  transcriptPath?: string;
  invocation?: AgentInvocation;
} {
  if (!value || typeof value !== "object") return false;
  const record = value as Record<string, unknown>;
  if (!isSafePersistedString(record.id, 256)
    || !isSafePersistedString(record.type, 256)
    || !isSafePersistedString(record.description, 4096)
    || typeof record.status !== "string"
    || !RESTORABLE_STATUSES.has(record.status as RestorableAgentStatus)
    || !isFiniteTimestamp(record.startedAt)
    || !isFiniteTimestamp(record.completedAt)
    || record.completedAt < record.startedAt) return false;
  if (record.result !== undefined && !isSafePersistedString(record.result, 2_000_000)) return false;
  if (record.error !== undefined && !isSafePersistedString(record.error, 64_000)) return false;
  if (record.toolUses !== undefined && (!Number.isInteger(record.toolUses) || (record.toolUses as number) < 0)) return false;
  if (record.lifetimeUsage !== undefined && !isValidUsage(record.lifetimeUsage)) return false;
  if (record.transcriptPath !== undefined && !isSafeTranscriptLocator(record.transcriptPath)) return false;
  if (record.invocation !== undefined && !isSafeInvocation(record.invocation)) return false;
  return true;
}

/**
 * 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;
  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;
  /**
   * 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. */
  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;
  /** Called synchronously after the record exists, before it is queued or started. */
  onSpawned?: (id: string) => void;
}

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;
  /** 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,
    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;
  }

  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.
   */
  spawn(
    pi: ExtensionAPI,
    ctx: ExtensionContext,
    type: SubagentType,
    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 = {
      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);
    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 };

    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);
      this.recoveryCwds.delete(id);
      removeAgentRecoveryCheckpoint(ctx.cwd, 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();
    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
    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,
      // 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;
        const model = session.model;
        record.invocation = {
          ...(record.invocation ?? {}),
          ...(model && { effectiveModelName: model.name ?? model.id }),
          effectiveThinking: session.thinkingLevel,
        };
        // 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);
        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
          // (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}\`)` : ""}`;
          }
        }
        this.checkpoint(record);

        // 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.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";
        }
        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 */ }
        }
        this.checkpoint(record);

        // 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.finishBackground(id);
          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.checkpoint(record);
        this.onComplete?.(record);
      }
    }
  }

  /**
   * 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,
    ctx: ExtensionContext,
    type: SubagentType,
    prompt: string,
    options: Omit<SpawnOptions, "isBackground">,
    onSpawned?: (id: string) => void,
  ): Promise<{ id: string; record: AgentRecord }> {
    const id = this.spawn(pi, ctx, type, prompt, {
      ...options,
      isBackground: false,
      onSpawned,
    });
    const record = this.agents.get(id)!;
    await record.promise;
    return { id, 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;
    const resumedModel = record.session.model;
    record.invocation = {
      ...(record.invocation ?? {}),
      ...(resumedModel && { effectiveModelName: resumedModel.name ?? resumedModel.id }),
      effectiveThinking: record.session.thinkingLevel,
    };
    this.checkpoint(record);

    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();
      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;
  }

  /**
   * 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);
  }

  listAgents(): AgentRecord[] {
    return [...this.agents.values()].sort(
      (a, b) => b.startedAt - a.startedAt,
    );
  }

  /** Restore terminal records persisted by a parent branch without runtime handles. */
  restoreCompleted(records: readonly unknown[]): void {
    const latest = new Map<string, ReturnType<typeof this.createRestoredRecord>>();
    for (const value of records) {
      if (isRestorableAgentRecord(value)) {
        latest.set(value.id, this.createRestoredRecord(value));
      }
    }

    for (const [id, restored] of latest) {
      const existing = this.agents.get(id);
      if (existing?.status === "running" || existing?.status === "queued") continue;
      this.agents.set(id, restored);
    }
  }

  private createRestoredRecord(record: {
    id: string;
    type: string;
    description: string;
    status: RestorableAgentStatus;
    startedAt: number;
    completedAt: number;
    result?: string;
    error?: string;
    toolUses?: number;
    lifetimeUsage?: { input: number; output: number; cacheWrite: number };
    transcriptPath?: string;
    invocation?: AgentInvocation;
    compactionCount?: number;
  }): AgentRecord {
    return {
      id: record.id,
      type: record.type,
      description: record.description,
      status: record.status,
      result: record.result,
      error: record.error,
      toolUses: record.toolUses ?? 0,
      startedAt: record.startedAt,
      completedAt: record.completedAt,
      transcriptPath: record.transcriptPath,
      invocation: cloneInvocation(record.invocation),
      lifetimeUsage: record.lifetimeUsage
        ? { ...record.lifetimeUsage }
        : { input: 0, output: 0, cacheWrite: 0 },
      compactionCount: record.compactionCount ?? 0,
    };
  }

  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();
      this.checkpoint(record);
      return true;
    }

    if (record.status !== "running") return false;
    record.abortController?.abort();
    record.status = "stopped";
    record.completedAt = Date.now();
    this.flushOutput(record);
    this.checkpoint(record);
    this.finishBackground(id);
    this.drainQueue();
    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;
      // A durable transcript is the source of truth for history. Release the
      // live session after the TTL, but retain a lightweight record so opening
      // history again in this session does not silently lose its identity or
      // locator. Records without durable storage remain eligible for eviction.
      if (record.transcriptPath) {
        try { record.session?.dispose?.(); } catch { /* ignore cleanup failures */ }
        record.session = undefined;
        try { record.outputCleanup?.(); } catch { /* ignore cleanup failures */ }
        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);
    }
  }

  /**
   * 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): 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);
    }
  }

  /** 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();
        this.checkpoint(record);
        count++;
      }
    }
    this.queue = [];
    // 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++;
      }
    }
    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();
    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
    // exit with in-flight agents would otherwise leave stale registrations there.
    for (const repo of this.worktreeRepos) {
      try { pruneWorktrees(repo); } catch { /* ignore */ }
    }
  }
}