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..fac437b3f78 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 @@ -506,15 +526,40 @@ export class BlockExecutor { const policy = resolveBlockRetryPolicy(block) if (!policy) return invoke() + const shouldAccumulateFunctionCost = block.metadata?.id === BlockType.FUNCTION + let accumulatedFunctionCost: TrustedExecutionCost | undefined let tries = 0 try { for (;;) { tries++ try { - return await invoke() + const output = await invoke() + if (!shouldAccumulateFunctionCost || !accumulatedFunctionCost || !isRecordLike(output)) { + return output + } + + const totalCost = addTrustedExecutionCosts( + accumulatedFunctionCost, + readTrustedExecutionCost(output) + ) + if (!totalCost) return output + + const outputWithCost = { ...output, cost: totalCost } + attachTrustedExecutionCost(outputWithCost, totalCost) + return outputWithCost as T } catch (error) { + if (shouldAccumulateFunctionCost) { + accumulatedFunctionCost = addTrustedExecutionCosts( + accumulatedFunctionCost, + readTrustedExecutionCost(error) + ) + } + const isFinalTry = tries >= policy.maxTries - if (isFinalTry || ctx.abortSignal?.aborted || !isRetryableBlockError(error)) throw error + if (isFinalTry || ctx.abortSignal?.aborted || !isRetryableBlockError(error)) { + attachTrustedExecutionCost(error, accumulatedFunctionCost) + throw error + } this.execLogger.warn('Block failed; retrying', { blockId: block.id, @@ -528,7 +573,10 @@ export class BlockExecutor { if (policy.waitBetweenTriesMs > 0) await sleep(policy.waitBetweenTriesMs) /** `sleep` is not abort-aware, so a run stopped mid-wait must not start another try. */ - if (ctx.abortSignal?.aborted) throw error + if (ctx.abortSignal?.aborted) { + attachTrustedExecutionCost(error, accumulatedFunctionCost) + throw error + } } } } finally { @@ -620,8 +668,10 @@ export class BlockExecutor { return softOutput } + const trustedExecutionCost = readTrustedExecutionCost(error) const errorOutput: NormalizedBlockOutput = { error: errorMessage, + ...(trustedExecutionCost ? { cost: trustedExecutionCost } : {}), } // Keep any answer text already drained before timeout/failure so logs match 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/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..be2d41a6fd7 100644 --- a/apps/sim/executor/handlers/pi/local/sim-tools.ts +++ b/apps/sim/executor/handlers/pi/local/sim-tools.ts @@ -47,6 +47,10 @@ type PiToolResultProjection = | { safe: true; result: PiToolResult } | { safe: false; result: PiToolResult } +export interface PiFunctionToolCostAccumulator { + total: number +} + function projectToolResult( result: ToolResponse, registry: ResolvedSecretTraceRegistry | undefined @@ -107,7 +111,8 @@ function buildSimToolSpec( ctx: ExecutionContext, inputTools: ToolInput[], provider: ProviderToolConfig, - toolIndex: number + toolIndex: number, + functionToolCost?: PiFunctionToolCostAccumulator ): PiToolSpec { const toolId = provider.canonicalId ?? provider.id const preseededParams = provider.params || {} @@ -170,6 +175,20 @@ function buildSimToolSpec( resolvedSecretTraceRegistry: toolCallRegistry, } ) + const resultCost = result.output?.cost + const resultCostTotal = + resultCost && typeof resultCost === 'object' + ? (resultCost as Record).total + : undefined + if ( + toolId === 'function_execute' && + functionToolCost && + typeof resultCostTotal === 'number' && + Number.isFinite(resultCostTotal) && + resultCostTotal > 0 + ) { + functionToolCost.total += resultCostTotal + } const projection = projectToolResult(result, toolCallRegistry?.forkForPropagatedEntries()) if (projection.safe && registry && toolCallRegistry?.isComplete()) { registry.mergeToolCallRegistry(toolCallRegistry) @@ -199,7 +218,8 @@ function buildSimToolSpec( */ export async function buildSimToolSpecs( ctx: ExecutionContext, - inputTools: unknown + inputTools: unknown, + functionToolCost?: PiFunctionToolCostAccumulator ): Promise { if (!Array.isArray(inputTools)) return [] @@ -243,6 +263,6 @@ export async function buildSimToolSpecs( await annotateDuplicateToolBindings(ctx, providers) assignProviderToolIdentities(providers) return configuredTools.map(({ provider, toolIndex }) => - buildSimToolSpec(ctx, inputTools, provider, toolIndex) + buildSimToolSpec(ctx, inputTools, provider, toolIndex, functionToolCost) ) } diff --git a/apps/sim/executor/handlers/pi/pi-handler.test.ts b/apps/sim/executor/handlers/pi/pi-handler.test.ts index 063b319c6ff..aee9b0fd0f6 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', () => ({ @@ -153,6 +155,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 +278,19 @@ 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('routes Create PR to the cloud backend and surfaces PR output', async () => { const output = (await handler.execute(ctx(), block, { mode: 'cloud', @@ -942,6 +958,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 +987,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..d152fa52cf4 100644 --- a/apps/sim/executor/handlers/pi/pi-handler.ts +++ b/apps/sim/executor/handlers/pi/pi-handler.ts @@ -45,7 +45,10 @@ import { resolvePiSearchKey, } from '@/executor/handlers/pi/core/keys' import { runLocalPi } from '@/executor/handlers/pi/local/backend' -import { buildSimToolSpecs } from '@/executor/handlers/pi/local/sim-tools' +import { + buildSimToolSpecs, + type PiFunctionToolCostAccumulator, +} from '@/executor/handlers/pi/local/sim-tools' import { buildPiSearchToolSpec } from '@/executor/handlers/pi/search/tool' import type { BlockHandler, @@ -267,7 +270,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 functionToolCost: PiFunctionToolCostAccumulator = { total: 0 } + const tools = await buildSimToolSpecs(ctx, inputs.tools, functionToolCost) const params: PiLocalRunParams = { ...contextualBase, mode: 'local', @@ -282,7 +286,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, functionToolCost) } const owner = asOptString(inputs.owner) @@ -473,10 +477,20 @@ export class PiBlockHandler implements BlockHandler { model: string, isBYOK: boolean, startTime: number, - startTimeISO: string + startTimeISO: string, + functionToolCost = 0 ): NormalizedBlockOutput { const { totals } = result const endTime = Date.now() + const modelCost = computePiCost(model, totals.inputTokens, totals.outputTokens, isBYOK) + const cost = + functionToolCost > 0 + ? { + ...modelCost, + toolCost: functionToolCost, + total: modelCost.total + functionToolCost, + } + : modelCost return { content: totals.finalText, model, @@ -505,7 +519,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 +533,8 @@ export class PiBlockHandler implements BlockHandler { block: SerializedBlock, backend: PiBackendRun

, params: P, - memoryConfig?: PiMemoryConfig + memoryConfig?: PiMemoryConfig, + functionToolCost?: PiFunctionToolCostAccumulator ): Promise { const startTime = Date.now() const startTimeISO = new Date(startTime).toISOString() @@ -561,7 +576,8 @@ export class PiBlockHandler implements BlockHandler { params.model, params.isBYOK, startTime, - startTimeISO + startTimeISO, + functionToolCost?.total ) ) if (memoryConfig) { @@ -610,7 +626,8 @@ export class PiBlockHandler implements BlockHandler { params.model, params.isBYOK, startTime, - startTimeISO + startTimeISO, + functionToolCost?.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 0f4271033ef..842a7919a56 100644 --- a/apps/sim/lib/execution/remote-sandbox/conformance.test.ts +++ b/apps/sim/lib/execution/remote-sandbox/conformance.test.ts @@ -118,12 +118,20 @@ 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_PROCESS_OUTPUT_BYTES, MAX_SANDBOX_STREAMED_OUTPUT_TAIL_BYTES, + readTrustedSandboxOutputCost, } from '@/lib/execution/remote-sandbox/output-limits' import { PI_SANDBOX_MIN_LIFETIME_MS, @@ -134,6 +142,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 @@ -330,6 +359,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 () => { @@ -354,10 +424,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 () => { @@ -386,13 +458,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) }) @@ -495,11 +573,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 () => { @@ -816,11 +896,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 () => { @@ -922,18 +1008,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() }) @@ -986,15 +1078,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 () => { @@ -1313,6 +1486,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')) @@ -2253,10 +2478,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', @@ -2280,9 +2507,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 0df2ae61443..4f36328c9ba 100644 --- a/apps/sim/lib/execution/remote-sandbox/daytona.ts +++ b/apps/sim/lib/execution/remote-sandbox/daytona.ts @@ -41,6 +41,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)) @@ -294,6 +299,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 @@ -332,6 +338,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 @@ -653,7 +660,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)) { @@ -674,6 +688,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 { @@ -775,6 +795,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) { @@ -790,11 +811,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 db387fd93e7..e7cb1564aa5 100644 --- a/apps/sim/lib/execution/remote-sandbox/e2b.ts +++ b/apps/sim/lib/execution/remote-sandbox/e2b.ts @@ -95,6 +95,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 = @@ -365,7 +370,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, @@ -375,6 +380,7 @@ class E2BSandboxHandle implements SandboxHandle { value: E2B_PROVIDER_LIMIT_ERROR, traceback: E2B_PROVIDER_LIMIT_ERROR, }, + providerFailure: result.providerFailure, } } return processCodeFailure(result) @@ -534,7 +540,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 { @@ -842,6 +853,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) { @@ -860,7 +872,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 } : {}), @@ -868,6 +882,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 7055a88ec2e..bcb4af59343 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,6 +15,7 @@ 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, MAX_SANDBOX_OUTPUT_BYTES, @@ -29,12 +35,15 @@ import type { CreateSandboxOptions, SandboxCodeResult, SandboxCommandResult, + SandboxExecutionCost, SandboxExecutionRequest, SandboxExecutionResult, SandboxFile, SandboxHandle, SandboxKind, SandboxPrivateInput, + SandboxProvider, + SandboxProviderId, SandboxShellExecutionRequest, } from '@/lib/execution/remote-sandbox/types' @@ -48,14 +57,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 } : {}), + } } /** @@ -72,10 +108,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 @@ -155,10 +192,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) { @@ -200,6 +240,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 } +} + /** * Materializes sandbox input files before user code runs. `content` entries are written inline; * `url` entries are fetched from inside the sandbox via `curl` — their bytes never pass through the @@ -406,14 +459,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) { @@ -425,7 +478,7 @@ async function inspectSandboxOutputFileSize( logger.warn('Failed to inspect requested sandbox output file', { sandboxId: sandbox.sandboxId, }) - return undefined + throw error } } @@ -451,7 +504,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) @@ -547,7 +599,7 @@ async function executeInSandboxWithinBudget( }) throwIfAborted(signal) - const sandbox = await createSelectedSandbox( + const created = await createSelectedSandbox( kind, { language, @@ -555,10 +607,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) @@ -602,12 +658,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 @@ -636,19 +694,32 @@ async function executeInSandboxWithinBudget( } } - const { exportedFiles, exportedFileContent } = await collectExportedFiles(sandbox, req, { - signal, - }) - throwIfAborted(signal) - - return { + billableResult = { result: extraction.result, stdout: cleanedStdout, sandboxId, - exportedFileContent, - exportedFiles, } + try { + const { exportedFiles, exportedFileContent } = await collectExportedFiles(sandbox, req, { + signal, + }) + throwIfAborted(signal) + billableResult.exportedFileContent = exportedFileContent + billableResult.exportedFiles = exportedFiles + } catch (error) { + if (isSandboxOutputLimitError(error) || isSandboxOutputFileError(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() } @@ -677,14 +748,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) @@ -735,7 +810,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 @@ -744,19 +821,32 @@ async function executeShellInSandboxWithinBudget( const extraction = extractSimResult(stdout) const parsed = extraction.parseFailed ? extraction.rawPayload : extraction.result - const { exportedFiles, exportedFileContent } = await collectExportedFiles(sandbox, req, { - signal, - }) - throwIfAborted(signal) - - return { + billableResult = { result: parsed, stdout: extraction.cleanedStdout, sandboxId, - exportedFileContent, - exportedFiles, } + try { + const { exportedFiles, exportedFileContent } = await collectExportedFiles(sandbox, req, { + signal, + }) + throwIfAborted(signal) + billableResult.exportedFileContent = exportedFileContent + billableResult.exportedFiles = exportedFiles + } catch (error) { + if (isSandboxOutputLimitError(error) || isSandboxOutputFileError(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() } @@ -818,7 +908,7 @@ export async function withPiSandbox( ): Promise { const lifetimeMs = options.lifetimeMs !== undefined ? options.lifetimeMs : resolvePiSandboxLifetimeMs() - const sandbox = await createSandbox('pi', { lifetimeMs }) + const { sandbox } = await createSandbox('pi', { lifetimeMs }) logger.info('Started Pi sandbox', { sandboxId: sandbox.sandboxId, lifetimeMs }) const runner: PiSandboxRunner = { diff --git a/apps/sim/lib/execution/remote-sandbox/output-limits.ts b/apps/sim/lib/execution/remote-sandbox/output-limits.ts index 91fc5cb3616..3f5dbe9b1e9 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 /** @@ -53,6 +55,21 @@ export function appendStreamedSandboxOutput(current: string, chunk: string): str export const SANDBOX_OUTPUT_LIMIT_CODE = 'sandbox_output_limit_exceeded' as const export const SANDBOX_OUTPUT_FILE_INVALID_CODE = 'sandbox_output_file_invalid' as const +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/types.ts b/apps/sim/lib/execution/remote-sandbox/types.ts index c2f4e7b2b28..a1ee43a1bf8 100644 --- a/apps/sim/lib/execution/remote-sandbox/types.ts +++ b/apps/sim/lib/execution/remote-sandbox/types.ts @@ -53,6 +53,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 { @@ -76,6 +78,14 @@ 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 } export interface SandboxExecutionResult { @@ -85,6 +95,7 @@ export interface SandboxExecutionResult { error?: string exportedFileContent?: string exportedFiles?: Record + cost?: SandboxExecutionCost } /** Result of one command run inside a sandbox. */ @@ -94,6 +105,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' } /** @@ -117,6 +130,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 { @@ -199,6 +214,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 } /** @@ -286,5 +303,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 2fd1b35eaa5..a6fcbc3f0df 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)) { @@ -281,6 +286,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( @@ -291,6 +297,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' }], }, @@ -301,6 +309,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) }) @@ -320,6 +329,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"', @@ -375,6 +491,7 @@ describe('Function execution request', () => { expect.objectContaining({ language, sandboxKind: 'mothership', + meterUsage: false, }) ) expect(mockExecuteInIsolatedVM).not.toHaveBeenCalled() @@ -396,7 +513,7 @@ describe('Function execution request', () => { expect(response.status).toBe(200) expect(mockExecuteShellInSandbox).toHaveBeenCalledWith( - expect.objectContaining({ sandboxKind: 'mothership' }) + expect.objectContaining({ sandboxKind: 'mothership', meterUsage: false }) ) }) @@ -697,6 +814,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}', @@ -707,6 +825,8 @@ describe('Function execution request', () => { code: 'print("done")', language: 'python', workspaceId: 'workspace-1', + workflowId: 'workflow-1', + executionId: 'execution-1', outputs: { files: [ { @@ -753,6 +873,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' }), @@ -1331,9 +1452,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")', @@ -1354,6 +1476,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() }) @@ -1375,15 +1498,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}', @@ -1397,6 +1546,8 @@ describe('Function execution request', () => { code: 'print("done")', language: 'python', workspaceId: 'workspace-1', + workflowId: 'workflow-1', + executionId: 'execution-1', outputs: { files: [ { @@ -1419,6 +1570,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() }) @@ -1725,6 +1877,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 }, }) @@ -1733,6 +1886,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', @@ -1754,6 +1909,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 443a3fdc0e1..e1926d3d949 100644 --- a/apps/sim/lib/function-execution/execute-request.ts +++ b/apps/sim/lib/function-execution/execute-request.ts @@ -77,6 +77,7 @@ import { isSandboxOutputFileError, isSandboxOutputLimitError, MAX_SANDBOX_OUTPUT_BYTES, + readTrustedSandboxOutputCost, } from '@/lib/execution/remote-sandbox/output-limits' import { isExecutionResourceLimitError } from '@/lib/execution/resource-errors' import { @@ -112,6 +113,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 @@ -1402,10 +1409,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 } ) } @@ -1428,6 +1445,7 @@ async function maybeExportSandboxFileToWorkspace(args: { exportedFileContent?: string stdout: string executionTime: number + cost?: FunctionExecutionCost }) { const { routeContext, @@ -1443,6 +1461,7 @@ async function maybeExportSandboxFileToWorkspace(args: { exportedFileContent, stdout, executionTime, + cost, } = args if (!outputSandboxPath) return null @@ -1452,7 +1471,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 ) } @@ -1464,7 +1484,8 @@ async function maybeExportSandboxFileToWorkspace(args: { 'Workspace context required to save sandbox file to workspace', 400, stdout, - executionTime + executionTime, + cost ) } @@ -1473,7 +1494,8 @@ async function maybeExportSandboxFileToWorkspace(args: { `Sandbox file "${outputSandboxPath}" was not found or could not be read`, 500, stdout, - executionTime + executionTime, + cost ) } @@ -1491,7 +1513,8 @@ async function maybeExportSandboxFileToWorkspace(args: { `Sandbox output files exceed ${MAX_SANDBOX_OUTPUT_BYTES} bytes total`, 400, stdout, - executionTime + executionTime, + cost ) } const fileBuffer = isBinary @@ -1565,6 +1588,7 @@ async function maybeExportSandboxFileToWorkspace(args: { }, stdout: cleanStdout(stdout), executionTime, + ...(cost ? { cost } : {}), }, resources: [{ type: 'file', id: written.id, title: written.name, path: written.vfsPath }], }) @@ -1573,7 +1597,8 @@ async function maybeExportSandboxFileToWorkspace(args: { getErrorMessage(error, 'Failed to export sandbox file'), workspaceFileExportErrorStatus(error), stdout, - executionTime + executionTime, + cost ) } } @@ -1588,6 +1613,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 @@ -1596,7 +1622,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 ) } @@ -1617,6 +1644,7 @@ async function maybeExportSandboxFilesToWorkspace(args: { args.exportedFileContent, stdout: args.stdout, executionTime: args.executionTime, + cost: args.cost, }) } @@ -1628,7 +1656,8 @@ async function maybeExportSandboxFilesToWorkspace(args: { 'Workspace context required to save sandbox files to workspace', 400, args.stdout, - args.executionTime + args.executionTime, + args.cost ) } @@ -1642,7 +1671,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 @@ -1659,7 +1689,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') @@ -1708,7 +1739,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( @@ -1719,7 +1751,8 @@ async function maybeExportSandboxFilesToWorkspace(args: { `Duplicate sandbox output destination: ${duplicateDestination}`, 400, args.stdout, - args.executionTime + args.executionTime, + args.cost ) } @@ -1775,7 +1808,8 @@ async function maybeExportSandboxFilesToWorkspace(args: { getErrorMessage(error, 'Failed to export sandbox files'), workspaceFileExportErrorStatus(error), args.stdout, - args.executionTime + args.executionTime, + args.cost ) } @@ -1814,6 +1848,7 @@ async function maybeExportSandboxFilesToWorkspace(args: { }, stdout: cleanStdout(args.stdout), executionTime: args.executionTime, + ...(args.cost ? { cost: args.cost } : {}), }, resources: writtenFiles.map((file) => ({ type: 'file', @@ -1927,6 +1962,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' }, @@ -1985,6 +2022,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 + ) + } const executionParams = { ...params } executionParams._context = undefined @@ -2133,6 +2187,7 @@ export async function executeFunctionRequest( error: shellError, exportedFileContent, exportedFiles, + cost: shellCost, } = await executeShellInSandbox({ code: resolvedCode, envs: shellEnvs, @@ -2147,6 +2202,7 @@ export async function executeFunctionRequest( ? { sandboxKind: 'mothership' as const } : {}), signal: executionSignal, + meterUsage: meterRemoteSandboxUsage, }) const executionTime = Date.now() - execStart @@ -2161,7 +2217,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 } @@ -2179,6 +2240,7 @@ export async function executeFunctionRequest( exportedFileContent, stdout: shellStdout, executionTime, + cost: shellCost, }) if (fileExportResponse) { return appendResolvedSecretNames(fileExportResponse, routeContext) @@ -2188,7 +2250,12 @@ export async function executeFunctionRequest( return functionJsonResponse( { success: true, - output: { result: shellResult ?? null, stdout: cleanStdout(shellStdout), executionTime }, + output: { + result: shellResult ?? null, + stdout: cleanStdout(shellStdout), + executionTime, + ...(shellCost ? { cost: shellCost } : {}), + }, }, routeContext ) @@ -2300,6 +2367,7 @@ export async function executeFunctionRequest( error: e2bError, exportedFileContent, exportedFiles, + cost: sandboxCost, } = await executeInSandbox({ code: codeForE2B, language: CodeLanguage.JavaScript, @@ -2315,6 +2383,7 @@ export async function executeFunctionRequest( ? { sandboxKind: 'mothership' as const } : {}), signal: executionSignal, + meterUsage: meterRemoteSandboxUsage, }) const executionTime = Date.now() - execStart stdout += e2bStdout @@ -2340,7 +2409,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 } @@ -2358,6 +2432,7 @@ export async function executeFunctionRequest( exportedFileContent, stdout, executionTime, + cost: sandboxCost, }) if (fileExportResponse) { return appendResolvedSecretNames(fileExportResponse, routeContext) @@ -2367,7 +2442,12 @@ export async function executeFunctionRequest( return functionJsonResponse( { success: true, - output: { result: e2bResult ?? null, stdout: cleanStdout(stdout), executionTime }, + output: { + result: e2bResult ?? null, + stdout: cleanStdout(stdout), + executionTime, + ...(sandboxCost ? { cost: sandboxCost } : {}), + }, }, routeContext ) @@ -2391,6 +2471,7 @@ export async function executeFunctionRequest( error: e2bError, exportedFileContent, exportedFiles, + cost: sandboxCost, } = await executeInSandbox({ code: codeForE2B, language: CodeLanguage.Python, @@ -2405,6 +2486,7 @@ export async function executeFunctionRequest( ? { sandboxKind: 'mothership' as const } : {}), signal: executionSignal, + meterUsage: meterRemoteSandboxUsage, }) const executionTime = Date.now() - execStart stdout += e2bStdout @@ -2430,7 +2512,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 } @@ -2448,6 +2535,7 @@ export async function executeFunctionRequest( exportedFileContent, stdout, executionTime, + cost: sandboxCost, }) if (fileExportResponse) { return appendResolvedSecretNames(fileExportResponse, routeContext) @@ -2457,7 +2545,12 @@ export async function executeFunctionRequest( return functionJsonResponse( { success: true, - output: { result: e2bResult ?? null, stdout: cleanStdout(stdout), executionTime }, + output: { + result: e2bResult ?? null, + stdout: cleanStdout(stdout), + executionTime, + ...(sandboxCost ? { cost: sandboxCost } : {}), + }, }, routeContext ) @@ -2637,10 +2730,16 @@ export async function executeFunctionRequest( ) } if (isSandboxOutputLimitError(error) || isSandboxOutputFileError(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 de9b27a3121..bd42772505f 100644 --- a/apps/sim/tools/function/execute.ts +++ b/apps/sim/tools/function/execute.ts @@ -201,6 +201,7 @@ export const functionExecuteTool: InternalToolConfig { 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 30acd56bcd3..f872d4a2193 100644 --- a/apps/sim/tools/index.ts +++ b/apps/sim/tools/index.ts @@ -1565,11 +1565,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 @@ -1581,7 +1582,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') } @@ -2219,9 +2219,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 @@ -2382,6 +2387,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) } @@ -2406,12 +2438,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