fix(falai): bound pricing lookups and preserve result receipts (#9706)

Bound and coalesce account-specific fal pricing lookups, honor Retry-After cooldowns, and validate exact receipt identity.

Preserve exact receipts through cancellation, terminal errors, and stalled retrieval while distinguishing explicit generation failures from transport failures. Keep error output bounded and redacted, and support generated images up to the existing buffered-transfer limit.

Validated with focused regressions, the full repository suite, type checks, lint, audits, and CI.

Author: waleedlatif1

apps/sim/lib/internal/image/operations.test.ts

@@ -6,14 +6,16 @@ import {
   inputValidationMock,
   inputValidationMockFns,
 } from '@sim/testing/mocks/input-validation.mock'
+import { getMockLogger } from '@sim/testing/mocks/logger.mock'
 import { uploadsCopilotMock, uploadsCopilotMockFns } from '@sim/testing/mocks/uploads-copilot.mock'
 import { uploadsExecutionMock } from '@sim/testing/mocks/uploads-execution.mock'
 import { utilsHelpersMock, utilsHelpersMockFns } from '@sim/testing/mocks/utils-helpers.mock'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 const mocks = vi.hoisted(() => ({
   fetch: vi.fn(),
-  getFalAICostMetadata: vi.fn(),
+  getFalAIEstimatedCostMetadata: vi.fn(),
+  getFalAIBillingReceipt: vi.fn(),
 }))
 
 vi.mock('@/lib/core/security/input-validation.server', () => inputValidationMock)
@@ -23,7 +25,9 @@ vi.mock('@sim/utils/helpers', () => utilsHelpersMock)
 vi.mock('@/lib/core/execution-limits', () => executionLimitsMock)
 
 vi.mock('@/lib/tools/falai-pricing', () => ({
-  getFalAICostMetadata: mocks.getFalAICostMetadata,
+  FALAI_COST_LOOKUP_TIMEOUT_MS: 10_000,
+  getFalAIEstimatedCostMetadata: mocks.getFalAIEstimatedCostMetadata,
+  getFalAIBillingReceipt: mocks.getFalAIBillingReceipt,
 }))
 
 vi.mock('@/lib/uploads/contexts/copilot', () => uploadsCopilotMock)
@@ -33,6 +37,7 @@ vi.mock('@/lib/uploads/contexts/execution', () => uploadsExecutionMock)
 const { mockSecureFetchWithPinnedIP, mockValidateUrlWithDNS } = inputValidationMockFns
 
 import { executeImageGeneration } from '@/lib/internal/image/operations'
+import { observeServiceCosts } from '@/lib/mothership/billing/service-observer'
 
 const { mockInterruptibleSleep } = utilsHelpersMockFns
 
@@ -50,6 +55,8 @@ const falInput = {
 describe('image operations', () => {
   beforeEach(() => {
     mocks.fetch.mockReset()
+    mocks.getFalAIEstimatedCostMetadata.mockReset()
+    mocks.getFalAIBillingReceipt.mockReset().mockResolvedValue(undefined)
     mockSecureFetchWithPinnedIP.mockReset()
     vi.stubGlobal('fetch', mocks.fetch)
     mockInterruptibleSleep.mockResolvedValue(undefined)
@@ -182,15 +189,17 @@ describe('image operations', () => {
   it('retains the completed fal receipt when delivery is interrupted', async () => {
     const controller = new AbortController()
     mocks.fetch.mockResolvedValueOnce(Response.json({ request_id: 'paid-job' }))
-    mocks.getFalAICostMetadata.mockImplementation(async ({ signal }: { signal?: AbortSignal }) => {
-      signal?.throwIfAborted()
-      return {
-        requestId: 'paid-job',
-        endpointId: 'fal-ai/nano-banana-2',
-        costDollars: 0.1,
-        source: 'billing_events',
+    mocks.getFalAIBillingReceipt.mockImplementation(
+      async ({ signal }: { signal?: AbortSignal }) => {
+        signal?.throwIfAborted()
+        return {
+          requestId: 'paid-job',
+          endpointId: 'fal-ai/nano-banana-2',
+          costDollars: 0.1,
+          source: 'billing_events',
+        }
       }
-    })
+    )
     mockSecureFetchWithPinnedIP.mockResolvedValueOnce(Response.json({ status: 'COMPLETED' }))
     mockSecureFetchWithPinnedIP.mockImplementationOnce(async () => {
       controller.abort()
@@ -207,9 +216,130 @@ describe('image operations', () => {
         }
       )
     ).rejects.toMatchObject({ name: 'AbortError' })
-    expect(mocks.getFalAICostMetadata).toHaveBeenCalledWith(
+    expect(mocks.getFalAIBillingReceipt).toHaveBeenCalledWith(
       expect.objectContaining({ requestId: 'paid-job' })
     )
-    expect(mocks.getFalAICostMetadata.mock.calls[0][0].signal).toBeUndefined()
+    expect(mocks.getFalAIBillingReceipt.mock.calls[0][0].signal).toBeUndefined()
+  })
+  it.each([422, 500])(
+    'reports HTTP %i generation failures without speculative delivery receipts',
+    async (status) => {
+      mocks.fetch.mockResolvedValueOnce(Response.json({ request_id: 'rejected-job' }))
+      mockSecureFetchWithPinnedIP
+        .mockResolvedValueOnce(Response.json({ status: 'COMPLETED' }))
+        .mockResolvedValueOnce(
+          Response.json(
+            {
+              detail: [
+                { type: 'no_media_generated', msg: 'secret-prompt', input: 'secret-prompt' },
+              ],
+            },
+            { status }
+          )
+        )
+      mocks.getFalAIEstimatedCostMetadata.mockResolvedValue({
+        requestId: 'rejected-job',
+        source: 'fallback_floor',
+        costDollars: 0.05,
+      })
+      const receipts: string[] = []
+      const response = await observeServiceCosts(
+        async () => {},
+        () =>
+          executeImageGeneration(
+            { ...falInput, useHostedCostTracking: true },
+            { userId: 'user', requestId: 'rejected' }
+          ),
+        async (receipt) => {
+          receipts.push(receipt)
+        }
+      )
+      expect(response.status).toBe(status === 422 ? 422 : 502)
+      const body = await response.json()
+      expect(body.error).toContain('no_media_generated')
+      expect(JSON.stringify(body)).not.toContain('secret-prompt')
+      expect(JSON.stringify(getMockLogger('ImageOperations').error.mock.calls)).not.toContain(
+        'secret-prompt'
+      )
+      expect(receipts).toEqual([])
+      expect(mocks.getFalAIEstimatedCostMetadata).not.toHaveBeenCalled()
+      expect(mockUploadCopilotFile).not.toHaveBeenCalled()
+    }
+  )
+
+  it.each([
+    { bytes: 32 * 1024 * 1024, declared: true, expected: 200 },
+    { bytes: 101 * 1024 * 1024, declared: true, expected: 413 },
+    { bytes: 101 * 1024 * 1024, declared: false, expected: 413 },
+  ])(
+    'bounds remote generated images ($bytes bytes, declared=$declared)',
+    async ({ bytes, declared, expected }) => {
+      mocks.fetch.mockResolvedValueOnce(Response.json({ request_id: 'large-image' }))
+      let remaining = bytes
+      const chunk = new Uint8Array(1024 * 1024)
+      const binary = new Response(
+        new ReadableStream({
+          pull(controller) {
+            if (remaining <= 0) {
+              controller.close()
+              return
+            }
+            controller.enqueue(chunk)
+            remaining -= chunk.length
+          },
+        }),
+        {
+          headers: {
+            'Content-Type': 'image/png',
+            ...(declared ? { 'Content-Length': String(bytes) } : {}),
+          },
+        }
+      )
+      mockSecureFetchWithPinnedIP
+        .mockResolvedValueOnce(Response.json({ status: 'COMPLETED' }))
+        .mockResolvedValueOnce(
+          Response.json({ images: [{ url: 'https://cdn.example/image.png' }] })
+        )
+        .mockResolvedValueOnce(binary)
+      const response = await executeImageGeneration(falInput, {
+        userId: 'user',
+        requestId: 'large',
+      })
+      expect(response.status).toBe(expected)
+      if (expected === 200) {
+        expect(mockUploadCopilotFile.mock.calls[0][0].buffer.byteLength).toBe(bytes)
+      } else {
+        expect(mockUploadCopilotFile).not.toHaveBeenCalled()
+      }
+    }
+  )
+  it('reconciles exact image receipts when a completed queue status contains an error', async () => {
+    mocks.fetch.mockResolvedValueOnce(Response.json({ request_id: 'image-queue-error' }))
+    mockSecureFetchWithPinnedIP.mockResolvedValueOnce(
+      Response.json({ status: 'COMPLETED', error: 'sensitive provider error' })
+    )
+    mocks.getFalAIBillingReceipt.mockResolvedValueOnce({
+      requestId: 'image-queue-error',
+      endpointId: 'fal-ai/nano-banana-2',
+      source: 'billing_events',
+      costDollars: 0.1,
+    })
+    const receipts: string[] = []
+    const response = await observeServiceCosts(
+      async () => {},
+      () =>
+        executeImageGeneration(
+          { ...falInput, useHostedCostTracking: true },
+          { userId: 'user', requestId: 'queue-error' }
+        ),
+      async (receipt) => {
+        receipts.push(receipt)
+      }
+    )
+    expect(response.ok).toBe(false)
+    expect(JSON.stringify(await response.json())).not.toContain('sensitive provider error')
+    expect(receipts).toHaveLength(1)
+    expect(receipts[0]).toContain('0.1 (billing_events)')
+    expect(mockSecureFetchWithPinnedIP).toHaveBeenCalledTimes(1)
   })
 })

apps/sim/lib/internal/image/operations.ts

@@ -19,10 +19,12 @@ 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 { FalAIResultError, fetchFalAIResult } from '@/lib/media/falai-result'
 import { recordServiceMeteringFailure } from '@/lib/mothership/billing/service-observer'
-import { type FalAICostMetadata, getFalAICostMetadata } from '@/lib/tools/falai-pricing'
+import type { FalAICostMetadata } from '@/lib/tools/falai-pricing'
 import { uploadCopilotFile } from '@/lib/uploads/contexts/copilot'
 import { uploadExecutionFile } from '@/lib/uploads/contexts/execution'
+import { MAX_BUFFERED_TRANSFER_BYTES } from '@/lib/uploads/shared/types'
 import { resolveStoredFileMetadata } from '@/lib/uploads/utils/stored-file-metadata'
 
 const logger = createLogger('ImageOperations')
@@ -78,7 +80,7 @@ export async function executeImageGeneration(
   context: ImageOperationContext
 ): Promise<Response> {
   const { requestId } = context
-  let completedFalCost: FalAICostMetadata | undefined
+  let undeliveredFalCost: FalAICostMetadata | undefined
   logger.info(`[${requestId}] Image generation request started`)
 
   try {
@@ -108,7 +110,7 @@ export async function executeImageGeneration(
         requestId,
         logger,
         (cost) => {
-          completedFalCost = cost
+          undeliveredFalCost = cost
         },
         context.signal
       )
@@ -126,16 +128,16 @@ export async function executeImageGeneration(
 
     return Response.json(storedImage)
   } catch (error) {
-    if (completedFalCost) {
-      logger.warn('Hosted fal image completed but delivery failed', {
+    if (undeliveredFalCost) {
+      logger.warn('Hosted fal image result delivery interrupted', {
         requestId,
-        falaiBilling: completedFalCost,
+        falaiBilling: undeliveredFalCost,
       })
-      const { requestId: jobId, endpointId, costDollars } = completedFalCost
+      const { requestId: jobId, endpointId, costDollars } = undeliveredFalCost
       await recordServiceMeteringFailure(
-        `Completed fal image ${jobId} (${endpointId}): USD ${costDollars}; delivery failed`
+        `Undelivered fal image ${jobId} (${endpointId}): USD ${costDollars} (${undeliveredFalCost.source}); reconciliation required`
       ).catch((meteringError) => {
-        logger.error('Failed to record completed fal image for reconciliation', {
+        logger.error('Failed to record undelivered fal image for reconciliation', {
           requestId,
           error: getErrorMessage(meteringError),
         })
@@ -146,7 +148,14 @@ export async function executeImageGeneration(
     const errorMessage = getUserFacingErrorMessage(error, 'Unknown error')
     return Response.json(
       { error: errorMessage },
-      { status: isPayloadSizeLimitError(error) ? 413 : 500 }
+      {
+        status:
+          error instanceof FalAIResultError
+            ? error.status
+            : isPayloadSizeLimitError(error)
+              ? 413
+              : 500,
+      }
     )
   }
 }
@@ -207,7 +216,7 @@ async function bufferFromImageUrl(
   const imageResponse = await secureFetchWithPinnedIP(url, urlValidation.resolvedIP, {
     profile: 'contentFetch',
     method: 'GET',
-    maxResponseBytes: MAX_IMAGE_BYTES,
+    maxResponseBytes: MAX_BUFFERED_TRANSFER_BYTES,
     signal,
   })
   if (!imageResponse.ok) {
@@ -221,7 +230,7 @@ async function bufferFromImageUrl(
 
   const contentType = imageResponse.headers.get('content-type') || 'image/png'
   const buffer = await readResponseToBufferWithLimit(imageResponse, {
-    maxBytes: MAX_IMAGE_BYTES,
+    maxBytes: MAX_BUFFERED_TRANSFER_BYTES,
     label: 'generated image download',
     signal,
   })
@@ -500,46 +509,31 @@ async function generateWithFalAI(
 
     const status = getStringProperty(statusData, 'status')
     if (status === 'COMPLETED') {
-      const statusError = statusData.error
-      if (statusError) {
-        throw new Error(`Fal.ai generation failed: ${getFalAIErrorMessage(statusError)}`)
-      }
-
-      const falaiCost = body.useHostedCostTracking
-        ? await getFalAICostMetadata({ apiKey, endpointId: endpoint, requestId: falRequestId })
-        : undefined
-      if (falaiCost) onFalCost(falaiCost)
-      signal?.throwIfAborted()
-
-      const resultUrl = falQueueUrl(getStringProperty(statusData, 'response_url') || responseUrl)
-      const resultValidation = await validateUrlWithDNS(resultUrl, 'resultUrl', 'contentFetch')
-      if (!resultValidation.isValid) {
-        throw new Error(resultValidation.error)
-      }
-      const resultResponse = await secureFetchWithPinnedIP(resultUrl, resultValidation.resolvedIP, {
-        profile: 'contentFetch',
-        maxRedirects: 0,
-        headers: {
-          Authorization: `Key ${apiKey}`,
-        },
+      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 (!resultResponse.ok) {
-        await readResponseTextWithLimit(resultResponse, {
-          maxBytes: DEFAULT_MAX_ERROR_BODY_BYTES,
-          label: 'Fal.ai result error response',
-          signal,
-        }).catch(() => '')
-        throw new Error(`Failed to fetch Fal.ai result: ${resultResponse.status}`)
-      }
-
-      const resultData = await readResponseJsonWithLimit(resultResponse, {
-        maxBytes: MAX_IMAGE_JSON_BYTES,
-        label: 'Fal.ai result response',
-        signal,
-      })
       if (!isRecordLike(resultData)) {
         throw new Error('Invalid Fal.ai result response')
       }