fix(falai): bound pricing lookups and correct result handling

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')
       }