From 4e6baed5f05a37a8811fc7fb3a58e9ac032da2f0 Mon Sep 17 00:00:00 2001 From: alexion Date: Sat, 1 Aug 2026 14:13:44 -0400 Subject: [PATCH] feat(pi): harden subagent lifecycle --- .../agents/pi/extensions/subagents/index.ts | 7 ++ .../agents/pi/extensions/subagents/runner.ts | 44 ++++++- .../extensions/subagents/supervisor.test.ts | 111 ++++++++++++++++++ .../pi/extensions/subagents/supervisor.ts | 87 ++++++++++++-- 4 files changed, 236 insertions(+), 13 deletions(-) create mode 100644 modules/agents/pi/extensions/subagents/supervisor.test.ts diff --git a/modules/agents/pi/extensions/subagents/index.ts b/modules/agents/pi/extensions/subagents/index.ts index ac63f49..2be89dc 100644 --- a/modules/agents/pi/extensions/subagents/index.ts +++ b/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; diff --git a/modules/agents/pi/extensions/subagents/runner.ts b/modules/agents/pi/extensions/subagents/runner.ts index 7683827..8c88112 100644 --- a/modules/agents/pi/extensions/subagents/runner.ts +++ b/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(); constructor( @@ -46,11 +49,12 @@ class RpcChildHandle implements ChildHandle { } async cancel(): Promise { + 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(); - if (!this.child.killed) this.child.kill("SIGTERM"); + this.terminate(); } private onStdout(chunk: string) { @@ -97,14 +101,38 @@ class RpcChildHandle implements ChildHandle { } private async finish() { - if (this.settled) return; - this.settled = true; + if (this.settled || this.finishing) return; + 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 { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + function childEnvironment(): NodeJS.ProcessEnv { const env = { ...process.env }; delete env.PI_SESSION_ID; diff --git a/modules/agents/pi/extensions/subagents/supervisor.test.ts b/modules/agents/pi/extensions/subagents/supervisor.test.ts new file mode 100644 index 0000000..69f1bc8 --- /dev/null +++ b/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 { + 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 { + 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); +}); diff --git a/modules/agents/pi/extensions/subagents/supervisor.ts b/modules/agents/pi/extensions/subagents/supervisor.ts index 6f45ab7..2129536 100644 --- a/modules/agents/pi/extensions/subagents/supervisor.ts +++ b/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; + runTimer?: ReturnType; } +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(); @@ -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(); - 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"; + this.completeWithoutResult(child, "cancelled", "cancelled"); return cloneStatus(child.record.status); } async shutdown(): Promise { 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();