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 aff0735736e..7f92e3c2a7e 100644 --- a/server/src/services/native-runtime/native-session-executor.test.ts +++ b/server/src/services/native-runtime/native-session-executor.test.ts @@ -393,6 +393,7 @@ import { cancelNativeSession, closeWarmNativeSessionsForEnvironment, claimWarmNativeInstructionCopy, + warmSessionReleaseWait, nativeSessionWorkspaceScope, reserveWarmNativeInstructionDirectory, closeIdleWarmNativeSessionsForRestart, @@ -10425,6 +10426,72 @@ describe("native warm session supervision", () => { } }, 90_000); + it.each([["finishes", true], ["outlasts the wait", false]] as const)( + "a follow-up admitted while the previous turn's warm release is still running waits for it when the release %s", + async (_label, releaseFinishes) => { + const identity = `warm-release-wait-${releaseFinishes}`; + const warmExecution = { ...execution, binding: { ...execution.binding, runId: identity, executionWorkspaceId: `${identity}-workspace` }, + session: { ...execution.session, normalizedSessionId: identity, lifecyclePolicy: { mode: "warm" as const, idleTimeoutMs: 60_000 } }, + } as NativeExecutionInputV1; + const session = { close: vi.fn(async () => undefined) }; + let armPostTurnHold = false, holdPostTurnCheck = false; + state.execute.mockReset().mockImplementation(async (options) => { + await options.onSession?.(session); + // Only checks after the provider turn are held: that is the release window. + if (armPostTurnHold) holdPostTurnCheck = true; + return { result: { summary: "completed" }, terminal: { runTerminalState: "succeeded" }, + turnId: identity, normalizedSessionId: identity, providerSessionId: identity, + driverKind: "test", driverVersion: "1", nativeEventCount: 1, highestContiguousSourceSeq: 1, usage: null }; + }); + let releaseGate!: () => void, settling!: () => void; + const gate = new Promise(resolve => { releaseGate = resolve; }); + const reachedSettle = new Promise(resolve => { settling = resolve; }); + const copy = (runId: string, slow = false) => ({ runId, root: `/tmp/${identity}`, preparationKey: "same", + // The previous turn's post-turn check runs while it still holds the session. + hasChanges: async () => { if (slow && holdPostTurnCheck) { settling(); await gate; } return false; }, + collectStopped: async () => undefined }); + const original = { ...warmSessionReleaseWait }; + warmSessionReleaseWait.timeoutMs = releaseFinishes ? 5_000 : 200; + try { + await executePaperclipNativeSession({ db: leaseDb(warmExecution), execution: warmExecution, runnerInstanceId: identity, + runnerExecutionTarget: { kind: "remote", transport: "sandbox", environmentId: identity, remoteCwd: `/tmp/${identity}` }, + instructionWorkingCopy: copy(identity) }); + const second = { ...warmExecution, binding: { ...warmExecution.binding, runId: `${identity}-second` } }; + const secondCopy = copy(second.binding.runId, true); + expect(await claimWarmNativeInstructionCopy({ priorExecution: warmExecution, + companyId: second.binding.companyId, agentId: second.binding.agentId, + executionWorkspaceId: second.binding.executionWorkspaceId, workspace: second.workspace, + environmentId: identity, runId: second.binding.runId, preparationKey: "same", adopt: async () => secondCopy })).toBeTypeOf("function"); + armPostTurnHold = true; + const previous = executePaperclipNativeSession({ db: leaseDb(second), execution: second, runnerInstanceId: identity, + runnerExecutionTarget: { kind: "remote", transport: "sandbox", environmentId: identity, remoteCwd: `/tmp/${identity}` }, + instructionWorkingCopy: secondCopy }); + await reachedSettle; + // A follow-up wakes 43 ms after the turn settled, while its release runs. + const followUpRunId = `${identity}-follow-up`; + const followUp = claimWarmNativeInstructionCopy({ priorExecution: second, + companyId: second.binding.companyId, agentId: second.binding.agentId, + executionWorkspaceId: second.binding.executionWorkspaceId, workspace: second.workspace, + environmentId: identity, runId: followUpRunId, preparationKey: "same", adopt: async () => copy(followUpRunId) }); + if (releaseFinishes) { + setTimeout(releaseGate, 100); + const release = await followUp; + expect(release).toBeTypeOf("function"); + await previous; + // Abandon the claimed preparation so the warm owner is retired. + await release!(); + } else { + await expect(followUp).rejects.toThrow("native_session_supervisor_busy"); + releaseGate(); + } + await previous; + } finally { + Object.assign(warmSessionReleaseWait, original); + releaseGate(); + await closeWarmNativeSessionsForEnvironment({ environmentId: identity, reason: "test cleanup" }); + } + }); + it.each(["idle", "edit", "add", "remove", "projectless"])("retains the actual agent_files copy across warm turns, then retires for %s", async ending => { const root = await realpath(await mkdtemp(join(tmpdir(), "warm-agent-files-"))); const localRoot = join(root, "run", "live"); diff --git a/server/src/services/native-runtime/native-session-executor.ts b/server/src/services/native-runtime/native-session-executor.ts index 5a3c30b02cd..dba388d18ba 100644 --- a/server/src/services/native-runtime/native-session-executor.ts +++ b/server/src/services/native-runtime/native-session-executor.ts @@ -738,6 +738,25 @@ export function readNativeSessionIdleWork(): IdleLocalWork { } registerIdleLocalWorkProbe(readNativeSessionIdleWork); +/** + * How long a new run's admission waits for the warm session to stop being busy. + * A follow-up comment can wake the next run within milliseconds of the previous + * run settling, while that run's warm release is still finishing. Waiting out + * the release lets the follow-up reuse the session instead of failing. A + * session still busy at the deadline is a real concurrent owner and fails as + * before. + */ +export const warmSessionReleaseWait = { timeoutMs: 30_000, pollMs: 25 }; + +async function waitForWarmNativeSessionIdle(scope: string, entry: WarmNativeSession): Promise { + const deadline = Date.now() + warmSessionReleaseWait.timeoutMs; + // Wait only while the owner's turn is over and it is releasing the session. + while (warmNativeSessions.get(scope) === entry && entry.busy && settlingNativeSessionScopes.has(scope) + && Date.now() < deadline) { + await new Promise(resolve => setTimeout(resolve, warmSessionReleaseWait.pollMs)); + } +} + /** Reserve the existing physical copy before heartbeat composes AGENT_HOME and * task guidance. The final execution/configuration fence still owns admission. */ export async function claimWarmNativeInstructionCopy(input: { @@ -759,6 +778,9 @@ export async function claimWarmNativeInstructionCopy(input: { binding: { ...prior.binding, runId: input.runId, executionWorkspaceId: input.executionWorkspaceId }, workspace: input.workspace, })) || prior.workspace.cwd !== input.workspace.cwd) return null; const scope = nativeSessionScopeKey(prior); + const busyEntry = warmNativeSessions.get(scope); + if (busyEntry?.busy) await waitForWarmNativeSessionIdle(scope, busyEntry); + // Re-read after the wait: the release may have retired or replaced the session. const entry = warmNativeSessions.get(scope); if (!entry || !entry.instructionWorkingCopy?.runId || entry.environmentId !== input.environmentId) return null; if (entry.busy) throw new Error("native_session_supervisor_busy"); @@ -1982,6 +2004,9 @@ const initializingSessionToolAuthorities = new Set(); const executingRunnerdSessionScopes = new Map(); const executingNativeOwnerScopes = new Map(); const executingNativeIngressAllocations = new Map(); +/** Session scope -> the run releasing it after its turn (warm release and + * hand-back), which a follow-up's instruction claim may wait for. */ +const settlingNativeSessionScopes = new Map(); export function nativeSessionScopeDigestIsBusy(digest: string): boolean { const scopes = [ @@ -9440,6 +9465,8 @@ export async function executePaperclipNativeSession(input: { nativeSessionStartups.delete(runId); setNativeExecutionActive(runId, false); } + const settledScope = nativeSessionScopeKey(input.execution); + if (settlingNativeSessionScopes.get(settledScope) === runId) settlingNativeSessionScopes.delete(settledScope); } if (startupUnsettled) throw new NativeStartupCleanupPendingError(startup.executionContext!.id, startup.executionContext!.deadlineAtMs, startup.leaseOwner!, startup.executionContext!.pluginSupported, @@ -12009,6 +12036,9 @@ async function executePaperclipNativeSessionWithinScope( await timeController("github_stop", () => githubAccess!.stop()); } } + // The provider turn is over: what follows is release and hand-back, which a + // follow-up's admission may wait for (`warmSessionReleaseWait`). + settlingNativeSessionScopes.set(nativeSessionScopeKey(input.execution), input.execution.binding.runId); let postExecutionFailed = false; let postExecutionCleanupFailed = false; try {