fix(billing): address sandbox usage review feedback

Author: Bill Leoutsakos

apps/sim/app/api/function/execute/route.test.ts

@@ -15,6 +15,11 @@ import {
 } from '@sim/testing'
 import { NextRequest } from 'next/server'
 import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
+import {
+  BILLING_ATTRIBUTION_HEADER,
+  type BillingAttributionSnapshot,
+  serializeBillingAttributionHeader,
+} from '@/lib/billing/core/billing-attribution'
 import { INTERNAL_EXECUTION_DEADLINE_HEADER } from '@/lib/execution/execution-deadline-header'
 import {
   MOUNTED_WORKSPACE_FILES_PROVENANCE_KEY,
@@ -39,6 +44,22 @@ function grantedAccess(workspaceId: string) {
   }
 }
 
+function billingAttributionHeaders(workspaceId = 'workspace-1'): Record<string, string> {
+  const attribution: BillingAttributionSnapshot = {
+    actorUserId: 'user-123',
+    workspaceId,
+    organizationId: null,
+    billedAccountUserId: 'user-123',
+    billingEntity: { type: 'user', id: 'user-123' },
+    billingPeriod: {
+      start: '2026-08-01T00:00:00.000Z',
+      end: '2026-09-01T00:00:00.000Z',
+    },
+    payerSubscription: null,
+  }
+  return { [BILLING_ATTRIBUTION_HEADER]: serializeBillingAttributionHeader(attribution) }
+}
+
 const {
   mockExecuteInSandbox,
   mockExecuteInIsolatedVM,
@@ -318,26 +339,8 @@ describe('Function Execute API Route', () => {
       expect(mockWriteWorkspaceFileByPath).not.toHaveBeenCalled()
     })
 
-    it('rejects an export whose workspace is derived from a body-supplied workflowId', async () => {
+    it('fails closed before deriving an export workspace from incomplete workflow context', async () => {
       envFlagsMock.isRemoteSandboxEnabled = true
-      mockExecuteInSandbox.mockResolvedValueOnce({
-        result: 'done',
-        stdout: 'ok',
-        sandboxId: 'sandbox-123',
-        exportedFiles: { '/tmp/out.txt': 'owned by attacker' },
-      })
-      workflowsUtilsMock.getWorkflowById.mockResolvedValueOnce({
-        id: 'workflow-victim',
-        workspaceId: 'workspace-victim',
-      })
-      mockResolveWorkspaceAccess.mockResolvedValue({
-        exists: true,
-        hasAccess: false,
-        canWrite: false,
-        canAdmin: false,
-        workspace: { id: 'workspace-victim' },
-        permission: null,
-      })
 
       const req = createMockRequest('POST', {
         code: 'print("done")',
@@ -350,7 +353,12 @@ describe('Function Execute API Route', () => {
 
       const response = await POST(req)
 
-      expect(response.status).toBe(403)
+      expect(response.status).toBe(503)
+      await expect(response.json()).resolves.toMatchObject({
+        retryable: false,
+        code: 'sandbox_usage_attribution_invalid',
+      })
+      expect(mockExecuteInSandbox).not.toHaveBeenCalled()
       expect(mockWriteWorkspaceFileByPath).not.toHaveBeenCalled()
     })
 
@@ -537,6 +545,85 @@ describe('Function Execute API Route', () => {
       expect(mockExecuteInIsolatedVM).not.toHaveBeenCalled()
     })
 
+    it.each([
+      {
+        label: 'execution ID',
+        body: { workflowId: 'workflow-1', workspaceId: 'workspace-1' },
+      },
+      {
+        label: 'trusted billing attribution',
+        body: {
+          workflowId: 'workflow-1',
+          workspaceId: 'workspace-1',
+          executionId: 'execution-1',
+        },
+      },
+    ])(
+      'fails before remote sandbox creation when workflow usage lacks $label',
+      async ({ body }) => {
+        envFlagsMock.isRemoteSandboxEnabled = true
+
+        const response = await POST(
+          createMockRequest('POST', { code: 'print("ready")', language: 'python', ...body })
+        )
+
+        expect(response.status).toBe(503)
+        await expect(response.json()).resolves.toMatchObject({
+          success: false,
+          retryable: false,
+          code: 'sandbox_usage_attribution_invalid',
+        })
+        expect(mockExecuteInSandbox).not.toHaveBeenCalled()
+      }
+    )
+
+    it('passes complete trusted workflow attribution to remote sandbox billing', async () => {
+      envFlagsMock.isRemoteSandboxEnabled = true
+
+      const response = await POST(
+        createMockRequest(
+          'POST',
+          {
+            code: 'print("ready")',
+            language: 'python',
+            workflowId: 'workflow-1',
+            workspaceId: 'workspace-1',
+            executionId: 'execution-1',
+          },
+          billingAttributionHeaders()
+        )
+      )
+
+      expect(response.status).toBe(200)
+      expect(mockExecuteInSandbox).toHaveBeenCalledWith(
+        expect.objectContaining({
+          usageContext: expect.objectContaining({
+            workspaceId: 'workspace-1',
+            workflowId: 'workflow-1',
+            executionId: 'execution-1',
+            billingAttribution: expect.objectContaining({ actorUserId: 'user-123' }),
+          }),
+        })
+      )
+    })
+
+    it('allows a non-workflow remote call without sandbox usage attribution', async () => {
+      envFlagsMock.isRemoteSandboxEnabled = true
+
+      const response = await POST(
+        createMockRequest('POST', {
+          code: 'print("ready")',
+          language: 'python',
+          workspaceId: 'workspace-1',
+        })
+      )
+
+      expect(response.status).toBe(200)
+      expect(mockExecuteInSandbox).toHaveBeenCalledWith(
+        expect.objectContaining({ usageContext: undefined })
+      )
+    })
+
     it('forces import-free JavaScript into the remote runtime when a Sim sandbox is selected', async () => {
       envFlagsMock.isRemoteSandboxEnabled = true
 

apps/sim/lib/billing/core/usage-log.test.ts

@@ -86,7 +86,7 @@ describe('recordUsage', () => {
   })
 
   it('commits canonical usage rows with deterministic event keys and billing scope', async () => {
-    await recordUsage({
+    const insertedCost = await recordUsage({
       userId: 'external-actor',
       workspaceId: 'workspace-1',
       billingEntity: { type: 'organization', id: 'workspace-org' },
@@ -124,9 +124,34 @@ describe('recordUsage', () => {
     expect(mockOnConflictDoNothing.mock.calls[0][0]).toMatchObject({
       target: usageLog.eventKey,
     })
+    expect(insertedCost).toBeCloseTo(0.3, 8)
     expect(mockGetHighestPrioritySubscription).not.toHaveBeenCalled()
   })
 
+  it('returns zero when idempotency skips a duplicate event', async () => {
+    mockReturning.mockResolvedValueOnce([])
+
+    const insertedCost = await recordUsage({
+      userId: 'user-1',
+      billingEntity: { type: 'user', id: 'user-1' },
+      billingPeriod: {
+        start: new Date('2026-05-01T00:00:00.000Z'),
+        end: new Date('2026-06-01T00:00:00.000Z'),
+      },
+      entries: [
+        {
+          category: 'tool',
+          source: 'workflow',
+          description: 'Code sandbox',
+          cost: 0.1,
+          eventKey: 'sandbox-event',
+        },
+      ],
+    })
+
+    expect(insertedCost).toBe(0)
+  })
+
   it('uses pre-resolved billing context without loading subscriptions', async () => {
     await recordUsage({
       userId: 'user-1',

apps/sim/lib/billing/core/usage-log.ts

@@ -381,7 +381,7 @@ export async function getStampedPeriodRangeUsageCostByUser(
  * as the pre-cutover period baseline and for low-frequency billing trackers,
  * but usage writes no longer contend on the user_stats row.
  */
-export async function recordUsage(params: RecordUsageParams): Promise<void> {
+export async function recordUsage(params: RecordUsageParams): Promise<number> {
   // The usage ledger is written regardless of BILLING_ENABLED so it is the
   // single, universal source of truth for cost (including self-hosted, where
   // it powers the logs-page cost display). Billing *enforcement* (Stripe /
@@ -400,7 +400,7 @@ export async function recordUsage(params: RecordUsageParams): Promise<void> {
   const validEntries = entries.filter((e) => e.cost > 0)
 
   if (validEntries.length === 0) {
-    return
+    return 0
   }
 
   if (workspaceId && (!billingEntity || !billingPeriod)) {
@@ -473,6 +473,8 @@ export async function recordUsage(params: RecordUsageParams): Promise<void> {
     entryCount: validEntries.length,
     sources: [...new Set(validEntries.map((e) => e.source))],
   })
+
+  return insertedCost
 }
 
 /**

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

@@ -81,14 +81,28 @@ async function runHandler(value: SandboxUsageOutboxPayloadV1): Promise<void> {
   await sandboxUsageOutboxHandlers[SANDBOX_USAGE_OUTBOX_EVENT_TYPE](value, context)
 }
 
+function costTotalUpdate(index: number): { text: string; params: unknown[] } {
+  const updates = dbChainMockFns.set.mock.calls
+    .map(([value]) => value)
+    .filter(
+      (value): value is { costTotal: { strings: string[]; values: unknown[] } } =>
+        typeof value === 'object' && value !== null && 'costTotal' in value
+    )
+  const update = updates[index].costTotal
+  return {
+    text: update.strings.join(''),
+    params: update.values.filter((value) => typeof value === 'string' || typeof value === 'number'),
+  }
+}
+
 afterAll(resetDbChainMock)
 
 describe('sandbox usage outbox finalizer', () => {
   beforeEach(() => {
     vi.clearAllMocks()
     resetDbChainMock()
     queueTableRows(usageLog, [{ cost: '0.02' }])
-    mockRecordUsage.mockResolvedValue(undefined)
+    mockRecordUsage.mockResolvedValue(0.00046)
     mockTerminateById.mockResolvedValue('terminated')
   })
 
@@ -114,6 +128,10 @@ describe('sandbox usage outbox finalizer', () => {
     )
     expect(dbChainMockFns.execute).toHaveBeenCalledOnce()
     expect(dbChainMockFns.update).toHaveBeenCalled()
+    expect(costTotalUpdate(0)).toEqual({
+      text: 'GREATEST(COALESCE(, 0) + ::numeric, ::numeric)',
+      params: ['workflowExecutionLogs.costTotal', '0.00046', '0.02'],
+    })
   })
 
   it('terminates and checkpoints a sandbox whose terminal timestamp is missing', async () => {
@@ -143,11 +161,20 @@ describe('sandbox usage outbox finalizer', () => {
   })
 
   it('uses the same ledger event key when finalization is replayed', async () => {
+    mockRecordUsage.mockResolvedValueOnce(0.00046).mockResolvedValueOnce(0)
+    queueTableRows(usageLog, [{ cost: '0.02' }])
+
     await runHandler(payload())
     await runHandler(payload())
 
     const firstEventKey = mockRecordUsage.mock.calls[0][0].entries[0].eventKey
     const secondEventKey = mockRecordUsage.mock.calls[1][0].entries[0].eventKey
     expect(firstEventKey).toBe(secondEventKey)
+    expect(costTotalUpdate(0).params).toEqual([
+      'workflowExecutionLogs.costTotal',
+      '0.00046',
+      '0.02',
+    ])
+    expect(costTotalUpdate(1).params).toEqual(['workflowExecutionLogs.costTotal', '0', '0.02'])
   })
 })

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

@@ -277,7 +277,7 @@ async function finalizeSandboxUsage(
 
   await db.transaction(async (tx) => {
     await tx.execute(sql`select pg_advisory_xact_lock(hashtextextended(${payload.executionId}, 0))`)
-    await recordUsage({
+    const insertedCost = await recordUsage({
       userId: payload.billingAttribution.actorUserId,
       entries: [
         {
@@ -323,7 +323,7 @@ async function finalizeSandboxUsage(
     await tx
       .update(workflowExecutionLogs)
       .set({
-        costTotal: sql`GREATEST(COALESCE(${workflowExecutionLogs.costTotal}, 0), ${ledgerCost})`,
+        costTotal: sql`GREATEST(COALESCE(${workflowExecutionLogs.costTotal}, 0) + ${insertedCost.toString()}::numeric, ${ledgerCost}::numeric)`,
       })
       .where(eq(workflowExecutionLogs.executionId, payload.executionId))
   })

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

@@ -39,6 +39,19 @@ export class SandboxUsagePersistenceError extends NonRetryableExecutionError {
   }
 }
 
+/** A workflow sandbox request lacks the trusted context required for durable billing. */
+export class SandboxUsageAttributionError extends NonRetryableExecutionError {
+  readonly code = 'sandbox_usage_attribution_invalid' as const
+
+  constructor(options?: ErrorOptions) {
+    super(
+      'Sim could not establish trusted billing attribution for this Function sandbox. The sandbox was not created.',
+      options
+    )
+    this.name = 'SandboxUsageAttributionError'
+  }
+}
+
 export function isNonRetryableExecutionError(error: unknown): boolean {
   return Boolean(
     findCause(error, (cause): cause is NonRetryableExecutionError => {
@@ -64,3 +77,9 @@ export function isSandboxUsagePersistenceError(
 ): error is SandboxUsagePersistenceError {
   return Boolean(findCause(error, (cause) => cause instanceof SandboxUsagePersistenceError))
 }
+
+export function isSandboxUsageAttributionError(
+  error: unknown
+): error is SandboxUsageAttributionError {
+  return Boolean(findCause(error, (cause) => cause instanceof SandboxUsageAttributionError))
+}

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

@@ -650,8 +650,17 @@ describe.each(PROVIDERS)('sandbox conformance [%s]', (provider) => {
     }
 
     await expect(
-      executeInSandbox({ code: 'sleep()', language: CodeLanguage.Python, timeoutMs: 1000 })
+      executeInSandbox({
+        code: 'sleep()',
+        language: CodeLanguage.Python,
+        timeoutMs: 1000,
+        usageContext,
+      })
     ).rejects.toMatchObject({ name: 'AbortError', message: 'timeout' })
+    expect(mockReleaseAndProcessSandboxUsage).toHaveBeenCalledWith(
+      'sandbox-usage-event',
+      expect.objectContaining({ outcome: 'timeout' })
+    )
   })
 
   it('normalizes JavaScript code budget expiry to a typed timeout abort', async () => {
@@ -994,8 +1003,12 @@ describe.each(PROVIDERS)('sandbox conformance [%s]', (provider) => {
     }
 
     await expect(
-      executeShellInSandbox({ code: 'sleep infinity', timeoutMs: 1000 })
+      executeShellInSandbox({ code: 'sleep infinity', timeoutMs: 1000, usageContext })
     ).rejects.toMatchObject({ name: 'AbortError', message: 'timeout' })
+    expect(mockReleaseAndProcessSandboxUsage).toHaveBeenCalledWith(
+      'sandbox-usage-event',
+      expect.objectContaining({ outcome: 'timeout' })
+    )
   })
 
   it('preserves a user process exit code 124 as an ordinary failure', async () => {

apps/sim/lib/execution/remote-sandbox/index.ts

@@ -686,6 +686,7 @@ async function executeInSandboxWithinBudget(
       throw error
     }
     throwIfAborted(signal)
+    if (execution.timedOut) outcome = 'timeout'
     throwIfSandboxTimedOut(execution)
 
     if (execution.error) {
@@ -822,6 +823,7 @@ async function executeShellInSandboxWithinBudget(
       throw error
     }
     throwIfAborted(signal)
+    if (result.timedOut) outcome = 'timeout'
     throwIfSandboxTimedOut(result)
 
     const stdout = [result.stdout, result.stderr].filter(Boolean).join('\n')