diff --git a/apps/sim/lib/media/falai-audio.ts b/apps/sim/lib/media/falai-audio.ts index da2e0a55233..f55b27d633d 100644 --- a/apps/sim/lib/media/falai-audio.ts +++ b/apps/sim/lib/media/falai-audio.ts @@ -1,5 +1,12 @@ -import { downloadFalMedia, extractFalMediaUrl, getFalApiKey, runFalQueue } from '@/lib/media/falai' -import { type FalAICostMetadata, getFalAICostMetadata } from '@/lib/tools/falai-pricing' +import { + downloadFalMedia, + extractFalMediaUrl, + type FalCostRecorder, + getFalApiKey, + recordFalCost, + runFalQueue, +} from '@/lib/media/falai' +import type { FalAICostMetadata } from '@/lib/tools/falai-pricing' export type AudioType = 'speech' | 'music' | 'sfx' @@ -28,6 +35,8 @@ export interface GenerateFalAudioParams { voiceSampleDataUri?: string /** Ends the generation: the queued request is cancelled and every wait stops. */ signal?: AbortSignal + /** Bills the generation once fal completes it, even if a Stop then ends the download. */ + onCost: FalCostRecorder } export interface GeneratedAudio { @@ -97,15 +106,10 @@ export async function generateFalAudio(params: GenerateFalAudioParams): Promise< model_type: 'F5-TTS', } const { requestId, data } = await runFalQueue(model, input, apiKey, params.signal) + const cost = await recordFalCost(apiKey, 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) - const cost = await getFalAICostMetadata({ - apiKey, - endpointId: model, - requestId, - signal: params.signal, - }) return { buffer, contentType: contentType.startsWith('audio/') ? contentType : 'audio/mpeg', @@ -121,16 +125,11 @@ export async function generateFalAudio(params: GenerateFalAudioParams): Promise< // For fal audio models the model ID is the queue endpoint. const { requestId, data } = await runFalQueue(model, input, apiKey, params.signal) + const cost = await recordFalCost(apiKey, 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') const { buffer, contentType } = await downloadFalMedia(url, params.signal) - const cost = await getFalAICostMetadata({ - apiKey, - endpointId: model, - requestId, - signal: params.signal, - }) return { buffer, diff --git a/apps/sim/lib/media/falai-video.ts b/apps/sim/lib/media/falai-video.ts index 9e749af623a..cdea37ea4cc 100644 --- a/apps/sim/lib/media/falai-video.ts +++ b/apps/sim/lib/media/falai-video.ts @@ -2,11 +2,13 @@ import { isRecordLike } from '@sim/utils/object' import { downloadFalMedia, extractFalMediaUrl, + type FalCostRecorder, getFalApiKey, getNumberProp, + recordFalCost, runFalQueue, } from '@/lib/media/falai' -import { type FalAICostMetadata, getFalAICostMetadata } from '@/lib/tools/falai-pricing' +import type { FalAICostMetadata } from '@/lib/tools/falai-pricing' type DurationFormat = 'number' | 'seconds' | 'string' @@ -112,6 +114,8 @@ export interface GenerateFalVideoParams { imageDataUri?: string /** Ends the generation: the queued request is cancelled and every wait stops. */ signal?: AbortSignal + /** Bills the generation once fal completes it, even if a Stop then ends the download. */ + onCost: FalCostRecorder } export interface GeneratedVideo { @@ -174,17 +178,12 @@ export async function generateFalVideo(params: GenerateFalVideoParams): Promise< } const { requestId, data } = await runFalQueue(endpoint, input, apiKey, params.signal) + const cost = await recordFalCost(apiKey, endpoint, requestId, params.onCost) const url = extractFalMediaUrl(data, ['video', 'output']) if (!url) throw new Error('No video URL in Fal.ai response') const videoNode = isRecordLike(data.video) ? data.video : undefined const { buffer, contentType } = await downloadFalMedia(url, params.signal) - const cost = await getFalAICostMetadata({ - apiKey, - endpointId: endpoint, - requestId, - signal: params.signal, - }) return { buffer, diff --git a/apps/sim/lib/media/falai.test.ts b/apps/sim/lib/media/falai.test.ts index acd9d4b8a90..5407736b4d0 100644 --- a/apps/sim/lib/media/falai.test.ts +++ b/apps/sim/lib/media/falai.test.ts @@ -1,3 +1,4 @@ +import { createDeferred } from '@sim/testing/helpers/deferred' import { jsonResponse } from '@sim/testing/helpers/http' import { inputValidationMock, @@ -8,7 +9,9 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' vi.mock('@/lib/core/security/input-validation.server', () => inputValidationMock) -import { downloadFalMedia, runFalQueue } from '@/lib/media/falai' +import { downloadFalMedia, type FalCostRecorder, runFalQueue } from '@/lib/media/falai' +import { generateFalAudio } from '@/lib/media/falai-audio' +import { generateFalVideo } from '@/lib/media/falai-video' 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' @@ -210,3 +213,189 @@ describe('downloadFalMedia', () => { await expect(download).rejects.toThrow('watchdog') }) }) + +const MEDIA_URL = 'https://media.example/out.bin' +const BILLED_NANO_USD = 40_000_000 + +/** + * A fal queue whose request completes on the first poll. The billing-event lookup answers at + * once, answers only once `answered` settles (whatever its caller does), or (`hangs`) never. + */ +function stubCompletedFalQueue( + billing: { answered?: Promise; hangs?: boolean } = {}, + result: Record = { audio: { url: MEDIA_URL }, video: { url: MEDIA_URL } } +) { + const billingLookups: string[] = [] + vi.stubGlobal( + 'fetch', + vi.fn(async (input: string | URL, init: RequestInit = {}) => { + const url = String(input) + if (url === STATUS_URL) return jsonResponse({ status: 'COMPLETED' }) + if (url.endsWith('/requests/req-1/response')) return jsonResponse(result) + if (!url.startsWith('https://api.fal.ai/v1/models/billing-events')) + return jsonResponse({ request_id: 'req-1', status_url: STATUS_URL }) + billingLookups.push(url) + if (billing.hangs) + return new Promise((_resolve, reject) => { + init.signal?.addEventListener('abort', () => reject(init.signal?.reason), { once: true }) + }) + await billing.answered + return jsonResponse({ + billing_events: [ + { + request_id: 'req-1', + endpoint_id: 'fal-ai/model', + cost_estimate_nano_usd: BILLED_NANO_USD, + }, + ], + }) + }) + ) + return { billingLookups } +} + +/** The generated media's host resolves; the download then answers with `download`. */ +function stubMediaHost(download: (options: { signal?: AbortSignal }) => Promise) { + inputValidationMockFns.mockValidateUrlWithDNS.mockResolvedValueOnce({ + isValid: true, + resolvedIP: '203.0.113.1', + }) + inputValidationMockFns.mockSecureFetchWithPinnedIP.mockImplementationOnce( + (_url: string, _ip: string, options: { signal?: AbortSignal }) => download(options) + ) +} + +const downloadAudio = async () => + new Response('audio', { headers: { 'content-type': 'audio/mpeg' } }) + +const startAudio = (onCost: FalCostRecorder, signal?: AbortSignal) => + generateFalAudio({ prompt: 'p', model: 'fal-ai/model', signal, onCost }) + +const generators = [ + ['audio', startAudio], + [ + 'video', + (onCost: FalCostRecorder, signal?: AbortSignal) => + generateFalVideo({ prompt: 'p', signal, onCost }), + ], +] as const + +describe('billing a fal generation', () => { + beforeEach(() => { + vi.useFakeTimers() + vi.stubEnv('FALAI_API_KEY', 'key') + }) + afterEach(() => { + vi.useRealTimers() + inputValidationMockFns.mockValidateUrlWithDNS.mockReset() + inputValidationMockFns.mockSecureFetchWithPinnedIP.mockReset() + }) + + it.each(generators)( + 'bills a completed %s generation once when a Stop ends its download', + async (_kind, generate) => { + stubCompletedFalQueue() + stubMediaHost( + ({ signal }) => + new Promise((_resolve, reject) => { + signal?.addEventListener('abort', () => reject(signal.reason), { once: true }) + }) + ) + const onCost = vi.fn().mockResolvedValue(undefined) + const caller = new AbortController() + const generation = generate(onCost, caller.signal).catch((error: unknown) => error) + + await vi.advanceTimersByTimeAsync(3_000) + expect(inputValidationMockFns.mockSecureFetchWithPinnedIP).toHaveBeenCalled() + caller.abort(new Error('watchdog')) + + expect(await generation).toEqual(new Error('watchdog')) + expect(onCost).toHaveBeenCalledTimes(1) + expect(onCost).toHaveBeenCalledWith( + expect.objectContaining({ source: 'billing_events', costDollars: 0.04 }) + ) + } + ) + + it.each(generators)( + 'bills a completed %s generation once when its result has no media URL', + async (_kind, generate) => { + stubCompletedFalQueue({}, { unexpected: true }) + const onCost = vi.fn().mockResolvedValue(undefined) + const generation = generate(onCost).catch((error: unknown) => String(error)) + + await vi.advanceTimersByTimeAsync(3_000) + + expect(await generation).toMatch(/No (audio|video) URL/) + expect(onCost).toHaveBeenCalledTimes(1) + expect(onCost).toHaveBeenCalledWith(expect.objectContaining({ source: 'billing_events' })) + } + ) + + it('finishes the cost lookup of a completed generation a Stop lands in, then stops', async () => { + const billingAnswer = createDeferred() + const fal = stubCompletedFalQueue({ answered: billingAnswer.promise }) + stubMediaHost(downloadAudio) + const onCost = vi.fn().mockResolvedValue(undefined) + const caller = new AbortController() + const generation = startAudio(onCost, caller.signal).catch((error: unknown) => error) + + await vi.advanceTimersByTimeAsync(3_000) + expect(fal.billingLookups).toHaveLength(1) + caller.abort(new Error('watchdog')) + billingAnswer.resolve() + + expect(await generation).toEqual(new Error('watchdog')) + expect(onCost).toHaveBeenCalledTimes(1) + expect(onCost).toHaveBeenCalledWith( + expect.objectContaining({ source: 'billing_events', costDollars: 0.04 }) + ) + expect(inputValidationMockFns.mockSecureFetchWithPinnedIP).not.toHaveBeenCalled() + }) + + it('bills the fallback floor once when the cost lookup outlasts its own bound', async () => { + stubCompletedFalQueue({ hangs: true }) + stubMediaHost(downloadAudio) + const onCost = vi.fn().mockResolvedValue(undefined) + const caller = new AbortController() + const generation = startAudio(onCost, caller.signal).catch((error: unknown) => error) + + await vi.advanceTimersByTimeAsync(3_000) + caller.abort(new Error('watchdog')) + await vi.advanceTimersByTimeAsync(9_999) + expect(onCost).not.toHaveBeenCalled() + await vi.advanceTimersByTimeAsync(1) + + expect(await generation).toEqual(new Error('watchdog')) + expect(onCost).toHaveBeenCalledTimes(1) + expect(onCost).toHaveBeenCalledWith(expect.objectContaining({ source: 'fallback_floor' })) + }) + + it('bills a generation that completes and downloads exactly once', async () => { + stubCompletedFalQueue() + stubMediaHost(downloadAudio) + const onCost = vi.fn().mockResolvedValue(undefined) + const generation = startAudio(onCost) + + await vi.advanceTimersByTimeAsync(3_000) + const audio = await generation + + expect(audio.buffer.toString()).toBe('audio') + expect(onCost).toHaveBeenCalledTimes(1) + expect(onCost).toHaveBeenCalledWith(audio.cost) + }) + + it('bills nothing for a generation stopped before fal completes it', async () => { + const fal = stubStuckFalQueue() + const onCost = vi.fn().mockResolvedValue(undefined) + const caller = new AbortController() + const generation = startAudio(onCost, caller.signal).catch((error: unknown) => error) + + await vi.advanceTimersByTimeAsync(3_500) + caller.abort(new Error('watchdog')) + + expect(await generation).toEqual(new Error('watchdog')) + expect(fal.cancels()).toHaveLength(1) + expect(onCost).not.toHaveBeenCalled() + }) +}) diff --git a/apps/sim/lib/media/falai.ts b/apps/sim/lib/media/falai.ts index 9b1f2c0e7da..40db1bd5f6e 100644 --- a/apps/sim/lib/media/falai.ts +++ b/apps/sim/lib/media/falai.ts @@ -13,6 +13,7 @@ import { readResponseTextWithLimit, readResponseToBufferWithLimit, } from '@/lib/core/utils/stream-limits' +import { type FalAICostMetadata, getFalAICostMetadata } from '@/lib/tools/falai-pricing' const logger = createLogger('FalMediaClient') @@ -252,6 +253,25 @@ 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 + +/** + * Bills a request fal has completed. fal charges for it whatever happens next, so the lookup + * ignores the caller's Stop (it has its own bound) and runs before anything that can still fail + * the generation on our side: reading the media URL and the download a Stop can end. + */ +export async function recordFalCost( + apiKey: string, + endpointId: string, + requestId: string, + onCost: FalCostRecorder +): Promise { + const cost = await getFalAICostMetadata({ apiKey, endpointId, requestId }) + await onCost(cost) + return cost +} + /** * Pull the output media URL out of a Fal.ai result, tolerating the various * shapes different models return (string url, { url }, nested arrays). 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 5c9b87bfd35..2a2d03dd8cc 100644 --- a/apps/sim/lib/mothership/tools/server/media/generate-audio.ts +++ b/apps/sim/lib/mothership/tools/server/media/generate-audio.ts @@ -131,10 +131,9 @@ export const generateAudioServerTool: BaseServerTool recordServiceCost('falai_audio', cost.costDollars), }) - await recordServiceCost('falai_audio', result.cost.costDollars) - const outputFile = params.outputs?.files?.[0] const ext = audioExtFromContentType(result.contentType) const outputPath = outputFile?.path || `files/generated-audio.${ext}` 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 5b4114772af..e563fa702b6 100644 --- a/apps/sim/lib/mothership/tools/server/media/generate-video.ts +++ b/apps/sim/lib/mothership/tools/server/media/generate-video.ts @@ -108,10 +108,9 @@ export const generateVideoServerTool: BaseServerTool recordServiceCost('falai_video', cost.costDollars), }) - await recordServiceCost('falai_video', result.cost.costDollars) - const outputFile = params.outputs?.files?.[0] const outputPath = outputFile?.path || 'files/generated-video.mp4' const mode = outputFile?.mode ?? 'create' 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 9a678a5da88..96f62494678 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 @@ -36,6 +36,20 @@ import { generateAudioServerTool } from '@/lib/mothership/tools/server/media/gen import { generateVideoServerTool } from '@/lib/mothership/tools/server/media/generate-video' import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' +interface FalGenerationParams { + signal?: AbortSignal + onCost: (cost: { costDollars: number }) => Promise +} + +/** 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) + return { ...media, cost } + } +} + function contextWithSecrets( entries: Array<{ name: string; plaintext: string }> ): ServerToolContext { @@ -67,19 +81,17 @@ describe('Mothership media model boundaries', () => { { content: { parts: [{ inlineData: { data: 'aW1hZ2U=', mimeType: 'image/png' } }] } }, ], }) - mockGenerateFalVideo.mockResolvedValue({ - buffer: Buffer.from('video'), - contentType: 'video/mp4', - model: 'veo-3.1-fast', - cost: { costDollars: 0.1 }, - }) - mockGenerateFalAudio.mockResolvedValue({ - buffer: Buffer.from('audio'), - contentType: 'audio/mpeg', - type: 'music', - model: 'music-model', - cost: { costDollars: 0.1 }, - }) + mockGenerateFalVideo.mockImplementation( + billsThen({ buffer: Buffer.from('video'), contentType: 'video/mp4', model: 'veo-3.1-fast' }) + ) + mockGenerateFalAudio.mockImplementation( + billsThen({ + buffer: Buffer.from('audio'), + contentType: 'audio/mpeg', + type: 'music', + model: 'music-model', + }) + ) }) it('meters image, video and audio at their provider boundary before file persistence', async () => { @@ -121,6 +133,27 @@ describe('Mothership media model boundaries', () => { } ) + it.each([ + ['audio', 'falai_audio', mockGenerateFalAudio, generateAudioServerTool], + ['video', 'falai_video', mockGenerateFalVideo, generateVideoServerTool], + ] as const)( + '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 }) + throw new Error('Tool timed out') + }) + const charges = vi.fn().mockResolvedValue(undefined) + + const result = await observeServiceCosts(charges, () => + tool.execute({ prompt: 'p' }, contextWithSecrets([])) + ) + + expect(result).toMatchObject({ success: false }) + expect(charges.mock.calls).toEqual([[service, 0.2]]) + } + ) + it('preserves image prompts that merely collide with ambient secret plaintext', async () => { const context = contextWithSecrets([{ name: 'PROMPT', plaintext: 'private prompt' }]) diff --git a/apps/sim/lib/tools/falai-pricing.test.ts b/apps/sim/lib/tools/falai-pricing.test.ts index 1cb90c7fd8d..65d8c3187df 100644 --- a/apps/sim/lib/tools/falai-pricing.test.ts +++ b/apps/sim/lib/tools/falai-pricing.test.ts @@ -47,6 +47,26 @@ describe('getFalAICostMetadata fallback floor', () => { expect(result.costDollars).toBe(expected) }) + it('bills the floor when the pricing estimate answers with an unreadable body', async () => { + vi.stubGlobal( + 'fetch', + vi.fn(async (url: string | URL) => + String(url).includes('pricing/estimate') + ? new Response('not json', { status: 200 }) + : new Response('error', { status: 500 }) + ) + ) + + const result = await getFalAICostMetadata({ + apiKey: 'test-key', + endpointId: 'fal-ai/veo-3.1', + requestId: 'req_123', + }) + + expect(result.source).toBe('fallback_floor') + expect(result.costDollars).toBe(FALAI_VIDEO_FALLBACK_PROVIDER_COST_DOLLARS) + }) + it('never bills an audio clip at the video floor', () => { expect(FALAI_AUDIO_FALLBACK_PROVIDER_COST_DOLLARS).toBeLessThan( FALAI_VIDEO_FALLBACK_PROVIDER_COST_DOLLARS diff --git a/apps/sim/lib/tools/falai-pricing.ts b/apps/sim/lib/tools/falai-pricing.ts index 0695a7bf0d1..b31a1cbf428 100644 --- a/apps/sim/lib/tools/falai-pricing.ts +++ b/apps/sim/lib/tools/falai-pricing.ts @@ -2,6 +2,7 @@ import { createLogger } from '@sim/logger' import { getErrorMessage } from '@sim/utils/errors' import { interruptibleSleep } from '@sim/utils/helpers' import { isRecordLike } from '@sim/utils/object' +import { DeadlineExceededError, withinDeadline } from '@/lib/core/utils/deadline' import { DEFAULT_MAX_ERROR_BODY_BYTES, readResponseJsonWithLimit, @@ -15,6 +16,11 @@ export const FALAI_VIDEO_FALLBACK_PROVIDER_COST_DOLLARS = 0.25 const FALAI_BILLING_EVENT_ATTEMPTS = 2 const FALAI_BILLING_EVENT_RETRY_MS = 500 const FALAI_PRICING_RESPONSE_MAX_BYTES = 2 * 1024 * 1024 +/** + * Longest a cost lookup may take, covering every billing-event attempt and the estimate. The + * generation is already paid for by then, so a lookup that runs out falls back to the floor. + */ +const FALAI_COST_LOOKUP_TIMEOUT_MS = 10_000 const logger = createLogger('FalAIPricing') export interface FalAICostMetadata { @@ -167,10 +173,16 @@ async function estimateFalAICallCost( return { error: `Fal.ai pricing estimate failed: ${response.status} ${error}` } } - const data = await readResponseJsonWithLimit(response, { - maxBytes: FALAI_PRICING_RESPONSE_MAX_BYTES, - label: 'Fal.ai pricing estimate response', - }) + let data: unknown + try { + data = await readResponseJsonWithLimit(response, { + maxBytes: FALAI_PRICING_RESPONSE_MAX_BYTES, + label: 'Fal.ai pricing estimate response', + }) + } catch (error) { + signal?.throwIfAborted() + return { error: getErrorMessage(error, 'Unknown error') } + } const totalCost = isRecordLike(data) ? getNumber(data.total_cost) : undefined if (totalCost === undefined) { return { error: 'Fal.ai pricing estimate missing total_cost' } @@ -179,18 +191,57 @@ async function estimateFalAICallCost( return { costDollars: totalCost } } -export async function getFalAICostMetadata({ - apiKey, - endpointId, - requestId, - signal, -}: { +interface FalAICostLookup { apiKey: string endpointId: string requestId: string + /** Ends the lookup with an error. Without one, only {@link FALAI_COST_LOOKUP_TIMEOUT_MS} ends it. */ signal?: AbortSignal -}): Promise { - signal?.throwIfAborted() +} + +/** + * The cost of a completed fal request. The lookup is bounded on its own: past + * {@link FALAI_COST_LOOKUP_TIMEOUT_MS} it returns the fallback floor rather than nothing. + */ +export async function getFalAICostMetadata(lookup: FalAICostLookup): Promise { + const { endpointId, requestId } = lookup + try { + return await withinDeadline( + (signal) => lookUpFalAICost({ ...lookup, signal }), + Date.now() + FALAI_COST_LOOKUP_TIMEOUT_MS, + lookup.signal + ) + } catch (error) { + if (!(error instanceof DeadlineExceededError)) throw error + logger.warn('Fal.ai cost lookup timed out after generation completed', { + endpointId, + requestId, + }) + return fallbackFloorCost(endpointId, requestId, error.message) + } +} + +function fallbackFloorCost( + endpointId: string, + requestId: string, + error: string | undefined +): FalAICostMetadata { + return { + endpointId, + requestId, + costDollars: getFalAIFallbackProviderCostDollars(endpointId), + source: 'fallback_floor', + currency: 'USD', + error, + } +} + +async function lookUpFalAICost({ + apiKey, + endpointId, + requestId, + signal, +}: FalAICostLookup): Promise { for (let attempt = 0; attempt < FALAI_BILLING_EVENT_ATTEMPTS; attempt++) { const event = await fetchFalAIBillingEvent(apiKey, requestId, signal) if (event) { @@ -229,12 +280,5 @@ export async function getFalAICostMetadata({ error: estimate.error, }) - return { - endpointId, - requestId, - costDollars: getFalAIFallbackProviderCostDollars(endpointId), - source: 'fallback_floor', - currency: 'USD', - error: estimate.error, - } + return fallbackFloorCost(endpointId, requestId, estimate.error) }