fix(mothership): bound every handed tool call by the deadline its frame names (#9090)

* fix(mothership): bound every handed tool call by the deadline its frame names

* fix(mothership): start nothing for a tool call once its deadline has passed

* test(mothership): type the client deadline waiter mock

* fix(mothership): report a past-deadline client call as not run only when nothing claimed it

* fix(media): submit no fal request for a caller that already gave up

* fix(media): flag a refused fal cancellation and stop an aborted inline download

Author: waleedlatif1

apps/sim/lib/media/falai-audio.ts

@@ -26,6 +26,8 @@ export interface GenerateFalAudioParams {
   instrumental?: boolean
   /** When set, clones the voice from this reference sample (data URI) via a zero-shot clone model. */
   voiceSampleDataUri?: string
+  /** Ends the generation: the queued request is cancelled and every wait stops. */
+  signal?: AbortSignal
 }
 
 export interface GeneratedAudio {
@@ -94,11 +96,16 @@ export async function generateFalAudio(params: GenerateFalAudioParams): Promise<
       ref_audio_url: params.voiceSampleDataUri,
       model_type: 'F5-TTS',
     }
-    const { requestId, data } = await runFalQueue(model, input, apiKey)
+    const { requestId, data } = await runFalQueue(model, input, apiKey, params.signal)
     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)
-    const cost = await getFalAICostMetadata({ apiKey, endpointId: model, requestId })
+    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',
@@ -113,12 +120,17 @@ export async function generateFalAudio(params: GenerateFalAudioParams): Promise<
   const input = buildInput(type, params, model)
 
   // For fal audio models the model ID is the queue endpoint.
-  const { requestId, data } = await runFalQueue(model, input, apiKey)
+  const { requestId, data } = await runFalQueue(model, input, apiKey, params.signal)
   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)
-  const cost = await getFalAICostMetadata({ apiKey, endpointId: model, requestId })
+  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

@@ -110,6 +110,8 @@ export interface GenerateFalVideoParams {
   promptOptimizer?: boolean
   /** Optional start-frame image as a data URI; when set, routes to the model's image-to-video endpoint. */
   imageDataUri?: string
+  /** Ends the generation: the queued request is cancelled and every wait stops. */
+  signal?: AbortSignal
 }
 
 export interface GeneratedVideo {
@@ -171,13 +173,18 @@ export async function generateFalVideo(params: GenerateFalVideoParams): Promise<
     input.prompt_optimizer = params.promptOptimizer
   }
 
-  const { requestId, data } = await runFalQueue(endpoint, input, apiKey)
+  const { requestId, data } = await runFalQueue(endpoint, input, apiKey, params.signal)
   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)
-  const cost = await getFalAICostMetadata({ apiKey, endpointId: endpoint, requestId })
+  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

