fix(agents): fence yield abort lock release

steipete · steipete · commit 73bd1b1f46f9 · 2026-05-27T02:10:08.000+01:00
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
@@ -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
@@ -749,32 +749,31 @@ export async function createEmbeddedAttemptSessionLockController(params: {
 
   const noopLock: SessionLock = { release: async () => {} };
 
+  async function releaseHeldLockWithFence(): Promise<void> {
+    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<void> {
-      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<void> {
-      if (!heldLock) {
-        return;
-      }
-      const lock = heldLock;
-      heldLock = undefined;
-      await lock.release();
+      await releaseHeldLockWithFence();
     },
     refreshAfterOwnedSessionWrite(): void {
       if (fenceActive && !takeoverDetected) {
