fix(sessions): prevent cross-agent lifecycle deadlocks (#114524)

Co-authored-by: Peter Steinberger <steipete@gmail.com>

src/gateway/server.sessions.create.test.ts

@@ -3119,6 +3119,138 @@ test("sessions.create resolves an agent-qualified fork from the parent store", a
   }
 });
 
+test("sessions.create completes simultaneous opposite-direction cross-agent forks", async () => {
+  const { dir } = await createSessionStoreDir();
+  const storeTemplate = path.join(dir, "{agentId}", "sessions.json");
+  const mainStorePath = storeTemplate.replace("{agentId}", "main");
+  const workStorePath = storeTemplate.replace("{agentId}", "work");
+  testState.sessionStorePath = storeTemplate;
+  testState.sessionConfig = { scope: "per-sender" };
+  testState.agentsConfig = { list: [{ id: "main", default: true }, { id: "work" }] };
+
+  try {
+    const mainDir = path.dirname(mainStorePath);
+    const workDir = path.dirname(workStorePath);
+    await Promise.all([
+      fs.mkdir(mainDir, { recursive: true }),
+      fs.mkdir(workDir, { recursive: true }),
+    ]);
+    const [mainParent, workParent] = await Promise.all([
+      createCheckpointFixture(mainDir),
+      createCheckpointFixture(workDir),
+    ]);
+    await Promise.all([
+      writeSessionStore({
+        storePath: mainStorePath,
+        agentId: "main",
+        entries: {
+          main: sessionStoreEntry(mainParent.sessionId, {
+            sessionFile: mainParent.sessionFile,
+          }),
+        },
+      }),
+      writeSessionStore({
+        storePath: workStorePath,
+        agentId: "work",
+        entries: {
+          main: sessionStoreEntry(workParent.sessionId, {
+            sessionFile: workParent.sessionFile,
+          }),
+        },
+      }),
+    ]);
+    await Promise.all([
+      seedSessionTranscript({
+        agentId: "main",
+        sessionId: mainParent.sessionId,
+        sessionKey: "agent:main:main",
+        storePath: mainStorePath,
+        messages: [{ role: "user", content: "main parent context" }],
+      }),
+      seedSessionTranscript({
+        agentId: "work",
+        sessionId: workParent.sessionId,
+        sessionKey: "agent:work:main",
+        storePath: workStorePath,
+        messages: [{ role: "user", content: "work parent context" }],
+      }),
+    ]);
+
+    const requests = Array.from({ length: 12 }, (_, index) =>
+      index % 2 === 0
+        ? {
+            agentId: "main",
+            parentSessionKey: "agent:work:main",
+            parentSessionId: workParent.sessionId,
+            storePath: mainStorePath,
+          }
+        : {
+            agentId: "work",
+            parentSessionKey: "agent:main:main",
+            parentSessionId: mainParent.sessionId,
+            storePath: workStorePath,
+          },
+    );
+    const created = await Promise.all(
+      requests.map((request) =>
+        directSessionReq<{
+          key: string;
+          sessionId: string;
+          entry: {
+            parentSessionKey?: string;
+            forkSource?: { sessionKey: string; sessionId: string };
+            forkedFromParent?: boolean;
+          };
+        }>("sessions.create", {
+          agentId: request.agentId,
+          parentSessionKey: request.parentSessionKey,
+          fork: true,
+        }),
+      ),
+    );
+
+    expect(
+      created.every((result) => result.ok),
+      JSON.stringify(created.filter((result) => !result.ok)),
+    ).toBe(true);
+    expect(new Set(created.map((result) => result.payload?.key)).size).toBe(requests.length);
+    expect(new Set(created.map((result) => result.payload?.sessionId)).size).toBe(requests.length);
+    for (const [index, result] of created.entries()) {
+      const request = requests[index];
+      if (!request) {
+        throw new Error(`missing cross-agent fork request ${index}`);
+      }
+      expect(result.payload?.entry).toMatchObject({
+        forkSource: {
+          sessionKey: request.parentSessionKey,
+          sessionId: request.parentSessionId,
+        },
+        forkedFromParent: true,
+        parentSessionKey: request.parentSessionKey,
+      });
+      const key = requireNonEmptyString(result.payload?.key, "cross-agent fork session key");
+      expect(
+        loadSessionEntry({
+          agentId: request.agentId,
+          sessionKey: key,
+          storePath: request.storePath,
+        }),
+      ).toMatchObject({
+        forkSource: {
+          sessionKey: request.parentSessionKey,
+          sessionId: request.parentSessionId,
+        },
+        parentSessionKey: request.parentSessionKey,
+        sessionId: result.payload?.sessionId,
+      });
+    }
+  } finally {
+    testState.sessionStorePath = undefined;
+    testState.sessionConfig = undefined;
+    testState.agentsConfig = undefined;
+  }
+});
+
 test("sessions.create can start the first agent turn from an initial task", async () => {
   await createSessionStoreDir();
   // Register "ops" so the deleted-agent guard added in #65986 does not

src/gateway/session-create-service.ts

@@ -1154,14 +1154,12 @@ export async function createGatewaySession(params: {
     };
   };
 
-  const runWithCreationTargetLock = async () =>
-    await runExclusiveSessionLifecycleMutation({
+  const lifecycleTargets = [
+    {
       scope: creationTarget.storePath,
       identities: [creationTarget.canonicalKey],
-      run: createChildSession,
-    });
-
-  let result: CreateGatewaySessionResult;
+    },
+  ];
   if (
     canonicalParentSessionKey &&
     parentSessionEntry?.sessionId &&
@@ -1170,39 +1168,17 @@ export async function createGatewaySession(params: {
       params.fork === true ||
       params.authorizedPluginId !== undefined)
   ) {
-    if (parentSessionTarget.storePath === creationTarget.storePath) {
-      result = await runExclusiveSessionLifecycleMutation({
-        scope: creationTarget.storePath,
-        identities: [
-          canonicalParentSessionKey,
-          parentSessionEntry.sessionId,
-          creationTarget.canonicalKey,
-        ],
-        run: createChildSession,
-      });
-    } else {
-      const runWithParentLock = async (run: () => Promise<CreateGatewaySessionResult>) =>
-        await runExclusiveSessionLifecycleMutation({
-          scope: parentSessionTarget.storePath,
-          identities: [canonicalParentSessionKey, parentSessionEntry.sessionId],
-          run,
-        });
-      // Cross-agent forks touch two stores. Acquire their locks in canonical
-      // store order so simultaneous opposite-direction forks cannot deadlock.
-      result =
-        parentSessionTarget.storePath < creationTarget.storePath
-          ? await runWithParentLock(runWithCreationTargetLock)
-          : await runExclusiveSessionLifecycleMutation({
-              scope: creationTarget.storePath,
-              identities: [creationTarget.canonicalKey],
-              run: async () => await runWithParentLock(createChildSession),
-            });
-    }
-  } else {
-    // Keyed creates must observe and adopt the winning row under the same
-    // lifecycle fence; otherwise concurrent callers mint divergent session IDs.
-    result = await runWithCreationTargetLock();
+    lifecycleTargets.push({
+      scope: parentSessionTarget.storePath,
+      identities: [canonicalParentSessionKey, parentSessionEntry.sessionId],
+    });
   }
+  // Generated, keyed, same-store, and cross-agent creations all share the
+  // lifecycle owner's canonical identity order and one active mutation fence.
+  const result = await runExclusiveSessionLifecycleMutation({
+    targets: lifecycleTargets,
+    run: createChildSession,
+  });
   if (result.ok && !result.resetExisting && createdContext) {
     if (createdNewEntry) {
       recordSessionCreated({

src/sessions/session-lifecycle-admission.test.ts

@@ -15,6 +15,7 @@ import {
   getActiveSessionWorkAdmissionCount,
   hasOnlySessionLifecycleMutationKindActive,
   interruptSessionWorkAdmissions,
+  isSessionLifecycleMutationActive,
   isSessionWorkAdmissionActive,
   runExclusiveSessionLifecycleMutation,
 } from "./session-lifecycle-admission.js";
@@ -141,6 +142,185 @@ it("counts one multi-identity lifecycle mutation once across module instances",
   expect(second.getActiveSessionLifecycleMutationCount()).toBe(0);
 });
 
+it("counts a cross-store lifecycle mutation once and fences every target", async () => {
+  const mutationStarted = createDeferred();
+  const releaseMutation = createDeferred();
+  const mutation = runExclusiveSessionLifecycleMutation({
+    targets: [
+      {
+        scope: "store-cross-count-b",
+        identities: ["agent:work:main", "session-cross-count-b"],
+      },
+      {
+        scope: "store-cross-count-a",
+        identities: ["agent:main:main", "session-cross-count-a"],
+      },
+      {
+        scope: "store-cross-count-a",
+        identities: ["session-cross-count-a", undefined],
+      },
+    ],
+    run: async () => {
+      mutationStarted.resolve();
+      await releaseMutation.promise;
+    },
+  });
+  await mutationStarted.promise;
+
+  try {
+    expect(getActiveSessionLifecycleMutationCount()).toBe(1);
+    expect(isSessionLifecycleMutationActive("store-cross-count-a", ["agent:main:main"])).toBe(true);
+    expect(isSessionLifecycleMutationActive("store-cross-count-b", ["session-cross-count-b"])).toBe(
+      true,
+    );
+    expect(isSessionLifecycleMutationActive("store-cross-count-b", ["agent:main:main"])).toBe(
+      false,
+    );
+  } finally {
+    releaseMutation.resolve();
+    await mutation;
+  }
+
+  expect(getActiveSessionLifecycleMutationCount()).toBe(0);
+  expect(isSessionLifecycleMutationActive("store-cross-count-a", ["session-cross-count-a"])).toBe(
+    false,
+  );
+  expect(isSessionLifecycleMutationActive("store-cross-count-b", ["session-cross-count-b"])).toBe(
+    false,
+  );
+});
+
+it("serializes opposite-direction cross-store lifecycle mutations", async () => {
+  const main = {
+    scope: "store-cross-order-a",
+    identities: ["agent:main:main", "session-cross-order-a"],
+  };
+  const work = {
+    scope: "store-cross-order-b",
+    identities: ["agent:work:main", "session-cross-order-b"],
+  };
+  let activeMutations = 0;
+  let maximumActiveMutations = 0;
+  let completedMutations = 0;
+
+  await Promise.all(
+    Array.from({ length: 48 }, async (_, index) =>
+      runExclusiveSessionLifecycleMutation({
+        targets: index % 2 === 0 ? [main, work] : [work, main],
+        run: async () => {
+          activeMutations += 1;
+          maximumActiveMutations = Math.max(maximumActiveMutations, activeMutations);
+          try {
+            await Promise.resolve();
+            completedMutations += 1;
+          } finally {
+            activeMutations -= 1;
+          }
+        },
+      }),
+    ),
+  );
+
+  expect(completedMutations).toBe(48);
+  expect(maximumActiveMutations).toBe(1);
+  expect(activeMutations).toBe(0);
+  expect(getActiveSessionLifecycleMutationCount()).toBe(0);
+});
+
+it("interrupts admitted work in both stores before a cross-store mutation", async () => {
+  const mainTarget = {
+    scope: "store-cross-interrupt-a",
+    identities: ["agent:main:main", "session-cross-interrupt-a"],
+  };
+  const workTarget = {
+    scope: "store-cross-interrupt-b",
+    identities: ["agent:work:main", "session-cross-interrupt-b"],
+  };
+  let mainInterrupted = false;
+  let workInterrupted = false;
+  const mainAdmission = await beginSessionWorkAdmission({
+    ...mainTarget,
+    assertAllowed: () => {},
+    onInterrupt: () => {
+      mainInterrupted = true;
+      mainAdmission.release();
+    },
+  });
+  const workAdmission = await beginSessionWorkAdmission({
+    ...workTarget,
+    assertAllowed: () => {},
+    onInterrupt: () => {
+      workInterrupted = true;
+      workAdmission.release();
+    },
+  });
+
+  try {
+    await runExclusiveSessionLifecycleMutation({
+      targets: [workTarget, mainTarget],
+      prepare: async () => {
+        const interrupted = await Promise.all([
+          interruptSessionWorkAdmissions({ ...mainTarget, timeoutMs: 1_000 }),
+          interruptSessionWorkAdmissions({ ...workTarget, timeoutMs: 1_000 }),
+        ]);
+        expect(interrupted).toEqual([true, true]);
+      },
+      run: async () => {
+        expect(mainInterrupted).toBe(true);
+        expect(workInterrupted).toBe(true);
+        expect(isSessionWorkAdmissionActive(mainTarget.scope, mainTarget.identities)).toBe(false);
+        expect(isSessionWorkAdmissionActive(workTarget.scope, workTarget.identities)).toBe(false);
+      },
+    });
+  } finally {
+    mainAdmission.release();
+    workAdmission.release();
+  }
+});
+
+it("cancels an opposite-direction cross-store mutation before activation", async () => {
+  const main = {
+    scope: "store-cross-cancel-a",
+    identities: ["agent:main:main", "session-cross-cancel-a"],
+  };
+  const work = {
+    scope: "store-cross-cancel-b",
+    identities: ["agent:work:main", "session-cross-cancel-b"],
+  };
+  const mutationStarted = createDeferred();
+  const releaseMutation = createDeferred();
+  const first = runExclusiveSessionLifecycleMutation({
+    targets: [main, work],
+    run: async () => {
+      mutationStarted.resolve();
+      await releaseMutation.promise;
+    },
+  });
+  await mutationStarted.promise;
+
+  const controller = new AbortController();
+  const abortError = new Error("cancel queued cross-store lifecycle mutation");
+  let cancelledMutationRan = false;
+  const cancelled = runExclusiveSessionLifecycleMutation({
+    targets: [work, main],
+    signal: controller.signal,
+    run: async () => {
+      cancelledMutationRan = true;
+    },
+  });
+  controller.abort(abortError);
+
+  try {
+    await expect(cancelled).rejects.toBe(abortError);
+  } finally {
+    releaseMutation.resolve();
+    await first;
+  }
+
+  expect(cancelledMutationRan).toBe(false);
+  expect(getActiveSessionLifecycleMutationCount()).toBe(0);
+});
+
 it("rejects an admission that resumes after suspension closes the async gap", async () => {
   resetGatewayWorkAdmission();
   const mutationStarted = createDeferred();

src/sessions/session-lifecycle-admission.ts

@@ -40,6 +40,18 @@ type SessionLifecycleAdmissionState = {
 
 type SessionLifecycleMutationKind = "compaction";
 
+type SessionLifecycleMutationTarget = {
+  scope: string;
+  identities: Iterable<string | undefined>;
+};
+
+type SessionLifecycleMutationParams<T> = {
+  kind?: SessionLifecycleMutationKind;
+  prepare?: () => Promise<void>;
+  run: () => Promise<T>;
+  signal?: AbortSignal;
+} & (SessionLifecycleMutationTarget | { targets: Iterable<SessionLifecycleMutationTarget> });
+
 // Runtime chunks can load separate module instances while still coordinating
 // the same sessions. One shared state keeps every lock and admission visible.
 const SESSION_LIFECYCLE_ADMISSION_STATE = resolveGlobalSingleton(
@@ -72,35 +84,19 @@ async function runWithSessionIdentityLocks<T>(
   identities: readonly string[],
   index: number,
   run: () => Promise<T>,
+  kind: "lifecycle" | "mutation" = "lifecycle",
 ): Promise<T> {
   const identity = identities[index];
   if (!identity) {
     return await run();
   }
   return await runQueuedStoreWrite({
-    queues: SESSION_LIFECYCLE_QUEUES,
+    queues: kind === "mutation" ? SESSION_LIFECYCLE_MUTATION_QUEUES : SESSION_LIFECYCLE_QUEUES,
     storePath: identity,
-    label: "runExclusiveSessionLifecycle",
+    label:
+      kind === "mutation" ? "runExclusiveSessionLifecycleMutation" : "runExclusiveSessionLifecycle",
     reentrant: true,
-    fn: async () => await runWithSessionIdentityLocks(identities, index + 1, run),
-  });
-}
-
-async function runWithSessionMutationIdentityLocks<T>(
-  identities: readonly string[],
-  index: number,
-  run: () => Promise<T>,
-): Promise<T> {
-  const identity = identities[index];
-  if (!identity) {
-    return await run();
-  }
-  return await runQueuedStoreWrite({
-    queues: SESSION_LIFECYCLE_MUTATION_QUEUES,
-    storePath: identity,
-    label: "runExclusiveSessionLifecycleMutation",
-    reentrant: true,
-    fn: async () => await runWithSessionMutationIdentityLocks(identities, index + 1, run),
+    fn: async () => await runWithSessionIdentityLocks(identities, index + 1, run, kind),
   });
 }
 
@@ -195,22 +191,29 @@ async function runExclusiveSessionLifecycle<T>(params: {
   }
 }
 
-export async function runExclusiveSessionLifecycleMutation<T>(params: {
-  scope: string;
-  identities: Iterable<string | undefined>;
-  kind?: SessionLifecycleMutationKind;
-  prepare?: () => Promise<void>;
-  run: () => Promise<T>;
-  signal?: AbortSignal;
-}): Promise<T> {
-  const identities = normalizeSessionIdentities(params.scope, params.identities);
+export async function runExclusiveSessionLifecycleMutation<T>(
+  params: SessionLifecycleMutationParams<T>,
+): Promise<T> {
+  // Normalize every store and session into one globally ordered identity set.
+  // Cross-agent mutations then acquire one fence and count as one active run,
+  // instead of nesting store locks in caller-selected order.
+  const identities =
+    "targets" in params
+      ? Array.from(
+          new Set(
+            Array.from(params.targets, (target) =>
+              normalizeSessionIdentities(target.scope, target.identities),
+            ).flat(),
+          ),
+        ).toSorted()
+      : normalizeSessionIdentities(params.scope, params.identities);
   const signal = params.signal;
   signal?.throwIfAborted();
   const callerAdmissions = new Set(CURRENT_SESSION_WORK_ADMISSIONS.getStore());
   const mutationRun = {};
   let mutationActivated = false;
   let removeAbortListener = () => {};
-  const mutation = runWithSessionMutationIdentityLocks(
+  const mutation = runWithSessionIdentityLocks(
     identities,
     0,
     async () =>
@@ -268,6 +271,7 @@ export async function runExclusiveSessionLifecycleMutation<T>(params: {
           });
         }
       }),
+    "mutation",
   );
   if (!signal) {
     return await mutation;