From 58d08dedd39b849ef8093d12ca8f7c4a0cf28226 Mon Sep 17 00:00:00 2001
From: Vincent Koc <vincentkoc@ieee.org>
Date: Fri, 7 Aug 2026 09:11:23 +0200
Subject: [PATCH] fix(acp): capture caller lifecycle before setup

---
 src/auto-reply/reply/dispatch-acp.test.ts | 67 +++++++++++++++++++++++
 src/auto-reply/reply/dispatch-acp.ts      | 27 ++++++++-
 2 files changed, 91 insertions(+), 3 deletions(-)

diff --git a/src/auto-reply/reply/dispatch-acp.test.ts b/src/auto-reply/reply/dispatch-acp.test.ts
index b933ee8e9215..5237442dc4f4 100644
--- a/src/auto-reply/reply/dispatch-acp.test.ts
+++ b/src/auto-reply/reply/dispatch-acp.test.ts
@@ -788,6 +788,73 @@ describe("tryDispatchAcpReply", () => {
     clearAgentRunContext(runId, replacementGeneration);
   });
 
+  it("captures caller-owned ACP lifecycle before lazy dispatch setup yields", async () => {
+    setReadyAcpResolution();
+    const runId = "caller-reused-during-setup";
+    const admittedGeneration = getAgentEventLifecycleGeneration();
+    claimAgentRunContext(runId, {
+      agentId: "caller",
+      lifecycleGeneration: admittedGeneration,
+      sessionKey,
+    });
+    const admittedToken = getAgentRunContextLifecycleToken(runId, admittedGeneration);
+    let startGeneration: string | undefined;
+    let startToken: object | undefined;
+    let callbackGeneration: string | undefined;
+    let callbackToken: object | undefined;
+    let terminalGeneration: string | undefined;
+    let terminalToken: object | undefined;
+
+    auditMocks.emitAcpLifecycleStart.mockImplementationOnce(
+      (params: { lifecycleGeneration?: string }) => {
+        startGeneration = params.lifecycleGeneration;
+        startToken = getAgentRunExecutionContextLifecycleToken(runId);
+      },
+    );
+    auditMocks.emitAcpRuntimeEvent.mockImplementationOnce(() => {
+      callbackGeneration = getAgentRunExecutionLifecycleGeneration();
+      callbackToken = getAgentRunExecutionContextLifecycleToken(runId);
+    });
+    auditMocks.emitAcpLifecycleEnd.mockImplementationOnce(() => {
+      terminalGeneration = getAgentRunExecutionLifecycleGeneration();
+      terminalToken = getAgentRunExecutionContextLifecycleToken(runId);
+    });
+    managerMocks.runTurn.mockImplementationOnce(
+      async ({ onEvent }: { onEvent?: (event: unknown) => Promise<void> }) => {
+        if (onEvent) {
+          await emitToolLifecycleEvents(onEvent, "tool-after-setup-reuse");
+        }
+      },
+    );
+
+    const dispatch = runDispatch({ bodyForAgent: "audit setup race", runId });
+    const replacementGeneration = rotateAgentEventLifecycleGeneration();
+    claimAgentRunContext(runId, {
+      agentId: "replacement",
+      lifecycleGeneration: replacementGeneration,
+      sessionKey: "agent:replacement:main",
+    });
+    const replacementToken = getAgentRunContextLifecycleToken(runId, replacementGeneration);
+
+    await dispatch;
+
+    expect(admittedToken).toBeTypeOf("object");
+    expect(replacementToken).toBeTypeOf("object");
+    expect(replacementToken).not.toBe(admittedToken);
+    expect(startGeneration).toBe(admittedGeneration);
+    expect(startToken).toBe(admittedToken);
+    expect(callbackGeneration).toBe(admittedGeneration);
+    expect(callbackToken).toBe(admittedToken);
+    expect(terminalGeneration).toBe(admittedGeneration);
+    expect(terminalToken).toBe(admittedToken);
+    expect(getAgentRunContext(runId)).toMatchObject({
+      agentId: "replacement",
+      lifecycleGeneration: replacementGeneration,
+      sessionKey: "agent:replacement:main",
+    });
+    clearAgentRunContext(runId, replacementGeneration);
+  });
+
   it("keeps audit run ids unique when channel message ids repeat", async () => {
     setReadyAcpResolution();
 
diff --git a/src/auto-reply/reply/dispatch-acp.ts b/src/auto-reply/reply/dispatch-acp.ts
index 9a992ca6c8b1..f45d13c30daa 100644
--- a/src/auto-reply/reply/dispatch-acp.ts
+++ b/src/auto-reply/reply/dispatch-acp.ts
@@ -22,8 +22,10 @@ import {
   captureAgentRunLifecycleGeneration,
   withAgentRunLifecycleGeneration,
 } from "../../infra/agent-events.js";
+import { captureAgentRunExecutionContextLifecycleToken } from "../../infra/agent-run-execution-context.js";
 import {
   claimAgentRunContext,
+  getAgentRunContextLifecycleToken,
   releaseAgentRunContext,
   retainActiveAgentRunContext,
 } from "../../infra/agent-run-registry.js";
@@ -455,6 +457,17 @@ export async function tryDispatchAcpReply(params: {
     return null;
   }
 
+  const existingRunId = normalizeOptionalString(params.runId);
+  // Lazy ACP setup can yield while a caller-owned run id is rotated and reused.
+  // Snapshot both ownership identities now so later audit callbacks cannot rebound.
+  const admittedAuditLifecycleGeneration = existingRunId
+    ? captureAgentRunLifecycleGeneration(existingRunId)
+    : undefined;
+  const admittedAuditContextLifecycleToken =
+    existingRunId && admittedAuditLifecycleGeneration
+      ? getAgentRunContextLifecycleToken(existingRunId, admittedAuditLifecycleGeneration)
+      : undefined;
+
   const { getAcpSessionManager } = await loadDispatchAcpManagerRuntime();
   const acpManager = getAcpSessionManager();
   const acpResolution = acpManager.resolveSession({
@@ -581,14 +594,13 @@ export async function tryDispatchAcpReply(params: {
       markIdle: params.markIdle,
     });
   const requestId = resolveAcpRequestId(params.ctx);
-  const existingRunId = normalizeOptionalString(params.runId);
   const auditOnly = existingRunId === undefined;
   const auditRunId = existingRunId ?? generateSecureUuid();
   const auditRuntime = await loadDispatchAcpAuditRuntime();
   const auditToolTracker = auditRuntime.createAcpToolLifecycleTracker();
   let auditStarted = false;
   let auditFinished = false;
-  let auditLifecycleGeneration: string | undefined;
+  let auditLifecycleGeneration = admittedAuditLifecycleGeneration;
   let auditContextOwnerToken: string | undefined;
   let releaseAuditContextLease: (() => void) | undefined;
   let auditTerminalOutcome: "blocked" | undefined;
@@ -639,7 +651,16 @@ export async function tryDispatchAcpReply(params: {
   const runWithAuditLifecycle = <T>(run: () => T): T => {
     claimAuditContext();
     return auditLifecycleGeneration
-      ? withAgentRunLifecycleGeneration(auditLifecycleGeneration, run)
+      ? withAgentRunLifecycleGeneration(auditLifecycleGeneration, () => {
+          if (admittedAuditContextLifecycleToken) {
+            captureAgentRunExecutionContextLifecycleToken(
+              auditRunId,
+              auditLifecycleGeneration,
+              admittedAuditContextLifecycleToken,
+            );
+          }
+          return run();
+        })
       : run();
   };
   const releaseAuditContext = () => {