diff --git a/server/src/services/native-runtime/native-session-executor.test.ts b/server/src/services/native-runtime/native-session-executor.test.ts index 0dd73006c1d..8090b3cab6f 100644 --- a/server/src/services/native-runtime/native-session-executor.test.ts +++ b/server/src/services/native-runtime/native-session-executor.test.ts @@ -10116,6 +10116,69 @@ describe("native warm session supervision", () => { await closeWarmNativeSessionsForEnvironment({ environmentId: f.target.environmentId, reason: "test" }); expect(nextCopy.collectStopped).toHaveBeenCalledOnce(); }); + it("retires its own reserved directory owner when the session identity rotates", async () => { + // A follow-up run that cannot resume its saved session (for example after + // a configuration change) reserves the prior owner's directory, then + // starts a new session identity. Retiring that prior owner is this run's + // own reservation, not another executor's busy turn. + const name = "managed-identity-rotation"; + const f = await start(name); + const reservation = await f.reserve(async () => true); + expect(reservation?.reuseRunId).toBe(f.current.binding.runId); + const nextCopy = { ...f.copy, runId: `${name}-two`, collectStopped: vi.fn(async () => undefined) }; + reservation!.adopt(nextCopy.root, nextCopy.collectStopped); + const replacement = { close: vi.fn(async () => undefined) }; + state.execute.mockImplementationOnce(async options => { + expect(options.existingSession).toBeUndefined(); + expect(f.close).toHaveBeenCalledOnce(); + expect(nextCopy.collectStopped).not.toHaveBeenCalled(); + await options.onSession?.(replacement); return result; + }); + const next = { ...f.current, binding: { ...f.current.binding, runId: nextCopy.runId }, + session: { ...f.current.session, normalizedSessionId: `${name}-rotated` } } as NativeExecutionInputV1; + await expect(executePaperclipNativeSession({ db: leaseDb(next), execution: next, runnerInstanceId: name, + runnerExecutionTarget: f.target, instructionWorkingCopy: nextCopy })).resolves.toMatchObject({ exitCode: 0 }); + expect(state.execute).toHaveBeenCalledTimes(2); + await reservation!.release(); + expect(replacement.close).not.toHaveBeenCalled(); + expect(f.copy.collectStopped).not.toHaveBeenCalled(); + expect(nextCopy.collectStopped).not.toHaveBeenCalled(); + await closeWarmNativeSessionsForEnvironment({ environmentId: f.target.environmentId, reason: "test" }); + expect(replacement.close).toHaveBeenCalledOnce(); + expect(nextCopy.collectStopped).toHaveBeenCalledOnce(); + }); + it("retires its own reserved directory owner after the Runner claim when the session identity rotates", async () => { + // Production Runner path: retirement happens after the execution claim. + // This fixture has no remote Runner, so the turn itself fails later; the + // regression is that admission no longer rejects its own reservation. + const name = "managed-identity-rotation-runnerd"; + const f = await start(name); + const reservation = await f.reserve(async () => true); + const nextCopy = { ...f.copy, runId: `${name}-two`, collectStopped: vi.fn(async () => undefined) }; + reservation!.adopt(nextCopy.root, nextCopy.collectStopped); + const next = { ...f.current, binding: { ...f.current.binding, runId: nextCopy.runId }, + session: { ...f.current.session, normalizedSessionId: `${name}-rotated` } } as NativeExecutionInputV1; + const outcome = await executePaperclipNativeSession({ db: leaseDb(next), execution: next, runnerInstanceId: name, + runnerExecutionTarget: f.target, instructionWorkingCopy: nextCopy, useRunnerd: true }).then(() => null, (error: Error) => error); + expect(outcome?.message ?? "").not.toContain("native_session_supervisor_busy"); + expect(f.close).toHaveBeenCalledOnce(); + expect(f.copy.collectStopped).not.toHaveBeenCalled(); + expect(nextCopy.collectStopped).not.toHaveBeenCalled(); + await reservation!.release(); + }); + it("still rejects a session identity rotation over another run's directory reservation", async () => { + const f = await start("managed-identity-foreign"); + const reservation = await f.reserve(async () => true); + const next = { ...f.current, binding: { ...f.current.binding, runId: "managed-identity-foreign-three" }, + session: { ...f.current.session, normalizedSessionId: "managed-identity-foreign-rotated" } } as NativeExecutionInputV1; + state.execute.mockClear(); + await expect(executePaperclipNativeSession({ db: leaseDb(next), execution: next, runnerInstanceId: "foreign", + runnerExecutionTarget: f.target, instructionWorkingCopy: { ...f.copy, runId: next.binding.runId } })) + .rejects.toThrow("native_session_supervisor_busy"); + expect(state.execute).not.toHaveBeenCalled(); + expect(f.close).not.toHaveBeenCalled(); + await reservation!.release(); + }); it("cannot reuse a provider with a different registered directory", async () => { const f = await start("managed-root-replaced"); state.execute.mockImplementationOnce(async options => { expect(options.existingSession).toBeUndefined(); return result; }); diff --git a/server/src/services/native-runtime/native-session-executor.ts b/server/src/services/native-runtime/native-session-executor.ts index 1d1febfb4f6..5d183df60ae 100644 --- a/server/src/services/native-runtime/native-session-executor.ts +++ b/server/src/services/native-runtime/native-session-executor.ts @@ -2015,6 +2015,7 @@ async function retireSupersededWarmNativeSessions( ownerScope: string, nextSessionScope: string, nextSessionId: string, + runId: string, ingress?: { db: Db; claim: NativeWarmIngressClaim }, ): Promise { for (const [scope, entry] of warmNativeSessions) { @@ -2023,7 +2024,12 @@ async function retireSupersededWarmNativeSessions( const priorIngress = entry.ingressOwner; const physicalOwner = ingress && priorIngress && nativeWarmIngressAllocation(priorIngress.ingress) === nativeWarmIngressAllocation(ingress.claim.ingress); if (!logicalOwner && !physicalOwner) continue; - if (entry.busy || entry.preparingRunId || executingRunnerdSessionScopes.has(scope)) { + // This run's own instruction-directory reservation on the prior owner of + // the same task is not contention. A follow-up that cannot resume the + // saved session reserves that directory, then must retire the old provider + // while keeping the directory it adopted. Any other reservation is busy. + const ownReservation = logicalOwner && entry.preparingRunId === runId; + if (entry.busy || (entry.preparingRunId && !ownReservation) || executingRunnerdSessionScopes.has(scope)) { throw new Error("native_session_supervisor_busy"); } if (!logicalOwner) { @@ -2038,7 +2044,10 @@ async function retireSupersededWarmNativeSessions( try { await closeWarmNativeSession(entry, entry.closeOnReleaseReason, { discardUnsettledSession: logicalOwner && entry.normalizedSessionId !== nextSessionId, + ...(ownReservation ? { preserveInstructionsForRunId: runId } : {}), }); + // A failed close keeps the reservation, so its release retries retirement. + if (ownReservation) entry.preparingRunId = undefined; forgetWarmNativeSession(scope, entry); if (!logicalOwner) await assertNativeWarmIngressRetirement(ingress!.db, ingress!.claim, priorIngress!); } finally { @@ -8536,7 +8545,7 @@ export async function executePaperclipNativeSession(input: { } executingNativeOwnerScopes.set(ownerScope, ownerToken); if (!input.useRunnerd) await retireSupersededWarmNativeSessions( - ownerScope, nativeSessionScopeKey(input.execution), nativeSessionKey(input.execution), + ownerScope, nativeSessionScopeKey(input.execution), nativeSessionKey(input.execution), runId, ); if ((input.execution.schema === "paperclip.native-execution-input.v6" || input.execution.schema === "paperclip.native-execution-input.v8") && isCloudManagedInstance() && (!input.useRunnerd || input.runnerExecutionTarget?.kind !== "remote" @@ -9181,7 +9190,7 @@ async function executePaperclipNativeSessionWithinScope( const ingress = nativeWarmIngressIdentity(input.execution, input.runnerExecutionTarget); if (input.useRunnerd && !claimedStopCleanup) await retireSupersededWarmNativeSessions( nativeSessionOwnerScope(input.execution, input.runnerExecutionTarget?.environmentId ?? null), - nativeSessionScopeKey(input.execution), nativeSessionKey(input.execution), + nativeSessionScopeKey(input.execution), nativeSessionKey(input.execution), input.execution.binding.runId, ingress ? { db: input.db, claim: { ...input.execution.binding, ingress, leaseOwner, controllerBootId: controller.bootId, controllerGeneration } } : undefined, );