deepline 0.3.285 → 0.3.286

package/dist/bundling-sources/shared_libs/play-runtime/scheduler-backends/postgres.ts

@@ -7400,7 +7400,7 @@ export async function markPostgresSchedulerAttemptWaiting(
   );
   const waiting = await client.query<{ run_id: string; event_id: string }>(
     `
-      WITH waiting_attempt AS (
+      WITH ${runtimeSandboxCustodyLocksSql('run.run_id = $1', options)}, waiting_attempt AS (
         UPDATE ${tableName('attempts', options)}
         SET status = 'waiting',
             resource_json = CASE
@@ -7436,6 +7436,7 @@ export async function markPostgresSchedulerAttemptWaiting(
             END,
             heartbeat_at = CASE WHEN $8::boolean THEN heartbeat_at ELSE now() END,
             updated_at = now()
+        FROM (SELECT count(*) FROM custody_attempts) AS attempt_locks
         WHERE run_id = $1
           AND attempt = $2
           AND status = 'running'
@@ -8164,13 +8165,14 @@ export async function deferPostgresSchedulerAttemptForCapacity(
 ): Promise<boolean> {
   const deferred = await client.query<{ run_id: string }>(
     `
-      WITH deferred_attempt AS (
+      WITH ${runtimeSandboxCustodyLocksSql('run.run_id = $1', options)}, deferred_attempt AS (
         UPDATE ${tableName('attempts', options)}
         SET status = 'stale',
             lease_expires_at = NULL,
             heartbeat_at = now(),
             error_json = $4::jsonb,
             updated_at = now()
+        FROM (SELECT count(*) FROM custody_attempts) AS attempt_locks
         WHERE run_id = $1
           AND attempt = $2
           AND status IN ('claimed', 'running', 'waiting')
@@ -8232,9 +8234,9 @@ export async function deferPostgresSchedulerAttemptForExecutorTokenOutage(
     deferred: boolean;
   }>(
     `
-      WITH current_attempt AS (
+      WITH ${runtimeSandboxCustodyLocksSql('run.run_id = $1', options)}, current_attempt AS (
         SELECT created_at
-        FROM ${tableName('attempts', options)}
+        FROM custody_attempts
         WHERE run_id = $1
           AND attempt = $2
           AND lease_token = $3
@@ -8267,6 +8269,7 @@ export async function deferPostgresSchedulerAttemptForExecutorTokenOutage(
             heartbeat_at = now(),
             error_json = $4::jsonb,
             updated_at = now()
+        FROM (SELECT count(*) FROM custody_attempts) AS attempt_locks
         WHERE run_id = $1
           AND attempt = $2
           AND status IN ('claimed', 'running')
@@ -8326,10 +8329,10 @@ export async function claimPostgresSchedulerWaitingDetachedAttemptForTerminal(
 ): Promise<boolean> {
   const claimed = await client.query<{ run_id: string }>(
     `
-      WITH candidate AS MATERIALIZED (
+      WITH ${runtimeSandboxCustodyLocksSql('run.run_id = $1', options)}, candidate AS MATERIALIZED (
         SELECT runs.run_id, attempts.attempt
-        FROM ${tableName('runs', options)} AS runs
-        JOIN ${tableName('attempts', options)} AS attempts
+        FROM custody_runs AS runs
+        JOIN custody_attempts AS attempts
           ON attempts.run_id = runs.run_id
          AND attempts.attempt = (runs.suspension_json ->> 'runnerAttempt')::integer
         WHERE runs.run_id = $1
@@ -8349,7 +8352,6 @@ export async function claimPostgresSchedulerWaitingDetachedAttemptForTerminal(
               AND newer.attempt > attempts.attempt
               AND newer.status IN ('claimed', 'running', 'waiting', 'completed', 'failed')
           )
-        FOR UPDATE OF runs, attempts
       ),
       running_attempt AS (
         UPDATE ${tableName('attempts', options)} AS target
@@ -8588,19 +8590,18 @@ export async function deferPostgresSchedulerAttemptForRecovery(
     deadline_at_ms: string | null;
   }>(
     `
-      WITH candidate AS (
+      WITH ${runtimeSandboxCustodyLocksSql('run.run_id = $1', options)}, candidate AS (
         SELECT runs.run_id,
                runs.recovery_json,
                attempts.run_id AS attempt_run_id
-        FROM ${tableName('runs', options)} AS runs
-        JOIN ${tableName('attempts', options)} AS attempts
+        FROM custody_runs AS runs
+        JOIN custody_attempts AS attempts
           ON attempts.run_id = runs.run_id
         WHERE attempts.run_id = $1
           AND attempts.attempt = $2
           AND attempts.lease_token = $3
           AND attempts.status IN ('claimed', 'running', 'waiting')
           AND runs.status IN ('queued', 'running', 'waiting')
-        FOR UPDATE OF runs, attempts
       ), computed AS (
         SELECT run_id,
                coalesce((recovery_json ->> 'deferCount')::integer, 0) + 1
@@ -9322,20 +9323,24 @@ export async function recordPostgresSchedulerManagedSandboxOperationLive(
     'production';
   const recorded = await client.query<{ operation_key: string }>(
     `
-      WITH eligible_attempt AS (
+      WITH ${runtimeSandboxCustodyLocksSql('run.run_id = $1', options)}, eligible_attempt AS (
         SELECT run_id
-        FROM ${tableName('attempts', options)}
+        FROM custody_attempts
         WHERE run_id = $1
           AND attempt = $2
           AND lease_token = $3
           AND status IN ('claimed', 'running', 'waiting')
+          AND EXISTS (SELECT 1 FROM custody_runs WHERE status IN ('queued', 'running', 'waiting'))
+          AND NOT EXISTS (SELECT 1 FROM custody_attempts AS newer WHERE newer.attempt > $2)
       ), eligible_lease AS MATERIALIZED (
         SELECT lease_id
-        FROM ${tableName('runtime_sandbox_capacity_leases', options)}
+        FROM custody_leases
+        CROSS JOIN (SELECT count(*) FROM custody_operations) AS operation_locks
         WHERE lease_id = $5
+          AND run_id = $1 AND provider = $10
+          AND (sandbox_id IS NULL OR (sandbox_id = $6 AND routing_domain IS NOT DISTINCT FROM $7))
           AND released_at IS NULL
           AND expires_at > clock_timestamp()
-        FOR UPDATE
       ), live_operation AS (
         UPDATE ${tableName('runtime_sandbox_operations', options)} AS operation
         SET sandbox_id = $6,
@@ -9378,7 +9383,6 @@ export async function recordPostgresSchedulerManagedSandboxOperationLive(
         JOIN recorded_attempt ON TRUE
         WHERE lease.lease_id = live_operation.operation_key
           AND lease.released_at IS NULL
-          AND lease.expires_at > clock_timestamp()
         RETURNING lease.lease_id
       )
       SELECT live_operation.operation_key
@@ -9413,42 +9417,52 @@ export async function recordPostgresSchedulerManagedSandboxOperationLive(
     return false;
   }
   const updated = await client.query<{ run_id: string }>(
-    `UPDATE ${tableName('attempts', options)} AS attempt
-     SET resource_json = (
-           SELECT jsonb_agg(
-             CASE
-               WHEN resource->>'sandboxCapacityLeaseId' = $5
-                AND resource->>'sandboxId' = $6
-                AND resource->>'kind' = $4::jsonb->>'kind'
-               THEN resource || jsonb_build_object(
-                 'billingEndedAt', $4::jsonb->'billingEndedAt'
-               )
-               ELSE resource
-             END ORDER BY ordinal
-           )
-           FROM jsonb_array_elements(attempt.resource_json)
-             WITH ORDINALITY AS resources(resource, ordinal)
-         ),
-         updated_at = clock_timestamp()
-     FROM ${tableName('runtime_sandbox_operations', options)} AS operation
-     JOIN ${tableName('runtime_sandbox_capacity_leases', options)} AS lease
-       ON lease.lease_id = operation.operation_key
-     WHERE attempt.run_id = $1
-       AND attempt.attempt = $2
-       AND attempt.lease_token = $3
-       AND attempt.status IN ('claimed', 'running', 'waiting')
-       AND operation.operation_key = $5
-       AND operation.run_id = attempt.run_id
-       AND operation.attempt = attempt.attempt
-       AND operation.sandbox_id = $6
-       AND operation.provider = $7
-       AND operation.state = $8
-       AND operation.environment = $9
-       AND operation.routing_domain IS NOT DISTINCT FROM $10
-       AND lease.sandbox_id = $6
-       AND lease.released_at IS NULL
-       AND lease.expires_at > clock_timestamp()
-     RETURNING attempt.run_id`,
+    `WITH ${runtimeSandboxCustodyLocksSql('run.run_id = $1', options)}, eligible_retirement AS MATERIALIZED (
+       SELECT operation.operation_key, operation.run_id, operation.attempt
+       FROM custody_operations AS operation
+       JOIN custody_attempts AS attempt ON attempt.run_id = operation.run_id AND attempt.attempt = operation.attempt
+       JOIN custody_leases AS lease ON lease.lease_id = operation.operation_key
+       WHERE attempt.run_id = $1 AND attempt.attempt = $2 AND attempt.lease_token = $3
+         AND attempt.status IN ('claimed', 'running', 'waiting')
+         AND NOT EXISTS (SELECT 1 FROM custody_attempts AS newer WHERE newer.attempt > $2)
+         AND EXISTS (SELECT 1 FROM custody_runs WHERE status IN ('queued', 'running', 'waiting'))
+         AND operation.operation_key = $5 AND operation.sandbox_id = $6
+         AND operation.provider = $7 AND operation.state IN ($8, 'deleting')
+         AND operation.environment = $9 AND operation.routing_domain IS NOT DISTINCT FROM $10
+         AND lease.sandbox_id = $6 AND lease.released_at IS NULL
+         AND EXISTS (
+           SELECT 1 FROM jsonb_array_elements(attempt.resource_json) AS fact
+           WHERE fact->>'sandboxCapacityLeaseId' = $5 AND fact->>'sandboxId' = $6
+             AND fact->>'kind' = $4::jsonb->>'kind'
+         )
+     ), retiring_operation AS (
+       UPDATE ${tableName('runtime_sandbox_operations', options)} AS operation
+       SET state = 'deleting', updated_at = clock_timestamp()
+       FROM eligible_retirement
+       WHERE operation.operation_key = eligible_retirement.operation_key
+       RETURNING operation.operation_key
+     ), closed_resource AS (
+       UPDATE ${tableName('attempts', options)} AS attempt
+       SET resource_json = (
+             SELECT jsonb_agg(
+               CASE
+                 WHEN resource->>'sandboxCapacityLeaseId' = $5 AND resource->>'sandboxId' = $6
+                  AND resource->>'kind' = $4::jsonb->>'kind'
+                  AND (resource->'billingEndedAt' IS NULL OR resource->'billingEndedAt' = 'null'::jsonb)
+                 THEN resource || jsonb_build_object('billingEndedAt', $4::jsonb->'billingEndedAt')
+                 ELSE resource END ORDER BY ordinal
+             ) FROM jsonb_array_elements(attempt.resource_json) WITH ORDINALITY AS resources(resource, ordinal)
+           ), updated_at = clock_timestamp()
+       FROM retiring_operation
+       WHERE attempt.run_id = $1 AND attempt.attempt = $2 AND attempt.lease_token = $3
+       RETURNING attempt.run_id
+     ), queued AS (
+       INSERT INTO ${tableName('sandbox_cleanup_jobs_v2', options)} (
+         provider, sandbox_id, run_id, attempt, environment, routing_domain
+       ) SELECT $7, $6, $1, $2, $9, $10 FROM closed_resource
+       ON CONFLICT (provider, sandbox_id, routing_domain_key) DO NOTHING
+     )
+     SELECT run_id FROM closed_resource`,
     [
       input.runId,
       input.attempt,
@@ -9492,14 +9506,17 @@ export async function readPostgresSchedulerManagedSandboxOperationAlreadyLive(
     operationKey: string;
     provider: RuntimeSandboxProvider;
     sandboxId: string;
+    environment?: 'preview' | 'production';
+    routingDomain?: string | null;
   },
   options?: Pick<PostgresSchedulerOptions, 'schema'>,
 ): Promise<boolean> {
   const found = await client.query<{ operation_key: string }>(
     `
+      WITH ${runtimeSandboxCustodyLocksSql('run.run_id = $1', options)}
       SELECT operation.operation_key
-      FROM ${tableName('runtime_sandbox_operations', options)} AS operation
-      JOIN ${tableName('attempts', options)} AS attempt
+      FROM custody_operations AS operation
+      JOIN custody_attempts AS attempt
         ON attempt.run_id = operation.run_id
        AND attempt.attempt = operation.attempt
       WHERE operation.operation_key = $4
@@ -9508,8 +9525,12 @@ export async function readPostgresSchedulerManagedSandboxOperationAlreadyLive(
         AND operation.provider = $5
         AND operation.state = $6
         AND operation.sandbox_id = $7
+        AND ($8::text IS NULL OR operation.environment = $8)
+        AND (NOT $9::boolean OR operation.routing_domain IS NOT DISTINCT FROM $10::text)
         AND attempt.lease_token = $3
         AND attempt.status IN ('claimed', 'running', 'waiting')
+        AND NOT EXISTS (SELECT 1 FROM custody_attempts AS newer WHERE newer.attempt > $2)
+        AND EXISTS (SELECT 1 FROM custody_runs WHERE status IN ('queued', 'running', 'waiting'))
     `,
     [
       input.runId,
@@ -9519,6 +9540,9 @@ export async function readPostgresSchedulerManagedSandboxOperationAlreadyLive(
       input.provider,
       RuntimeSandboxOperationState.Live,
       input.sandboxId,
+      input.environment ?? null,
+      input.routingDomain !== undefined,
+      input.routingDomain?.trim() || null,
     ],
   );
   return found.rows.length === 1;
@@ -9671,9 +9695,10 @@ export async function resolvePostgresSchedulerAmbiguousModalCreateIntentFromProv
 ): Promise<boolean> {
   const resolved = await client.query<{ run_id: string }>(
     `
-      WITH resolved AS (
+      WITH ${runtimeSandboxCustodyLocksSql('run.run_id = $1', options)}, resolved AS (
         UPDATE ${tableName('ambiguous_modal_creates', options)}
         SET status = 'resolved', resolved_at = clock_timestamp()
+        FROM (SELECT count(*) FROM custody_operations) AS operation_locks
         WHERE run_id = $1
           AND attempt = $2
           AND modal_app_id = $3
@@ -9726,9 +9751,10 @@ export async function recordPostgresSchedulerReconciledModalResourcesForCleanup(
   const payload = stringifyPostgresJson(resources);
   const recorded = await client.query<{ run_id: string }>(
     `
-      WITH claimed_intent AS (
+      WITH ${runtimeSandboxCustodyLocksSql('run.run_id = $1', options)}, claimed_intent AS (
         UPDATE ${tableName('ambiguous_modal_creates', options)}
         SET status = 'resolved', resolved_at = clock_timestamp()
+        FROM (SELECT count(*) FROM custody_operations) AS operation_locks
         WHERE run_id = $1
           AND attempt = $2
           AND modal_app_id = $3
@@ -9951,16 +9977,16 @@ export async function recordPostgresSchedulerLateManagedSandboxOperationForClean
     'production';
   const result = await client.query<{ operation_key: string }>(
     `
-      WITH eligible_lease AS MATERIALIZED (
+      WITH ${runtimeSandboxCustodyLocksSql('run.run_id = $1', options)}, eligible_lease AS MATERIALIZED (
         SELECT lease_id
-        FROM ${tableName('runtime_sandbox_capacity_leases', options)}
+        FROM custody_leases
+        CROSS JOIN (SELECT count(*) FROM custody_operations) AS operation_locks
         WHERE lease_id = $4
           AND released_at IS NULL
           AND expires_at > clock_timestamp()
-        FOR UPDATE
       ), eligible_attempt AS (
         SELECT run_id
-        FROM ${tableName('attempts', options)}
+        FROM custody_attempts
         WHERE run_id = $1
           AND attempt = $2
           AND lease_token = $3
@@ -10215,12 +10241,14 @@ export async function settlePostgresSchedulerRuntimeResources(
 ): Promise<number> {
   const stamped = await client.query<{ run_id: string }>(
     `
+      WITH ${runtimeSandboxCustodyLocksSql('run.run_id = $1', options)}
       UPDATE ${tableName('attempts', options)}
       SET resource_json = (
             SELECT COALESCE(
               jsonb_agg(
                 CASE
                   WHEN elem->>'kind' IN ('daytona_sandbox', 'modal_sandbox')
+                   AND NOT ${protectedManagedSandboxSql({ runId: '$1', provider: "(CASE elem->>'kind' WHEN 'modal_sandbox' THEN 'modal' ELSE 'daytona' END)", sandboxId: "elem->>'sandboxId'", routingDomain: "(CASE elem->>'kind' WHEN 'modal_sandbox' THEN NULLIF(elem->>'modalAppId', '') ELSE NULLIF(elem->>'daytonaOrganizationId', '') END)" }, options, true)}
                    AND (
                      elem->'billingEndedAt' IS NULL
                      OR elem->'billingEndedAt' = 'null'::jsonb
@@ -10241,6 +10269,7 @@ export async function settlePostgresSchedulerRuntimeResources(
             FROM jsonb_array_elements(resource_json) AS elem
           ),
           updated_at = now()
+      FROM (SELECT count(*) FROM custody_operations) AS operation_locks
       WHERE run_id = $1
         AND ($4::integer IS NULL OR attempt <= $4)
         AND resource_json IS NOT NULL
@@ -11539,6 +11568,79 @@ export async function completePostgresSchedulerComputeBillingSessionBootstrap(
   return completed.rows.length > 0;
 }
 
+function runtimeSandboxCustodyLocksSql(
+  runPredicate: string,
+  options?: Pick<PostgresSchedulerOptions, 'schema'>,
+  skipLocked = false,
+): string {
+  return `custody_runs AS MATERIALIZED (
+    SELECT run.* FROM ${tableName('runs', options)} AS run
+    WHERE ${runPredicate}
+    ORDER BY run.run_id FOR NO KEY UPDATE ${skipLocked ? 'SKIP LOCKED' : ''}
+  ), custody_attempts AS MATERIALIZED (
+    SELECT attempt.* FROM ${tableName('attempts', options)} AS attempt
+    WHERE attempt.run_id IN (SELECT run_id FROM custody_runs)
+    ORDER BY attempt.run_id, attempt.attempt FOR UPDATE
+  ), custody_leases AS MATERIALIZED (
+    SELECT lease.* FROM ${tableName('runtime_sandbox_capacity_leases', options)} AS lease
+    CROSS JOIN (SELECT count(*) FROM custody_attempts) AS attempt_locks
+    WHERE lease.run_id IN (SELECT run_id FROM custody_runs)
+    ORDER BY lease.run_id, lease.lease_id FOR UPDATE OF lease
+  ), custody_operations AS MATERIALIZED (
+    SELECT operation.* FROM ${tableName('runtime_sandbox_operations', options)} AS operation
+    CROSS JOIN (SELECT count(*) FROM custody_leases) AS lease_locks
+    WHERE operation.run_id IN (SELECT run_id FROM custody_runs)
+    ORDER BY operation.run_id, operation.operation_key FOR UPDATE OF operation
+  )`;
+}
+
+function protectedManagedSandboxSql(
+  identity: {
+    runId: string;
+    provider: string;
+    sandboxId: string;
+    routingDomain: string;
+  },
+  options?: Pick<PostgresSchedulerOptions, 'schema'>,
+  locked = false,
+): string {
+  const runs = locked ? 'custody_runs' : tableName('runs', options);
+  const attempts = locked ? 'custody_attempts' : tableName('attempts', options);
+  const operations = locked
+    ? 'custody_operations'
+    : tableName('runtime_sandbox_operations', options);
+  const leases = locked
+    ? 'custody_leases'
+    : tableName('runtime_sandbox_capacity_leases', options);
+  return `EXISTS (
+    SELECT 1 FROM ${runs} AS managed_run
+    JOIN ${operations} AS managed_operation ON managed_operation.run_id = managed_run.run_id
+    JOIN ${leases} AS managed_lease ON managed_lease.lease_id = managed_operation.operation_key
+    WHERE managed_run.run_id = ${identity.runId}
+      AND managed_run.status IN ('queued', 'running', 'waiting')
+      AND managed_operation.provider = ${identity.provider}
+      AND COALESCE(managed_operation.sandbox_id, managed_lease.sandbox_id) = ${identity.sandboxId}
+      AND COALESCE(managed_operation.routing_domain, managed_lease.routing_domain)
+          IS NOT DISTINCT FROM ${identity.routingDomain}
+      AND managed_lease.released_at IS NULL
+      AND NOT EXISTS (SELECT 1 FROM ${tableName('control', options)} AS cancel
+                      WHERE cancel.key = 'runtime_cancel:' || managed_run.run_id)
+      AND (
+        (managed_operation.state = 'reserved' AND managed_lease.expires_at > clock_timestamp())
+        OR (managed_operation.state = 'live' AND EXISTS (
+          SELECT 1 FROM ${attempts} AS owner_attempt
+          WHERE owner_attempt.run_id = managed_run.run_id
+            AND owner_attempt.attempt = managed_operation.attempt
+            AND owner_attempt.status IN ('claimed', 'running', 'waiting')
+            AND owner_attempt.runner_result_json IS NULL
+            AND NOT EXISTS (SELECT 1 FROM ${attempts} AS newer
+                            WHERE newer.run_id = owner_attempt.run_id AND newer.attempt > owner_attempt.attempt)
+            AND (owner_attempt.lease_expires_at IS NULL OR owner_attempt.lease_expires_at > clock_timestamp())
+        ))
+      )
+  )`;
+}
+
 /** Deletion needs terminal/cancellation/retired-resource authority. An old
  * worker error (including a billing end) is not evidence the runner stopped. */
 function protectedDetachedSandboxSql(
@@ -11549,10 +11651,11 @@ function protectedDetachedSandboxSql(
     routingDomain: string;
   },
   options?: Pick<PostgresSchedulerOptions, 'schema'>,
+  locked = false,
 ): string {
-  return `EXISTS (
-    SELECT 1 FROM ${tableName('runs', options)} AS detached_run
-    JOIN ${tableName('attempts', options)} AS detached_attempt
+  return `(${protectedManagedSandboxSql(identity, options, locked)} OR EXISTS (
+    SELECT 1 FROM ${locked ? 'custody_runs' : tableName('runs', options)} AS detached_run
+    JOIN ${locked ? 'custody_attempts' : tableName('attempts', options)} AS detached_attempt
       ON detached_attempt.run_id = detached_run.run_id
     CROSS JOIN LATERAL jsonb_array_elements(
       CASE WHEN jsonb_typeof(detached_attempt.resource_json) = 'array'
@@ -11580,18 +11683,19 @@ function protectedDetachedSandboxSql(
             AND detached_resource -> 'executionDetachedAt' IS NOT NULL
             AND detached_resource -> 'executionDetachedAt' <> 'null'::jsonb)
       )
-      AND NOT EXISTS (SELECT 1 FROM ${tableName('attempts', options)} AS newer
+      AND NOT EXISTS (SELECT 1 FROM ${locked ? 'custody_attempts' : tableName('attempts', options)} AS newer
                       WHERE newer.run_id = detached_run.run_id
                         AND newer.attempt > detached_attempt.attempt
                         AND newer.status IN ('claimed', 'running', 'waiting', 'completed', 'failed'))
       AND NOT EXISTS (SELECT 1 FROM ${tableName('control', options)} AS cancel
                       WHERE cancel.key = 'runtime_cancel:' || detached_run.run_id)
-  )`;
+  ))`;
 }
 
 function protectedDetachedResourceSql(
   runId: string,
   options?: Pick<PostgresSchedulerOptions, 'schema'>,
+  locked = false,
 ): string {
   return protectedDetachedSandboxSql(
     {
@@ -11603,6 +11707,7 @@ function protectedDetachedResourceSql(
         "(CASE resource->>'kind' WHEN 'modal_sandbox' THEN NULLIF(resource->>'modalAppId', '') ELSE NULLIF(resource->>'daytonaOrganizationId', '') END)",
     },
     options,
+    locked,
   );
 }
 
@@ -11632,9 +11737,10 @@ export async function enqueuePostgresSchedulerSandboxCleanupForAttempt(
 ): Promise<number> {
   const enqueued = await client.query<{ sandbox_id: string }>(
     `
-      WITH attempt_resources AS (
+      WITH ${runtimeSandboxCustodyLocksSql('run.run_id = $1', options)}, attempt_resources AS (
         SELECT run_id, attempt, resource_json
-        FROM ${tableName('attempts', options)}
+        FROM custody_attempts
+        CROSS JOIN (SELECT count(*) FROM custody_operations) AS operation_locks
         WHERE run_id = $1
           AND attempt = $2
           AND resource_json IS NOT NULL
@@ -11669,7 +11775,7 @@ export async function enqueuePostgresSchedulerSandboxCleanupForAttempt(
         ) AS resource
         WHERE resource->>'kind' IN ('daytona_sandbox', 'modal_sandbox')
           AND NULLIF(resource->>'sandboxId', '') IS NOT NULL
-          AND NOT ${protectedDetachedResourceSql('$1', options)}
+          AND NOT ${protectedDetachedResourceSql('$1', options, true)}
         ON CONFLICT (provider, sandbox_id, routing_domain_key) DO NOTHING
         RETURNING sandbox_id
       )
@@ -11689,6 +11795,7 @@ export async function enqueuePostgresSchedulerSandboxCleanupForRun(
 ): Promise<void> {
   await client.query(
     `
+      WITH ${runtimeSandboxCustodyLocksSql('run.run_id = $1', options)}
       INSERT INTO ${tableName('sandbox_cleanup_jobs_v2', options)} (
         provider, sandbox_id, run_id, attempt, environment, routing_domain
       )
@@ -11709,7 +11816,8 @@ export async function enqueuePostgresSchedulerSandboxCleanupForRun(
           WHEN 'modal_sandbox' THEN NULLIF(resource->>'modalAppId', '')
           ELSE NULLIF(resource->>'daytonaOrganizationId', '')
         END
-      FROM ${tableName('attempts', options)} AS attempt
+      FROM custody_attempts AS attempt
+      CROSS JOIN (SELECT count(*) FROM custody_operations) AS operation_locks
       CROSS JOIN LATERAL jsonb_array_elements(
         CASE
           WHEN jsonb_typeof(attempt.resource_json) = 'array'
@@ -11722,16 +11830,16 @@ export async function enqueuePostgresSchedulerSandboxCleanupForRun(
         AND attempt.resource_json IS NOT NULL
         AND resource->>'kind' IN ('daytona_sandbox', 'modal_sandbox')
         AND NULLIF(resource->>'sandboxId', '') IS NOT NULL
-          AND NOT ${protectedDetachedResourceSql('$1', options)}
+          AND NOT ${protectedDetachedResourceSql('$1', options, true)}
       UNION ALL
       SELECT operation.provider, operation.sandbox_id, operation.run_id,
              operation.attempt, operation.environment, operation.routing_domain
-      FROM ${tableName('runtime_sandbox_operations', options)} AS operation
+      FROM custody_operations AS operation
       WHERE operation.run_id = $1
         AND ($2::integer IS NULL OR operation.attempt <= $2)
         AND operation.state IN ('live', 'deleting')
         AND operation.sandbox_id IS NOT NULL
-        AND NOT ${protectedDetachedSandboxSql({ runId: 'operation.run_id', provider: 'operation.provider', sandboxId: 'operation.sandbox_id', routingDomain: 'operation.routing_domain' }, options)}
+        AND NOT ${protectedDetachedSandboxSql({ runId: 'operation.run_id', provider: 'operation.provider', sandboxId: 'operation.sandbox_id', routingDomain: 'operation.routing_domain' }, options, true)}
       ON CONFLICT (provider, sandbox_id, routing_domain_key) DO NOTHING
     `,
     [input.runId, input.throughAttempt ?? null],
@@ -11746,9 +11854,10 @@ export async function abandonUnstartedModalCancellationReservations(
   options?: Pick<PostgresSchedulerOptions, 'schema'>,
 ): Promise<void> {
   await client.query(
-    `WITH abandoned AS (
+    `WITH ${runtimeSandboxCustodyLocksSql('run.run_id = $1', options)}, abandoned AS (
        UPDATE ${tableName('runtime_sandbox_operations', options)} AS operation
        SET state = 'abandoned', updated_at = clock_timestamp()
+       FROM (SELECT count(*) FROM custody_operations) AS operation_locks
        WHERE operation.run_id = $1
          AND operation.provider = 'modal'
          AND operation.state = 'reserved'
@@ -11890,7 +11999,21 @@ export async function claimPostgresSchedulerSandboxCleanupJobs(
     diagnostic_json: unknown | null;
   }>(
     `
-      WITH candidates AS (
+      WITH candidate_jobs AS MATERIALIZED (
+        SELECT cleanup.provider, cleanup.sandbox_id, cleanup.routing_domain_key, cleanup.run_id
+        FROM ${tableName('sandbox_cleanup_jobs_v2', options)} AS cleanup
+        WHERE cleanup.status <> 'deleted' AND cleanup.blocked_at IS NULL
+          AND (cleanup.available_at <= now() OR cleanup.diagnostic_json IS NOT NULL)
+          AND (cleanup.lease_expires_at IS NULL OR cleanup.lease_expires_at <= now())
+          AND ($4::text IS NULL OR cleanup.environment = $4)
+          AND ($5::text[] IS NULL OR cleanup.provider = ANY($5::text[]))
+          AND ($8::text IS NULL OR cleanup.run_id = $8)
+          AND ($9::integer IS NULL OR cleanup.attempt <= $9)
+          AND NOT ${protectedDetachedSandboxSql({ runId: 'cleanup.run_id', provider: 'cleanup.provider', sandboxId: 'cleanup.sandbox_id', routingDomain: 'cleanup.routing_domain' }, options)}
+        ORDER BY CASE WHEN cleanup.diagnostic_json IS NOT NULL THEN 0 ELSE 1 END,
+                 cleanup.available_at, cleanup.created_at, cleanup.provider, cleanup.sandbox_id, cleanup.routing_domain_key
+        LIMIT $2
+      ), ${runtimeSandboxCustodyLocksSql('EXISTS (SELECT 1 FROM candidate_jobs WHERE candidate_jobs.run_id = run.run_id)', options, true)}, candidates AS (
         SELECT
           cleanup.provider,
           cleanup.sandbox_id,
@@ -11898,9 +12021,14 @@ export async function claimPostgresSchedulerSandboxCleanupJobs(
           cleanup.routing_domain_key,
           run.status AS terminal_status
         FROM ${tableName('sandbox_cleanup_jobs_v2', options)} AS cleanup
-        JOIN ${tableName('runs', options)} AS run
+        JOIN custody_runs AS run
           ON run.run_id = cleanup.run_id
-        WHERE cleanup.status <> 'deleted'
+        CROSS JOIN (SELECT count(*) FROM custody_operations) AS operation_locks
+        WHERE EXISTS (SELECT 1 FROM candidate_jobs AS selected
+                        WHERE selected.provider = cleanup.provider
+                          AND selected.sandbox_id = cleanup.sandbox_id
+                          AND selected.routing_domain_key = cleanup.routing_domain_key)
+          AND cleanup.status <> 'deleted'
           AND cleanup.blocked_at IS NULL
           AND (
             cleanup.available_at <= now()
@@ -11911,7 +12039,7 @@ export async function claimPostgresSchedulerSandboxCleanupJobs(
           AND ($5::text[] IS NULL OR cleanup.provider = ANY($5::text[]))
           AND ($8::text IS NULL OR cleanup.run_id = $8)
           AND ($9::integer IS NULL OR cleanup.attempt <= $9)
-          AND NOT ${protectedDetachedSandboxSql({ runId: 'cleanup.run_id', provider: 'cleanup.provider', sandboxId: 'cleanup.sandbox_id', routingDomain: 'cleanup.routing_domain' }, options)}
+          AND NOT ${protectedDetachedSandboxSql({ runId: 'cleanup.run_id', provider: 'cleanup.provider', sandboxId: 'cleanup.sandbox_id', routingDomain: 'cleanup.routing_domain' }, options, true)}
         ORDER BY
           CASE WHEN cleanup.diagnostic_json IS NOT NULL THEN 0 ELSE 1 END ASC,
           cleanup.available_at ASC,
@@ -11948,12 +12076,28 @@ export async function claimPostgresSchedulerSandboxCleanupJobs(
       operation_claimed AS (
         UPDATE ${tableName('runtime_sandbox_operations', options)} AS operation
         SET state = $6,
+            sandbox_id = claimed.sandbox_id,
+            routing_domain = claimed.routing_domain,
+            provider_created_at = COALESCE(operation.provider_created_at, clock_timestamp()),
             updated_at = clock_timestamp()
         FROM claimed
-        WHERE operation.provider = claimed.provider
-          AND operation.sandbox_id = claimed.sandbox_id
-          AND operation.routing_domain IS NOT DISTINCT FROM claimed.routing_domain
-          AND operation.state = $7
+        WHERE operation.run_id = claimed.run_id
+          AND operation.provider = claimed.provider
+          AND (
+            (operation.state = $7
+             AND operation.sandbox_id = claimed.sandbox_id
+             AND operation.routing_domain IS NOT DISTINCT FROM claimed.routing_domain)
+            OR (operation.state = 'reserved'
+                AND operation.sandbox_id IS NULL
+                AND EXISTS (
+                  SELECT 1 FROM custody_leases AS lease
+                  WHERE lease.lease_id = operation.operation_key
+                    AND lease.run_id = claimed.run_id
+                    AND lease.provider = claimed.provider
+                    AND lease.sandbox_id = claimed.sandbox_id
+                    AND lease.routing_domain IS NOT DISTINCT FROM claimed.routing_domain
+                ))
+          )
         RETURNING operation.operation_key
       ),
       claim_events AS (
@@ -12165,12 +12309,14 @@ export async function completePostgresSchedulerSandboxCleanupJob(
 > {
   const completed = await client.query<{ outcome: string }>(
     `
-      WITH current_job AS MATERIALIZED (
+      WITH ${runtimeSandboxCustodyLocksSql('run.run_id = $3', options)}, current_job AS MATERIALIZED (
         SELECT status, attempt_count, lease_owner, routing_domain
-        FROM ${tableName('sandbox_cleanup_jobs_v2', options)}
+        FROM ${tableName('sandbox_cleanup_jobs_v2', options)} AS sandbox_cleanup_jobs_v2
+        CROSS JOIN (SELECT count(*) FROM custody_operations) AS operation_locks
         WHERE provider = $8
           AND sandbox_id = $1
           AND routing_domain IS NOT DISTINCT FROM $9
+        FOR UPDATE OF sandbox_cleanup_jobs_v2
       ),
       completed AS (
         UPDATE ${tableName('sandbox_cleanup_jobs_v2', options)} AS target
@@ -12246,7 +12392,8 @@ export async function completePostgresSchedulerSandboxCleanupJob(
         SET resource_json = (
               SELECT jsonb_agg(
                 CASE
-                  WHEN resource->>'kind' = 'modal_sandbox'
+                  WHEN resource->>'kind' = CASE $8::text WHEN 'modal' THEN 'modal_sandbox' ELSE 'daytona_sandbox' END
+                   AND (CASE $8::text WHEN 'modal' THEN NULLIF(resource->>'modalAppId', '') ELSE NULLIF(resource->>'daytonaOrganizationId', '') END) IS NOT DISTINCT FROM $9::text
                    AND resource->>'sandboxId' = $1
                    AND resource->>'lateAcquired' = 'true'
                    AND (
@@ -12287,7 +12434,8 @@ export async function completePostgresSchedulerSandboxCleanupJob(
                 ELSE jsonb_build_array(attempt.resource_json)
               END
             ) AS candidate
-            WHERE candidate->>'kind' = 'modal_sandbox'
+            WHERE candidate->>'kind' = CASE $8::text WHEN 'modal' THEN 'modal_sandbox' ELSE 'daytona_sandbox' END
+              AND (CASE $8::text WHEN 'modal' THEN NULLIF(candidate->>'modalAppId', '') ELSE NULLIF(candidate->>'daytonaOrganizationId', '') END) IS NOT DISTINCT FROM $9::text
               AND candidate->>'sandboxId' = $1
               AND candidate->>'lateAcquired' = 'true'
           )
@@ -13014,7 +13162,7 @@ export async function completePostgresSchedulerAttempt(
   );
   const completed = await client.query<{ run_id: string }>(
     `
-      WITH completed AS (
+      WITH ${runtimeSandboxCustodyLocksSql('run.run_id = $1', options)}, completed AS (
         UPDATE ${tableName('attempts', options)}
         SET status = 'completed',
             terminalization_mode = 'compatibility_wake',
@@ -13051,6 +13199,7 @@ export async function completePostgresSchedulerAttempt(
               )
             END,
             updated_at = now()
+        FROM (SELECT count(*) FROM custody_attempts) AS attempt_locks
         WHERE run_id = $1
           AND attempt = $2
           AND status = 'running'
@@ -13291,7 +13440,7 @@ export async function failPostgresSchedulerAttempt(
   );
   const failed = await client.query<{ run_id: string }>(
     `
-      WITH failed AS (
+      WITH ${runtimeSandboxCustodyLocksSql('run.run_id = $1', options)}, failed AS (
         UPDATE ${tableName('attempts', options)}
         SET status = 'failed',
             lease_expires_at = NULL,
@@ -13338,6 +13487,7 @@ export async function failPostgresSchedulerAttempt(
               )
             END,
             updated_at = now()
+        FROM (SELECT count(*) FROM custody_attempts) AS attempt_locks
         WHERE run_id = $1
           AND attempt = $2
           AND status = 'running'
@@ -13625,7 +13775,7 @@ export async function cancelPostgresSchedulerRun(
     launch_json: PlaySchedulerSubmitInput;
   }>(
     `
-      WITH cancelled_run AS (
+      WITH ${runtimeSandboxCustodyLocksSql('run.run_id = $1', options)}, cancelled_run AS (
         UPDATE ${tableName('runs', options)} AS target
         SET status = 'cancelled',
             terminal_result_json = NULL,
@@ -13634,6 +13784,7 @@ export async function cancelPostgresSchedulerRun(
             wait_deadline = NULL,
             finished_at = now(),
             updated_at = now()
+        FROM (SELECT count(*) FROM custody_attempts) AS attempt_locks
         WHERE target.run_id = $1
           AND target.status IN ('queued', 'running', 'waiting')
           ${strandedEligibility}