fix(media): bill a completed fal generation even when a Stop ends its download

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,

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,

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<void>; hangs?: boolean } = {},
+  result: Record<string, unknown> = { 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<Response>((_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<Response>) {
+  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<void>()
+    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()
+  })
+})

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<void>
+
+/**
+ * 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<FalAICostMetadata> {
+  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).

apps/sim/lib/mothership/tools/server/media/generate-audio.ts

@@ -131,10 +131,9 @@ export const generateAudioServerTool: BaseServerTool<GenerateAudioArgs, Generate
         instrumental: params.instrumental,
         voiceSampleDataUri,
         signal: context.abortSignal,
+        onCost: (cost) => 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}`

apps/sim/lib/mothership/tools/server/media/generate-video.ts

@@ -108,10 +108,9 @@ export const generateVideoServerTool: BaseServerTool<GenerateVideoArgs, Generate
         promptOptimizer: params.promptOptimizer,
         imageDataUri,
         signal: context.abortSignal,
+        onCost: (cost) => 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'