@@ -0,0 +1,212 @@
+import { jsonResponse } from '@sim/testing/helpers/http'
+import {
+  inputValidationMock,
+  inputValidationMockFns,
+} from '@sim/testing/mocks/input-validation.mock'
+import { getMockLogger } from '@sim/testing/mocks/logger.mock'
+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'
+
+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'
+
+interface FalCall {
+  url: string
+  method: string
+  headers: Record<string, string>
+}
+
+/**
+ * A fal queue whose request is accepted and then never finishes: every status poll reports
+ * IN_PROGRESS, or (`statusHangs`) never answers at all until its signal aborts.
+ */
+function stubStuckFalQueue(options: { statusHangs?: boolean; cancelStatus?: number } = {}) {
+  const calls: FalCall[] = []
+  vi.stubGlobal(
+    'fetch',
+    vi.fn(async (url: string, init: RequestInit = {}) => {
+      calls.push({
+        url,
+        method: init.method ?? 'GET',
+        headers: (init.headers ?? {}) as Record<string, string>,
+      })
+      if (init.method === 'POST')
+        return jsonResponse({
+          request_id: 'req-1',
+          status_url: STATUS_URL,
+          cancel_url: CANCEL_URL,
+        })
+      if (init.method === 'PUT')
+        return jsonResponse({ status: 'CANCELLATION_REQUESTED' }, options.cancelStatus ?? 202)
+      if (options.statusHangs)
+        return new Promise<Response>((_resolve, reject) => {
+          init.signal?.addEventListener('abort', () => reject(init.signal?.reason), { once: true })
+        })
+      return jsonResponse({ status: 'IN_PROGRESS' })
+    })
+  )
+  const statusPolls = () => calls.filter((call) => call.url === STATUS_URL).length
+  const cancels = () => calls.filter((call) => call.method === 'PUT')
+  return { calls, statusPolls, cancels }
+}
+
+describe('runFalQueue on a request that never finishes', () => {
+  beforeEach(() => {
+    vi.useFakeTimers()
+  })
+  afterEach(() => {
+    vi.useRealTimers()
+  })
+
+  it('stops polling and cancels the request once its caller gives up', async () => {
+    const fal = stubStuckFalQueue()
+    const caller = new AbortController()
+    const generation = runFalQueue('fal-ai/model', { prompt: 'p' }, 'key', caller.signal)
+    const outcome = generation.then(
+      () => 'resolved',
+      () => 'rejected'
+    )
+
+    await vi.advanceTimersByTimeAsync(9_500)
+    expect(fal.statusPolls()).toBe(3)
+    caller.abort(new Error('watchdog'))
+    await vi.advanceTimersByTimeAsync(0)
+
+    expect(await outcome).toBe('rejected')
+    expect(fal.cancels()).toEqual([expect.objectContaining({ url: CANCEL_URL })])
+    await vi.advanceTimersByTimeAsync(60_000)
+    expect(fal.statusPolls()).toBe(3)
+  })
+
+  it('ends a poll that never answers when its caller gives up', async () => {
+    const fal = stubStuckFalQueue({ statusHangs: true })
+    const caller = new AbortController()
+    const generation = runFalQueue('fal-ai/model', { prompt: 'p' }, 'key', caller.signal)
+    const outcome = generation.then(
+      () => 'resolved',
+      () => 'rejected'
+    )
+
+    await vi.advanceTimersByTimeAsync(3_000)
+    expect(fal.statusPolls()).toBe(1)
+    caller.abort(new Error('watchdog'))
+    await vi.advanceTimersByTimeAsync(0)
+
+    expect(await outcome).toBe('rejected')
+    expect(fal.cancels()).toHaveLength(1)
+  })
+
+  it('fails and cancels after its own ceiling even when nothing aborts it', async () => {
+    const fal = stubStuckFalQueue()
+    const generation = runFalQueue('fal-ai/model', { prompt: 'p' }, 'key')
+    const outcome = generation.then(
+      () => 'resolved',
+      (error: unknown) => String(error)
+    )
+
+    await vi.advanceTimersByTimeAsync(30 * 60 * 1000)
+
+    expect(await outcome).toContain('Fal.ai generation timed out')
+    expect(fal.cancels()).toHaveLength(1)
+  })
+
+  it('warns when fal refuses to cancel a request', async () => {
+    const fal = stubStuckFalQueue({ cancelStatus: 500 })
+    const caller = new AbortController()
+    const generation = runFalQueue('fal-ai/model', { prompt: 'p' }, 'key', caller.signal).catch(
+      () => undefined
+    )
+    await vi.advanceTimersByTimeAsync(3_500)
+    caller.abort()
+    await generation
+
+    expect(fal.cancels()).toHaveLength(1)
+    expect(getMockLogger('FalMediaClient').warn).toHaveBeenCalledWith(
+      'Fal.ai request cancellation failed',
+      expect.objectContaining({ requestId: 'req-1' })
+    )
+  })
+
+  it('submits nothing for a caller that already gave up', async () => {
+    const fal = stubStuckFalQueue()
+
+    await expect(
+      runFalQueue('fal-ai/model', { prompt: 'p' }, 'key', AbortSignal.abort())
+    ).rejects.toThrow()
+    expect(fal.calls).toEqual([])
+  })
+
+  it('cancels a request its caller gave up on while submitting it', async () => {
+    const calls: string[] = []
+    let accept: (response: Response) => void = () => {}
+    vi.stubGlobal(
+      'fetch',
+      vi.fn((url: string, init: RequestInit = {}) => {
+        calls.push(`${init.method ?? 'GET'} ${url}`)
+        if (init.method === 'POST')
+          return new Promise<Response>((resolve, reject) => {
+            accept = resolve
+            init.signal?.addEventListener('abort', () => reject(init.signal?.reason), {
+              once: true,
+            })
+          })
+        return Promise.resolve(jsonResponse({ status: 'CANCELLATION_REQUESTED' }, 202))
+      })
+    )
+    const caller = new AbortController()
+    const generation = runFalQueue('fal-ai/model', { prompt: 'p' }, 'key', caller.signal).catch(
+      () => 'rejected'
+    )
+
+    await vi.advanceTimersByTimeAsync(0)
+    caller.abort(new Error('watchdog'))
+    accept(jsonResponse({ request_id: 'req-1', cancel_url: CANCEL_URL }))
+
+    expect(await generation).toBe('rejected')
+    expect(calls).toContain(`PUT ${CANCEL_URL}`)
+  })
+
+  it('asks fal never to start the request past that ceiling', async () => {
+    const fal = stubStuckFalQueue()
+    const caller = new AbortController()
+    const generation = runFalQueue('fal-ai/model', { prompt: 'p' }, 'key', caller.signal).catch(
+      () => undefined
+    )
+    await vi.advanceTimersByTimeAsync(0)
+    caller.abort()
+    await generation
+
+    expect(fal.calls[0]).toMatchObject({
+      method: 'POST',
+      headers: expect.objectContaining({ 'X-Fal-Request-Timeout': '1800' }),
+    })
+  })
+})
+
+describe('downloadFalMedia', () => {
+  it('returns nothing for a caller that already gave up, even inline media', async () => {
+    await expect(
+      downloadFalMedia('data:audio/mpeg;base64,AAAA', AbortSignal.abort())
+    ).rejects.toThrow()
+  })
+
+  it('ends a download still resolving its host once its caller gives up', async () => {
+    inputValidationMockFns.mockValidateUrlWithDNS.mockImplementationOnce(
+      (_url: string, _name: string, _profile: string, options?: { signal?: AbortSignal }) =>
+        new Promise((_resolve, reject) => {
+          options?.signal?.addEventListener('abort', () => reject(options.signal?.reason), {
+            once: true,
+          })
+        })
+    )
+    const caller = new AbortController()
+    const download = downloadFalMedia('https://media.example/out.mp3', caller.signal)
+
+    caller.abort(new Error('watchdog'))
+
+    await expect(download).rejects.toThrow('watchdog')
+  })
+})