feat(billing): track E2B and Daytona Function sandbox usage

apps/sim/executor/execution/block-executor.retry.test.ts

@@ -5,11 +5,13 @@
  * client has already seen and cannot re-run the deterministic post-processing.
  */
 import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { buildTraceSpans } from '@/lib/logs/execution/trace-spans/trace-spans'
 import { BlockType, EDGE } from '@/executor/constants'
 import type { DAGNode } from '@/executor/dag/builder'
 import { BlockExecutor } from '@/executor/execution/block-executor'
 import { ExecutionState } from '@/executor/execution/state'
 import type { BlockHandler, ExecutionContext } from '@/executor/types'
+import { attachTrustedExecutionCost } from '@/executor/utils/errors'
 import { VariableResolver } from '@/executor/variables/resolver'
 import type { SerializedBlock, SerializedWorkflow } from '@/serializer/types'
 
@@ -137,6 +139,68 @@ describe('BlockExecutor retry', () => {
     expect(ctx.blockLogs[0]?.tries).toBe(2)
   })
 
+  it('adds the trusted cost of failed Function tries to the successful result', async () => {
+    const block = createBlock(enabled)
+    const firstFailure = new Error('first attempt failed')
+    attachTrustedExecutionCost(firstFailure, { input: 0, output: 0, total: 0.125 })
+    const successfulOutput = {
+      result: 'done',
+      cost: { input: 0, output: 0, total: 0.25 },
+    }
+    attachTrustedExecutionCost(successfulOutput, successfulOutput.cost)
+    const execute = vi
+      .fn()
+      .mockRejectedValueOnce(firstFailure)
+      .mockResolvedValueOnce(successfulOutput)
+    const state = new ExecutionState()
+    const ctx = createContext(state)
+    const executor = buildExecutor(block, { canHandle: () => true, execute }, state)
+
+    const output = await executor.execute(ctx, createNode(block), block)
+
+    expect(execute).toHaveBeenCalledTimes(2)
+    expect(output.cost).toEqual({ input: 0, output: 0, total: 0.375 })
+    expect(ctx.blockLogs[0]?.output?.cost).toEqual(output.cost)
+  })
+
+  it('keeps earlier trusted Function costs when the final try is an infrastructure error', async () => {
+    const block = createBlock(enabled)
+    const firstFailure = new Error('first Function attempt failed')
+    const secondFailure = new Error('second Function attempt failed')
+    const finalFailure = new Error('provider unavailable')
+    attachTrustedExecutionCost(firstFailure, { input: 0, output: 0, total: 0.125 })
+    attachTrustedExecutionCost(secondFailure, { input: 0, output: 0, total: 0.25 })
+    const execute = vi
+      .fn()
+      .mockRejectedValueOnce(firstFailure)
+      .mockRejectedValueOnce(secondFailure)
+      .mockRejectedValueOnce(finalFailure)
+    const state = new ExecutionState()
+    const ctx = createContext(state)
+    const executor = buildExecutor(block, { canHandle: () => true, execute }, state)
+
+    await expect(executor.execute(ctx, createNode(block), block)).rejects.toThrow(
+      'provider unavailable'
+    )
+
+    expect(execute).toHaveBeenCalledTimes(3)
+    expect(ctx.blockLogs[0]?.output).toEqual({
+      error: 'provider unavailable',
+      cost: { input: 0, output: 0, total: 0.375 },
+    })
+
+    const { traceSpans } = buildTraceSpans({
+      success: false,
+      output: { error: 'provider unavailable' },
+      error: 'provider unavailable',
+      logs: ctx.blockLogs,
+    })
+    expect(traceSpans[0]).toMatchObject({
+      status: 'error',
+      cost: { input: 0, output: 0, total: 0.375 },
+    })
+  })
+
   it('stops at maxTries and rethrows the final error unchanged', async () => {
     const block = createBlock(enabled)
     const failure = new Error('still failing')

apps/sim/executor/execution/block-executor.ts

@@ -48,7 +48,13 @@ import {
   type StreamingExecution,
 } from '@/executor/types'
 import { streamingResponseFormatProcessor } from '@/executor/utils'
-import { buildBlockExecutionError, normalizeError } from '@/executor/utils/errors'
+import {
+  attachTrustedExecutionCost,
+  buildBlockExecutionError,
+  normalizeError,
+  readTrustedExecutionCost,
+  type TrustedExecutionCost,
+} from '@/executor/utils/errors'
 import {
   buildUnifiedParentIterations,
   getIterationContext,
@@ -76,6 +82,20 @@ import { SYSTEM_SUBBLOCK_IDS } from '@/triggers/constants'
 
 const logger = createLogger('BlockExecutor')
 
+function addTrustedExecutionCosts(
+  accumulated: TrustedExecutionCost | undefined,
+  current: TrustedExecutionCost | undefined
+): TrustedExecutionCost | undefined {
+  if (!accumulated) return current
+  if (!current) return accumulated
+
+  return {
+    input: accumulated.input + current.input,
+    output: accumulated.output + current.output,
+    total: accumulated.total + current.total,
+  }
+}
+
 export class BlockExecutor {
   private execLogger: Logger
 
@@ -229,6 +249,17 @@ export class BlockExecutor {
     cleanupSelfReference?.()
 
     let streamingPartialOutput: Record<string, any> | undefined
+    /**
+     * Cost of a handler that already finished, kept for the catch below.
+     *
+     * A Function block's sandbox is paid for the moment it completes, but the
+     * steps after the handler returns — base64 hydration, and large-value
+     * redaction that deliberately throws rather than emit unredacted data — can
+     * still fail the block. The error those raise carries no cost of its own, so
+     * without holding it here the completed sandbox would go unbilled. Hoisted
+     * for the same reason `streamingPartialOutput` above is.
+     */
+    let completedHandlerCost: TrustedExecutionCost | undefined
     try {
       /**
        * Only the handler call is retried. A streaming handler returns before any
@@ -241,6 +272,8 @@ export class BlockExecutor {
           : handler.execute(blockCtx, block, resolvedInputs, nodeMetadata)
       )
 
+      completedHandlerCost = readTrustedExecutionCost(output)
+
       const isStreamingExecution =
         output && typeof output === 'object' && 'stream' in output && 'execution' in output
 
@@ -416,7 +449,8 @@ export class BlockExecutor {
           inputDisplayRegistry,
           isSentinel,
           'execution',
-          streamingPartialOutput
+          streamingPartialOutput,
+          completedHandlerCost
         )
       } finally {
         commitBlockRegistry()
@@ -506,15 +540,40 @@ export class BlockExecutor {
     const policy = resolveBlockRetryPolicy(block)
     if (!policy) return invoke()
 
+    const shouldAccumulateFunctionCost = block.metadata?.id === BlockType.FUNCTION
+    let accumulatedFunctionCost: TrustedExecutionCost | undefined
     let tries = 0
     try {
       for (;;) {
         tries++
         try {
-          return await invoke()
+          const output = await invoke()
+          if (!shouldAccumulateFunctionCost || !accumulatedFunctionCost || !isRecordLike(output)) {
+            return output
+          }
+
+          const totalCost = addTrustedExecutionCosts(
+            accumulatedFunctionCost,
+            readTrustedExecutionCost(output)
+          )
+          if (!totalCost) return output
+
+          const outputWithCost = { ...output, cost: totalCost }
+          attachTrustedExecutionCost(outputWithCost, totalCost)
+          return outputWithCost as T
         } catch (error) {
+          if (shouldAccumulateFunctionCost) {
+            accumulatedFunctionCost = addTrustedExecutionCosts(
+              accumulatedFunctionCost,
+              readTrustedExecutionCost(error)
+            )
+          }
+
           const isFinalTry = tries >= policy.maxTries
-          if (isFinalTry || ctx.abortSignal?.aborted || !isRetryableBlockError(error)) throw error
+          if (isFinalTry || ctx.abortSignal?.aborted || !isRetryableBlockError(error)) {
+            attachTrustedExecutionCost(error, accumulatedFunctionCost)
+            throw error
+          }
 
           this.execLogger.warn('Block failed; retrying', {
             blockId: block.id,
@@ -528,7 +587,10 @@ export class BlockExecutor {
           if (policy.waitBetweenTriesMs > 0) await sleep(policy.waitBetweenTriesMs)
 
           /** `sleep` is not abort-aware, so a run stopped mid-wait must not start another try. */
-          if (ctx.abortSignal?.aborted) throw error
+          if (ctx.abortSignal?.aborted) {
+            attachTrustedExecutionCost(error, accumulatedFunctionCost)
+            throw error
+          }
         }
       }
     } finally {
@@ -548,7 +610,8 @@ export class BlockExecutor {
     inputDisplayRegistry: ResolvedSecretTraceRegistry | undefined,
     isSentinel: boolean,
     phase: 'input_resolution' | 'execution',
-    streamingPartialOutput?: Record<string, any>
+    streamingPartialOutput?: Record<string, any>,
+    completedHandlerCost?: TrustedExecutionCost
   ): Promise<NormalizedBlockOutput> {
     const endedAt = new Date().toISOString()
     const duration = performance.now() - startTime
@@ -620,8 +683,10 @@ export class BlockExecutor {
       return softOutput
     }
 
+    const trustedExecutionCost = readTrustedExecutionCost(error) ?? completedHandlerCost
     const errorOutput: NormalizedBlockOutput = {
       error: errorMessage,
+      ...(trustedExecutionCost ? { cost: trustedExecutionCost } : {}),
     }
 
     // Keep any answer text already drained before timeout/failure so logs match

apps/sim/executor/handlers/function/function-handler.test.ts

@@ -1,9 +1,11 @@
 import { beforeEach, describe, expect, it, type Mock, vi } from 'vitest'
 import { createTimeoutAbortController } from '@/lib/core/execution-limits'
 import { DEFAULT_EXECUTION_TIMEOUT_MS } from '@/lib/execution/constants'
+import { NonRetryableExecutionError } from '@/lib/execution/non-retryable-error'
 import { BlockType } from '@/executor/constants'
 import { FunctionBlockHandler } from '@/executor/handlers/function/function-handler'
 import type { ExecutionContext } from '@/executor/types'
+import { readTrustedExecutionCost } from '@/executor/utils/errors'
 import {
   FUNCTION_BLOCK_CONTEXT_VARS_KEY,
   FUNCTION_BLOCK_DISPLAY_CODE_KEY,
@@ -254,6 +256,45 @@ describe('FunctionBlockHandler', () => {
     expect(mockExecuteTool).toHaveBeenCalled()
   })
 
+  it.each([
+    { retryable: true, nonRetryable: false },
+    { retryable: false, nonRetryable: true },
+  ])(
+    'attaches trusted cost to a failed execution when retryable is $retryable',
+    async ({ retryable, nonRetryable }) => {
+      const cost = { input: 0, output: 0, total: 0.125 }
+      mockExecuteTool.mockResolvedValue({
+        success: false,
+        error: 'Remote Function failed',
+        retryable,
+        output: { result: null, stdout: '', cost },
+      })
+
+      let thrown: unknown
+      try {
+        await handler.execute(mockContext, mockBlock, { code: 'throw new Error("failed")' })
+      } catch (error) {
+        thrown = error
+      }
+
+      expect(thrown).toBeInstanceOf(Error)
+      expect(thrown instanceof NonRetryableExecutionError).toBe(nonRetryable)
+      expect(readTrustedExecutionCost(thrown)).toEqual(cost)
+    }
+  )
+
+  it('attaches trusted cost to a successful execution for retry aggregation', async () => {
+    const cost = { input: 0, output: 0, total: 0.25 }
+    mockExecuteTool.mockResolvedValue({
+      success: true,
+      output: { result: 42, stdout: '', cost },
+    })
+
+    const output = await handler.execute(mockContext, mockBlock, { code: 'return 42' })
+
+    expect(readTrustedExecutionCost(output)).toEqual(cost)
+  })
+
   it('should pass runtime context variables to function_execute', async () => {
     const contextVariables = { __blockRef_0: { result: 'from-block' } }
 

apps/sim/executor/handlers/function/function-handler.ts

@@ -12,6 +12,7 @@ import { mergeFileKeys, mergeLargeValueKeys } from '@/lib/execution/payloads/acc
 import { BlockType } from '@/executor/constants'
 import type { BlockHandler, ExecutionContext } from '@/executor/types'
 import { collectBlockData } from '@/executor/utils/block-data'
+import { attachTrustedExecutionCost } from '@/executor/utils/errors'
 import {
   FUNCTION_BLOCK_CONTEXT_VARS_KEY,
   FUNCTION_BLOCK_DISPLAY_CODE_KEY,
@@ -111,15 +112,18 @@ export class FunctionBlockHandler implements BlockHandler {
     const result = await executeTool('function_execute', toolParams, { executionContext: ctx })
 
     if (!result.success) {
-      if (result.retryable === false) {
-        throw new NonRetryableExecutionError(result.error || 'Function execution is indeterminate')
-      }
-      throw new Error(result.error || 'Function execution failed')
+      const error =
+        result.retryable === false
+          ? new NonRetryableExecutionError(result.error || 'Function execution is indeterminate')
+          : new Error(result.error || 'Function execution failed')
+      attachTrustedExecutionCost(error, result.output?.cost)
+      throw error
     }
 
     mergeLargeValueKeys(ctx, result.largeValueKeys ?? [])
     mergeFileKeys(ctx, result.fileKeys ?? [])
 
+    attachTrustedExecutionCost(result.output, result.output?.cost)
     return result.output
   }
 }

apps/sim/executor/handlers/pi/cloud/authoring/backend.ts

@@ -435,7 +435,10 @@ async function runCloudAuthoringPi(
   const lifetimeMs = resolvePiRunLifetimeMs(context.signal)
   const piTimeoutMs = resolvePiTimeoutMs(lifetimeMs)
 
-  const authored = await withPiSandbox<AuthoringPhaseResult>({ lifetimeMs }, async (runner) => {
+  // Bound to a local so the call stays on one line: inlining the second option
+  // reflows this whole callback body and buries the change in re-indentation.
+  const sandboxOptions = { lifetimeMs, cost: context.sandboxCost }
+  const authored = await withPiSandbox<AuthoringPhaseResult>(sandboxOptions, async (runner) => {
     try {
       const clone = await raceAbort(
         runner.run(params.mode === 'cloud' ? CREATE_PR_CLONE_SCRIPT : UPDATE_BRANCH_CLONE_SCRIPT, {

apps/sim/executor/handlers/pi/cloud/babysit/backend.ts

@@ -784,7 +784,7 @@ export async function runBabysitPiWithOptions(
     const lifetimeMs = resolvePiRunLifetimeMs(context.signal)
     const piTimeoutMs = resolvePiTimeoutMs(lifetimeMs)
 
-    return await withPiSandbox({ lifetimeMs }, async (runner) => {
+    return await withPiSandbox({ lifetimeMs, cost: context.sandboxCost }, async (runner) => {
       const clone = await raceAbort(
         runner.run(BABYSIT_CLONE_SCRIPT, {
           envs: {

apps/sim/executor/handlers/pi/cloud/plan/backend.ts

@@ -80,7 +80,7 @@ export const runCloudPlanPi: PiBackendRun<PiCloudPlanRunParams> = async (params,
   const thinking = mapThinkingLevel(params.thinkingLevel) ?? 'medium'
   const lifetimeMs = resolvePiRunLifetimeMs(context.signal)
 
-  return withPiSandbox({ lifetimeMs }, async (runner) => {
+  return withPiSandbox({ lifetimeMs, cost: context.sandboxCost }, async (runner) => {
     try {
       const clone = await raceAbort(
         runner.run(PLAN_CLONE_SCRIPT, {

apps/sim/executor/handlers/pi/cloud/review/backend.ts

@@ -218,7 +218,7 @@ export const runCloudReviewPi: PiBackendRun<PiCloudReviewRunParams> = async (par
     const lifetimeMs = resolvePiRunLifetimeMs(context.signal)
 
     try {
-      return await withPiSandbox({ lifetimeMs }, async (runner) => {
+      return await withPiSandbox({ lifetimeMs, cost: context.sandboxCost }, async (runner) => {
         await runner.writeFile(GIT_ASKPASS_PATH, GIT_ASKPASS_SCRIPT)
         const fetched = await raceAbort(
           runner.run(FETCH_PR_SCRIPT, {

apps/sim/executor/handlers/pi/core/backend.ts

@@ -9,6 +9,7 @@
  */
 
 import type { TSchema } from 'typebox'
+import type { SandboxCostSink } from '@/lib/execution/remote-sandbox/types'
 import type { SSHConnectionConfig } from '@/lib/internal/ssh/client'
 import type { Message } from '@/executor/handlers/agent/types'
 import type { PiEvent, PiRunTotals } from '@/executor/handlers/pi/core/events'
@@ -172,6 +173,20 @@ export type PiRunParams =
 export interface PiRunContext {
   onEvent: (event: PiEvent) => void
   signal?: AbortSignal
+  /**
+   * Where a backend reports the cost of Sim-provisioned compute it used.
+   *
+   * Both modes can fill it, from different sources. Cloud modes run the agent in
+   * a Sim-paid sandbox and report that session. Local mode drives the caller's
+   * own machine over SSH, so the agent itself costs Sim nothing — but the Sim
+   * tools it calls still run here, and a `function_execute` among them bills its
+   * own remote sandbox into the same total.
+   *
+   * The handler folds whatever lands here into the block's `toolCost`, which is
+   * what keeps a BYOK Pi run — model unbilled by definition — from reporting no
+   * cost at all for compute Sim actually paid for.
+   */
+  sandboxCost?: SandboxCostSink
 }
 
 /** Final result of a Pi run. */