From 30121697eef59e3ddca91271bb63d1777b5e346a Mon Sep 17 00:00:00 2001 From: alexion Date: Sat, 1 Aug 2026 16:30:06 -0400 Subject: [PATCH] feat(pi): add subagent wait tool --- .../agents/pi/extensions/subagents/index.ts | 46 +++++++++ .../extensions/subagents/supervisor.test.ts | 95 +++++++++++++++++++ .../pi/extensions/subagents/supervisor.ts | 83 +++++++++++++++- .../agents/pi/extensions/subagents/types.ts | 12 +++ 4 files changed, 235 insertions(+), 1 deletion(-) diff --git a/modules/agents/pi/extensions/subagents/index.ts b/modules/agents/pi/extensions/subagents/index.ts index 7394769..5ea2958 100644 --- a/modules/agents/pi/extensions/subagents/index.ts +++ b/modules/agents/pi/extensions/subagents/index.ts @@ -129,6 +129,24 @@ export default function subagents(pi: ExtensionAPI) { }, }); + pi.registerTool({ + name: "subagent_wait", + label: "Wait for subagents", + description: "Block until multiple subagents are terminal or a timeout expires. Prefer setting timeoutMs so the parent turn cannot hang forever", + parameters: Type.Object({ + ids: Type.Array(Type.String({ description: "Subagent id returned by subagent_spawn or subagent_batch" })), + timeoutMs: Type.Optional(Type.Number({ description: "Maximum milliseconds to wait. Omit or use 0 to wait indefinitely" })), + mode: Type.Optional(Type.Union([Type.Literal("all"), Type.Literal("any")], { description: "Wait for all ids by default, or return after any id is terminal" })), + }), + async execute(_toolCallId, params, signal, _onUpdate, ctx) { + const input = params as { ids?: unknown; timeoutMs?: unknown; mode?: unknown }; + const ids = Array.isArray(input.ids) ? input.ids.map(String) : []; + const timeoutMs = typeof input.timeoutMs === "number" && Number.isFinite(input.timeoutMs) ? input.timeoutMs : undefined; + const mode = input.mode === "any" ? "any" : "all"; + return textResult(await getSupervisor(ctx).wait(ids, { timeoutMs, mode, signal })); + }, + }); + pi.registerTool({ name: "subagent_cancel", label: "Cancel subagent", @@ -182,6 +200,14 @@ export default function subagents(pi: ExtensionAPI) { }, }); + pi.registerCommand("subagent-wait", { + description: "Wait for subagent ids separated by spaces", + handler: async (args, ctx) => { + const { ids, timeoutMs, mode } = parseWaitArgs(args); + ctx.ui.notify(JSON.stringify(await getSupervisor(ctx).wait(ids, { timeoutMs, mode }), null, 2), "info"); + }, + }); + pi.registerCommand("subagent-ui", { description: "Toggle the bundled subagent status inspector", handler: async (_args, ctx) => { @@ -231,6 +257,26 @@ function parseSpawnArgs(args: string): SpawnRequest { return { ...request, prompt: parts.join(" ") || args } as SpawnRequest; } +function parseWaitArgs(args: string): { ids: string[]; timeoutMs?: number; mode?: "all" | "any" } { + const parts = args.trim().split(/\s+/u).filter(Boolean); + let timeoutMs: number | undefined; + let mode: "all" | "any" | undefined; + const ids: string[] = []; + while (parts.length > 0) { + const part = parts.shift(); + if (!part) continue; + if (part === "--timeout-ms" && parts[0]) { + const parsed = Number(parts.shift()); + if (Number.isFinite(parsed)) timeoutMs = parsed; + } else if (part === "--mode" && (parts[0] === "all" || parts[0] === "any")) { + mode = parts.shift() as "all" | "any"; + } else { + ids.push(part); + } + } + return { ids, timeoutMs, mode }; +} + function isProjectTrusted(ctx: ExtensionContext): boolean { const value = (ctx as unknown as { isProjectTrusted?: () => boolean }).isProjectTrusted?.(); return value === true; diff --git a/modules/agents/pi/extensions/subagents/supervisor.test.ts b/modules/agents/pi/extensions/subagents/supervisor.test.ts index 5eba4ea..7cff49e 100644 --- a/modules/agents/pi/extensions/subagents/supervisor.test.ts +++ b/modules/agents/pi/extensions/subagents/supervisor.test.ts @@ -139,3 +139,98 @@ test("maxConcurrent preserves queued records", async () => { assert.equal(runner.starts.length, 2); }); + +test("wait blocks until multiple subagents are terminal", async () => { + const runner = new FakeRunner(); + const supervisor = new Supervisor(runner, "/tmp"); + const first = await spawnStarted(supervisor, "one"); + const second = await spawnStarted(supervisor, "two"); + + const waiting = supervisor.wait([first.id, second.id], { timeoutMs: 100 }); + runner.starts[0].events.completed("one done", "agent_settled"); + await sleep(0); + + assert.equal(await Promise.race([waiting.then(() => "done"), sleep(10).then(() => "pending")]), "pending"); + + runner.starts[1].events.failed("two failed"); + const result = await waiting; + + assert.equal(result.timedOut, false); + assert.equal(result.ready, true); + assert.deepEqual(result.ids, [first.id, second.id]); + assert.equal(result.pending.length, 0); + assert.deepEqual(result.results.map((item) => item.state), ["completed", "failed"]); + assert.equal(result.results[0].result, "one done"); + assert.equal(result.results[1].error, "two failed"); +}); + +test("wait returns pending statuses on timeout", async () => { + const runner = new FakeRunner(); + const supervisor = new Supervisor(runner, "/tmp"); + const first = await spawnStarted(supervisor, "one"); + const second = await spawnStarted(supervisor, "two"); + + runner.starts[0].events.completed("one done", "agent_settled"); + const result = await supervisor.wait([first.id, second.id], { timeoutMs: 5 }); + + assert.equal(result.timedOut, true); + assert.equal(result.ready, false); + assert.deepEqual(result.results.map((item) => item.state), ["completed", "running"]); + assert.deepEqual(result.pending.map((item) => item.id), [second.id]); +}); + +test("wait any returns after the first terminal subagent", async () => { + const runner = new FakeRunner(); + const supervisor = new Supervisor(runner, "/tmp"); + const first = await spawnStarted(supervisor, "one"); + const second = await spawnStarted(supervisor, "two"); + + const waiting = supervisor.wait([first.id, second.id], { mode: "any", timeoutMs: 100 }); + runner.starts[1].events.completed("two done", "agent_settled"); + const result = await waiting; + + assert.equal(result.timedOut, false); + assert.equal(result.ready, true); + assert.deepEqual(result.results.map((item) => item.state), ["running", "completed"]); + assert.deepEqual(result.pending.map((item) => item.id), [first.id]); +}); + +test("wait rejects unknown and empty id sets", async () => { + const runner = new FakeRunner(); + const supervisor = new Supervisor(runner, "/tmp"); + + await assert.rejects(() => supervisor.wait([]), /at least one subagent id is required/); + await assert.rejects(() => supervisor.wait(["missing"]), /unknown subagent id: missing/); +}); + +test("wait abort rejects without cancelling child", async () => { + const runner = new FakeRunner(); + const supervisor = new Supervisor(runner, "/tmp"); + const accepted = await spawnStarted(supervisor, "one"); + const controller = new AbortController(); + + const waiting = supervisor.wait([accepted.id], { signal: controller.signal }); + controller.abort(); + + await assert.rejects(waiting, /subagent wait aborted/); + assert.equal(runner.starts[0].handle.cancelCalls, 0); +}); + +test("wait follows queued subagents through queue start and completion", async () => { + const runner = new FakeRunner(); + const supervisor = new Supervisor(runner, "/tmp", { maxConcurrent: 1 }); + const batch = supervisor.spawnBatch([{ prompt: "one" }, { prompt: "two" }]); + await sleep(0); + + const waiting = supervisor.wait([batch.accepted[1].id], { timeoutMs: 100 }); + assert.equal(await Promise.race([waiting.then(() => "done"), sleep(10).then(() => "pending")]), "pending"); + + runner.starts[0].events.completed("one done", "agent_settled"); + await sleep(0); + runner.starts[1].events.completed("two done", "agent_settled"); + const result = await waiting; + + assert.equal(result.timedOut, false); + assert.equal(result.ready, true); + assert.deepEqual(result.results.map((item) => item.result), ["two done"]); +}); diff --git a/modules/agents/pi/extensions/subagents/supervisor.ts b/modules/agents/pi/extensions/subagents/supervisor.ts index 86f3b32..743bfde 100644 --- a/modules/agents/pi/extensions/subagents/supervisor.ts +++ b/modules/agents/pi/extensions/subagents/supervisor.ts @@ -1,4 +1,16 @@ -import type { ChildHandle, ChildRecord, ChildRunner, ContextMode, RunnerEvents, SpawnAccepted, SpawnRequest, SubagentResult, SubagentStatus } from "./types.ts"; +import type { + ChildHandle, + ChildRecord, + ChildRunner, + ContextMode, + RunnerEvents, + SpawnAccepted, + SpawnRequest, + SubagentResult, + SubagentStatus, + SubagentWaitMode, + SubagentWaitResult, +} from "./types.ts"; import { cloneResult, cloneStatus, toAccepted } from "./status.ts"; interface RunningChild { @@ -34,6 +46,7 @@ export class Supervisor { private nextChild = 0; private readonly children = new Map(); private readonly queue: RunningChild[] = []; + private readonly waiters = new Set<() => void>(); constructor( private readonly runner: ChildRunner, @@ -76,6 +89,41 @@ export class Supervisor { return cloneResult(this.require(id).record); } + async wait( + ids: string[], + options: { timeoutMs?: number; signal?: AbortSignal; mode?: SubagentWaitMode } = {}, + ): Promise { + const uniqueIds = [...new Set(ids.map((id) => id.trim()).filter(Boolean))]; + if (uniqueIds.length === 0) throw new Error("at least one subagent id is required"); + for (const id of uniqueIds) this.require(id); + + const startedAt = Date.now(); + const mode = options.mode ?? "all"; + if (mode !== "all" && mode !== "any") throw new Error(`unknown wait mode: ${mode}`); + const deadline = options.timeoutMs && options.timeoutMs > 0 ? startedAt + options.timeoutMs : undefined; + let timedOut = false; + + while (!this.waitReady(uniqueIds, mode)) { + if (options.signal?.aborted) throw new Error("subagent wait aborted"); + const remainingMs = deadline === undefined ? undefined : deadline - Date.now(); + if (remainingMs !== undefined && remainingMs <= 0) { + timedOut = true; + break; + } + await this.nextChange(remainingMs, options.signal).catch((error) => { + if (error instanceof Error && error.message === "subagent wait timed out") timedOut = true; + else throw error; + }); + if (timedOut) break; + } + + const results = uniqueIds.map((id) => this.result(id)); + const pending = uniqueIds + .map((id) => this.status(id)) + .filter((status) => !isTerminal(status.state)); + return { ids: uniqueIds, mode, ready: this.waitReady(uniqueIds, mode), results, pending, timedOut, elapsedMs: Date.now() - startedAt }; + } + async cancel(id: string): Promise { const child = this.require(id); if (isTerminal(child.record.status.state)) return cloneStatus(child.record.status); @@ -291,6 +339,39 @@ export class Supervisor { private emitChange() { this.options.onChange?.(this.list()); + for (const waiter of this.waiters) waiter(); + } + + private waitReady(ids: string[], mode: SubagentWaitMode): boolean { + const terminal = (id: string) => isTerminal(this.require(id).record.status.state); + return mode === "all" ? ids.every(terminal) : ids.some(terminal); + } + + private nextChange(timeoutMs: number | undefined, signal: AbortSignal | undefined): Promise { + return new Promise((resolve, reject) => { + let timer: ReturnType | undefined; + const cleanup = () => { + this.waiters.delete(resolveOnce); + if (timer) clearTimeout(timer); + signal?.removeEventListener("abort", abort); + }; + const resolveOnce = () => { + cleanup(); + resolve(); + }; + const abort = () => { + cleanup(); + reject(new Error("subagent wait aborted")); + }; + this.waiters.add(resolveOnce); + signal?.addEventListener("abort", abort, { once: true }); + if (timeoutMs !== undefined) { + timer = setTimeout(() => { + cleanup(); + reject(new Error("subagent wait timed out")); + }, timeoutMs); + } + }); } private allocateId(): string { diff --git a/modules/agents/pi/extensions/subagents/types.ts b/modules/agents/pi/extensions/subagents/types.ts index ac6f407..14a5143 100644 --- a/modules/agents/pi/extensions/subagents/types.ts +++ b/modules/agents/pi/extensions/subagents/types.ts @@ -69,6 +69,18 @@ export interface SubagentResult { elapsedMs: number; } +export type SubagentWaitMode = "all" | "any"; + +export interface SubagentWaitResult { + ids: string[]; + mode: SubagentWaitMode; + ready: boolean; + results: SubagentResult[]; + pending: SubagentStatus[]; + timedOut: boolean; + elapsedMs: number; +} + export interface ChildRecord { status: SubagentStatus; result?: string;