Defer spawn_agent worktree cleanup while the session is retained

A retained or interrupt_agent keep-alive session still needs its cwd for
followup_task. Cleaning the worktree in the run's finally removed that
path too early; reclaim it from the close handle instead, matching
run.ts's persisting gate.

Author: TheGreatAxios

src/subagent/agent-fleet.ts

@@ -592,6 +592,26 @@ export function createSpawnAgentTool(deps: AgentFleetDeps): AgentTool {
           }
         : undefined;
 
+      // Aligns with run.ts: (persist && turnSucceeded) || interruptedKeepAlive.
+      // When true, the worktree stays until close_agent / eviction calls the
+      // wrapped close below; otherwise the run's finally reclaims it immediately.
+      let keepWorktreeAlive = false;
+      const reclaimWorktree = async (): Promise<void> => {
+        if (worktreeCwd === undefined) return;
+        const path = worktreeCwd;
+        worktreeCwd = undefined;
+        try {
+          await cleanupSubAgentWorktree(deps.cwd, path, {
+            stashBaseline: worktreeStashBaseline,
+            ...(worktreeHeadAtCreate !== undefined ? { headAtCreate: worktreeHeadAtCreate } : {}),
+          });
+        } catch (err: unknown) {
+          log.error("spawn_agent worktree cleanup failed: {error}", {
+            error: err instanceof Error ? err.message : String(err),
+          });
+        }
+      };
+
       const params: RunSubAgentParams = {
         // Name the trace directory after the session-store id so the
         // descendant-scoping check behind read_agent_trace can resolve this
@@ -636,9 +656,19 @@ export function createSpawnAgentTool(deps: AgentFleetDeps): AgentTool {
           : {}),
         // Keep the session open after a clean completion, and hand the
         // store a bounded close for close_agent to call later.
+        // Worktree cleanup is deferred until that close when the session
+        // stays alive for followup (agentRetained / interrupt keep-alive) —
+        // matching run.ts's persisting gate so followup_task does not hit a
+        // removed cwd.
         persist: true,
         onAgentReady: ({ close, interrupt, followup, deliver }) => {
-          deps.sessions.registerClose(session.id, close);
+          deps.sessions.registerClose(session.id, async (deadlineMs) => {
+            try {
+              await close(deadlineMs);
+            } finally {
+              await reclaimWorktree();
+            }
+          });
           deps.sessions.registerInterrupt(session.id, interrupt);
           deps.sessions.registerFollowup(session.id, followup);
           deps.sessions.registerDeliver(session.id, deliver);
@@ -662,6 +692,7 @@ export function createSpawnAgentTool(deps: AgentFleetDeps): AgentTool {
           // "completed" status. Still terminalize fleetRecords so a waiter
           // that never saw interrupt_agent (or raced it) cannot hang.
           if (result.interrupted === true) {
+            keepWorktreeAlive = true;
             deps.fleetRecords.interrupt(session.id, result.report);
             return;
           }
@@ -676,8 +707,10 @@ export function createSpawnAgentTool(deps: AgentFleetDeps): AgentTool {
           // its agent first, so the store must not treat it as resumable
           // just because retained:true was requested at spawn.
           // complete() no-ops when status is already cancelled.
+          const agentRetained = result.agentRetained === true;
+          if (agentRetained) keepWorktreeAlive = true;
           deps.sessions.complete(session.id, result.report, {
-            agentRetained: result.agentRetained === true,
+            agentRetained,
             ...(result.stopReason !== undefined ? { stopReason: result.stopReason } : {}),
           });
         })
@@ -695,15 +728,7 @@ export function createSpawnAgentTool(deps: AgentFleetDeps): AgentTool {
             status: deps.sessions.get(session.id)?.status ?? "completed",
             duration_ms: Date.now() - startedAt,
           });
-          if (worktreeCwd === undefined) return;
-          void cleanupSubAgentWorktree(deps.cwd, worktreeCwd, {
-            stashBaseline: worktreeStashBaseline,
-            ...(worktreeHeadAtCreate !== undefined ? { headAtCreate: worktreeHeadAtCreate } : {}),
-          }).catch((err: unknown) => {
-            log.error("spawn_agent worktree cleanup failed: {error}", {
-              error: err instanceof Error ? err.message : String(err),
-            });
-          });
+          if (!keepWorktreeAlive) void reclaimWorktree();
         });
 
       return fleetResult(call.id, JSON.stringify({ agent_id: session.id, status: "running" }));

src/subagent/spawn-agent-worktree.test.ts

@@ -1,14 +1,14 @@
 import { afterEach, describe, expect, test } from "bun:test";
 import { execFile } from "node:child_process";
-import { mkdtemp, rm, writeFile } from "node:fs/promises";
+import { access, mkdtemp, rm, writeFile } from "node:fs/promises";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
 import { promisify } from "node:util";
 
 import { createFleetRecords, createSpawnAgentTool } from "./agent-fleet.js";
 import { createSubAgentSessionStore } from "./session-store.js";
 import { createPermissionGate } from "../permission/gate.js";
-import type { RunSubAgentParams } from "./types.js";
+import type { RunSubAgentParams, RunSubAgentResult } from "./types.js";
 
 const run = promisify(execFile);
 
@@ -44,6 +44,26 @@ async function makeRepo(): Promise<string> {
   return dir;
 }
 
