feat(pi): harden subagent lifecycle

modules/agents/pi/extensions/subagents/index.ts

@@ -104,6 +104,13 @@ export default function subagents(pi: ExtensionAPI) {
     },
   });
 
+  pi.registerCommand("subagent-cancel", {
+    description: "Cancel a running subagent by id",
+    handler: async (args, ctx) => {
+      ctx.ui.notify(JSON.stringify(await getSupervisor(ctx).cancel(args.trim()), null, 2), "info");
+    },
+  });
+
   pi.on("session_shutdown", async () => {
     await supervisor?.shutdown();
     supervisor = undefined;

modules/agents/pi/extensions/subagents/runner.ts

@@ -21,6 +21,9 @@ class RpcChildHandle implements ChildHandle {
   private buffer = "";
   private nextRequest = 0;
   private settled = false;
+  private finishing = false;
+  private cancelling = false;
+  private killed = false;
   private readonly pending = new Map<string, PendingResponse>();
 
   constructor(
@@ -46,11 +49,12 @@ class RpcChildHandle implements ChildHandle {
   }
 
   async cancel(): Promise<void> {
+    if (this.cancelling) return;
+    this.cancelling = true;
     try {
-      await this.send("abort", {});
+      await Promise.race([this.send("abort", {}), delay(200)]);
     } catch {}
-    this.child.stdin.end();
+    this.terminate();
-    if (!this.child.killed) this.child.kill("SIGTERM");
   }
 
   private onStdout(chunk: string) {
@@ -97,14 +101,38 @@ class RpcChildHandle implements ChildHandle {
   }
 
   private async finish() {
-    if (this.settled) return;
+    if (this.settled || this.finishing) return;
-    this.settled = true;
+    this.finishing = true;
     this.events.settling();
     const result = await this.send("get_last_assistant_text", {});
     const text = typeof result === "string" ? result : result && typeof result === "object" && "text" in result ? String((result as { text: unknown }).text) : "";
+    this.settled = true;
     this.events.completed(text, "agent_settled");
+    this.terminate();
+  }
+
+  private terminate() {
+    if (this.killed) return;
+    this.killed = true;
     this.child.stdin.end();
-    if (!this.child.killed) this.child.kill("SIGTERM");
+    if (this.child.killed) return;
+    if (process.platform !== "win32" && this.child.pid) {
+      try {
+        process.kill(-this.child.pid, "SIGTERM");
+      } catch {
+        this.child.kill("SIGTERM");
+      }
+      setTimeout(() => {
+        if (this.child.killed || !this.child.pid) return;
+        try {
+          process.kill(-this.child.pid, "SIGKILL");
+        } catch {
+          this.child.kill("SIGKILL");
+        }
+      }, 2_000).unref();
+      return;
+    }
+    this.child.kill("SIGTERM");
   }
 
   private fail(error: string) {
@@ -142,6 +170,10 @@ export class SubprocessRpcRunner implements ChildRunner {
   }
 }
 
+function delay(ms: number): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
 function childEnvironment(): NodeJS.ProcessEnv {
   const env = { ...process.env };
   delete env.PI_SESSION_ID;

modules/agents/pi/extensions/subagents/supervisor.test.ts

@@ -0,0 +1,111 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+import { Supervisor } from "./supervisor.ts";
+import type { ChildHandle, ChildRunner, RunnerEvents, SpawnRequest } from "./types.ts";
+
+class FakeHandle implements ChildHandle {
+  cancelCalls = 0;
+
+  async cancel(): Promise<void> {
+    this.cancelCalls += 1;
+  }
+}
+
+class FakeRunner implements ChildRunner {
+  starts: Array<{ id: string; request: SpawnRequest; events: RunnerEvents; handle: FakeHandle }> = [];
+  autoAccept = true;
+
+  async start(id: string, request: SpawnRequest, _cwd: string, events: RunnerEvents): Promise<ChildHandle> {
+    const handle = new FakeHandle();
+    this.starts.push({ id, request, events, handle });
+    if (this.autoAccept) events.accepted(`session-${id}`);
+    return handle;
+  }
+}
+
+const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
+
+async function spawnStarted(supervisor: Supervisor, prompt = "work") {
+  const accepted = supervisor.spawn({ prompt });
+  await sleep(0);
+  return accepted;
+}
+
+test("cancel is idempotent and reaches cancelled", async () => {
+  const runner = new FakeRunner();
+  const supervisor = new Supervisor(runner, "/tmp");
+  const accepted = await spawnStarted(supervisor);
+
+  const first = await supervisor.cancel(accepted.id);
+  const second = await supervisor.cancel(accepted.id);
+
+  assert.equal(first.state, "cancelled");
+  assert.equal(second.state, "cancelled");
+  assert.equal(runner.starts[0].handle.cancelCalls, 1);
+});
+
+test("startup timeout reaches timed_out", async () => {
+  const runner = new FakeRunner();
+  runner.autoAccept = false;
+  const supervisor = new Supervisor(runner, "/tmp", { timeouts: { startMs: 5 } });
+  const accepted = await spawnStarted(supervisor);
+
+  await sleep(20);
+
+  const status = supervisor.status(accepted.id);
+  assert.equal(status.state, "timed_out");
+  assert.equal(status.stopReason, "start_timeout");
+  assert.equal(runner.starts[0].handle.cancelCalls, 1);
+});
+
+test("runtime timeout reaches timed_out", async () => {
+  const runner = new FakeRunner();
+  const supervisor = new Supervisor(runner, "/tmp", { timeouts: { runMs: 5 } });
+  const accepted = await spawnStarted(supervisor);
+
+  await sleep(20);
+
+  const status = supervisor.status(accepted.id);
+  assert.equal(status.state, "timed_out");
+  assert.equal(status.stopReason, "run_timeout");
+  assert.equal(runner.starts[0].handle.cancelCalls, 1);
+});
+
+test("process failure reaches failed with diagnostics", async () => {
+  const runner = new FakeRunner();
+  const supervisor = new Supervisor(runner, "/tmp");
+  const accepted = await spawnStarted(supervisor);
+
+  runner.starts[0].events.failed("process closed with code 1");
+
+  const status = supervisor.status(accepted.id);
+  assert.equal(status.state, "failed");
+  assert.equal(status.error, "process closed with code 1");
+});
+
+test("shutdown cancels running children", async () => {
+  const runner = new FakeRunner();
+  const supervisor = new Supervisor(runner, "/tmp");
+  const accepted = await spawnStarted(supervisor);
+
+  await supervisor.shutdown();
+
+  const status = supervisor.status(accepted.id);
+  assert.equal(status.state, "cancelled");
+  assert.equal(status.stopReason, "shutdown");
+  assert.equal(runner.starts[0].handle.cancelCalls, 1);
+});
+
+test("completed children ignore later cancel", async () => {
+  const runner = new FakeRunner();
+  const supervisor = new Supervisor(runner, "/tmp");
+  const accepted = await spawnStarted(supervisor);
+
+  runner.starts[0].events.completed("done", "agent_settled");
+  await supervisor.cancel(accepted.id);
+
+  const result = supervisor.result(accepted.id);
+  assert.equal(result.state, "completed");
+  assert.equal(result.result, "done");
+  assert.equal(runner.starts[0].handle.cancelCalls, 0);
+});

modules/agents/pi/extensions/subagents/supervisor.ts

@@ -4,8 +4,22 @@ import { cloneResult, cloneStatus, toAccepted } from "./status.ts";
 interface RunningChild {
   record: ChildRecord;
   handle?: ChildHandle;
+  startTimer?: ReturnType<typeof setTimeout>;
+  runTimer?: ReturnType<typeof setTimeout>;
 }
 
+interface SupervisorOptions {
+  timeouts?: {
+    startMs?: number;
+    runMs?: number;
+  };
+}
+
+const DEFAULT_TIMEOUTS = {
+  startMs: 30_000,
+  runMs: 0,
+};
+
 export class Supervisor {
   private nextChild = 0;
   private readonly children = new Map<string, RunningChild>();
@@ -13,6 +27,7 @@ export class Supervisor {
   constructor(
     private readonly runner: ChildRunner,
     private readonly cwd: string,
+    private readonly options: SupervisorOptions = {},
   ) {}
 
   spawn(request: SpawnRequest): SpawnAccepted {
@@ -41,6 +56,7 @@ export class Supervisor {
     const child: RunningChild = { record: { status } };
     this.children.set(id, child);
     this.setState(child.record.status, "starting", "starting");
+    this.armStartTimer(child);
 
     setTimeout(() => {
       if (isTerminal(child.record.status.state)) return;
@@ -48,6 +64,7 @@ export class Supervisor {
         .start(id, { ...request, prompt, context: status.context, tools: status.tools }, this.cwd, this.eventsFor(child.record))
         .then((handle) => {
           child.handle = handle;
+          if (isTerminal(child.record.status.state)) void handle.cancel();
         })
         .catch((error) => {
           this.fail(child.record, error instanceof Error ? error.message : String(error));
@@ -73,19 +90,17 @@ export class Supervisor {
     const child = this.require(id);
     if (isTerminal(child.record.status.state)) return cloneStatus(child.record.status);
     await child.handle?.cancel();
-    const now = new Date().toISOString();
+    this.completeWithoutResult(child, "cancelled", "cancelled");
-    child.record.status.state = "cancelled";
-    child.record.status.completedAt = now;
-    child.record.status.lastEvent = "cancelled";
-    child.record.status.lastEventAt = now;
-    child.record.status.stopReason = "cancelled";
     return cloneStatus(child.record.status);
   }
 
   async shutdown(): Promise<void> {
     await Promise.allSettled(
       [...this.children.values()].map(async (child) => {
-        if (!isTerminal(child.record.status.state)) await child.handle?.cancel();
+        if (!isTerminal(child.record.status.state)) {
+          await child.handle?.cancel();
+          this.completeWithoutResult(child, "cancelled", "shutdown");
+        }
       }),
     );
   }
@@ -93,6 +108,11 @@ export class Supervisor {
   private eventsFor(record: ChildRecord): RunnerEvents {
     return {
       accepted: (childSession) => {
+        const child = this.findChild(record);
+        if (child) {
+          this.clearTimer(child, "startTimer");
+          this.armRunTimer(child);
+        }
         if (childSession) record.status.childSession = childSession;
         this.setState(record.status, "running", "prompt accepted");
       },
@@ -104,6 +124,8 @@ export class Supervisor {
       },
       completed: (result, stopReason) => {
         const now = new Date().toISOString();
+        const child = this.findChild(record);
+        if (child) this.clearTimers(child);
         record.result = result;
         record.status.state = "completed";
         record.status.completedAt = now;
@@ -118,6 +140,8 @@ export class Supervisor {
 
   private fail(record: ChildRecord, error: string) {
     if (isTerminal(record.status.state)) return;
+    const child = this.findChild(record);
+    if (child) this.clearTimers(child);
     const now = new Date().toISOString();
     record.status.state = "failed";
     record.status.completedAt = now;
@@ -127,6 +151,55 @@ export class Supervisor {
     record.status.stopReason = "failed";
   }
 
+  private completeWithoutResult(child: RunningChild, state: "cancelled" | "timed_out", reason: string) {
+    if (isTerminal(child.record.status.state)) return;
+    this.clearTimers(child);
+    const now = new Date().toISOString();
+    child.record.status.state = state;
+    child.record.status.completedAt = now;
+    child.record.status.lastEvent = state;
+    child.record.status.lastEventAt = now;
+    child.record.status.stopReason = reason;
+  }
+
+  private armStartTimer(child: RunningChild) {
+    const timeout = this.options.timeouts?.startMs ?? DEFAULT_TIMEOUTS.startMs;
+    if (timeout <= 0) return;
+    child.startTimer = setTimeout(() => {
+      this.timeout(child, "start_timeout");
+    }, timeout);
+  }
+
+  private armRunTimer(child: RunningChild) {
+    const timeout = this.options.timeouts?.runMs ?? DEFAULT_TIMEOUTS.runMs;
+    if (timeout <= 0) return;
+    child.runTimer = setTimeout(() => {
+      this.timeout(child, "run_timeout");
+    }, timeout);
+  }
+
+  private timeout(child: RunningChild, reason: string) {
+    if (isTerminal(child.record.status.state)) return;
+    void child.handle?.cancel();
+    this.completeWithoutResult(child, "timed_out", reason);
+  }
+
+  private clearTimers(child: RunningChild) {
+    this.clearTimer(child, "startTimer");
+    this.clearTimer(child, "runTimer");
+  }
+
+  private clearTimer(child: RunningChild, key: "startTimer" | "runTimer") {
+    const timer = child[key];
+    if (!timer) return;
+    clearTimeout(timer);
+    child[key] = undefined;
+  }
+
+  private findChild(record: ChildRecord): RunningChild | undefined {
+    return [...this.children.values()].find((child) => child.record === record);
+  }
+
   private setState(status: SubagentStatus, state: SubagentStatus["state"], event: string) {
     if (isTerminal(status.state)) return;
     const now = new Date().toISOString();