feat(pi): add subagent wait tool

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;

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"]);
+});

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<string, RunningChild>();
   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<SubagentWaitResult> {
+    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<SubagentStatus> {
     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<void> {
+    return new Promise((resolve, reject) => {
+      let timer: ReturnType<typeof setTimeout> | 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 {

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;