diff --git a/src/gateway/server.sessions.create.test.ts b/src/gateway/server.sessions.create.test.ts index 17f29d063fe1..fe3af8d21f15 100644 --- a/src/gateway/server.sessions.create.test.ts +++ b/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 diff --git a/src/gateway/session-create-service.ts b/src/gateway/session-create-service.ts index ac0928d03e30..a30817bd30dc 100644 --- a/src/gateway/session-create-service.ts +++ b/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) => - 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({ diff --git a/src/sessions/session-lifecycle-admission.test.ts b/src/sessions/session-lifecycle-admission.test.ts index 6b92ca0ce84e..f56e860e2ee5 100644 --- a/src/sessions/session-lifecycle-admission.test.ts +++ b/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(); diff --git a/src/sessions/session-lifecycle-admission.ts b/src/sessions/session-lifecycle-admission.ts index 9d1a72237e89..0187ca3fecb4 100644 --- a/src/sessions/session-lifecycle-admission.ts +++ b/src/sessions/session-lifecycle-admission.ts @@ -40,6 +40,18 @@ type SessionLifecycleAdmissionState = { type SessionLifecycleMutationKind = "compaction"; +type SessionLifecycleMutationTarget = { + scope: string; + identities: Iterable; +}; + +type SessionLifecycleMutationParams = { + kind?: SessionLifecycleMutationKind; + prepare?: () => Promise; + run: () => Promise; + signal?: AbortSignal; +} & (SessionLifecycleMutationTarget | { targets: Iterable }); + // 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( identities: readonly string[], index: number, run: () => Promise, + kind: "lifecycle" | "mutation" = "lifecycle", ): Promise { 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( - identities: readonly string[], - index: number, - run: () => Promise, -): Promise { - 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(params: { } } -export async function runExclusiveSessionLifecycleMutation(params: { - scope: string; - identities: Iterable; - kind?: SessionLifecycleMutationKind; - prepare?: () => Promise; - run: () => Promise; - signal?: AbortSignal; -}): Promise { - const identities = normalizeSessionIdentities(params.scope, params.identities); +export async function runExclusiveSessionLifecycleMutation( + params: SessionLifecycleMutationParams, +): Promise { + // 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(params: { }); } }), + "mutation", ); if (!signal) { return await mutation;