diff --git a/apps/sim/background/falai-cost-settlement.ts b/apps/sim/background/falai-cost-settlement.ts new file mode 100644 index 00000000000..7543765b79e --- /dev/null +++ b/apps/sim/background/falai-cost-settlement.ts @@ -0,0 +1,37 @@ +import { createLogger } from '@sim/logger' +import { task } from '@trigger.dev/sdk' +import { + FALAI_COST_SETTLEMENT_TASK_ID, + type FalAICostSettlementPayload, + settleFalAICost, +} from '@/lib/tools/falai-settlement' + +const logger = createLogger('FalAICostSettlementTask') + +/** + * Settles a hosted fal generation billed on an estimate to fal's exact charge. A receipt fal + * has not recorded yet fails the attempt so the retries keep looking for about an hour; past + * that the estimate stands. + */ +export const falaiCostSettlementTask = task({ + id: FALAI_COST_SETTLEMENT_TASK_ID, + retry: { + maxAttempts: 8, + factor: 2, + minTimeoutInMs: 60_000, + maxTimeoutInMs: 15 * 60_000, + }, + queue: { concurrencyLimit: 5 }, + run: async (payload: FalAICostSettlementPayload) => { + const outcome = await settleFalAICost(payload) + if (outcome === 'pending') { + throw new Error(`fal has not recorded the charge for request ${payload.requestId} yet`) + } + logger.info('fal cost settlement finished', { + requestId: payload.requestId, + endpointId: payload.endpointId, + outcome, + }) + return { outcome } + }, +}) diff --git a/apps/sim/lib/internal/image/operations.test.ts b/apps/sim/lib/internal/image/operations.test.ts index 7021860ffb1..1a3a3d92d37 100644 --- a/apps/sim/lib/internal/image/operations.test.ts +++ b/apps/sim/lib/internal/image/operations.test.ts @@ -102,7 +102,7 @@ describe('image operations', () => { ]) }) - it('cancels polling without resubmitting or storing an image', async () => { + it('cancels the fal request when polling stops, without resubmitting or storing an image', async () => { const controller = new AbortController() mocks.fetch.mockResolvedValueOnce( Response.json({ @@ -121,12 +121,16 @@ describe('image operations', () => { }) ).rejects.toMatchObject({ name: 'AbortError' }) - expect(mocks.fetch).toHaveBeenCalledTimes(1) + expect(mocks.fetch).toHaveBeenCalledTimes(2) + expect(mocks.fetch).toHaveBeenLastCalledWith( + expect.stringMatching(/^https:\/\/queue\.fal\.run\/.+\/requests\/job-2\/cancel$/), + expect.objectContaining({ method: 'PUT' }) + ) expect(mockSecureFetchWithPinnedIP).not.toHaveBeenCalled() expect(mockUploadCopilotFile).not.toHaveBeenCalled() }) - it('rejects foreign fal queue hosts before sending credentials', async () => { + it('rejects foreign fal queue hosts before sending credentials and cancels on the fal queue', async () => { mocks.fetch.mockResolvedValueOnce( Response.json({ request_id: 'foreign-job', @@ -143,6 +147,10 @@ describe('image operations', () => { expect(response.status).toBe(500) expect(await response.json()).toMatchObject({ error: 'Invalid Fal.ai queue URL' }) expect(mockSecureFetchWithPinnedIP).not.toHaveBeenCalled() + expect(mocks.fetch).toHaveBeenLastCalledWith( + expect.stringMatching(/^https:\/\/queue\.fal\.run\/.+\/requests\/foreign-job\/cancel$/), + expect.objectContaining({ method: 'PUT' }) + ) }) it('retrieves nested fal jobs from the documented application request path', async () => { diff --git a/apps/sim/lib/internal/image/operations.ts b/apps/sim/lib/internal/image/operations.ts index bc505761e00..de2202cba7f 100644 --- a/apps/sim/lib/internal/image/operations.ts +++ b/apps/sim/lib/internal/image/operations.ts @@ -19,6 +19,7 @@ import { import { MAX_REMOTE_IMAGE_BYTES } from '@/lib/internal/image/fetch' import { buildImageGenerationRequest } from '@/lib/internal/image/request' import type { ImageGenerationInput, ImageProvider } from '@/lib/internal/image/schema' +import { cancelFalRequest } from '@/lib/media/falai' import { FalAIResultError, fetchFalAIResult } from '@/lib/media/falai-result' import { recordServiceMeteringFailure } from '@/lib/mothership/billing/service-observer' import type { FalAICostMetadata } from '@/lib/tools/falai-pricing' @@ -456,134 +457,149 @@ async function generateWithFalAI( throw new Error('Fal.ai queue response missing request_id') } - const statusUrl = falQueueUrl( - getStringProperty(createData, 'status_url') || - buildFalAIQueueUrl(endpoint, falRequestId, 'status') - ) - const responseUrl = falQueueUrl( - getStringProperty(createData, 'response_url') || - buildFalAIQueueUrl(endpoint, falRequestId, 'response') - ) + /** Built from the endpoint, never a provider URL, so a rejected one can still be cancelled. */ + const cancelUrl = `${buildFalAIQueueUrl(endpoint, falRequestId, 'response')}/cancel` logger.info('Fal.ai image request created', { requestId, falRequestId }) - const statusValidation = await validateUrlWithDNS(statusUrl, 'statusUrl', 'contentFetch') - if (!statusValidation.isValid) { - throw new Error(statusValidation.error) - } - - const pollIntervalMs = 3000 - const maxAttempts = Math.ceil(getMaxExecutionTimeout() / pollIntervalMs) - let attempts = 0 - - while (attempts < maxAttempts) { - await interruptibleSleep(pollIntervalMs, signal) - signal?.throwIfAborted() - - const statusResponse = await secureFetchWithPinnedIP(statusUrl, statusValidation.resolvedIP, { - profile: 'contentFetch', - maxRedirects: 0, - headers: { - Authorization: `Key ${apiKey}`, - }, - maxResponseBytes: MAX_IMAGE_JSON_BYTES, - signal, - }) + let reachedTerminalStatus = false + try { + const statusUrl = falQueueUrl( + getStringProperty(createData, 'status_url') || + buildFalAIQueueUrl(endpoint, falRequestId, 'status') + ) + const responseUrl = falQueueUrl( + getStringProperty(createData, 'response_url') || + buildFalAIQueueUrl(endpoint, falRequestId, 'response') + ) - if (!statusResponse.ok) { - await readResponseTextWithLimit(statusResponse, { - maxBytes: DEFAULT_MAX_ERROR_BODY_BYTES, - label: 'Fal.ai status error response', - signal, - }).catch(() => '') - throw new Error(`Fal.ai status check failed: ${statusResponse.status}`) + const statusValidation = await validateUrlWithDNS(statusUrl, 'statusUrl', 'contentFetch') + if (!statusValidation.isValid) { + throw new Error(statusValidation.error) } - const statusData = await readResponseJsonWithLimit(statusResponse, { - maxBytes: MAX_IMAGE_JSON_BYTES, - label: 'Fal.ai status response', - signal, - }) - if (!isRecordLike(statusData)) { - throw new Error('Invalid Fal.ai status response') - } + const pollIntervalMs = 3000 + const maxAttempts = Math.ceil(getMaxExecutionTimeout() / pollIntervalMs) + let attempts = 0 - const status = getStringProperty(statusData, 'status') - if (status === 'COMPLETED') { - const { data: resultData, cost: falaiCost } = await fetchFalAIResult({ - apiKey, - endpointId: endpoint, - requestId: falRequestId, - queueError: statusData.error, + while (attempts < maxAttempts) { + await interruptibleSleep(pollIntervalMs, signal) + signal?.throwIfAborted() + + const statusResponse = await secureFetchWithPinnedIP(statusUrl, statusValidation.resolvedIP, { + profile: 'contentFetch', + maxRedirects: 0, + headers: { + Authorization: `Key ${apiKey}`, + }, maxResponseBytes: MAX_IMAGE_JSON_BYTES, - trackCost: body.useHostedCostTracking, - onCost: onFalCost, signal, - fetchResponse: async (resultSignal) => { - const resultUrl = falQueueUrl( - getStringProperty(statusData, 'response_url') || responseUrl - ) - const resultValidation = await validateUrlWithDNS(resultUrl, 'resultUrl', 'contentFetch') - if (!resultValidation.isValid) throw new Error(resultValidation.error) - return secureFetchWithPinnedIP(resultUrl, resultValidation.resolvedIP, { - profile: 'contentFetch', - maxRedirects: 0, - headers: { Authorization: `Key ${apiKey}` }, - maxResponseBytes: MAX_IMAGE_JSON_BYTES, - signal: resultSignal, - }) - }, }) - if (!isRecordLike(resultData)) { - throw new Error('Invalid Fal.ai result response') + if (!statusResponse.ok) { + await readResponseTextWithLimit(statusResponse, { + maxBytes: DEFAULT_MAX_ERROR_BODY_BYTES, + label: 'Fal.ai status error response', + signal, + }).catch(() => '') + throw new Error(`Fal.ai status check failed: ${statusResponse.status}`) } - const firstImage = - prepared.model.output === 'image' && isRecordLike(resultData.image) - ? resultData.image - : firstRecord(resultData.images) - const imageUrl = - getStringProperty(firstImage, 'url') || - getStringProperty(firstImage, 'data') || - getStringProperty(firstImage, 'content') - if (!imageUrl) { - throw new Error('No image URL in Fal.ai response') + const statusData = await readResponseJsonWithLimit(statusResponse, { + maxBytes: MAX_IMAGE_JSON_BYTES, + label: 'Fal.ai status response', + signal, + }) + if (!isRecordLike(statusData)) { + throw new Error('Invalid Fal.ai status response') } - const downloaded = await bufferFromImageUrl(imageUrl, signal) - const contentType = - getStringProperty(firstImage, 'content_type') || - getStringProperty(firstImage, 'contentType') || - downloaded.contentType - const fileName = - getStringProperty(firstImage, 'file_name') || - getStringProperty(firstImage, 'fileName') || - `falai-${model}.${extensionFromContentType(contentType)}` - - return { - buffer: downloaded.buffer, - contentType, - fileName, - provider: 'falai', - model, - sourceUrl: imageUrl.startsWith('data:') ? undefined : imageUrl, - description: getStringProperty(resultData, 'description'), - revisedPrompt: getStringProperty(resultData, 'revised_prompt'), - seed: getNumberProperty(resultData, 'seed'), - jobId: falRequestId, - falaiCost, + const status = getStringProperty(statusData, 'status') + if (status === 'COMPLETED') { + reachedTerminalStatus = true + const { data: resultData, cost: falaiCost } = await fetchFalAIResult({ + apiKey, + endpointId: endpoint, + requestId: falRequestId, + queueError: statusData.error, + maxResponseBytes: MAX_IMAGE_JSON_BYTES, + trackCost: body.useHostedCostTracking, + onCost: onFalCost, + signal, + fetchResponse: async (resultSignal) => { + const resultUrl = falQueueUrl( + getStringProperty(statusData, 'response_url') || responseUrl + ) + const resultValidation = await validateUrlWithDNS( + resultUrl, + 'resultUrl', + 'contentFetch' + ) + if (!resultValidation.isValid) throw new Error(resultValidation.error) + return secureFetchWithPinnedIP(resultUrl, resultValidation.resolvedIP, { + profile: 'contentFetch', + maxRedirects: 0, + headers: { Authorization: `Key ${apiKey}` }, + maxResponseBytes: MAX_IMAGE_JSON_BYTES, + signal: resultSignal, + }) + }, + }) + + if (!isRecordLike(resultData)) { + throw new Error('Invalid Fal.ai result response') + } + + const firstImage = + prepared.model.output === 'image' && isRecordLike(resultData.image) + ? resultData.image + : firstRecord(resultData.images) + const imageUrl = + getStringProperty(firstImage, 'url') || + getStringProperty(firstImage, 'data') || + getStringProperty(firstImage, 'content') + if (!imageUrl) { + throw new Error('No image URL in Fal.ai response') + } + + const downloaded = await bufferFromImageUrl(imageUrl, signal) + const contentType = + getStringProperty(firstImage, 'content_type') || + getStringProperty(firstImage, 'contentType') || + downloaded.contentType + const fileName = + getStringProperty(firstImage, 'file_name') || + getStringProperty(firstImage, 'fileName') || + `falai-${model}.${extensionFromContentType(contentType)}` + + return { + buffer: downloaded.buffer, + contentType, + fileName, + provider: 'falai', + model, + sourceUrl: imageUrl.startsWith('data:') ? undefined : imageUrl, + description: getStringProperty(resultData, 'description'), + revisedPrompt: getStringProperty(resultData, 'revised_prompt'), + seed: getNumberProperty(resultData, 'seed'), + jobId: falRequestId, + falaiCost, + } } - } - if (['ERROR', 'FAILED', 'CANCELLED'].includes(status || '')) { - throw new Error(`Fal.ai generation failed: ${getFalAIErrorMessage(statusData.error)}`) + if (['ERROR', 'FAILED', 'CANCELLED'].includes(status || '')) { + reachedTerminalStatus = true + throw new Error(`Fal.ai generation failed: ${getFalAIErrorMessage(statusData.error)}`) + } + + attempts += 1 } - attempts += 1 + throw new Error('Fal.ai image generation timed out') + } catch (error) { + if (!reachedTerminalStatus) await cancelFalRequest(cancelUrl, apiKey, falRequestId) + throw error } - - throw new Error('Fal.ai image generation timed out') } async function storeGeneratedImage( diff --git a/apps/sim/lib/internal/video/client.ts b/apps/sim/lib/internal/video/client.ts index 6ef27ce4565..9f8a8137352 100644 --- a/apps/sim/lib/internal/video/client.ts +++ b/apps/sim/lib/internal/video/client.ts @@ -17,7 +17,7 @@ import { import { buildDirectVideoRequest } from '@/lib/internal/video/direct-request' import { buildFalVideoRequest, getFalAIValidationError } from '@/lib/internal/video/fal-request' import { getFalVideoModel } from '@/lib/media/fal-video-catalog' -import { downloadFalMedia } from '@/lib/media/falai' +import { cancelFalRequest, downloadFalMedia } from '@/lib/media/falai' import { fetchFalAIResult } from '@/lib/media/falai-result' import type { FalAICostMetadata } from '@/lib/tools/falai-pricing' import { downloadFileFromStorage } from '@/lib/uploads/utils/file.server' @@ -393,100 +393,113 @@ async function generateWithFalAI( const app = endpoint.split('/', 2).join('/') const requestUrl = `https://queue.fal.run/${app}/requests/${encodeURIComponent(requestIdFal)}` - const statusUrl = falQueueUrl( - getStringProperty(createData, 'status_url') || `${requestUrl}/status` - ) - const responseUrl = falQueueUrl(getStringProperty(createData, 'response_url') || requestUrl) + /** Built from the endpoint, never a provider URL, so a rejected one can still be cancelled. */ + const cancelUrl = `${requestUrl}/cancel` logger.info(`Fal.ai request created: ${requestIdFal}`, { requestId }) - const maxAttempts = Math.ceil(getMaxExecutionTimeout() / POLL_INTERVAL_MS) - let attempts = 0 - - while (attempts < maxAttempts) { - await waitForProvider(context) - - const statusResponse = await fetch(statusUrl, { - redirect: 'error', - headers: { - Authorization: `Key ${apiKey}`, - }, - signal, - }) + let reachedTerminalStatus = false + try { + const statusUrl = falQueueUrl( + getStringProperty(createData, 'status_url') || `${requestUrl}/status` + ) + const responseUrl = falQueueUrl(getStringProperty(createData, 'response_url') || requestUrl) - if (!statusResponse.ok) { - await readVideoErrorText(statusResponse, 'Fal.ai status error response') - throw new Error(`Fal.ai status check failed: ${statusResponse.status}`) - } + const maxAttempts = Math.ceil(getMaxExecutionTimeout() / POLL_INTERVAL_MS) + let attempts = 0 - const statusData = await readVideoJson(statusResponse, 'Fal.ai status response') - if (!isRecordLike(statusData)) { - throw new Error('Invalid Fal.ai status response') - } + while (attempts < maxAttempts) { + await waitForProvider(context) - if (getStringProperty(statusData, 'status') === 'COMPLETED') { - const { data: resultData, cost: falaiCost } = await fetchFalAIResult({ - apiKey, - endpointId: endpoint, - requestId: requestIdFal, - queueError: statusData.error, - maxResponseBytes: MAX_VIDEO_JSON_BYTES, - trackCost: useHostedCostTracking, - onCost: context.onFalCost, + const statusResponse = await fetch(statusUrl, { + redirect: 'error', + headers: { + Authorization: `Key ${apiKey}`, + }, signal, - fetchResponse: (resultSignal) => - fetch(falQueueUrl(getStringProperty(statusData, 'response_url') || responseUrl), { - redirect: 'error', - headers: { Authorization: `Key ${apiKey}` }, - signal: resultSignal, - }), }) - if (!isRecordLike(resultData)) { - throw new Error('Invalid Fal.ai result response') + if (!statusResponse.ok) { + await readVideoErrorText(statusResponse, 'Fal.ai status error response') + throw new Error(`Fal.ai status check failed: ${statusResponse.status}`) } - const videoOutput = isRecordLike(resultData.video) ? resultData.video : undefined - const fallbackOutput = isRecordLike(resultData.output) ? resultData.output : undefined - const videoUrl = - getStringProperty(videoOutput, 'url') || getStringProperty(fallbackOutput, 'url') - if (!videoUrl) { - throw new Error('No video URL in response') + const statusData = await readVideoJson(statusResponse, 'Fal.ai status response') + if (!isRecordLike(statusData)) { + throw new Error('Invalid Fal.ai status response') } - const { buffer } = await downloadFalMedia(videoUrl, signal) - const catalogModel = getFalVideoModel(model) - const legacyMetadata = !catalogModel || catalogModel.legacyId === model - let width = getNumberProperty(videoOutput, 'width') - let height = getNumberProperty(videoOutput, 'height') - if (legacyMetadata && !width && aspectRatio?.includes(':')) { - const dims = getVideoDimensions(aspectRatio, resolution || '1080p') - width = dims.width - height = dims.height + if (getStringProperty(statusData, 'status') === 'COMPLETED') { + reachedTerminalStatus = true + const { data: resultData, cost: falaiCost } = await fetchFalAIResult({ + apiKey, + endpointId: endpoint, + requestId: requestIdFal, + queueError: statusData.error, + maxResponseBytes: MAX_VIDEO_JSON_BYTES, + trackCost: useHostedCostTracking, + onCost: context.onFalCost, + signal, + fetchResponse: (resultSignal) => + fetch(falQueueUrl(getStringProperty(statusData, 'response_url') || responseUrl), { + redirect: 'error', + headers: { Authorization: `Key ${apiKey}` }, + signal: resultSignal, + }), + }) + + if (!isRecordLike(resultData)) { + throw new Error('Invalid Fal.ai result response') + } + + const videoOutput = isRecordLike(resultData.video) ? resultData.video : undefined + const fallbackOutput = isRecordLike(resultData.output) ? resultData.output : undefined + const videoUrl = + getStringProperty(videoOutput, 'url') || getStringProperty(fallbackOutput, 'url') + if (!videoUrl) { + throw new Error('No video URL in response') + } + + const { buffer } = await downloadFalMedia(videoUrl, signal) + const catalogModel = getFalVideoModel(model) + const legacyMetadata = !catalogModel || catalogModel.legacyId === model + let width = getNumberProperty(videoOutput, 'width') + let height = getNumberProperty(videoOutput, 'height') + if (legacyMetadata && !width && aspectRatio?.includes(':')) { + const dims = getVideoDimensions(aspectRatio, resolution || '1080p') + width = dims.width + height = dims.height + } + + return { + buffer, + width: width ?? (legacyMetadata ? 1920 : undefined), + height: height ?? (legacyMetadata ? 1080 : undefined), + jobId: requestIdFal, + duration: + getNumberProperty(videoOutput, 'duration') ?? + getNumberProperty(resultData, 'duration') ?? + duration ?? + (legacyMetadata ? 5 : undefined), + falaiCost, + } } - return { - buffer, - width: width ?? (legacyMetadata ? 1920 : undefined), - height: height ?? (legacyMetadata ? 1080 : undefined), - jobId: requestIdFal, - duration: - getNumberProperty(videoOutput, 'duration') ?? - getNumberProperty(resultData, 'duration') ?? - duration ?? - (legacyMetadata ? 5 : undefined), - falaiCost, + if ( + ['ERROR', 'FAILED', 'CANCELLED'].includes(getStringProperty(statusData, 'status') || '') + ) { + reachedTerminalStatus = true + throw new Error(`Fal.ai generation failed: ${getFalAIErrorMessage(statusData.error)}`) } - } - if (['ERROR', 'FAILED', 'CANCELLED'].includes(getStringProperty(statusData, 'status') || '')) { - throw new Error(`Fal.ai generation failed: ${getFalAIErrorMessage(statusData.error)}`) + attempts++ } - attempts++ + throw new Error('Fal.ai generation timed out') + } catch (error) { + if (!reachedTerminalStatus) await cancelFalRequest(cancelUrl, apiKey, requestIdFal) + throw error } - - throw new Error('Fal.ai generation timed out') } function getVideoDimensions( diff --git a/apps/sim/lib/media/falai-audio.ts b/apps/sim/lib/media/falai-audio.ts index e9705d9e89e..719eb68187e 100644 --- a/apps/sim/lib/media/falai-audio.ts +++ b/apps/sim/lib/media/falai-audio.ts @@ -159,7 +159,7 @@ export async function generateFalAudio(params: GenerateFalAudioParams): Promise< { tool: 'generate_audio', envVarName }, params.signal ) - const cost = await recordFalCost(apiKey, model, requestId, params.onCost) + const cost = await recordFalCost({ apiKey, envVarName }, model, requestId, params.onCost) const url = extractFalMediaUrl(data, ['audio', 'audio_url', 'audio_file', 'output']) if (!url) throw new Error('No audio URL in Fal.ai clone response') const { buffer, contentType } = await downloadFalMedia(url, params.signal) @@ -185,7 +185,7 @@ export async function generateFalAudio(params: GenerateFalAudioParams): Promise< { tool: 'generate_audio', envVarName }, params.signal ) - const cost = await recordFalCost(apiKey, model, requestId, params.onCost) + const cost = await recordFalCost({ apiKey, envVarName }, model, requestId, params.onCost) const url = extractFalMediaUrl(data, ['audio', 'audio_url', 'audio_file', 'output']) if (!url) throw new Error('No audio URL in Fal.ai response') diff --git a/apps/sim/lib/media/falai-video.ts b/apps/sim/lib/media/falai-video.ts index ddeb7e56b6d..b1eb49f1923 100644 --- a/apps/sim/lib/media/falai-video.ts +++ b/apps/sim/lib/media/falai-video.ts @@ -184,7 +184,7 @@ export async function generateFalVideo(params: GenerateFalVideoParams): Promise< { tool: 'generate_video', envVarName }, params.signal ) - const cost = await recordFalCost(apiKey, endpoint, requestId, params.onCost) + const cost = await recordFalCost({ apiKey, envVarName }, endpoint, requestId, params.onCost) const url = extractFalMediaUrl(data, ['video', 'output']) if (!url) throw new Error('No video URL in Fal.ai response') diff --git a/apps/sim/lib/media/falai.test.ts b/apps/sim/lib/media/falai.test.ts index 410ec48b575..0bde5560f7c 100644 --- a/apps/sim/lib/media/falai.test.ts +++ b/apps/sim/lib/media/falai.test.ts @@ -14,6 +14,7 @@ import { generateFalAudio } from '@/lib/media/falai-audio' import { generateFalVideo } from '@/lib/media/falai-video' const FAL_KEY_USE = { tool: 'generate_video', envVarName: 'FALAI_API_KEY' } as const +const HOSTED_KEY = { envVarName: FAL_KEY_USE.envVarName } const CANCEL_URL = 'https://queue.fal.run/fal-ai/model/requests/req-1/cancel' const STATUS_URL = 'https://queue.fal.run/fal-ai/model/requests/req-1/status' @@ -337,7 +338,8 @@ describe('billing a fal generation', () => { expect(await generation).toEqual(new Error('watchdog')) expect(onCost).toHaveBeenCalledTimes(1) expect(onCost).toHaveBeenCalledWith( - expect.objectContaining({ source: 'billing_events', costDollars: 0.04 }) + expect.objectContaining({ source: 'billing_events', costDollars: 0.04 }), + HOSTED_KEY ) } ) @@ -353,7 +355,10 @@ describe('billing a fal generation', () => { expect(await generation).toMatch(/No (audio|video) URL/) expect(onCost).toHaveBeenCalledTimes(1) - expect(onCost).toHaveBeenCalledWith(expect.objectContaining({ source: 'billing_events' })) + expect(onCost).toHaveBeenCalledWith( + expect.objectContaining({ source: 'billing_events' }), + HOSTED_KEY + ) } ) @@ -373,7 +378,8 @@ describe('billing a fal generation', () => { expect(await generation).toEqual(new Error('watchdog')) expect(onCost).toHaveBeenCalledTimes(1) expect(onCost).toHaveBeenCalledWith( - expect.objectContaining({ source: 'billing_events', costDollars: 0.04 }) + expect.objectContaining({ source: 'billing_events', costDollars: 0.04 }), + HOSTED_KEY ) expect(inputValidationMockFns.mockSecureFetchWithPinnedIP).not.toHaveBeenCalled() }) @@ -393,7 +399,10 @@ describe('billing a fal generation', () => { expect(await generation).toEqual(new Error('watchdog')) expect(onCost).toHaveBeenCalledTimes(1) - expect(onCost).toHaveBeenCalledWith(expect.objectContaining({ source: 'fallback_floor' })) + expect(onCost).toHaveBeenCalledWith( + expect.objectContaining({ source: 'fallback_floor' }), + HOSTED_KEY + ) }) it('bills a generation that completes and downloads exactly once', async () => { @@ -407,7 +416,7 @@ describe('billing a fal generation', () => { expect(audio.buffer.toString()).toBe('audio') expect(onCost).toHaveBeenCalledTimes(1) - expect(onCost).toHaveBeenCalledWith(audio.cost) + expect(onCost).toHaveBeenCalledWith(audio.cost, HOSTED_KEY) }) it('bills nothing for a generation stopped before fal completes it', async () => { diff --git a/apps/sim/lib/media/falai.ts b/apps/sim/lib/media/falai.ts index e92dadac42b..28ed99a0687 100644 --- a/apps/sim/lib/media/falai.ts +++ b/apps/sim/lib/media/falai.ts @@ -257,7 +257,7 @@ async function pollFalQueue( * Asks fal to drop a request nobody will read: a queued request never runs, and a running * one is signalled to stop. Best effort; the generation already failed either way. */ -async function cancelFalRequest(cancelUrl: string, apiKey: string, requestId: string) { +export async function cancelFalRequest(cancelUrl: string, apiKey: string, requestId: string) { try { // boundary-raw-fetch: cancels a request on the external Fal.ai queue const response = await fetch(cancelUrl, { @@ -277,8 +277,14 @@ async function cancelFalRequest(cancelUrl: string, apiKey: string, requestId: st } } -/** Records a completed generation's cost; called at most once per generation. */ -export type FalCostRecorder = (cost: FalAICostMetadata) => Promise +/** + * Records a completed generation's cost; called at most once per generation. `envVarName` names + * the hosted key that ran it, so a later exact-cost settlement reads the same key. + */ +export type FalCostRecorder = ( + cost: FalAICostMetadata, + hostedKey: { envVarName: string } +) => Promise /** * Bills a request fal has completed. fal charges for it whatever happens next, so the lookup @@ -286,13 +292,13 @@ export type FalCostRecorder = (cost: FalAICostMetadata) => Promise * the generation on our side: reading the media URL and the download a Stop can end. */ export async function recordFalCost( - apiKey: string, + hostedKey: { apiKey: string; envVarName: string }, endpointId: string, requestId: string, onCost: FalCostRecorder ): Promise { - const cost = await getFalAICostMetadata({ apiKey, endpointId, requestId }) - await onCost(cost) + const cost = await getFalAICostMetadata({ apiKey: hostedKey.apiKey, endpointId, requestId }) + await onCost(cost, { envVarName: hostedKey.envVarName }) return cost } diff --git a/apps/sim/lib/mothership/tools/server/media/generate-audio.ts b/apps/sim/lib/mothership/tools/server/media/generate-audio.ts index fc57fab405d..61afdeede3c 100644 --- a/apps/sim/lib/mothership/tools/server/media/generate-audio.ts +++ b/apps/sim/lib/mothership/tools/server/media/generate-audio.ts @@ -17,6 +17,10 @@ import { validateCopilotWorkspaceFileWriteTarget, writeCopilotWorkspaceFileByPath, } from '@/lib/mothership/vfs/resource-writer' +import { + MOTHERSHIP_FAL_AT_COST_MULTIPLIER, + scheduleFalAICostSettlement, +} from '@/lib/tools/falai-settlement' import { createWorkspaceFileSecretProvenanceFromRegistry } from '@/lib/uploads/contexts/workspace/file-secret-provenance' import { fileOperations } from '@/lib/workspace-files/application/operations' import { readWorkspaceFileContent } from '@/lib/workspace-files/application/read-workspace-file-content' @@ -142,7 +146,16 @@ export const generateAudioServerTool: BaseServerTool recordServiceCost('falai_audio', cost.costDollars), + onCost: async (cost, { envVarName }) => { + await recordServiceCost('falai_audio', cost.costDollars) + await scheduleFalAICostSettlement({ + cost, + envVarName, + markupMultiplier: MOTHERSHIP_FAL_AT_COST_MULTIPLIER, + source: 'copilot', + billingAttribution: context.billingAttribution, + }) + }, }) const ext = audioExtFromContentType(result.contentType) diff --git a/apps/sim/lib/mothership/tools/server/media/generate-video.ts b/apps/sim/lib/mothership/tools/server/media/generate-video.ts index 71e2f216e40..d6fc3c76c8c 100644 --- a/apps/sim/lib/mothership/tools/server/media/generate-video.ts +++ b/apps/sim/lib/mothership/tools/server/media/generate-video.ts @@ -17,6 +17,10 @@ import { validateCopilotWorkspaceFileWriteTarget, writeCopilotWorkspaceFileByPath, } from '@/lib/mothership/vfs/resource-writer' +import { + MOTHERSHIP_FAL_AT_COST_MULTIPLIER, + scheduleFalAICostSettlement, +} from '@/lib/tools/falai-settlement' import { createWorkspaceFileSecretProvenanceFromRegistry } from '@/lib/uploads/contexts/workspace/file-secret-provenance' import { fileOperations } from '@/lib/workspace-files/application/operations' import { readWorkspaceFileContent } from '@/lib/workspace-files/application/read-workspace-file-content' @@ -119,7 +123,16 @@ export const generateVideoServerTool: BaseServerTool recordServiceCost('falai_video', cost.costDollars), + onCost: async (cost, { envVarName }) => { + await recordServiceCost('falai_video', cost.costDollars) + await scheduleFalAICostSettlement({ + cost, + envVarName, + markupMultiplier: MOTHERSHIP_FAL_AT_COST_MULTIPLIER, + source: 'copilot', + billingAttribution: context.billingAttribution, + }) + }, }) assertServerToolNotAborted(context) diff --git a/apps/sim/lib/mothership/tools/server/media/model-boundaries.test.ts b/apps/sim/lib/mothership/tools/server/media/model-boundaries.test.ts index d1e9d401050..a43b77a66f0 100644 --- a/apps/sim/lib/mothership/tools/server/media/model-boundaries.test.ts +++ b/apps/sim/lib/mothership/tools/server/media/model-boundaries.test.ts @@ -39,14 +39,16 @@ import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-tr interface FalGenerationParams { signal?: AbortSignal - onCost: (cost: { costDollars: number }) => Promise + onCost: (cost: { costDollars: number }, hostedKey: { envVarName: string }) => Promise } +const HOSTED_KEY = { envVarName: 'FALAI_API_KEY_1' } + /** A fal generation that bills 0.1 through `onCost`, then returns `media` with that cost. */ function billsThen(media: Record) { return async ({ onCost }: FalGenerationParams) => { const cost = { costDollars: 0.1 } - await onCost(cost) + await onCost(cost, HOSTED_KEY) return { ...media, cost } } } @@ -142,7 +144,7 @@ describe('Mothership media model boundaries', () => { 'bills a completed %s generation once when the tool call is aborted after it', async (_kind, service, generator, tool) => { generator.mockImplementationOnce(async ({ onCost }: FalGenerationParams) => { - await onCost({ costDollars: 0.2 }) + await onCost({ costDollars: 0.2 }, HOSTED_KEY) throw new Error('Tool timed out') }) const charges = vi.fn().mockResolvedValue(undefined) diff --git a/apps/sim/lib/mothership/tools/server/media/write-target-precheck.integration.ts b/apps/sim/lib/mothership/tools/server/media/write-target-precheck.integration.ts index f8505521287..94ef39026ac 100644 --- a/apps/sim/lib/mothership/tools/server/media/write-target-precheck.integration.ts +++ b/apps/sim/lib/mothership/tools/server/media/write-target-precheck.integration.ts @@ -64,9 +64,13 @@ afterAll(async () => { /** A fal generation that bills through `onCost` before returning, as the real client does. */ function billsThen(media: Record) { - return async ({ onCost }: { onCost: (cost: { costDollars: number }) => Promise }) => { + return async ({ + onCost, + }: { + onCost: (cost: { costDollars: number }, hostedKey: { envVarName: string }) => Promise + }) => { const cost = { costDollars: 0.1 } - await onCost(cost) + await onCost(cost, { envVarName: 'FALAI_API_KEY_1' }) return { ...media, cost } } } diff --git a/apps/sim/lib/tools/falai-pricing.ts b/apps/sim/lib/tools/falai-pricing.ts index c5f2e54dd75..3b3a9cb22da 100644 --- a/apps/sim/lib/tools/falai-pricing.ts +++ b/apps/sim/lib/tools/falai-pricing.ts @@ -64,6 +64,16 @@ export interface FalAICostMetadata { error?: string } +const FALAI_COST_SOURCES = new Set([ + 'billing_events', + 'historical_estimate', + 'fallback_floor', +]) + +export function isFalAICostSource(value: unknown): value is FalAICostMetadata['source'] { + return typeof value === 'string' && FALAI_COST_SOURCES.has(value) +} + interface FalAIBillingEvent { request_id: string endpoint_id: string diff --git a/apps/sim/lib/tools/falai-settlement.integration.ts b/apps/sim/lib/tools/falai-settlement.integration.ts new file mode 100644 index 00000000000..1d0154df588 --- /dev/null +++ b/apps/sim/lib/tools/falai-settlement.integration.ts @@ -0,0 +1,175 @@ +/** + * The fal exact-cost settlement against a real usage ledger: a shortfall is billed once however + * often the settlement runs, and nothing is billed until fal records the charge. Only fal's + * billing-events API is stubbed. + */ +import { db } from '@sim/db' +import { organization, usageLog, user, userStats } from '@sim/db/schema' +import { generateId } from '@sim/utils/id' +import { eq, inArray } from 'drizzle-orm' +import { afterAll, afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { type FalAICostSettlementPayload, settleFalAICost } from '@/lib/tools/falai-settlement' + +const HOSTED_KEY_ENV_VAR = 'FALAI_API_KEY_1' +const userIds: string[] = [] +const organizationIds: string[] = [] + +async function createUser(): Promise { + const id = `falai-settlement-user-${generateId()}` + const now = new Date() + await db.insert(user).values({ + id, + name: 'Fal Settlement Test', + email: `${id}@falai-settlement.test`, + emailVerified: true, + createdAt: now, + updatedAt: now, + }) + await db.insert(userStats).values({ id: generateId(), userId: id }) + userIds.push(id) + return id +} + +async function createOrganization(): Promise { + const id = `falai-settlement-org-${generateId()}` + await db.insert(organization).values({ id, name: 'Fal Settlement Org', slug: id }) + organizationIds.push(id) + return id +} + +/** A settlement billed to an organization payer, the snapshot shape for a run outside a workspace. */ +function payloadFor( + userId: string, + organizationId: string, + billedProviderCostDollars: number, + billingPeriod = { start: '2026-11-01T00:00:00.000Z', end: '2026-12-01T00:00:00.000Z' } +): FalAICostSettlementPayload { + return { + envVarName: HOSTED_KEY_ENV_VAR, + endpointId: 'bytedance/seedance-2.5/text-to-video', + requestId: `fal-request-${generateId()}`, + billedProviderCostDollars, + markupMultiplier: 1.5, + ledger: { + source: 'workflow', + billingAttribution: { + actorUserId: userId, + workspaceId: null, + organizationId, + billedAccountUserId: userId, + billingEntity: { type: 'organization', id: organizationId }, + billingPeriod, + payerSubscription: null, + }, + }, + } +} + +/** fal's billing-events API answering with the charge for `requestId`, or with none yet. */ +function stubFalBillingEvents(charge: { requestId: string; costDollars: number } | null) { + const fetchMock = vi.fn(async () => + Response.json({ + billing_events: charge + ? [ + { + request_id: charge.requestId, + endpoint_id: 'bytedance/seedance-2.5/text-to-video', + output_units: 700, + unit_price: 0.0214, + percent_discount: null, + cost_estimate_nano_usd: Math.round(charge.costDollars * 1_000_000_000), + }, + ] + : [], + }) + ) + vi.stubGlobal('fetch', fetchMock) + return fetchMock +} + +async function settlementRows(userId: string) { + return db + .select({ + cost: usageLog.cost, + description: usageLog.description, + source: usageLog.source, + billingPeriodEnd: usageLog.billingPeriodEnd, + }) + .from(usageLog) + .where(eq(usageLog.userId, userId)) +} + +beforeEach(() => { + vi.stubEnv(HOSTED_KEY_ENV_VAR, 'hosted-fixture-key') +}) + +afterEach(() => { + vi.unstubAllGlobals() + vi.unstubAllEnvs() +}) + +afterAll(async () => { + if (userIds.length > 0) { + await db.delete(usageLog).where(inArray(usageLog.userId, userIds)) + await db.delete(user).where(inArray(user.id, userIds)) + } + if (organizationIds.length > 0) { + await db.delete(organization).where(inArray(organization.id, organizationIds)) + } +}) + +describe('settleFalAICost', () => { + it('bills the shortfall once when the exact charge exceeds the estimate, however often it runs', async () => { + const userId = await createUser() + const organizationId = await createOrganization() + const payload = payloadFor(userId, organizationId, 6.49) + stubFalBillingEvents({ requestId: payload.requestId, costDollars: 15 }) + + expect(await settleFalAICost(payload)).toBe('charged') + expect(await settleFalAICost(payload)).toBe('charged') + + const rows = await settlementRows(userId) + expect(rows).toHaveLength(1) + expect(Number(rows[0].cost)).toBeCloseTo((15 - 6.49) * 1.5, 6) + expect(rows[0]).toMatchObject({ + source: 'workflow', + description: 'fal exact cost settlement: bytedance/seedance-2.5/text-to-video', + }) + }) + + it('bills a settlement that outlives its billing period to the payer current period', async () => { + const userId = await createUser() + const organizationId = await createOrganization() + const endedPeriod = { start: '2026-01-01T00:00:00.000Z', end: '2026-02-01T00:00:00.000Z' } + const payload = payloadFor(userId, organizationId, 6.49, endedPeriod) + stubFalBillingEvents({ requestId: payload.requestId, costDollars: 15 }) + + expect(await settleFalAICost(payload)).toBe('charged') + + const rows = await settlementRows(userId) + expect(rows).toHaveLength(1) + expect(rows[0].billingPeriodEnd?.getTime()).toBeGreaterThan(Date.now()) + }) + + it('bills nothing while fal has not recorded the charge', async () => { + const userId = await createUser() + const organizationId = await createOrganization() + stubFalBillingEvents(null) + + expect(await settleFalAICost(payloadFor(userId, organizationId, 0.25))).toBe('pending') + + expect(await settlementRows(userId)).toHaveLength(0) + }) + + it('refuses to settle when the hosted key that ran the request is not configured', async () => { + const userId = await createUser() + const organizationId = await createOrganization() + vi.stubEnv(HOSTED_KEY_ENV_VAR, '') + const fetchMock = stubFalBillingEvents(null) + + await expect(settleFalAICost(payloadFor(userId, organizationId, 6.49))).rejects.toThrow( + 'Hosted fal key FALAI_API_KEY_1 is not configured' + ) + expect(fetchMock).not.toHaveBeenCalled() + }) +}) diff --git a/apps/sim/lib/tools/falai-settlement.ts b/apps/sim/lib/tools/falai-settlement.ts new file mode 100644 index 00000000000..819e763c04c --- /dev/null +++ b/apps/sim/lib/tools/falai-settlement.ts @@ -0,0 +1,176 @@ +import { createLogger } from '@sim/logger' +import { getErrorMessage } from '@sim/utils/errors' +import { tasks } from '@trigger.dev/sdk' +import { + type BillingAttributionSnapshot, + refreshAttributionPeriod, +} from '@/lib/billing/core/attribution' +import { recordUsage, stableEventKey } from '@/lib/billing/core/usage-log' +import { resolveTriggerRegion } from '@/lib/core/async-jobs/region' +import { env } from '@/lib/core/config/env' +import { isBillingEnabled, isHosted, isTriggerDevEnabled } from '@/lib/core/config/env-flags' +import { type FalAICostMetadata, getFalAIBillingReceipt } from '@/lib/tools/falai-pricing' + +const logger = createLogger('FalAICostSettlement') + +export const FALAI_COST_SETTLEMENT_TASK_ID = 'falai-cost-settlement' as const + +/** fal records a request's billing event after it completes; the first look waits this long. */ +const FALAI_SETTLEMENT_DELAY = '2m' + +/** Mothership bills fal media at cost, so a Chat generation settles without markup. */ +export const MOTHERSHIP_FAL_AT_COST_MULTIPLIER = 1 + +/** Ledger sources a hosted fal charge can be billed under. */ +type FalAISettlementSource = 'workflow' | 'copilot' | 'api-tool' + +/** + * Everything a settlement needs once the generation's run is gone. The hosted key travels as + * its env var name, never its value, so the task reads the same key that created the request. + */ +export interface FalAICostSettlementPayload { + envVarName: string + endpointId: string + requestId: string + /** The provider cost billed at completion, before markup. */ + billedProviderCostDollars: number + markupMultiplier: number + ledger: { + source: FalAISettlementSource + workflowId?: string + executionId?: string + /** The run's payer as admitted; its period is refreshed when settlement outlives it. */ + billingAttribution: BillingAttributionSnapshot + } +} + +/** The outcome of one settlement attempt. `pending` means fal has not recorded the charge yet. */ +type FalAICostSettlementOutcome = 'pending' | 'charged' | 'exact' | 'overbilled' + +interface ScheduleFalAICostSettlementParams { + cost: Pick + envVarName: string + markupMultiplier: number + source: FalAISettlementSource + /** The run's payer; a run without one is billed on its estimate. */ + billingAttribution: BillingAttributionSnapshot | undefined + workflowId?: string + executionId?: string +} + +/** + * Schedules the exact-cost settlement for a hosted fal generation billed on an estimate. A + * generation billed from its own receipt is already exact and schedules nothing. Settlement only + * adds what the estimate missed, so a failure here leaves the estimate billed and never fails + * the generation. + */ +export async function scheduleFalAICostSettlement( + params: ScheduleFalAICostSettlementParams +): Promise { + const { cost, billingAttribution } = params + if (!isHosted || !isBillingEnabled) return + if (cost.source === 'billing_events' || !billingAttribution) return + if (!isTriggerDevEnabled || !env.TRIGGER_SECRET_KEY) { + logger.warn('Trigger.dev not configured; fal charge stays on its estimate', { + endpointId: cost.endpointId, + requestId: cost.requestId, + }) + return + } + + const payload: FalAICostSettlementPayload = { + envVarName: params.envVarName, + endpointId: cost.endpointId, + requestId: cost.requestId, + billedProviderCostDollars: cost.costDollars, + markupMultiplier: params.markupMultiplier, + ledger: { + source: params.source, + workflowId: params.workflowId, + executionId: params.executionId, + billingAttribution, + }, + } + + try { + await tasks.trigger(FALAI_COST_SETTLEMENT_TASK_ID, payload, { + delay: FALAI_SETTLEMENT_DELAY, + idempotencyKey: `falai-cost-settlement-${cost.requestId}`, + region: await resolveTriggerRegion(), + }) + } catch (error) { + logger.warn('Failed to schedule the fal exact-cost settlement; the estimate stays billed', { + endpointId: cost.endpointId, + requestId: cost.requestId, + error: getErrorMessage(error), + }) + } +} + +/** + * Bills the difference between fal's exact charge for a request and the estimate billed at + * completion. The ledger row's event key is the request, so a retried settlement never bills + * twice. An estimate above the exact charge is left as billed: the ledger records charges only. + * A settlement that outlives the run's billing period bills the payer's current period, since + * the ended one may already be closed and invoiced. + */ +export async function settleFalAICost( + payload: FalAICostSettlementPayload +): Promise { + const { envVarName, endpointId, requestId, billedProviderCostDollars, markupMultiplier, ledger } = + payload + const apiKey = process.env[envVarName] + if (!apiKey) throw new Error(`Hosted fal key ${envVarName} is not configured`) + + const receipt = await getFalAIBillingReceipt({ apiKey, endpointId, requestId }) + if (!receipt) return 'pending' + + const shortfallDollars = (receipt.costDollars - billedProviderCostDollars) * markupMultiplier + if (shortfallDollars < 0) { + logger.info('fal estimate exceeded the exact charge', { + endpointId, + requestId, + overbilledDollars: -shortfallDollars, + }) + return 'overbilled' + } + if (shortfallDollars === 0) return 'exact' + + const attribution = + Date.parse(ledger.billingAttribution.billingPeriod.end) <= Date.now() + ? await refreshAttributionPeriod(ledger.billingAttribution) + : ledger.billingAttribution + await recordUsage({ + userId: attribution.actorUserId, + workspaceId: attribution.workspaceId ?? undefined, + workflowId: ledger.workflowId, + executionId: ledger.executionId, + billingEntity: { ...attribution.billingEntity }, + billingPeriod: { + start: new Date(attribution.billingPeriod.start), + end: new Date(attribution.billingPeriod.end), + }, + entries: [ + { + category: 'tool', + source: ledger.source, + description: `fal exact cost settlement: ${endpointId}`, + cost: shortfallDollars, + eventKey: stableEventKey({ kind: 'falai-cost-settlement', requestId }), + metadata: { + endpointId, + requestId, + billedProviderCostDollars, + exactProviderCostDollars: receipt.costDollars, + markupMultiplier, + }, + }, + ], + }) + logger.info('Settled fal charge to its exact cost', { + endpointId, + requestId, + shortfallDollars, + }) + return 'charged' +} diff --git a/apps/sim/tools/image/generate.ts b/apps/sim/tools/image/generate.ts index 0c9257c6e51..7226e12b41a 100644 --- a/apps/sim/tools/image/generate.ts +++ b/apps/sim/tools/image/generate.ts @@ -127,6 +127,7 @@ export const imageGenerateTool: ImageGenerationToolConfig = { envKeyPrefix: 'FALAI_API_KEY', apiKeyParam: 'apiKey', byokProviderId: 'falai', + exactCostSettlement: 'falai', pricing: { type: 'custom', getCost: (_params, output) => { diff --git a/apps/sim/tools/index.test.ts b/apps/sim/tools/index.test.ts index fd6f413a055..b4495864fdf 100644 --- a/apps/sim/tools/index.test.ts +++ b/apps/sim/tools/index.test.ts @@ -116,8 +116,10 @@ const { mockGetInternalToolOperationHandler, mockExecuteInternalToolOperation, mockPayerHasPaidAccess, + mockScheduleFalAICostSettlement, } = vi.hoisted(() => ({ mockPayerHasPaidAccess: vi.fn(), + mockScheduleFalAICostSettlement: vi.fn(), mockGetToolAsync: vi.fn(), mockRateLimiterFns: { acquireKey: vi.fn(), @@ -183,6 +185,10 @@ vi.mock('@/lib/billing/core/chat-access', () => ({ payerHasPaidAccess: mockPayerHasPaidAccess, })) +vi.mock('@/lib/tools/falai-settlement', () => ({ + scheduleFalAICostSettlement: mockScheduleFalAICostSettlement, +})) + vi.mock('@/lib/billing/core/attribution', () => billingAttributionMock) vi.mock('@/lib/uploads/contexts/workspace/file-manager', () => workspaceFileManagerMock) @@ -6564,6 +6570,74 @@ describe('Hosted Key Injection', () => { } }) + it('schedules the exact-cost settlement for a hosted fal charge estimated at zero', async () => { + setEnvFlags({ isHosted: true }) + mockGetBYOKKey.mockResolvedValue(null) + mockRateLimiterFns.acquireKey.mockResolvedValue({ + success: true, + key: 'platform-fixture-key', + keyIndex: 0, + envVarName: 'FALAI_API_KEY_1', + }) + mockValidateUrlWithDNS.mockResolvedValue({ isValid: true, resolvedIP: '93.184.216.34' }) + mockSecureFetchWithPinnedIP.mockResolvedValue({ + ok: true, + status: 200, + statusText: 'OK', + headers: { get: () => null, toRecord: () => ({}) }, + json: async () => ({ ok: true }), + text: async () => '{"ok":true}', + }) + tools.test_fal_zero_estimate = { + id: 'test_fal_zero_estimate', + name: 'Zero-estimate fal fixture', + description: 'Hosted fal charge whose estimate is zero', + version: '1.0.0', + params: {}, + request: { url: 'https://example.com/fixture', method: 'GET', headers: () => ({}) }, + hosting: { + envKeyPrefix: 'FALAI_API_KEY', + apiKeyParam: 'apiKey', + byokProviderId: 'falai', + exactCostSettlement: 'falai', + pricing: { + type: 'custom', + getCost: () => ({ + cost: 0, + metadata: { + endpointId: 'fal-ai/fixture/text-to-video', + requestId: 'fal-request-1', + source: 'historical_estimate', + providerCostDollars: 0, + markupMultiplier: 1.5, + }, + }), + }, + rateLimit: { mode: 'per_request', requestsPerMinute: 100 }, + }, + transformResponse: async () => ({ success: true, output: {} }), + } + try { + const result = await executeTool( + 'test_fal_zero_estimate', + {}, + { executionContext: createToolExecutionContext({ userId: 'user-1' }) } + ) + + expect(result.success).toBe(true) + expect(mockScheduleFalAICostSettlement).toHaveBeenCalledWith( + expect.objectContaining({ + cost: expect.objectContaining({ requestId: 'fal-request-1', costDollars: 0 }), + envVarName: 'FALAI_API_KEY_1', + source: 'workflow', + }) + ) + } finally { + Reflect.deleteProperty(tools, 'test_fal_zero_estimate') + resetEnvFlagsMock() + } + }) + describe('requiresPaidPlan', () => { const PAID_ONLY_TOOL_ID = 'test_paid_only_hosting' diff --git a/apps/sim/tools/index.ts b/apps/sim/tools/index.ts index 05dff63008e..cfc9708c3cc 100644 --- a/apps/sim/tools/index.ts +++ b/apps/sim/tools/index.ts @@ -1,5 +1,6 @@ import { createLogger } from '@sim/logger' import { isLoopbackIp, unwrapIpv6Brackets } from '@sim/security/ssrf' +import { toNumberOrNull, toStringOrNull } from '@sim/utils/coerce' import { describeError, getErrorMessage, toError } from '@sim/utils/errors' import { sleep } from '@sim/utils/helpers' import { isPlainRecord, isRecordLike, toRecord } from '@sim/utils/object' @@ -99,6 +100,8 @@ import { } from '@/lib/mothership/billing/service-observer' import { CredentialRevokedError } from '@/lib/oauth/credential-revoked' import type { CredentialTokenPayload } from '@/lib/oauth/token-resolution' +import { isFalAICostSource } from '@/lib/tools/falai-pricing' +import { scheduleFalAICostSettlement } from '@/lib/tools/falai-settlement' import { resolveWorkspaceFileReference } from '@/lib/uploads/contexts/workspace/file-manager' import { markWorkspaceFileSecretProvenanceUnknown } from '@/lib/uploads/contexts/workspace/file-secret-provenance' import { MAX_FILE_SIZE } from '@/lib/uploads/utils/validation' @@ -954,7 +957,7 @@ async function processHostedKeyCost( const { cost, metadata } = calculateToolCost(tool.hosting.pricing, params, response, tool.id) - if (cost <= 0) return { cost: 0 } + if (cost <= 0) return { cost: 0, metadata } logger.debug(`Hosted key cost for ${tool.id}: $${cost}`, metadata ? { metadata } : {}, { requestId, @@ -1028,6 +1031,49 @@ export function postProcessToolOutput(toolId: string, output: Record | undefined, + scope: ToolExecutionScope, + envVarName: string | undefined +): Promise { + if (tool.hosting?.exactCostSettlement !== 'falai' || !envVarName) return + const endpointId = toStringOrNull(metadata?.endpointId) + const falRequestId = toStringOrNull(metadata?.requestId) + const costSource = metadata?.source + const providerCostDollars = toNumberOrNull(metadata?.providerCostDollars) + const markupMultiplier = toNumberOrNull(metadata?.markupMultiplier) + if ( + !endpointId || + !falRequestId || + !isFalAICostSource(costSource) || + providerCostDollars === null || + markupMultiplier === null + ) { + return + } + + const { billingAttribution, copilotToolExecution, workflowId, executionId } = scope + await scheduleFalAICostSettlement({ + cost: { + endpointId, + requestId: falRequestId, + costDollars: providerCostDollars, + source: costSource, + }, + envVarName, + markupMultiplier, + source: copilotToolExecution ? 'copilot' : executionId ? 'workflow' : 'api-tool', + billingAttribution, + workflowId, + executionId, + }) +} + /** * Apply post-execution hosted-key cost tracking to a successful tool result. * Reports custom dimension usage, calculates cost, and merges it into the output. @@ -1080,8 +1126,12 @@ async function applyHostedKeyCostToResult( hostedKeyMetrics.recordUsed({ provider, tool: tool.id, key }) hostedKeyMetrics.recordCostCharged(hostedKeyCost, { provider, tool: tool.id }) + const scope = resolveToolScope(params, executionContext) + // A zero estimate is not proof the generation was free; settlement bills the receipt. + await scheduleExactCostSettlement(tool, metadata, scope, envVarName) + if (hostedKeyCost > 0) { - const { copilotToolExecution } = resolveToolScope(params, executionContext) + const { copilotToolExecution } = scope if (copilotToolExecution) await recordServiceCost(provider, hostedKeyCost) finalResult.output = { ...finalResult.output, diff --git a/apps/sim/tools/types.ts b/apps/sim/tools/types.ts index 17f2bff890b..c013e7b6121 100644 --- a/apps/sim/tools/types.ts +++ b/apps/sim/tools/types.ts @@ -505,6 +505,11 @@ export interface ToolHostingConfig

> { * or a key the user passes still runs; only Sim's own key is refused. */ requiresPaidPlan?: boolean + /** + * Bills the hosted charge again from the provider's exact receipt once it is recorded, for a + * provider whose receipt arrives after the result. `falai` settles through fal's billing events. + */ + exactCostSettlement?: 'falai' /** Pricing when using hosted key */ pricing: ToolHostingPricing

/** Hosted key rate limit configuration (required for hosted key distribution) */ diff --git a/apps/sim/tools/video/falai.ts b/apps/sim/tools/video/falai.ts index 1c8907eb184..be34f58fd0f 100644 --- a/apps/sim/tools/video/falai.ts +++ b/apps/sim/tools/video/falai.ts @@ -71,6 +71,7 @@ export const falaiVideoTool: InternalToolConfig = { envKeyPrefix: 'FALAI_API_KEY', apiKeyParam: 'apiKey', byokProviderId: 'falai', + exactCostSettlement: 'falai', requiresPaidPlan: true, pricing: { type: 'custom',