@cr1ms0n/pi-subagent 0.8.1

package/src/registry.ts

@@ -0,0 +1,632 @@
+import { Buffer } from "node:buffer";
+import { randomUUID } from "node:crypto";
+import type { Message } from "@earendil-works/pi-ai";
+import type { SubagentConfig } from "./config.js";
+import type { PersistenceAdapter, PersistedResult } from "./persistence.js";
+import { PersistenceLayer } from "./persistence.js";
+import type { ProcessLockManager } from "./process-lock.js";
+import type { RunMode, RunSnapshot, RunState, TaskResult, TaskSpec } from "./types.js";
+import { emptyUsage } from "./types.js";
+
+export interface LiveRun {
+  id: string;
+  sessionKey: string;
+  mode: RunMode;
+  state: RunState;
+  startedAt: number;
+  endedAt?: number;
+  taskPreviews: string[];
+  taskSpecs: TaskSpec[];
+  results: TaskResult[];
+  summary?: string;
+  delivered: boolean;
+  promise: Promise<unknown>;
+  controller: AbortController;
+  childSessionIds: Set<string>;
+  lastProgressCheckpoint: number;
+}
+
+export interface RunLookupResult {
+  status: "found" | "not-found" | "ambiguous";
+  run?: LiveRun | RunSnapshot;
+  matches?: string[];
+}
+
+export interface SessionRuntime {
+  sessionKey: string;
+  runs: Map<string, LiveRun>;
+  snapshots: Map<string, RunSnapshot>;
+  activeResumes: Map<string, string>;
+  shuttingDown: boolean;
+}
+
+export type RegistryEvent =
+  | { type: "changed"; sessionKey: string; runId: string }
+  | { type: "terminal"; sessionKey: string; runId: string; state: RunState };
+
+const terminalStates = new Set<RunState>(["completed", "partial", "failed", "cancelled", "lost", "timeout"]);
+
+/** Trailing coalesce window for high-frequency "changed" events. */
+const EMIT_COALESCE_MS = 100;
+
+function finalText(messages: Message[], fallback?: string): string | undefined {
+  for (let i = messages.length - 1; i >= 0; i--) {
+    const message = messages[i];
+    if (message?.role !== "assistant" || !Array.isArray(message.content)) continue;
+    const text = message.content
+      .filter((part: any) => part?.type === "text" && typeof part.text === "string")
+      .map((part: any) => part.text)
+      .join("");
+    if (text) return text;
+  }
+  return fallback;
+}
+
+function utf8Prefix(value: string | undefined, maxBytes: number): string | undefined {
+  if (!value) return undefined;
+  const buffer = Buffer.from(value, "utf8");
+  if (buffer.length <= maxBytes) return value;
+  let end = maxBytes;
+  while (end > 0 && (buffer[end] & 0xc0) === 0x80) end--;
+  return buffer.subarray(0, end).toString("utf8");
+}
+
+/**
+ * Full result projection: capped transcripts included. Used for terminal
+ * persistence and UI snapshots — the single converter for both paths.
+ */
+export function toPersistedResult(result: TaskResult): PersistedResult {
+  return {
+    label: result.label,
+    task: result.task.slice(0, 1_000),
+    state: result.state,
+    exitCode: result.exitCode,
+    stopReason: result.stopReason,
+    timeoutPhase: result.timeoutPhase,
+    errorMessage: result.errorMessage,
+    usage: result.usage,
+    model: result.model,
+    thinking: result.thinking,
+    profile: result.profile,
+    backend: result.backend,
+    canWrite: result.canWrite,
+    outputFile: result.outputFile,
+    outputMode: result.outputMode,
+    sessionId: result.sessionId,
+    process: result.process,
+    finalOutput: utf8Prefix(finalText(result.messages, result.liveText), 16_384),
+    transcript: utf8Prefix(result.transcript, 32_768),
+    worktree: result.worktree,
+    wrappedUp: result.wrappedUp,
+    stalledSince: result.stalledSince,
+    attempts: result.attempts,
+    attemptedModels: result.attemptedModels,
+    structuredOutput: result.structuredOutput,
+    structuredError: result.structuredError,
+  };
+}
+
+/**
+ * Memoized per-result projection for live snapshots. High-frequency emitters
+ * (footer refresh, streamed tool updates) re-snapshot the whole run on every
+ * event; only the task that actually changed should pay the projection cost
+ * (message scan + capped-string allocation).
+ */
+const projectionCache = new WeakMap<TaskResult, { fingerprint: string; projected: PersistedResult }>();
+
+function resultFingerprint(result: TaskResult): string {
+  return [
+    result.state,
+    result.usage.turns,
+    result.usage.cost,
+    result.sessionId ?? "",
+    result.messages.length,
+    result.liveText?.length ?? 0,
+    result.transcript?.length ?? 0,
+    result.errorMessage?.length ?? 0,
+    result.stalledSince ?? 0,
+    result.attempts ?? 0,
+    result.worktree ? 1 : 0,
+    result.structuredOutput !== undefined ? 1 : 0,
+    result.structuredError?.length ?? 0,
+  ].join("|");
+}
+
+function toPersistedResultCached(result: TaskResult): PersistedResult {
+  const fingerprint = resultFingerprint(result);
+  const cached = projectionCache.get(result);
+  if (cached && cached.fingerprint === fingerprint) return cached.projected;
+  const projected = toPersistedResult(result);
+  projectionCache.set(result, { fingerprint, projected });
+  return projected;
+}
+
+/**
+ * Lightweight result projection for checkpoint events: state + usage +
+ * pointers only. Keeps checkpoint entries small so the parent session file
+ * does not bloat during long runs. Transcripts are persisted once, at terminal.
+ */
+export function toCheckpointResult(result: TaskResult): PersistedResult {
+  return {
+    label: result.label,
+    task: result.task.slice(0, 200),
+    state: result.state,
+    exitCode: result.exitCode,
+    stopReason: result.stopReason,
+    timeoutPhase: result.timeoutPhase,
+    errorMessage: utf8Prefix(result.errorMessage, 1_000),
+    usage: result.usage,
+    model: result.model,
+    thinking: result.thinking,
+    profile: result.profile,
+    backend: result.backend,
+    canWrite: result.canWrite,
+    outputFile: result.outputFile,
+    outputMode: result.outputMode,
+    sessionId: result.sessionId,
+    process: result.process,
+    worktree: result.worktree,
+    wrappedUp: result.wrappedUp,
+    stalledSince: result.stalledSince,
+    attempts: result.attempts,
+    attemptedModels: result.attemptedModels,
+  };
+}
+
+/**
+ * One shared LiveRun → RunSnapshot projection with capped transcripts.
+ * Unchanged task results reuse their cached projection (see toPersistedResultCached).
+ */
+export function snapshotFromLiveRun(run: LiveRun): RunSnapshot {
+  return {
+    schemaVersion: 1,
+    id: run.id,
+    sessionKey: run.sessionKey,
+    mode: run.mode,
+    state: run.state,
+    startedAt: run.startedAt,
+    endedAt: run.endedAt,
+    taskPreviews: run.taskPreviews,
+    summary: run.summary,
+    delivered: run.delivered,
+    results: run.results.map(toPersistedResultCached),
+  };
+}
+
+/** Session-owned live state plus immutable, bounded terminal snapshots. */
+export class SessionScopedRunRegistry {
+  private readonly runtimes = new Map<string, SessionRuntime>();
+  private readonly persistence: PersistenceLayer;
+  private readonly listeners = new Set<(event: RegistryEvent) => void>();
+  private readonly pendingEmits = new Map<string, NodeJS.Timeout>();
+  private readonly locks?: ProcessLockManager;
+
+  constructor(
+    private readonly config: SubagentConfig,
+    persistenceAdapter: PersistenceAdapter,
+    locks?: ProcessLockManager,
+  ) {
+    this.persistence = new PersistenceLayer(persistenceAdapter, config);
+    this.locks = locks;
+  }
+
+  allocateRunId(): string {
+    return randomUUID();
+  }
+
+  subscribe(listener: (event: RegistryEvent) => void): () => void {
+    this.listeners.add(listener);
+    return () => this.listeners.delete(listener);
+  }
+
+  private emit(event: RegistryEvent): void {
+    for (const listener of this.listeners) listener(event);
+  }
+
+  /**
+   * Coalesce per-run "changed" bursts (live-text ticks can arrive per stdout
+   * chunk) into at most one listener notification per window. Terminal and
+   * structural events always flush immediately.
+   */
+  private emitChanged(sessionKey: string, runId: string, immediate = false): void {
+    const key = `${sessionKey}\u0000${runId}`;
+    if (immediate) {
+      const pending = this.pendingEmits.get(key);
+      if (pending) {
+        clearTimeout(pending);
+        this.pendingEmits.delete(key);
+      }
+      this.emit({ type: "changed", sessionKey, runId });
+      return;
+    }
+    if (this.pendingEmits.has(key)) return;
+    const timer = setTimeout(() => {
+      this.pendingEmits.delete(key);
+      this.emit({ type: "changed", sessionKey, runId });
+    }, EMIT_COALESCE_MS);
+    timer.unref?.();
+    this.pendingEmits.set(key, timer);
+  }
+
+  private clearPendingEmits(sessionKey?: string): void {
+    for (const [key, timer] of this.pendingEmits) {
+      if (sessionKey && !key.startsWith(`${sessionKey}\u0000`)) continue;
+      clearTimeout(timer);
+      this.pendingEmits.delete(key);
+    }
+  }
+
+  private getOrCreateRuntime(sessionKey: string): SessionRuntime {
+    let runtime = this.runtimes.get(sessionKey);
+    if (!runtime) {
+      runtime = {
+        sessionKey,
+        runs: new Map(),
+        snapshots: this.persistence.rebuild(sessionKey),
+        activeResumes: new Map(),
+        shuttingDown: false,
+      };
+      this.runtimes.set(sessionKey, runtime);
+    }
+    return runtime;
+  }
+
+  getSessionRuntime(sessionKey: string): SessionRuntime | undefined {
+    return this.runtimes.get(sessionKey);
+  }
+
+  getLiveRuns(sessionKey: string): LiveRun[] {
+    return [...this.getOrCreateRuntime(sessionKey).runs.values()];
+  }
+
+  getSnapshots(sessionKey: string): RunSnapshot[] {
+    return [...this.getOrCreateRuntime(sessionKey).snapshots.values()];
+  }
+
+  /** Live cwds of worktree-isolated tasks; used to protect them from sweeps. */
+  getLiveWorktreeCwds(sessionKey: string): Set<string> {
+    const cwds = new Set<string>();
+    for (const run of this.getOrCreateRuntime(sessionKey).runs.values()) {
+      for (const result of run.results) if (result.worktree?.cwd) cwds.add(result.worktree.cwd);
+      for (const spec of run.taskSpecs) if (spec.isolation === "worktree" && spec.cwd) cwds.add(spec.cwd);
+    }
+    return cwds;
+  }
+
+  /** Rebuild terminal history after active-branch navigation. Live runs stay session-owned. */
+  refreshSnapshots(sessionKey: string): void {
+    const runtime = this.getOrCreateRuntime(sessionKey);
+    runtime.snapshots = this.persistence.rebuild(sessionKey);
+    this.capSnapshots(runtime);
+    this.emitChanged(sessionKey, "branch", true);
+  }
+
+  lookup(idOrPrefix: string, sessionKey: string): RunLookupResult {
+    if (!idOrPrefix) return { status: "not-found" };
+    const runtime = this.getOrCreateRuntime(sessionKey);
+    const exact = runtime.runs.get(idOrPrefix) ?? runtime.snapshots.get(idOrPrefix);
+    if (exact) return { status: "found", run: exact };
+
+    const matches = new Map<string, LiveRun | RunSnapshot>();
+    for (const [id, run] of runtime.runs) if (id.startsWith(idOrPrefix)) matches.set(id, run);
+    for (const [id, run] of runtime.snapshots) if (id.startsWith(idOrPrefix)) matches.set(id, run);
+    if (matches.size === 0) return { status: "not-found" };
+    if (matches.size > 1) return { status: "ambiguous", matches: [...matches.keys()].sort() };
+    return { status: "found", run: [...matches.values()][0] };
+  }
+
+  start(
+    sessionKey: string,
+    mode: RunMode,
+    specs: TaskSpec[],
+    controller: AbortController,
+    promise: Promise<unknown>,
+    labels: string[] = [],
+    id = this.allocateRunId(),
+  ): string {
+    const runtime = this.getOrCreateRuntime(sessionKey);
+    if (runtime.shuttingDown) throw new Error("Cannot start a subagent while the parent session is shutting down");
+    if (runtime.runs.has(id) || runtime.snapshots.has(id)) throw new Error(`Duplicate run id ${id}`);
+
+    const startedAt = Date.now();
+    const taskPreviews = specs.map((spec, i) => `${labels[i] || `task-${i + 1}`}: ${spec.task.slice(0, 120)}`);
+    const results: TaskResult[] = specs.map((spec, i) => ({
+      label: labels[i] || `task-${i + 1}`,
+      task: spec.task,
+      state: "queued",
+      exitCode: null,
+      messages: [],
+      stderr: "",
+      usage: emptyUsage(),
+      outputFile: spec.output,
+      outputMode: spec.outputMode,
+      thinking: spec.thinking,
+      profile: spec.profile,
+      backend: spec.backend,
+      canWrite: spec.canWrite,
+      protocol: {
+        headerSeen: false,
+        assistantEndSeen: false,
+        agentEndSeen: false,
+        agentSettledSeen: false,
+        validEvents: 0,
+        parseErrors: 0,
+      },
+    }));
+
+    runtime.runs.set(id, {
+      id,
+      sessionKey,
+      mode,
+      state: "queued",
+      startedAt,
+      taskPreviews,
+      taskSpecs: [...specs],
+      results,
+      delivered: false,
+      promise,
+      controller,
+      childSessionIds: new Set(),
+      lastProgressCheckpoint: 0,
+    });
+    this.persistence.persist(id, sessionKey, "start", {
+      mode,
+      state: "queued",
+      startedAt,
+      taskPreviews,
+      results: results.map(toCheckpointResult),
+    });
+    this.emitChanged(sessionKey, id, true);
+    return id;
+  }
+
+  checkpoint(
+    id: string,
+    sessionKey: string,
+    updates: {
+      childSessionId?: string;
+      progress?: string;
+      turn?: number;
+      resultIndex?: number;
+      resultUpdate?: Partial<TaskResult>;
+      state?: RunState;
+    },
+  ): boolean {
+    const runtime = this.runtimes.get(sessionKey);
+    const run = runtime?.runs.get(id);
+    if (!runtime || !run || run.sessionKey !== sessionKey || runtime.shuttingDown) return false;
+
+    const index = updates.resultIndex ?? 0;
+    const result = run.results[index];
+    const previousTurns = result?.usage.turns ?? 0;
+    const previousCost = result?.usage.cost ?? 0;
+    const previousRunState = run.state;
+    if (result && updates.resultUpdate) Object.assign(result, updates.resultUpdate);
+    const usageAdvanced = !!result && (result.usage.turns > previousTurns || result.usage.cost > previousCost);
+    if (updates.state) run.state = updates.state;
+    else if (run.state === "queued") run.state = "running";
+    const stateChanged = run.state !== previousRunState;
+
+    const childSessionId = updates.childSessionId ?? updates.resultUpdate?.sessionId;
+    let newChildSession = false;
+    if (childSessionId) {
+      newChildSession = !run.childSessionIds.has(childSessionId);
+      run.childSessionIds.add(childSessionId);
+      if (result) result.sessionId = childSessionId;
+      // Crash recovery requirement: session ids are never throttled.
+      if (newChildSession) {
+        this.persistence.persist(id, sessionKey, "checkpoint", {
+          state: run.state,
+          resultIndex: index,
+          childSessionId,
+          results: run.results.map(toCheckpointResult),
+        });
+      }
+    }
+
+    // Persist lightweight checkpoints (state + usage + pointers, never
+    // transcripts) only when billed usage advanced, or on a throttled progress
+    // beat. Full transcripts are written exactly once, in the terminal event.
+    const now = Date.now();
+    if (usageAdvanced || ((updates.progress || updates.turn !== undefined) && now - run.lastProgressCheckpoint >= 500)) {
+      run.lastProgressCheckpoint = now;
+      this.persistence.persist(id, sessionKey, "checkpoint", {
+        state: run.state,
+        resultIndex: index,
+        progress: utf8Prefix(updates.progress, 200),
+        turn: updates.turn,
+        results: run.results.map(toCheckpointResult),
+      });
+    }
+    // Structural changes flush immediately; live-text ticks coalesce.
+    this.emitChanged(sessionKey, id, usageAdvanced || stateChanged || newChildSession);
+    return true;
+  }
+
+  complete(
+    id: string,
+    sessionKey: string,
+    finalState: RunState,
+    summary?: string,
+    finalResults?: TaskResult[],
+  ): boolean {
+    const runtime = this.runtimes.get(sessionKey);
+    const run = runtime?.runs.get(id);
+    if (!runtime || !run || run.sessionKey !== sessionKey) return false;
+
+    const endedAt = Date.now();
+    const results = finalResults ?? run.results;
+    const snapshot: RunSnapshot = {
+      schemaVersion: 1,
+      id,
+      sessionKey,
+      mode: run.mode,
+      state: terminalStates.has(finalState) ? finalState : "failed",
+      startedAt: run.startedAt,
+      endedAt,
+      taskPreviews: run.taskPreviews,
+      summary,
+      delivered: run.delivered,
+      results: results.map(toPersistedResult),
+    };
+    runtime.snapshots.set(id, snapshot);
+    runtime.runs.delete(id);
+    this.releaseLocksForRun(runtime, id);
+    this.persistence.persist(id, sessionKey, "terminal", {
+      mode: snapshot.mode,
+      state: snapshot.state,
+      startedAt: snapshot.startedAt,
+      endedAt,
+      taskPreviews: snapshot.taskPreviews,
+      summary,
+      delivered: snapshot.delivered,
+      results: snapshot.results,
+    });
+    this.capSnapshots(runtime);
+    this.clearPendingEmitsForRun(sessionKey, id);
+    this.emit({ type: "terminal", sessionKey, runId: id, state: snapshot.state });
+    return true;
+  }
+
+  private clearPendingEmitsForRun(sessionKey: string, runId: string): void {
+    const key = `${sessionKey}\u0000${runId}`;
+    const pending = this.pendingEmits.get(key);
+    if (pending) {
+      clearTimeout(pending);
+      this.pendingEmits.delete(key);
+    }
+  }
+
+  /** Returns false when this result was already delivered. */
+  markDelivered(id: string, sessionKey: string): boolean {
+    const runtime = this.getOrCreateRuntime(sessionKey);
+    const run = runtime.runs.get(id) ?? runtime.snapshots.get(id);
+    if (!run || run.delivered) return false;
+    run.delivered = true;
+    this.persistence.markDelivered(id, sessionKey);
+    this.emitChanged(sessionKey, id, true);
+    return true;
+  }
+
+  markDismissed(id: string, sessionKey: string): boolean {
+    return this.markDelivered(id, sessionKey);
+  }
+
+  acquireResumeLocks(childSessionIds: string[], runId: string, sessionKey: string): {
+    ok: boolean;
+    conflict?: { sessionId: string; runId: string };
+  } {
+    const runtime = this.getOrCreateRuntime(sessionKey);
+    const unique = [...new Set(childSessionIds.filter(Boolean))];
+    // Block resume of runs whose ownership is not yet proven dead.
+    for (const snapshot of runtime.snapshots.values()) {
+      if (!snapshot.resumeBlocked) continue;
+      for (const result of snapshot.results) {
+        if (result.sessionId && unique.includes(result.sessionId)) {
+          return {
+            ok: false,
+            conflict: {
+              sessionId: result.sessionId,
+              runId: snapshot.id,
+            },
+          };
+        }
+      }
+    }
+    // In-memory lock first (fast path within one process).
+    for (const sessionId of unique) {
+      const holder = runtime.activeResumes.get(sessionId);
+      if (holder && holder !== runId) return { ok: false, conflict: { sessionId, runId: holder } };
+    }
+    // Durable, machine-wide lock — survival across parent crashes and cross-process contention.
+    const durableHeld: string[] = [];
+    if (this.locks) {
+      for (const sessionId of unique) {
+        const acquired = this.locks.acquireSessionLock(sessionId, {
+          ownerId: `${sessionKey}:${runId}`,
+          runId,
+          parentSessionKey: sessionKey,
+        });
+        if (!acquired.ok) {
+          for (const held of durableHeld) this.locks.releaseSessionLock(held, runId);
+          return {
+            ok: false,
+            conflict: { sessionId: acquired.conflict.childSessionId, runId: acquired.conflict.runId },
+          };
+        }
+        durableHeld.push(sessionId);
+      }
+    }
+    for (const sessionId of unique) runtime.activeResumes.set(sessionId, runId);
+    return { ok: true };
+  }
+
+  acquireResumeLock(childSessionId: string, runId: string, sessionKey: string, isFork = false): boolean {
+    if (isFork) return true;
+    return this.acquireResumeLocks([childSessionId], runId, sessionKey).ok;
+  }
+
+  releaseResumeLock(childSessionId: string, sessionKey: string, runId?: string): void {
+    const runtime = this.runtimes.get(sessionKey);
+    if (!runtime) return;
+    if (!runId || runtime.activeResumes.get(childSessionId) === runId) runtime.activeResumes.delete(childSessionId);
+    if (this.locks) this.locks.releaseSessionLock(childSessionId, runId);
+  }
+
+  private releaseLocksForRun(runtime: SessionRuntime, runId: string): void {
+    for (const [sessionId, holder] of [...runtime.activeResumes]) {
+      if (holder === runId) {
+        runtime.activeResumes.delete(sessionId);
+        this.locks?.releaseSessionLock(sessionId, runId);
+      }
+    }
+  }
+
+  /** Clear the resume-blocked flag after orphan reconcile has proven the child dead. */
+  clearResumeBlock(id: string, sessionKey: string): void {
+    const runtime = this.getOrCreateRuntime(sessionKey);
+    const snapshot = runtime.snapshots.get(id);
+    if (!snapshot || !snapshot.resumeBlocked) return;
+    snapshot.resumeBlocked = false;
+    this.persistence.persist(id, sessionKey, "checkpoint", { resumeBlocked: false, state: snapshot.state });
+    this.emitChanged(sessionKey, id, true);
+  }
+
+  async shutdown(sessionKey: string, graceMs = 8_000): Promise<void> {
+    const runtime = this.runtimes.get(sessionKey);
+    if (!runtime) return;
+    runtime.shuttingDown = true;
+    const live = [...runtime.runs.values()];
+    for (const run of live) run.controller.abort();
+
+    let timer: NodeJS.Timeout | undefined;
+    await Promise.race([
+      Promise.allSettled(live.map((run) => run.promise)),
+      new Promise<void>((resolve) => {
+        timer = setTimeout(resolve, graceMs);
+        timer.unref?.();
+      }),
+    ]);
+    if (timer) clearTimeout(timer);
+
+    // Any orchestration promise that did not call complete is snapshotted as cancelled.
+    for (const run of [...runtime.runs.values()]) {
+      this.complete(run.id, sessionKey, "cancelled", "Cancelled when the parent session shut down", run.results);
+    }
+    runtime.activeResumes.clear();
+    runtime.shuttingDown = false;
+    this.clearPendingEmits(sessionKey);
+  }
+
+  planSessionRetention(referencedSessionIds = new Set<string>()): { keep: string[]; candidates: string[] } {
+    return this.persistence.planRetention(referencedSessionIds);
+  }
+
+  private capSnapshots(runtime: SessionRuntime): void {
+    while (runtime.snapshots.size > this.config.maxCompletedInMemory) {
+      const oldest = [...runtime.snapshots.values()].sort((a, b) => a.startedAt - b.startedAt)[0];
+      if (!oldest) break;
+      runtime.snapshots.delete(oldest.id);
+    }
+  }
+}