fix(billing): harden sandbox usage failure recovery

Author: Bill Leoutsakos

apps/sim/lib/billing/sandbox-usage-outbox.test.ts

@@ -95,6 +95,13 @@ function costTotalUpdate(index: number): { text: string; params: unknown[] } {
   }
 }
 
+function executedSqlIndex(substring: string): number {
+  return dbChainMockFns.execute.mock.calls.findIndex(([query]) => {
+    const strings = (query as { strings?: readonly string[] } | null)?.strings
+    return Array.isArray(strings) && strings.some((value) => value.includes(substring))
+  })
+}
+
 afterAll(resetDbChainMock)
 
 describe('sandbox usage outbox finalizer', () => {
@@ -126,7 +133,9 @@ describe('sandbox usage outbox finalizer', () => {
         ],
       })
     )
-    expect(dbChainMockFns.execute).toHaveBeenCalledOnce()
+    expect(dbChainMockFns.execute).toHaveBeenCalledTimes(2)
+    expect(executedSqlIndex('lock_timeout')).toBe(0)
+    expect(executedSqlIndex('pg_advisory_xact_lock')).toBe(1)
     expect(dbChainMockFns.update).toHaveBeenCalled()
     expect(costTotalUpdate(0)).toEqual({
       text: 'GREATEST(COALESCE(, 0) + ::numeric, ::numeric)',
@@ -160,6 +169,18 @@ describe('sandbox usage outbox finalizer', () => {
     expect(mockRecordUsage).not.toHaveBeenCalled()
   })
 
+  it('propagates advisory lock timeouts so the generic outbox retries', async () => {
+    const lockTimeout = Object.assign(new Error('canceling statement due to lock timeout'), {
+      code: '55P03',
+    })
+    dbChainMockFns.execute.mockResolvedValueOnce([]).mockRejectedValueOnce(lockTimeout)
+
+    await expect(runHandler(payload())).rejects.toBe(lockTimeout)
+    expect(executedSqlIndex('lock_timeout')).toBe(0)
+    expect(executedSqlIndex('pg_advisory_xact_lock')).toBe(1)
+    expect(mockRecordUsage).not.toHaveBeenCalled()
+  })
+
   it('uses the same ledger event key when finalization is replayed', async () => {
     mockRecordUsage.mockResolvedValueOnce(0.00046).mockResolvedValueOnce(0)
     queueTableRows(usageLog, [{ cost: '0.02' }])

apps/sim/lib/billing/sandbox-usage-outbox.ts

@@ -19,7 +19,7 @@ import {
 import { checkAndBillPayerOverageThreshold } from '@/lib/billing/threshold-billing'
 import { isBillingEnabled } from '@/lib/core/config/env-flags'
 import {
-  enqueueOutboxEvent,
+  enqueueOutboxEventIfAbsent,
   type OutboxEventContext,
   type OutboxHandlerRegistry,
   patchAndReleasePendingOutboxEvent,
@@ -36,6 +36,7 @@ const logger = createLogger('SandboxUsageOutbox')
 
 export const SANDBOX_USAGE_OUTBOX_EVENT_TYPE = 'sandbox.usage.finalize'
 const CRASH_RECOVERY_GRACE_MS = 60_000
+const SANDBOX_USAGE_LOCK_TIMEOUT_MS = 10_000
 
 export type SandboxUsageCleanupStatus = 'active' | 'terminated' | 'pending_reconciliation'
 
@@ -199,7 +200,7 @@ export async function beginSandboxUsage(params: BeginSandboxUsageParams): Promis
     pricing: createSandboxPricingSnapshot(params.provider),
   }
 
-  await enqueueOutboxEvent(db, SANDBOX_USAGE_OUTBOX_EVENT_TYPE, payload, {
+  await enqueueOutboxEventIfAbsent(db, SANDBOX_USAGE_OUTBOX_EVENT_TYPE, payload, {
     id: eventId,
     availableAt: new Date(params.providerExpiresAt.getTime() + CRASH_RECOVERY_GRACE_MS),
   })
@@ -276,6 +277,9 @@ async function finalizeSandboxUsage(
   const billingContext = toBillingContext(payload.billingAttribution)
 
   await db.transaction(async (tx) => {
+    await tx.execute(
+      sql`select set_config('lock_timeout', ${`${SANDBOX_USAGE_LOCK_TIMEOUT_MS}ms`}, true)`
+    )
     await tx.execute(sql`select pg_advisory_xact_lock(hashtextextended(${payload.executionId}, 0))`)
     const insertedCost = await recordUsage({
       userId: payload.billingAttribution.actorUserId,

apps/sim/lib/core/outbox/service.test.ts

@@ -28,6 +28,7 @@ import {
   deferOutboxHandler,
   enqueueOrReschedulePendingOutboxEvent,
   enqueueOutboxEvent,
+  enqueueOutboxEventIfAbsent,
   enqueueOutboxEvents,
   outboxEventHasSourceOperationId,
   outboxPayloadHasSourceOperationId,
@@ -98,6 +99,19 @@ describe('enqueueOutboxEvent', () => {
     )
   })
 
+  it('ensures a stable event exists without failing on a duplicate ID', async () => {
+    const id = await enqueueOutboxEventIfAbsent(
+      dbChainMock.db,
+      'test.event',
+      { foo: 'bar' },
+      { id: 'stable-event-id' }
+    )
+
+    expect(id).toBe('stable-event-id')
+    expect(dbChainMockFns.values.mock.calls[0][0]).toMatchObject({ id: 'stable-event-id' })
+    expect(dbChainMockFns.onConflictDoNothing).toHaveBeenCalledWith({ target: outboxEvent.id })
+  })
+
   it('inserts a bounded event batch in one statement', async () => {
     const ids = await enqueueOutboxEvents(dbChainMock.db, 'test.event', [
       { sequence: 0 },

apps/sim/lib/core/outbox/service.ts

@@ -174,6 +174,34 @@ export async function enqueueOutboxEvent<T>(
   return id
 }
 
+/**
+ * Ensures an outbox event with a caller-owned stable ID exists.
+ *
+ * This is reserved for recovery after an insert returned an indeterminate
+ * result: PostgreSQL may have committed the first insert before the client saw
+ * a connection failure, so retrying must collapse on the event ID instead of
+ * surfacing a duplicate-key error.
+ */
+export async function enqueueOutboxEventIfAbsent<T>(
+  executor: Pick<typeof db, 'insert'>,
+  eventType: string,
+  payload: T,
+  options: EnqueueOptions & { id: string }
+): Promise<string> {
+  await executor
+    .insert(outboxEvent)
+    .values({
+      id: options.id,
+      eventType,
+      payload: payload as never,
+      maxAttempts: options.maxAttempts ?? DEFAULT_MAX_ATTEMPTS,
+      availableAt: options.availableAt ?? new Date(),
+    })
+    .onConflictDoNothing({ target: outboxEvent.id })
+  logger.info('Ensured outbox event exists', { id: options.id, eventType })
+  return options.id
+}
+
 export async function enqueueOutboxEvents<T>(
   executor: Pick<typeof db, 'insert'>,
   eventType: string,

apps/sim/lib/execution/non-retryable-error.ts

@@ -32,7 +32,7 @@ export class SandboxUsagePersistenceError extends NonRetryableExecutionError {
 
   constructor(provider: string, options?: ErrorOptions) {
     super(
-      `${provider} created a Function sandbox, but Sim could not persist its usage record. The sandbox was stopped before user code ran.`,
+      `${provider} created a Function sandbox, but Sim could not persist its usage record. User code was not run, and Sim initiated sandbox cleanup before returning.`,
       options
     )
     this.name = 'SandboxUsagePersistenceError'

apps/sim/lib/execution/remote-sandbox/conformance.test.ts

@@ -153,6 +153,16 @@ import type { SandboxUsageContext } from '@/lib/execution/remote-sandbox/types'
 type Provider = 'e2b' | 'daytona'
 const PROVIDERS: Provider[] = ['e2b', 'daytona']
 
+function deferred<T>() {
+  let resolve!: (value: T | PromiseLike<T>) => void
+  let reject!: (reason?: unknown) => void
+  const promise = new Promise<T>((resolvePromise, rejectPromise) => {
+    resolve = resolvePromise
+    reject = rejectPromise
+  })
+  return { promise, resolve, reject }
+}
+
 const usageContext: SandboxUsageContext = {
   workspaceId: 'ws-1',
   workflowId: 'wf-1',
@@ -405,7 +415,7 @@ describe.each(PROVIDERS)('sandbox conformance [%s]', (provider) => {
     )
   })
 
-  it('kills without running user code when initial usage persistence fails', async () => {
+  it('kills and recovers usage without running user code when initial persistence fails', async () => {
     mockBeginSandboxUsage.mockRejectedValueOnce(new Error('database unavailable'))
 
     await expect(
@@ -424,7 +434,94 @@ describe.each(PROVIDERS)('sandbox conformance [%s]', (provider) => {
     expect(
       provider === 'e2b' ? mockE2BCommandsRun : mockExecuteSessionCommand
     ).not.toHaveBeenCalled()
+    expect(mockBeginSandboxUsage).toHaveBeenCalledTimes(2)
+    expect(mockReleaseAndProcessSandboxUsage).toHaveBeenCalledWith(
+      'sandbox-usage-event',
+      expect.objectContaining({
+        outcome: 'infrastructure_error',
+        cleanupStatus: 'terminated',
+      })
+    )
+  })
+
+  it('preserves usage persistence failure when cancellation races admission', async () => {
+    const controller = new AbortController()
+    const persistence = deferred<string>()
+    mockBeginSandboxUsage.mockImplementationOnce(() => persistence.promise)
+
+    const execution = executeInSandbox({
+      code: 'x',
+      language: CodeLanguage.Python,
+      timeoutMs: 1000,
+      signal: controller.signal,
+      usageContext,
+    })
+    const rejection = expect(execution).rejects.toMatchObject({
+      name: 'SandboxUsagePersistenceError',
+      retryable: false,
+    })
+
+    await vi.waitFor(() => expect(mockBeginSandboxUsage).toHaveBeenCalledOnce())
+    controller.abort(new DOMException('cancelled', 'AbortError'))
+    persistence.reject(new Error('database unavailable'))
+
+    await rejection
+    expect(
+      provider === 'e2b' ? mockE2BCommandsRun : mockExecuteSessionCommand
+    ).not.toHaveBeenCalled()
+  })
+
+  it('durably releases cleanup when persistence and live teardown initially fail', async () => {
+    mockBeginSandboxUsage.mockRejectedValueOnce(new Error('database unavailable'))
+    const teardown = provider === 'e2b' ? mockE2BKill : mockDelete
+    teardown.mockRejectedValue(new Error('provider unavailable'))
+
+    await expect(
+      executeInSandbox({
+        code: 'x',
+        language: CodeLanguage.Python,
+        timeoutMs: 1000,
+        usageContext,
+      })
+    ).rejects.toMatchObject({ name: 'SandboxUsagePersistenceError', retryable: false })
+
+    expect(teardown).toHaveBeenCalledTimes(2)
+    expect(mockBeginSandboxUsage).toHaveBeenCalledTimes(2)
+    expect(mockReleaseAndProcessSandboxUsage).toHaveBeenCalledWith(
+      'sandbox-usage-event',
+      expect.objectContaining({
+        outcome: 'infrastructure_error',
+        cleanupStatus: 'pending_reconciliation',
+      })
+    )
+    expect(
+      provider === 'e2b' ? mockE2BCommandsRun : mockExecuteSessionCommand
+    ).not.toHaveBeenCalled()
+  })
+
+  it('falls back to termination by ID when persistence recovery also fails', async () => {
+    mockBeginSandboxUsage.mockRejectedValue(new Error('database unavailable'))
+    const teardown = provider === 'e2b' ? mockE2BKill : mockDelete
+    teardown.mockRejectedValue(new Error('provider unavailable'))
+
+    await expect(
+      executeInSandbox({
+        code: 'x',
+        language: CodeLanguage.Python,
+        timeoutMs: 1000,
+        usageContext,
+      })
+    ).rejects.toMatchObject({ name: 'SandboxUsagePersistenceError', retryable: false })
+
+    if (provider === 'e2b') {
+      expect(mockE2BStaticKill).toHaveBeenCalledWith('sb_1', { apiKey: 'test-key' })
+    } else {
+      expect(mockDaytonaGet).toHaveBeenCalledWith('sb_1')
+    }
     expect(mockReleaseAndProcessSandboxUsage).not.toHaveBeenCalled()
+    expect(
+      provider === 'e2b' ? mockE2BCommandsRun : mockExecuteSessionCommand
+    ).not.toHaveBeenCalled()
   })
 
   it('releases returned user errors as billable usage', async () => {