fix: claim sandbox ownership before resume and preserve legacy launch profiles

Co-Authored-By: Paperclip <noreply@paperclip.ing>

server/src/__tests__/environment-runtime.test.ts

@@ -5855,6 +5855,43 @@ describeEmbeddedPostgres("environmentRuntimeService", () => {
     expect(result.value.lease.metadata?.sandboxLeaseAcquisition).toEqual({ outcome: "resumed" });
     expect(workerManager.call).toHaveBeenCalledOnce();
     expect((await environmentService(db).getLeaseById(first.lease.id))?.status).toBe("expired");
+
+    // A second server may arrive while the provider resume RPC is still in
+    // flight. It must see the new database owner before it can touch the disk.
+    await db.update(heartbeatRuns).set({ status: "succeeded" }).where(eq(heartbeatRuns.id, nextRunId));
+    await environmentService(db).releaseLease(result.value.lease.id, "released");
+    const raceIds = [randomUUID(), randomUUID()];
+    await db.insert(heartbeatRuns).values(raceIds.map((id) => ({
+      id, companyId: seeded.companyId, agentId: seeded.agentId, status: "running",
+    })));
+    let finishResume!: () => void;
+    let startedResume!: () => void;
+    const resumeStarted = new Promise<void>((resolve) => { startedResume = resolve; });
+    const resumeGate = new Promise<void>((resolve) => { finishResume = resolve; });
+    vi.mocked(workerManager.call).mockClear();
+    vi.mocked(workerManager.call).mockImplementation(async (_pluginId, method) => {
+      if (method !== "environmentResumeLease") throw new Error(`Unexpected replacement: ${method}`);
+      startedResume();
+      await resumeGate;
+      return { providerLeaseId: first.lease.providerLeaseId, metadata: {
+        provider: "fake-plugin", image: "fake:test", timeoutMs: 1234, reuseLease: true,
+      } };
+    });
+    const winner = acquire(raceIds[0]!, taskId);
+    // If acquisition regresses before the RPC, surface that rejection too.
+    await Promise.race([resumeStarted, winner.then(() => undefined)]);
+    try {
+      const otherServer = environmentRuntimeService(db, { pluginWorkerManager: workerManager });
+      await expect(otherServer.acquireRunLease({
+        companyId: seeded.companyId, environment: seeded.environment, issueId: taskId,
+        agentId: seeded.agentId, heartbeatRunId: raceIds[1]!, adapterType,
+        persistedExecutionWorkspace: { id: seeded.executionWorkspaceId, mode: "shared_workspace" },
+      })).rejects.toThrow("the lease was preserved and no replacement was created");
+      expect(workerManager.call).toHaveBeenCalledOnce();
+    } finally {
+      finishResume();
+      await winner;
+    }
   });
 
   it.each([false, true])("preserves active reusable sandbox leases held by another running run (same task: %s)", async (sameTask) => {

server/src/services/environment-runtime.ts

@@ -1295,6 +1295,31 @@ function createSandboxEnvironmentDriver(
   );
   const environmentsSvc = environmentService(db);
 
+  async function claimReusableLeaseBeforeResume(
+    lease: EnvironmentLease,
+    input: Parameters<EnvironmentRuntimeDriver["acquireRunLease"]>[0],
+  ): Promise<EnvironmentLease> {
+    if (!input.heartbeatRunId || lease.heartbeatRunId === input.heartbeatRunId) return lease;
+    // Transfer ownership atomically before touching the provider. Two server
+    // processes can observe the same released lease; only the winner of this
+    // conditional update/insert may resume its sandbox. A failed resume leaves
+    // the claimed row with this run for normal retention/recovery cleanup.
+    return environmentsSvc.acquireLease({
+      companyId: input.companyId,
+      environmentId: input.environment.id,
+      executionWorkspaceId: input.executionWorkspaceId,
+      issueId: input.issueId,
+      heartbeatRunId: input.heartbeatRunId,
+      assertCompanyBinding: input.assertCompanyBinding,
+      leasePolicy: lease.leasePolicy,
+      provider: lease.provider,
+      providerLeaseId: lease.providerLeaseId,
+      expiresAt: lease.expiresAt ? new Date(lease.expiresAt) : null,
+      metadata: lease.metadata,
+      replacesReusableLeaseId: lease.id,
+    });
+  }
+
   // A live sandbox whose teardown failed needs a durable `pending_cleanup` row,
   // so a sweep can find and release it. When every synchronous write attempt
   // fails, the database is down but the process still runs. So the driver keeps
@@ -1941,7 +1966,7 @@ function createSandboxEnvironmentDriver(
           ? (reusableExistingLeases.find((lease) => lease.metadata?.workFolderRecoveryRequired === true)?.providerLeaseId
             ?? findReusableSandboxLeaseId({ config: storedConfig, leases: reusableExistingLeases }))
           : null;
-        const reusableLease = reusableProviderLeaseId
+        let reusableLease = reusableProviderLeaseId
           ? reusableExistingLeases.find((lease) => lease.providerLeaseId === reusableProviderLeaseId)
           : null;
 
@@ -1971,6 +1996,8 @@ function createSandboxEnvironmentDriver(
               providerLeaseId: reusableLease.providerLeaseId,
             });
           }
+          reusableLease = await claimReusableLeaseBeforeResume(reusableLease, input);
+          if (!reusableLease.providerLeaseId) throw new Error("Reusable sandbox claim lost its provider identity");
           try {
             const resumeDeadline = Date.now() + 60_000;
             const configuredResumeTimeoutMs =
@@ -2308,10 +2335,11 @@ function createSandboxEnvironmentDriver(
           ? (reusableExistingLeases.find((lease) => lease.metadata?.workFolderRecoveryRequired === true)?.providerLeaseId
             ?? findReusableSandboxLeaseId({ config: parsed.config, leases: reusableExistingLeases }))
         : null;
-      const reusableLease = reusableProviderLeaseId
+      let reusableLease = reusableProviderLeaseId
         ? reusableExistingLeases.find((lease) => lease.providerLeaseId === reusableProviderLeaseId)
         : null;
 
+      if (reusableLease) reusableLease = await claimReusableLeaseBeforeResume(reusableLease, input);
       let providerLease;
       try {
         if (reusableLease && (reusableLease.metadata?.workFolderRecoveryRequired === true || hasLegacySandboxWorkspace(reusableLease))) {