Skip to content

Commit b73d52b

Browse files
waleedlatif1waleedlatif1
committed
fix(falai): bound pricing lookups and correct result handling
1 parent 6c21f82 commit b73d52b

10 files changed

Lines changed: 692 additions & 119 deletions

File tree

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

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import {
66
inputValidationMock,
77
inputValidationMockFns,
88
} from '@sim/testing/mocks/input-validation.mock'
9+
import { getMockLogger } from '@sim/testing/mocks/logger.mock'
910
import { uploadsCopilotMock, uploadsCopilotMockFns } from '@sim/testing/mocks/uploads-copilot.mock'
1011
import { uploadsExecutionMock } from '@sim/testing/mocks/uploads-execution.mock'
1112
import { utilsHelpersMock, utilsHelpersMockFns } from '@sim/testing/mocks/utils-helpers.mock'
@@ -14,6 +15,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'
1415
const mocks = vi.hoisted(() => ({
1516
fetch: vi.fn(),
1617
getFalAICostMetadata: vi.fn(),
18+
getFalAIBillingReceipt: vi.fn(),
1719
}))
1820

1921
vi.mock('@/lib/core/security/input-validation.server', () => inputValidationMock)
@@ -24,6 +26,7 @@ vi.mock('@/lib/core/execution-limits', () => executionLimitsMock)
2426

2527
vi.mock('@/lib/tools/falai-pricing', () => ({
2628
getFalAICostMetadata: mocks.getFalAICostMetadata,
29+
getFalAIBillingReceipt: mocks.getFalAIBillingReceipt,
2730
}))
2831

2932
vi.mock('@/lib/uploads/contexts/copilot', () => uploadsCopilotMock)
@@ -33,6 +36,7 @@ vi.mock('@/lib/uploads/contexts/execution', () => uploadsExecutionMock)
3336
const { mockSecureFetchWithPinnedIP, mockValidateUrlWithDNS } = inputValidationMockFns
3437

3538
import { executeImageGeneration } from '@/lib/internal/image/operations'
39+
import { observeServiceCosts } from '@/lib/mothership/billing/service-observer'
3640

3741
const { mockInterruptibleSleep } = utilsHelpersMockFns
3842

