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 index 8b3c6cf1252..6c0258e87d5 100644 --- a/src/agents/pi-embedded-runner/run/attempt.session-lock.test.ts +++ b/src/agents/pi-embedded-runner/run/attempt.session-lock.test.ts @@ -128,6 +128,26 @@ describe("embedded attempt session lock lifecycle", () => { expect(events).toEqual(["prep-release", "yield-cleanup-write", "cleanup-release"]); }); + it("keeps the session fence active after releasing for sessions_yield abort cleanup", async () => { + const sessionFile = await createTempSessionFile(); + const release = vi.fn(async () => {}); + const acquireSessionWriteLock = vi.fn(async () => ({ release })); + const controller = await createEmbeddedAttemptSessionLockController({ + acquireSessionWriteLock, + lockOptions: { ...lockOptions, sessionFile }, + }); + + await controller.releaseHeldLockForAbort(); + await fs.appendFile(sessionFile, '{"type":"message","id":"abort-takeover"}\n', "utf8"); + + await expect(controller.withSessionWriteLock(() => "yield-cleanup")).rejects.toBeInstanceOf( + EmbeddedAttemptSessionTakeoverError, + ); + expect(controller.hasSessionTakeover()).toBe(true); + expect(acquireSessionWriteLock).toHaveBeenCalledTimes(2); + expect(release).toHaveBeenCalledTimes(2); + }); + it("runs post-prompt transcript writes under a short reacquired lock", async () => { const events: string[] = []; const acquireSessionWriteLock = vi diff --git a/src/agents/pi-embedded-runner/run/attempt.session-lock.ts b/src/agents/pi-embedded-runner/run/attempt.session-lock.ts index 670ee749576..9a09ca9c976 100644 --- a/src/agents/pi-embedded-runner/run/attempt.session-lock.ts +++ b/src/agents/pi-embedded-runner/run/attempt.session-lock.ts @@ -749,32 +749,31 @@ export async function createEmbeddedAttemptSessionLockController(params: { const noopLock: SessionLock = { release: async () => {} }; + async function releaseHeldLockWithFence(): Promise { + if (!heldLock) { + return; + } + const lock = heldLock; + heldLock = undefined; + const fingerprint = await readSessionFileFingerprint(params.lockOptions.sessionFile); + const ownedWrite = ownedSessionFileWrites.get(sessionFileFenceKey); + const trustedGeneration = trustSessionFileState(sessionFileFenceKey, fingerprint); + fenceFingerprint = fingerprint; + fenceSnapshot = await readSessionFileFenceSnapshot(params.lockOptions.sessionFile); + fenceGeneration = + ownedWrite && sameSessionFileFingerprint(ownedWrite.fingerprint, fingerprint) + ? ownedWrite.generation + : (trustedGeneration ?? fenceGeneration); + fenceActive = true; + await lock.release(); + } + return { async releaseForPrompt(): Promise { - if (!heldLock) { - return; - } - const lock = heldLock; - heldLock = undefined; - const fingerprint = await readSessionFileFingerprint(params.lockOptions.sessionFile); - const ownedWrite = ownedSessionFileWrites.get(sessionFileFenceKey); - const trustedGeneration = trustSessionFileState(sessionFileFenceKey, fingerprint); - fenceFingerprint = fingerprint; - fenceSnapshot = await readSessionFileFenceSnapshot(params.lockOptions.sessionFile); - fenceGeneration = - ownedWrite && sameSessionFileFingerprint(ownedWrite.fingerprint, fingerprint) - ? ownedWrite.generation - : (trustedGeneration ?? fenceGeneration); - fenceActive = true; - await lock.release(); + await releaseHeldLockWithFence(); }, async releaseHeldLockForAbort(): Promise { - if (!heldLock) { - return; - } - const lock = heldLock; - heldLock = undefined; - await lock.release(); + await releaseHeldLockWithFence(); }, refreshAfterOwnedSessionWrite(): void { if (fenceActive && !takeoverDetected) {