fix(billing): settle hosted fal generations to their exact cost

Author: waleedlatif1

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 }
+  },
+})

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,7 +121,11 @@ 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()
   })

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'
@@ -465,125 +466,142 @@ async function generateWithFalAI(
       buildFalAIQueueUrl(endpoint, falRequestId, 'response')
   )
 
-  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 cancelUrl = falQueueUrl(
+    getStringProperty(createData, 'cancel_url') ||
+      `${buildFalAIQueueUrl(endpoint, falRequestId, 'response')}/cancel`
+  )
 
-  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,
-    })
+  logger.info('Fal.ai image request created', { requestId, falRequestId })
 
-    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}`)
+  let reachedTerminalStatus = false
+  try {
+    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 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 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 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,
+        }
       }
 
-      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 || '')) {
+        reachedTerminalStatus = true
+        throw new Error(`Fal.ai generation failed: ${getFalAIErrorMessage(statusData.error)}`)
       }
-    }
 
-    if (['ERROR', 'FAILED', 'CANCELLED'].includes(status || '')) {
-      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(