@pify/swarm 0.4.0 → 0.5.0

package/extensions/swarm.ts

@@ -25,6 +25,7 @@ import { Text } from "@earendil-works/pi-tui";
 import { Type } from "typebox";
 
 import { BUILTIN_AGENTS } from "../src/builtin.ts";
+import { LiveChildren, cancelNote, type CancelReason } from "../src/cancel.ts";
 import { createIsolationWorktree, isolationNote, removeIfUnchanged } from "../src/isolate.ts";
 import { formatInbox, mailboxDir, mailboxPrompt, postMessage, readInbox } from "../src/mailbox.ts";
 import { parseAgentFile } from "../src/frontmatter.ts";
@@ -77,6 +78,8 @@ function loadDefs(cwd: string, agentDir: string): Map<string, AgentDef> {
 export default function swarm(pi: ExtensionAPI) {
   let defs = new Map<string, AgentDef>();
   const runs = new Map<string, SwarmRun>();
+  /** Live child sessions per run, so a stop actually reaches the children. */
+  const live = new LiveChildren();
   let activeRun: SwarmRun | null = null;
   let runCounter = 0;
   let lastUiCtx: UiContext | null = null;
@@ -86,7 +89,7 @@ export default function swarm(pi: ExtensionAPI) {
     lastUiCtx = ctx;
     const run = activeRun;
     const now = Date.now();
-    if (!run || (run.status === "done" && (run.finishedAt ?? 0) < now - 15_000)) {
+    if (!run || (run.status !== "running" && (run.finishedAt ?? 0) < now - 15_000)) {
       ctx.ui.setWidget("swarm", undefined);
       return;
     }
@@ -151,6 +154,7 @@ export default function swarm(pi: ExtensionAPI) {
 
   async function runItem(
     ctx: UiContext,
+    runId: string,
     def: AgentDef,
     item: ItemState,
     context: string,
@@ -161,6 +165,7 @@ export default function swarm(pi: ExtensionAPI) {
     renderWidget();
     let session: AgentSession | null = null;
     let unsubscribe: (() => void) | null = null;
+    let releaseLive: (() => void) | null = null;
     try {
       let model = ctx.model ?? null;
       if (def.model) {
@@ -198,6 +203,7 @@ export default function swarm(pi: ExtensionAPI) {
         }),
       });
       session = created.session;
+      releaseLive = live.register(runId, session);
 
       unsubscribe = session.subscribe((event) => {
         if (event.type === "message_end" && (event as { message?: { role?: string } }).message?.role === "assistant") {
@@ -232,6 +238,7 @@ export default function swarm(pi: ExtensionAPI) {
       item.status = "error";
       item.error = err instanceof Error ? err.message : String(err);
     } finally {
+      if (releaseLive) releaseLive();
       if (unsubscribe) {
         try {
           unsubscribe();
@@ -264,6 +271,9 @@ export default function swarm(pi: ExtensionAPI) {
     const queue = [...run.items];
     const workers = Array.from({ length: Math.min(DEFAULT_CONCURRENCY, queue.length) }, async () => {
       for (;;) {
+        // A cancelled run stops taking new items; the ones already in flight
+        // were aborted by cancelRun.
+        if (run.status === "cancelled") return;
         const item = queue.shift();
         if (!item) return;
         const def = routeItem(item.item, defs, fixed);
@@ -271,24 +281,43 @@ export default function swarm(pi: ExtensionAPI) {
         if (isolate) {
           try {
             const iso = createIsolationWorktree(ctx.cwd, run.runId + "-i" + (item.index + 1));
-            await runItem(ctx, def, item, context, iso.path, mailbox);
+            await runItem(ctx, run.runId, def, item, context, iso.path, mailbox);
             if (item.result !== null) item.result = `${item.result}\n\n${isolationNote(iso)}`;
           } catch (err) {
             item.status = "error";
             item.error = err instanceof Error ? err.message : String(err);
           }
         } else {
-          await runItem(ctx, def, item, context, undefined, mailbox);
+          await runItem(ctx, run.runId, def, item, context, undefined, mailbox);
         }
       }
     });
     await Promise.all(workers);
-    run.status = "done";
+    if (run.status !== "cancelled") run.status = "done";
     run.finishedAt = Date.now();
     pi.appendEntry(RUN_ENTRY, run);
     renderWidget();
   }
 
+  /**
+   * Stop a run and every child it started. Both meanings of "stop" — the
+   * user's abort and session teardown — come through here.
+   */
+  function cancelRun(run: SwarmRun, reason: CancelReason): void {
+    const stopped = live.abortRun(run.runId);
+    if (run.status === "running") {
+      run.status = "cancelled";
+      run.finishedAt = Date.now();
+    }
+    for (const item of run.items) {
+      if (item.status === "running" || item.status === "queued") {
+        item.status = "aborted";
+        item.error = cancelNote(reason, stopped);
+      }
+    }
+    renderWidget();
+  }
+
   // ── Tools ────────────────────────────────────────────────────────────
 
   pi.registerTool({
@@ -323,7 +352,7 @@ export default function swarm(pi: ExtensionAPI) {
         isolation?: string;
         mailbox?: boolean;
       },
-      _signal,
+      signal,
       _onUpdate,
       ctx,
     ) {
@@ -359,6 +388,18 @@ export default function swarm(pi: ExtensionAPI) {
       activeRun = run;
       renderWidget(uiCtx);
 
+      // Esc must reach the children. A background run outlives this tool call
+      // by design, so its signal is not its cancel button.
+      let stopListening: (() => void) | null = null;
+      if (signal && !run.background) {
+        const onAbort = () => cancelRun(run, "user-abort");
+        if (signal.aborted) onAbort();
+        else {
+          signal.addEventListener("abort", onAbort, { once: true });
+          stopListening = () => signal.removeEventListener("abort", onAbort);
+        }
+      }
+
       if (run.background) {
         void executeRun(uiCtx, run, params.context ?? "", params.agent, params.isolation === "worktree", params.mailbox === true).then(() => {
           notify(uiCtx, `swarm ${run.runId} finished — collect with swarm_status`, "info");
@@ -371,7 +412,11 @@ export default function swarm(pi: ExtensionAPI) {
         };
       }
 
-      await executeRun(uiCtx, run, params.context ?? "", params.agent, params.isolation === "worktree", params.mailbox === true);
+      try {
+        await executeRun(uiCtx, run, params.context ?? "", params.agent, params.isolation === "worktree", params.mailbox === true);
+      } finally {
+        if (stopListening) stopListening();
+      }
       return {
         content: [{ type: "text", text: buildReport(run) }],
         details: { runId: run.runId },
@@ -389,7 +434,13 @@ export default function swarm(pi: ExtensionAPI) {
     async execute(_id, params: { runId?: string }) {
       const run = params.runId ? runs.get(params.runId.trim()) : activeRun ?? [...runs.values()].pop();
       if (!run) throw new Error("No swarm runs this session.");
-      const text = run.status === "done" ? buildReport(run) : buildStatusLine(run);
+      const text =
+        run.status === "cancelled"
+          ? "This run was cancelled before it finished. Below is what the items that did complete produced.\n" +
+            buildReport(run)
+          : run.status === "done"
+            ? buildReport(run)
+            : buildStatusLine(run);
       return { content: [{ type: "text", text }], details: { runId: run.runId, status: run.status } };
     },
   });
@@ -404,7 +455,7 @@ export default function swarm(pi: ExtensionAPI) {
       const e = entry as { type?: string; customType?: string; data?: unknown };
       if (e.type !== "custom" || e.customType !== RUN_ENTRY || !isRecord(e.data)) continue;
       const run = e.data as unknown as SwarmRun;
-      if (typeof run.runId === "string" && run.status === "done") {
+      if (typeof run.runId === "string" && run.status !== "running") {
         runs.set(run.runId, run);
         const n = Number.parseInt(run.runId.slice(1), 10);
         if (Number.isFinite(n) && n > runCounter) runCounter = n;
@@ -414,6 +465,10 @@ export default function swarm(pi: ExtensionAPI) {
   });
 
   pi.on("session_shutdown", async (_event, ctx) => {
+    // A run cannot outlive the session that owns it.
+    for (const run of runs.values()) {
+      if (run.status === "running") cancelRun(run, "session-switch");
+    }
     if (ctx.hasUI) ctx.ui.setWidget("swarm", undefined);
   });
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@pify/swarm",
-  "version": "0.4.0",
+  "version": "0.5.0",
   "description": "Coordinate multiple pi agents in parallel: swarm_run fan-out with per-item auto-routing, concurrency queue, aggregated reports",
   "keywords": [
     "pi-package",

package/src/cancel.ts

@@ -0,0 +1,104 @@
+/**
+ * Stopping means stopping the children too.
+ *
+ * A run in this package is not one process: it is a tree of child agent
+ * sessions, each with its own provider connection. The tool that started them
+ * is handed an AbortSignal and the extension is told when the session goes
+ * away — and until now neither reached the children. Pressing Esc, or
+ * switching sessions with a run in flight, marked a record "aborted" while the
+ * children kept talking to the provider on the user's money, writing into a
+ * conversation nobody was reading.
+ *
+ * So every live child registers here, and the two places that mean "stop"
+ * abort all of them. (The rule is FradSer-adjacent prior art: @zhushanwen's
+ * subagent-workflow terminates running runs on session switch or shutdown
+ * rather than letting them outlive the session that owns them.)
+ */
+
+/** The part of a child agent session this module needs. */
+export interface Abortable {
+  abort(): unknown;
+}
+
+export type CancelReason = "user-abort" | "session-switch" | "timeout";
+
+/**
+ * Live child sessions, grouped by the run that owns them. Registration
+ * returns its own release, so a child that finishes normally leaves no trace
+ * and cannot be aborted twice.
+ */
+export class LiveChildren {
+  private byRun = new Map<string, Set<Abortable>>();
+
+  register(runId: string, child: Abortable): () => void {
+    let set = this.byRun.get(runId);
+    if (!set) {
+      set = new Set();
+      this.byRun.set(runId, set);
+    }
+    set.add(child);
+    return () => {
+      const current = this.byRun.get(runId);
+      if (!current) return;
+      current.delete(child);
+      if (current.size === 0) this.byRun.delete(runId);
+    };
+  }
+
+  /** How many children of this run are still live. */
+  count(runId: string): number {
+    return this.byRun.get(runId)?.size ?? 0;
+  }
+
+  /** Total live children across every run. */
+  total(): number {
+    let sum = 0;
+    for (const set of this.byRun.values()) sum += set.size;
+    return sum;
+  }
+
+  /**
+   * Abort every live child of one run and return how many were stopped. A
+   * child that throws from abort() is still counted and still dropped: the
+   * point is that nothing is left holding a connection, and one stubborn
+   * child must not spare the others.
+   */
+  abortRun(runId: string): number {
+    const set = this.byRun.get(runId);
+    if (!set) return 0;
+    let stopped = 0;
+    for (const child of [...set]) {
+      try {
+        const result = child.abort();
+        // abort() is async in pi; a rejection here is not ours to surface.
+        void Promise.resolve(result).catch(() => {});
+      } catch {
+        // already gone
+      }
+      stopped++;
+    }
+    this.byRun.delete(runId);
+    return stopped;
+  }
+
+  /** Abort every live child of every run. */
+  abortAll(): number {
+    let stopped = 0;
+    for (const runId of [...this.byRun.keys()]) stopped += this.abortRun(runId);
+    return stopped;
+  }
+}
+
+/** One line for the run log, naming who stopped it and what that cost. */
+export function cancelNote(reason: CancelReason, stopped: number): string {
+  const children =
+    stopped === 0 ? "no child agents were running" : `${stopped} child agent${stopped === 1 ? "" : "s"} stopped`;
+  switch (reason) {
+    case "user-abort":
+      return `Cancelled by the user — ${children}. Work already finished is kept; the run itself did not complete.`;
+    case "session-switch":
+      return `The session went away, so the run was terminated — ${children}. Tokens already spent are not recoverable; start a new run for a result.`;
+    case "timeout":
+      return `The run exceeded its time limit — ${children}.`;
+  }
+}

package/src/types.ts

@@ -61,7 +61,11 @@ export interface ItemState {
   error: string | null;
 }
 
-export type RunStatus = "running" | "done";
+/**
+ * "cancelled" is its own outcome, not a completion: someone stopped the run,
+ * and calling it done would report results nobody produced.
+ */
+export type RunStatus = "running" | "done" | "cancelled";
 
 export interface SwarmRun {
   runId: string;