diff --git a/apps/sim/executor/execution/block-executor.retry.test.ts b/apps/sim/executor/execution/block-executor.retry.test.ts index 54811c71f8f..e22c93104eb 100644 --- a/apps/sim/executor/execution/block-executor.retry.test.ts +++ b/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') diff --git a/apps/sim/executor/execution/block-executor.ts b/apps/sim/executor/execution/block-executor.ts index bc4241208fd..1695cde05b4 100644 --- a/apps/sim/executor/execution/block-executor.ts +++ b/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 | 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 + streamingPartialOutput?: Record, + completedHandlerCost?: TrustedExecutionCost ): Promise { 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 diff --git a/apps/sim/executor/handlers/function/function-handler.test.ts b/apps/sim/executor/handlers/function/function-handler.test.ts index c914d2f05f9..72fe2ff4262 100644 --- a/apps/sim/executor/handlers/function/function-handler.test.ts +++ b/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' } } diff --git a/apps/sim/executor/handlers/function/function-handler.ts b/apps/sim/executor/handlers/function/function-handler.ts index aefb5ab39d4..22d0b3b9938 100644 --- a/apps/sim/executor/handlers/function/function-handler.ts +++ b/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 } } diff --git a/apps/sim/executor/handlers/pi/cloud/authoring/backend.ts b/apps/sim/executor/handlers/pi/cloud/authoring/backend.ts index d874c6ab477..40d19e23385 100644 --- a/apps/sim/executor/handlers/pi/cloud/authoring/backend.ts +++ b/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({ 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(sandboxOptions, async (runner) => { try { const clone = await raceAbort( runner.run(params.mode === 'cloud' ? CREATE_PR_CLONE_SCRIPT : UPDATE_BRANCH_CLONE_SCRIPT, { diff --git a/apps/sim/executor/handlers/pi/cloud/babysit/backend.ts b/apps/sim/executor/handlers/pi/cloud/babysit/backend.ts index 6860dee6a4e..08d4af65ac5 100644 --- a/apps/sim/executor/handlers/pi/cloud/babysit/backend.ts +++ b/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: { diff --git a/apps/sim/executor/handlers/pi/cloud/plan/backend.ts b/apps/sim/executor/handlers/pi/cloud/plan/backend.ts index 27444a4493b..3997fac8eee 100644 --- a/apps/sim/executor/handlers/pi/cloud/plan/backend.ts +++ b/apps/sim/executor/handlers/pi/cloud/plan/backend.ts @@ -80,7 +80,7 @@ export const runCloudPlanPi: PiBackendRun = 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, { diff --git a/apps/sim/executor/handlers/pi/cloud/review/backend.ts b/apps/sim/executor/handlers/pi/cloud/review/backend.ts index cfd25533952..289416ce96b 100644 --- a/apps/sim/executor/handlers/pi/cloud/review/backend.ts +++ b/apps/sim/executor/handlers/pi/cloud/review/backend.ts @@ -218,7 +218,7 @@ export const runCloudReviewPi: PiBackendRun = 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, { diff --git a/apps/sim/executor/handlers/pi/core/backend.ts b/apps/sim/executor/handlers/pi/core/backend.ts index 4fb3b3ee3d2..488a0d50286 100644 --- a/apps/sim/executor/handlers/pi/core/backend.ts +++ b/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. */ diff --git a/apps/sim/executor/handlers/pi/local/sim-tools.test.ts b/apps/sim/executor/handlers/pi/local/sim-tools.test.ts index 60b99bb95ac..6427fc3cbf7 100644 --- a/apps/sim/executor/handlers/pi/local/sim-tools.test.ts +++ b/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' }) diff --git a/apps/sim/executor/handlers/pi/local/sim-tools.ts b/apps/sim/executor/handlers/pi/local/sim-tools.ts index 483eb87f0a6..876bb7e99ee 100644 --- a/apps/sim/executor/handlers/pi/local/sim-tools.ts +++ b/apps/sim/executor/handlers/pi/local/sim-tools.ts @@ -9,6 +9,7 @@ import { createLogger } from '@sim/logger' import { getErrorMessage } from '@sim/utils/errors' +import type { SandboxCostSink } from '@/lib/execution/remote-sandbox/types' import { readWorkflowInputFieldsForTool, readWorkflowMetadataForTool, @@ -107,7 +108,8 @@ function buildSimToolSpec( ctx: ExecutionContext, inputTools: ToolInput[], provider: ProviderToolConfig, - toolIndex: number + toolIndex: number, + sandboxCost?: SandboxCostSink ): PiToolSpec { const toolId = provider.canonicalId ?? provider.id const preseededParams = provider.params || {} @@ -170,6 +172,20 @@ function buildSimToolSpec( resolvedSecretTraceRegistry: toolCallRegistry, } ) + const resultCost = result.output?.cost + const resultCostTotal = + resultCost && typeof resultCost === 'object' + ? (resultCost as Record).total + : undefined + if ( + toolId === 'function_execute' && + sandboxCost && + typeof resultCostTotal === 'number' && + Number.isFinite(resultCostTotal) && + resultCostTotal > 0 + ) { + sandboxCost.total += resultCostTotal + } const projection = projectToolResult(result, toolCallRegistry?.forkForPropagatedEntries()) if (projection.safe && registry && toolCallRegistry?.isComplete()) { registry.mergeToolCallRegistry(toolCallRegistry) @@ -199,7 +215,8 @@ function buildSimToolSpec( */ export async function buildSimToolSpecs( ctx: ExecutionContext, - inputTools: unknown + inputTools: unknown, + sandboxCost?: SandboxCostSink ): Promise { if (!Array.isArray(inputTools)) return [] @@ -243,6 +260,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, sandboxCost) ) } diff --git a/apps/sim/executor/handlers/pi/pi-handler.test.ts b/apps/sim/executor/handlers/pi/pi-handler.test.ts index 063b319c6ff..d6dcc62fb41 100644 --- a/apps/sim/executor/handlers/pi/pi-handler.test.ts +++ b/apps/sim/executor/handlers/pi/pi-handler.test.ts @@ -20,6 +20,7 @@ const { mockResolveSearchKey, mockBuildSearchTool, mockAssertPermissionsAllowed, + mockBuildSimToolSpecs, MockToolNotAllowedError, } = vi.hoisted(() => ({ mockRunLocal: vi.fn(), @@ -38,6 +39,7 @@ const { mockResolveSearchKey: vi.fn(), mockBuildSearchTool: vi.fn(), mockAssertPermissionsAllowed: vi.fn(), + mockBuildSimToolSpecs: vi.fn(), MockToolNotAllowedError: class ToolNotAllowedError extends Error {}, })) @@ -64,7 +66,7 @@ vi.mock('@/executor/handlers/pi/core/context', () => ({ appendPiMemory: mockAppendMemory, })) vi.mock('@/executor/handlers/pi/local/sim-tools', () => ({ - buildSimToolSpecs: vi.fn().mockResolvedValue([]), + buildSimToolSpecs: mockBuildSimToolSpecs, })) vi.mock('@/executor/handlers/pi/local/backend', () => ({ runLocalPi: mockRunLocal })) vi.mock('@/executor/handlers/pi/cloud/authoring/backend', () => ({ @@ -109,8 +111,10 @@ vi.mock('@/blocks/utils', () => ({ }, })) +import type { PiRunContext } from '@/executor/handlers/pi/core/backend' import { PiBlockHandler, parsePiReviewMentions } from '@/executor/handlers/pi/pi-handler' import type { ExecutionContext, StreamingExecution } from '@/executor/types' +import { readTrustedExecutionCost } from '@/executor/utils/errors' import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' import type { SerializedBlock } from '@/serializer/types' @@ -153,6 +157,7 @@ describe('PiBlockHandler', () => { mockResolveSearchKey.mockReturnValue('search-key') mockBuildSearchTool.mockReturnValue({ name: 'web_search' }) mockAssertPermissionsAllowed.mockResolvedValue(undefined) + mockBuildSimToolSpecs.mockResolvedValue([]) mockResolveSkills.mockResolvedValue([]) mockLoadMemory.mockResolvedValue([]) mockAppendMemory.mockResolvedValue(undefined) @@ -275,6 +280,83 @@ describe('PiBlockHandler', () => { expect((output as Record).content).toBe('hi') }) + it('adds successful Function tool cost once to a non-streaming Local Dev result', async () => { + mockBuildSimToolSpecs.mockImplementation( + async (_ctx: unknown, _tools: unknown, functionToolCost: { total: number }) => { + functionToolCost.total += 0.125 + return [] + } + ) + + const output = (await handler.execute(ctx(), block, localInputs())) as { cost: unknown } + + expect(output.cost).toEqual({ input: 0, output: 0, toolCost: 0.125, total: 0.125 }) + }) + + it('bills the cloud sandbox a Pi session ran in, even when the model is BYOK', async () => { + // The regression this guards: the agent's own sandbox runs on Sim's provider + // account, so a BYOK run whose model cost is zero by definition would + // otherwise report no cost at all for tens of minutes of paid compute. + mockRunCloud.mockImplementation(async (_params: unknown, context: PiRunContext) => { + if (context.sandboxCost) context.sandboxCost.total += 0.0842 + return { totals: { finalText: 'done', inputTokens: 0, outputTokens: 0 } } + }) + + const output = (await handler.execute(ctx(), block, { + mode: 'cloud', + task: 'do it', + model: 'claude', + owner: 'o', + repo: 'r', + githubToken: 'ghp', + })) as { cost: unknown } + + expect(output.cost).toEqual({ input: 0, output: 0, toolCost: 0.0842, total: 0.0842 }) + }) + + it('keeps the sandbox charge on a cloud session whose agent reported an error', async () => { + // The backend returned, so the sandbox was billed and the sink holds the + // charge — but this path throws instead of reaching buildOutput, which is + // what would otherwise have published it. + mockRunCloud.mockImplementation(async (_params: unknown, context: PiRunContext) => { + if (context.sandboxCost) context.sandboxCost.total += 0.0631 + return { + totals: { finalText: '', inputTokens: 0, outputTokens: 0, errorMessage: 'agent gave up' }, + } + }) + + const error = await handler + .execute(ctx(), block, { + mode: 'cloud', + task: 'do it', + model: 'claude', + owner: 'o', + repo: 'r', + githubToken: 'ghp', + }) + .catch((thrown: unknown) => thrown) + + expect(error).toBeInstanceOf(Error) + // `toolCost` is absent by design: the trusted envelope validates exactly the + // three numeric fields it will let cross the handler boundary. `total` is + // what the ledger bills on, and it carries the sandbox charge intact. + expect(readTrustedExecutionCost(error)).toEqual({ input: 0, output: 0, total: 0.0631 }) + }) + + it('leaves a cloud run that provisioned no sandbox uncharged', async () => { + const output = (await handler.execute(ctx(), block, { + mode: 'cloud', + task: 'do it', + model: 'claude', + owner: 'o', + repo: 'r', + githubToken: 'ghp', + })) as { cost: Record } + + expect(output.cost.toolCost).toBeUndefined() + expect(output.cost.total).toBe(0) + }) + it('routes Create PR to the cloud backend and surfaces PR output', async () => { const output = (await handler.execute(ctx(), block, { mode: 'cloud', @@ -942,6 +1024,12 @@ describe('PiBlockHandler', () => { }) it('streams text when the block is selected for streaming output', async () => { + mockBuildSimToolSpecs.mockImplementation( + async (_ctx: unknown, _tools: unknown, functionToolCost: { total: number }) => { + functionToolCost.total += 0.25 + return [] + } + ) mockRunLocal.mockImplementation(async (_params, runCtx) => { runCtx.onEvent({ type: 'text', text: 'streamed' }) return { totals: { finalText: 'streamed', inputTokens: 0, outputTokens: 0, toolCalls: [] } } @@ -965,6 +1053,12 @@ describe('PiBlockHandler', () => { } expect(text).toContain('streamed') expect(result.execution.output.content).toBe('streamed') + expect(result.execution.output.cost).toEqual({ + input: 0, + output: 0, + toolCost: 0.25, + total: 0.25, + }) }) it('streams only the canonical final document for Plan mode', async () => { diff --git a/apps/sim/executor/handlers/pi/pi-handler.ts b/apps/sim/executor/handlers/pi/pi-handler.ts index 2f42719b71c..97f674ed1c8 100644 --- a/apps/sim/executor/handlers/pi/pi-handler.ts +++ b/apps/sim/executor/handlers/pi/pi-handler.ts @@ -8,6 +8,7 @@ import { createLogger } from '@sim/logger' import { projectResolvedModelInput } from '@/lib/execution/model-input-provenance' +import type { SandboxCostSink } from '@/lib/execution/remote-sandbox/types' import type { BlockOutput } from '@/blocks/types' import { parseOptionalNumberInput } from '@/blocks/utils' import { @@ -36,6 +37,7 @@ import { type PiMemoryConfig, resolvePiSkills, } from '@/executor/handlers/pi/core/context' +import type { PiRunTotals } from '@/executor/handlers/pi/core/events' import { streamTextForEvent } from '@/executor/handlers/pi/core/events' import { computePiCost, @@ -53,7 +55,9 @@ import type { NormalizedBlockOutput, StreamingExecution, } from '@/executor/types' +import { attachTrustedExecutionCost } from '@/executor/utils/errors' import { refuseResolvedSecretProjection } from '@/executor/utils/resolved-secret-projection-refusal' +import type { ModelCost } from '@/providers/cost-policy' import { isPiSupportedProvider, resolvePiModelId } from '@/providers/pi-providers' import { getProviderFromModel } from '@/providers/utils' import type { SerializedBlock } from '@/serializer/types' @@ -154,6 +158,34 @@ export function parsePiReviewMentions(value: unknown): string[] { return mentions } +/** + * What a Pi block charges: its model tokens plus the Sim-paid sandbox compute. + * + * Sandbox cost rides in `toolCost` so it survives a BYOK run — the model side is + * zero by definition there, and the ledger bills a model row on `total > 0`. + * Folding it in is what makes a BYOK Pi session bill for the provider time it + * actually consumed instead of nothing at all. + * + * Shared with the failure path deliberately: an agent that ran and then reported + * an error consumed exactly the same tokens and sandbox seconds as one that + * succeeded, so both have to arrive at the same number. + */ +function buildPiCost( + model: string, + isBYOK: boolean, + totals: PiRunTotals, + sandboxCost: number +): ModelCost { + const modelCost = computePiCost(model, totals.inputTokens, totals.outputTokens, isBYOK) + if (sandboxCost <= 0) return modelCost + + return { + ...modelCost, + toolCost: sandboxCost, + total: modelCost.total + sandboxCost, + } +} + export class PiBlockHandler implements BlockHandler { canHandle(block: SerializedBlock): boolean { return block.metadata?.id === BlockType.PI @@ -267,7 +299,8 @@ export class PiBlockHandler implements BlockHandler { } const usePrivateKey = inputs.authMethod === 'privateKey' const port = parseOptionalNumberInput(inputs.port, 'port', { integer: true, min: 1 }) ?? 22 - const tools = await buildSimToolSpecs(ctx, inputs.tools) + const sandboxCost: SandboxCostSink = { total: 0 } + const tools = await buildSimToolSpecs(ctx, inputs.tools, sandboxCost) const params: PiLocalRunParams = { ...contextualBase, mode: 'local', @@ -282,7 +315,7 @@ export class PiBlockHandler implements BlockHandler { passphrase: usePrivateKey ? asRawString(inputs.passphrase) : undefined, }, } - return this.runPi(ctx, block, runLocalPi, params, memoryConfig) + return this.runPi(ctx, block, runLocalPi, params, memoryConfig, sandboxCost) } const owner = asOptString(inputs.owner) @@ -473,10 +506,12 @@ export class PiBlockHandler implements BlockHandler { model: string, isBYOK: boolean, startTime: number, - startTimeISO: string + startTimeISO: string, + sandboxCost = 0 ): NormalizedBlockOutput { const { totals } = result const endTime = Date.now() + const cost = buildPiCost(model, isBYOK, totals, sandboxCost) return { content: totals.finalText, model, @@ -505,7 +540,7 @@ export class PiBlockHandler implements BlockHandler { output: totals.outputTokens, total: totals.inputTokens + totals.outputTokens, }, - cost: computePiCost(model, totals.inputTokens, totals.outputTokens, isBYOK), + cost, providerTiming: { startTime: startTimeISO, endTime: new Date(endTime).toISOString(), @@ -519,7 +554,15 @@ export class PiBlockHandler implements BlockHandler { block: SerializedBlock, backend: PiBackendRun

, params: P, - memoryConfig?: PiMemoryConfig + memoryConfig?: PiMemoryConfig, + /** + * One sink for every Sim-paid sandbox this block touches. Local mode fills it + * from the Function tools it runs host-side; cloud modes fill it from the + * sandbox the agent itself runs in. They are mutually exclusive in practice, + * and sharing one total means neither can be forgotten at the point the cost + * is folded into the block's output. + */ + sandboxCost: SandboxCostSink = { total: 0 } ): Promise { const startTime = Date.now() const startTimeISO = new Date(startTime).toISOString() @@ -545,9 +588,15 @@ export class PiBlockHandler implements BlockHandler { if (text) controller.enqueue(encoder.encode(text)) }, signal: ctx.abortSignal, + sandboxCost, }) if (result.totals.errorMessage) { - controller.error(new Error(result.totals.errorMessage)) + const error = new Error(result.totals.errorMessage) + attachTrustedExecutionCost( + error, + buildPiCost(params.model, params.isBYOK, result.totals, sandboxCost.total) + ) + controller.error(error) return } if (params.mode === 'cloud_plan' && result.totals.finalText) { @@ -561,7 +610,8 @@ export class PiBlockHandler implements BlockHandler { params.model, params.isBYOK, startTime, - startTimeISO + startTimeISO, + sandboxCost.total ) ) if (memoryConfig) { @@ -592,9 +642,25 @@ export class PiBlockHandler implements BlockHandler { } } - const result = await backend(params, { onEvent: () => {}, signal: ctx.abortSignal }) + const result = await backend(params, { + onEvent: () => {}, + signal: ctx.abortSignal, + sandboxCost, + }) if (result.totals.errorMessage) { - throw new Error(result.totals.errorMessage) + /* + * The backend returned, so the sandbox was billed and the sink holds the + * charge — but this throw skips `buildOutput`, which is what would have + * published it. Carrying the cost on the error is what keeps a session + * whose agent reported a failure from being run for free, the same way the + * Function handler carries its tool cost onto the error it raises. + */ + const error = new Error(result.totals.errorMessage) + attachTrustedExecutionCost( + error, + buildPiCost(params.model, params.isBYOK, result.totals, sandboxCost.total) + ) + throw error } if (memoryConfig) { await appendPiMemory( @@ -610,7 +676,8 @@ export class PiBlockHandler implements BlockHandler { params.model, params.isBYOK, startTime, - startTimeISO + startTimeISO, + sandboxCost.total ) } } diff --git a/apps/sim/executor/utils/errors.ts b/apps/sim/executor/utils/errors.ts index deb1306e6ac..e3f0a9ee9b8 100644 --- a/apps/sim/executor/utils/errors.ts +++ b/apps/sim/executor/utils/errors.ts @@ -49,6 +49,22 @@ export function attachExecutionResult(error: Error, executionResult: ExecutionRe */ const attemptedExecutionIds = new WeakMap() +/** Cost emitted by a trusted execution boundary and safe to project into a block trace. */ +export interface TrustedExecutionCost { + readonly input: number + readonly output: number + readonly total: number +} + +/** + * Trusted execution costs, keyed by the value crossing the handler boundary. + * + * Cost stays in a side table until the executor deliberately copies it into block output. This + * prevents arbitrary properties on provider errors (or user-thrown values) from becoming billed + * trace data while still allowing a handler to preserve cost when it throws. + */ +const trustedExecutionCosts = new WeakMap() + /** * Names the run a failure belongs to once dispatch has been attempted. * @@ -72,6 +88,46 @@ export function readAttemptedExecutionId(error: unknown): string | undefined { return isRecordedThrown(error) ? attemptedExecutionIds.get(error) : undefined } +/** Attaches a validated, Sim-produced execution cost to an object crossing the handler boundary. */ +export function attachTrustedExecutionCost(subject: unknown, cost: unknown): void { + if (!isRecordedThrown(subject)) return + + const normalizedCost = normalizeTrustedExecutionCost(cost) + if (!normalizedCost) return + + trustedExecutionCosts.set(subject, normalizedCost) +} + +/** Reads execution cost only when a trusted caller previously attached it. */ +export function readTrustedExecutionCost(subject: unknown): TrustedExecutionCost | undefined { + return isRecordedThrown(subject) ? trustedExecutionCosts.get(subject) : undefined +} + +function normalizeTrustedExecutionCost(cost: unknown): TrustedExecutionCost | undefined { + if (!cost || typeof cost !== 'object' || Array.isArray(cost)) return undefined + + const candidate = cost as Record + if ( + typeof candidate.input !== 'number' || + !Number.isFinite(candidate.input) || + candidate.input < 0 || + typeof candidate.output !== 'number' || + !Number.isFinite(candidate.output) || + candidate.output < 0 || + typeof candidate.total !== 'number' || + !Number.isFinite(candidate.total) || + candidate.total < 0 + ) { + return undefined + } + + return { + input: candidate.input, + output: candidate.output, + total: candidate.total, + } +} + /** * Any non-null object, not only an `Error`. * diff --git a/apps/sim/lib/billing/sandbox-pricing.test.ts b/apps/sim/lib/billing/sandbox-pricing.test.ts new file mode 100644 index 00000000000..e523120bc32 --- /dev/null +++ b/apps/sim/lib/billing/sandbox-pricing.test.ts @@ -0,0 +1,36 @@ +import { describe, expect, it } from 'vitest' +import { createSandboxPricing, priceSandboxUsage } from '@/lib/billing/sandbox-pricing' + +describe('sandbox pricing', () => { + it.each([ + ['e2b', 0.1656], + ['daytona', 0.16668], + ] as const)('prices one hour of the Function profile on %s', (provider, expected) => { + const pricing = createSandboxPricing(provider, 1) + + expect(priceSandboxUsage(pricing, 3_600_000, 3_600_000).rawCost).toBeCloseTo(expected, 8) + }) + + it('applies the multiplier once and rounds the final cost to eight decimals', () => { + const pricing = createSandboxPricing('e2b', 1.75) + + expect(priceSandboxUsage(pricing, 1234, 10_000).billedCost).toBe(0.00009934) + }) + + it('caps duration at the provider lifetime', () => { + const pricing = createSandboxPricing('daytona', 1) + + expect(priceSandboxUsage(pricing, 90_000, 60_000).durationMs).toBe(60_000) + }) + + it('allows a zero multiplier and rejects invalid multipliers', () => { + const freePricing = createSandboxPricing('e2b', 0) + + expect(priceSandboxUsage(freePricing, 1000, 1000).billedCost).toBe(0) + expect(() => createSandboxPricing('e2b', -1)).toThrow('finite nonnegative') + expect(() => createSandboxPricing('e2b', Number.NaN)).toThrow('finite nonnegative') + expect(() => createSandboxPricing('e2b', Number.POSITIVE_INFINITY)).toThrow( + 'finite nonnegative' + ) + }) +}) diff --git a/apps/sim/lib/billing/sandbox-pricing.ts b/apps/sim/lib/billing/sandbox-pricing.ts new file mode 100644 index 00000000000..828fb83204a --- /dev/null +++ b/apps/sim/lib/billing/sandbox-pricing.ts @@ -0,0 +1,103 @@ +import { getCostMultiplier } from '@/lib/core/config/env-flags' +import { + FUNCTION_DAYTONA_DISK_GB, + FUNCTION_SANDBOX_CPU_COUNT, + FUNCTION_SANDBOX_MEMORY_GB, +} from '@/lib/execution/remote-sandbox/function-resources' +import type { SandboxProviderId } from '@/lib/execution/remote-sandbox/types' + +const E2B_CPU_USD_PER_VCPU_SECOND = 0.000014 +const E2B_MEMORY_USD_PER_GIB_SECOND = 0.0000045 +const DAYTONA_CPU_USD_PER_VCPU_SECOND = 0.0504 / 3600 +const DAYTONA_MEMORY_USD_PER_GIB_SECOND = 0.0162 / 3600 +/** + * Sim prices the full provisioned disk at the marginal list rate; provider free allowances, + * credits, and discounts are intentionally not subtracted. + */ +const DAYTONA_DISK_USD_PER_GIB_SECOND = 0.000108 / 3600 + +export interface SandboxPricing { + provider: SandboxProviderId + multiplier: number + resources: { + vcpu: number + memoryGiB: number + diskGiB: number + } + rates: { + cpuUsdPerVcpuSecond: number + memoryUsdPerGiBSecond: number + diskUsdPerGiBSecond: number + } +} + +export interface PricedSandboxUsage { + durationMs: number + rawCost: number + billedCost: number +} + +const PRICING_BY_PROVIDER: Record< + SandboxProviderId, + Pick +> = { + e2b: { + resources: { + vcpu: FUNCTION_SANDBOX_CPU_COUNT, + memoryGiB: FUNCTION_SANDBOX_MEMORY_GB, + diskGiB: 0, + }, + rates: { + cpuUsdPerVcpuSecond: E2B_CPU_USD_PER_VCPU_SECOND, + memoryUsdPerGiBSecond: E2B_MEMORY_USD_PER_GIB_SECOND, + diskUsdPerGiBSecond: 0, + }, + }, + daytona: { + resources: { + vcpu: FUNCTION_SANDBOX_CPU_COUNT, + memoryGiB: FUNCTION_SANDBOX_MEMORY_GB, + diskGiB: FUNCTION_DAYTONA_DISK_GB, + }, + rates: { + cpuUsdPerVcpuSecond: DAYTONA_CPU_USD_PER_VCPU_SECOND, + memoryUsdPerGiBSecond: DAYTONA_MEMORY_USD_PER_GIB_SECOND, + diskUsdPerGiBSecond: DAYTONA_DISK_USD_PER_GIB_SECOND, + }, + }, +} + +export function createSandboxPricing( + provider: SandboxProviderId, + multiplier = getCostMultiplier() +): SandboxPricing { + if (!Number.isFinite(multiplier) || multiplier < 0) { + throw new Error('Sandbox pricing multiplier must be a finite nonnegative number') + } + const pricing = PRICING_BY_PROVIDER[provider] + return { + provider, + multiplier, + resources: { ...pricing.resources }, + rates: { ...pricing.rates }, + } +} + +export function priceSandboxUsage( + pricing: SandboxPricing, + observedDurationMs: number, + providerLifetimeMs: number +): PricedSandboxUsage { + const durationMs = Math.max(0, Math.min(observedDurationMs, providerLifetimeMs)) + const seconds = durationMs / 1000 + const rawCost = + seconds * pricing.resources.vcpu * pricing.rates.cpuUsdPerVcpuSecond + + seconds * pricing.resources.memoryGiB * pricing.rates.memoryUsdPerGiBSecond + + seconds * pricing.resources.diskGiB * pricing.rates.diskUsdPerGiBSecond + + return { + durationMs, + rawCost, + billedCost: Number.parseFloat((rawCost * pricing.multiplier).toFixed(8)), + } +} diff --git a/apps/sim/lib/execution/remote-sandbox/conformance.test.ts b/apps/sim/lib/execution/remote-sandbox/conformance.test.ts index 0ea967c1742..9bf4f80649d 100644 --- a/apps/sim/lib/execution/remote-sandbox/conformance.test.ts +++ b/apps/sim/lib/execution/remote-sandbox/conformance.test.ts @@ -123,12 +123,21 @@ import { SIM_RESULT_PREFIX, withPiSandbox, } from '@/lib/execution/remote-sandbox' -import { daytonaProvider } from '@/lib/execution/remote-sandbox/daytona' -import { E2B_MAX_SANDBOX_LIFETIME_MS, e2bProvider } from '@/lib/execution/remote-sandbox/e2b' +import { + daytonaProvider, + resolveDaytonaSandboxLifetimeMs, +} from '@/lib/execution/remote-sandbox/daytona' +import { + E2B_MAX_SANDBOX_LIFETIME_MS, + e2bProvider, + resolveE2BSandboxLifetimeMs, +} from '@/lib/execution/remote-sandbox/e2b' import { MAX_SANDBOX_OUTPUT_BYTES, + MAX_SANDBOX_OUTPUT_FILES, MAX_SANDBOX_PROCESS_OUTPUT_BYTES, MAX_SANDBOX_STREAMED_OUTPUT_TAIL_BYTES, + readTrustedSandboxOutputCost, } from '@/lib/execution/remote-sandbox/output-limits' import { PI_SANDBOX_MIN_LIFETIME_MS, @@ -139,6 +148,27 @@ import { resolveProvider } from '@/lib/execution/remote-sandbox/provider' type Provider = 'e2b' | 'daytona' const PROVIDERS: Provider[] = ['e2b', 'daytona'] +describe('provider-effective sandbox lifetimes', () => { + it('matches E2B second and Daytona minute rounding', () => { + expect(resolveE2BSandboxLifetimeMs(1001)).toBe(2000) + expect(resolveDaytonaSandboxLifetimeMs(1001)).toBe(60_000) + }) + + it.each(PROVIDERS)('reports the %s SDK create dispatch time', async (provider) => { + useProvider(provider) + const onProviderRequestStarted = vi.fn() + const createMock = provider === 'e2b' ? mockE2BCreate : mockDaytonaCreate + + await resolveProvider().create('code', { lifetimeMs: 1000, onProviderRequestStarted }) + + expect(onProviderRequestStarted).toHaveBeenCalledOnce() + expect(onProviderRequestStarted).toHaveBeenCalledWith(expect.any(Number)) + expect(onProviderRequestStarted.mock.invocationCallOrder[0]).toBeLessThan( + createMock.mock.invocationCallOrder[0] + ) + }) +}) + /** Points the shared layer at one provider via the SANDBOX_PROVIDER env var. */ function useProvider(provider: Provider) { mockEnv.SANDBOX_PROVIDER = provider @@ -337,6 +367,47 @@ describe.each(PROVIDERS)('sandbox conformance [%s]', (provider) => { expect(res.result).toEqual({ ok: true }) expect(res.stdout).toBe('hello') expect(res.error).toBeUndefined() + expect(res.cost).toBeUndefined() + }) + + it('adds provider cost to a metered successful code result', async () => { + stubCodeRun(provider, `${SIM_RESULT_PREFIX}{"ok":true}`) + let now = 1_800_000_000_000 + const nowSpy = vi.spyOn(Date, 'now').mockImplementation(() => ++now) + + try { + const res = await executeInSandbox({ + code: 'x', + language: CodeLanguage.Python, + timeoutMs: 1000, + meterUsage: true, + }) + + expect(res.cost).toEqual({ input: 0, output: 0, total: expect.any(Number) }) + expect(res.cost?.total).toBeGreaterThan(0) + } finally { + nowSpy.mockRestore() + } + }) + + it('adds provider cost to a metered successful shell result', async () => { + stubShellCommand(provider, 'ok', '', 0) + let now = 1_800_000_000_000 + const nowSpy = vi.spyOn(Date, 'now').mockImplementation(() => ++now) + + try { + const res = await executeShellInSandbox({ + code: 'echo ok', + envs: {}, + timeoutMs: 1000, + meterUsage: true, + }) + + expect(res.cost).toEqual({ input: 0, output: 0, total: expect.any(Number) }) + expect(res.cost?.total).toBeGreaterThan(0) + } finally { + nowSpy.mockRestore() + } }) it('takes the LAST marker so user output cannot shadow the real result', async () => { @@ -361,10 +432,12 @@ describe.each(PROVIDERS)('sandbox conformance [%s]', (provider) => { code: 'x', language: CodeLanguage.Python, timeoutMs: 1000, + meterUsage: true, }) expect(res.result).toBeNull() expect(res.error).toContain('corrupted in transport') + expect(res.cost).toBeUndefined() }) it('survives a large single-line payload without chunk corruption', async () => { @@ -393,13 +466,19 @@ describe.each(PROVIDERS)('sandbox conformance [%s]', (provider) => { }) } - await expect( - executeInSandbox({ code: 'x', language: CodeLanguage.Python, timeoutMs: 1000 }) - ).rejects.toMatchObject({ + const error = await executeInSandbox({ + code: 'x', + language: CodeLanguage.Python, + timeoutMs: 1000, + meterUsage: true, + }).catch((error: unknown) => error) + + expect(error).toMatchObject({ code: 'sandbox_output_limit_exceeded', outputKind: 'process', limitBytes: MAX_SANDBOX_PROCESS_OUTPUT_BYTES, }) + expect(readTrustedSandboxOutputCost(error)).toBeUndefined() expect(provider === 'e2b' ? mockE2BKill : mockDelete).toHaveBeenCalledTimes(1) }) @@ -502,11 +581,13 @@ describe.each(PROVIDERS)('sandbox conformance [%s]', (provider) => { code: 'raise ValueError("boom")', language: CodeLanguage.Python, timeoutMs: 1000, + meterUsage: true, }) expect(res.error).toBe('ValueError: boom') expect(res.stdout).toContain('ValueError: boom') expect(res.result).toBeNull() + expect(res.cost).toEqual({ input: 0, output: 0, total: expect.any(Number) }) }) it('normalizes Python code budget expiry to a typed timeout abort', async () => { @@ -633,6 +714,34 @@ describe.each(PROVIDERS)('sandbox conformance [%s]', (provider) => { } } + it('bills a completed run whose harvest produced more files than it can export', async () => { + // The sandbox executed and was paid for; the refusal is about what the code + // wrote, so it belongs with the post-completion export failures rather than + // the provider failures the policy absorbs. + stubCodeRun(provider, `${SIM_RESULT_PREFIX}null`) + stubOutputDirListing( + Array.from({ length: MAX_SANDBOX_OUTPUT_FILES + 1 }, (_, index) => ({ + path: `/tmp/sim/outputs/file-${index}.txt`, + size: 1, + })) + ) + + const error = await executeInSandbox({ + code: 'x', + language: CodeLanguage.Python, + timeoutMs: 1000, + outputSandboxDir: '/tmp/sim/outputs', + meterUsage: true, + }).catch((error: unknown) => error) + + expect(error).toMatchObject({ code: 'sandbox_output_not_exportable' }) + expect(readTrustedSandboxOutputCost(error)).toEqual({ + input: 0, + output: 0, + total: expect.any(Number), + }) + }) + it('creates the output directory before user code runs', async () => { stubCodeRun(provider, `__SIM_RESULT__=${JSON.stringify('done')}`) stubOutputDirListing([]) @@ -988,11 +1097,17 @@ describe.each(PROVIDERS)('sandbox conformance [%s]', (provider) => { */ stubShellCommand(provider, provider === 'daytona' ? 'boom detail' : '', 'boom detail', 3) - const res = await executeShellInSandbox({ code: 'false', envs: {}, timeoutMs: 1000 }) + const res = await executeShellInSandbox({ + code: 'false', + envs: {}, + timeoutMs: 1000, + meterUsage: true, + }) expect(res.result).toBeNull() expect(res.error).toContain('boom detail') expect(res.stdout).toContain('boom detail') + expect(res.cost).toEqual({ input: 0, output: 0, total: expect.any(Number) }) }) it('terminates shell execution when streamed process output exceeds the byte budget', async () => { @@ -1094,18 +1209,24 @@ describe.each(PROVIDERS)('sandbox conformance [%s]', (provider) => { stubCodeRun(provider, `${SIM_RESULT_PREFIX}null`) stubOutputFileSizes(provider, MAX_SANDBOX_OUTPUT_BYTES + 1) - await expect( - executeInSandbox({ - code: 'x', - language: CodeLanguage.Python, - timeoutMs: 1000, - outputSandboxPath: '/out/report.txt', - }) - ).rejects.toMatchObject({ + const error = await executeInSandbox({ + code: 'x', + language: CodeLanguage.Python, + timeoutMs: 1000, + outputSandboxPath: '/out/report.txt', + meterUsage: true, + }).catch((error: unknown) => error) + + expect(error).toMatchObject({ code: 'sandbox_output_limit_exceeded', attemptedBytes: MAX_SANDBOX_OUTPUT_BYTES + 1, limitBytes: MAX_SANDBOX_OUTPUT_BYTES, }) + expect(readTrustedSandboxOutputCost(error)).toEqual({ + input: 0, + output: 0, + total: expect.any(Number), + }) expect(provider === 'e2b' ? mockE2BFilesRead : mockDownloadFileStream).not.toHaveBeenCalled() }) @@ -1158,15 +1279,96 @@ describe.each(PROVIDERS)('sandbox conformance [%s]', (provider) => { mockGetFileDetails.mockResolvedValueOnce({ size: 1, isDir: false, mode: 'prw-r--r--' }) } - await expect( - executeInSandbox({ - code: 'x', - language: CodeLanguage.Python, + const error = await executeInSandbox({ + code: 'x', + language: CodeLanguage.Python, + timeoutMs: 1000, + outputSandboxPath: '/out/link.txt', + meterUsage: true, + }).catch((error: unknown) => error) + + expect(error).toMatchObject({ code: 'sandbox_output_file_invalid' }) + expect(readTrustedSandboxOutputCost(error)).toEqual({ + input: 0, + output: 0, + total: expect.any(Number), + }) + expect(provider === 'e2b' ? mockE2BFilesRead : mockDownloadFileStream).not.toHaveBeenCalled() + }) + + it.each(['oversized', 'non-regular'] as const)( + 'retains metered shell cost for a completed execution with %s output', + async (failure) => { + stubShellCommand(provider, '', '', 0) + if (failure === 'oversized') { + stubOutputFileSizes(provider, MAX_SANDBOX_OUTPUT_BYTES + 1) + } else if (provider === 'e2b') { + mockE2BFilesGetInfo.mockResolvedValueOnce({ size: 1, type: 'symlink' }) + } else { + mockGetFileDetails.mockResolvedValueOnce({ size: 1, isDir: false, mode: 'prw-r--r--' }) + } + + const error = await executeShellInSandbox({ + code: 'echo done', + envs: {}, timeoutMs: 1000, - outputSandboxPath: '/out/link.txt', + outputSandboxPath: '/out/result.txt', + meterUsage: true, + }).catch((error: unknown) => error) + + expect(error).toMatchObject({ + code: + failure === 'oversized' ? 'sandbox_output_limit_exceeded' : 'sandbox_output_file_invalid', }) - ).rejects.toMatchObject({ code: 'sandbox_output_file_invalid' }) - expect(provider === 'e2b' ? mockE2BFilesRead : mockDownloadFileStream).not.toHaveBeenCalled() + expect(readTrustedSandboxOutputCost(error)).toEqual({ + input: 0, + output: 0, + total: expect.any(Number), + }) + } + ) + + it('does not attach cost to a generic provider failure during output collection', async () => { + stubCodeRun(provider, `${SIM_RESULT_PREFIX}null`) + stubOutputFileSizes(provider, 1, 1) + const failure = new Error('provider file read failed') + if (provider === 'e2b') { + mockE2BFilesRead.mockRejectedValueOnce(failure) + } else { + mockDownloadFileStream.mockRejectedValueOnce(failure) + } + + const error = await executeInSandbox({ + code: 'x', + language: CodeLanguage.Python, + timeoutMs: 1000, + outputSandboxPath: '/out/result.txt', + meterUsage: true, + }).catch((error: unknown) => error) + + expect(error).toBe(failure) + expect(readTrustedSandboxOutputCost(error)).toBeUndefined() + }) + + it('does not attach cost to a generic provider failure during output inspection', async () => { + stubCodeRun(provider, `${SIM_RESULT_PREFIX}null`) + const failure = new Error('provider file metadata failed') + if (provider === 'e2b') { + mockE2BFilesGetInfo.mockRejectedValueOnce(failure) + } else { + mockGetFileDetails.mockRejectedValueOnce(failure) + } + + const error = await executeInSandbox({ + code: 'x', + language: CodeLanguage.Python, + timeoutMs: 1000, + outputSandboxPath: '/out/result.txt', + meterUsage: true, + }).catch((error: unknown) => error) + + expect(error).toBe(failure) + expect(readTrustedSandboxOutputCost(error)).toBeUndefined() }) it('does not return code results when cancellation arrives during output collection', async () => { @@ -1485,6 +1687,58 @@ describe('provider stream recovery', () => { expect(mockGetSessionCommand).toHaveBeenCalledWith(expect.any(String), 'cmd_1') }) + it('fails an at-most-once Daytona run closed when final status has no exit code', async () => { + mockGetSessionCommand.mockResolvedValueOnce({}) + const sandbox = await daytonaProvider.create('code') + + await expect( + sandbox.runCommand('node function.js', { timeoutMs: 1000, atMostOnce: true }) + ).rejects.toMatchObject({ + name: 'SandboxLaunchIndeterminateError', + retryable: false, + code: 'sandbox_launch_indeterminate', + }) + }) + + it('fails an at-most-once Daytona run closed when its readiness handshake never completes', async () => { + mockGetSessionCommandLogs.mockResolvedValueOnce(undefined) + mockGetSessionCommand.mockResolvedValueOnce({ exitCode: 78 }) + const sandbox = await daytonaProvider.create('code') + + await expect( + sandbox.runCommand('node function.js', { timeoutMs: 1000, atMostOnce: true }) + ).rejects.toMatchObject({ + name: 'SandboxLaunchIndeterminateError', + retryable: false, + code: 'sandbox_launch_indeterminate', + }) + expect(mockSendSessionCommandInput).not.toHaveBeenCalled() + }) + + it('fails an at-most-once Daytona run closed when final status lookup fails', async () => { + mockGetSessionCommand.mockRejectedValueOnce(new Error('control plane unavailable')) + const sandbox = await daytonaProvider.create('code') + + await expect( + sandbox.runCommand('node function.js', { timeoutMs: 1000, atMostOnce: true }) + ).rejects.toMatchObject({ + name: 'SandboxLaunchIndeterminateError', + retryable: false, + code: 'sandbox_launch_indeterminate', + }) + }) + + it('preserves a pre-dispatch Daytona failure for at-most-once runs', async () => { + const failure = new Error('session unavailable') + mockCreateSession.mockRejectedValueOnce(failure) + const sandbox = await daytonaProvider.create('code') + + await expect( + sandbox.runCommand('node function.js', { timeoutMs: 1000, atMostOnce: true }) + ).rejects.toBe(failure) + expect(mockExecuteSessionCommand).not.toHaveBeenCalled() + }) + it('keeps the original Daytona deadline while recovering a disconnected stream', async () => { mockGetSessionCommandLogs .mockRejectedValueOnce(new Error('stream disconnected')) @@ -2425,10 +2679,12 @@ describe('Pi sandbox lifetime', () => { code: 'x', language: CodeLanguage.Python, timeoutMs: 7 * 24 * 60 * 60 * 1000, + meterUsage: true, }) expect(result.error).toContain('E2B reached its 24-hour limit') expect(result.error).toContain('workflow timeout may be longer') + expect(result.cost).toBeUndefined() expect(mockRecordSandboxProviderLimit).toHaveBeenCalledWith({ provider: 'e2b', operation: 'code', @@ -2452,9 +2708,11 @@ describe('Pi sandbox lifetime', () => { const result = await executeShellInSandbox({ code: 'sleep infinity', timeoutMs: 7 * 24 * 60 * 60 * 1000, + meterUsage: true, }) expect(result.error).toContain('E2B reached its 24-hour limit') + expect(result.cost).toBeUndefined() expect(mockRecordSandboxProviderLimit).toHaveBeenCalledWith({ provider: 'e2b', operation: 'command', diff --git a/apps/sim/lib/execution/remote-sandbox/daytona.ts b/apps/sim/lib/execution/remote-sandbox/daytona.ts index 875fcb06856..11df31c48f5 100644 --- a/apps/sim/lib/execution/remote-sandbox/daytona.ts +++ b/apps/sim/lib/execution/remote-sandbox/daytona.ts @@ -43,6 +43,11 @@ const logger = createLogger('DaytonaSandboxProvider') const DAYTONA_DEFAULT_SANDBOX_TTL_MS = 24 * 60 * 60 * 1000 const DAYTONA_STREAM_READY_MARKER = '__SIM_DAYTONA_STREAM_READY__' +/** Daytona expresses sandbox TTLs as whole minutes. */ +export function resolveDaytonaSandboxLifetimeMs(lifetimeMs: number): number { + return Math.max(1, Math.ceil(lifetimeMs / 60_000)) * 60_000 +} + /** Daytona expresses every timeout in seconds; the rest of Sim works in milliseconds. */ function toSeconds(timeoutMs: number): number { return Math.max(1, Math.ceil(timeoutMs / 1000)) @@ -296,6 +301,7 @@ class DaytonaSandboxHandle implements SandboxHandle { // must never have. const finalStdout = () => (retainStdout ? stdout : tailStreamedSandboxOutput(stdout)) const finalStderr = () => (retainStderr ? stderr : tailStreamedSandboxOutput(stderr)) + let commandDispatched = false try { await this.sandbox.process.createSession(sessionId) sessionCreated = true @@ -334,6 +340,7 @@ class DaytonaSandboxHandle implements SandboxHandle { if (typeof commandId !== 'string' || commandId.length === 0) { throw new SandboxLaunchIndeterminateError('Daytona') } + commandDispatched = true // Accumulate the streamed chunks as well as forwarding them: callers read // markers out of stdout (the Pi cloud flow parses __BASE_SHA__/__CHANGED__) // and format failures from stderr, so returning empty strings here would @@ -655,7 +662,14 @@ class DaytonaSandboxHandle implements SandboxHandle { } const finished = await this.sandbox.process.getSessionCommand(sessionId, commandId) - const exitCode = finished.exitCode ?? 0 + if (options.atMostOnce && !releaseRequested) { + throw new SandboxLaunchIndeterminateError('Daytona') + } + const exitCode = finished.exitCode + if (typeof exitCode !== 'number' || !Number.isFinite(exitCode)) { + if (options.atMostOnce) throw new SandboxLaunchIndeterminateError('Daytona') + return { stdout: finalStdout(), stderr: finalStderr(), exitCode: 0 } + } return { stdout: finalStdout(), stderr: finalStderr(), exitCode } } catch (error) { if (isSandboxOutputLimitError(error)) { @@ -676,6 +690,12 @@ class DaytonaSandboxHandle implements SandboxHandle { timedOut: true, } } + if (options.atMostOnce) { + if (commandDispatched) { + throw new SandboxLaunchIndeterminateError('Daytona', { cause: error }) + } + throw error + } if (operation === 'code') throw error return { stdout: finalStdout(), stderr: finalStderr() || getErrorMessage(error), exitCode: 1 } } finally { @@ -795,6 +815,7 @@ function shellQuote(value: string): string { export const daytonaProvider: SandboxProvider = { id: 'daytona', dependencyStrategy: 'runtime', + resolveLifetimeMs: resolveDaytonaSandboxLifetimeMs, async create(kind: SandboxKind, options?: CreateSandboxOptions): Promise { const apiKey = env.DAYTONA_API_KEY if (!apiKey) { @@ -810,11 +831,11 @@ export const daytonaProvider: SandboxProvider = { snapshot, language: toDaytonaLanguage(language), ephemeral: true, - ttlMinutes: Math.max( - 1, - Math.ceil((options?.lifetimeMs ?? DAYTONA_DEFAULT_SANDBOX_TTL_MS) / 60_000) - ), + ttlMinutes: + resolveDaytonaSandboxLifetimeMs(options?.lifetimeMs ?? DAYTONA_DEFAULT_SANDBOX_TTL_MS) / + 60_000, } + options?.onProviderRequestStarted?.(Date.now()) const sandbox = await daytona.create(createOptions) return new DaytonaSandboxHandle(sandbox, language) diff --git a/apps/sim/lib/execution/remote-sandbox/e2b.ts b/apps/sim/lib/execution/remote-sandbox/e2b.ts index 6a294713f6e..c8922b2f3d6 100644 --- a/apps/sim/lib/execution/remote-sandbox/e2b.ts +++ b/apps/sim/lib/execution/remote-sandbox/e2b.ts @@ -97,6 +97,11 @@ export const E2B_SANDBOX_MATERIALIZER_REVISION = FUNCTION_SANDBOX_MATERIALIZER_R /** Maximum continuous sandbox lifetime supported by E2B. */ export const E2B_MAX_SANDBOX_LIFETIME_MS = 24 * 60 * 60 * 1000 +/** E2B sends sandbox lifetimes as whole seconds. */ +export function resolveE2BSandboxLifetimeMs(lifetimeMs: number): number { + return Math.min(Math.ceil(lifetimeMs / 1000) * 1000, E2B_MAX_SANDBOX_LIFETIME_MS) +} + const E2B_PROVIDER_LIMIT_ERROR = 'E2B reached its 24-hour limit for a single sandbox execution. The workflow timeout may be longer, but this Function call must finish within 24 hours.' const E2B_TIMEOUT_MESSAGE_PATTERN = @@ -367,7 +372,7 @@ class E2BSandboxHandle implements SandboxHandle { return { text: '', stdout: result.stdout, stderr: result.stderr, timedOut: true } } if (result.exitCode !== 0) { - if (result.stderr === E2B_PROVIDER_LIMIT_ERROR) { + if (result.providerFailure === 'provider_limit') { return { text: '', stdout: result.stdout, @@ -377,6 +382,7 @@ class E2BSandboxHandle implements SandboxHandle { value: E2B_PROVIDER_LIMIT_ERROR, traceback: E2B_PROVIDER_LIMIT_ERROR, }, + providerFailure: result.providerFailure, } } return processCodeFailure(result) @@ -536,7 +542,12 @@ class E2BSandboxHandle implements SandboxHandle { if (isNonRetryableExecutionError(error)) throw error if (reachedE2BProviderLimit(error, this.providerLimitAtMs, options.signal)) { recordSandboxProviderLimit({ provider: 'e2b', operation }) - return { stdout: '', stderr: E2B_PROVIDER_LIMIT_ERROR, exitCode: 1 } + return { + stdout: '', + stderr: E2B_PROVIDER_LIMIT_ERROR, + exitCode: 1, + providerFailure: 'provider_limit', + } } // The SDK throws on non-zero exit; callers want the streams, not a throw. const failure = error as { @@ -863,6 +874,7 @@ export const e2bProvider: SandboxProvider = { id: 'e2b', dependencyStrategy: 'prebuilt', images: e2bImages, + resolveLifetimeMs: resolveE2BSandboxLifetimeMs, async create(kind: SandboxKind, options?: CreateSandboxOptions): Promise { const apiKey = env.E2B_API_KEY if (!apiKey) { @@ -881,7 +893,9 @@ export const e2bProvider: SandboxProvider = { // default — longer than the lifetime it asked for, which is the opposite of // what it requested. const effectiveLifetimeMs = - options?.lifetimeMs !== undefined ? e2bTimeoutMs(options.lifetimeMs) : undefined + options?.lifetimeMs !== undefined + ? resolveE2BSandboxLifetimeMs(options.lifetimeMs) + : undefined const createOptions = { apiKey, ...(effectiveLifetimeMs !== undefined ? { timeoutMs: effectiveLifetimeMs } : {}), @@ -889,6 +903,7 @@ export const e2bProvider: SandboxProvider = { const { Sandbox } = await import('@e2b/code-interpreter') const lifetimeStartedAtMs = Date.now() + options?.onProviderRequestStarted?.(lifetimeStartedAtMs) const sandbox = await Sandbox.create(templateName, createOptions) return new E2BSandboxHandle( diff --git a/apps/sim/lib/execution/remote-sandbox/function-resources.ts b/apps/sim/lib/execution/remote-sandbox/function-resources.ts index afbe013c32f..9061d043da9 100644 --- a/apps/sim/lib/execution/remote-sandbox/function-resources.ts +++ b/apps/sim/lib/execution/remote-sandbox/function-resources.ts @@ -2,6 +2,7 @@ export const FUNCTION_SANDBOX_CPU_COUNT = 2 export const FUNCTION_SANDBOX_MEMORY_GB = 4 export const FUNCTION_SANDBOX_MEMORY_MB = FUNCTION_SANDBOX_MEMORY_GB * 1024 +export const FUNCTION_DAYTONA_DISK_GB = 10 /** Bump when custom dependency-layer rendering changes without a semantic spec change. */ export const FUNCTION_SANDBOX_MATERIALIZER_REVISION = 2 diff --git a/apps/sim/lib/execution/remote-sandbox/index.ts b/apps/sim/lib/execution/remote-sandbox/index.ts index 7bf9f5bff1f..ec5a82cb1e0 100644 --- a/apps/sim/lib/execution/remote-sandbox/index.ts +++ b/apps/sim/lib/execution/remote-sandbox/index.ts @@ -1,6 +1,11 @@ import { createLogger } from '@sim/logger' import { getErrorMessage } from '@sim/utils/errors' import { generateShortId } from '@sim/utils/id' +import { + createSandboxPricing, + priceSandboxUsage, + type SandboxPricing, +} from '@/lib/billing/sandbox-pricing' import { createTimeoutAbortController, getRemainingExecutionMs, @@ -10,8 +15,10 @@ import { recordSandboxTeardownFailure } from '@/lib/core/execution-limits/metric import { buildJavaScriptRuntimeBindingsSource } from '@/lib/execution/code-placeholders/javascript-runtime' import { SANDBOX_SYSTEM_PATH } from '@/lib/execution/remote-sandbox/cli-tools.server' import { + attachTrustedSandboxOutputCost, isSandboxOutputFileError, isSandboxOutputLimitError, + isSandboxOutputNotExportableError, MAX_SANDBOX_OUTPUT_BYTES, MAX_SANDBOX_OUTPUT_FILES, MAX_SANDBOX_PROCESS_OUTPUT_BYTES, @@ -39,17 +46,22 @@ import type { SandboxCodeResult, SandboxCollectedFile, SandboxCommandResult, + SandboxCostSink, SandboxDirectoryEntry, + SandboxExecutionCost, SandboxExecutionRequest, SandboxExecutionResult, SandboxFile, SandboxHandle, SandboxKind, SandboxPrivateInput, + SandboxProvider, + SandboxProviderId, SandboxShellExecutionRequest, } from '@/lib/execution/remote-sandbox/types' export type { + SandboxCostSink, SandboxExecutionRequest, SandboxExecutionResult, SandboxFile, @@ -59,14 +71,41 @@ export type { const logger = createLogger('RemoteSandbox') +interface CreatedSandbox { + sandbox: SandboxHandle + providerId: SandboxProviderId + startedAtMs: number + effectiveLifetimeMs?: number + pricing?: SandboxPricing +} + async function createSandbox( kind: SandboxKind, - options?: CreateSandboxOptions -): Promise { - const provider = resolveProvider() - const sandbox = await provider.create(kind, options) + options?: CreateSandboxOptions, + meterUsage = false, + provider: SandboxProvider = resolveProvider() +): Promise { + const effectiveLifetimeMs = + options?.lifetimeMs !== undefined ? provider.resolveLifetimeMs(options.lifetimeMs) : undefined + if (meterUsage && effectiveLifetimeMs === undefined) { + throw new Error('Metered sandbox execution requires a provider lifetime') + } + const pricing = meterUsage ? createSandboxPricing(provider.id) : undefined + let startedAtMs = Date.now() + const providerOptions = { + ...options, + ...(effectiveLifetimeMs !== undefined ? { lifetimeMs: effectiveLifetimeMs } : {}), + ...(meterUsage ? { onProviderRequestStarted: (value: number) => (startedAtMs = value) } : {}), + } + const sandbox = await provider.create(kind, providerOptions) logger.info('Created sandbox', { provider: provider.id, kind, sandboxId: sandbox.sandboxId }) - return sandbox + return { + sandbox, + providerId: provider.id, + startedAtMs, + ...(effectiveLifetimeMs !== undefined ? { effectiveLifetimeMs } : {}), + ...(pricing ? { pricing } : {}), + } } /** @@ -83,10 +122,11 @@ async function createSelectedSandbox( kind: SandboxKind, options: CreateSandboxOptions, selected: ResolvedSandbox | null, - signal: AbortSignal -): Promise { + signal: AbortSignal, + meterUsage = false +): Promise { try { - return await createSandbox(kind, options) + return await createSandbox(kind, options, meterUsage) } catch (error) { throwIfAborted(signal) if (!selected) throw error @@ -166,10 +206,13 @@ function throwIfSandboxTimedOut(result: { timedOut?: boolean }): void { if (result.timedOut) throw new DOMException('timeout', 'AbortError') } -function bindSandboxAbort(sandbox: SandboxHandle, signal?: AbortSignal) { +function bindSandboxAbort( + sandbox: SandboxHandle, + provider: SandboxProviderId, + signal?: AbortSignal +) { let killed = false let killPromise: Promise | null = null - const provider = resolveProvider().id const kill = (reason: 'cleanup' | 'cancellation' | 'timeout'): Promise => { if (killed) return Promise.resolve() if (!killPromise) { @@ -211,6 +254,19 @@ function bindSandboxAbort(sandbox: SandboxHandle, signal?: AbortSignal) { } } +function calculateSandboxCost( + created: CreatedSandbox, + cleanupStartedAtMs: number +): SandboxExecutionCost | undefined { + if (!created.pricing || created.effectiveLifetimeMs === undefined) return undefined + const usage = priceSandboxUsage( + created.pricing, + cleanupStartedAtMs - created.startedAtMs, + created.effectiveLifetimeMs + ) + return { input: 0, output: 0, total: usage.billedCost } +} + /** * Fetches one URL mount inside the sandbox, bounded by MAX_BYTES. * @@ -458,14 +514,14 @@ async function readSandboxOutputFile( logger.warn('Failed to read requested sandbox output file', { sandboxId: sandbox.sandboxId, }) - return undefined + throw error } } async function inspectSandboxOutputFileSize( sandbox: SandboxHandle, outputSandboxPath: string -): Promise { +): Promise { try { const size = await sandbox.getFileSize(outputSandboxPath) if (!Number.isSafeInteger(size) || size < 0) { @@ -477,7 +533,7 @@ async function inspectSandboxOutputFileSize( logger.warn('Failed to inspect requested sandbox output file', { sandboxId: sandbox.sandboxId, }) - return undefined + throw error } } @@ -576,7 +632,6 @@ async function collectExportedFiles( for (const outputSandboxPath of requestedOutputSandboxPaths(req)) { const size = await inspectSandboxOutputFileSize(sandbox, outputSandboxPath) remainingSandboxBudgetMs(options.signal) - if (size === undefined) continue totalOutputBytes += size if (totalOutputBytes > MAX_SANDBOX_OUTPUT_BYTES) { throw new SandboxOutputLimitError(totalOutputBytes) @@ -753,7 +808,7 @@ async function executeInSandboxWithinBudget( }) throwIfAborted(signal) - const sandbox = await createSelectedSandbox( + const created = await createSelectedSandbox( kind, { language, @@ -761,10 +816,14 @@ async function executeInSandboxWithinBudget( lifetimeMs: remainingSandboxBudgetMs(signal), }, selected, - signal + signal, + req.meterUsage ) + const sandbox = created.sandbox const sandboxId = sandbox.sandboxId - const abortBinding = bindSandboxAbort(sandbox, signal) + const abortBinding = bindSandboxAbort(sandbox, created.providerId, signal) + let billableResult: SandboxExecutionResult | undefined + let billableOutputError: unknown try { throwIfAborted(signal) @@ -809,12 +868,14 @@ async function executeInSandboxWithinBudget( sandboxId, hasTraceback: Boolean(execution.error.traceback), }) - return { + const executionResult = { result: null, stdout: execution.error.traceback || errorMessage, error: errorMessage, sandboxId, } + if (execution.providerFailure !== 'provider_limit') billableResult = executionResult + return executionResult } // Distinct sources (final-expression text, stdout, stderr) join with '\n' so @@ -843,22 +904,47 @@ async function executeInSandboxWithinBudget( } } - const { exportedFiles, exportedFileContent, collectedFiles } = await collectExportedFiles( - sandbox, - req, - { signal } - ) - throwIfAborted(signal) - - return { + billableResult = { result: extraction.result, stdout: cleanedStdout, sandboxId, - exportedFileContent, - exportedFiles, - collectedFiles, } + try { + const { exportedFiles, exportedFileContent, collectedFiles } = await collectExportedFiles( + sandbox, + req, + { signal } + ) + throwIfAborted(signal) + billableResult.exportedFileContent = exportedFileContent + billableResult.exportedFiles = exportedFiles + billableResult.collectedFiles = collectedFiles + } catch (error) { + /* + * A harvest that cannot return what the run produced — too many files, too + * deep, or an output directory the code deleted — is the caller's to fix + * and arrives only after the sandbox has already executed. It belongs with + * the other post-completion export failures the policy bills, not with the + * provider failures it absorbs; leaving it out let a completed run whose + * code wrote one file too many go free. + */ + if ( + isSandboxOutputLimitError(error) || + isSandboxOutputFileError(error) || + isSandboxOutputNotExportableError(error) + ) { + billableOutputError = error + } + throw error + } + return billableResult } finally { + const cleanupStartedAtMs = Date.now() + const cost = calculateSandboxCost(created, cleanupStartedAtMs) + if (cost && billableResult) billableResult.cost = cost + if (cost && billableOutputError) { + attachTrustedSandboxOutputCost(billableOutputError, cost) + } abortBinding.detach() await abortBinding.cleanup() } @@ -887,14 +973,18 @@ async function executeShellInSandboxWithinBudget( }) throwIfAborted(signal) - const sandbox = await createSelectedSandbox( + const created = await createSelectedSandbox( kind, { imageRef: selected?.imageRef, lifetimeMs: remainingSandboxBudgetMs(signal) }, selected, - signal + signal, + req.meterUsage ) + const sandbox = created.sandbox const sandboxId = sandbox.sandboxId - const abortBinding = bindSandboxAbort(sandbox, signal) + const abortBinding = bindSandboxAbort(sandbox, created.providerId, signal) + let billableResult: SandboxExecutionResult | undefined + let billableOutputError: unknown try { throwIfAborted(signal) @@ -946,7 +1036,9 @@ async function executeShellInSandboxWithinBudget( sandboxId, exitCode: result.exitCode, }) - return { result: null, stdout, error: errorMessage, sandboxId } + const executionResult = { result: null, stdout, error: errorMessage, sandboxId } + if (result.providerFailure !== 'provider_limit') billableResult = executionResult + return executionResult } // Shell scripts have no wrapper: any __SIM_RESULT__ line is user-authored @@ -955,22 +1047,47 @@ async function executeShellInSandboxWithinBudget( const extraction = extractSimResult(stdout) const parsed = extraction.parseFailed ? extraction.rawPayload : extraction.result - const { exportedFiles, exportedFileContent, collectedFiles } = await collectExportedFiles( - sandbox, - req, - { signal } - ) - throwIfAborted(signal) - - return { + billableResult = { result: parsed, stdout: extraction.cleanedStdout, sandboxId, - exportedFileContent, - exportedFiles, - collectedFiles, } + try { + const { exportedFiles, exportedFileContent, collectedFiles } = await collectExportedFiles( + sandbox, + req, + { signal } + ) + throwIfAborted(signal) + billableResult.exportedFileContent = exportedFileContent + billableResult.exportedFiles = exportedFiles + billableResult.collectedFiles = collectedFiles + } catch (error) { + /* + * A harvest that cannot return what the run produced — too many files, too + * deep, or an output directory the code deleted — is the caller's to fix + * and arrives only after the sandbox has already executed. It belongs with + * the other post-completion export failures the policy bills, not with the + * provider failures it absorbs; leaving it out let a completed run whose + * code wrote one file too many go free. + */ + if ( + isSandboxOutputLimitError(error) || + isSandboxOutputFileError(error) || + isSandboxOutputNotExportableError(error) + ) { + billableOutputError = error + } + throw error + } + return billableResult } finally { + const cleanupStartedAtMs = Date.now() + const cost = calculateSandboxCost(created, cleanupStartedAtMs) + if (cost && billableResult) billableResult.cost = cost + if (cost && billableOutputError) { + attachTrustedSandboxOutputCost(billableOutputError, cost) + } abortBinding.detach() await abortBinding.cleanup() } @@ -1027,12 +1144,13 @@ export interface PiSandboxRunner { * caller's sandbox body, which would have buried the change in whitespace. */ export async function withPiSandbox( - options: { lifetimeMs?: number }, + options: { lifetimeMs?: number; cost?: SandboxCostSink }, fn: (runner: PiSandboxRunner) => Promise ): Promise { const lifetimeMs = options.lifetimeMs !== undefined ? options.lifetimeMs : resolvePiSandboxLifetimeMs() - const sandbox = await createSandbox('pi', { lifetimeMs }) + const created = await createSandbox('pi', { lifetimeMs }, Boolean(options.cost)) + const { sandbox } = created logger.info('Started Pi sandbox', { sandboxId: sandbox.sandboxId, lifetimeMs }) const runner: PiSandboxRunner = { @@ -1049,9 +1167,31 @@ export async function withPiSandbox( writeFile: (path, content) => sandbox.writeFile(path, content), } + let sessionCompleted = false try { - return await fn(runner) + const result = await fn(runner) + sessionCompleted = true + return result } finally { + /* + * Charged only for a session that ran to completion, which is the same rule + * the Function path applies to its own outcomes: a run whose sandbox never + * delivered is not billed, because a charge nobody can tie to delivered work + * is not one worth defending. A session that ends by throwing — a provider + * crash, a lifetime limit, a cancellation — is absorbed, and a create that + * throws never reaches here at all. + * + * A command exiting non-zero is not a failure by this rule. `fn` returns + * normally there, the agent produced its answer, and the Function path bills + * its own non-zero exits for the same reason. + * + * Measured up to teardown rather than to the last command, so the window + * covers the whole time the provider held the sandbox. + */ + if (sessionCompleted) { + const cost = calculateSandboxCost(created, Date.now()) + if (cost && options.cost) options.cost.total += cost.total + } try { await sandbox.kill() } catch { diff --git a/apps/sim/lib/execution/remote-sandbox/output-limits.ts b/apps/sim/lib/execution/remote-sandbox/output-limits.ts index 0110bf57215..e260660aa19 100644 --- a/apps/sim/lib/execution/remote-sandbox/output-limits.ts +++ b/apps/sim/lib/execution/remote-sandbox/output-limits.ts @@ -1,3 +1,5 @@ +import type { SandboxExecutionCost } from '@/lib/execution/remote-sandbox/types' + export const MAX_SANDBOX_OUTPUT_BYTES = 50 * 1024 * 1024 /** @@ -103,6 +105,21 @@ export class SandboxOutputDepthError extends Error { } } +const trustedSandboxOutputCosts = new WeakMap() + +/** Associates Sim-calculated cost with a trusted post-execution output error. */ +export function attachTrustedSandboxOutputCost(error: unknown, cost: SandboxExecutionCost): void { + if (typeof error !== 'object' || error === null) return + trustedSandboxOutputCosts.set(error, cost) +} + +/** Reads cost only when the sandbox lifecycle attached it after a completed execution. */ +export function readTrustedSandboxOutputCost(error: unknown): SandboxExecutionCost | undefined { + return typeof error === 'object' && error !== null + ? trustedSandboxOutputCosts.get(error) + : undefined +} + export class SandboxOutputFileError extends Error { readonly code = SANDBOX_OUTPUT_FILE_INVALID_CODE diff --git a/apps/sim/lib/execution/remote-sandbox/pi-sandbox-billing.smoke.test.ts b/apps/sim/lib/execution/remote-sandbox/pi-sandbox-billing.smoke.test.ts new file mode 100644 index 00000000000..648eb64e491 --- /dev/null +++ b/apps/sim/lib/execution/remote-sandbox/pi-sandbox-billing.smoke.test.ts @@ -0,0 +1,96 @@ +/** + * @vitest-environment node + * + * Checks that a Pi session's sandbox is actually metered against a real provider. + * + * The handler-level test mocks the backend and writes into the sink by hand, so + * it proves the wiring from a backend to the block's cost and nothing else. It + * would still pass if `withPiSandbox` never metered at all — which is exactly + * the bug this path had. Only a real Pi sandbox shows that creation is metered, + * that teardown reports, and that the amount tracks the session's real lifetime. + * + * Enable with `SANDBOX_BILLING_SMOKE=1`, against whichever provider + * `SANDBOX_PROVIDER` selects. Needs that provider's Pi image configured + * (`E2B_PI_TEMPLATE_ID` / `DAYTONA_PI_SNAPSHOT_ID`). + */ +import { describe, expect, it } from 'vitest' +import { createSandboxPricing } from '@/lib/billing/sandbox-pricing' +import { withPiSandbox } from '@/lib/execution/remote-sandbox' +import { resolveProvider } from '@/lib/execution/remote-sandbox/provider' +import type { SandboxCostSink } from '@/lib/execution/remote-sandbox/types' + +const smokeEnabled = process.env.SANDBOX_BILLING_SMOKE === '1' +const CASE_TIMEOUT_MS = 5 * 60_000 + +/** Long enough that provisioning jitter cannot dominate the measured session. */ +const SLEEP_SECONDS = 5 +/** Well under any provider ceiling, so the lifetime cap never clamps the charge. */ +const LIFETIME_MS = 10 * 60_000 + +describe.skipIf(!smokeEnabled)('pi sandbox billing smoke', () => { + it( + 'bills the session a Pi sandbox was held for', + async () => { + const pricing = createSandboxPricing(resolveProvider().id) + const usdPerBilledSecond = + (pricing.resources.vcpu * pricing.rates.cpuUsdPerVcpuSecond + + pricing.resources.memoryGiB * pricing.rates.memoryUsdPerGiBSecond + + pricing.resources.diskGiB * pricing.rates.diskUsdPerGiBSecond) * + pricing.multiplier + + const sandboxCost: SandboxCostSink = { total: 0 } + const wallClockStartedAtMs = Date.now() + const exitCode = await withPiSandbox( + { lifetimeMs: LIFETIME_MS, cost: sandboxCost }, + async (runner) => { + const result = await runner.run(`sleep ${SLEEP_SECONDS}; echo held`, { + envs: {}, + timeoutMs: CASE_TIMEOUT_MS, + }) + return result.exitCode + } + ) + const wallClockMs = Date.now() - wallClockStartedAtMs + + expect(exitCode).toBe(0) + expect(sandboxCost.total).toBeGreaterThanOrEqual(SLEEP_SECONDS * usdPerBilledSecond) + expect(sandboxCost.total).toBeLessThanOrEqual((wallClockMs / 1000) * usdPerBilledSecond) + }, + CASE_TIMEOUT_MS + ) + + it( + 'bills nothing for a session that ended by throwing', + async () => { + // Mirrors the Function path: a sandbox that never delivered is absorbed + // rather than charged. Covers a provider crash, a lifetime limit, and a + // cancellation alike, since all three reach here the same way. + const sandboxCost: SandboxCostSink = { total: 0 } + + await expect( + withPiSandbox({ lifetimeMs: LIFETIME_MS, cost: sandboxCost }, async (runner) => { + await runner.run('echo started', { envs: {}, timeoutMs: CASE_TIMEOUT_MS }) + throw new Error('session failed after the sandbox was provisioned') + }) + ).rejects.toThrow('session failed after the sandbox was provisioned') + + expect(sandboxCost.total).toBe(0) + }, + CASE_TIMEOUT_MS + ) + + it( + 'bills nothing when no sink is supplied', + async () => { + // The mothership and any other internal caller must stay free, and the + // absence of a sink is the whole mechanism keeping them that way. + const held = await withPiSandbox({ lifetimeMs: LIFETIME_MS }, async (runner) => { + const result = await runner.run('echo held', { envs: {}, timeoutMs: CASE_TIMEOUT_MS }) + return result.exitCode + }) + + expect(held).toBe(0) + }, + CASE_TIMEOUT_MS + ) +}) diff --git a/apps/sim/lib/execution/remote-sandbox/sandbox-billing.smoke.test.ts b/apps/sim/lib/execution/remote-sandbox/sandbox-billing.smoke.test.ts new file mode 100644 index 00000000000..6b5ebfa63ee --- /dev/null +++ b/apps/sim/lib/execution/remote-sandbox/sandbox-billing.smoke.test.ts @@ -0,0 +1,78 @@ +/** + * @vitest-environment node + * + * Checks the metered amount against a real provider run. + * + * `sandbox-pricing.test.ts` pins the arithmetic and the conformance suite proves + * a cost is produced, attached, and routed — but that suite stubs the provider + * and mocks `Date.now()`, so its clock advances one millisecond per call. Under + * those conditions `total > 0` is the strongest claim available, and it would + * hold just as well if the metered window measured the wrong instants. Only a + * real run can show that the window tracks the sandbox's actual lifetime. + * + * Enable with `SANDBOX_BILLING_SMOKE=1`. Runs against whichever provider + * `SANDBOX_PROVIDER` selects, so point it at each in turn to cover both. + */ +import { describe, expect, it } from 'vitest' +import { createSandboxPricing } from '@/lib/billing/sandbox-pricing' +import { CodeLanguage } from '@/lib/execution/languages' +import { executeInSandbox } from '@/lib/execution/remote-sandbox' +import { resolveProvider } from '@/lib/execution/remote-sandbox/provider' + +const smokeEnabled = process.env.SANDBOX_BILLING_SMOKE === '1' +const CASE_TIMEOUT_MS = 5 * 60_000 +const RUN_TIMEOUT_MS = 4 * 60_000 + +/** Long enough that provisioning jitter cannot dominate the measured runtime. */ +const SLEEP_SECONDS = 5 + +describe.skipIf(!smokeEnabled)('sandbox billing smoke', () => { + it( + 'bills the sandbox lifetime at the provider rate', + async () => { + const pricing = createSandboxPricing(resolveProvider().id) + const usdPerSecond = + pricing.resources.vcpu * pricing.rates.cpuUsdPerVcpuSecond + + pricing.resources.memoryGiB * pricing.rates.memoryUsdPerGiBSecond + + pricing.resources.diskGiB * pricing.rates.diskUsdPerGiBSecond + const usdPerBilledSecond = usdPerSecond * pricing.multiplier + + const wallClockStartedAtMs = Date.now() + const result = await executeInSandbox({ + code: `import time\ntime.sleep(${SLEEP_SECONDS})\nprint("slept")`, + language: CodeLanguage.Python, + timeoutMs: RUN_TIMEOUT_MS, + meterUsage: true, + }) + const wallClockMs = Date.now() - wallClockStartedAtMs + + expect(result.cost).toEqual({ input: 0, output: 0, total: expect.any(Number) }) + const billed = result.cost?.total ?? 0 + + /** + * The window opens immediately before the provider create call and closes + * before teardown, so it has to cover the sleep and cannot exceed the whole + * call measured from out here. A rate error, a wrong resource constant, or a + * window anchored to the wrong instant all land outside these bounds — which + * an `expect.any(Number)` assertion cannot see. + */ + expect(billed).toBeGreaterThanOrEqual(SLEEP_SECONDS * usdPerBilledSecond) + expect(billed).toBeLessThanOrEqual((wallClockMs / 1000) * usdPerBilledSecond) + }, + CASE_TIMEOUT_MS + ) + + it( + 'bills nothing when the caller did not ask for metering', + async () => { + const result = await executeInSandbox({ + code: 'print("unmetered")', + language: CodeLanguage.Python, + timeoutMs: RUN_TIMEOUT_MS, + }) + + expect(result.cost).toBeUndefined() + }, + CASE_TIMEOUT_MS + ) +}) diff --git a/apps/sim/lib/execution/remote-sandbox/types.ts b/apps/sim/lib/execution/remote-sandbox/types.ts index 2010c3cee2e..10732ad35ae 100644 --- a/apps/sim/lib/execution/remote-sandbox/types.ts +++ b/apps/sim/lib/execution/remote-sandbox/types.ts @@ -72,6 +72,8 @@ export interface SandboxExecutionRequest { sandboxId?: string /** Cancels the provider sandbox when the caller's execution budget expires. */ signal?: AbortSignal + /** Adds the remote provider cost to a completed, billable Function outcome. */ + meterUsage?: boolean } export interface SandboxShellExecutionRequest { @@ -97,6 +99,26 @@ export interface SandboxShellExecutionRequest { sandboxId?: string /** Cancels the provider sandbox when the caller's execution budget expires. */ signal?: AbortSignal + /** Adds the remote provider cost to a completed, billable Function outcome. */ + meterUsage?: boolean +} + +export interface SandboxExecutionCost { + input: number + output: number + total: number +} + +/** + * Running total a caller accumulates sandbox charges into. + * + * A long-lived sandbox reports its cost when it is torn down, which is after the + * value its caller cares about has already been returned. Handing the layer a + * sink lets the charge land without reshaping every return type between here and + * the block that owns the bill. + */ +export interface SandboxCostSink { + total: number } export interface SandboxExecutionResult { @@ -116,6 +138,7 @@ export interface SandboxExecutionResult { * sequence, and the byte budget is enforced on the decoded length. */ collectedFiles?: SandboxCollectedFile[] + cost?: SandboxExecutionCost } /** One harvested output file, carried as base64 with its decoded length. */ @@ -133,6 +156,8 @@ export interface SandboxCommandResult { exitCode: number /** The provider stopped the command because its supplied execution budget elapsed. */ timedOut?: boolean + /** The provider ended execution for an infrastructure reason, not a user-process outcome. */ + providerFailure?: 'provider_limit' } /** @@ -156,6 +181,8 @@ export interface SandboxCodeResult { error?: SandboxCodeError /** The provider stopped the code runner because its supplied execution budget elapsed. */ timedOut?: boolean + /** The provider ended execution for an infrastructure reason, not a user-program outcome. */ + providerFailure?: 'provider_limit' } export interface RunCommandOptions { @@ -279,6 +306,8 @@ export interface CreateSandboxOptions { * and creates the sandbox as ephemeral. */ lifetimeMs?: number + /** Reports the instant immediately before the provider SDK create request is dispatched. */ + onProviderRequestStarted?: (startedAtMs: number) => void } /** @@ -366,5 +395,7 @@ export interface SandboxProvider { readonly dependencyStrategy: SandboxDependencyStrategy /** Present exactly when {@link dependencyStrategy} is `prebuilt`. */ readonly images?: SandboxImageBuilder + /** Resolves the provider's rounded lifetime for both creation and metering. */ + resolveLifetimeMs(lifetimeMs: number): number create(kind: SandboxKind, options?: CreateSandboxOptions): Promise } diff --git a/apps/sim/lib/function-execution/execute-request.test.ts b/apps/sim/lib/function-execution/execute-request.test.ts index cdba883bb08..dab70ebe860 100644 --- a/apps/sim/lib/function-execution/execute-request.test.ts +++ b/apps/sim/lib/function-execution/execute-request.test.ts @@ -23,6 +23,7 @@ import { PRIVATE_SECRET_PROVENANCE_HEADER, } from '@/lib/execution/private-tool-metadata' import { + attachTrustedSandboxOutputCost, MAX_SANDBOX_OUTPUT_BYTES, SandboxOutputFileError, SandboxOutputLimitError, @@ -84,7 +85,11 @@ vi.mock('@/lib/copilot/request/tools/files', () => ({ md: 'text/markdown', html: 'text/html', }, - normalizeOutputWorkspaceFileName: vi.fn((p: string) => p.replace(/^files\//, '')), + normalizeOutputWorkspaceFileName: vi.fn((p: string) => { + const normalized = p.trim().replace(/^\/+|\/+$/g, '') + if (!normalized) throw new Error('Output path must include a file name') + return normalized.replace(/^files\//, '') + }), resolveOutputFormat: vi.fn(() => 'json'), getOutputFileDeclarations: vi.fn((params: Record) => { if (Array.isArray(params.outputs?.files)) { @@ -328,6 +333,7 @@ describe('Function execution request', () => { result: 'done', stdout: 'ok', sandboxId: 'sandbox-123', + cost: { input: 0, output: 0, total: 0.00012345 }, exportedFiles: { '/tmp/out.txt': 'owned by attacker' }, }) mockWriteWorkspaceFileByPath.mockRejectedValueOnce( @@ -338,6 +344,8 @@ describe('Function execution request', () => { code: 'print("done")', language: 'python', workspaceId: 'workspace-victim', + workflowId: 'workflow-1', + executionId: 'execution-1', outputs: { files: [{ path: 'files/README.md', mode: 'overwrite', sandboxPath: '/tmp/out.txt' }], }, @@ -348,6 +356,7 @@ describe('Function execution request', () => { expect(response.status).toBe(403) expect(data).toHaveProperty('error', 'Insufficient workspace permissions') + expect(data.output.cost).toEqual({ input: 0, output: 0, total: 0.00012345 }) expect(mockWriteWorkspaceFileByPath).toHaveBeenCalledTimes(1) }) @@ -367,6 +376,113 @@ describe('Function execution request', () => { expect(mockExecuteShellInSandbox).not.toHaveBeenCalled() }) + it.each([ + { language: 'python', code: 'return 42', kind: 'code' }, + { language: 'shell', code: 'echo ready', kind: 'shell' }, + ])( + 'meters a standard workflow Function $kind sandbox and preserves its cost', + async ({ language, code, kind }) => { + envFlagsMock.isRemoteSandboxEnabled = true + const cost = { input: 0, output: 0, total: 0.00012345 } + const executeSandbox = kind === 'shell' ? mockExecuteShellInSandbox : mockExecuteInSandbox + executeSandbox.mockResolvedValueOnce({ + result: 42, + stdout: 'ready', + sandboxId: `sandbox-${kind}`, + cost, + }) + + const response = await POST( + createMockRequest('POST', { + code, + language, + workflowId: 'workflow-1', + workspaceId: 'workspace-1', + executionId: 'execution-1', + }) + ) + const data = await response.json() + + expect(response.status).toBe(200) + expect(executeSandbox).toHaveBeenCalledWith(expect.objectContaining({ meterUsage: true })) + expect(data.output.cost).toEqual(cost) + } + ) + + it.each([ + { + language: 'javascript', + code: 'import "node:path"\nthrow new Error("boom")', + kind: 'code', + }, + { language: 'python', code: 'raise ValueError("boom")', kind: 'code' }, + { language: 'shell', code: 'exit 1', kind: 'shell' }, + ])( + 'preserves sandbox cost in a failed remote $language Function response', + async ({ language, code, kind }) => { + envFlagsMock.isRemoteSandboxEnabled = true + const cost = { input: 0, output: 0, total: 0.00012345 } + const executeSandbox = kind === 'shell' ? mockExecuteShellInSandbox : mockExecuteInSandbox + executeSandbox.mockResolvedValueOnce({ + result: null, + stdout: 'boom', + error: 'boom', + sandboxId: `sandbox-${language}`, + cost, + }) + + const response = await POST( + createMockRequest('POST', { + code, + language, + workflowId: 'workflow-1', + workspaceId: 'workspace-1', + executionId: 'execution-1', + }) + ) + const data = await response.json() + + expect(response.status).toBe(422) + expect(executeSandbox).toHaveBeenCalledWith(expect.objectContaining({ meterUsage: true })) + expect(data.output.cost).toEqual(cost) + } + ) + + it('does not meter a non-workflow remote Function call', async () => { + envFlagsMock.isRemoteSandboxEnabled = true + + const response = await POST( + createMockRequest('POST', { + code: 'import path from "node:path"\nreturn path.sep', + language: 'javascript', + }) + ) + + expect(response.status).toBe(200) + expect(mockExecuteInSandbox).toHaveBeenCalledWith( + expect.objectContaining({ meterUsage: false }) + ) + }) + + it('keeps a custom Function tool local even when workflow context is present', async () => { + envFlagsMock.isRemoteSandboxEnabled = true + + const response = await POST( + createMockRequest('POST', { + code: 'return 42', + language: 'python', + workflowId: 'workflow-1', + workspaceId: 'workspace-1', + executionId: 'execution-1', + isCustomTool: true, + }) + ) + + expect(response.status).toBe(200) + expect(mockExecuteInIsolatedVM).toHaveBeenCalledOnce() + expect(mockExecuteInSandbox).not.toHaveBeenCalled() + }) + it('does not accept a Mothership sandbox profile from the request body', async () => { const req = createMockRequest('POST', { code: 'return "test"', @@ -422,6 +538,7 @@ describe('Function execution request', () => { expect.objectContaining({ language, sandboxKind: 'mothership', + meterUsage: false, }) ) expect(mockExecuteInIsolatedVM).not.toHaveBeenCalled() @@ -443,7 +560,7 @@ describe('Function execution request', () => { expect(response.status).toBe(200) expect(mockExecuteShellInSandbox).toHaveBeenCalledWith( - expect.objectContaining({ sandboxKind: 'mothership' }) + expect.objectContaining({ sandboxKind: 'mothership', meterUsage: false }) ) }) @@ -744,6 +861,7 @@ describe('Function execution request', () => { result: 'done', stdout: 'ok', sandboxId: 'sandbox-123', + cost: { input: 0, output: 0, total: 0.00023456 }, exportedFiles: { '/home/user/chart.png': 'iVBORw0KGgo=', '/home/user/summary.json': '{"ok":true}', @@ -754,6 +872,8 @@ describe('Function execution request', () => { code: 'print("done")', language: 'python', workspaceId: 'workspace-1', + workflowId: 'workflow-1', + executionId: 'execution-1', outputs: { files: [ { @@ -800,6 +920,7 @@ describe('Function execution request', () => { }) ) expect(data.output.result.files).toHaveLength(2) + expect(data.output.cost).toEqual({ input: 0, output: 0, total: 0.00023456 }) expect(data.resources).toEqual([ expect.objectContaining({ path: 'files/reports/chart.png' }), expect.objectContaining({ path: 'files/reports/summary.json' }), @@ -1378,9 +1499,10 @@ describe('Function execution request', () => { it('preserves output-limit classification from provider-side size inspection', async () => { envFlagsMock.isRemoteSandboxEnabled = true - mockExecuteInSandbox.mockRejectedValueOnce( - new SandboxOutputLimitError(MAX_SANDBOX_OUTPUT_BYTES + 1) - ) + const error = new SandboxOutputLimitError(MAX_SANDBOX_OUTPUT_BYTES + 1) + const cost = { input: 0, output: 0, total: 0.00023456 } + attachTrustedSandboxOutputCost(error, cost) + mockExecuteInSandbox.mockRejectedValueOnce(error) const req = createMockRequest('POST', { code: 'print("done")', @@ -1401,6 +1523,7 @@ describe('Function execution request', () => { expect(response.status).toBe(400) expect(data.error).toBe(`Sandbox output files exceed ${MAX_SANDBOX_OUTPUT_BYTES} bytes total`) + expect(data.output.cost).toEqual(cost) expect(mockWriteWorkspaceFileByPath).not.toHaveBeenCalled() }) @@ -1422,15 +1545,41 @@ describe('Function execution request', () => { expect(response.status).toBe(400) expect(data.error).toContain('must reference a regular file') + expect(data.output.cost).toBeUndefined() expect(mockWriteWorkspaceFileByPath).not.toHaveBeenCalled() }) + it.each(['/', '///', ' / '])( + 'rejects malformed workspace output destination %j before sandbox execution', + async (path) => { + envFlagsMock.isRemoteSandboxEnabled = true + + const response = await POST( + createMockRequest('POST', { + code: 'print("done")', + language: 'python', + workspaceId: 'workspace-1', + outputs: { + files: [{ path, sandboxPath: '/out/report.json' }], + }, + }) + ) + const data = await response.json() + + expect(response.status).toBe(400) + expect(data.error).toBe('Output path must include a file name') + expect(mockExecuteInSandbox).not.toHaveBeenCalled() + expect(mockExecuteShellInSandbox).not.toHaveBeenCalled() + } + ) + it('prevalidates all sandbox output destinations before writing any files', async () => { envFlagsMock.isRemoteSandboxEnabled = true mockExecuteInSandbox.mockResolvedValueOnce({ result: 'done', stdout: 'ok', sandboxId: 'sandbox-123', + cost: { input: 0, output: 0, total: 0.00023456 }, exportedFiles: { '/home/user/first.json': '{"first":true}', '/home/user/second.json': '{"second":true}', @@ -1444,6 +1593,8 @@ describe('Function execution request', () => { code: 'print("done")', language: 'python', workspaceId: 'workspace-1', + workflowId: 'workflow-1', + executionId: 'execution-1', outputs: { files: [ { @@ -1466,6 +1617,7 @@ describe('Function execution request', () => { expect(response.status).toBe(400) expect(data.success).toBe(false) expect(data.error).toContain('Directory not yet created') + expect(data.output.cost).toEqual({ input: 0, output: 0, total: 0.00023456 }) expect(mockWriteWorkspaceFileByPath).not.toHaveBeenCalled() }) @@ -1986,6 +2138,7 @@ describe('Function execution request', () => { result: null, stdout: 'generated 1 preview', sandboxId: 'sandbox-123', + cost: { input: 0, output: 0, total: 0.00034567 }, exportedFiles: { '/tmp/fellows-previews.zip': archiveBase64 }, }) @@ -1994,6 +2147,8 @@ describe('Function execution request', () => { code: source, language: 'python', workspaceId: 'workspace-1', + workflowId: 'workflow-1', + executionId: 'execution-1', sandboxId: 'fellows-sandbox', envVars: { AIRTABLE_PAT: 'stub-airtable-token', @@ -2015,6 +2170,9 @@ describe('Function execution request', () => { ) expect(response.status).toBe(200) + await expect(response.clone().json()).resolves.toMatchObject({ + output: { cost: { input: 0, output: 0, total: 0.00034567 } }, + }) const sandboxRequest = mockExecuteInSandbox.mock.calls[0][0] expect(sandboxRequest.code).toContain("['bq', 'query'") expect(sandboxRequest.code).toContain('__sim_exec_globals__["__name__"] = "__main__"') diff --git a/apps/sim/lib/function-execution/execute-request.ts b/apps/sim/lib/function-execution/execute-request.ts index 73ef60c28d5..1392221fbd0 100644 --- a/apps/sim/lib/function-execution/execute-request.ts +++ b/apps/sim/lib/function-execution/execute-request.ts @@ -82,6 +82,7 @@ import { isSandboxOutputLimitError, isSandboxOutputNotExportableError, MAX_SANDBOX_OUTPUT_BYTES, + readTrustedSandboxOutputCost, } from '@/lib/execution/remote-sandbox/output-limits' import { MAX_BLOCK_MOUNTED_FILES, @@ -128,6 +129,12 @@ const MAX_SANDBOX_OUTPUT_FILES = 20 const MAX_PRIVATE_FILE_SECRET_MATCH_EVENTS = 1_000_000 const SANDBOX_RUNTIME_PAYLOAD_PATH_ENV = '__SIM_RUNTIME_PAYLOAD_PATH' +interface FunctionExecutionCost { + input: number + output: number + total: number +} + interface SandboxRuntimePayload { params: Record environmentVariables: Record @@ -1432,10 +1439,20 @@ function exportFailure( error: string, status: number, stdout: string, - executionTime: number + executionTime: number, + cost: FunctionExecutionCost | undefined ): NextResponse { return NextResponse.json( - { success: false, error, output: { result: null, stdout: cleanStdout(stdout), executionTime } }, + { + success: false, + error, + output: { + result: null, + stdout: cleanStdout(stdout), + executionTime, + ...(cost ? { cost } : {}), + }, + }, { status } ) } @@ -1458,6 +1475,7 @@ async function maybeExportSandboxFileToWorkspace(args: { exportedFileContent?: string stdout: string executionTime: number + cost?: FunctionExecutionCost }) { const { routeContext, @@ -1473,6 +1491,7 @@ async function maybeExportSandboxFileToWorkspace(args: { exportedFileContent, stdout, executionTime, + cost, } = args if (!outputSandboxPath) return null @@ -1482,7 +1501,8 @@ async function maybeExportSandboxFileToWorkspace(args: { 'outputSandboxPath requires outputPath. Set outputPath to the destination workspace file, e.g. "files/result.csv".', 400, stdout, - executionTime + executionTime, + cost ) } @@ -1494,7 +1514,8 @@ async function maybeExportSandboxFileToWorkspace(args: { 'Workspace context required to save sandbox file to workspace', 400, stdout, - executionTime + executionTime, + cost ) } @@ -1503,7 +1524,8 @@ async function maybeExportSandboxFileToWorkspace(args: { `Sandbox file "${outputSandboxPath}" was not found or could not be read`, 500, stdout, - executionTime + executionTime, + cost ) } @@ -1521,7 +1543,8 @@ async function maybeExportSandboxFileToWorkspace(args: { `Sandbox output files exceed ${MAX_SANDBOX_OUTPUT_BYTES} bytes total`, 400, stdout, - executionTime + executionTime, + cost ) } const fileBuffer = isBinary @@ -1595,6 +1618,7 @@ async function maybeExportSandboxFileToWorkspace(args: { }, stdout: cleanStdout(stdout), executionTime, + ...(cost ? { cost } : {}), }, resources: [{ type: 'file', id: written.id, title: written.name, path: written.vfsPath }], }) @@ -1603,7 +1627,8 @@ async function maybeExportSandboxFileToWorkspace(args: { getErrorMessage(error, 'Failed to export sandbox file'), workspaceFileExportErrorStatus(error), stdout, - executionTime + executionTime, + cost ) } } @@ -1618,6 +1643,7 @@ async function maybeExportSandboxFilesToWorkspace(args: { exportedFileContent?: string stdout: string executionTime: number + cost?: FunctionExecutionCost }) { const sandboxFiles = args.outputFiles.filter((file) => file.sandboxPath) if (sandboxFiles.length === 0) return null @@ -1626,7 +1652,8 @@ async function maybeExportSandboxFilesToWorkspace(args: { `Too many sandbox output files requested (${sandboxFiles.length}). Maximum is ${MAX_SANDBOX_OUTPUT_FILES}.`, 400, args.stdout, - args.executionTime + args.executionTime, + args.cost ) } @@ -1647,6 +1674,7 @@ async function maybeExportSandboxFilesToWorkspace(args: { args.exportedFileContent, stdout: args.stdout, executionTime: args.executionTime, + cost: args.cost, }) } @@ -1658,7 +1686,8 @@ async function maybeExportSandboxFilesToWorkspace(args: { 'Workspace context required to save sandbox files to workspace', 400, args.stdout, - args.executionTime + args.executionTime, + args.cost ) } @@ -1672,7 +1701,8 @@ async function maybeExportSandboxFilesToWorkspace(args: { `Sandbox file "${sandboxPath}" was not found or could not be read`, 500, args.stdout, - args.executionTime + args.executionTime, + args.cost ) } const outputPath = file.formatPath ?? file.path @@ -1689,7 +1719,8 @@ async function maybeExportSandboxFilesToWorkspace(args: { `Sandbox output files exceed ${MAX_SANDBOX_OUTPUT_BYTES} bytes total`, 400, args.stdout, - args.executionTime + args.executionTime, + args.cost ) } const scanBuffer = isBinary ? Buffer.from(content, 'base64') : Buffer.from(content, 'utf-8') @@ -1738,7 +1769,8 @@ async function maybeExportSandboxFilesToWorkspace(args: { getErrorMessage(error, 'Invalid sandbox output destination'), workspaceFileExportErrorStatus(error), args.stdout, - args.executionTime + args.executionTime, + args.cost ) } const duplicateDestination = validationPaths.find( @@ -1749,7 +1781,8 @@ async function maybeExportSandboxFilesToWorkspace(args: { `Duplicate sandbox output destination: ${duplicateDestination}`, 400, args.stdout, - args.executionTime + args.executionTime, + args.cost ) } @@ -1805,7 +1838,8 @@ async function maybeExportSandboxFilesToWorkspace(args: { getErrorMessage(error, 'Failed to export sandbox files'), workspaceFileExportErrorStatus(error), args.stdout, - args.executionTime + args.executionTime, + args.cost ) } @@ -1844,6 +1878,7 @@ async function maybeExportSandboxFilesToWorkspace(args: { }, stdout: cleanStdout(args.stdout), executionTime: args.executionTime, + ...(args.cost ? { cost: args.cost } : {}), }, resources: writtenFiles.map((file) => ({ type: 'file', @@ -1933,6 +1968,7 @@ async function collectExecutionOutputFiles(args: { collectedFiles: SandboxCollectedFile[] stdout: string executionTime: number + cost?: FunctionExecutionCost }): Promise<{ files: UserFile[] } | { response: NextResponse }> { const { routeContext, collectedFiles } = args if (collectedFiles.length === 0) return { files: [] } @@ -1949,7 +1985,8 @@ async function collectExecutionOutputFiles(args: { 'Workspace, workflow, and execution context are required to return files from the sandbox.', 400, args.stdout, - args.executionTime + args.executionTime, + args.cost ), } } @@ -1989,7 +2026,8 @@ async function collectExecutionOutputFiles(args: { `Sandbox output file "${name}" contains a resolved secret value and was not returned. Write the file without embedding secret values, or export it to a workspace file where its provenance can be recorded.`, 400, args.stdout, - args.executionTime + args.executionTime, + args.cost ), } } @@ -2131,6 +2169,8 @@ export async function executeFunctionRequest( _sandboxFiles, } = body + const meterRemoteSandboxUsage = Boolean(workflowId && !isCustomTool && !usesMothershipSandbox) + if (selectedSandboxId && !isRemoteSandboxEnabled) { return NextResponse.json( { success: false, error: 'The Function code sandbox is not configured' }, @@ -2189,6 +2229,23 @@ export async function executeFunctionRequest( privateResolvedSecretNamesMetadataType ) } + try { + for (const file of outputFiles) { + normalizeOutputWorkspaceFileName(file.formatPath ?? file.path) + } + } catch (error) { + return appendPrivateResolvedSecretNames( + NextResponse.json( + { + success: false, + error: getErrorMessage(error, 'Invalid sandbox output destination'), + }, + { status: 400 } + ), + includePrivateResolvedSecretNames ? [] : null, + privateResolvedSecretNamesMetadataType + ) + } // Planned before the runtime is chosen because it is pure: it decides whether // this execution needs a sandbox filesystem at all, without spending a presign @@ -2500,6 +2557,7 @@ export async function executeFunctionRequest( exportedFileContent, exportedFiles, collectedFiles: shellCollectedFiles, + cost: shellCost, } = await executeShellInSandbox({ code: resolvedCode, envs: shellEnvs, @@ -2515,6 +2573,7 @@ export async function executeFunctionRequest( ? { sandboxKind: 'mothership' as const } : {}), signal: executionSignal, + meterUsage: meterRemoteSandboxUsage, }) const executionTime = Date.now() - execStart @@ -2529,7 +2588,12 @@ export async function executeFunctionRequest( { success: false, error: scrubInternalIdentifiers(shellError, compilerInternalIdentifiers), - output: { result: null, stdout: cleanStdout(shellStdout), executionTime }, + output: { + result: null, + stdout: cleanStdout(shellStdout), + executionTime, + ...(shellCost ? { cost: shellCost } : {}), + }, }, routeContext, { status: 422 } @@ -2547,6 +2611,7 @@ export async function executeFunctionRequest( exportedFileContent, stdout: shellStdout, executionTime, + cost: shellCost, }) if (fileExportResponse) { return appendResolvedSecretNames(fileExportResponse, routeContext) @@ -2562,6 +2627,7 @@ export async function executeFunctionRequest( collectedFiles: shellCollectedFiles ?? [], stdout: shellStdout, executionTime, + cost: shellCost, }) if ('response' in shellOutputFiles) { return appendResolvedSecretNames(shellOutputFiles.response, routeContext) @@ -2575,6 +2641,7 @@ export async function executeFunctionRequest( stdout: cleanStdout(shellStdout), executionTime, files: shellOutputFiles.files, + ...(shellCost ? { cost: shellCost } : {}), }, }, routeContext @@ -2637,6 +2704,7 @@ export async function executeFunctionRequest( exportedFileContent, exportedFiles, collectedFiles: jsCollectedFiles, + cost: sandboxCost, } = await executeInSandbox({ code: codeForE2B, language: CodeLanguage.JavaScript, @@ -2653,6 +2721,7 @@ export async function executeFunctionRequest( ? { sandboxKind: 'mothership' as const } : {}), signal: executionSignal, + meterUsage: meterRemoteSandboxUsage, }) const executionTime = Date.now() - execStart stdout += e2bStdout @@ -2678,7 +2747,12 @@ export async function executeFunctionRequest( { success: false, error: formattedError, - output: { result: null, stdout: cleanedOutput, executionTime }, + output: { + result: null, + stdout: cleanedOutput, + executionTime, + ...(sandboxCost ? { cost: sandboxCost } : {}), + }, }, routeContext, { status: 422 } @@ -2696,6 +2770,7 @@ export async function executeFunctionRequest( exportedFileContent, stdout, executionTime, + cost: sandboxCost, }) if (fileExportResponse) { return appendResolvedSecretNames(fileExportResponse, routeContext) @@ -2711,6 +2786,7 @@ export async function executeFunctionRequest( collectedFiles: jsCollectedFiles ?? [], stdout, executionTime, + cost: sandboxCost, }) if ('response' in jsOutputFiles) { return appendResolvedSecretNames(jsOutputFiles.response, routeContext) @@ -2724,6 +2800,7 @@ export async function executeFunctionRequest( stdout: cleanStdout(stdout), executionTime, files: jsOutputFiles.files, + ...(sandboxCost ? { cost: sandboxCost } : {}), }, }, routeContext @@ -2749,6 +2826,7 @@ export async function executeFunctionRequest( exportedFileContent, exportedFiles, collectedFiles: pythonCollectedFiles, + cost: sandboxCost, } = await executeInSandbox({ code: codeForE2B, language: CodeLanguage.Python, @@ -2764,6 +2842,7 @@ export async function executeFunctionRequest( ? { sandboxKind: 'mothership' as const } : {}), signal: executionSignal, + meterUsage: meterRemoteSandboxUsage, }) const executionTime = Date.now() - execStart stdout += e2bStdout @@ -2789,7 +2868,12 @@ export async function executeFunctionRequest( { success: false, error: formattedError, - output: { result: null, stdout: cleanedOutput, executionTime }, + output: { + result: null, + stdout: cleanedOutput, + executionTime, + ...(sandboxCost ? { cost: sandboxCost } : {}), + }, }, routeContext, { status: 422 } @@ -2807,6 +2891,7 @@ export async function executeFunctionRequest( exportedFileContent, stdout, executionTime, + cost: sandboxCost, }) if (fileExportResponse) { return appendResolvedSecretNames(fileExportResponse, routeContext) @@ -2822,6 +2907,7 @@ export async function executeFunctionRequest( collectedFiles: pythonCollectedFiles ?? [], stdout, executionTime, + cost: sandboxCost, }) if ('response' in pythonOutputFiles) { return appendResolvedSecretNames(pythonOutputFiles.response, routeContext) @@ -2835,6 +2921,7 @@ export async function executeFunctionRequest( stdout: cleanStdout(stdout), executionTime, files: pythonOutputFiles.files, + ...(sandboxCost ? { cost: sandboxCost } : {}), }, }, routeContext @@ -3019,10 +3106,16 @@ export async function executeFunctionRequest( isSandboxOutputFileError(error) || isSandboxOutputNotExportableError(error) ) { + const cost = readTrustedSandboxOutputCost(error) const outputLimitResponse = { success: false, error: error.message, - output: { result: null, stdout: cleanStdout(stdout), executionTime }, + output: { + result: null, + stdout: cleanStdout(stdout), + executionTime, + ...(cost ? { cost } : {}), + }, } return routeContext ? functionJsonResponse(outputLimitResponse, routeContext, { status: 400 }) diff --git a/apps/sim/providers/cost-policy.test.ts b/apps/sim/providers/cost-policy.test.ts index 5d7ebd9869e..48ced0fffcd 100644 --- a/apps/sim/providers/cost-policy.test.ts +++ b/apps/sim/providers/cost-policy.test.ts @@ -226,6 +226,23 @@ describe('installStreamingCostPolicy', () => { expect(output.cost).toMatchObject({ input: 0, output: 0, total: 0.75, toolCost: 0.75 }) }) + it('adds late failed Function cost once without applying the model multiplier', () => { + const output = { + cost: { input: 1, output: 2, total: 3.25, toolCost: 0.25 }, + } as NormalizedBlockOutput + const failedFunctionToolCost = { total: 0 } + installStreamingCostPolicy( + output, + { billable: false, multiplier: 0 }, + () => failedFunctionToolCost.total + ) + + failedFunctionToolCost.total = 0.125 + + expect(output.cost).toMatchObject({ input: 0, output: 0, total: 0.375, toolCost: 0.375 }) + expect(output.cost).toMatchObject({ total: 0.375, toolCost: 0.375 }) + }) + it('zeroes model cost written by a provider for a model Sim does not host', () => { const output = { cost: { input: 0, output: 0, total: 0 } } as NormalizedBlockOutput installStreamingCostPolicy(output, resolveModelCostPolicy(SELF_KEYED_MODEL)) diff --git a/apps/sim/providers/cost-policy.ts b/apps/sim/providers/cost-policy.ts index 7c67d82dee8..b116a90c472 100644 --- a/apps/sim/providers/cost-policy.ts +++ b/apps/sim/providers/cost-policy.ts @@ -266,12 +266,23 @@ export function resolveProxiedModelCost(cost: unknown): ModelCost { */ export function installStreamingCostPolicy( output: NormalizedBlockOutput, - policy: ModelCostPolicy + policy: ModelCostPolicy, + additionalToolCost?: () => number ): void { let raw = output.cost as ModelCost | undefined Object.defineProperty(output, 'cost', { - get: () => applyModelCostPolicy(raw, policy), + get: () => { + const projected = applyModelCostPolicy(raw, policy) + const additional = additionalToolCost?.() ?? 0 + if (!Number.isFinite(additional) || additional <= 0) return projected + + return { + ...projected, + toolCost: roundCost((projected.toolCost ?? 0) + additional), + total: roundCost(projected.total + additional), + } + }, set: (value: ModelCost | undefined) => { raw = value }, diff --git a/apps/sim/providers/index.test.ts b/apps/sim/providers/index.test.ts index 22efa428d51..f6f962b71e5 100644 --- a/apps/sim/providers/index.test.ts +++ b/apps/sim/providers/index.test.ts @@ -278,6 +278,33 @@ describe('executeProviderRequest — BYOK regression', () => { expect(result.cost?.total).toBeCloseTo(0.00675, 8) }) + it('adds failed Function cost once alongside successful tool results', async () => { + mockGetApiKeyWithBYOK.mockResolvedValue({ apiKey: 'sk-byok', isBYOK: true }) + mockExecuteTool.mockResolvedValueOnce({ + success: false, + output: { cost: { total: 0.004 } }, + error: 'execution failed', + }) + mockExecuteRequest.mockImplementationOnce(async () => { + const execution = await executeProviderTool('function_execute', {}) + expect(execution.rawResponse.success).toBe(false) + return { + ...makeAnthropicResponse(), + toolResults: [{ cost: { total: 0.005 } }], + } as ProviderResponse + }) + + const result = (await executeProviderRequest('anthropic', { + model: 'claude-opus-4-6', + workspaceId: 'ws-1', + tools: [makeProviderTool('function_execute', 'credential')], + })) as ProviderResponse + + expect(result.cost).toMatchObject({ input: 0, output: 0 }) + expect(result.cost?.toolCost).toBeCloseTo(0.009, 8) + expect(result.cost?.total).toBeCloseTo(0.009, 8) + }) + /** * Gemini hands the same cost object to its response and its model segment. * Adding tool cost by mutation would charge it to the segment too. diff --git a/apps/sim/providers/index.ts b/apps/sim/providers/index.ts index d8ecbb04388..f4db8ee535a 100644 --- a/apps/sim/providers/index.ts +++ b/apps/sim/providers/index.ts @@ -152,14 +152,18 @@ function isReadableStream(response: any): response is ReadableStream { * stream drain — long after this function returns — so the policy is installed * on the live output object rather than applied to a value. */ -function applyStreamingCostPolicy(response: StreamingExecution, policy: ModelCostPolicy): void { +function applyStreamingCostPolicy( + response: StreamingExecution, + policy: ModelCostPolicy, + additionalToolCost?: () => number +): void { const output = response.execution?.output if (!output || typeof output !== 'object') { logger.warn('Streaming output unavailable at intercept time; cost policy not applied') return } - installStreamingCostPolicy(output, policy) + installStreamingCostPolicy(output, policy, additionalToolCost) const segments = output.providerTiming?.timeSegments if (Array.isArray(segments)) { @@ -224,16 +228,19 @@ export async function executeProviderRequest( const provenanceSafeRequest = await omitUnsafeProviderFileAttachments(sanitizedRequest) const modelSafeRequest = provenanceSafeRequest const toolIdentities = assignProviderToolIdentities(modelSafeRequest.tools) - const requestRuntimeContext = - toolIdentities.toolIdByWireId.size > 0 + const failedFunctionToolCost = { total: 0 } + const requestRuntimeContext: ProviderRuntimeContext = { + ...runtimeContext, + failedFunctionToolCost, + ...(toolIdentities.toolIdByWireId.size > 0 ? { - ...runtimeContext, toolIdByWireId: new Map([ ...(runtimeContext?.toolIdByWireId ?? []), ...toolIdentities.toolIdByWireId, ]), } - : runtimeContext + : {}), + } if (modelSafeRequest.responseFormat) { const structuredOutputInstructions = generateStructuredOutputInstructions( @@ -254,7 +261,11 @@ export async function executeProviderRequest( if (isStreamingExecution(response)) { logger.info('Provider returned StreamingExecution', { isBYOK }) - applyStreamingCostPolicy(response, resolveModelCostPolicy(sanitizedRequest.model, isBYOK)) + applyStreamingCostPolicy( + response, + resolveModelCostPolicy(sanitizedRequest.model, isBYOK), + () => failedFunctionToolCost.total + ) projectStreamingExecutionToolIdentities(response, toolIdentities) return response } @@ -300,7 +311,7 @@ export async function executeProviderRequest( applySegmentCostPolicy(response.timing.timeSegments, costPolicy) } - const toolCost = sumToolCosts(response.toolResults) + const toolCost = sumToolCosts(response.toolResults) + failedFunctionToolCost.total if (toolCost > 0 && response.cost) { // Replaced rather than mutated: a provider-supplied cost can be the same // object it also handed to a time segment, and tool cost belongs only to diff --git a/apps/sim/providers/runtime-context.test.ts b/apps/sim/providers/runtime-context.test.ts index c4761eea0b6..6a7c718b2da 100644 --- a/apps/sim/providers/runtime-context.test.ts +++ b/apps/sim/providers/runtime-context.test.ts @@ -147,6 +147,35 @@ describe('provider runtime context', () => { ) }) + it('accumulates cost only for failed canonical Function results', async () => { + const failedFunctionToolCost = { total: 0 } + const context = { + failedFunctionToolCost, + toolIdByWireId: new Map([['function_execute__sim_2', 'function_execute']]), + } + + mockExecuteTool + .mockResolvedValueOnce({ + success: false, + output: { cost: { total: 0.125 } }, + error: 'execution failed', + }) + .mockResolvedValueOnce({ success: true, output: { cost: { total: 4 } } }) + .mockResolvedValueOnce({ + success: false, + output: { cost: { total: 8 } }, + error: 'other tool failed', + }) + + await runWithProviderRuntimeContext(context, () => + executeProviderTool('function_execute__sim_2', {}) + ) + await runWithProviderRuntimeContext(context, () => executeProviderTool('function_execute', {})) + await runWithProviderRuntimeContext(context, () => executeProviderTool('exa_search', {})) + + expect(failedFunctionToolCost.total).toBe(0.125) + }) + it('rebinds a prompt-exposed environment placeholder for the exact tool call', async () => { const sourceRegistry = new ResolvedSecretTraceRegistry([ { diff --git a/apps/sim/providers/runtime-context.ts b/apps/sim/providers/runtime-context.ts index 2e602a83ba0..3cfea3e1eab 100644 --- a/apps/sim/providers/runtime-context.ts +++ b/apps/sim/providers/runtime-context.ts @@ -19,6 +19,8 @@ export interface ProviderRuntimeContext { executionContext?: ExecutionContext /** Request-scoped provider wire ids mapped back to canonical tool registry ids. */ toolIdByWireId?: ReadonlyMap + /** Failed canonical Function cost omitted from provider tool-result collections. */ + failedFunctionToolCost?: { total: number } } export type ExecuteProviderToolOptions = ExecuteToolOptions @@ -88,6 +90,20 @@ function withoutChildTraceHandle(response: ToolResponse): ToolResponse { } } +function accumulateFailedFunctionToolCost( + toolId: string, + result: ToolResponse, + accumulator: ProviderRuntimeContext['failedFunctionToolCost'] +): void { + if (toolId !== 'function_execute' || result.success || !accumulator) return + if (!isRecordLike(result.output) || !isRecordLike(result.output.cost)) return + + const total = result.output.cost.total + if (typeof total === 'number' && Number.isFinite(total) && total > 0) { + accumulator.total += total + } +} + export async function executeProviderTool( toolId: string, params: Parameters[1], @@ -120,6 +136,11 @@ export async function executeProviderTool( ...(executionContext ? { executionContext } : {}), resolvedSecretTraceRegistry: toolCallRegistry, }) + accumulateFailedFunctionToolCost( + executionToolId, + result, + runtimeContext?.failedFunctionToolCost + ) if (!registry || !toolCallRegistry) { return { rawResponse: result, modelResponse: withoutChildTraceHandle(result) } } diff --git a/apps/sim/scripts/build-function-daytona-snapshot.ts b/apps/sim/scripts/build-function-daytona-snapshot.ts index 5dfae90c3cf..2c6b85b0be6 100644 --- a/apps/sim/scripts/build-function-daytona-snapshot.ts +++ b/apps/sim/scripts/build-function-daytona-snapshot.ts @@ -20,6 +20,7 @@ import { isImmutableDaytonaSnapshotRef, } from '@sim/utils/sandbox-references' import { + FUNCTION_DAYTONA_DISK_GB, FUNCTION_SANDBOX_CPU_COUNT, FUNCTION_SANDBOX_MEMORY_GB, } from '@/lib/execution/remote-sandbox/function-resources' @@ -53,7 +54,7 @@ const APT_INSTALL = 'DEBIAN_FRONTEND=noninteractive apt-get install -y --no-inst const RESOURCES = { cpu: FUNCTION_SANDBOX_CPU_COUNT, memory: FUNCTION_SANDBOX_MEMORY_GB, - disk: 10, + disk: FUNCTION_DAYTONA_DISK_GB, } as const export function createFunctionImage(manifest: FunctionSandboxParityManifest) { diff --git a/apps/sim/tools/function/execute.test.ts b/apps/sim/tools/function/execute.test.ts index 053ce3c3372..f4dc0573265 100644 --- a/apps/sim/tools/function/execute.test.ts +++ b/apps/sim/tools/function/execute.test.ts @@ -120,4 +120,41 @@ describe('Function Execute Tool', () => { expect(body[PRIVATE_SECRET_PROVENANCE_FIELD]).toEqual(bundle) expect(JSON.stringify(body)).not.toContain('plaintext') }) + + it('preserves sandbox cost in a successful Function result', async () => { + const cost = { input: 0, output: 0, total: 0.00012345 } + const result = await functionExecuteTool.transformResponse?.( + Response.json({ + success: true, + output: { result: 42, stdout: 'done', cost }, + }), + { code: 'return 42' } + ) + + expect(result).toMatchObject({ + success: true, + output: { result: 42, stdout: 'done', cost }, + }) + }) + + it('preserves sandbox cost in a failed Function result', async () => { + const cost = { input: 0, output: 0, total: 0.00012345 } + const result = await functionExecuteTool.transformResponse?.( + Response.json( + { + success: false, + error: 'boom', + output: { result: null, stdout: 'trace', cost }, + }, + { status: 422 } + ), + { code: 'throw new Error("boom")' } + ) + + expect(result).toMatchObject({ + success: false, + output: { result: null, stdout: 'trace', cost }, + error: 'boom', + }) + }) }) diff --git a/apps/sim/tools/function/execute.ts b/apps/sim/tools/function/execute.ts index 8da668e12a9..ab2cde32a1d 100644 --- a/apps/sim/tools/function/execute.ts +++ b/apps/sim/tools/function/execute.ts @@ -244,6 +244,7 @@ To return a file, write it to ${SANDBOX_OUTPUT_DIR}. Everything there comes back // Always an array, never undefined: a declared `file[]` output that is // missing warns on every call, and this branch runs for every failure. files: result.output?.files ?? [], + ...(result.output?.cost ? { cost: result.output.cost } : {}), }, error: result.error, retryable: result.retryable, @@ -259,6 +260,7 @@ To return a file, write it to ${SANDBOX_OUTPUT_DIR}. Everything there comes back result: result.output.result, stdout: result.output.stdout, files: result.output.files ?? [], + ...(result.output.cost ? { cost: result.output.cost } : {}), }, resources: result.resources, largeValueKeys: result.largeValueKeys, diff --git a/apps/sim/tools/function/types.ts b/apps/sim/tools/function/types.ts index 9c590a7fa5c..9dd38298582 100644 --- a/apps/sim/tools/function/types.ts +++ b/apps/sim/tools/function/types.ts @@ -86,5 +86,10 @@ export interface CodeExecutionOutput extends ToolResponse { stdout: string /** Files harvested from the sandbox output directory, already persisted. */ files: UserFile[] + cost?: { + input: number + output: number + total: number + } } } diff --git a/apps/sim/tools/index.test.ts b/apps/sim/tools/index.test.ts index a16de831d9f..e9d2b43c50b 100644 --- a/apps/sim/tools/index.test.ts +++ b/apps/sim/tools/index.test.ts @@ -800,7 +800,7 @@ describe('executeTool Function', () => { cleanupEnvVars() }) - it('executes trusted Function calls in process without dropping resolved execution context', async () => { + it('stamps standard Function identity and preserves trusted execution context', async () => { const fetchSpy = vi.fn() global.fetch = Object.assign(fetchSpy, { preconnect: vi.fn() }) as typeof fetch @@ -856,7 +856,7 @@ describe('executeTool Function', () => { workspaceId: 'workspace-456', body: { code: 'return [{{API_KEY}}, __blockRef_0.field]', - isCustomTool: true, + isCustomTool: false, inputs: { location: 'San Francisco' }, envVars: { API_KEY: 'resolved-secret' }, contextVariables: { @@ -1744,6 +1744,7 @@ describe('executeTool Function', () => { it('does not log plaintext or runtime aliases from Function errors', async () => { const secret = 'function-error-secret-value' const runtimeAlias = '__var_API_KEY' + const cost = { input: 0, output: 0, total: 0.00012345 } const registry = new ResolvedSecretTraceRegistry([ { name: 'API_KEY', @@ -1756,6 +1757,7 @@ describe('executeTool Function', () => { JSON.stringify({ success: false, error: `Execution failed with ${secret} via ${runtimeAlias}`, + output: { result: null, stdout: 'trace', cost }, __resolvedSecretNames: ['API_KEY'], }), { @@ -1782,6 +1784,7 @@ describe('executeTool Function', () => { expect(result.success).toBe(false) expect(result.error).toContain(secret) + expect(result.output?.cost).toEqual(cost) expect(JSON.stringify(mockToolsLogger.error.mock.calls)).not.toContain(secret) expect(JSON.stringify(mockToolsLogger.error.mock.calls)).not.toContain(runtimeAlias) expect(JSON.stringify(projectToolResultForCopilot(result, registry))).not.toContain(secret) @@ -1824,6 +1827,32 @@ describe('executeTool Function', () => { expect(JSON.stringify(mockToolsLogger.error.mock.calls)).not.toContain(runtimeAlias) }) + it('does not lift an invalid sandbox cost from a Function error response', async () => { + mockExecuteFunction.mockResolvedValueOnce( + Response.json( + { + success: false, + error: 'boom', + output: { + result: null, + stdout: 'trace', + cost: { input: 0, output: 0, total: -1 }, + }, + }, + { status: 422 } + ) + ) + + const result = await executeTool( + 'function_execute', + { code: 'throw new Error("boom")' }, + { executionContext: createToolExecutionContext({ userId: 'user-1' }) } + ) + + expect(result.success).toBe(false) + expect(result.output).not.toHaveProperty('cost') + }) + it('does not log a secret-bearing non-OK response stream error', async () => { const secret = 'function-body-stream-secret-value' const streamError = `${secret} __var_API_KEY __sim_code_0_binding_0` diff --git a/apps/sim/tools/index.ts b/apps/sim/tools/index.ts index 847ac33a4e2..72c8ba6f605 100644 --- a/apps/sim/tools/index.ts +++ b/apps/sim/tools/index.ts @@ -1629,11 +1629,12 @@ async function executeToolImplementation( const startTime = new Date() const startTimeISO = startTime.toISOString() const requestId = generateRequestId() + const normalizedToolId = normalizeToolId(toolId) const privateToolMetadataPolicy = resolvedSecretTraceRegistry ? getPrivateToolMetadataPolicy(toolId) : undefined const structuralOnlyToolLogs = - normalizeToolId(toolId) === 'function_execute' || + normalizedToolId === 'function_execute' || isCustomTool(toolId) || privateToolMetadataPolicy !== undefined @@ -1645,7 +1646,6 @@ async function executeToolImplementation( let tool: ExecutableToolConfig | undefined // Preserve direct-call compatibility with legacy resource-suffixed tool ids. - const normalizedToolId = normalizeToolId(toolId) if (internalSandboxProfile && normalizedToolId !== 'function_execute') { throw new Error('An internal sandbox profile may only be used with function_execute') } @@ -2283,9 +2283,14 @@ async function executeToolImplementation( const rawResponseData = error instanceof Error && 'data' in error ? (error as { data?: unknown }).data : undefined const responseData = isRecordLike(rawResponseData) ? rawResponseData : undefined + const functionSandboxCost = + normalizedToolId === 'function_execute' ? readFunctionSandboxCost(responseData) : undefined return { success: false, - output: errorDetails, + output: { + ...errorDetails, + ...(functionSandboxCost ? { cost: functionSandboxCost } : {}), + }, error: errorMessage, ...(responseData?.retryable === false ? { retryable: false } : {}), // Sim's own status (hosted-key 429/503) survives the flattening from a @@ -2446,6 +2451,33 @@ function isFunctionExecuteBody(value: unknown): value is FunctionExecuteBody { return isPlainRecord(value) && typeof value.code === 'string' } +interface FunctionSandboxCost { + input: number + output: number + total: number +} + +function readFunctionSandboxCost(value: unknown): FunctionSandboxCost | undefined { + if (!isRecordLike(value) || !isRecordLike(value.output) || !isRecordLike(value.output.cost)) { + return undefined + } + const { input, output, total } = value.output.cost + if ( + typeof input !== 'number' || + !Number.isFinite(input) || + input < 0 || + typeof output !== 'number' || + !Number.isFinite(output) || + output < 0 || + typeof total !== 'number' || + !Number.isFinite(total) || + total < 0 + ) { + return undefined + } + return { input, output, total } +} + function isToolResponse(value: unknown): value is ToolResponse { return isRecordLike(value) && typeof value.success === 'boolean' && isRecordLike(value.output) } @@ -2470,12 +2502,16 @@ async function executeDeclaredInternalOperation({ const operationParams = projectToolModelInputParams(tool, params, resolvedSecretTraceRegistry) let operationInput = tool.operation.input(operationParams) - const isFunctionOperation = toolId === 'function_execute' || isCustomTool(toolId) - if (isFunctionOperation && !isFunctionExecuteBody(operationInput)) { - throw new Error('Function operation input must be an object') + const isRegisteredCustomTool = isCustomTool(toolId) + const isFunctionOperation = toolId === 'function_execute' || isRegisteredCustomTool + if (isFunctionOperation) { + if (!isFunctionExecuteBody(operationInput)) { + throw new Error('Function operation input must be an object') + } + operationInput = { ...operationInput, isCustomTool: isRegisteredCustomTool } } if ( - isCustomTool(toolId) && + isRegisteredCustomTool && isFunctionExecuteBody(operationInput) && 'schema' in operationInput && 'params' in operationInput