pi-agent-fleet 0.5.0 → 0.7.0

package/src/edits.ts

@@ -8,14 +8,14 @@ import type { NodeStatus, ThinkingLevelName } from "./types.js";
 import { THINKING_LEVELS } from "./types.js";
 
 export type NodeEditKey = "model" | "effort" | "task";
-export type ConfigEditKey = "max_concurrent" | "warn_cost_usd" | "model" | "effort";
+export type ConfigEditKey = "max_concurrent" | "warn_cost_usd" | "max_cost_usd" | "worker_extensions" | "model" | "effort";
 
 export interface EditResult {
   ok: boolean;
   message: string;
 }
 
-const EDITABLE_NODE_STATUSES: ReadonlySet<NodeStatus> = new Set(["pending", "ready", "failed", "contract_failed", "killed"]);
+const EDITABLE_NODE_STATUSES: ReadonlySet<NodeStatus> = new Set(["pending", "ready", "failed", "contract_failed", "killed", "blocked"]);
 
 export async function editNode(
   fleet: ActiveFleet,
@@ -28,8 +28,8 @@ export async function editNode(
   const node = fleet.state.nodes[nodeId];
   if (!worker || !node) return { ok: false, message: `unknown node "${nodeId}"` };
   if (!EDITABLE_NODE_STATUSES.has(node.status)) {
-      return { ok: false, message: `node "${nodeId}" is ${node.status}; only pending, failed, contract_failed, or killed nodes can be edited` };
-    }
+    return { ok: false, message: `node "${nodeId}" is ${node.status}; only pending, blocked, failed, contract_failed, or killed nodes can be edited` };
+  }
   switch (key) {
     case "model": {
       const r = resolveModelReference(registry, value);
@@ -85,6 +85,25 @@ export async function editConfig(
       fleet.spec.config.warn_cost_usd = n;
       break;
     }
+    case "max_cost_usd": {
+      const n = Number(value);
+      if (!Number.isFinite(n) || n < 0) return { ok: false, message: "max_cost_usd must be a number >= 0" };
+      fleet.spec.config.max_cost_usd = n;
+      break;
+    }
+    case "worker_extensions": {
+      // accept JSON array or comma-separated string
+      let list: string[];
+      try {
+        const parsed: unknown = JSON.parse(value);
+        list = Array.isArray(parsed) ? parsed.map(String) : [value];
+      } catch {
+        list = value.split(",").map((s) => s.trim()).filter(Boolean);
+      }
+      if (list.some((s) => s.length === 0)) return { ok: false, message: "worker_extensions must be non-empty strings" };
+      fleet.spec.config.worker_extensions = list;
+      break;
+    }
     case "model": {
       const r = resolveModelReference(registry, value);
       if (!r.ok) return { ok: false, message: r.error };
@@ -99,7 +118,7 @@ export async function editConfig(
       break;
     }
     default:
-      return { ok: false, message: `unknown config key "${String(key)}" (keys: max_concurrent, warn_cost_usd, model, effort)` };
+      return { ok: false, message: `unknown config key "${String(key)}" (keys: max_concurrent, warn_cost_usd, max_cost_usd, worker_extensions, model, effort)` };
   }
   await persistFleetJson(fleet);
   return { ok: true, message: `config.${key} updated` };

package/src/fleet-recovery.ts

@@ -1,7 +1,7 @@
 import { readdir, readFile, stat } from "node:fs/promises";
 import { join } from "node:path";
 import type { ActiveFleet } from "./controller.js";
-import { readState } from "./state.js";
+import { initFleetState, readState } from "./state.js";
 import type { FleetSpec, FleetState } from "./types.js";
 
 export interface FleetRootInfo {
@@ -13,7 +13,16 @@ export interface FleetRootInfo {
 
 export async function readDiskFleet(fleetRoot: string): Promise<ActiveFleet> {
   const spec = JSON.parse(await readFile(join(fleetRoot, "fleet.json"), "utf-8")) as FleetSpec;
-  const state = await readState(fleetRoot);
+  // Bare fleet.json-only root (e.g. hand-copied spec, never planned/launched):
+  // synthesize a fresh pending state so the canvas can still render the DAG.
+  let state: FleetState;
+  try {
+    state = await readState(fleetRoot);
+  } catch (e: unknown) {
+    const err = e as { code?: string };
+    if (err.code !== "ENOENT") throw e;
+    state = initFleetState(spec);
+  }
   return {
     spec,
     fleetRoot,
@@ -23,6 +32,7 @@ export async function readDiskFleet(fleetRoot: string): Promise<ActiveFleet> {
     running: false,
     sessions: new Map(),
     killedNodes: new Set(),
+    relaunchRequests: new Set(),
   };
 }
 

package/src/prompts.ts

@@ -17,6 +17,17 @@ export function buildWorkerPrompt(opts: {
   const dependents = getDependents(spec, workerId);
   const out: string[] = [];
 
+  out.push(
+    "## Autonomy contract (read first)",
+    "",
+    "You are an unattended fleet worker — there is no human watching this session and nobody will answer questions.",
+    "- Everything in this prompt is PRE-APPROVED. Do not ask for approval; do the work.",
+    "- Do NOT invoke any skill or workflow with a human approval gate (e.g. brainstorming hard-gates). Skip gated steps and execute the task directly.",
+    "- Do not end your turn with a question, a plan awaiting approval, or an 'Approve?' prompt. End your turn only when every REQUIRED output below exists on disk.",
+    "- Ambiguity is yours to resolve: decide, record the decision in your output, and continue.",
+    "",
+  );
+
   const fleetTs = basename(fleetRoot);
 
   out.push(`# Fleet worker: ${workerId}`, "", `Type: ${worker.type}`, "", `## Task`, "", worker.task, "");

package/src/runner.ts

@@ -1,4 +1,4 @@
-import { createAgentSession, SessionManager, type CreateAgentSessionOptions } from "@earendil-works/pi-coding-agent";
+import { createAgentSession, SessionManager, DefaultResourceLoader, type CreateAgentSessionOptions } from "@earendil-works/pi-coding-agent";
 import type { Api, Model } from "@earendil-works/pi-ai";
 import type { ThinkingLevelName, WorkerSpec } from "./types.js";
 import { WORKER_TYPE_TOOLS } from "./types.js";
@@ -25,16 +25,52 @@ export interface SessionOpts {
   tools: string[];
   model?: string;
   thinkingLevel?: ThinkingLevelName;
+  extensionAllowlist?: string[];
 }
 
 export type SessionFactory = (opts: SessionOpts) => Promise<AgentSessionLike>;
 
+// Fleet workers must not load user extensions/skills: reading global skill files
+// burns tokens mid-task, and extensions like pi-web-access run background async work
+// that can crash the pi process after the session is disposed.
+// Exception: some model providers are registered BY extensions (e.g. opencode-pi
+// registers the opencode-cli provider). worker_extensions allowlists exactly those.
+export function filterExtensionsByAllowlist<T extends { path: string }>(
+  extensions: T[],
+  allowlist: string[] | undefined,
+): T[] {
+  if (!allowlist || allowlist.length === 0) return extensions;
+  return extensions.filter((e) => allowlist.some((a) => e.path.includes(a)));
+}
+
+async function createLeanResourceLoader(cwd: string, extensionAllowlist?: string[]): Promise<DefaultResourceLoader> {
+  const allow = extensionAllowlist && extensionAllowlist.length > 0;
+  const loader = new DefaultResourceLoader({
+    cwd,
+    agentDir: cwd,
+    noExtensions: !allow,
+    extensionsOverride: allow
+      ? (base) => ({ ...base, extensions: filterExtensionsByAllowlist(base.extensions, extensionAllowlist) })
+      : undefined,
+    noSkills: true,
+    noPromptTemplates: true,
+    noThemes: true,
+    noContextFiles: true,
+    systemPromptOverride: () => undefined,
+    appendSystemPromptOverride: () => [],
+  });
+  await loader.reload();
+  return loader;
+}
+
 export const defaultSessionFactory: SessionFactory = async (opts) => {
+  const resourceLoader = await createLeanResourceLoader(opts.cwd, opts.extensionAllowlist);
   const { session } = await createAgentSession({
     cwd: opts.cwd,
     tools: opts.tools,
     sessionManager: SessionManager.create(opts.cwd, opts.sessionDir),
     thinkingLevel: opts.thinkingLevel as ThinkingLevelOption,
+    resourceLoader,
   });
   return session as unknown as AgentSessionLike;
 };
@@ -49,6 +85,7 @@ export interface RunWorkerOpts {
   onSession?: (session: AgentSessionLike) => void;
   sessionFactory?: SessionFactory;
   thinkingLevel?: ThinkingLevelName;
+  extensionAllowlist?: string[];
 }
 
 export interface RunWorkerResult {
@@ -69,6 +106,7 @@ export async function runWorker(opts: RunWorkerOpts): Promise<RunWorkerResult> {
       tools: WORKER_TYPE_TOOLS[opts.worker.type],
       model: opts.worker.model,
       thinkingLevel: opts.thinkingLevel,
+      extensionAllowlist: opts.extensionAllowlist,
     });
   } catch (err) {
     const message = err instanceof Error ? err.message : String(err);
@@ -79,13 +117,29 @@ export async function runWorker(opts: RunWorkerOpts): Promise<RunWorkerResult> {
   let turns = 0;
   let tokens = 0;
   let cost = 0;
+  // pi resolves session.prompt() normally even when the model run ends in an error
+  // (e.g. provider usage limit) — the failure is only visible on the final assistant
+  // message's stopReason/errorMessage. Track the last assistant message so a resolved
+  // prompt that actually errored is reported as ok:false instead of running the output
+  // contract against files that were never written.
+  let lastStopReason: string | undefined;
+  let lastErrorMessage: string | undefined;
   const unsub = session.subscribe((e) => {
     if (e.type === "turn_end") {
       turns++;
       opts.onEvent({ type: "turn", nodeId: opts.nodeId, turns });
     }
     if (e.type === "message_end") {
-      const msg = e.message as { role?: string; usage?: { totalTokens?: number; cost?: { total?: number } } } | undefined;
+      const msg = e.message as {
+        role?: string;
+        stopReason?: string;
+        errorMessage?: string;
+        usage?: { totalTokens?: number; cost?: { total?: number } };
+      } | undefined;
+      if (msg?.role === "assistant") {
+        lastStopReason = msg.stopReason;
+        lastErrorMessage = msg.errorMessage;
+      }
       if (msg?.role === "assistant" && msg.usage?.totalTokens) {
         tokens += msg.usage.totalTokens;
         opts.onEvent({ type: "tokens", nodeId: opts.nodeId, tokens });
@@ -98,6 +152,11 @@ export async function runWorker(opts: RunWorkerOpts): Promise<RunWorkerResult> {
   });
   try {
     await session.prompt(opts.prompt);
+    if (lastStopReason === "error") {
+      const message = lastErrorMessage ?? "worker model run ended with an error";
+      opts.onEvent({ type: "error", nodeId: opts.nodeId, message });
+      return { ok: false, turns, tokens, cost, error: message };
+    }
     opts.onEvent({ type: "done", nodeId: opts.nodeId });
     return { ok: true, turns, tokens, cost };
   } catch (err) {
@@ -112,12 +171,14 @@ export async function runWorker(opts: RunWorkerOpts): Promise<RunWorkerResult> {
 
 export function sessionFactoryForModel(model: Model<Api>): SessionFactory {
   return async (opts) => {
+    const resourceLoader = await createLeanResourceLoader(opts.cwd, opts.extensionAllowlist);
     const { session } = await createAgentSession({
       cwd: opts.cwd,
       tools: opts.tools,
       sessionManager: SessionManager.create(opts.cwd, opts.sessionDir),
       model,
       thinkingLevel: opts.thinkingLevel as ThinkingLevelOption,
+      resourceLoader,
     });
     return session as unknown as AgentSessionLike;
   };

package/src/scheduler.ts

@@ -1,7 +1,7 @@
 import { mkdir, rm } from "node:fs/promises";
 import { join } from "node:path";
 import { contractFailureNote, verifyOutputs } from "./contracts.js";
-import { archiveIteration, initFleetState, patchNode, resetForIteration, snapshotIteration, writeState } from "./state.js";
+import { archiveIteration, initFleetState, patchNode, relaunchResetIds, resetForIteration, snapshotIteration, writeState } from "./state.js";
 import { TERMINAL_NODE_STATUSES } from "./types.js";
 import type { FleetSpec, FleetState, IterationSnapshot, NodeState, Verdict, WorkerSpec } from "./types.js";
 import { commitWorktree, createWorktree, prepareIntegratorWorktree } from "./worktree.js";
@@ -20,10 +20,20 @@ export interface RunFleetOpts {
   killSwitch?: { killed: boolean };
   pauseSwitch?: { paused: boolean };
   nodeKills?: ReadonlySet<string>;
+  relaunchRequests?: Set<string>;
   resumeFrom?: FleetState;
   continuePass?: boolean;
   onIterationEnd?: (snap: IterationSnapshot) => void;
   prepareIteration?: (n: number, state: FleetState) => Promise<void>;
+  retryDelayMs?: (attempt: number) => number;
+}
+
+const MAX_RETRY_ATTEMPTS = 3;
+const DEFAULT_RETRY_DELAY_MS = (attempt: number): number => 1000 * 2 ** attempt;
+const RETRYABLE_ERROR_RE = /usage limit|rate.?limit|429|overloaded|502|503|timed? ?out|timeout|econnreset|etimedout|econnrefused|fetch failed/i;
+
+export function isRetryableError(message: string): boolean {
+  return RETRYABLE_ERROR_RE.test(message);
 }
 
 const FAILED: ReadonlySet<string> = new Set(["failed", "contract_failed", "killed", "blocked"]);
@@ -51,6 +61,19 @@ export async function runFleet(opts: RunFleetOpts): Promise<FleetState> {
   let state: FleetState;
   if (opts.resumeFrom) {
     state = { ...opts.resumeFrom, paused: false, status: "running" };
+    // Crash recovery: nodes left "running" on disk by a dead process have no live
+    // session in this runFleet call — reset them to pending so they get dispatched.
+    // Without this, the scheduler can neither dispatch nor terminate them and spins
+    // forever (nothing runnable, not all nodes terminal).
+    for (const w of spec.workers) {
+      if (state.nodes[w.id]?.status === "running") {
+        state = patchNode(fleetRoot, state, w.id, {
+          status: "pending",
+          started_at: undefined,
+          status_note: "recovered stale running state from a previous crashed/killed process",
+        });
+      }
+    }
     await writeState(fleetRoot, state);
     if (allNodesTerminal(state, spec) && !opts.continuePass) {
       state = resetForIteration(state, spec);
@@ -70,6 +93,15 @@ export async function runFleet(opts: RunFleetOpts): Promise<FleetState> {
   };
 
   const running = new Set<Promise<void>>();
+  const circuitErrors = new Map<string, Set<string>>();
+  let circuitOpen = false;
+  function recordCircuitError(nodeId: string, error: string | undefined) {
+    if (!error) return;
+    const set = circuitErrors.get(error) ?? new Set<string>();
+    set.add(nodeId);
+    circuitErrors.set(error, set);
+    if (set.size >= 2) circuitOpen = true;
+  }
 
   const repoCwdFor = (nodeId: string): string =>
     typeof opts.repoCwd === "function" ? opts.repoCwd(nodeId) : opts.repoCwd;
@@ -111,6 +143,31 @@ export async function runFleet(opts: RunFleetOpts): Promise<FleetState> {
 
   const runPass = async (): Promise<void> => {
     while (true) {
+      // apply queued relaunch requests (lost-wakeup fix, issue #1 bug 2)
+      if (opts.relaunchRequests && opts.relaunchRequests.size > 0) {
+        for (const id of [...opts.relaunchRequests]) {
+          opts.relaunchRequests.delete(id);
+          const n = state.nodes[id];
+          if (!n || !FAILED.has(n.status)) continue;
+          for (const rid of relaunchResetIds(spec, state, id)) {
+            const rn = state.nodes[rid];
+            if (!rn) continue;
+            if (rid === id && !FAILED.has(rn.status)) continue;
+            if (rid !== id && rn.status !== "blocked") continue;
+            await patch(rid, {
+              status: "pending",
+              started_at: undefined,
+              ended_at: undefined,
+              turns: 0,
+              tokens: 0,
+              cost_usd_estimate: 0,
+              produced_outputs: [],
+              contract_result: undefined,
+              status_note: undefined,
+            });
+          }
+        }
+      }
       // auto-initialize workers inserted into the spec after the run started
       for (const w of spec.workers) {
         if (!state.nodes[w.id]) {
@@ -144,6 +201,32 @@ export async function runFleet(opts: RunFleetOpts): Promise<FleetState> {
         }
         break;
       }
+      // honor kill requests for not-yet-running nodes (incl. blocked, issue #1 bug 3)
+      for (const w of spec.workers) {
+        const n = state.nodes[w.id];
+        if (!n) continue;
+        if (!opts.nodeKills?.has(w.id)) continue;
+        if (n.status === "pending" || n.status === "ready" || n.status === "blocked") {
+          await patch(w.id, { status: "killed", ended_at: new Date().toISOString() });
+        }
+      }
+      // circuit breaker: two distinct nodes failed with the identical error
+      if (circuitOpen) {
+        const [tripError, tripSet] = [...circuitErrors.entries()].find(([_, ids]) => ids.size >= 2) ?? ["", new Set<string>()];
+        const tripCount = tripSet.size;
+        const note = `circuit breaker: ${tripCount} nodes failed with identical error: ${tripError.slice(0, 80)}`;
+        for (const w of spec.workers) {
+          const n = state.nodes[w.id];
+          if (!n || (n.status !== "pending" && n.status !== "ready")) continue;
+          await patch(w.id, { status: "blocked", ended_at: new Date().toISOString(), status_note: note });
+        }
+        if (running.size > 0) {
+          await Promise.race(running);
+          continue;
+        }
+        break;
+      }
+
       // dispatch ready
       const activeCount = running.size;
       let slots = spec.config.max_concurrent - activeCount;
@@ -214,17 +297,57 @@ export async function runFleet(opts: RunFleetOpts): Promise<FleetState> {
           }
         }
 
+        const dispatchMs = Date.now();
         await patch(w.id, { status: "running", started_at: new Date().toISOString() });
-        const p = opts.spawn(w.id).then(async (res) => {
+        const p = (async () => {
+          let totals = { turns: 0, tokens: 0, cost: 0 };
+          let lastRes: Awaited<ReturnType<typeof opts.spawn>> | undefined;
+          let finalError: string | undefined;
+
+          for (let attempt = 0; attempt < MAX_RETRY_ATTEMPTS; attempt++) {
+            if (opts.killSwitch?.killed) return;
+            if (opts.nodeKills?.has(w.id)) {
+              await patch(w.id, { status: "killed", ended_at: new Date().toISOString(), turns: totals.turns, tokens: totals.tokens, cost_usd_estimate: totals.cost });
+              return;
+            }
+            const res = await opts.spawn(w.id);
+            lastRes = res;
+            totals.turns += res.turns;
+            totals.tokens += res.tokens;
+            totals.cost += res.cost ?? 0;
+            if (opts.killSwitch?.killed) return;
+            if (opts.nodeKills?.has(w.id)) {
+              await patch(w.id, { status: "killed", ended_at: new Date().toISOString(), turns: totals.turns, tokens: totals.tokens, cost_usd_estimate: totals.cost });
+              return;
+            }
+            if (res.ok) {
+              finalError = undefined;
+              break;
+            }
+            finalError = res.error;
+            if (!res.error || !isRetryableError(res.error) || attempt === MAX_RETRY_ATTEMPTS - 1) {
+              break;
+            }
+            await patch(w.id, { status_note: `retry ${attempt + 1}/${MAX_RETRY_ATTEMPTS - 1} after: ${res.error}` });
+            const delayMs = opts.retryDelayMs ? opts.retryDelayMs(attempt) : DEFAULT_RETRY_DELAY_MS(attempt);
+            await new Promise((r) => setTimeout(r, delayMs));
+          }
+
+          const res = lastRes;
+          if (!res) return;
+
           if (opts.killSwitch?.killed) return;
           if (opts.nodeKills?.has(w.id)) {
-            await patch(w.id, { status: "killed", ended_at: new Date().toISOString(), turns: res.turns, tokens: res.tokens, cost_usd_estimate: res.cost ?? 0 });
+            await patch(w.id, { status: "killed", ended_at: new Date().toISOString(), turns: totals.turns, tokens: totals.tokens, cost_usd_estimate: totals.cost });
             return;
           }
+
           if (!res.ok) {
-            await patch(w.id, { status: "failed", ended_at: new Date().toISOString(), turns: res.turns, tokens: res.tokens, cost_usd_estimate: res.cost ?? 0 });
+            await patch(w.id, { status: "failed", ended_at: new Date().toISOString(), turns: totals.turns, tokens: totals.tokens, cost_usd_estimate: totals.cost });
+            recordCircuitError(w.id, finalError ?? res.error);
             return;
           }
+
           if (w.worktree) {
             try {
               await commitWorktree({
@@ -238,9 +361,9 @@ export async function runFleet(opts: RunFleetOpts): Promise<FleetState> {
               await patch(w.id, {
                 status: "failed",
                 ended_at: new Date().toISOString(),
-                turns: res.turns,
-                tokens: res.tokens,
-                cost_usd_estimate: res.cost ?? 0,
+                turns: totals.turns,
+                tokens: totals.tokens,
+                cost_usd_estimate: totals.cost,
                 status_note: `commit failed: ${msg}`,
               });
               return;
@@ -250,22 +373,25 @@ export async function runFleet(opts: RunFleetOpts): Promise<FleetState> {
             workerDir: `${fleetRoot}/workers/${w.id}`,
             repoCwd: repoCwdFor(w.id),
             outputs: w.outputs,
+            notBeforeMs: dispatchMs,
           });
+          const costUnknown = res.tokens > 0 && res.cost === 0;
+          const costNote = costUnknown ? `cost unavailable: no pricing for model (${res.tokens} tokens used)` : undefined;
           await patch(w.id, {
             status: contract.ok ? "completed" : "contract_failed",
             ended_at: new Date().toISOString(),
-            turns: res.turns,
-            tokens: res.tokens,
-            cost_usd_estimate: res.cost ?? 0,
+            turns: totals.turns,
+            tokens: totals.tokens,
+            cost_usd_estimate: totals.cost,
             contract_result: contract,
             produced_outputs: contract.checks.filter((c) => c.ok || c.actualPath).map((c) => c.path),
-            status_note: contract.ok ? undefined : contractFailureNote(contract.checks),
+            status_note: contract.ok ? costNote : [contractFailureNote(contract.checks), costNote].filter(Boolean).join(" · "),
           });
           if (contract.ok) {
             const note = await opts.onNodeCompleted?.(w.id);
             if (note) await patch(w.id, { status_note: note });
           }
-        }).finally(() => running.delete(p));
+        })().finally(() => running.delete(p));
         running.add(p);
       }
       if (running.size > 0) {
@@ -275,13 +401,46 @@ export async function runFleet(opts: RunFleetOpts): Promise<FleetState> {
         return !!n && TERMINAL_NODE_STATUSES.has(n.status);
       })) {
         break;
+      } else {
+        // Defensive: nothing in flight, not all terminal, and nothing dispatchable
+        // would mean a tight infinite loop (busy-spin, event-loop starvation).
+        // Fail the stuck nodes instead of hanging.
+        const dispatchable = spec.workers.some((w) => {
+          const n = state.nodes[w.id];
+          if (!n || (n.status !== "pending" && n.status !== "ready")) return false;
+          if (opts.nodeKills?.has(w.id)) return true; // will be killed next pass
+          return w.depends_on.every((d) => state.nodes[d]?.status === "completed");
+        });
+        const blockingDepsFailed = spec.workers.some((w) => {
+          const n = state.nodes[w.id];
+          if (!n || (n.status !== "pending" && n.status !== "ready")) return false;
+          return w.depends_on.some((d) => FAILED.has(state.nodes[d]?.status ?? ""));
+        });
+        if (!dispatchable && !blockingDepsFailed) {
+          for (const w of spec.workers) {
+            const n = state.nodes[w.id];
+            if (!n || TERMINAL_NODE_STATUSES.has(n.status)) continue;
+            await patch(w.id, {
+              status: "failed",
+              ended_at: new Date().toISOString(),
+              status_note: `stuck in non-terminal status "${n.status}" with no dispatchable path; failed to avoid scheduler hang`,
+            });
+          }
+          break;
+        }
       }
     }
     await Promise.allSettled([...running]);
   };
 
+  const runPassUntilDrained = async () => {
+    do {
+      await runPass();
+    } while (opts.relaunchRequests && opts.relaunchRequests.size > 0);
+  };
+
   if (!loop) {
-    await runPass();
+    await runPassUntilDrained();
     const anyFailed = spec.workers.some((w) =>
       ["failed", "contract_failed"].includes(state.nodes[w.id]?.status ?? ""));
     const finalStatus = opts.killSwitch?.killed ? "killed" : anyFailed ? "failed" : "completed";
@@ -311,7 +470,7 @@ export async function runFleet(opts: RunFleetOpts): Promise<FleetState> {
 
     await opts.prepareIteration?.(n, state);
 
-    await runPass();
+    await runPassUntilDrained();
 
     let verdict: Verdict | null = null;
     let verdictBody: string | null = null;

package/src/state.ts

@@ -32,6 +32,8 @@ export async function readState(fleetRoot: string): Promise<FleetState> {
     lgtm_streak: parsed.lgtm_streak ?? 0,
     paused: parsed.paused ?? false,
     iterations: parsed.iterations ?? [],
+    pid: parsed.pid,
+    heartbeat_at: parsed.heartbeat_at,
   };
 }
 
@@ -121,34 +123,34 @@ export function patchNode(
   return { ...state, nodes, cost_usd_estimate: cost };
 }
 
-export function resetForRelaunch(state: FleetState, spec: FleetSpec, nodeId: string): FleetState {
+export function relaunchResetIds(spec: FleetSpec, state: FleetState, nodeId: string): string[] {
   if (!state.nodes[nodeId]) throw new Error(`unknown node "${nodeId}"`);
-
   const dependents: Record<string, string[]> = {};
   for (const w of spec.workers) {
     for (const dep of w.depends_on) {
       (dependents[dep] ??= []).push(w.id);
     }
   }
-
-  const downstream = new Set<string>();
-  const walk = (id: string) => {
-    for (const d of dependents[id] ?? []) {
-      if (!downstream.has(d)) {
-        downstream.add(d);
-        walk(d);
-      }
+  const out = [nodeId];
+  const seen = new Set(out);
+  const queue = [nodeId];
+  while (queue.length) {
+    for (const d of dependents[queue.shift()!] ?? []) {
+      if (seen.has(d)) continue;
+      seen.add(d);
+      if (state.nodes[d]?.status === "blocked") out.push(d);
+      queue.push(d);
     }
-  };
-  walk(nodeId);
+  }
+  return out;
+}
 
+export function resetForRelaunch(state: FleetState, spec: FleetSpec, nodeId: string): FleetState {
   const fresh: NodeState = { status: "pending", turns: 0, tokens: 0, cost_usd_estimate: 0, produced_outputs: [] };
   const nodes = { ...state.nodes };
-  nodes[nodeId] = fresh;
-  for (const id of downstream) {
-    if (state.nodes[id].status === "blocked") {
-      nodes[id] = fresh;
-    }
+  for (const id of relaunchResetIds(spec, state, nodeId)) {
+    if (id !== nodeId && state.nodes[id]?.status !== "blocked") continue;
+    nodes[id] = fresh;
   }
 
   const cost = fleetCost({ ...state, nodes });