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
 
@@ -506,15 +526,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 +573,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 {
@@ -620,8 +668,10 @@ export class BlockExecutor {
       return softOutput
     }
 
+    const trustedExecutionCost = readTrustedExecutionCost(error)
     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/local/sim-tools.test.ts

@@ -225,6 +225,55 @@ describe('buildSimToolSpecs', () => {
     })
   })
 
+  it('accumulates cost from canonical Function results while preserving failures', async () => {
+    mockTransformBlockTool
+      .mockResolvedValueOnce({
+        id: 'function_execute',
+        name: 'Function Execute',
+        description: 'Execute code',
+        params: {},
+        parameters: { type: 'object', properties: {} },
+      })
+      .mockResolvedValueOnce({
+        id: 'exa_search',
+        name: 'Exa Search',
+        description: 'Search the web',
+        params: {},
+        parameters: { type: 'object', properties: {} },
+      })
+    const functionToolCost = { total: 0 }
+    const [functionSpec, searchSpec] = await buildSimToolSpecs(
+      executionContext(undefined),
+      [
+        { type: 'function', operation: 'execute', usageControl: 'auto' },
+        { type: 'exa', operation: 'exa_search', usageControl: 'auto' },
+      ],
+      functionToolCost
+    )
+
+    mockExecuteTool
+      .mockResolvedValueOnce({
+        success: true,
+        output: { result: 'ok', cost: { total: 0.125 } },
+      })
+      .mockResolvedValueOnce({
+        success: true,
+        output: { result: 'search result', cost: { total: 4 } },
+      })
+      .mockResolvedValueOnce({
+        success: false,
+        output: { cost: { total: 8 } },
+        error: 'execution failed',
+      })
+
+    await functionSpec.execute({})
+    await searchSpec.execute({})
+    const failedResult = await functionSpec.execute({})
+
+    expect(functionToolCost.total).toBe(8.125)
+    expect(failedResult).toEqual({ text: 'execution failed', isError: true })
+  })
+
   it('projects named provenance in successful Sim tool output', async () => {
     mockToolAdapter({ apiKey: 'secret-value' })
     encryptionMockFns.mockDecryptSecret.mockResolvedValue({ decrypted: 'secret-value' })

apps/sim/executor/handlers/pi/local/sim-tools.ts

@@ -47,6 +47,10 @@ type PiToolResultProjection =
   | { safe: true; result: PiToolResult }
   | { safe: false; result: PiToolResult }
 
+export interface PiFunctionToolCostAccumulator {
+  total: number
+}
+
 function projectToolResult(
   result: ToolResponse,
   registry: ResolvedSecretTraceRegistry | undefined
@@ -107,7 +111,8 @@ function buildSimToolSpec(
   ctx: ExecutionContext,
   inputTools: ToolInput[],
   provider: ProviderToolConfig,
-  toolIndex: number
+  toolIndex: number,
+  functionToolCost?: PiFunctionToolCostAccumulator
 ): PiToolSpec {
   const toolId = provider.canonicalId ?? provider.id
   const preseededParams = provider.params || {}
@@ -170,6 +175,20 @@ function buildSimToolSpec(
             resolvedSecretTraceRegistry: toolCallRegistry,
           }
         )
+        const resultCost = result.output?.cost
+        const resultCostTotal =
+          resultCost && typeof resultCost === 'object'
+            ? (resultCost as Record<string, unknown>).total
+            : undefined
+        if (
+          toolId === 'function_execute' &&
+          functionToolCost &&
+          typeof resultCostTotal === 'number' &&
+          Number.isFinite(resultCostTotal) &&
+          resultCostTotal > 0
+        ) {
+          functionToolCost.total += resultCostTotal
+        }
         const projection = projectToolResult(result, toolCallRegistry?.forkForPropagatedEntries())
         if (projection.safe && registry && toolCallRegistry?.isComplete()) {
           registry.mergeToolCallRegistry(toolCallRegistry)
@@ -199,7 +218,8 @@ function buildSimToolSpec(
  */
 export async function buildSimToolSpecs(
   ctx: ExecutionContext,
-  inputTools: unknown
+  inputTools: unknown,
+  functionToolCost?: PiFunctionToolCostAccumulator
 ): Promise<PiToolSpec[]> {
   if (!Array.isArray(inputTools)) return []
 
@@ -243,6 +263,6 @@ export async function buildSimToolSpecs(
   await annotateDuplicateToolBindings(ctx, providers)
   assignProviderToolIdentities(providers)
   return configuredTools.map(({ provider, toolIndex }) =>
-    buildSimToolSpec(ctx, inputTools, provider, toolIndex)
+    buildSimToolSpec(ctx, inputTools, provider, toolIndex, functionToolCost)
   )
 }