@staix/agent-hub 0.3.2 → 0.5.0

package/src/hub/budget.ts

@@ -76,6 +76,9 @@ export class Budget {
   /** `ahub budget resume`: the owner overrode a pause; readings over the gate are ignored for that peer until then. */
   private readonly ignoreUntil = new Map<PeerId, number>();
   private closed = false;
+  private recoveryHeld = false;
+  private deferred = false;
+  private readonly deferredHard = new Set<PeerId>();
 
   constructor(
     dbPath: string,
@@ -88,11 +91,33 @@ export class Budget {
       reason TEXT NOT NULL, summary TEXT, handed_off INTEGER NOT NULL DEFAULT 0, moved TEXT NOT NULL DEFAULT '[]')`);
   }
 
+  /** Freeze pause/resume and handoff transitions while a controlled restart is fenced. Readings continue to update. */
+  setRecoveryHold(held: boolean): void {
+    this.recoveryHeld = held;
+    if (!held && this.deferred && !this.closed) {
+      this.deferred = false;
+      this.tick();
+      for (const peer of this.readings.keys()) {
+        const hard = this.deferredHard.delete(peer);
+        void this.evaluate(peer, hard).catch((e: Error) => this.closed || this.hooks.notify(`budget: evaluating ${peer} failed: ${e.message}`));
+      }
+    }
+  }
+
+  get recoverySettled(): boolean {
+    return this.pausing.size === 0 && this.handingOff.size === 0;
+  }
+
   record(peer: PeerId): PauseRecord | undefined {
     const r = this.db.query("SELECT * FROM budget_pauses WHERE peer = ?").get(peer) as Record<string, any> | null;
     return r ? { peer: r.peer, since: r.since, resetsAt: r.resets_at, reason: r.reason, summary: r.summary, handedOff: !!r.handed_off, moved: JSON.parse(r.moved) } : undefined;
   }
 
+  /** Recovery integrity input: persisted pause records only, excluding live usage readings and private summaries. */
+  persistedPauseDigestRows(): { peer: PeerId; since: number; resetsAt: number; reason: string }[] {
+    return this.records().map(({ peer, since, resetsAt, reason }) => ({ peer, since, resetsAt, reason }));
+  }
+
   private records(): PauseRecord[] {
     return (this.db.query("SELECT peer FROM budget_pauses").all() as { peer: string }[]).map((r) => this.record(r.peer)!);
   }
@@ -122,6 +147,7 @@ export class Budget {
     const mine = this.readings.get(peer) ?? new Map<string, Reading>();
     this.readings.set(peer, mine);
     for (const w of windows) mine.set(w.id, { ...w, used: Math.max(0, Math.min(1, w.used)), at });
+    if (this.recoveryHeld) { this.deferred = true; if (hard) this.deferredHard.add(peer); return; }
     this.evaluate(peer, hard).catch((e: Error) => this.closed || this.hooks.notify(`budget: evaluating ${peer} failed: ${e.message}`));
   }
 
@@ -141,6 +167,7 @@ export class Budget {
   }
 
   private async evaluate(peer: PeerId, hard: boolean): Promise<void> {
+    if (this.recoveryHeld) { this.deferred = true; if (hard) this.deferredHard.add(peer); return; }
     const fresh = this.fresh(peer);
     const over = fresh.filter((r) => r.used >= this.cfg.gate);
     const open = this.record(peer);
@@ -160,6 +187,7 @@ export class Budget {
       const reason = `${worst.id} window at ${Math.round(worst.used * 100)}% (${worst.source})`;
       // Checkpoint first, pause second: a paused peer receives nothing, and a turn in flight is never cut by the coordinator.
       const summary = hard ? undefined : await this.hooks.requestCheckpoint(peer).catch(() => undefined);
+      if (this.recoveryHeld) { this.deferred = true; return; }
       if (this.closed) return; // the hub is shutting down: nothing is recorded, the next run sees the reading again
       this.db.query("INSERT OR REPLACE INTO budget_pauses (peer, since, resets_at, reason, summary) VALUES (?, ?, ?, ?, ?)").run(peer, this.now(), resetsAt, reason, summary ?? null);
       this.hooks.pause(peer);
@@ -171,10 +199,12 @@ export class Budget {
   }
 
   private async finishHandoff(peer: PeerId, summary: string | undefined): Promise<void> {
-    if (this.closed || this.handingOff.has(peer) || !this.hooks.canHandOff(peer)) return; // retried by the next tick
+    if (this.closed || this.recoveryHeld) { this.deferred = true; return; }
+    if (this.handingOff.has(peer) || !this.hooks.canHandOff(peer)) return; // retried by the next tick
     this.handingOff.add(peer);
     try {
       await this.handOff(peer, summary);
+      if (this.recoveryHeld) { this.deferred = true; return; }
     } finally {
       this.handingOff.delete(peer);
     }
@@ -202,6 +232,7 @@ export class Budget {
 
   /** Called on a timer. Resumes every peer whose window has reset. */
   tick(): void {
+    if (this.recoveryHeld) { this.deferred = true; return; }
     for (const r of this.records()) {
       if (this.now() >= r.resetsAt + 60_000) {
         if (this.hooks.attached(r.peer)) this.resume(r.peer, "the window has reset"); // else: wait, the notice needs somebody to receive it

package/src/hub/bus.ts

@@ -17,6 +17,17 @@ export interface BusOptions {
   /** Optional: rewrite a delivery before it goes out (M6 digest condensation). Must return its input on any failure. */
   condense?: (envs: Envelope[]) => Promise<Envelope[]>;
 }
+
+/** Serializable delivery state used by the controlled restart coordinator. Bodies stay in the private daemon file. */
+export interface BusSnapshot {
+  schemaVersion: 1;
+  queues: Record<string, Envelope[]>;
+  prefaces: Record<string, Envelope>;
+  seen: Envelope[];
+  attempts: Record<string, number>;
+  withdrawn: string[];
+}
+
 export const DEFAULT_BUS: BusOptions = { retryMs: 1000, batchMax: 3, batchMs: 15_000, queueCap: 200 };
 
 const SEEN_CAP = 2048;
@@ -42,6 +53,9 @@ export class Bus {
   /** What each peer was last handed, next to what it stands for: an adapter that reports a failure later hands back the former. */
   private readonly lastDelivery = new Map<PeerId, { out: Envelope[]; originals: Envelope[] }>();
   private readonly attempts = new Map<string, number>(); // `${peer}:${envelope id}` -> failed deliveries
+  private recoveryHeld = false;
+  private steering = 0;
+  private condensing = 0;
 
   constructor(opts: Partial<BusOptions> = {}) {
     this.opts = { ...DEFAULT_BUS, ...opts };
@@ -66,6 +80,57 @@ export class Bus {
     };
   }
 
+  /** Stop new deliveries while a coordinator takes a stable snapshot. Existing adapter turns are left alone. */
+  setRecoveryHold(held: boolean): void {
+    this.recoveryHeld = held;
+    if (!held) for (const id of this.peers.keys()) void this.drain(id);
+  }
+
+  get isRecoveryHeld(): boolean {
+    return this.recoveryHeld;
+  }
+
+  /** Wait until a hold has fenced condensation, steering and queue drains already in flight. */
+  async fenceRecovery(): Promise<void> {
+    this.recoveryHeld = true;
+    while (this.draining.size || this.steering || this.condensing) await Bun.sleep(0);
+  }
+
+  /** A bounded, JSON-safe representation of queued work and retry/dedupe state. */
+  snapshot(): BusSnapshot {
+    return {
+      schemaVersion: 1,
+      queues: Object.fromEntries([...this.queues].map(([id, queue]) => [id, queue.map((e) => ({ ...e, ...(e.to ? { to: [...e.to] } : {}) }))])),
+      prefaces: Object.fromEntries([...this.prefaces].map(([id, e]) => [id, { ...e, ...(e.to ? { to: [...e.to] } : {}) }])),
+      seen: [...this.seen.values()].map((e) => ({ ...e, ...(e.to ? { to: [...e.to] } : {}) })),
+      attempts: Object.fromEntries(this.attempts),
+      withdrawn: [...this.withdrawn],
+    };
+  }
+
+  /** Restore state before peers attach. Queues remain held until the coordinator calls setRecoveryHold(false). */
+  restore(snapshot: BusSnapshot): void {
+    if (snapshot.schemaVersion !== 1) throw new Error("unsupported bus recovery snapshot");
+    this.queues.clear();
+    for (const [id, queue] of Object.entries(snapshot.queues ?? {})) this.queues.set(id, queue.map((e) => ({ ...e, ...(e.to ? { to: [...e.to] } : {}) })));
+    this.prefaces.clear();
+    for (const [id, e] of Object.entries(snapshot.prefaces ?? {})) this.prefaces.set(id, { ...e, ...(e.to ? { to: [...e.to] } : {}) });
+    this.seen.clear();
+    for (const e of snapshot.seen ?? []) this.seen.set(e.id, { ...e, ...(e.to ? { to: [...e.to] } : {}) });
+    this.attempts.clear();
+    for (const [key, n] of Object.entries(snapshot.attempts ?? {})) if (Number.isInteger(n) && n > 0) this.attempts.set(key, n);
+    this.withdrawn.clear();
+    for (const id of snapshot.withdrawn ?? []) this.withdrawn.add(id);
+  }
+
+  queueIds(id: PeerId): string[] {
+    return (this.queues.get(id) ?? []).map((e) => e.id);
+  }
+
+  attemptState(): Record<string, number> {
+    return Object.fromEntries(this.attempts);
+  }
+
   tap(fn: (e: BusEvent) => void): () => void {
     this.taps.add(fn);
     return () => this.taps.delete(fn);
@@ -90,6 +155,13 @@ export class Bus {
 
   /** Context the hub wants the peer to see once: it rides in front of the next delivery instead of costing a turn of its own. */
   preface(id: PeerId, body: string): void {
+    const existing = this.prefaces.get(id);
+    if (existing) {
+      // Recovery and session recall may both contribute context. Keep the original id/hop so a queued
+      // preface remains deduplicated and append the new, separately generated context.
+      this.prefaces.set(id, { ...existing, body: `${existing.body}\n\n${body}` });
+      return;
+    }
     this.prefaces.set(id, newEnvelope(HUB, body, { to: [id], kind: "presence" }));
   }
 
@@ -107,9 +179,10 @@ export class Bus {
     const targets = (env.to ?? [...this.peers.keys()]).filter((id) => id !== env.from && this.peers.has(id));
     for (const id of targets) {
       const peer = this.peers.get(id)!;
-      if (env.priority === "important" && peer.steer && this.stateOf(id) === "busy") {
+      if (!this.recoveryHeld && env.priority === "important" && peer.steer && this.stateOf(id) === "busy") {
         // Not queued while the steer is in flight, or an idle transition would deliver it a second time.
-        peer.steer([env]).catch(() => this.enqueue(id, env, true));
+        this.steering++;
+        peer.steer([env]).catch(() => this.enqueue(id, env, true)).finally(() => { this.steering--; });
       } else this.enqueue(id, env);
     }
     return targets;
@@ -165,7 +238,7 @@ export class Bus {
     try {
       const peer = this.peers.get(id)!;
       const queue = this.queues.get(id)!;
-      while (queue.length && this.stateOf(id) === "idle") {
+      while (!this.recoveryHeld && queue.length && this.stateOf(id) === "idle") {
         const wait = this.wait(queue);
         if (wait > 0) {
           this.arm(id, wait);
@@ -178,8 +251,10 @@ export class Bus {
         // What goes out may be condensed; what comes back on failure is always the originals. An important envelope
         // is why the queue became ready, so it never waits on a model call.
         const mayCondense = this.opts.condense && !delivery.some((e) => e.priority === "important");
+        this.condensing += mayCondense ? 1 : 0;
         const out = mayCondense ? await this.opts.condense!(delivery).catch(() => delivery) : delivery;
-        if (mayCondense && this.stateOf(id) !== "idle") {
+        this.condensing -= mayCondense ? 1 : 0;
+        if (this.recoveryHeld || (mayCondense && this.stateOf(id) !== "idle")) {
           // The peer got busy while the delivery was prepared. Nothing was attempted: back to the head of the queue,
           // without counting against the envelopes.
           if (preface) this.prefaces.set(id, preface);

package/src/hub/child-process.ts

@@ -0,0 +1,55 @@
+import type { ChildProcess } from "node:child_process";
+
+/**
+ * Recovery authority belongs to the short-lived ahub wrapper/launcher. Native
+ * agents and their long-lived children must never inherit it, or their later
+ * lifecycle calls could bypass the machine recovery lock.
+ */
+export function childEnv(source: NodeJS.ProcessEnv = process.env): NodeJS.ProcessEnv {
+  const env = { ...source };
+  delete env.AGENTHUB_RECOVERY_OPERATION;
+  return env;
+}
+
+/**
+ * Stop a child owned by this hub and confirm that it exited. A timeout is an
+ * incomplete shutdown, even after SIGKILL, because callers must not reuse a
+ * port or claim ownership while the old process may still be alive.
+ */
+export async function stopOwnedProcess(proc: ChildProcess, termMs = 1_000, killMs = 2_000): Promise<void> {
+  if (proc.exitCode !== null || proc.signalCode !== null || proc.pid === undefined) return;
+
+  const waitForExit = (timeoutMs: number): Promise<boolean> =>
+    new Promise((resolve) => {
+      let timer: ReturnType<typeof setTimeout> | undefined;
+      const done = (exited: boolean) => {
+        if (timer) clearTimeout(timer);
+        proc.off("exit", onExit);
+        proc.off("error", onError);
+        resolve(exited);
+      };
+      const onExit = () => done(true);
+      const onError = () => done(true); // spawn errors have no owned process to wait for
+      proc.once("exit", onExit);
+      proc.once("error", onError);
+      timer = setTimeout(() => done(false), timeoutMs);
+      timer.unref?.();
+    });
+
+  try {
+    proc.kill("SIGTERM");
+  } catch {
+    if (proc.exitCode !== null || proc.signalCode !== null) return;
+    throw new Error(`owned child ${proc.pid} could not be terminated`);
+  }
+  if (proc.exitCode !== null || proc.signalCode !== null || (await waitForExit(termMs))) return;
+
+  try {
+    proc.kill("SIGKILL");
+  } catch {
+    if (proc.exitCode !== null || proc.signalCode !== null) return;
+    throw new Error(`owned child ${proc.pid} could not be killed`);
+  }
+  if (proc.exitCode !== null || proc.signalCode !== null || (await waitForExit(killMs))) return;
+  throw new Error(`owned child ${proc.pid} shutdown incomplete after SIGKILL`);
+}

package/src/hub/control-client.ts

@@ -1,25 +1,41 @@
 import { readFileSync } from "node:fs";
 import { join } from "node:path";
+import { projectContext } from "./project.ts";
 
 export function stateDirFor(cwd: string): string {
-  return process.env.AGENTHUB_STATE_DIR ?? join(cwd, ".agenthub", "state");
+  return projectContext(cwd).stateDir;
 }
 
 /** Control WS wire version. 2 = `deliver` carries `envs` (digests); 3 = `tools` role and task messages; 4 = budget messages and `hub_checkpoint`; 5 = `ask`; 6 = console-only `ui` session bootstrap. The plugin is installed apart from the daemon, so they can drift. */
-export const PROTOCOL = 6;
+export const PROTOCOL = 8; // controlled daemon restart/recovery RPC and project/instance identity
 
 export interface Hello {
   /** `tools`: acts for `peer` (task tools, hub_send) without being a delivery target: the MCP server Kimi and Codex run. */
   role: "peer" | "console" | "tools";
   peer?: string;
+  projectId?: string;
+  instanceId?: string;
+  projectRoot?: string;
+}
+
+export interface ControlDescriptor {
+  url: string;
+  token: string;
+  projectId?: string;
+  instanceId?: string;
+  cwd?: string;
+  protocol?: number;
+  pid?: number;
 }
 
 /** Where the daemon of this state dir listens. Re-read before every connect: port and token change per run. */
-export function readControl(stateDir: string): { url: string; token: string } | undefined {
+export function readControl(stateDir: string): ControlDescriptor | undefined {
   try {
     const status = JSON.parse(readFileSync(join(stateDir, "status.json"), "utf8"));
     const token = readFileSync(join(stateDir, "control-token"), "utf8").trim();
-    return { url: `ws://127.0.0.1:${status.controlPort}`, token };
+    if (!Number.isInteger(status.controlPort) || status.controlPort < 1 || status.controlPort > 65535 || !token) return undefined;
+    return { url: `ws://127.0.0.1:${status.controlPort}`, token, projectId: status.projectId,
+      instanceId: status.instanceId, cwd: status.cwd, protocol: status.protocol, pid: status.pid };
   } catch {
     return undefined;
   }
@@ -35,14 +51,27 @@ export class ControlClient {
 
   private constructor(private readonly ws: WebSocket) {}
 
-  static connect(stateDir: string, hello: Hello): Promise<ControlClient> {
+  static connect(stateDir: string, hello: Hello, timeoutMs = 3000): Promise<ControlClient> {
     const control = readControl(stateDir);
     if (!control) return Promise.reject(new Error(`no hub running for ${stateDir} (run: ahub up)`));
+    if (control.protocol !== undefined && control.protocol !== PROTOCOL) {
+      return Promise.reject(Object.assign(new Error(`wire version mismatch: hub speaks ${control.protocol}, CLI speaks ${PROTOCOL}; stop it with its matching CLI, then upgrade and restart`), { code: 4426 }));
+    }
+    if ((hello.projectId && hello.projectId !== control.projectId) ||
+        (hello.instanceId && hello.instanceId !== control.instanceId) ||
+        (hello.projectRoot && hello.projectRoot !== control.cwd)) {
+      return Promise.reject(Object.assign(new Error("hub project or instance does not match the selected project"), { code: 4404 }));
+    }
+    const expected = { projectId: hello.projectId ?? control.projectId, instanceId: hello.instanceId ?? control.instanceId,
+      projectRoot: hello.projectRoot ?? control.cwd };
     return new Promise((resolve, reject) => {
       const ws = new WebSocket(control.url);
       const client = new ControlClient(ws);
-      ws.onerror = () => reject(new Error(`cannot reach hub at ${control.url}`));
+      const timer = setTimeout(() => refuse(new Error(`hub connection timed out at ${control.url}`)), timeoutMs);
+      const refuse = (error: Error) => { clearTimeout(timer); reject(error); ws.close(); };
+      ws.onerror = () => refuse(new Error(`cannot reach hub at ${control.url}`));
       ws.onclose = (ev) => {
+        clearTimeout(timer);
         // A refused hello carries its reason (bad token, wire version, peer id); say that, not a guess.
         reject(Object.assign(new Error(`hub closed the connection: ${ev.reason || "stale token?"}`), { code: ev.code }));
         for (const done of client.pending.values()) done({ ok: false, error: "hub connection closed" });
@@ -50,18 +79,30 @@ export class ControlClient {
         client.onClose(ev.code, ev.reason);
       };
       ws.onmessage = (ev) => {
-        const msg = JSON.parse(String(ev.data));
+        let msg: any;
+        try { msg = JSON.parse(String(ev.data)); } catch { return refuse(new Error("invalid hub response")); }
+        if (!msg || typeof msg !== "object" || Array.isArray(msg)) return refuse(new Error("invalid hub response"));
         const done = msg.rid !== undefined ? client.pending.get(msg.rid) : undefined;
         if (!done) return client.onPush(msg);
         client.pending.delete(msg.rid);
         done(msg);
       };
-      ws.onopen = () => void client.request({ t: "hello", v: PROTOCOL, token: control.token, ...hello }).then(() => resolve(client));
+      ws.onopen = () => void client.request({ t: "hello", v: PROTOCOL, token: control.token, ...hello, ...expected }, timeoutMs).then((reply) => {
+        if (reply.t !== "welcome" || reply.ok === false) return refuse(new Error(reply.error ?? "hub refused handshake"));
+        if ((expected.projectId && reply.projectId !== expected.projectId) ||
+            (expected.instanceId && reply.instanceId !== expected.instanceId) ||
+            (expected.projectRoot && reply.cwd !== expected.projectRoot)) {
+          return refuse(Object.assign(new Error("connected hub has a different project or instance"), { code: 4404 }));
+        }
+        clearTimeout(timer);
+        resolve(client);
+      });
     });
   }
 
   /** `timeoutMs`: for calls a person waits on. The hub has its own budget; this one only makes sure the CLI never hangs. */
-  request(msg: Record<string, unknown>, timeoutMs?: number): Promise<any> {
+  request(msg: Record<string, unknown>, timeoutMs = 30_000): Promise<any> {
+    if (this.ws.readyState !== WebSocket.OPEN) return Promise.resolve({ ok: false, error: "hub connection is not open" });
     const rid = this.nextRid++;
     return new Promise((resolve) => {
       const timer = timeoutMs ? setTimeout(() => (this.pending.delete(rid), resolve({ ok: false, error: `no answer from the hub within ${Math.round(timeoutMs / 1000)} s` })), timeoutMs) : undefined;