diff --git a/server/src/__tests__/environment-runtime.test.ts b/server/src/__tests__/environment-runtime.test.ts index 79cc2b987b1..34c9b18aa09 100644 --- a/server/src/__tests__/environment-runtime.test.ts +++ b/server/src/__tests__/environment-runtime.test.ts @@ -10,7 +10,7 @@ import { mkdir, mkdtemp, readdir, rm, writeFile } from "node:fs/promises"; import os from "node:os"; import path from "node:path"; import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; -import { eq } from "drizzle-orm"; +import { eq, sql } from "drizzle-orm"; import { buildSshEnvLabFixtureConfig, getSshEnvLabSupport, @@ -1126,6 +1126,34 @@ describeEmbeddedPostgres("environmentRuntimeService", () => { }, ); + it("renews the stop-only cleanup claim while the provider stop is still running", async () => { + const seeded = await seedReusablePluginSandboxLease("paperclip_runner"); + const lease = seeded.reusableLease; + await db.update(environmentLeases).set({ leasePolicy: "ephemeral", metadata: { ...lease.metadata, reuseLease: false } }) + .where(eq(environmentLeases.id, lease.id)); + await db.update(heartbeatRuns).set({ status: "failed" }).where(eq(heartbeatRuns.id, seeded.runId)); + const started = Promise.withResolvers(), finish = Promise.withResolvers(); + const call = vi.fn(async (_id: string, method: string) => { + if (method === "environmentStopLease") { started.resolve(); await finish.promise; } + return { providerLeaseId: lease.providerLeaseId, state: "stopped" }; + }); + const runtime = environmentRuntimeService(db, { pluginWorkerManager: { + isRunning: () => true, call, + getWorker: () => ({ supportedMethods: ["environmentReleaseLease", "environmentDestroyLease", "environmentStopLease"] }), + } as unknown as PluginWorkerManager }); + vi.useFakeTimers({ toFake: ["setInterval", "clearInterval"] }); + const releasing = runtime.releaseRunLeases(seeded.runId, "failed", undefined, "stop_and_retain"); + try { + await started.promise; + const expiry = async () => Number((await environmentService(db).getLeaseById(lease.id))?.metadata?.pendingCleanupLeaseExpiresAtMs); + await db.update(environmentLeases).set({ metadata: sql`${environmentLeases.metadata} || '{"pendingCleanupLeaseExpiresAtMs":1}'::jsonb` }) + .where(eq(environmentLeases.id, lease.id)); + await vi.advanceTimersByTimeAsync(30_000); + await vi.waitFor(async () => expect(await expiry()).toBeGreaterThan(Date.now() + 14 * 60_000)); + } finally { finish.resolve(); await releasing; vi.useRealTimers(); } + expect(await environmentService(db).getLeaseById(lease.id)).toMatchObject({ status: "released", cleanupStatus: "success" }); + }); + it("reports an unwired manager separately from a stopped sandbox worker", async () => { const { companyId, environment, runId } = await seedReusablePluginSandboxLease(); const input = { diff --git a/server/src/services/environment-runtime.ts b/server/src/services/environment-runtime.ts index ba7c75b7391..2a9895ece08 100644 --- a/server/src/services/environment-runtime.ts +++ b/server/src/services/environment-runtime.ts @@ -3697,6 +3697,26 @@ export function environmentRuntimeService( return driver; } + function renewCleanupClaim(leaseId: string, attemptId: string, failure: string): ReturnType { + let renewalInFlight = false; + const renewal = setInterval(() => { + if (renewalInFlight) return; + renewalInFlight = true; + void db.update(environmentLeases).set({ + metadata: sql`${environmentLeases.metadata} || ${JSON.stringify({ + pendingCleanupLeaseExpiresAtMs: Date.now() + 15 * 60_000, + })}::jsonb`, + }).where(and(eq(environmentLeases.id, leaseId), eq(environmentLeases.status, "pending_cleanup"), + sql`${environmentLeases.metadata}->>'pendingCleanupAttemptId' = ${attemptId}`, + sql`${environmentLeases.metadata}->>'pendingCleanupInFlight' = 'true'`, + )).then(() => {}).catch(() => { + logger.warn({ leaseId }, failure); + }).finally(() => { renewalInFlight = false; }); + }, 30_000); + renewal.unref(); + return renewal; + } + async function destroyClaimedSandboxLease(input: { environment: Environment; leaseRow: typeof environmentLeases.$inferSelect; @@ -3747,22 +3767,7 @@ export function environmentRuntimeService( if (!claimed) return null; crashPoint("environment.teardown.after_claim_before_destroy"); const lease = toEnvironmentLeaseSnapshot(claimed); - let renewalInFlight = false; - const renewal = setInterval(() => { - if (renewalInFlight) return; - renewalInFlight = true; - void db.update(environmentLeases).set({ - metadata: sql`${environmentLeases.metadata} || ${JSON.stringify({ - pendingCleanupLeaseExpiresAtMs: Date.now() + 15 * 60_000, - })}::jsonb`, - }).where(and(eq(environmentLeases.id, lease.id), eq(environmentLeases.status, "pending_cleanup"), - sql`${environmentLeases.metadata}->>'pendingCleanupAttemptId' = ${attemptId}`, - sql`${environmentLeases.metadata}->>'pendingCleanupInFlight' = 'true'`, - )).then(() => {}).catch(() => { - logger.warn({ leaseId: lease.id }, "scoped cleanup ownership renewal failed"); - }).finally(() => { renewalInFlight = false; }); - }, 30_000); - renewal.unref(); + const renewal = renewCleanupClaim(lease.id, attemptId, "scoped cleanup ownership renewal failed"); try { const driver = getDriver(getLeaseDriverKey(lease, input.environment)); if (!driver?.destroyRunLease) return lease; @@ -3984,16 +3989,21 @@ export function environmentRuntimeService( if (!claimed) continue; crashPoint("environment.stop_retain.after_claim_before_stop"); const snapshot = toEnvironmentLeaseSnapshot(claimed); + const renewal = renewCleanupClaim(snapshot.id, attemptId, "stop-only cleanup ownership renewal failed"); try { const driver = requireDriverKey(getLeaseDriverKey(snapshot, environment)); if (!driver.retryPendingSandboxTeardown) throw new Error("Workspace export recovery requires stop-only cleanup."); const receipt = await driver.retryPendingSandboxTeardown({ environment, lease: snapshot }); + clearInterval(renewal); const lease = await settleStopOnlyCleanup(db, snapshot, { attemptId, receipt }); if (lease && environment) released.push({ environment, lease: toEnvironmentLeaseSnapshot(lease), leaseContext: { executionWorkspaceId: lease.executionWorkspaceId, executionWorkspaceMode: null } }); } catch (error) { + clearInterval(renewal); await settleStopOnlyCleanup(db, snapshot, { attemptId }); throw error; + } finally { + clearInterval(renewal); } continue; }