🤖 fix: session disposal does not wait for a stopping migration

src/node/services/backgroundProcessManager.test.ts

@@ -2998,6 +2998,27 @@ describe("BackgroundProcessManager", () => {
       expect(ws2Processes.length).toBeGreaterThanOrEqual(1);
       expect(ws2Processes.some((p) => p.status === "running")).toBe(true);
     });
+
+    it("lets session disposal finish once a migration it waits for becomes stopping", async () => {
+      // #5589: disposal can start while a migration is still on its way to registering. If the
+      // migration then fails and its command's exit is never confirmed, it never settles; the
+      // disposal drain must notice the stopping transition instead of waiting on it forever.
+      const migration = manager.beginMigration(testWorkspaceId);
+      expect(migration.admitted).toBe(true);
+      const disposal = manager.cleanup(testWorkspaceId).then(() => "finished");
+      await Bun.sleep(50);
+      migration.markStopping();
+      expect(await Promise.race([disposal, Bun.sleep(2_000).then(() => "waiting")])).toBe(
+        "finished"
+      );
+      // Removal still waits for it (and fails closed at its drain deadline).
+      const removal = manager
+        .cleanup(testWorkspaceId, { failClosedAfterDrainTimeout: true })
+        .then(() => "finished");
+      expect(await Promise.race([removal, Bun.sleep(200).then(() => "waiting")])).toBe("waiting");
+      migration[Symbol.dispose]();
+      expect(await removal).toBe("finished");
+    });
   });
 
   describe("cleanup with a hung spawn", () => {

src/node/services/backgroundProcessManager.ts

@@ -516,6 +516,14 @@ export class BackgroundProcessManager extends EventEmitter<BackgroundProcessMana
    * afterwards.
    */
   private readonly pendingAdmissions = new Map<string, Set<Promise<void>>>();
+  /**
+   * Pending entries of migrations whose command is being terminated (refused or failed): they
+   * never register, and they stay pending until the command's exit is confirmed, which may be
+   * never (#5589). Only cleanups that delete the checkout afterwards wait for them.
+   */
+  private readonly stoppingAdmissions = new WeakSet<Promise<void>>();
+  /** Resolved and replaced on each markStopping(), so a disposal drain already waiting re-checks. */
+  private stoppingMarked = Promise.withResolvers<void>();
   /**
    * Open admission seals per workspace (#4967), counted: while any is held, beginMigration() and
    * spawn() refuse. cleanup() holds one for its own duration; a removal or archive holds one
@@ -1456,11 +1464,25 @@ export class BackgroundProcessManager extends EventEmitter<BackgroundProcessMana
    * While the workspace is sealed (sealAdmissions) the migration is not `admitted`: the caller
    * must terminate the command as on a failed migration. It is still tracked until disposed, so
    * a running cleanup() also waits for that termination.
+   *
+   * `markStopping()` says the command is being terminated and will never register: a removal or
+   * archive cleanup keeps waiting for it (and fails closed at its deadline), but session disposal,
+   * which deletes nothing, does not (#5589).
    */
-  beginMigration(workspaceId: string): Disposable & { readonly admitted: boolean } {
+  beginMigration(
+    workspaceId: string
+  ): Disposable & { readonly admitted: boolean; markStopping(): void } {
     const admitted = !this.admissionSeals.has(workspaceId);
     const pending = this.trackPendingAdmission(workspaceId);
-    return { admitted, [Symbol.dispose]: () => pending[Symbol.dispose]() };
+    return {
+      admitted,
+      markStopping: () => {
+        this.stoppingAdmissions.add(pending.settled);
+        this.stoppingMarked.resolve();
+        this.stoppingMarked = Promise.withResolvers<void>();
+      },
+      [Symbol.dispose]: () => pending[Symbol.dispose](),
+    };
   }
 
   /**
@@ -1475,7 +1497,9 @@ export class BackgroundProcessManager extends EventEmitter<BackgroundProcessMana
   }
 
   /** Counts a command as pending in `workspaceId` until the returned handle is disposed. */
-  private trackPendingAdmission(workspaceId: string): Disposable {
+  private trackPendingAdmission(
+    workspaceId: string
+  ): Disposable & { readonly settled: Promise<void> } {
     const settled = Promise.withResolvers<void>();
     let pending = this.pendingAdmissions.get(workspaceId);
     if (pending === undefined) {
@@ -1484,6 +1508,7 @@ export class BackgroundProcessManager extends EventEmitter<BackgroundProcessMana
     }
     pending.add(settled.promise);
     return {
+      settled: settled.promise,
       [Symbol.dispose]: () => {
         const current = this.pendingAdmissions.get(workspaceId);
         current?.delete(settled.promise);
@@ -1518,6 +1543,9 @@ export class BackgroundProcessManager extends EventEmitter<BackgroundProcessMana
    * Waits until no migration or spawn is pending in `workspaceId`, including later ones. With a
    * `deadline` (epoch ms), throws once it passes (#5477), so a caller that deletes the checkout
    * afterwards (archive, removal) fails closed instead of hanging with a stuck runtime call.
+   * Without one (session disposal, which deletes nothing), it skips migrations whose command is
+   * stopping: their exit may never be confirmed, and waiting would hang disposal with its admission
+   * seal held (#5589). They stay pending, so removal and archive still wait for them.
    */
   private async drainPendingAdmissions(
     workspaceId: string,
@@ -1529,7 +1557,11 @@ export class BackgroundProcessManager extends EventEmitter<BackgroundProcessMana
         pending !== undefined;
         pending = this.pendingAdmissions.get(workspaceId)
       ) {
-        await Promise.all([...pending]);
+        const registering = [...pending].filter((entry) => !this.stoppingAdmissions.has(entry));
+        if (registering.length === 0) return;
+        // A migration may become stopping while this waits (disposal can start before a migration
+        // fails): wake up and re-check rather than wait on an entry that may never settle.
+        await Promise.race([Promise.all(registering), this.stoppingMarked.promise]);
       }
       return;
     }

src/node/services/backgroundProcessesFormalRepro.test.ts

@@ -428,8 +428,7 @@ describe("#5465 case 1: a refused migration whose command outlives the kill join
   ) {
     const ws = uniqueWorkspace(tag);
     const manager = new BackgroundProcessManager(path.dirname(localBgWorkspaceDir(ws)));
-    // A rejected exit observation keeps the migration pending for good, so cleanup() would hang.
-    if (!rejectExit) cleanups.push(() => manager.cleanup(ws));
+    cleanups.push(() => manager.cleanup(ws));
     const dir = await tempDir(tag);
     const pidFile = path.join(dir, "pid");
     const releaseKill = Promise.withResolvers<void>();
@@ -527,12 +526,24 @@ describe("#5465 case 1: a refused migration whose command outlives the kill join
     }, 20_000);
   }
 
-  test("a migration whose exit observation fails stays pending", async () => {
+  test("a migration whose exit is never confirmed blocks removal but not session disposal", async () => {
     const { manager, ws, releaseKill } = await refuseMigration("join-reject", true, false, true);
     // The kill takes effect, but the runtime reports an error instead of the exit: that does not
-    // confirm the stop, so a removal's cleanup must not finish.
+    // confirm the stop, so the migration stays pending.
     releaseKill();
-    const cleanup = manager.cleanup(ws).then(() => "finished");
-    expect(await Promise.race([cleanup, Bun.sleep(500).then(() => "waiting")])).toBe("waiting");
+    // Session disposal (agentSession.ts) deletes nothing: it finishes and releases its admission
+    // seal, so the workspace can send commands to the background again (#5589).
+    const disposal = manager.cleanup(ws).then(() => "finished");
+    expect(await Promise.race([disposal, Bun.sleep(2_000).then(() => "waiting")])).toBe("finished");
+    const later = manager.beginMigration(ws);
+    expect(later.admitted).toBe(true);
+    later[Symbol.dispose]();
+    // A removal's cleanup (workspaceService.ts) deletes the checkout once it returns, so it keeps
+    // waiting for the unconfirmed stop (and fails closed at its drain deadline).
+    const removal = manager.cleanup(ws, { failClosedAfterDrainTimeout: true }).then(
+      () => "finished",
+      () => "failed closed"
+    );
+    expect(await Promise.race([removal, Bun.sleep(500).then(() => "waiting")])).toBe("waiting");
   }, 20_000);
 });

src/node/services/tools/bash.ts

@@ -1587,8 +1587,10 @@ ${scriptWithEnv}`;
             // has not taken effect yet (#5522).
             // A rejected exitCode (e.g. a remote transport error) does not confirm the stop, so the
             // migration then stays pending for the session: removal and archive keep failing
-            // closed rather than deleting the checkout under a command that may still run.
+            // closed rather than deleting the checkout under a command that may still run. Session
+            // disposal, which deletes nothing, does not wait for a stopping migration (#5589).
             migrationHandedToExit = true;
+            migration?.markStopping();
             void execStream.exitCode.then(
               () => migration?.[Symbol.dispose](),
               () => undefined