@@ -50,6 +54,7 @@ const falInput = {
5054
describe('image operations', () => {
5155
beforeEach(() => {
5256
mocks.fetch.mockReset()
57+
mocks.getFalAICostMetadata.mockReset()
5358
mockSecureFetchWithPinnedIP.mockReset()
5459
vi.stubGlobal('fetch', mocks.fetch)
5560
mockInterruptibleSleep.mockResolvedValue(undefined)
@@ -212,4 +217,96 @@ describe('image operations', () => {
212217
)
213218
expect(mocks.getFalAICostMetadata.mock.calls[0][0].signal).toBeUndefined()
214219
})
220+
it.each([422, 500])(
221+
'reports HTTP %i generation failures without speculative delivery receipts',
222+
async (status) => {
223+
mocks.fetch.mockResolvedValueOnce(Response.json({ request_id: 'rejected-job' }))
224+
mockSecureFetchWithPinnedIP
225+
.mockResolvedValueOnce(Response.json({ status: 'COMPLETED' }))
226+
.mockResolvedValueOnce(
227+
Response.json(
228+
{
229+
detail: [
230+
{ type: 'no_media_generated', msg: 'secret-prompt', input: 'secret-prompt' },
231+
],
232+
},
233+
{ status }
234+
)
235+
)
236+
mocks.getFalAICostMetadata.mockResolvedValue({
237+
requestId: 'rejected-job',
238+
source: 'fallback_floor',
239+
costDollars: 0.05,
240+
})
241+
const receipts: string[] = []
242+
const response = await observeServiceCosts(
243+
async () => {},
244+
() =>
245+
executeImageGeneration(
246+
{ ...falInput, useHostedCostTracking: true },
247+
{ userId: 'user', requestId: 'rejected' }
248+
),
249+
async (receipt) => {
250+
receipts.push(receipt)
251+
}
252+
)
253+
expect(response.status).toBe(status === 422 ? 422 : 502)
254+
const body = await response.json()
255+
expect(body.error).toContain('no_media_generated')
256+
expect(JSON.stringify(body)).not.toContain('secret-prompt')
257+
expect(JSON.stringify(getMockLogger('ImageOperations').error.mock.calls)).not.toContain(
258+
'secret-prompt'
259+
)
260+
expect(receipts).toEqual([])
261+
expect(mocks.getFalAICostMetadata).not.toHaveBeenCalled()
262+
expect(mockUploadCopilotFile).not.toHaveBeenCalled()
263+
}
264+
)
265+
266+
it.each([
267+
{ bytes: 32 * 1024 * 1024, declared: true, expected: 200 },
268+
{ bytes: 101 * 1024 * 1024, declared: true, expected: 413 },
269+
{ bytes: 101 * 1024 * 1024, declared: false, expected: 413 },
270+
])(
271+
'bounds remote generated images ($bytes bytes, declared=$declared)',
272+
async ({ bytes, declared, expected }) => {
273+
mocks.fetch.mockResolvedValueOnce(Response.json({ request_id: 'large-image' }))
274+
let remaining = bytes
275+
const chunk = new Uint8Array(1024 * 1024)
276+
const binary = new Response(
277+
new ReadableStream({
278+
pull(controller) {
279+
if (remaining <= 0) {
280+
controller.close()
281+
return
282+
}
283+
controller.enqueue(chunk)
284+
remaining -= chunk.length
285+
},
286+
}),
287+
{
288+
headers: {
289+
'Content-Type': 'image/png',
290+
...(declared ? { 'Content-Length': String(bytes) } : {}),
291+
},
292+
}
293+
)
294+
mockSecureFetchWithPinnedIP
295+
.mockResolvedValueOnce(Response.json({ status: 'COMPLETED' }))
296+
.mockResolvedValueOnce(
297+
Response.json({ images: [{ url: 'https://cdn.example/image.png' }] })
298+
)
299+
.mockResolvedValueOnce(binary)
300+
const response = await executeImageGeneration(falInput, {
301+
userId: 'user',
302+
requestId: 'large',
303+
})
304+
expect(response.status).toBe(expected)
305+
if (expected === 200) {
306+
expect(mockUploadCopilotFile.mock.calls[0][0].buffer.byteLength).toBe(bytes)
307+
} else {
308+
expect(mockUploadCopilotFile).not.toHaveBeenCalled()
309+
}
310+
}
311+
)
215312
})

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

Lines changed: 41 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,12 @@ import {
1919
import { MAX_REMOTE_IMAGE_BYTES } from '@/lib/internal/image/fetch'
2020
import { buildImageGenerationRequest } from '@/lib/internal/image/request'
2121
import type { ImageGenerationInput, ImageProvider } from '@/lib/internal/image/schema'
22+
import { FalAIResultError, fetchFalAIResult } from '@/lib/media/falai-result'
2223
import { recordServiceMeteringFailure } from '@/lib/mothership/billing/service-observer'
23-
import { type FalAICostMetadata, getFalAICostMetadata } from '@/lib/tools/falai-pricing'
24+
import type { FalAICostMetadata } from '@/lib/tools/falai-pricing'
2425
import { uploadCopilotFile } from '@/lib/uploads/contexts/copilot'
2526
import { uploadExecutionFile } from '@/lib/uploads/contexts/execution'
27+
import { MAX_BUFFERED_TRANSFER_BYTES } from '@/lib/uploads/shared/types'
2628
import { resolveStoredFileMetadata } from '@/lib/uploads/utils/stored-file-metadata'
2729

2830
const logger = createLogger('ImageOperations')
@@ -78,7 +80,7 @@ export async function executeImageGeneration(
7880
context: ImageOperationContext
7981
): Promise<Response> {
8082
const { requestId } = context
81-
let completedFalCost: FalAICostMetadata | undefined
83+
let undeliveredFalCost: FalAICostMetadata | undefined
8284
logger.info(`[${requestId}] Image generation request started`)
8385

8486
try {
@@ -108,7 +110,7 @@ export async function executeImageGeneration(
108110
requestId,
109111
logger,
110112
(cost) => {
111-
completedFalCost = cost
113+
undeliveredFalCost = cost
112114
},
113115
context.signal
114116
)
@@ -126,16 +128,16 @@ export async function executeImageGeneration(
126128

127129
return Response.json(storedImage)
128130
} catch (error) {
129-
if (completedFalCost) {
130-
logger.warn('Hosted fal image completed but delivery failed', {
131+
if (undeliveredFalCost) {
132+
logger.warn('Hosted fal image result delivery interrupted', {
131133
requestId,
132-
falaiBilling: completedFalCost,
134+
falaiBilling: undeliveredFalCost,
133135
})
134-
const { requestId: jobId, endpointId, costDollars } = completedFalCost
136+
const { requestId: jobId, endpointId, costDollars } = undeliveredFalCost
135137
await recordServiceMeteringFailure(
136-
`Completed fal image ${jobId} (${endpointId}): USD ${costDollars}; delivery failed`
138+
`Undelivered fal image ${jobId} (${endpointId}): USD ${costDollars} (${undeliveredFalCost.source}); reconciliation required`
137139
).catch((meteringError) => {
138-
logger.error('Failed to record completed fal image for reconciliation', {
140+
logger.error('Failed to record undelivered fal image for reconciliation', {
139141
requestId,
140142
error: getErrorMessage(meteringError),
141143
})
@@ -146,7 +148,14 @@ export async function executeImageGeneration(
146148
const errorMessage = getUserFacingErrorMessage(error, 'Unknown error')
147149
return Response.json(
148150
{ error: errorMessage },
149-
{ status: isPayloadSizeLimitError(error) ? 413 : 500 }
151+
{
152+
status:
153+
error instanceof FalAIResultError
154+
? error.status
155+
: isPayloadSizeLimitError(error)
156+
? 413
157+
: 500,
158+
}
150159
)
151160
}
152161
}
@@ -207,7 +216,7 @@ async function bufferFromImageUrl(
207216
const imageResponse = await secureFetchWithPinnedIP(url, urlValidation.resolvedIP, {
208217
profile: 'contentFetch',
209218
method: 'GET',
210-
maxResponseBytes: MAX_IMAGE_BYTES,
219+
maxResponseBytes: MAX_BUFFERED_TRANSFER_BYTES,
211220
signal,
212221
})
213222
if (!imageResponse.ok) {
@@ -221,7 +230,7 @@ async function bufferFromImageUrl(
221230

222231
const contentType = imageResponse.headers.get('content-type') || 'image/png'
223232
const buffer = await readResponseToBufferWithLimit(imageResponse, {
224-
maxBytes: MAX_IMAGE_BYTES,
233+
maxBytes: MAX_BUFFERED_TRANSFER_BYTES,
225234
label: 'generated image download',
226235
signal,
227236
})
@@ -505,41 +514,30 @@ async function generateWithFalAI(
505514
throw new Error(`Fal.ai generation failed: ${getFalAIErrorMessage(statusError)}`)
506515
}
507516

508-
const falaiCost = body.useHostedCostTracking
509-
? await getFalAICostMetadata({ apiKey, endpointId: endpoint, requestId: falRequestId })
510-
: undefined
511-
if (falaiCost) onFalCost(falaiCost)
512-
signal?.throwIfAborted()
513-
514-
const resultUrl = falQueueUrl(getStringProperty(statusData, 'response_url') || responseUrl)
515-
const resultValidation = await validateUrlWithDNS(resultUrl, 'resultUrl', 'contentFetch')
516-
if (!resultValidation.isValid) {
517-
throw new Error(resultValidation.error)
518-
}
519-
const resultResponse = await secureFetchWithPinnedIP(resultUrl, resultValidation.resolvedIP, {
520-
profile: 'contentFetch',
521-
maxRedirects: 0,
522-
headers: {
523-
Authorization: `Key ${apiKey}`,
524-
},
517+
const { data: resultData, cost: falaiCost } = await fetchFalAIResult({
518+
apiKey,
519+
endpointId: endpoint,
520+
requestId: falRequestId,
525521
maxResponseBytes: MAX_IMAGE_JSON_BYTES,
522+
trackCost: body.useHostedCostTracking,
523+
onCost: onFalCost,
526524
signal,
525+
fetchResponse: async () => {
526+
const resultUrl = falQueueUrl(
527+
getStringProperty(statusData, 'response_url') || responseUrl
528+
)
529+
const resultValidation = await validateUrlWithDNS(resultUrl, 'resultUrl', 'contentFetch')
530+
if (!resultValidation.isValid) throw new Error(resultValidation.error)
531+
return secureFetchWithPinnedIP(resultUrl, resultValidation.resolvedIP, {
532+
profile: 'contentFetch',
533+
maxRedirects: 0,
534+
headers: { Authorization: `Key ${apiKey}` },
535+
maxResponseBytes: MAX_IMAGE_JSON_BYTES,
536+
signal,
537+
})
538+
},
527539
})
528540

529-
if (!resultResponse.ok) {
530-
await readResponseTextWithLimit(resultResponse, {
531-
maxBytes: DEFAULT_MAX_ERROR_BODY_BYTES,
532-
label: 'Fal.ai result error response',
533-
signal,
534-
}).catch(() => '')
535-
throw new Error(`Failed to fetch Fal.ai result: ${resultResponse.status}`)
536-
}
537-
538-
const resultData = await readResponseJsonWithLimit(resultResponse, {
539-
maxBytes: MAX_IMAGE_JSON_BYTES,
540-
label: 'Fal.ai result response',
541-
signal,
542-
})
543541
if (!isRecordLike(resultData)) {
544542
throw new Error('Invalid Fal.ai result response')
545543
}

‎apps/sim/lib/internal/video/client.test.ts‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -517,6 +517,97 @@ describe('Video provider client', () => {
517517
)
518518
})
519519

520+
it.each([422, 500])(
521+
'does not reconcile a failed fal result (HTTP %i) as delivered media',
522+
async (status) => {
523+
const paths: string[] = []
524+
vi.stubGlobal('fetch', async (url: string | URL) => {
525+
const path = new URL(url).pathname
526+
paths.push(path)
527+
if (path.endsWith('/status')) return jsonResponse({ status: 'COMPLETED' })
528+
if (path.includes('/requests/'))
529+
return jsonResponse(
530+
{
531+
detail: [
532+
{ type: 'no_media_generated', msg: 'secret-prompt', input: 'secret-prompt' },
533+
],
534+
},
535+
status
536+
)
537+
if (path.endsWith('/billing-events')) return jsonResponse({ billing_events: [] })
538+
if (path.endsWith('/estimate')) return jsonResponse({ total_cost: 0.25 })
539+
return jsonResponse({ request_id: 'failed-result' })
540+
})
541+
const pending: string[] = []
542+
const result = observeServiceCosts(
543+
async () => {},
544+
() =>
545+
executeVideoOperation(
546+
{
547+
provider: 'falai',
548+
model: 'veo-3.1',
549+
prompt: 'A sunrise',
550+
apiKey: 'key',
551+
useHostedCostTracking: true,
552+
},
553+
{ requestId: 'trace', userId: 'user', headers: new Headers() }
554+
),
555+
async (receipt) => {
556+
pending.push(receipt)
557+
}
558+
)
559+
const rejection = expect(result).rejects.toMatchObject({
560+
status: status === 422 ? 422 : 502,
561+
message: expect.stringContaining('no_media_generated'),
562+
})
563+
await Promise.all([rejection, vi.advanceTimersByTimeAsync(6000)])
564+
expect(pending).toEqual([])
565+
expect(paths.some((path) => path.endsWith('/estimate'))).toBe(false)
566+
}
567+
)
568+
569+
it('preserves an exact charge even when the provider returns a failed result', async () => {
570+
vi.stubGlobal('fetch', async (url: string | URL) => {
571+
const path = new URL(url).pathname
572+
if (path.endsWith('/status')) return jsonResponse({ status: 'COMPLETED' })
573+
if (path.includes('/requests/'))
574+
return jsonResponse({ detail: [{ type: 'no_media_generated' }] }, 422)
575+
if (path.endsWith('/billing-events'))
576+
return jsonResponse({
577+
billing_events: [
578+
{
579+
request_id: 'charged-failure',
580+
endpoint_id: 'fal-ai/veo3.1',
581+
cost_estimate_nano_usd: 400000000,
582+
},
583+
],
584+
})
585+
return jsonResponse({ request_id: 'charged-failure' })
586+
})
587+
const receipts: string[] = []
588+
const result = observeServiceCosts(
589+
async () => {},
590+
() =>
591+
executeVideoOperation(
592+
{
593+
provider: 'falai',
594+
model: 'veo-3.1',
595+
prompt: 'A sunrise',
596+
apiKey: 'key',
597+
useHostedCostTracking: true,
598+
},
599+
{ requestId: 'trace', userId: 'user', headers: new Headers() }
600+
),
601+
async (receipt) => {
602+
receipts.push(receipt)
603+
}
604+
)
605+
const rejection = expect(result).rejects.toMatchObject({ status: 422 })
606+
await Promise.all([rejection, vi.advanceTimersByTimeAsync(6000)])
607+
expect(receipts).toHaveLength(1)
608+
expect(receipts[0]).toContain('0.4 (billing_events)')
609+
})
610+
520611
it('preserves the completed hosted receipt for reconciliation when Stop interrupts delivery', async () => {
521612
const controller = new AbortController()
522613
const requestId = 'completed-video'

0 commit comments

Comments
 (0)