From 5df0ed3b9f5682ae87c246bccada342267778e65 Mon Sep 17 00:00:00 2001
From: Andy Ye <35905412+TurboTheTurtle@users.noreply.github.com>
Date: Sun, 31 May 2026 14:00:37 -0700
Subject: [PATCH] fix(agents): publish owned announcement session writes

Forward prompt-submission owned session write publication into the embedded session lock controller so same-process announcement/completion writes can advance the requester fence while external edits still trigger takeover protection.

Adds regression coverage for a second controller publishing an owned announcement write and for preserving rejection of a later unowned edit.

Closes #88703.

Thanks @TurboTheTurtle.
---
 .../run/attempt.session-lock.test.ts          | 69 +++++++++++++++++++
 .../embedded-agent-runner/run/attempt.ts      |  3 +-
 2 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/src/agents/embedded-agent-runner/run/attempt.session-lock.test.ts b/src/agents/embedded-agent-runner/run/attempt.session-lock.test.ts
index a817741853af..2debb441dbe9 100644
--- a/src/agents/embedded-agent-runner/run/attempt.session-lock.test.ts
+++ b/src/agents/embedded-agent-runner/run/attempt.session-lock.test.ts
@@ -922,6 +922,75 @@ describe("embedded attempt session lock lifecycle", () => {
     expect(releases).toEqual(["release", "release", "release"]);
   });
 
+  it("allows prompt-stream announcement writes from another controller but still rejects external edits", async () => {
+    const sessionFile = await createTempSessionFile();
+    const acquireSessionWriteLockAnnouncement = vi.fn(async () => ({ release: vi.fn() }));
+    const firstController = await createEmbeddedAttemptSessionLockController({
+      acquireSessionWriteLock: acquireSessionWriteLockAnnouncement,
+      lockOptions: { ...lockOptions, sessionFile },
+    });
+
+    await firstController.releaseForPrompt();
+
+    const sessionKey = "agent:main:imessage:requester";
+    const secondController = await createEmbeddedAttemptSessionLockController({
+      acquireSessionWriteLock: acquireSessionWriteLockAnnouncement,
+      lockOptions: { ...lockOptions, sessionFile },
+    });
+    const forwardedOptions: Array<{ publishOwnedWrite?: boolean } | undefined> = [];
+    const announceSession = {
+      agent: {
+        streamFn: vi.fn(async () => {
+          await runWithOwnedSessionTranscriptWritePublication(
+            { sessionFile, sessionKey },
+            async () => {
+              await fs.appendFile(
+                sessionFile,
+                '{"type":"message","id":"announcement-complete"}\n',
+                "utf8",
+              );
+            },
+          );
+        }),
+      },
+    };
+
+    installPromptSubmissionLockRelease({
+      session: announceSession,
+      waitForSessionEvents: (sessionToDrain) =>
+        secondController.waitForSessionEvents(sessionToDrain),
+      releaseForPrompt: () => secondController.releaseForPrompt(),
+      reacquireAfterPrompt: () => secondController.reacquireAfterPrompt(),
+      sessionFile,
+      sessionKey,
+      withSessionWriteLock: (run, options) => {
+        forwardedOptions.push(options);
+        return secondController.withSessionWriteLock(run, options);
+      },
+    });
+
+    await announceSession.agent.streamFn();
+    await expect(
+      firstController.withSessionWriteLock(async () => {
+        await fs.appendFile(sessionFile, '{"type":"message","id":"post-announcement"}\n', "utf8");
+        return "post-announcement";
+      }),
+    ).resolves.toBe("post-announcement");
+    expect(firstController.hasSessionTakeover()).toBe(false);
+
+    await fs.appendFile(
+      sessionFile,
+      '{"type":"message","id":"external-after-announcement"}\n',
+      "utf8",
+    );
+    await expect(firstController.withSessionWriteLock(() => "late")).rejects.toBeInstanceOf(
+      EmbeddedAttemptSessionTakeoverError,
+    );
+
+    expect(firstController.hasSessionTakeover()).toBe(true);
+    expect(forwardedOptions).toContainEqual({ publishOwnedWrite: true });
+  });
+
   it("rejects external edits interleaved while another controller holds cleanup lock", async () => {
     const sessionFile = await createTempSessionFile();
     const releases: string[] = [];
diff --git a/src/agents/embedded-agent-runner/run/attempt.ts b/src/agents/embedded-agent-runner/run/attempt.ts
index b6325ac4305c..c51846922257 100644
--- a/src/agents/embedded-agent-runner/run/attempt.ts
+++ b/src/agents/embedded-agent-runner/run/attempt.ts
@@ -3838,7 +3838,8 @@ export async function runEmbeddedAttempt(
               reacquireAfterPrompt: () => sessionLockController.reacquireAfterPrompt(),
               sessionKey: params.sessionKey,
               sessionFile: params.sessionFile,
-              withSessionWriteLock: (run) => sessionLockController.withSessionWriteLock(run),
+              withSessionWriteLock: (run, options) =>
+                sessionLockController.withSessionWriteLock(run, options),
             });
           }