pi-famulus 0.1.3 → 0.1.4

package/src/subagent/registry.ts

@@ -26,9 +26,17 @@ import type {
   ChildRunRequest,
   ChildRunner,
   ChildStatus,
+  ChildTokenUsage,
+  ChildWallUsage,
   DisposableChildHandle,
 } from "./types";
 
+async function awaitDisposals(disposals: Array<void | Promise<void>>): Promise<void> {
+  const results = await Promise.allSettled(disposals);
+  const errors = results.flatMap((result) => result.status === "rejected" ? [result.reason] : []);
+  if (errors.length > 0) throw new AggregateError(errors, "Child session cleanup failed");
+}
+
 // ---------------------------------------------------------------------------
 // Appendix B public record shape
 // ---------------------------------------------------------------------------
@@ -40,6 +48,8 @@ export interface RunRecord {
     childId: string;
     name: string;
     agent: string;
+    workKind?: string;
+    queueMs?: number;
     model?: string;
     status: ChildStatus;
     /** User-authored prompt (after chain interpolation), without injected preamble. */
@@ -51,6 +61,8 @@ export interface RunRecord {
     endedAt?: number;
     /** User turns so far: 1 for the launch, +1 per accepted resume. */
     turn?: number;
+    tokenUsage?: ChildTokenUsage;
+    wallUsage?: ChildWallUsage;
   }[];
   status: "running" | "completed" | "partial" | "failed" | "interrupted";
   createdAt: number;
@@ -64,7 +76,7 @@ export interface RunRegistry {
   findChild(runId: string, childIdOrName: string): ChildHandle | undefined;
   lineage(runIdA: string, runIdB: string): boolean; // v1: same runId
   onTransition(cb: (run: RunRecord) => void): void; // fleet widget / notifications
-  disposeRun(runId: string): void;
+  disposeRun(runId: string): Promise<void>;
 }
 
 // ---------------------------------------------------------------------------
@@ -94,6 +106,8 @@ export interface ActiveChildInfo {
   runId: string;
   name: string;
   agent: string;
+  workKind?: string;
+  queueMs?: number;
   model?: string;
   status: ChildStatus;
   startedAt: number;
@@ -105,6 +119,7 @@ interface InternalChild {
   runId: string;
   name: string;
   agent: string;
+  workKind?: string;
   /** Best-effort model id for fleet / ls (set when startChild runs). */
   model?: string;
   status: ChildStatus;
@@ -115,7 +130,11 @@ interface InternalChild {
   endedAt?: number;
   turn: number;
   handle?: ChildHandle;
+  starting?: Promise<ChildHandle>;
   shouldStart?: () => boolean;
+  /** Initial admission wait while runner.start() has not returned its handle. */
+  queueStartedAt?: number;
+  queueMs: number;
 }
 
 interface InternalRun {
@@ -185,6 +204,14 @@ class FailedChildHandle implements ChildHandle {
   conversation() {
     return [];
   }
+
+  tokenUsage() {
+    return { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 };
+  }
+
+  wallUsage(): ChildWallUsage {
+    return { llmMs: 0, toolMs: 0, queueMs: 0, otherMs: 0, approximate: false };
+  }
 }
 
 export class SubagentRegistry implements RunRegistry {
@@ -229,16 +256,17 @@ export class SubagentRegistry implements RunRegistry {
     };
     this.runs.set(run.runId, run);
     this.emit(run);
-    return snapshot(run);
+    return snapshot(run, this.now());
   }
 
   get(runId: string): RunRecord | undefined {
     const run = this.runs.get(runId);
-    return run ? snapshot(run) : undefined;
+    return run ? snapshot(run, this.now()) : undefined;
   }
 
   list(): RunRecord[] {
-    return [...this.runs.values()].map(snapshot);
+    const now = this.now();
+    return [...this.runs.values()].map((run) => snapshot(run, now));
   }
 
   handle(childId: string): ChildHandle | undefined {
@@ -262,7 +290,7 @@ export class SubagentRegistry implements RunRegistry {
     this.transitionCbs.add(cb);
   }
 
-  disposeRun(runId: string): void {
+  async disposeRun(runId: string): Promise<void> {
     const run = this.runs.get(runId);
     if (!run) return;
     const now = this.now();
@@ -271,6 +299,10 @@ export class SubagentRegistry implements RunRegistry {
     // "running" record (ghost agents in ls).
     for (const child of run.children) {
       if (child.status === "pending" || child.status === "running") {
+        if (child.status === "pending" && child.queueStartedAt !== undefined) {
+          child.queueMs = Math.max(0, now - child.queueStartedAt);
+          child.queueStartedAt = undefined;
+        }
         child.status = "interrupted";
         child.endedAt = now;
         child.result = {
@@ -284,17 +316,26 @@ export class SubagentRegistry implements RunRegistry {
     this.recomputeRunStatus(run);
     this.emit(run);
     this.runs.delete(runId);
+    const disposals: Array<void | Promise<void>> = [];
     for (const child of run.children) {
       this.children.delete(child.childId);
       const handle = child.handle as DisposableChildHandle | undefined;
-      if (!handle) continue;
+      if (!handle) {
+        // Custom runners may not expose a handle until startup completes.
+        if (child.starting) disposals.push(child.starting.then(async (lateHandle) => {
+          await lateHandle.interrupt().catch(() => {});
+          await (lateHandle as DisposableChildHandle).dispose?.();
+        }, () => {}));
+        continue;
+      }
       void handle.interrupt().catch(() => {});
       try {
-        handle.dispose?.();
-      } catch {
-        // ignore
+        disposals.push(handle.dispose?.());
+      } catch (error) {
+        disposals.push(Promise.reject(error));
       }
     }
+    await awaitDisposals(disposals);
   }
 
   // -------------------------------------------------------------------------
@@ -302,7 +343,7 @@ export class SubagentRegistry implements RunRegistry {
   // -------------------------------------------------------------------------
 
   /** Allocate a pending child inside a run. Returns the child id. */
-  addChild(runId: string, info: { name: string; agent: string }): string {
+  addChild(runId: string, info: { name: string; agent: string; workKind?: string }): string {
     const run = this.runs.get(runId);
     if (!run) throw new Error(`unknown run ${runId}`);
     const child: InternalChild = {
@@ -310,9 +351,12 @@ export class SubagentRegistry implements RunRegistry {
       runId,
       name: info.name,
       agent: info.agent,
+      workKind: info.workKind ?? "other",
       status: "pending",
       startedAt: this.now(),
+      queueStartedAt: this.now(),
       turn: 1,
+      queueMs: 0,
     };
     run.children.push(child);
     this.children.set(child.childId, child);
@@ -329,6 +373,7 @@ export class SubagentRegistry implements RunRegistry {
     const child = this.children.get(req.childId);
     if (!child) throw new Error(`unknown child ${req.childId} (addChild first)`);
     child.shouldStart = opts?.shouldStart;
+    child.workKind = req.workKind ?? child.workKind ?? "other";
     child.prompt = req.taskPrompt ?? req.prompt;
     child.preamble = req.agent.systemPrompt || undefined;
     child.model = req.model ?? req.agent.model;
@@ -343,12 +388,28 @@ export class SubagentRegistry implements RunRegistry {
       return this.failWithoutSession(child, (err as Error).message);
     }
 
-    const handle = await runner.start(req);
+    child.queueStartedAt = this.now();
+    let handle: ChildHandle;
+    try {
+      child.starting = runner.start(req, (startingHandle) => { child.handle = startingHandle; });
+      handle = await child.starting;
+    } catch (error) {
+      this.closeInitialQueue(child, this.now());
+      throw error;
+    } finally {
+      child.starting = undefined;
+    }
+    // Teardown already owns this handle and its asynchronous cleanup.
+    if (!this.children.has(child.childId)) return handle;
+    this.closeInitialQueue(child, this.now());
     child.handle = handle;
-    // Prefer the actually resolved provider/id over the request-time spec
-    // (inherits parent model when neither param nor agent.model is set).
+    // Preserve the session's effective model and thinking setting, not a raw
+    // request suffix that may have fallen back to a different parent model.
     const resolved = handle.resolvedModel();
-    if (resolved) child.model = resolved;
+    if (resolved) {
+      const thinking = handle.effectiveThinkingLevel?.();
+      child.model = `${resolved}${thinking ? `:${thinking}` : ""}`;
+    }
     // Generation-1 result wiring (later generations are wired in admitChild,
     // where the handle is already visible).
     handle.result.then((result) => this.settleChild(child.childId, result));
@@ -370,7 +431,10 @@ export class SubagentRegistry implements RunRegistry {
    * and wires result settlement for resumed generations. Returns the slot
    * releaser.
    */
-  async admitChild(childId: string, ticket?: { current: () => boolean }): Promise<() => void> {
+  async reserveChildSlot(
+    childId: string,
+    ticket?: { current: () => boolean },
+  ): Promise<{ admit: () => void; release: () => void }> {
     await this.acquireSlot();
     let released = false;
     const release = () => {
@@ -383,27 +447,47 @@ export class SubagentRegistry implements RunRegistry {
       release();
       throw new Error(`unknown child ${childId}`);
     }
-    if (child.shouldStart && !child.shouldStart()) {
-      release();
-      throw new ChildCancelledError();
-    }
-    // The generation that queued this request settled while it waited
-    // (interrupt/dispose; maybe a newer resume is queued behind other
-    // children): hand the slot on, without touching the child's status.
-    // Without a ticket, fall back to the handle's status.
-    const stale = ticket ? !ticket.current() : (child.handle?.status() ?? "pending") !== "pending";
-    if (stale) {
+    const checkCurrent = () => {
+      if (child.shouldStart && !child.shouldStart()) throw new ChildCancelledError();
+      // A request settled while waiting (interrupt/dispose or a newer resume
+      // queued behind other children) must hand this local reservation on.
+      const stale = ticket ? !ticket.current() : (child.handle?.status() ?? "pending") !== "pending";
+      if (stale) throw new ChildCancelledError("settled while queued");
+    };
+    try {
+      checkCurrent();
+    } catch (error) {
+      this.closeInitialQueue(child, this.now());
       release();
-      throw new ChildCancelledError("settled while queued");
+      throw error;
     }
-    this.transitionChild(child, "running");
-    const handle = child.handle;
-    if (handle) {
-      // Resume generations: the handle (and its fresh result promise) already
-      // exists, so wire settlement here. Generation 1 is wired by startChild.
-      handle.result.then((result) => this.settleChild(childId, result));
+    let admitted = false;
+    return {
+      release,
+      admit: () => {
+        if (admitted) return;
+        checkCurrent();
+        this.closeInitialQueue(child, this.now());
+        admitted = true;
+        this.transitionChild(child, "running");
+        const handle = child.handle;
+        if (handle) {
+          // Resume generations: wire settlement only once global admission also succeeds.
+          handle.result.then((result) => this.settleChild(childId, result));
+        }
+      },
+    };
+  }
+
+  async admitChild(childId: string, ticket?: { current: () => boolean }): Promise<() => void> {
+    const reservation = await this.reserveChildSlot(childId, ticket);
+    try {
+      reservation.admit();
+      return reservation.release;
+    } catch (error) {
+      reservation.release();
+      throw error;
     }
-    return release;
   }
 
   /**
@@ -428,6 +512,8 @@ export class SubagentRegistry implements RunRegistry {
       result: child.result,
       startedAt: child.startedAt,
       endedAt: child.endedAt,
+      queueMs: child.queueMs,
+      queueStartedAt: child.queueStartedAt,
     };
     const accepted = handle.resume(message, opts);
     // Mark the child queued before awaiting: with a free slot, admitChild's
@@ -436,6 +522,8 @@ export class SubagentRegistry implements RunRegistry {
     child.result = undefined;
     child.endedAt = undefined;
     child.startedAt = this.now();
+    child.queueMs = 0;
+    child.queueStartedAt = child.startedAt;
     child.turn += 1;
     try {
       await accepted;
@@ -448,6 +536,8 @@ export class SubagentRegistry implements RunRegistry {
         child.result = previous.result;
         child.startedAt = previous.startedAt;
         child.endedAt = previous.endedAt;
+        child.queueMs = previous.queueMs;
+        child.queueStartedAt = previous.queueStartedAt;
       }
       throw err;
     }
@@ -477,12 +567,17 @@ export class SubagentRegistry implements RunRegistry {
   finalizeRun(runId: string, error: string): void {
     const run = this.runs.get(runId);
     if (!run) return;
+    const now = this.now();
     let changed = false;
     for (const child of run.children) {
       if (child.status === "pending" && !child.handle) {
+        if (child.queueStartedAt !== undefined) {
+          child.queueMs = (child.queueMs ?? 0) + Math.max(0, now - child.queueStartedAt);
+          child.queueStartedAt = undefined;
+        }
         child.status = "interrupted";
         child.result = { status: "interrupted", text: "", error, durationMs: 0 };
-        child.endedAt = this.now();
+        child.endedAt = now;
         changed = true;
       }
     }
@@ -502,6 +597,11 @@ export class SubagentRegistry implements RunRegistry {
         runId: child.runId,
         name: child.name,
         agent: child.agent,
+        workKind: child.workKind,
+        queueMs:
+          child.status === "pending" && child.queueStartedAt !== undefined
+            ? (child.queueMs ?? 0) + Math.max(0, this.now() - child.queueStartedAt)
+            : child.queueMs,
         model: child.model,
         status: child.status,
         startedAt: child.startedAt,
@@ -513,10 +613,8 @@ export class SubagentRegistry implements RunRegistry {
   }
 
   /** Dispose every run (session shutdown). */
-  disposeAll(): void {
-    for (const runId of [...this.runs.keys()]) {
-      this.disposeRun(runId);
-    }
+  async disposeAll(): Promise<void> {
+    await awaitDisposals([...this.runs.keys()].map((runId) => this.disposeRun(runId)));
   }
 
   // -------------------------------------------------------------------------
@@ -535,9 +633,15 @@ export class SubagentRegistry implements RunRegistry {
     const child = this.children.get(childId);
     if (!child) return;
     if (child.status !== "pending" && child.status !== "running") return; // already terminal
+    const settledAt = this.now();
+    if (child.status === "pending" && child.queueStartedAt !== undefined) {
+      child.queueMs = Math.max(0, settledAt - child.queueStartedAt);
+      child.queueStartedAt = undefined;
+    }
+    if (child.queueMs === undefined && result.queueMs !== undefined) child.queueMs = result.queueMs;
     child.status = result.status;
     child.result = result;
-    child.endedAt = this.now();
+    child.endedAt = settledAt;
     child.shouldStart = undefined;
     const run = this.runs.get(child.runId);
     if (run) {
@@ -550,8 +654,11 @@ export class SubagentRegistry implements RunRegistry {
     if (child.status === status) return;
     child.status = status;
     if (status === "running") {
+      const startedAt = this.now();
+      child.queueMs = Math.max(0, startedAt - (child.queueStartedAt ?? child.startedAt));
+      child.queueStartedAt = undefined;
       // (Re-)start: elapsed time and any previous generation's result reset.
-      child.startedAt = this.now();
+      child.startedAt = startedAt;
       child.result = undefined;
       child.endedAt = undefined;
     }
@@ -579,9 +686,15 @@ export class SubagentRegistry implements RunRegistry {
     else run.status = "partial";
   }
 
+  private closeInitialQueue(child: InternalChild, at: number): void {
+    if (child.queueStartedAt === undefined) return;
+    child.queueMs += Math.max(0, at - child.queueStartedAt);
+    child.queueStartedAt = undefined;
+  }
+
   private emit(run: InternalRun): void {
     if (this.transitionCbs.size === 0) return;
-    const record = snapshot(run);
+    const record = snapshot(run, this.now());
     for (const cb of this.transitionCbs) {
       try {
         cb(record);
@@ -630,7 +743,7 @@ export class SubagentRegistry implements RunRegistry {
   }
 }
 
-function snapshot(run: InternalRun): RunRecord {
+function snapshot(run: InternalRun, now: number): RunRecord {
   return {
     runId: run.runId,
     kind: run.kind,
@@ -640,6 +753,8 @@ function snapshot(run: InternalRun): RunRecord {
       childId: c.childId,
       name: c.name,
       agent: c.agent,
+      ...(c.workKind !== undefined ? { workKind: c.workKind } : {}),
+      ...(c.queueMs !== undefined ? { queueMs: c.queueMs } : {}),
       ...(c.model !== undefined ? { model: c.model } : {}),
       status: c.status,
       ...(c.prompt !== undefined ? { prompt: c.prompt } : {}),
@@ -648,6 +763,14 @@ function snapshot(run: InternalRun): RunRecord {
       startedAt: c.startedAt,
       endedAt: c.endedAt,
       turn: c.turn,
+      tokenUsage: c.handle?.tokenUsage() ?? { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+      wallUsage: c.handle?.wallUsage() ?? {
+        llmMs: 0,
+        toolMs: 0,
+        queueMs: c.queueMs + (c.queueStartedAt === undefined ? 0 : Math.max(0, now - c.queueStartedAt)),
+        otherMs: c.status === "running" ? Math.max(0, now - c.startedAt) : 0,
+        approximate: c.status === "running",
+      },
     })),
   };
 }