+async function pathExists(path: string): Promise<boolean> {
+  try {
+    await access(path);
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+function deferred<T>(): {
+  promise: Promise<T>;
+  resolve: (v: T) => void;
+} {
+  let resolve!: (v: T) => void;
+  const promise = new Promise<T>((res) => {
+    resolve = res;
+  });
+  return { promise, resolve };
+}
+
 describe("spawn_agent worktree isolation", () => {
   test("propagates a fresh worktree path as the worker cwd", async () => {
     const repo = await makeRepo();
@@ -113,4 +133,145 @@ describe("spawn_agent worktree isolation", () => {
     expect(result.isError).toBe(true);
     expect(ran).toBe(false);
   });
+
+  test("defers worktree cleanup while the session is retained for followup", async () => {
+    const repo = await makeRepo();
+    tempDirs.push(repo);
+    const workdirBase = await mkdtemp(join(tmpdir(), "corbits-workdir-"));
+    tempDirs.push(workdirBase);
+
+    const settle = deferred<RunSubAgentResult>();
+    let workerCwd: string | undefined;
+    const sessions = createSubAgentSessionStore();
+    const tool = createSpawnAgentTool({
+      permissionGate: testPermissionGate,
+      cwd: repo,
+      getWorkdirBase: () => workdirBase,
+      provider,
+      useWorktree: true,
+      run: async (params) => {
+        workerCwd = params.cwd;
+        params.onAgentReady?.({
+          close: async () => {},
+          interrupt: () => {},
+          followup: async () => "",
+          deliver: () => {},
+        });
+        return settle.promise;
+      },
+      sessions,
+      fleetRecords: createFleetRecords(),
+    });
+    if (tool.kind !== "full") throw new Error("expected full tool");
+    const spawned = await tool.handler(
+      {
+        id: "retain-wt",
+        name: "spawn_agent",
+        arguments: { description: "keep alive", prompt: "Do the work", intent: "explore" },
+      },
+      new AbortController().signal,
+    );
+    const content = typeof spawned.content === "string" ? spawned.content : "";
+    const agentId = (JSON.parse(content) as { agent_id: string }).agent_id;
+
+    settle.resolve({ report: "## Summary\nDone.", agentRetained: true });
+    await new Promise((resolve) => setTimeout(resolve, 50));
+
+    expect(workerCwd).toBeDefined();
+    expect(await pathExists(workerCwd!)).toBe(true);
+
+    await sessions.closeOne(agentId, 1000);
+    await new Promise((resolve) => setTimeout(resolve, 50));
+    expect(await pathExists(workerCwd!)).toBe(false);
+  });
+
+  test("defers worktree cleanup while the session is interrupted for followup", async () => {
+    const repo = await makeRepo();
+    tempDirs.push(repo);
+    const workdirBase = await mkdtemp(join(tmpdir(), "corbits-workdir-"));
+    tempDirs.push(workdirBase);
+
+    const settle = deferred<RunSubAgentResult>();
+    let workerCwd: string | undefined;
+    const sessions = createSubAgentSessionStore();
+    const tool = createSpawnAgentTool({
+      permissionGate: testPermissionGate,
+      cwd: repo,
+      getWorkdirBase: () => workdirBase,
+      provider,
+      useWorktree: true,
+      run: async (params) => {
+        workerCwd = params.cwd;
+        params.onAgentReady?.({
+          close: async () => {},
+          interrupt: () => {},
+          followup: async () => "",
+          deliver: () => {},
+        });
+        return settle.promise;
+      },
+      sessions,
+      fleetRecords: createFleetRecords(),
+    });
+    if (tool.kind !== "full") throw new Error("expected full tool");
+    const spawned = await tool.handler(
+      {
+        id: "interrupt-wt",
+        name: "spawn_agent",
+        arguments: { description: "interrupt me", prompt: "Do the work", intent: "explore" },
+      },
+      new AbortController().signal,
+    );
+    const content = typeof spawned.content === "string" ? spawned.content : "";
+    const agentId = (JSON.parse(content) as { agent_id: string }).agent_id;
+
+    settle.resolve({
+      report: "## Summary\nStopped.\n## Findings\npartial\n## Blockers\ninterrupted\n## Paths\n",
+      interrupted: true,
+    });
+    await new Promise((resolve) => setTimeout(resolve, 50));
+
+    expect(workerCwd).toBeDefined();
+    expect(await pathExists(workerCwd!)).toBe(true);
+
+    await sessions.closeOne(agentId, 1000);
+    await new Promise((resolve) => setTimeout(resolve, 50));
+    expect(await pathExists(workerCwd!)).toBe(false);
+  });
+
+  test("reclaims the worktree immediately when the agent is not retained", async () => {
+    const repo = await makeRepo();
+    tempDirs.push(repo);
+    const workdirBase = await mkdtemp(join(tmpdir(), "corbits-workdir-"));
+    tempDirs.push(workdirBase);
+
+    let workerCwd: string | undefined;
+    const tool = createSpawnAgentTool({
+      permissionGate: testPermissionGate,
+      cwd: repo,
+      getWorkdirBase: () => workdirBase,
+      provider,
+      useWorktree: true,
+      run: async (params) => {
+        workerCwd = params.cwd;
+        // Salvage / non-persist path: no agentRetained flag.
+        return { report: "## Summary\nSalvaged." };
+      },
+      sessions: createSubAgentSessionStore(),
+      fleetRecords: createFleetRecords(),
+    });
+    if (tool.kind !== "full") throw new Error("expected full tool");
+    await tool.handler(
+      {
+        id: "no-retain-wt",
+        name: "spawn_agent",
+        arguments: { description: "one shot", prompt: "Do the work", intent: "explore" },
+      },
+      new AbortController().signal,
+    );
+    await new Promise((resolve) => setTimeout(resolve, 50));
+
+    expect(workerCwd).toBeDefined();
+    expect(await pathExists(workerCwd!)).toBe(false);
+  });
 });