@d3ara1n/pi-subagent 1.1.0 → 1.1.1

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@d3ara1n/pi-subagent",
-  "version": "1.1.0",
+  "version": "1.1.1",
   "type": "module",
   "description": "Role-based subagent orchestration for pi — delegates tasks to specialized pi child processes with configurable model roles",
   "main": "src/index.ts",

package/src/index.ts

@@ -104,6 +104,16 @@ export default function subagentExtension(pi: ExtensionAPI) {
   const collectedRuns = new Map<string, CollectedRun>();
   let runCounter = 0;
 
+  // ── Live-run reaping ─────────────────────────────────────────
+  // Every in-flight run (foreground and background alike), removed once
+  // settled. session_shutdown aborts whatever is still here so no child
+  // process outlives the parent — quit, reload, or session replacement.
+  const liveRuns = new Set<RunHandle>();
+  function trackRun(run: RunHandle): void {
+    liveRuns.add(run);
+    void run.promise.then(() => liveRuns.delete(run));
+  }
+
   // Mutable guidelines array — rebuilt in session_start to reflect agentOverrides
   const guidelines: string[] = [];
 
@@ -219,6 +229,15 @@ export default function subagentExtension(pi: ExtensionAPI) {
     }
   });
 
+  // Fires before the extension runtime is torn down (quit, reload, or
+  // session replacement). Aborting funnels through the standard abort path:
+  // children get SIGTERM → their own handlers kill grandchildren, aborted
+  // runs are audited to history, gates release. Without this, background
+  // children would burn tokens as unwaitable orphans after /reload or /new.
+  pi.on("session_shutdown", () => {
+    for (const run of liveRuns) run.abort("session shutdown");
+  });
+
   pi.registerTool({
     name: "subagent_delegate",
     label: "Delegate to subagent",
@@ -307,6 +326,7 @@ export default function subagentExtension(pi: ExtensionAPI) {
         getRolesApi: getModelRolesAPI,
         getSessionId: () => ctx.sessionManager?.getSessionId(),
       });
+      trackRun(run);
 
       // ── Background: return the id immediately; the pipeline keeps running. ──
       if (params.background) {

package/src/run.test.ts

@@ -256,6 +256,63 @@ test("abort while queued fails the run and exposes thrown for the foreground pat
   gate.release();
 });
 
+test("handle.abort() reaps a queued background run (no caller signal)", async () => {
+  const gate = new AsyncSemaphore(1);
+  await gate.acquire();
+  const spawnImpl: SpawnImpl = async () => makeResult({ output: "never" });
+
+  const run = startSubagentRun(makeDeps({ gate, spawnImpl }));
+  run.abort("session shutdown");
+  const result = await run.promise;
+
+  assert.strictEqual(run.state, "failed");
+  assert.ok(run.thrown instanceof Error);
+  assert.match(result.errorMessage!, /cancelled while queued/);
+  assert.match(result.errorMessage!, /session shutdown/);
+  gate.release();
+});
+
+test("handle.abort(reason) fails a running run with the reason in the error message", async () => {
+  const signals: AbortSignal[] = [];
+  // Mirrors real spawn's abort handling: pre-aborted signals settle immediately
+  // (an "abort" listener alone would never fire — the event already happened).
+  const honoringSpawn: SpawnImpl = (_m, _t, options) =>
+    new Promise((_resolve, reject) => {
+      signals.push(options.signal!);
+      const die = () => reject(new Error("Subagent was aborted"));
+      if (options.signal?.aborted) die();
+      else options.signal?.addEventListener("abort", die, { once: true });
+    });
+
+  const run = startSubagentRun(makeDeps({ spawnImpl: honoringSpawn }));
+  run.abort("session shutdown");
+  const result = await run.promise;
+
+  assert.strictEqual(run.state, "failed");
+  assert.match(result.errorMessage!, /Subagent was aborted \(session shutdown\)/);
+  assert.ok(run.thrown instanceof Error);
+  // The internal controller the spawn honored is the same channel abort() used.
+  assert.ok(signals[0].aborted);
+});
+
+test("a pre-aborted caller signal chains into the run before spawn", async () => {
+  const controller = new AbortController();
+  controller.abort();
+  const spawnImpl: SpawnImpl = async (_m, _t, options) => {
+    if (options.signal?.aborted) throw new Error("Subagent was aborted");
+    return makeResult({ output: "done" });
+  };
+
+  const run = startSubagentRun(makeDeps({ signal: controller.signal, spawnImpl }));
+  const result = await run.promise;
+
+  assert.strictEqual(run.state, "failed");
+  assert.strictEqual(result.errorMessage, "Subagent was aborted");
+  // abort() after settle is a no-op — the terminal state never flips.
+  run.abort("session shutdown");
+  assert.strictEqual(run.state, "failed");
+});
+
 test("subscribers are notified on progress and terminal frames", async () => {
   let notifications = 0;
   const spawnImpl: SpawnImpl = async (_m, _t, options) => {

package/src/run.ts

@@ -9,6 +9,11 @@
  * exposed via `thrown`), and a subscriber list the `wait` tool uses to mirror
  * live progress into its own tool row.
  *
+ * Every run owns an AbortController. The foreground tool signal chains into
+ * it; runs started without a caller signal (background) are still abortable
+ * via handle.abort() — session_shutdown reaps every live run that way, so no
+ * child process outlives the parent.
+ *
  * All post-processing (fallback retry, output compression, summary
  * generation, history persistence) runs inside the pipeline, so background
  * runs finish exactly like foreground ones.
@@ -52,6 +57,8 @@ export interface RunHandle {
   readonly thrown: Error | undefined;
   /** Resolves with the terminal result once the run finishes (always succeeds). */
   readonly promise: Promise<SubagentResult>;
+  /** Abort the run — no-op after settle. Tool-cancellation and session-shutdown reaping both funnel here. */
+  abort(reason?: string): void;
   /** Get notified on every frame change. Returns an unsubscribe function. */
   subscribe(fn: () => void): () => void;
 }
@@ -69,7 +76,7 @@ export interface StartRunOptions {
   cwd: string;
   /** Nesting depth for the child (CURRENT_DEPTH + 1). */
   depth: number;
-  /** Foreground callers pass the tool's AbortSignal; background runs pass none and outlive the turn. */
+  /** Foreground callers chain the tool's AbortSignal in; background runs pass none and are aborted via handle.abort() instead. */
   signal?: AbortSignal;
   /** Per-call model override ('provider/model-id'), bypassing the role's configured model. */
   modelOverride?: string;
@@ -106,6 +113,14 @@ export function startSubagentRun(opts: StartRunOptions): RunHandle {
   let snapshot: SubagentResult = inputFrame(-1, true);
   let result: SubagentResult | undefined;
   let thrown: Error | undefined;
+  let settled = false;
+  let abortReason: string | undefined;
+  const controller = new AbortController();
+  const onCallerAbort = () => controller.abort();
+  if (opts.signal) {
+    if (opts.signal.aborted) controller.abort();
+    else opts.signal.addEventListener("abort", onCallerAbort, { once: true });
+  }
   let resolvePromise!: (r: SubagentResult) => void;
   const promise = new Promise<SubagentResult>((resolve) => {
     resolvePromise = resolve;
@@ -126,11 +141,14 @@ export function startSubagentRun(opts: StartRunOptions): RunHandle {
     notify();
   };
   const finish = (terminal: SubagentResult, error?: Error) => {
+    if (settled) return;
+    settled = true;
     result = terminal;
     snapshot = terminal;
     thrown = error;
     currentState = isFailedResult(terminal) ? "failed" : "finished";
     notify();
+    opts.signal?.removeEventListener("abort", onCallerAbort);
     resolvePromise(terminal);
   };
 
@@ -158,15 +176,22 @@ export function startSubagentRun(opts: StartRunOptions): RunHandle {
         listeners.delete(fn);
       };
     },
+    abort(reason?: string) {
+      if (settled) return;
+      if (reason) abortReason = reason;
+      controller.abort();
+    },
     promise,
   };
 
   (async () => {
     // ── Concurrency gate (abortable while queued) ──
     try {
-      await opts.gate.acquire(opts.signal);
+      await opts.gate.acquire(controller.signal);
     } catch {
-      const msg = "cancelled while queued for a concurrency slot";
+      const msg =
+        "cancelled while queued for a concurrency slot" +
+        (abortReason ? ` (${abortReason})` : "");
       finish({ ...inputFrame(1, false), errorMessage: msg }, new Error(msg));
       return;
     }
@@ -262,7 +287,7 @@ export function startSubagentRun(opts: StartRunOptions): RunHandle {
         maxTurns,
         maxCost,
         depth: opts.depth,
-        signal: opts.signal,
+        signal: controller.signal,
         onProgress: emitProgress,
       });
 
@@ -294,7 +319,7 @@ export function startSubagentRun(opts: StartRunOptions): RunHandle {
             maxTurns,
             maxCost,
             depth: opts.depth,
-            signal: opts.signal,
+            signal: controller.signal,
             onProgress: emitProgress,
           });
           runResult.fallbackFrom = fallbackFrom;
@@ -348,7 +373,9 @@ export function startSubagentRun(opts: StartRunOptions): RunHandle {
         activityLog: partial.activityLog,
         budgetMs: partial.budgetMs,
         elapsedMs: partial.startTime ? Date.now() - partial.startTime : undefined,
-        errorMessage: err?.message || String(err),
+        errorMessage: abortReason
+          ? `Subagent was aborted (${abortReason})`
+          : err?.message || String(err),
       };
       // The run spawned before throwing — audit it like any terminal state.
       // The partial output is raw (compression never ran on it).

package/src/spawn.ts

@@ -19,6 +19,29 @@ const INLINE_LIMIT = 8000;
 
 const PI_CODING_AGENT_PACKAGE = "@earendil-works/pi-coding-agent";
 
+// ── Parent-exit safety net ─────────────────────────────────────
+// process.on("exit") fires synchronously on every terminal path that goes
+// through process.exit — normal quit, signal-triggered graceful shutdown,
+// emergency terminal exit, uncaught crash. SIGTERM the live children so each
+// pi child runs its own cleanup (killing ITS tracked grandchildren) instead
+// of burning tokens as an orphan. This covers the paths where the graceful
+// session_shutdown reaping never fires; a SIGKILL'd parent is beyond help.
+const liveChildren = new Set<ChildProcess>();
+let exitHookInstalled = false;
+function reapChildrenOnExit(): void {
+  if (exitHookInstalled) return;
+  exitHookInstalled = true;
+  process.on("exit", () => {
+    for (const child of liveChildren) {
+      try {
+        child.kill("SIGTERM");
+      } catch {
+        /* already dead */
+      }
+    }
+  });
+}
+
 function isRunnableScript(filePath: string): boolean {
   try {
     if (!fs.existsSync(filePath)) return false;
@@ -508,6 +531,8 @@ export async function spawnSubagent(
         stdio: ["ignore", "pipe", "pipe"],
       });
       proc = p;
+      liveChildren.add(p);
+      reapChildrenOnExit();
 
       p.stdout.on("data", (data: Buffer) => {
         buffer += data.toString();
@@ -522,6 +547,7 @@ export async function spawnSubagent(
 
       p.on("exit", () => {
         processExited = true;
+        liveChildren.delete(p);
         clearEscalationTimer();
       });
 
@@ -548,6 +574,7 @@ export async function spawnSubagent(
 
       p.on("error", (err) => {
         processExited = true;
+        liveChildren.delete(p);
         if (timeoutHandle) clearTimeout(timeoutHandle);
         clearEscalationTimer();
         if (onAbort && options.signal) options.signal.removeEventListener("abort", onAbort);