fix(sandboxes): claim the image row and its reference check in one statement

Greptile P1. Reading references in one statement and deleting in another left a
window — a wide one, since a provider delete is a network call — where a second
workspace could declare the same package list, inherit the `ready` row, and have
its next run fail against a template already on its way out. Content addressing
is what makes that reachable: the image is shared, so one workspace's delete can
strand another's sandbox.

The reference check now lives in the conditional DELETE itself, so winning the
delete is the proof that nothing referenced the hash. A workspace that adopts the
hash first makes the delete match nothing and the release becomes a no-op.

Claiming the row before the provider call would otherwise strand a template
nothing points at if the provider then refused, so that path puts the row back
and the retention sweep inherits the retry — the same property the previous
ordering had.

The sweep is deliberately left as it is: its equivalent window needs a hash
unreferenced AND unused for 30 days, and its provider-first ordering encodes the
documented retry-on-refusal behaviour this path now reproduces explicitly.

No transaction is opened. The provider call sits between discrete statements
rather than inside one, so no pooled connection is held across it — which is why
this uses a conditional delete instead of the repo's `pg_advisory_xact_lock`
pattern, whose lock only releases at commit.

Co-Authored-By: Claude <noreply@anthropic.com>

Author: icecrasher321

apps/sim/lib/execution/remote-sandbox/image-registry.test.ts

