alexion/dotfiles
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(pi): harden subagent lifecycle

Browse Source
This commit is contained in:
alexion
2026-08-01 14:13:44 -04:00
parent 63da676b5a
commit 289ea1344c
4 changed files with 236 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
modules/agents/pi/extensions/subagents/index.ts
View File

@@ -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;

44
modules/agents/pi/extensions/subagents/runner.ts
View File

@@ -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();
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<void> {
return new Promise((resolve) => setTimeout(resolve, ms));
}
function childEnvironment(): NodeJS.ProcessEnv {
const env = { ...process.env };
delete env.PI_SESSION_ID;

111
modules/agents/pi/extensions/subagents/supervisor.test.ts Normal file
View File

@@ -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);
});

87
modules/agents/pi/extensions/subagents/supervisor.ts
View File

@@ -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();
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<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();