diff --git a/apps/sim/executor/execution/block-executor.ts b/apps/sim/executor/execution/block-executor.ts index 1259f1507b5..a859f3004f1 100644 --- a/apps/sim/executor/execution/block-executor.ts +++ b/apps/sim/executor/execution/block-executor.ts @@ -3,7 +3,12 @@ import { describeError, toError } from '@sim/utils/errors' import { sleep } from '@sim/utils/helpers' import { isRecordLike, toRecord } from '@sim/utils/object' import { findDatabaseQueryError } from '@/lib/core/errors/database-query-error' -import { logFailureOnce, markFailureKind, markFailureLogged } from '@/lib/core/errors/failure-log' +import { + classifyFailure, + logFailureOnce, + markFailureKind, + markFailureLogged, +} from '@/lib/core/errors/failure-log' import { isTimeoutAbortReason } from '@/lib/core/execution-limits/types' import { redactApiKeys } from '@/lib/core/security/redaction' import { normalizeStringArray } from '@/lib/core/utils/arrays' @@ -1327,10 +1332,13 @@ export class BlockExecutor { try { pumpResult = await pump.run() } catch (error) { - this.execLogger.error('Error reading stream for block', { - blockId, - ...projectStreamDiagnosticError(error), - }) + this.execLogger[classifyFailure(error) === 'user' ? 'info' : 'error']( + 'Error reading stream for block', + { + blockId, + ...projectStreamDiagnosticError(error), + } + ) if (onStreamPromise) { await onStreamPromise.catch(() => {}) } diff --git a/apps/sim/executor/execution/executor.ts b/apps/sim/executor/execution/executor.ts index 93c12e6a96b..d381c17c449 100644 --- a/apps/sim/executor/execution/executor.ts +++ b/apps/sim/executor/execution/executor.ts @@ -485,6 +485,7 @@ export class DAGExecutor { executionOrderCounter: { value: restoredBlockLogs.reduce((max, log) => Math.max(max, log.executionOrder ?? 0), 0), }, + unrecordedModelCost: this.contextExtensions.unrecordedModelCost ?? { total: 0 }, metadata: { ...this.contextExtensions.metadata, ...(this.contextExtensions.billingAttribution diff --git a/apps/sim/executor/execution/types.ts b/apps/sim/executor/execution/types.ts index c5aa218d906..905dc3ac6ec 100644 --- a/apps/sim/executor/execution/types.ts +++ b/apps/sim/executor/execution/types.ts @@ -11,6 +11,7 @@ import type { NormalizedBlockOutput, StartBlockRunMetadata, StreamingExecution, + UnrecordedModelCost, } from '@/executor/types' import type { ResolvedSecretTraceProvenanceV1, @@ -329,6 +330,8 @@ export interface ContextExtensions { * Takes precedence over `metadata.billingAttribution` when both are set. */ billingAttribution?: BillingAttributionSnapshot + /** The invoking run's {@link ExecutionContext.unrecordedModelCost}, for a child billed on its log. */ + unrecordedModelCost?: UnrecordedModelCost stream?: boolean selectedOutputs?: string[] edges?: Array<{ source: string; target: string }> diff --git a/apps/sim/executor/handlers/workflow/custom-block-tool-runner.ts b/apps/sim/executor/handlers/workflow/custom-block-tool-runner.ts index 99e6c6aa339..407b64e5f5d 100644 --- a/apps/sim/executor/handlers/workflow/custom-block-tool-runner.ts +++ b/apps/sim/executor/handlers/workflow/custom-block-tool-runner.ts @@ -6,7 +6,11 @@ import { isPlainRecord, isRecordLike } from '@sim/utils/object' import type { BillingAttributionSnapshot } from '@/lib/billing/core/attribution' import type { PiiBlockOutputRedaction } from '@/executor/execution/types' import { WorkflowBlockHandler } from '@/executor/handlers/workflow/handler' -import type { ExecutionContext, ExecutorDelegationOrigin } from '@/executor/types' +import type { + ExecutionContext, + ExecutorDelegationOrigin, + UnrecordedModelCost, +} from '@/executor/types' import { projectResolvedSecretDiagnosticContent } from '@/executor/utils/resolved-secret-content-projection' import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' import type { SerializedBlock } from '@/serializer/types' @@ -76,6 +80,8 @@ export function buildCustomBlockExecutionContext( principal?: WorkflowExecutionPrincipal /** The invoking run's in-flight block-output redaction policy. */ piiBlockOutputRedaction?: PiiBlockOutputRedaction + /** The invoking run's unrecorded model spend, which a child billed on its log adds to. */ + unrecordedModelCost?: UnrecordedModelCost } ): ExecutionContext { // Prefer the invoking agent run's ids so correlation and cancellation both @@ -98,6 +104,7 @@ export function buildCustomBlockExecutionContext( resolvedSecretTraceRegistry: options.resolvedSecretTraceRegistry, environmentVariables: options.environmentVariables, piiBlockOutputRedaction: options.piiBlockOutputRedaction, + unrecordedModelCost: options.unrecordedModelCost, blockStates: new Map(), executedBlocks: new Set(), blockLogs: [], diff --git a/apps/sim/executor/handlers/workflow/handler.test.ts b/apps/sim/executor/handlers/workflow/handler.test.ts index f738a2533d7..140bad47f89 100644 --- a/apps/sim/executor/handlers/workflow/handler.test.ts +++ b/apps/sim/executor/handlers/workflow/handler.test.ts @@ -398,6 +398,7 @@ describe('WorkflowBlockHandler', () => { workspaceId: 'workspace-parent', environmentVariables: { MY_API_KEY: 'parent-secret' }, piiBlockOutputRedaction, + unrecordedModelCost: { total: 1 }, } as unknown as ExecutionContext mockFetch.mockResolvedValueOnce({ @@ -425,6 +426,7 @@ describe('WorkflowBlockHandler', () => { expect(executorOptions[0].contextExtensions.piiBlockOutputRedaction).toBe( piiBlockOutputRedaction ) + expect(executorOptions[0].contextExtensions.unrecordedModelCost).toBe(ctx.unrecordedModelCost) expect(mockGetPersonalAndWorkspaceEnv).not.toHaveBeenCalled() }) @@ -439,6 +441,7 @@ describe('WorkflowBlockHandler', () => { ...mockContext, workspaceId: 'workspace-consumer', metadata: { ...mockContext.metadata, billingAttribution: consumerAttribution }, + unrecordedModelCost: { total: 1 }, } as unknown as ExecutionContext mockGetCustomBlockAuthority.mockResolvedValue({ @@ -506,6 +509,7 @@ describe('WorkflowBlockHandler', () => { }) expect(executorOptions).toHaveLength(1) expect(executorOptions[0].contextExtensions.billingAttribution).toBe(sourceAttribution) + expect(executorOptions[0].contextExtensions.unrecordedModelCost).toBeUndefined() expect(executorOptions[0].contextExtensions.userId).toBe('owner-9') expect(executorOptions[0].contextExtensions.workspaceId).toBe('workspace-source') expect(executorOptions[0].contextExtensions.executorDelegationOrigin).toEqual({ diff --git a/apps/sim/executor/handlers/workflow/handler.ts b/apps/sim/executor/handlers/workflow/handler.ts index b85f69f7864..bb357961de4 100644 --- a/apps/sim/executor/handlers/workflow/handler.ts +++ b/apps/sim/executor/handlers/workflow/handler.ts @@ -886,6 +886,9 @@ export class WorkflowBlockHandler implements BlockHandler { // internal tool calls (knowledge, guardrails, MCP, Mothership) can // attach the required billing attribution header. billingAttribution: childBillingAttribution, + // A custom block's child writes its own log and ledger rows; any other child's cost + // lands on this run's log, so it spends against this run's unrecorded total. + unrecordedModelCost: childExecutionId ? undefined : ctx.unrecordedModelCost, resolvedSecretTraceRegistry: childResolvedSecretTraceRegistry, // Fall back to the inherited metadata so a toggle-off intermediate // child still carries the trusted identity chain to deeper children. diff --git a/apps/sim/executor/handlers/workflow/tool-runner.test.ts b/apps/sim/executor/handlers/workflow/tool-runner.test.ts index 534743ec43b..623f067d449 100644 --- a/apps/sim/executor/handlers/workflow/tool-runner.test.ts +++ b/apps/sim/executor/handlers/workflow/tool-runner.test.ts @@ -70,6 +70,18 @@ describe('runWorkflowTool execution context', () => { expect(ctxArg.principal).toEqual(trustedPrincipal) expect(ctxArg.metadata.principal).toEqual(trustedPrincipal) }) + + it("hands the child the invoking run's unrecorded model total itself, not a copy", async () => { + const unrecordedModelCost = { total: 1 } + + await runWorkflowTool( + { workflowId: 'wf-child', _context: { workspaceId: 'ws-1' } }, + { environmentVariables: {}, unrecordedModelCost } + ) + + const [ctxArg] = mockExecute.mock.calls[0] + expect(ctxArg.unrecordedModelCost).toBe(unrecordedModelCost) + }) }) const CHILD_SECRET = 'sk-live-7f2c9a41d8be4055b6c3' diff --git a/apps/sim/executor/handlers/workflow/tool-runner.ts b/apps/sim/executor/handlers/workflow/tool-runner.ts index 6ce87ef86d1..c506fc97f92 100644 --- a/apps/sim/executor/handlers/workflow/tool-runner.ts +++ b/apps/sim/executor/handlers/workflow/tool-runner.ts @@ -12,7 +12,7 @@ import { type CustomBlockExecutorContext, } from '@/executor/handlers/workflow/custom-block-tool-runner' import { WorkflowBlockHandler } from '@/executor/handlers/workflow/handler' -import type { ExecutorDelegationOrigin } from '@/executor/types' +import type { ExecutorDelegationOrigin, UnrecordedModelCost } from '@/executor/types' import { classifyExecutionError } from '@/executor/utils/errors' import { parseJSON } from '@/executor/utils/json' import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' @@ -99,6 +99,7 @@ export async function runWorkflowTool( executorDelegationOrigin?: ExecutorDelegationOrigin principal?: WorkflowExecutionPrincipal piiBlockOutputRedaction?: PiiBlockOutputRedaction + unrecordedModelCost?: UnrecordedModelCost } ): Promise { if (!params.workflowId) { diff --git a/apps/sim/executor/types.ts b/apps/sim/executor/types.ts index fc45af27396..50e4b108773 100644 --- a/apps/sim/executor/types.ts +++ b/apps/sim/executor/types.ts @@ -668,6 +668,20 @@ export interface ExecutionContext { * Starts at 0 and increments for each block. Use getNextExecutionOrder() to access. */ executionOrderCounter?: { value: number } + + /** + * Billed model spend not yet in the usage ledger, which receives a run's cost only with the + * run's log. Shared by every block's context copy and by child workflows billed on this log. + */ + unrecordedModelCost?: UnrecordedModelCost +} + +/** A run's running model spend, and when its usage limit may next be read after a read failed. */ +export interface UnrecordedModelCost { + total: number + readDeferredUntil?: number + /** Agent turns whose journaled spend `total` already holds. */ + countedTurnIds?: Set } /** diff --git a/apps/sim/lib/billing/core/mid-run-usage.ts b/apps/sim/lib/billing/core/mid-run-usage.ts index cae9c27e6e0..f13d96ab97d 100644 --- a/apps/sim/lib/billing/core/mid-run-usage.ts +++ b/apps/sim/lib/billing/core/mid-run-usage.ts @@ -89,8 +89,30 @@ function periodHasEnded(attribution: BillingAttributionSnapshot): boolean { return Date.now() >= new Date(attribution.billingPeriod.end).getTime() } +/** + * The limit a run's not-yet-recorded spend takes it to, if any: the payer's pool, then the acting + * member's cap. + */ +function limitReachedWith( + usage: AttributedUsageLimitsResult, + unrecordedCost: number +): 'payer' | 'member' | undefined { + // Without unrecorded spend the gate's own verdict stands, so callers that pass none are unchanged. + if (unrecordedCost <= 0) return undefined + const { payerUsage, memberUsage } = usage + if (payerUsage && payerUsage.currentUsage + unrecordedCost >= payerUsage.limit) return 'payer' + if ( + memberUsage?.limit != null && + memberUsage.currentUsage + unrecordedCost >= memberUsage.limit + ) { + return 'member' + } + return undefined +} + async function readGateVerdict( - attribution: BillingAttributionSnapshot + attribution: BillingAttributionSnapshot, + unrecordedCost: number ): Promise { let usage: AttributedUsageLimitsResult try { @@ -101,7 +123,10 @@ async function readGateVerdict( }) return { status: 'unknown' } } - if (!usage.isExceeded) return { status: 'within' } + if (!usage.isExceeded) { + const scope = limitReachedWith(usage, unrecordedCost) + return scope ? { status: 'exceeded', scope } : { status: 'within' } + } if (usage.reason === 'billing_blocked') { return { status: 'blocked', ...(usage.message ? { message: usage.message } : {}) } } @@ -120,12 +145,17 @@ async function readGateVerdict( * next callback judges the next period. Any other payer's charges stay in the admitted * period (a reporting window, or the open default one), so that period is judged, even after it * ends. + * + * @param unrecordedCost - Spend the ledger does not hold yet, counted as recorded usage. */ export async function readMidRunUsageVerdict( - attribution: BillingAttributionSnapshot + attribution: BillingAttributionSnapshot, + unrecordedCost = 0 ): Promise { if (!isHosted || !isBillingEnabled) return { status: 'within' } - if (!rollsIntoCurrentPeriod(attribution.billingPeriod)) return readGateVerdict(attribution) + if (!rollsIntoCurrentPeriod(attribution.billingPeriod)) { + return readGateVerdict(attribution, unrecordedCost) + } let judged: BillingAttributionSnapshot try { judged = await currentAttribution(attribution) @@ -136,7 +166,7 @@ export async function readMidRunUsageVerdict( return { status: 'unknown' } } if (periodHasEnded(judged)) return { status: 'unknown' } - const verdict = await readGateVerdict(judged) + const verdict = await readGateVerdict(judged, unrecordedCost) // A period that ended during the read is judged at the next callback, which reloads it. return periodHasEnded(judged) ? { status: 'unknown' } : verdict } diff --git a/apps/sim/lib/memory/conversation-types.ts b/apps/sim/lib/memory/conversation-types.ts index 7da67f624a3..c4344dff68d 100644 --- a/apps/sim/lib/memory/conversation-types.ts +++ b/apps/sim/lib/memory/conversation-types.ts @@ -91,6 +91,7 @@ export interface AgentConversationSession { getFinalResponse(): AgentTurnState['final'] getFinalAssistantContent(): string | undefined readonly memoryId?: string + readonly turnId?: string getUsage(): ConversationUsageTotal recordContextUsage?(usage: ConversationUsageTotal): Promise captureStep(step: CapturedConversationStep): Promise diff --git a/apps/sim/providers/anthropic/streaming-tool-loop.ts b/apps/sim/providers/anthropic/streaming-tool-loop.ts index dc23fecd268..9c0aa3dcac6 100644 --- a/apps/sim/providers/anthropic/streaming-tool-loop.ts +++ b/apps/sim/providers/anthropic/streaming-tool-loop.ts @@ -585,8 +585,8 @@ export function createAnthropicStreamingToolLoopStream( aborted: loopAbortController.signal.aborted, consumerCancelled, error, - onUnexpectedError: (cause) => - logger.error('Anthropic streaming tool loop failed', { + onUnexpectedError: (cause, level) => + logger[level]('Anthropic streaming tool loop failed', { error: toError(cause).message, }), }) diff --git a/apps/sim/providers/bedrock/streaming-tool-loop.ts b/apps/sim/providers/bedrock/streaming-tool-loop.ts index d323aae8dd2..ea1e36b96fc 100644 --- a/apps/sim/providers/bedrock/streaming-tool-loop.ts +++ b/apps/sim/providers/bedrock/streaming-tool-loop.ts @@ -698,8 +698,8 @@ export function createBedrockStreamingToolLoopStream( aborted: loopAbortController.signal.aborted, consumerCancelled, error, - onUnexpectedError: (cause) => - logger.error('Bedrock streaming tool loop failed', { + onUnexpectedError: (cause, level) => + logger[level]('Bedrock streaming tool loop failed', { error: toError(cause).message, }), }) diff --git a/apps/sim/providers/conversation-generation.test.ts b/apps/sim/providers/conversation-generation.test.ts index 0407f1180b6..7f90a74eaa5 100644 --- a/apps/sim/providers/conversation-generation.test.ts +++ b/apps/sim/providers/conversation-generation.test.ts @@ -1,3 +1,8 @@ +import { resetEnvFlagsMock, setEnvFlags } from '@sim/testing' +import { + billingUsageGateCacheMock, + billingUsageGateCacheMockFns, +} from '@sim/testing/mocks/billing-usage-gate-cache.mock' import { providersConversationHistoryMock, providersConversationHistoryMockFns, @@ -29,6 +34,7 @@ const state = vi.hoisted(() => ({ enabled: true, historyTokens: 0 })) vi.mock('@/providers/conversation-history', () => providersConversationHistoryMock) vi.mock('@/providers/models', () => providersModelsMock) vi.mock('@/lib/tokenization/accurate', () => tokenizationAccurateMock) +vi.mock('@/lib/billing/core/usage-gate-cache', () => billingUsageGateCacheMock) providersModelsMockFns.mockGetMaxOutputTokensForModel.mockReturnValue(100) Object.assign(providersModelsMock.PROVIDER_DEFINITIONS, { @@ -611,3 +617,53 @@ describe('provider generation context boundary', () => { expect(payload.messages).toHaveLength(1) }) }) + +describe('usage limit across compaction', () => { + it('stops the generation when its compaction spends the run past the usage limit', async () => { + setEnvFlags({ isHosted: true, isBillingEnabled: true }) + billingUsageGateCacheMockFns.mockCheckExecutionUsageLimits.mockResolvedValue({ + isExceeded: false, + payerUsage: { currentUsage: 3, limit: 5 }, + }) + const unrecordedModelCost = { total: 1 } + providersConversationHistoryMockFns.mockGetConversationRequestContext.mockReturnValueOnce({ + agentConversation: {}, + agentMemoryContext: { historyTokens: 200 }, + executionContext: { + metadata: { + duration: 0, + billingAttribution: { + actorUserId: 'actor', + workspaceId: 'workspace', + organizationId: null, + billedAccountUserId: 'actor', + billingEntity: { type: 'user', id: 'actor' }, + billingPeriod: { start: '2026-10-01T00:00:00.000Z', end: '2099-01-01T00:00:00.000Z' }, + payerSubscription: null, + }, + }, + unrecordedModelCost, + }, + }) + const fixture = fixtures[0] + const input = request() + bindConversationGenerationCompactor(input, async () => { + unrecordedModelCost.total += 1.5 + return { role: 'user' as const, content: 'Derived prior context' } + }) + const messages = [ + fixture.prompt, + ...fixture.batch, + nativeText(fixture, 'old'.repeat(250)), + nativeText(fixture, 'latest receipt'), + ] + + try { + await expect( + prepareConversationGeneration(input, fixture.protocol, { [fixture.key]: messages }) + ).rejects.toThrow('usage limit') + } finally { + resetEnvFlagsMock() + } + }) +}) diff --git a/apps/sim/providers/conversation-generation.ts b/apps/sim/providers/conversation-generation.ts index 6aa5f0ac437..ac56db20abe 100644 --- a/apps/sim/providers/conversation-generation.ts +++ b/apps/sim/providers/conversation-generation.ts @@ -21,6 +21,7 @@ import { isConversationHistoryNotice, } from '@/providers/conversation-metadata' import { getConversationModelLimits } from '@/providers/conversation-model' +import { assertRunWithinUsageLimit, RunUsageLimitError } from '@/providers/run-usage-limit' import { isAbortError } from '@/providers/streaming-tool-loop-shared' import type { Message, ProviderRequest } from '@/providers/types' @@ -114,9 +115,16 @@ function contextError(message: string): Error { return new AgentContextProtocolError(message) } -/** Provider wrappers must preserve permanent context failures so fallback cannot restart them. */ +/** + * Provider wrappers must preserve the permanent failures this module raises, context faults and + * the run's usage-limit stop, so fallback cannot restart them and the stop keeps its billed cost. + */ export function isConversationContextError(error: unknown): boolean { - return error instanceof AgentContextLimitError || error instanceof AgentContextProtocolError + return ( + error instanceof AgentContextLimitError || + error instanceof AgentContextProtocolError || + error instanceof RunUsageLimitError + ) } function records(value: unknown): Record[] { @@ -283,6 +291,7 @@ export async function prepareConversationGeneration( payload: T ): Promise { const runtime = getConversationRequestContext(request) + await assertRunWithinUsageLimit(runtime) if (!runtime?.agentConversation) return payload request.abortSignal?.throwIfAborted() if (!isRecordLike(payload)) @@ -401,6 +410,8 @@ export async function prepareConversationGeneration( if (isAbortError(error) || request.abortSignal?.aborted) throw error logger.warn('Agent context compaction unavailable; retaining bounded history') } + // Summary calls spend too, and the compactor keeps going when one is stopped. + await assertRunWithinUsageLimit(runtime) } let selectedLength = 0 for (const group of selected) { diff --git a/apps/sim/providers/conversation-history.ts b/apps/sim/providers/conversation-history.ts index 6fbba36e33c..5812a04ba98 100644 --- a/apps/sim/providers/conversation-history.ts +++ b/apps/sim/providers/conversation-history.ts @@ -7,6 +7,7 @@ import type { ConversationProtocol, ConversationUsage } from '@/lib/memory/conve import { getConversationPrefixHash } from '@/providers/conversation-prefix' import { priceModelUsage, resolveModelCostPolicy } from '@/providers/cost-policy' import { providerHistoryAdapters } from '@/providers/history-adapters' +import { recordModelStepCost } from '@/providers/run-usage-limit' import { getProviderRuntimeContext, type ProviderRuntimeContext } from '@/providers/runtime-context' import type { ProviderId, ProviderRequest, ProviderToolConfig } from '@/providers/types' @@ -89,7 +90,18 @@ export function getConversationBinding(providerId: ProviderId, request: Provider .digest('hex') } -/** Called with a complete provider message, including streaming terminal responses. */ +function priceRequestUsage(request: ProviderRequest, usage: ConversationUsage) { + return priceModelUsage( + request.model, + usage, + resolveModelCostPolicy(request.model, request.isBYOK) + ) +} + +/** + * Called with a complete provider message, including streaming terminal responses. Its billed + * cost counts toward the run's usage limit whether or not the Agent keeps a conversation. + */ export async function captureProviderConversationStep( request: ProviderRequest, protocol: ConversationProtocol, @@ -102,6 +114,8 @@ export async function captureProviderConversationStep( } ): Promise { const runtime = getConversationRequestContext(request) + const priced = usage ? { usage, cost: priceRequestUsage(request, usage) } : undefined + if (priced) recordModelStepCost(runtime, priced.cost) if (!runtime?.agentConversation || !runtime.conversationProvider) return try { const captured = providerHistoryAdapters[protocol].capture(value) @@ -131,16 +145,7 @@ export async function captureProviderConversationStep( ? { prefixHash: getConversationPrefixHash(options.requestHistory) } : {}), }, - ...(usage - ? { - usage, - cost: priceModelUsage( - request.model, - usage, - resolveModelCostPolicy(request.model, request.isBYOK) - ), - } - : {}), + ...priced, }) } catch { logger.warn('Agent conversation capture unavailable') @@ -148,21 +153,21 @@ export async function captureProviderConversationStep( } /** - * Saves model usage outside captured exchanges, including discarded decisions at the tool limit. + * Saves model usage outside captured exchanges, including discarded decisions at the tool limit, + * and counts it toward the run's usage limit. * The current provider response already counts these tokens; only a later attempt adds this journal usage. */ export async function recordProviderConversationUsage( request: ProviderRequest, usage: ConversationUsage | undefined ): Promise { - const session = getConversationRequestContext(request)?.agentConversation - if (!usage || !session?.recordContextUsage) return + if (!usage) return + const runtime = getConversationRequestContext(request) + const cost = priceRequestUsage(request, usage) + recordModelStepCost(runtime, cost) + const session = runtime?.agentConversation + if (!session?.recordContextUsage) return try { - const cost = priceModelUsage( - request.model, - usage, - resolveModelCostPolicy(request.model, request.isBYOK) - ) await session.recordContextUsage({ tokens: usage, cost: { input: cost.input, output: cost.output, total: cost.total, toolCost: 0 }, diff --git a/apps/sim/providers/gemini/streaming-tool-loop.ts b/apps/sim/providers/gemini/streaming-tool-loop.ts index 223d0b93a8b..3e6b3dac549 100644 --- a/apps/sim/providers/gemini/streaming-tool-loop.ts +++ b/apps/sim/providers/gemini/streaming-tool-loop.ts @@ -683,8 +683,8 @@ export function createGeminiStreamingToolLoopStream( aborted: loopAbortController.signal.aborted, consumerCancelled, error, - onUnexpectedError: (cause) => - logger.error('Gemini streaming tool loop failed', { + onUnexpectedError: (cause, level) => + logger[level]('Gemini streaming tool loop failed', { error: toError(cause).message, }), }) diff --git a/apps/sim/providers/index.ts b/apps/sim/providers/index.ts index b380b327300..a5dc9e2e4d8 100644 --- a/apps/sim/providers/index.ts +++ b/apps/sim/providers/index.ts @@ -37,6 +37,7 @@ import { } from '@/providers/file-attachments.server' import { isEvaluationModel, isKnownModelId, supportsTemperature } from '@/providers/models' import { getProviderExecutor } from '@/providers/registry' +import { addSettledCost, recordJournaledModelCost } from '@/providers/run-usage-limit' import { type ProviderRuntimeContext, runWithProviderRuntimeContext, @@ -294,6 +295,8 @@ export async function executeProviderRequest( const modelSafeRequest = await prepareProviderFileAttachments(sanitizedRequest) const toolIdentities = assignProviderToolIdentities(modelSafeRequest.tools) const failedFunctionToolCost = { total: 0 } + // A conversation summary charges its parent request, so a stop after compaction bills it. + const settledCost = runtimeContext?.settledCost ?? { input: 0, output: 0, total: 0 } const requestRuntimeContext: ProviderRuntimeContext = { ...runtimeContext, ...(runtimeContext?.agentConversation @@ -305,6 +308,7 @@ export async function executeProviderRequest( } : {}), failedFunctionToolCost, + settledCost, ...(toolIdentities.toolIdByWireId.size > 0 ? { toolIdByWireId: new Map([ @@ -343,12 +347,19 @@ export async function executeProviderRequest( for (const registry of replayRegistries) { if (registry) await session.restoreProvenance?.(registry) } + // A continued call's cost is journaled, so the usage read below already counts it. + const settledBeforeContinuation = settledCost.total await continuePendingConversationCalls(modelSafeRequest, session) + settledCost.total = settledBeforeContinuation + failedFunctionToolCost.total = 0 const currentUserMessage = [...(modelSafeRequest.messages ?? [])] .reverse() .find((message) => message.role === 'user' && !isConversationHistoryNotice(message)) bindConversationGenerationPrompt(modelSafeRequest, currentUserMessage) priorConversationUsage = session.getUsage() + // A stop bills what success would: this journaled usage plus the request's own steps. + addSettledCost(settledCost, priorConversationUsage.cost) + recordJournaledModelCost(requestRuntimeContext, session.turnId, priorConversationUsage) bindConversationGenerationCompactor( modelSafeRequest, createAgentConversationCompactor( @@ -359,6 +370,7 @@ export async function executeProviderRequest( const summary = await executeProviderRequest(providerId, summaryRequest, { resolvedSecretTraceRegistry: requestRuntimeContext.resolvedSecretTraceRegistry, executionContext: requestRuntimeContext.executionContext, + settledCost, }) if (isStreamingExecution(summary) || isReadableStream(summary)) throw new Error('Conversation summary did not return a settled response') diff --git a/apps/sim/providers/openai-compat/streaming-tool-loop.ts b/apps/sim/providers/openai-compat/streaming-tool-loop.ts index a847c0dc246..232a50c7ceb 100644 --- a/apps/sim/providers/openai-compat/streaming-tool-loop.ts +++ b/apps/sim/providers/openai-compat/streaming-tool-loop.ts @@ -632,8 +632,8 @@ export function createOpenAICompatStreamingToolLoopStream( aborted: loopAbortController.signal.aborted, consumerCancelled, error, - onUnexpectedError: (cause) => - logger.error(`${providerName} streaming tool loop failed`, { + onUnexpectedError: (cause, level) => + logger[level](`${providerName} streaming tool loop failed`, { error: toError(cause).message, }), }) diff --git a/apps/sim/providers/openai/streaming-tool-loop.ts b/apps/sim/providers/openai/streaming-tool-loop.ts index 2bd1a602bbb..b96d7379101 100644 --- a/apps/sim/providers/openai/streaming-tool-loop.ts +++ b/apps/sim/providers/openai/streaming-tool-loop.ts @@ -611,8 +611,8 @@ export function createOpenAIResponsesStreamingToolLoopStream( aborted: loopAbortController.signal.aborted, consumerCancelled, error, - onUnexpectedError: (cause) => - logger.error(`Error in ${providerLabel} streaming tool loop`, { + onUnexpectedError: (cause, level) => + logger[level](`Error in ${providerLabel} streaming tool loop`, { providerId, error: cause, }), diff --git a/apps/sim/providers/run-usage-limit.test.ts b/apps/sim/providers/run-usage-limit.test.ts new file mode 100644 index 00000000000..3fd7e5f2a29 --- /dev/null +++ b/apps/sim/providers/run-usage-limit.test.ts @@ -0,0 +1,565 @@ +/** + * @vitest-environment node + */ +import { resetEnvFlagsMock, setEnvFlags } from '@sim/testing' +import { billingCoreMock } from '@sim/testing/mocks/billing-core.mock' +import { + billingUsageGateCacheMock, + billingUsageGateCacheMockFns, +} from '@sim/testing/mocks/billing-usage-gate-cache.mock' +import { toolsMock, toolsMockFns } from '@sim/testing/mocks/tools.mock' +import { getErrorMessage } from '@sim/utils/errors' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +const { mockExecuteRequest } = vi.hoisted(() => ({ mockExecuteRequest: vi.fn() })) + +vi.mock('@/providers/registry', () => ({ + getProviderExecutor: async () => ({ executeRequest: mockExecuteRequest }), +})) +vi.mock('@/lib/billing/core/usage-gate-cache', () => billingUsageGateCacheMock) +vi.mock('@/tools', () => toolsMock) +vi.mock('@/lib/billing/core/billing', () => billingCoreMock) +vi.mock('@/providers/cost-policy', async (importOriginal) => ({ + ...(await importOriginal()), + resolveModelCostPolicy: (_model: string, isBYOK = false) => ({ + billable: !isBYOK, + multiplier: 1, + }), + priceModelUsage: ( + _model: string, + usage: { input: number; output: number }, + policy: { billable: boolean } + ) => + policy.billable + ? { input: usage.input, output: usage.output, total: usage.input + usage.output } + : { input: 0, output: 0, total: 0 }, +})) + +import type { BillingAttributionSnapshot } from '@/lib/billing/core/attribution' +import { resetMidRunUsageCaches } from '@/lib/billing/core/mid-run-usage' +import { NEAR_LIMIT_USAGE_GATE_TTL_MS } from '@/lib/billing/core/usage-gate-cache' +import { markFailureKind } from '@/lib/core/errors/failure-log' +import { isNonRetryableExecutionError } from '@/lib/execution/non-retryable-error' +import { AgentTurnSession } from '@/lib/memory/agent-turn-session' +import { AgentTurnStateMachine } from '@/lib/memory/turn-state' +import type { ExecutionContext } from '@/executor/types' +import { readTrustedExecutionCost } from '@/executor/utils/errors' +import { executeProviderRequest } from '@/providers' +import { + isConversationContextError, + prepareConversationGeneration, +} from '@/providers/conversation-generation' +import { + bindConversationRequestContext, + captureProviderConversationStep, + getConfiguredConversationToolBinding, + recordProviderConversationUsage, +} from '@/providers/conversation-history' +import { executeProviderTool } from '@/providers/runtime-context' +import { providerFailureLogLevel } from '@/providers/streaming-tool-loop-shared' +import { ProviderError, type ProviderRequest, type ProviderToolConfig } from '@/providers/types' + +const { mockCheckExecutionUsageLimits } = billingUsageGateCacheMockFns +const { mockExecuteTool } = toolsMockFns + +const attribution: BillingAttributionSnapshot = { + actorUserId: 'actor', + workspaceId: 'workspace', + organizationId: null, + billedAccountUserId: 'actor', + billingEntity: { type: 'user', id: 'actor' }, + billingPeriod: { start: '2026-10-01T00:00:00.000Z', end: '2099-01-01T00:00:00.000Z' }, + payerSubscription: null, +} + +function startRun(): ExecutionContext { + return { + workflowId: 'workflow', + blockStates: new Map(), + executedBlocks: new Set(), + blockLogs: [], + metadata: { duration: 0, billingAttribution: attribution }, + environmentVariables: {}, + decisions: { router: new Map(), condition: new Map() }, + completedLoops: new Set(), + activeExecutionPath: new Set(), + unrecordedModelCost: { total: 0 }, + } +} + +/** One provider request inside the run, bound the way `executeProviderRequest` binds it. */ +function startRequest(executionContext: ExecutionContext, isBYOK = false): ProviderRequest { + const request: ProviderRequest = { model: 'model-a', isBYOK } + bindConversationRequestContext(request, { + executionContext, + settledCost: { input: 0, output: 0, total: 0 }, + }) + return request +} + +async function modelStep(request: ProviderRequest, input: number, output: number) { + await prepareConversationGeneration(request, 'anthropic', { messages: [] }) + await captureProviderConversationStep(request, 'anthropic', { content: [] }, { input, output }) +} + +beforeEach(() => { + setEnvFlags({ isHosted: true, isBillingEnabled: true }) + mockCheckExecutionUsageLimits.mockReset() + mockCheckExecutionUsageLimits.mockResolvedValue({ + isExceeded: false, + payerUsage: { currentUsage: 3, limit: 5 }, + }) + resetMidRunUsageCaches() +}) +afterEach(resetEnvFlagsMock) + +describe('run usage limit', () => { + it('stops before the model call that would spend past the payer limit, billing the steps it ran', async () => { + const run = startRun() + const request = startRequest(run) + await modelStep(request, 0.5, 0.5) + + await expect(modelStep(request, 0.5, 0.5)).resolves.toBeUndefined() + + const stop = prepareConversationGeneration(request, 'anthropic', { messages: [] }) + await expect(stop).rejects.toThrow('usage limit') + const error = await stop.catch((caught: unknown) => caught) + expect(isNonRetryableExecutionError(error)).toBe(true) + expect(readTrustedExecutionCost(error)).toEqual({ input: 1, output: 1, total: 2 }) + expect(providerFailureLogLevel(error, undefined)).toBe('info') + }) + + it('keeps the stop, unretried and billed, through the passthrough every provider wraps failures with', async () => { + mockExecuteRequest.mockImplementation(async (request: ProviderRequest) => { + try { + await modelStep(request, 1, 1) + await modelStep(request, 1, 1) + return { content: '', model: request.model } + } catch (error) { + if (isConversationContextError(error)) throw error + throw new ProviderError(getErrorMessage(error), { + startTime: '', + endTime: '', + duration: 0, + }) + } + }) + + const failure = await executeProviderRequest( + 'anthropic', + { model: 'model-a' }, + { executionContext: startRun() } + ).catch((caught: unknown) => caught) + + expect(isNonRetryableExecutionError(failure)).toBe(true) + expect(readTrustedExecutionCost(failure)).toEqual({ input: 1, output: 1, total: 2 }) + }) + + it("adds a nested request's steps to the settled cost its caller passes in", async () => { + const parentCost = { input: 0, output: 0, total: 0 } + mockExecuteRequest.mockImplementation(async (request: ProviderRequest) => { + await modelStep(request, 0.5, 0.25) + return { content: '', model: request.model } + }) + + await executeProviderRequest( + 'anthropic', + { model: 'model-a' }, + { executionContext: startRun(), settledCost: parentCost } + ) + + expect(parentCost).toEqual({ input: 0.5, output: 0.25, total: 0.75 }) + }) + + it('bills a stopped attempt for the journaled steps a successful one would have billed', async () => { + const session = new AgentTurnStateMachine( + { save: vi.fn() }, + { + version: 1, + steps: [ + { + id: 'failed-attempt-step', + assistant: { role: 'assistant', content: 'Partial answer' }, + calls: [], + results: [], + usage: { input: 1, output: 1 }, + cost: { input: 1, output: 1, total: 2 }, + }, + ], + } + ) + const run = startRun() + run.unrecordedModelCost = { total: 2 } + mockExecuteRequest.mockImplementation(async (request: ProviderRequest) => { + await prepareConversationGeneration(request, 'anthropic', { messages: [] }) + return { content: '', model: request.model } + }) + + const failure = await executeProviderRequest( + 'anthropic', + { model: 'model-a', messages: [{ role: 'user', content: 'Continue' }] }, + { executionContext: run, agentConversation: session } + ).catch((caught: unknown) => caught) + + expect(isNonRetryableExecutionError(failure)).toBe(true) + expect(readTrustedExecutionCost(failure)).toEqual({ input: 1, output: 1, total: 2 }) + }) + + it("counts a resumed turn's journaled spend toward the limit once, however often the run reopens it", async () => { + mockCheckExecutionUsageLimits.mockResolvedValue({ + isExceeded: false, + payerUsage: { currentUsage: 1, limit: 5 }, + }) + const reopenTurn = () => + new AgentTurnSession( + 'turn-1', + { save: vi.fn() }, + { + version: 1, + steps: [ + { + id: 'earlier-process-step', + assistant: { role: 'assistant', content: 'Partial answer' }, + calls: [], + results: [], + usage: { input: 1, output: 1 }, + cost: { input: 1, output: 1, total: 2 }, + }, + ], + } + ) + const run = startRun() + mockExecuteRequest.mockImplementation(async (request: ProviderRequest) => { + await prepareConversationGeneration(request, 'anthropic', { messages: request.messages }) + throw new Error('Model call sent') + }) + const request = { model: 'model-a', messages: [{ role: 'user' as const, content: 'Continue' }] } + + const attempt = () => + executeProviderRequest('anthropic', request, { + executionContext: run, + agentConversation: reopenTurn(), + }) + + await expect(attempt()).rejects.toThrow('Model call sent') + await expect(attempt()).rejects.toThrow('Model call sent') + expect(run.unrecordedModelCost?.total).toBe(2) + }) + + it('bills a retried pending tool call once, through the journal that records it', async () => { + const tool: ProviderToolConfig = { + id: 'search', + description: 'search', + params: {}, + parameters: { type: 'object', properties: {}, required: [] }, + } + const session = new AgentTurnStateMachine( + { save: vi.fn() }, + { + version: 1, + steps: [ + { + id: 'interrupted-step', + assistant: { role: 'assistant', content: '' }, + calls: [ + { + invocationId: 'call-1', + toolId: 'search', + arguments: '{}', + configuredToolBinding: getConfiguredConversationToolBinding(tool), + }, + ], + results: [], + usage: { input: 1, output: 1 }, + cost: { input: 1, output: 1, total: 2 }, + }, + ], + } + ) + mockExecuteTool.mockResolvedValueOnce({ success: true, output: { cost: { total: 0.25 } } }) + const run = startRun() + run.unrecordedModelCost = { total: 2 } + mockExecuteRequest.mockImplementation(async (request: ProviderRequest) => { + await prepareConversationGeneration(request, 'anthropic', { messages: [] }) + return { content: '', model: request.model } + }) + + const failure = await executeProviderRequest( + 'anthropic', + { + model: 'model-a', + workflowId: 'workflow-1', + messages: [{ role: 'user', content: 'Continue' }], + tools: [tool], + }, + { executionContext: run, agentConversation: session } + ).catch((caught: unknown) => caught) + + expect(isNonRetryableExecutionError(failure)).toBe(true) + expect(readTrustedExecutionCost(failure)).toEqual({ input: 1, output: 1, total: 2.25 }) + }) + + it("charges a continued Function call's failure once, through the journal that records it", async () => { + mockCheckExecutionUsageLimits.mockResolvedValue({ + isExceeded: false, + payerUsage: { currentUsage: 0, limit: 5 }, + }) + const tool: ProviderToolConfig = { + id: 'function_execute', + description: 'run code', + params: {}, + parameters: { type: 'object', properties: {}, required: [] }, + } + const session = new AgentTurnStateMachine( + { save: vi.fn() }, + { + version: 1, + steps: [ + { + id: 'interrupted-step', + assistant: { role: 'assistant', content: '' }, + calls: [ + { + invocationId: 'call-1', + toolId: 'function_execute', + arguments: '{}', + configuredToolBinding: getConfiguredConversationToolBinding(tool), + }, + ], + results: [], + }, + ], + } + ) + mockExecuteTool.mockResolvedValueOnce({ + success: false, + output: { cost: { total: 0.5 } }, + error: 'Function failed', + }) + mockExecuteRequest.mockResolvedValue({ + content: 'Done', + model: 'model-a', + cost: { input: 0, output: 0, total: 0 }, + }) + + const response = await executeProviderRequest( + 'anthropic', + { + model: 'model-a', + workflowId: 'workflow-1', + messages: [{ role: 'user', content: 'Continue' }], + tools: [tool], + }, + { executionContext: startRun(), agentConversation: session } + ) + + expect(response).toMatchObject({ cost: { total: 0.5 } }) + }) + + it('counts every block of the run, not only the request that is running', async () => { + const run = startRun() + await modelStep(startRequest(run), 1, 1) + const later = startRequest(run) + + const stop = prepareConversationGeneration(later, 'anthropic', { messages: [] }) + await expect(stop).rejects.toThrow('usage limit') + expect(readTrustedExecutionCost(await stop.catch((caught: unknown) => caught))).toBeUndefined() + }) + + it('never stops a run that has spent nothing, even with the payer over its limit', async () => { + mockCheckExecutionUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) + const request = startRequest(startRun()) + + await expect( + prepareConversationGeneration(request, 'anthropic', { messages: [] }) + ).resolves.toEqual({ messages: [] }) + expect(mockCheckExecutionUsageLimits).not.toHaveBeenCalled() + }) + + it('never counts steps billed to the workspace key', async () => { + const request = startRequest(startRun(), true) + await modelStep(request, 50, 50) + + await prepareConversationGeneration(request, 'anthropic', { messages: [] }) + + expect(mockCheckExecutionUsageLimits).not.toHaveBeenCalled() + }) + + it('stops a member whose own cap the run reaches before the pool', async () => { + mockCheckExecutionUsageLimits.mockResolvedValue({ + isExceeded: false, + payerUsage: { currentUsage: 10, limit: 1000 }, + memberUsage: { currentUsage: 9, limit: 10 }, + }) + const request = startRequest(startRun()) + await modelStep(request, 0.5, 0.5) + + await expect( + prepareConversationGeneration(request, 'anthropic', { messages: [] }) + ).rejects.toThrow('usage limit') + }) + + it('keeps going when the payer has no member cap and ample headroom', async () => { + mockCheckExecutionUsageLimits.mockResolvedValue({ + isExceeded: false, + payerUsage: { currentUsage: 10, limit: 1000 }, + memberUsage: { currentUsage: 9, limit: null }, + }) + const request = startRequest(startRun()) + await modelStep(request, 5, 5) + + await expect( + prepareConversationGeneration(request, 'anthropic', { messages: [] }) + ).resolves.toEqual({ messages: [] }) + }) + + it('stops when another run already took the payer past its limit', async () => { + const request = startRequest(startRun()) + await modelStep(request, 0.1, 0.1) + mockCheckExecutionUsageLimits.mockResolvedValue({ isExceeded: true, scope: 'payer' }) + + await expect( + prepareConversationGeneration(request, 'anthropic', { messages: [] }) + ).rejects.toThrow('usage limit') + }) + + it.each([ + [ + 'usage cannot be read', + () => mockCheckExecutionUsageLimits.mockRejectedValueOnce(new Error('ledger read timed out')), + ], + [ + 'the ledger is unavailable', + () => + mockCheckExecutionUsageLimits.mockResolvedValueOnce({ + isExceeded: true, + reason: 'usage_unavailable', + }), + ], + [ + 'the account is blocked', + () => + mockCheckExecutionUsageLimits.mockResolvedValueOnce({ + isExceeded: true, + reason: 'billing_blocked', + scope: 'payer', + }), + ], + ])('keeps the run going when %s, without reading again for a while', async (_case, answer) => { + vi.useFakeTimers({ toFake: ['Date'] }) + try { + const request = startRequest(startRun()) + await modelStep(request, 1, 1) + mockCheckExecutionUsageLimits.mockClear() + answer() + + await expect(modelStep(request, 1, 1)).resolves.toBeUndefined() + await expect(modelStep(request, 1, 1)).resolves.toBeUndefined() + expect(mockCheckExecutionUsageLimits).toHaveBeenCalledTimes(1) + + vi.advanceTimersByTime(NEAR_LIMIT_USAGE_GATE_TTL_MS) + await expect( + prepareConversationGeneration(request, 'anthropic', { messages: [] }) + ).rejects.toThrow('usage limit') + } finally { + vi.useRealTimers() + } + }) + + it('reads no usage for a paid payer, whose overage is billed', async () => { + const run = startRun() + run.metadata.billingAttribution = { + ...attribution, + payerSubscription: { + id: 'sub', + referenceId: 'actor', + plan: 'pro', + status: 'active', + seats: 1, + periodStart: null, + periodEnd: null, + }, + } + const request = startRequest(run) + await modelStep(request, 50, 50) + + await prepareConversationGeneration(request, 'anthropic', { messages: [] }) + + expect(mockCheckExecutionUsageLimits).not.toHaveBeenCalled() + }) + + it('reads no usage once the run is cancelled', async () => { + const controller = new AbortController() + const run = { ...startRun(), abortSignal: controller.signal } + const request = startRequest(run) + await modelStep(request, 5, 5) + controller.abort() + + await prepareConversationGeneration(request, 'anthropic', { messages: [] }) + + expect(mockCheckExecutionUsageLimits).not.toHaveBeenCalled() + }) + + it('keeps a cancellation that lands during the usage read a cancellation', async () => { + const controller = new AbortController() + const request = startRequest({ ...startRun(), abortSignal: controller.signal }) + await modelStep(request, 1, 1) + mockCheckExecutionUsageLimits.mockImplementationOnce(async () => { + controller.abort() + return { isExceeded: false, payerUsage: { currentUsage: 3, limit: 5 } } + }) + + await expect( + prepareConversationGeneration(request, 'anthropic', { messages: [] }) + ).resolves.toEqual({ messages: [] }) + }) + + it('counts a turn recorded outside a captured exchange', async () => { + const request = startRequest(startRun()) + await recordProviderConversationUsage(request, { input: 1, output: 1 }) + + await expect( + prepareConversationGeneration(request, 'anthropic', { messages: [] }) + ).rejects.toThrow('usage limit') + }) + + it('bills the stopped block for the tool calls it already settled', async () => { + mockExecuteTool + .mockResolvedValueOnce({ success: true, output: { cost: { total: 0.25 } } }) + .mockResolvedValueOnce({ success: false, output: { cost: { total: 0.5 } } }) + mockExecuteRequest.mockImplementation(async (request: ProviderRequest) => { + await modelStep(request, 1, 1) + await executeProviderTool('search', {}) + // A successful block never bills a failed workflow tool's cost, so neither does the stop. + await executeProviderTool('workflow_executor', {}) + await prepareConversationGeneration(request, 'anthropic', { messages: [] }) + return { content: '', model: request.model } + }) + + const failure = await executeProviderRequest( + 'anthropic', + { model: 'model-a' }, + { executionContext: startRun() } + ).catch((caught: unknown) => caught) + + expect(readTrustedExecutionCost(failure)).toEqual({ input: 1, output: 1, total: 2.25 }) + }) + + it("logs a run timeout at warn even when another block's failure attributed the reason to the user", () => { + const controller = new AbortController() + controller.abort(new DOMException('timeout', 'AbortError')) + markFailureKind(controller.signal.reason, 'user') + + expect(providerFailureLogLevel(controller.signal.reason, controller.signal)).toBe('warn') + }) + + it('reads no usage on a self-hosted deployment', async () => { + setEnvFlags({ isHosted: false }) + const request = startRequest(startRun()) + await modelStep(request, 5, 5) + + await prepareConversationGeneration(request, 'anthropic', { messages: [] }) + + expect(mockCheckExecutionUsageLimits).not.toHaveBeenCalled() + }) +}) diff --git a/apps/sim/providers/run-usage-limit.ts b/apps/sim/providers/run-usage-limit.ts new file mode 100644 index 00000000000..b971248d29c --- /dev/null +++ b/apps/sim/providers/run-usage-limit.ts @@ -0,0 +1,110 @@ +import { readMidRunUsageVerdict } from '@/lib/billing/core/mid-run-usage' +import { NEAR_LIMIT_USAGE_GATE_TTL_MS } from '@/lib/billing/core/usage-gate-cache' +import { isPaid } from '@/lib/billing/plan-helpers' +import { resolveUsageUpgradePayload } from '@/lib/billing/usage-upgrade' +import { markFailureKind } from '@/lib/core/errors/failure-log' +import { NonRetryableExecutionError } from '@/lib/execution/non-retryable-error' +import type { ConversationUsageTotal } from '@/lib/memory/conversation-types' +import { attachTrustedExecutionCost } from '@/executor/utils/errors' +import type { ModelCost } from '@/providers/cost-policy' +import type { TokenUsageCost } from '@/providers/pricing' +import type { ProviderRuntimeContext } from '@/providers/runtime-context' + +/** + * A run stopped at its payer's usage limit. Providers pass it through unwrapped (see + * `isConversationContextError`), so it keeps its billed cost and is never retried. + */ +export class RunUsageLimitError extends NonRetryableExecutionError { + constructor(message: string) { + super(message) + this.name = 'RunUsageLimitError' + } +} + +export function addSettledCost(settledCost: TokenUsageCost, cost: TokenUsageCost): void { + settledCost.input += cost.input + settledCost.output += cost.output + settledCost.total += cost.total +} + +/** + * Adds a resumed Agent turn's journaled model spend to its run's unrecorded total, once per turn. + * Steps an earlier attempt journaled in another process were never counted by this run, while a + * turn this run already opened (a fallback model, a block retry) holds steps it counted as they + * settled. Journaled tool costs are left out, as settled ones are. + */ +export function recordJournaledModelCost( + runtime: ProviderRuntimeContext, + turnId: string | undefined, + usage: ConversationUsageTotal +): void { + const unrecorded = runtime.executionContext?.unrecordedModelCost + if (!unrecorded || !turnId) return + unrecorded.countedTurnIds ??= new Set() + if (unrecorded.countedTurnIds.has(turnId)) return + unrecorded.countedTurnIds.add(turnId) + unrecorded.total += usage.cost.total - usage.cost.toolCost +} + +/** + * Adds a settled model step's billed cost to its request's and its run's running totals. Tool + * costs reach only the request's: a workflow tool's child already spends against the run's. + */ +export function recordModelStepCost( + runtime: ProviderRuntimeContext | undefined, + cost: ModelCost +): void { + if (!runtime) return + if (runtime.settledCost) addSettledCost(runtime.settledCost, cost) + const unrecorded = runtime.executionContext?.unrecordedModelCost + if (unrecorded) unrecorded.total += cost.total +} + +/** + * Stops a free-plan workflow run before its next model call once the payer's recorded usage plus + * the run's own unrecorded model spend reaches its limit. A run's cost reaches the ledger only + * with its log, so admission alone would let one run keep calling models long past the limit, + * and a free plan's spend past its limit is never billed. A paid payer's overshoot is billed as + * overage, so its runs are left to admission and never read here. + * + * Reads nothing until the run has spent, then the cached execution gate. A gate that cannot + * answer, or an account that is blocked, lets the run continue and is not read again for + * {@link NEAR_LIMIT_USAGE_GATE_TTL_MS}: the next admission refuses it. The stop is not retried, + * it carries the cost of the request's settled steps and tool calls so the stopped block is + * billed for them, and it is attributed to the user so it logs at info, not error. + */ +export async function assertRunWithinUsageLimit( + runtime: ProviderRuntimeContext | undefined +): Promise { + const executionContext = runtime?.executionContext + if (!runtime || !executionContext) return + const attribution = executionContext.metadata.billingAttribution + const unrecorded = executionContext.unrecordedModelCost + if ( + !attribution || + !unrecorded || + unrecorded.total <= 0 || + isPaid(attribution.payerSubscription?.plan) || + executionContext.abortSignal?.aborted || + Date.now() < (unrecorded.readDeferredUntil ?? 0) + ) { + return + } + + const verdict = await readMidRunUsageVerdict(attribution, unrecorded.total) + // A run cancelled during the read keeps its cancellation; the provider call then aborts. + if (verdict.status === 'within' || executionContext.abortSignal?.aborted) return + if (verdict.status !== 'exceeded') { + unrecorded.readDeferredUntil = Date.now() + NEAR_LIMIT_USAGE_GATE_TTL_MS + return + } + + const { message } = await resolveUsageUpgradePayload( + attribution.actorUserId, + verdict.payer ?? attribution, + verdict.scope + ) + const error = markFailureKind(new RunUsageLimitError(message), 'user') + if (runtime.settledCost?.total) attachTrustedExecutionCost(error, runtime.settledCost) + throw error +} diff --git a/apps/sim/providers/runtime-context.ts b/apps/sim/providers/runtime-context.ts index d3fad7846b4..4aa6bf9cc50 100644 --- a/apps/sim/providers/runtime-context.ts +++ b/apps/sim/providers/runtime-context.ts @@ -22,6 +22,7 @@ import type { ExecutionContext } from '@/executor/types' import { projectResolvedSecretModelContent } from '@/executor/utils/resolved-secret-content-projection' import type { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' import { stripCloneSuffixes } from '@/executor/utils/subflow-node-id-codec' +import type { TokenUsageCost } from '@/providers/pricing' import { isAbortError } from '@/providers/streaming-tool-loop-shared' import { getPreparedProviderToolInputProvenance } from '@/providers/tool-input-provenance' import type { ProviderId } from '@/providers/types' @@ -41,6 +42,8 @@ export interface ProviderRuntimeContext { toolIdByWireId?: ReadonlyMap /** Failed canonical Function cost omitted from provider tool-result collections. */ failedFunctionToolCost?: { total: number } + /** Billed cost of this request's settled model steps and tool calls, charged if the run stops it. */ + settledCost?: TokenUsageCost } export type ExecuteProviderToolOptions = ExecuteToolOptions @@ -115,17 +118,24 @@ function withoutChildTraceHandle(response: ToolResponse): ToolResponse { } } -function accumulateFailedFunctionToolCost( +function toolCostTotal(result: ToolResponse): number { + if (!isRecordLike(result.output) || !isRecordLike(result.output.cost)) return 0 + const total = result.output.cost.total + return typeof total === 'number' && Number.isFinite(total) && total > 0 ? total : 0 +} + +function accumulateToolCost( toolId: string, result: ToolResponse, - accumulator: ProviderRuntimeContext['failedFunctionToolCost'] + runtime: ProviderRuntimeContext | undefined ): 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 + const total = toolCostTotal(result) + if (total === 0 || !runtime) return + const failedFunction = toolId === 'function_execute' && !result.success + // A stop bills what success would: successful tool results plus failed Function runs. + if (runtime.settledCost && (result.success || failedFunction)) runtime.settledCost.total += total + if (failedFunction && runtime.failedFunctionToolCost) { + runtime.failedFunctionToolCost.total += total } } @@ -303,11 +313,7 @@ export async function executeProviderTool( ...(executionContext ? { executionContext } : {}), resolvedSecretTraceRegistry: toolCallRegistry, }) - accumulateFailedFunctionToolCost( - executionToolId, - result, - runtimeContext?.failedFunctionToolCost - ) + accumulateToolCost(executionToolId, result, runtimeContext) if (!registry || !toolCallRegistry) { return recordResult({ rawResponse: result, modelResponse: withoutChildTraceHandle(result) }) } diff --git a/apps/sim/providers/streaming-tool-loop-shared.ts b/apps/sim/providers/streaming-tool-loop-shared.ts index f76df2892cb..9354728efca 100644 --- a/apps/sim/providers/streaming-tool-loop-shared.ts +++ b/apps/sim/providers/streaming-tool-loop-shared.ts @@ -11,7 +11,7 @@ import { findCause, toError } from '@sim/utils/errors' import { isRecordLike } from '@sim/utils/object' import { APIUserAbortError as GroqUserAbortError } from 'groq-sdk/error' import { APIUserAbortError as OpenAIUserAbortError } from 'openai/error' -import { UserFailure } from '@/lib/core/errors/user-failure' +import { classifyFailure } from '@/lib/core/errors/failure-log' import { isTimeoutAbortReason } from '@/lib/core/execution-limits/types' import type { NormalizedBlockOutput } from '@/executor/types' import type { AgentStreamEvent, ToolCallEndStatus } from '@/providers/stream-events' @@ -74,17 +74,18 @@ export function isRunAbortError(error: unknown, abortSignal: AbortSignal | undef /** * The severity a provider logs its own request failure at, on the `failure-log` scale: a run - * the user stopped or a request the author must fix ({@link UserFailure}) is routine (info), one - * the run's time limit cut off is expected but worth seeing (warn), and anything else is an - * unattributed failure (error). + * the user stopped or a failure {@link classifyFailure} attributes to the user is routine + * (info), one the run's time limit cut off is expected but worth seeing (warn), and anything + * else is an unattributed failure (error). */ export function providerFailureLogLevel( error: unknown, abortSignal: AbortSignal | undefined ): 'info' | 'warn' | 'error' { - if (error instanceof UserFailure) return 'info' - if (!isRunAbortError(error, abortSignal)) return 'error' - return isTimeoutAbortReason(abortSignal?.reason) ? 'warn' : 'info' + if (isRunAbortError(error, abortSignal)) { + return isTimeoutAbortReason(abortSignal?.reason) ? 'warn' : 'info' + } + return classifyFailure(error) === 'user' ? 'info' : 'error' } /** Parse provider-supplied tool arguments without accepting non-object JSON values. */ @@ -131,8 +132,8 @@ interface TerminateToolLoopOptions { /** The stream consumer called `cancel()`; the controller is already closed. */ consumerCancelled: boolean error: unknown - /** Invoked only for a genuine failure, before the stream is errored. */ - onUnexpectedError?: (error: unknown) => void + /** Invoked only for a genuine failure, before the stream is errored, at its log severity. */ + onUnexpectedError?: (error: unknown, level: 'info' | 'warn' | 'error') => void } /** @@ -164,6 +165,6 @@ export function terminateToolLoop({ } settleOpenTools(controller, openTools, 'error') - onUnexpectedError?.(error) + onUnexpectedError?.(error, providerFailureLogLevel(error, undefined)) controller.error(toError(error)) } diff --git a/apps/sim/tools/index.ts b/apps/sim/tools/index.ts index d7ed8dff525..5fa80c13678 100644 --- a/apps/sim/tools/index.ts +++ b/apps/sim/tools/index.ts @@ -2085,6 +2085,7 @@ async function executeToolImplementation( resolvedSecretTraceRegistry, executorDelegationOrigin: executionContext?.executorDelegationOrigin, principal: executionContext?.principal, + unrecordedModelCost: executionContext?.unrecordedModelCost, // Trusted `executionContext`, never `_context` — that bag spreads // model-reachable `contextParams._context` first, so a model could otherwise // inject its own env map or disable redaction. diff --git a/packages/testing/src/mocks/billing-usage-gate-cache.mock.ts b/packages/testing/src/mocks/billing-usage-gate-cache.mock.ts index 75ebe916263..4c557558217 100644 --- a/packages/testing/src/mocks/billing-usage-gate-cache.mock.ts +++ b/packages/testing/src/mocks/billing-usage-gate-cache.mock.ts @@ -25,7 +25,8 @@ const USAGE_GATE_TTL_MS = 5 * 60 * 1000 /** * Static mock module for `@/lib/billing/core/usage-gate-cache`. Constants carry the real - * values (`USAGE_GATE_SETTLE_TIMEOUT_MS` = 2 x the 60 s ledger statement timeout + 15 s). + * values (`NEAR_LIMIT_USAGE_GATE_TTL_MS` = 10 s, `USAGE_GATE_SETTLE_TIMEOUT_MS` = 2 x the 60 s + * ledger statement timeout + 15 s). * * @example * ```ts @@ -34,6 +35,7 @@ const USAGE_GATE_TTL_MS = 5 * 60 * 1000 */ export const billingUsageGateCacheMock = { USAGE_GATE_TTL_MS, + NEAR_LIMIT_USAGE_GATE_TTL_MS: 10 * 1000, USAGE_GATE_SETTLE_TIMEOUT_MS: 2 * 60_000 + 15_000, checkIngestionUsageLimits: billingUsageGateCacheMockFns.mockCheckIngestionUsageLimits, checkSearchUsageLimits: billingUsageGateCacheMockFns.mockCheckSearchUsageLimits,