diff --git a/doc/execution-semantics.md b/doc/execution-semantics.md index 7cceed149b4..f9909012f43 100644 --- a/doc/execution-semantics.md +++ b/doc/execution-semantics.md @@ -3338,6 +3338,14 @@ run, it has a confirmed stop-and-retain receipt with state `stopped`, and no unsaved writer or pending native workspace export is bound to it. The backup stamp, single-owner, and later-provider-work checks above still apply. +A run can also start while its terminal predecessor still holds the warm lease +as `active`, because the predecessor releases the lease only after it saves its +instruction and workspace files. Before it selects a lease, acquisition waits up +to 20 seconds for that release when the matching lease belongs to a run that +ended in the last two minutes. It then resumes the warm sandbox. A predecessor +that is still running, or that ended earlier and never released its lease, +causes no wait; that run gets a new sandbox as before. + Reuse also skips a warm sandbox while its previous run's workspace export is still saving. A run that starts in that window gets a new sandbox, waits for the export, and realizes the saved files there. By the time it prepares the diff --git a/server/src/__tests__/environment-runtime.test.ts b/server/src/__tests__/environment-runtime.test.ts index 77fd0c8c0bb..3f73a1c6c98 100644 --- a/server/src/__tests__/environment-runtime.test.ts +++ b/server/src/__tests__/environment-runtime.test.ts @@ -4922,6 +4922,56 @@ describeEmbeddedPostgres("environmentRuntimeService", () => { }); }); + async function seedSettlingWarmLease(predecessorFinishedAt: Date, predecessorStatus: "succeeded" | "failed" = "succeeded", stopReplyDelayMs = 0) { + const seeded = await seedReusablePluginSandboxLease("paperclip_runner"); + await environmentService(db).releaseLease(seeded.reusableLease.id, "expired"); + const warmEnvironment = resolveRunnerEnvironmentForRun(seeded.environment, "paperclip_runner", { lifecycleMode: "warm" }, undefined, "selected-account:generation"); + let created = 0; + const call = vi.fn(async (_pluginId: string, method: string, input: any) => { + if (method === "environmentResumeLease") return { providerLeaseId: input.providerLeaseId, metadata: { remoteCwd: "/workspace" } }; + if (method === "environmentAcquireLease") return { providerLeaseId: `warm-${++created}`, metadata: { remoteCwd: "/workspace" } }; + if (method === "environmentStopLease") await new Promise(resolve => setTimeout(resolve, stopReplyDelayMs)); + if (method === "environmentDestroyLease" || method === "environmentStopLease" || method === "environmentReleaseLease") return { + providerLeaseId: input.providerLeaseId, state: method === "environmentDestroyLease" ? "destroyed" : "stopped", + }; + throw new Error(`Unexpected plugin method: ${method}`); + }); + const runtime = environmentRuntimeService(db, { pluginWorkerManager: { + isRunning: () => true, call, + getWorker: () => ({ supportedMethods: ["environmentResumeLease", "environmentReleaseLease", "environmentStopLease", "environmentDestroyLease"] }), + } as unknown as PluginWorkerManager }); + const input = { companyId: seeded.companyId, agentId: seeded.agentId, issueId: null, adapterType: "paperclip_runner", + environment: warmEnvironment, persistedExecutionWorkspace: { id: seeded.executionWorkspaceId, mode: "shared_workspace" as const } }; + const predecessorRunId = randomUUID(); + await db.insert(heartbeatRuns).values({ id: predecessorRunId, companyId: seeded.companyId, agentId: seeded.agentId, status: "running" }); + const predecessor = await runtime.acquireRunLease({ ...input, heartbeatRunId: predecessorRunId }); + await db.update(heartbeatRuns).set({ status: predecessorStatus, finishedAt: predecessorFinishedAt }).where(eq(heartbeatRuns.id, predecessorRunId)); + const nextRunId = randomUUID(); + await db.insert(heartbeatRuns).values({ id: nextRunId, companyId: seeded.companyId, agentId: seeded.agentId, status: "running" }); + return { runtime, call, input, predecessor, predecessorRunId, nextRunId }; + } + + it.each([["succeeded", 0], ["failed", 0], ["failed", 1_500]] as const)("resumes a warm lease that its %s predecessor is still releasing (stop reply after %i ms) instead of creating a fresh sandbox", async (predecessorStatus, stopReplyDelayMs) => { + const f = await seedSettlingWarmLease(new Date(), predecessorStatus, stopReplyDelayMs); + const released = new Promise(resolve => setTimeout(() => { + void f.runtime.releaseRunLeases(f.predecessorRunId, "released", undefined, "stop_and_retain").then(() => resolve()); + }, 400)); + const next = await f.runtime.acquireRunLease({ ...f.input, heartbeatRunId: f.nextRunId }); + await released; + expect(next.lease.providerLeaseId).toBe(f.predecessor.lease.providerLeaseId); + expect(next.lease.heartbeatRunId).toBe(f.nextRunId); + expect(f.call.mock.calls.filter(call => call[1] === "environmentAcquireLease")).toHaveLength(1); + }); + + it("does not wait for a warm lease that a predecessor finished long ago and never released", async () => { + const f = await seedSettlingWarmLease(new Date(Date.now() - 10 * 60_000)); + const startedAt = Date.now(); + const next = await f.runtime.acquireRunLease({ ...f.input, heartbeatRunId: f.nextRunId }); + expect(Date.now() - startedAt).toBeLessThan(5_000); + expect(next.lease.providerLeaseId).not.toBe(f.predecessor.lease.providerLeaseId); + await expect(environmentService(db).getLeaseById(f.predecessor.lease.id)).resolves.toMatchObject({ status: "active" }); + }); + it("does not retain or resume plugin-backed sandbox leases unless the provider opts in", async () => { const pluginId = randomUUID(); const { companyId, agentId, environment: baseEnvironment, runId } = await seedEnvironment(); diff --git a/server/src/services/environment-runtime.ts b/server/src/services/environment-runtime.ts index 8c84a559b04..b2fb67bb003 100644 --- a/server/src/services/environment-runtime.ts +++ b/server/src/services/environment-runtime.ts @@ -24,7 +24,7 @@ import { beginIdleTrackedWork } from "./task-admission.js"; import { historicalProviderCredentials } from "./historical-provider-credentials.js"; import { ownedWorkspaceRecoveryLease, type WorkspaceRecoverySourceIdentity } from "./workspace-recovery-source.js"; import type { NativeWorkspaceFinalizationOwnership } from "./native-runtime/native-workspace-finalization-ownership.js"; -import { hasStopOnlyCleanup, prepareSandboxStopAndRetain, readStopOnlyCleanup, settleStopOnlyCleanup, stopOnlyCleanupKey } from "./sandbox-stop-and-retain.js"; +import { hasStopOnlyCleanup, prepareSandboxStopAndRetain, readStopOnlyCleanup, SANDBOX_STOP_AND_RETAIN_KEY, settleStopOnlyCleanup, stopOnlyCleanupKey } from "./sandbox-stop-and-retain.js"; import { JsonRpcCallError, readEnvironmentAcquisitionDiagnostic, readEnvironmentCreationCleanupError } from "@paperclipai/plugin-sdk"; import { preserveEnvironmentSyncOutErrorDiagnostic } from "./environment-sync-out-error.js"; import { remoteTerminationReceipt } from "./remote-execution-termination.js"; @@ -1100,6 +1100,37 @@ function reusableSandboxLeaseScopeMatches(input: { }); } +const SETTLING_REUSABLE_LEASE_WAIT_MS = 20_000; +const SETTLING_REUSABLE_LEASE_RECENT_MS = 2 * 60_000; + +async function waitForSettlingReusableLease(db: Db, input: { + companyId: string; environmentId: string; executionWorkspaceId: string | null; issueId: string | null; + agentId: string; heartbeatRunId: string; +}) { + const deadline = Date.now() + SETTLING_REUSABLE_LEASE_WAIT_MS; + for (;;) { + const [settling] = await db.select({ id: environmentLeases.id }).from(environmentLeases) + .innerJoin(heartbeatRuns, and(eq(heartbeatRuns.id, environmentLeases.heartbeatRunId), + eq(heartbeatRuns.companyId, environmentLeases.companyId))) + .where(and( + eq(environmentLeases.companyId, input.companyId), + eq(environmentLeases.environmentId, input.environmentId), + eq(environmentLeases.leasePolicy, "reuse_by_environment"), + or(eq(environmentLeases.status, "active"), and(eq(environmentLeases.status, "pending_cleanup"), + sql`coalesce(${environmentLeases.metadata}, '{}'::jsonb) ? ${SANDBOX_STOP_AND_RETAIN_KEY}`)), + ne(environmentLeases.heartbeatRunId, input.heartbeatRunId), + input.executionWorkspaceId + ? eq(environmentLeases.executionWorkspaceId, input.executionWorkspaceId) + : and(sql`${environmentLeases.executionWorkspaceId} is null`, eq(environmentLeases.issueId, input.issueId!)), + sql`${environmentLeases.metadata}->>'agentId' = ${input.agentId}`, + inArray(heartbeatRuns.status, ["succeeded", "failed", "cancelled", "timed_out"]), + sql`${heartbeatRuns.finishedAt} > now() - make_interval(secs => ${SETTLING_REUSABLE_LEASE_RECENT_MS / 1000})`, + )).limit(1); + if (!settling || Date.now() >= deadline) return; + await new Promise(resolve => setTimeout(resolve, 250)); + } +} + function reusableLeaseCanBeResumed(input: { lease: Pick; heartbeatRunId: string | null; @@ -1748,6 +1779,12 @@ function createSandboxEnvironmentDriver( // We also filter out leases whose policy is not reuse_by_environment // and whose status is not reusable so non-reusable, cleanup-pending, // or terminal rows cannot be matched. + if (supportsReusableLeases && parsed.config.reuseLease && input.heartbeatRunId !== null + && (input.executionWorkspaceId !== null || input.issueId !== null) && input.agentId !== null) { + await waitForSettlingReusableLease(db, { companyId: input.companyId, environmentId: input.environment.id, + executionWorkspaceId: input.executionWorkspaceId, issueId: input.issueId, agentId: input.agentId, + heartbeatRunId: input.heartbeatRunId }); + } const reusableCandidateLeases = supportsReusableLeases && parsed.config.reuseLease && @@ -2119,6 +2156,12 @@ function createSandboxEnvironmentDriver( providerConfig: providerConfigForLease, }) : null; + if (supportsReusableLeases && parsed.config.reuseLease && input.heartbeatRunId !== null + && (input.executionWorkspaceId !== null || input.issueId !== null) && input.agentId !== null) { + await waitForSettlingReusableLease(db, { companyId: input.companyId, environmentId: input.environment.id, + executionWorkspaceId: input.executionWorkspaceId, issueId: input.issueId, agentId: input.agentId, + heartbeatRunId: input.heartbeatRunId }); + } const reusableCandidateLeases = supportsReusableLeases && parsed.config.reuseLease &&