Preserve active subagents across reload

test/integration/subagent-lifecycle.test.ts

@@ -30,6 +30,7 @@ import {
   uniqueId,
   trackTempFile,
   readScreen,
+  sendCommand,
   PI_TIMEOUT,
   type TestEnv,
 } from "./harness.ts";
@@ -103,6 +104,57 @@ for (const backend of backends) {
       }
     });
 
+    // ── Reload preservation ──
+
+    it("keeps a running subagent and fresh tool wrappers across parent extension reload", async () => {
+      const id = uniqueId();
+      const startedFile = `/tmp/pi-integ-reload-start-${id}.txt`;
+      const completedFile = `/tmp/pi-integ-reload-done-${id}.txt`;
+      const followupFile = `/tmp/pi-integ-reload-followup-${id}.txt`;
+      trackTempFile(env, startedFile);
+      trackTempFile(env, completedFile);
+      trackTempFile(env, followupFile);
+
+      const surface = createTrackedSurface(env, `reload-${id}`);
+      await sleep(1000);
+
+      const task = [
+        `Call the subagent tool once with these EXACT parameters:`,
+        `  name: "Reload-${id}"`,
+        `  agent: "test-echo"`,
+        `  task: "Run: echo 'START_${id}' > '${startedFile}'; sleep 10; echo 'DONE_${id}' > '${completedFile}'"`,
+        `Do not run the command yourself.`,
+        `After that subagent result arrives, call the subagent tool once more with these EXACT parameters:`,
+        `  name: "Reload-followup-${id}"`,
+        `  agent: "test-echo"`,
+        `  task: "Write 'FOLLOWUP_${id}' to '${followupFile}'"`,
+        `After the second subagent result arrives, say RELOAD_SURVIVED_${id}.`,
+      ].join("\n");
+
+      startPi(surface, env.dir, task);
+      await waitForFile(startedFile, PI_TIMEOUT, /START_/);
+
+      sendCommand(surface, "/reload");
+
+      const completed = await waitForFile(completedFile, PI_TIMEOUT, /DONE_/);
+      assert.ok(completed.includes(`DONE_${id}`), "child should finish after parent reload");
+      const followup = await waitForFile(followupFile, PI_TIMEOUT, /FOLLOWUP_/);
+      assert.ok(
+        followup.includes(`FOLLOWUP_${id}`),
+        "the result-triggered turn should spawn through the fresh tool wrapper",
+      );
+
+      const resultMarker = `RELOAD_SURVIVED_${id}`;
+      const responseLine = new RegExp(`^ ${resultMarker}\\s*$`, "m");
+      const screen = await waitForScreen(surface, responseLine, PI_TIMEOUT, 300);
+      assert.equal(
+        screen.match(new RegExp(`^ ${resultMarker}\\s*$`, "gm"))?.length,
+        1,
+        "the fresh runtime should receive the completion exactly once",
+      );
+      assert.doesNotMatch(screen, /Subagent cancelled|ctx is stale after session replacement or reload/i);
+    });
+
     // ── In-progress activity snapshots ──
 
     it("keeps a long active tool call from surfacing false stalled status", async () => {

test/test.ts

@@ -88,16 +88,22 @@ function createMockExtensionApi() {
   const registeredTools: Array<any> = [];
   const registeredCommands: Array<any> = [];
   const registeredMessageRenderers: Array<any> = [];
+  const eventHandlers = new Map<string, Array<(event: any, ctx: any) => any>>();
   const sentUserMessages: string[] = [];
   const sentMessages: Array<any> = [];
   return {
     registeredTools,
     registeredCommands,
     registeredMessageRenderers,
+    eventHandlers,
     sentUserMessages,
     sentMessages,
     api: {
-      on() {},
+      on(event: string, handler: (event: any, ctx: any) => any) {
+        const handlers = eventHandlers.get(event) ?? [];
+        handlers.push(handler);
+        eventHandlers.set(event, handlers);
+      },
       registerTool(tool: any) {
         registeredTools.push(tool);
       },
@@ -117,10 +123,24 @@ function createMockExtensionApi() {
       getAllTools() {
         return [];
       },
+      getActiveTools() {
+        return [];
+      },
     } as any,
   };
 }
 
+async function emitMockEvent(
+  mock: ReturnType<typeof createMockExtensionApi>,
+  event: string,
+  payload: Record<string, unknown>,
+  ctx: Record<string, unknown> = {},
+) {
+  for (const handler of mock.eventHandlers.get(event) ?? []) {
+    await handler({ type: event, ...payload }, ctx);
+  }
+}
+
 function restoreEnvVar(name: string, value: string | undefined) {
   if (value === undefined) {
     delete process.env[name];
@@ -1838,6 +1858,101 @@ describe("commands", () => {
   });
 });
 
+describe("reload lifecycle", () => {
+  const testApi = (subagentsModule as any).__test__;
+  const shared = testApi.sharedRuntime as {
+    runningSubagents: Map<string, any>;
+    latestCtx: any;
+    latestPi: any;
+    widgetInterval: ReturnType<typeof setInterval> | null;
+    statusInterval: ReturnType<typeof setInterval> | null;
+    pendingDeliveries: Array<(pi: any) => void>;
+    deliveryFlushScheduled: boolean;
+    activeRuntimeToken: symbol | null;
+  };
+  const ctx = { hasUI: false, ui: { setWidget() {} } };
+
+  function resetSharedRuntime() {
+    if (shared.widgetInterval) clearInterval(shared.widgetInterval);
+    if (shared.statusInterval) clearInterval(shared.statusInterval);
+    shared.widgetInterval = null;
+    shared.statusInterval = null;
+    shared.runningSubagents.clear();
+    shared.latestCtx = null;
+    shared.latestPi = null;
+    shared.pendingDeliveries.length = 0;
+    shared.deliveryFlushScheduled = false;
+    shared.activeRuntimeToken = null;
+  }
+
+  beforeEach(resetSharedRuntime);
+  after(resetSharedRuntime);
+
+  it("preserves a live subagent and delivers its queued result through the fresh runtime", async () => {
+    const oldRuntime = createMockExtensionApi();
+    (subagentsModule as any).default(oldRuntime.api);
+    await emitMockEvent(oldRuntime, "session_start", { reason: "startup" }, ctx);
+
+    const abortController = new AbortController();
+    const running: {
+      id: string;
+      name: string;
+      task: string;
+      surface: string;
+      startTime: number;
+      sessionFile: string;
+      statusState: ReturnType<typeof createStatusState>;
+      interactive: boolean;
+      abortController: AbortController;
+      suppressDelivery?: boolean;
+    } = {
+      id: "reload-child",
+      name: "reload-child",
+      task: "keep running",
+      surface: "surface-1",
+      startTime: Date.now(),
+      sessionFile: "/tmp/reload-child.jsonl",
+      statusState: createStatusState({ source: "pi", startTimeMs: Date.now() }),
+      interactive: false,
+      abortController,
+    };
+    shared.runningSubagents.set(running.id, running);
+
+    await emitMockEvent(oldRuntime, "session_shutdown", { reason: "reload" }, ctx);
+    assert.equal(abortController.signal.aborted, false, "reload must not cancel the child");
+    assert.equal(shared.runningSubagents.get(running.id), running);
+
+    testApi.sendWhenRuntimeReady((pi: any) => {
+      pi.sendMessage({ customType: "subagent_result", content: "done", display: true });
+    });
+    assert.equal(oldRuntime.sentMessages.length, 0, "the stale runtime must not receive completion");
+    assert.equal(shared.pendingDeliveries.length, 1);
+
+    const freshRuntime = createMockExtensionApi();
+    (subagentsModule as any).default(freshRuntime.api);
+    await emitMockEvent(freshRuntime, "session_start", { reason: "reload" }, ctx);
+    await new Promise<void>((resolve) => setImmediate(resolve));
+
+    assert.equal(freshRuntime.sentMessages.length, 1);
+    assert.equal(freshRuntime.sentMessages[0].message.customType, "subagent_result");
+    assert.equal(oldRuntime.sentMessages.length, 0);
+    assert.equal(shared.pendingDeliveries.length, 0);
+    assert.equal(shared.runningSubagents.get(running.id), running);
+
+    testApi.flushPendingDeliveries();
+    assert.equal(freshRuntime.sentMessages.length, 1, "a queued completion must flush exactly once");
+
+    await emitMockEvent(oldRuntime, "session_shutdown", { reason: "quit" }, ctx);
+    assert.equal(abortController.signal.aborted, false, "a stale runtime cannot cancel rebound state");
+    assert.equal(shared.runningSubagents.get(running.id), running);
+
+    await emitMockEvent(freshRuntime, "session_shutdown", { reason: "quit" }, ctx);
+    assert.equal(abortController.signal.aborted, true, "real shutdown still cancels the child");
+    assert.equal(running.suppressDelivery, true);
+    assert.equal(shared.runningSubagents.size, 0);
+  });
+});
+
 describe("tool registration", () => {
   it("always resumes subagents as autonomous (auto-exit, non-interactive tracking)", () => {
     const testApi = (subagentsModule as any).__test__;