@@ -8,18 +8,15 @@
  */
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
-const { mockSelect, mockDelete, mockInsert, mockDeleteImage, mockProviderStrategy } = vi.hoisted(
-  () => ({
-    mockSelect: vi.fn(),
-    mockDelete: vi.fn(),
-    mockInsert: vi.fn(),
-    mockDeleteImage: vi.fn(),
-    mockProviderStrategy: { current: 'prebuilt' as 'prebuilt' | 'runtime' },
-  })
-)
+const { mockDelete, mockInsert, mockDeleteImage, mockProviderStrategy } = vi.hoisted(() => ({
+  mockDelete: vi.fn(),
+  mockInsert: vi.fn(),
+  mockDeleteImage: vi.fn(),
+  mockProviderStrategy: { current: 'prebuilt' as 'prebuilt' | 'runtime' },
+}))
 
 vi.mock('@sim/db', () => ({
-  db: { select: mockSelect, delete: mockDelete, insert: mockInsert },
+  db: { delete: mockDelete, insert: mockInsert },
 }))
 
 vi.mock('@sim/db/schema', () => ({
@@ -43,6 +40,7 @@ vi.mock('drizzle-orm', () => ({
   eq: (...args: unknown[]) => args,
   inArray: (...args: unknown[]) => args,
   lt: (...args: unknown[]) => args,
+  notInArray: (...args: unknown[]) => args,
   or: (...args: unknown[]) => args,
   sql: (...args: unknown[]) => args,
 }))
@@ -67,16 +65,6 @@ import {
   releaseSandboxImage,
 } from '@/lib/execution/remote-sandbox/image-registry'
 
-/** Queues the rows each successive `db.select()` chain resolves to. */
-function queueSelects(...results: unknown[][]) {
-  mockSelect.mockReset()
-  for (const rows of results) {
-    mockSelect.mockReturnValueOnce({
-      from: () => ({ where: () => ({ limit: () => Promise.resolve(rows) }) }),
-    })
-  }
-}
-
 const READY_IMAGE = {
   id: 'img-1',
   status: 'ready',
@@ -88,13 +76,34 @@ const READY_IMAGE = {
 beforeEach(() => {
   vi.clearAllMocks()
   mockProviderStrategy.current = 'prebuilt'
-  mockDelete.mockReturnValue({ where: () => Promise.resolve() })
+  mockDelete.mockReturnValue({ where: () => ({ returning: () => Promise.resolve([]) }) })
+  mockInsert.mockReturnValue({ values: () => ({ onConflictDoNothing: () => Promise.resolve() }) })
   mockDeleteImage.mockResolvedValue(undefined)
 })
 
+/** Serializes the mocked predicate tree so a clause can be asserted by shape. */
+function predicateText(predicate: unknown): string {
+  if (predicate == null) return ''
+  if (Array.isArray(predicate)) return predicate.map(predicateText).join(' ')
+  if (typeof predicate === 'object') return JSON.stringify(predicate)
+  return String(predicate)
+}
+
+/** Captures the conditional-delete predicate and what the claim resolves to. */
+function stubClaim(rows: unknown[]): () => unknown {
+  let captured: unknown
+  mockDelete.mockReturnValue({
+    where: (predicate: unknown) => {
+      captured = predicate
+      return { returning: () => Promise.resolve(rows) }
+    },
+  })
+  return () => captured
+}
+
 describe('releaseSandboxImage', () => {
-  it('deletes the provider image and its row when nothing references the hash', async () => {
-    queueSelects([], [READY_IMAGE])
+  it('deletes the provider image once the row is claimed', async () => {
+    stubClaim([READY_IMAGE])
 
     await releaseSandboxImage('hash-1')
 
@@ -103,64 +112,73 @@ describe('releaseSandboxImage', () => {
       buildId: 'build-1',
       providerImageId: 'tmpl-1',
     })
-    expect(mockDelete).toHaveBeenCalledTimes(1)
   })
 
   /**
-   * The case that would break a bystander: two workspaces declaring the same
-   * package list share one build, so one workspace deleting its sandbox must not
-   * delete the image out from under the other.
+   * The bystander case: two workspaces declaring the same package list share one
+   * build, so one workspace's delete must not take the image out from under the
+   * other. The guard is the conditional delete itself — reading references in a
+   * separate statement left a window, spanning a provider network call, in which
+   * another workspace could adopt the hash between the check and the delete.
    */
-  it('leaves the image alone while another sandbox still declares that package list', async () => {
-    // A releasable image IS queued behind the reference lookup on purpose: without
-    // it, dropping the guard would fault on the missing second select and get
-    // swallowed, and this test would pass for the wrong reason.
-    queueSelects([{ id: 'sbx-other' }], [READY_IMAGE])
+  it('claims only when no sandbox references the hash, in one statement', async () => {
+    const read = stubClaim([READY_IMAGE])
 
     await releaseSandboxImage('hash-1')
 
-    expect(mockDeleteImage).not.toHaveBeenCalled()
-    expect(mockDelete).not.toHaveBeenCalled()
+    const clause = predicateText(read())
+    expect(clause).toContain('not exists')
+    expect(clause).toContain('workspace_sandbox')
   })
 
-  it('no-ops under a runtime provider, which has no images to release', async () => {
-    mockProviderStrategy.current = 'runtime'
-    queueSelects([], [READY_IMAGE])
+  it('excludes an in-flight build from the claim rather than racing it', async () => {
+    const read = stubClaim([READY_IMAGE])
+
+    await releaseSandboxImage('hash-1')
+
+    const clause = predicateText(read())
+    expect(clause).toContain('pending')
+    expect(clause).toContain('building')
+  })
+
+  it('touches the provider only when the claim actually took a row', async () => {
+    stubClaim([])
 
     await releaseSandboxImage('hash-1')
 
-    expect(mockSelect).not.toHaveBeenCalled()
     expect(mockDeleteImage).not.toHaveBeenCalled()
   })
 
-  it.each(['pending', 'building'])(
-    'leaves a %s build for the sweep rather than racing it',
-    async (status) => {
-      queueSelects([], [{ ...READY_IMAGE, status }])
+  it('no-ops under a runtime provider, which has no images to release', async () => {
+    mockProviderStrategy.current = 'runtime'
+    stubClaim([READY_IMAGE])
 
-      await releaseSandboxImage('hash-1')
+    await releaseSandboxImage('hash-1')
 
-      expect(mockDeleteImage).not.toHaveBeenCalled()
-      expect(mockDelete).not.toHaveBeenCalled()
-    }
-  )
+    expect(mockDelete).not.toHaveBeenCalled()
+    expect(mockDeleteImage).not.toHaveBeenCalled()
+  })
 
-  it('keeps the row when the provider refuses, so the sweep retries', async () => {
-    queueSelects([], [READY_IMAGE])
+  /**
+   * Claiming before the provider call means a refusal would otherwise strand a
+   * template nothing points at, so the row goes back and the sweep inherits it.
+   */
+  it('restores the claimed row when the provider refuses', async () => {
+    stubClaim([READY_IMAGE])
     mockDeleteImage.mockRejectedValue(new Error('E2B unreachable'))
 
     await expect(releaseSandboxImage('hash-1')).resolves.toBeUndefined()
 
-    expect(mockDelete).not.toHaveBeenCalled()
+    expect(mockInsert).toHaveBeenCalledTimes(1)
   })
 
-  it('does nothing when the hash has no build row at all', async () => {
-    queueSelects([], [])
+  it('skips the provider when the claimed row never had an image', async () => {
+    stubClaim([{ ...READY_IMAGE, imageRef: null }])
 
     await releaseSandboxImage('hash-1')
 
     expect(mockDeleteImage).not.toHaveBeenCalled()
-    expect(mockDelete).not.toHaveBeenCalled()
+    expect(mockInsert).not.toHaveBeenCalled()
   })
 })
 

apps/sim/lib/execution/remote-sandbox/image-registry.ts

@@ -1,11 +1,11 @@
 import { db } from '@sim/db'
-import { sandboxImage, workspaceSandbox } from '@sim/db/schema'
+import { sandboxImage } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
 import { getErrorMessage, toError } from '@sim/utils/errors'
 import { sleep } from '@sim/utils/helpers'
 import { generateId } from '@sim/utils/id'
 import { backoffWithJitter } from '@sim/utils/retry'
-import { and, eq, inArray, lt, or, sql } from 'drizzle-orm'
+import { and, eq, inArray, lt, notInArray, or, sql } from 'drizzle-orm'
 import { isTriggerDevEnabled } from '@/lib/core/config/env-flags'
 import { runDetached } from '@/lib/core/utils/background'
 import {
@@ -16,7 +16,7 @@ import {
 import { resolveProvider } from '@/lib/execution/remote-sandbox/provider'
 import { invalidateSandboxResolution } from '@/lib/execution/remote-sandbox/resolve'
 import type { SandboxSpec } from '@/lib/execution/remote-sandbox/sandbox-spec'
-import type { SandboxImageBuild } from '@/lib/execution/remote-sandbox/types'
+import type { SandboxImageBuild, SandboxImageStatus } from '@/lib/execution/remote-sandbox/types'
 
 const logger = createLogger('SandboxImageRegistry')
 
@@ -287,52 +287,76 @@ export async function runSandboxImageBuild(payload: SandboxImageBuildPayload): P
  * image sits in provider storage until the retention sweep, which is up to
  * `SANDBOX_IMAGE_RETENTION_DAYS` of paying to store something nothing can select.
  *
- * The reference check is what makes deleting this eagerly safe. Builds are keyed
- * by content, not by workspace, so two workspaces declaring the same package list
- * share one image, and deleting on the strength of one workspace's action would
- * break the other. An in-flight build is left alone rather than raced; the sweep
- * collects it once it settles.
+ * Builds are keyed by content, not by workspace, so two workspaces declaring the
+ * same package list share one image and deleting on the strength of one
+ * workspace's action would break the other. The reference check is therefore part
+ * of the same statement that removes the row: reading references first and
+ * deleting second left a window — wide, because a provider delete is a network
+ * call — in which another workspace could adopt the hash, inherit a `ready` row,
+ * and have its next run fail against a template already on its way out. Winning
+ * the conditional delete is what proves nothing referenced the hash.
+ *
+ * Claiming the row before the provider call means a provider that then refuses
+ * would strand a template nothing points at, so the row is put back and the
+ * retention sweep inherits the retry. An in-flight build is left alone rather
+ * than raced; the sweep collects it once it settles.
  *
  * Best-effort by contract: failures are logged and swallowed, because this runs
  * after the mutation it follows has already committed and must never turn a
- * successful delete into an error. The sweep stays the backstop.
+ * successful delete into an error.
  */
 export async function releaseSandboxImage(specHash: string): Promise<void> {
   const provider = resolveProvider()
   if (provider.dependencyStrategy !== 'prebuilt' || !provider.images) return
   const images = provider.images
 
   try {
-    const [referenced] = await db
-      .select({ id: workspaceSandbox.id })
-      .from(workspaceSandbox)
-      .where(eq(workspaceSandbox.specHash, specHash))
-      .limit(1)
-    if (referenced) return
-
-    const [image] = await db
-      .select({
+    const [claimed] = await db
+      .delete(sandboxImage)
+      .where(
+        and(
+          eq(sandboxImage.provider, provider.id),
+          eq(sandboxImage.specHash, specHash),
+          notInArray(sandboxImage.status, ['pending', 'building']),
+          sql`not exists (select 1 from workspace_sandbox ws where ws.spec_hash = ${specHash})`
+        )
+      )
+      .returning({
         id: sandboxImage.id,
+        spec: sandboxImage.spec,
         status: sandboxImage.status,
         imageRef: sandboxImage.imageRef,
         buildId: sandboxImage.buildId,
         providerImageId: sandboxImage.providerImageId,
       })
-      .from(sandboxImage)
-      .where(and(eq(sandboxImage.provider, provider.id), eq(sandboxImage.specHash, specHash)))
-      .limit(1)
-    if (!image) return
-    if (image.status === 'pending' || image.status === 'building') return
+    if (!claimed) return
+
+    invalidateSandboxResolution()
 
-    if (image.imageRef) {
+    if (!claimed.imageRef) return
+    try {
       await images.deleteImage({
-        imageRef: image.imageRef,
-        buildId: image.buildId ?? '',
-        providerImageId: image.providerImageId ?? undefined,
+        imageRef: claimed.imageRef,
+        buildId: claimed.buildId ?? '',
+        providerImageId: claimed.providerImageId ?? undefined,
       })
+    } catch (error) {
+      await db
+        .insert(sandboxImage)
+        .values({
+          id: claimed.id,
+          provider: provider.id,
+          specHash,
+          spec: claimed.spec,
+          status: claimed.status as SandboxImageStatus,
+          imageRef: claimed.imageRef,
+          buildId: claimed.buildId,
+          providerImageId: claimed.providerImageId,
+        })
+        .onConflictDoNothing()
+      throw error
     }
-    await db.delete(sandboxImage).where(eq(sandboxImage.id, image.id))
-    invalidateSandboxResolution()
+
     logger.info('Released unreferenced sandbox image', { specHash })
   } catch (error) {
     logger.warn('Failed to release sandbox image; the retention sweep will retry', {