fix(agents): restore controller.dispose() and clear prompt turn on teardown

ubehera · ubehera · commit 2606e52d7441 · 2026-05-25T15:17:23.000-06:00
ClawSweeper review on the previous force-push found that the rebase that dropped the kesslerio drain commit also removed unrelated infrastructure: the EmbeddedAttemptSessionLockController.dispose() method, its outer- finally call site in runEmbeddedAttempt, and two existing test cases on main that exercise it. Without dispose() the eager session lock leaks on post-prompt error paths (the original #86014 fix). The new activePromptSessionTurn introduced by this PR makes the gap worse — teardown after releaseForPrompt without going through reacquireAfterPrompt leaves the file-scoped queue tail unresolved and the next same-file controller waits forever (bot finding, confidence 0.9). Restored: - EmbeddedAttemptSessionLockController.dispose(): Promise<void> back on the type and implementation. Behavior: vacate activePromptSessionTurn first (new), then release heldLock if held (the original main behavior). Idempotent — both heldLock and activePromptSessionTurn are cleared after first call. - releaseRetainedSessionLock infrastructure in runEmbeddedAttempt: declaration before the try block, assignment after the controller constructor, and the outer-finally invocation with the log.error fallback. Restored from main as it was before this branch's rebase. - The two existing dispose tests from main (#86014): "releases the eagerly-held attempt lock on dispose when cleanup is skipped" and "dispose does not double-release a lock already handed to cleanup". New regression for the prompt-turn cleanup invariant: - "dispose-after-releaseForPrompt vacates the prompt turn so the next same-file controller does not hang". Verifies the case the bot flagged: controller A calls releaseForPrompt, then dispose without going through reacquireAfterPrompt. Controller B on the same session file must be able to releaseForPrompt without hanging. The test races releaseForPrompt against a 200 ms timeout and asserts the release resolves first. 80 cases across 2 test files green. Closes #85913.
diff --git a/src/agents/pi-embedded-runner/run/attempt.session-lock.test.ts b/src/agents/pi-embedded-runner/run/attempt.session-lock.test.ts
@@ -1431,4 +1431,45 @@ describe("embedded attempt session lock lifecycle", () => {
     await controllerB.releaseForPrompt();
     expect(controllerB.hasSessionTakeover()).toBe(false);
   });
+
+  it("dispose-after-releaseForPrompt vacates the prompt turn so the next same-file controller does not hang", async () => {
+    // Regression for the ClawSweeper P2 finding on #86067: after
+    // releaseForPrompt() opens the prompt window, heldLock is undefined but
+    // activePromptSessionTurn is still set until reacquireAfterPrompt clears
+    // it. If teardown runs in that window (mid-stream throw, etc.), dispose
+    // must release+clear the prompt turn or the file-scoped queue tail
+    // never resolves and the next same-file controller waits forever.
+    resetEmbeddedAttemptSessionFilePromptGuardsForTest();
+    const sessionFile = await createTempSessionFile();
+    const acquireSessionWriteLock = vi.fn(async () => ({
+      release: vi.fn(async () => {}),
+    }));
+    const controllerA = await createEmbeddedAttemptSessionLockController({
+      acquireSessionWriteLock,
+      lockOptions: { ...lockOptions, sessionFile },
+    });
+
+    await controllerA.releaseForPrompt();
+    // Teardown without going through reacquireAfterPrompt — the case the
+    // outer-finally `releaseRetainedSessionLock?.()` covers when the prompt
+    // stream throws.
+    await controllerA.dispose();
+    // Idempotent: second dispose must be safe (the outer finally can run
+    // even after acquireForCleanup has handed off the lock).
+    await controllerA.dispose();
+
+    // A new controller on the same file must be able to open its own
+    // prompt window without hanging on A's vacated turn.
+    const controllerB = await createEmbeddedAttemptSessionLockController({
+      acquireSessionWriteLock,
+      lockOptions: { ...lockOptions, sessionFile },
+    });
+    await expect(
+      Promise.race([
+        controllerB.releaseForPrompt().then(() => "released"),
+        new Promise<string>((resolve) => setTimeout(() => resolve("timeout"), 200)),
+      ]),
+    ).resolves.toBe("released");
+    expect(controllerB.hasSessionTakeover()).toBe(false);
+  });
 });
diff --git a/src/agents/pi-embedded-runner/run/attempt.session-lock.ts b/src/agents/pi-embedded-runner/run/attempt.session-lock.ts
@@ -981,6 +981,18 @@ export async function createEmbeddedAttemptSessionLockController(params: {
       return takeoverDetected;
     },
     async dispose(): Promise<void> {
+      // Vacate the prompt turn first, regardless of whether heldLock is set.
+      // After `releaseForPrompt`, heldLock is undefined but
+      // activePromptSessionTurn is still set until `reacquireAfterPrompt`
+      // clears it. If teardown runs in that window without going through
+      // reacquireAfterPrompt (e.g. mid-stream throw before the outer
+      // finally hits acquireForCleanup), the file-scoped queue tail never
+      // resolves and later same-file controllers block on it forever.
+      // Release the turn here so dispose-after-releaseForPrompt is safe.
+      if (activePromptSessionTurn) {
+        activePromptSessionTurn.release();
+        activePromptSessionTurn = undefined;
+      }
       if (!heldLock) {
         return;
       }
