Fix generation progress getting stuck at 1% after Comfy finishes.

Poll Comfy history and the job snapshot so the browser still gets the video when the live stream is buffered or the websocket misses completion.

Co-authored-by: Cursor <cursoragent@cursor.com>

nuxt.config.ts

@@ -38,6 +38,14 @@ export default defineNuxtConfig({
     }
   },
   nitro: {
-    preset: 'node-server'
+    preset: 'node-server',
+    routeRules: {
+      '/api/generate/**': {
+        headers: {
+          'Cache-Control': 'no-cache, no-store, no-transform',
+          'X-Accel-Buffering': 'no'
+        }
+      }
+    }
   }
 })

pages/index.vue

@@ -365,6 +365,8 @@ const currentPassword = ref('')
 const nextPassword = ref('')
 let timer: ReturnType<typeof setInterval> | null = null
 let source: EventSource | null = null
+let pollTimer: ReturnType<typeof setInterval> | null = null
+let settledUi = false
 
 const elapsedLabel = computed(() => {
   const seconds = Math.floor(elapsedMs.value / 1000)
@@ -593,6 +595,7 @@ async function generate() {
     return
   }
   busy.value = true
+  settledUi = false
   progress.value = 1
   videoUrl.value = ''
   awaitingReveal.value = false
@@ -629,16 +632,27 @@ async function generate() {
   }
 }
 
-function listen(id: string, hidden: boolean) {
+function stopListen() {
   source?.close()
-  source = new EventSource(`/api/generate/${id}/stream`)
+  source = null
-  source.onmessage = (event) => {
+  if (pollTimer) {
-    const payload = JSON.parse(event.data)
+    clearInterval(pollTimer)
+    pollTimer = null
+  }
+}
+
+function applyJobEvent(payload: Record<string, any>, hidden: boolean) {
   if (payload.message) statusMessage.value = payload.message
   if (typeof payload.progress === 'number') progress.value = payload.progress
   if (typeof payload.step === 'number') step.value = payload.step
   if (typeof payload.maxStep === 'number') maxStep.value = payload.maxStep
-    if (payload.type === 'complete') {
+
+  const failed = payload.type === 'error' || payload.status === 'error' || payload.status === 'cancelled'
+  const complete = payload.type === 'complete' || payload.status === 'complete'
+  if (settledUi) return
+
+  if (complete && (payload.clipId || payload.filename)) {
+    settledUi = true
     const clipUrl = payload.clipId
       ? `/api/library/clips/${payload.clipId}/video`
       : `/api/media?filename=${encodeURIComponent(payload.filename)}&subfolder=${encodeURIComponent(payload.subfolder || 'video')}&type=${encodeURIComponent(payload.mediaType || 'output')}`
@@ -646,20 +660,42 @@ function listen(id: string, hidden: boolean) {
     awaitingReveal.value = Boolean(payload.hideThumbnail || hidden)
     busy.value = false
     stopTimer()
-      source?.close()
+    stopListen()
     loadLibrary().catch(() => null)
+    return
   }
-    if (payload.type === 'error') {
+
+  if (failed) {
+    settledUi = true
     busy.value = false
     stopTimer()
-      source?.close()
+    stopListen()
     toast(payload.error || payload.message || 'Generation failed')
   }
 }
-  source.onerror = () => {
+
-    if (!busy.value) return
+function listen(id: string, hidden: boolean) {
-    toast('Progress stream disconnected. The ComfyUI job may still be running.')
+  stopListen()
+  source = new EventSource(`/api/generate/${id}/stream`)
+  source.onmessage = (event) => {
+    applyJobEvent(JSON.parse(event.data), hidden)
   }
+  const tick = async () => {
+    if (settledUi) {
+      stopListen()
+      return
+    }
+    try {
+      const snap = await $fetch<Record<string, any>>(`/api/generate/${id}`)
+      applyJobEvent(snap, hidden)
+    } catch {
+      // Job snapshot can 404 briefly after a restart; keep polling while busy.
+    }
+  }
+  void tick()
+  pollTimer = setInterval(() => {
+    void tick()
+  }, 1000)
 }
 
 async function cancelJob() {

server/api/generate.post.ts

@@ -95,9 +95,11 @@ async function runGeneration(
     length: number
   }
 ) {
+  const done = watchComfyJob(job)
   job.status = 'uploading'
   const uploaded = await uploadImage(params.image)
   emitJob(job, { type: 'status', message: 'Queueing MiniMax H3 job...', progress: 6 })
+  await waitForComfySocket(job, 4000)
 
   const graph = buildWorkflow({
     prompt: params.prompt,
@@ -114,5 +116,5 @@ async function runGeneration(
   job.promptId = queued.prompt_id
   job.status = 'running'
   emitJob(job, { type: 'status', message: 'Job queued on ComfyUI', progress: 8 })
-  await watchComfyJob(job)
+  await done
 }

server/api/generate/[id].get.ts

@@ -0,0 +1,8 @@
+export default defineEventHandler((event) => {
+  const id = getRouterParam(event, 'id')
+  const job = id ? getJob(id) : undefined
+  if (!job) {
+    throw createError({ statusCode: 404, statusMessage: 'Job not found' })
+  }
+  return jobSnapshot(job)
+})

server/api/generate/[id]/stream.get.ts

@@ -5,22 +5,18 @@ export default defineEventHandler(async (event) => {
     throw createError({ statusCode: 404, statusMessage: 'Job not found' })
   }
 
+  setResponseHeaders(event, {
+    'Cache-Control': 'no-cache, no-store, no-transform',
+    'X-Accel-Buffering': 'no',
+    Connection: 'keep-alive'
+  })
+
   const stream = createEventStream(event)
   const send = async (payload: unknown) => {
     await stream.push(JSON.stringify(payload))
   }
 
-  await send({
+  await send(jobSnapshot(job))
-    type: 'hello',
-    status: job.status,
-    message: job.message,
-    progress: job.progress,
-    step: job.step,
-    maxStep: job.maxStep,
-    promptId: job.promptId,
-    elapsedMs: Date.now() - job.startedAt
-  })
-
   for (const past of job.events) {
     await send(past)
   }
@@ -33,7 +29,12 @@ export default defineEventHandler(async (event) => {
     })
   })
 
+  const ping = setInterval(() => {
+    void send(jobSnapshot(job))
+  }, 1000)
+
   stream.onClosed(() => {
+    clearInterval(ping)
     unsubscribe()
   })
 

server/utils/comfy.ts

@@ -68,27 +68,62 @@ export async function fetchHistory(promptId: string) {
   return (await res.json()) as Record<string, unknown>
 }
 
-export function extractVideo(history: Record<string, unknown> | null, promptId: string) {
+function isVideoFile(item: { filename?: string; format?: string } | null | undefined) {
-  if (!history) return null
+  if (!item) return false
-  const entry = (history[promptId] || Object.values(history)[0]) as { outputs?: Record<string, Record<string, unknown>> } | undefined
+  const name = String(item.filename || '').toLowerCase()
-  const outputs = entry?.outputs || {}
+  const format = String(item.format || '').toLowerCase()
-  for (const node of Object.values(outputs)) {
+  return name.endsWith('.mp4') || name.endsWith('.webm') || format.includes('video') || format.includes('mp4')
-    for (const key of ['videos', 'gifs', 'images', 'files']) {
+}
-      const items = node[key]
+
-      if (!Array.isArray(items)) continue
+function normalizeVideo(item: { filename?: string; subfolder?: string; type?: string }) {
-      const match = items.find((item: { filename?: string }) => String(item.filename || '').toLowerCase().endsWith('.mp4'))
-        || items[0]
-      if (match?.filename) {
   return {
-          filename: String(match.filename),
+    filename: String(item.filename),
-          subfolder: String(match.subfolder || 'video'),
+    subfolder: String(item.subfolder || ''),
-          type: String(match.type || 'output')
+    type: String(item.type || 'output')
-        }
   }
 }
+
+function findVideo(value: unknown, depth = 0): { filename: string; subfolder: string; type: string } | null {
+  if (!value || typeof value !== 'object' || depth > 8) return null
+  if (Array.isArray(value)) {
+    const files = value.filter((item): item is { filename?: string; format?: string; subfolder?: string; type?: string } => Boolean(item && typeof item === 'object'))
+    const match = files.find(item => String(item.filename || '').toLowerCase().endsWith('.mp4'))
+      || files.find(isVideoFile)
+    if (match?.filename) return normalizeVideo(match)
+    for (const item of value) {
+      const nested = findVideo(item, depth + 1)
+      if (nested) return nested
     }
     return null
   }
+  const record = value as { filename?: string; format?: string; subfolder?: string; type?: string }
+  if (record.filename && isVideoFile(record)) return normalizeVideo(record)
+  for (const nested of Object.values(value as Record<string, unknown>)) {
+    const found = findVideo(nested, depth + 1)
+    if (found) return found
+  }
+  return null
+}
+
+export function extractVideo(history: Record<string, unknown> | null, promptId: string) {
+  if (!history) return null
+  const entry = (history[promptId] || Object.values(history)[0]) as { outputs?: Record<string, unknown> } | undefined
+  return findVideo(entry?.outputs || {})
+}
+
+export function inspectHistory(history: Record<string, unknown> | null, promptId: string) {
+  const video = extractVideo(history, promptId)
+  if (!history) return { video: null, completed: false, error: null as string | null }
+  const entry = (history[promptId] || Object.values(history)[0]) as {
+    status?: { status_str?: string; completed?: boolean; messages?: unknown[] }
+  } | undefined
+  const status = entry?.status
+  const error = status?.status_str === 'error'
+    ? 'ComfyUI reported an execution error'
+    : null
+  const completed = Boolean(video) || status?.completed === true || status?.status_str === 'success'
+  return { video, completed, error }
+}
 
 export async function probeComfy() {
   try {

server/utils/jobs.ts

@@ -42,6 +42,7 @@ export interface Job {
     thumb?: Buffer
   }
   error?: string
+  socketReady?: boolean
   events: JobEvent[]
   listeners: Set<(event: JobEvent) => void>
 }
@@ -74,6 +75,25 @@ export function getJob(id: string) {
   return jobs.get(id)
 }
 
+export function jobSnapshot(job: Job) {
+  return {
+    type: 'snapshot' as const,
+    status: job.status,
+    message: job.message,
+    progress: job.progress,
+    step: job.step,
+    maxStep: job.maxStep,
+    promptId: job.promptId,
+    elapsedMs: Date.now() - job.startedAt,
+    filename: job.video?.filename,
+    subfolder: job.video?.subfolder,
+    mediaType: job.video?.type,
+    clipId: job.clipId,
+    hideThumbnail: job.hideThumbnail,
+    error: job.error
+  }
+}
+
 export function emitJob(job: Job, event: JobEvent) {
   if (event.message) job.message = event.message
   if (typeof event.progress === 'number') job.progress = event.progress

server/utils/watch.ts

@@ -3,7 +3,7 @@ import type { Job } from '~/server/utils/jobs'
 
 function classifyError(message: string) {
   const lower = message.toLowerCase()
-  if (lower.includes('out of memory') || lower.includes('cuda') && lower.includes('alloc') || lower.includes('vram')) {
+  if (lower.includes('out of memory') || (lower.includes('cuda') && lower.includes('alloc')) || lower.includes('vram')) {
     return 'ComfyUI VRAM allocation failed. Try Turbo (8 steps) or a smaller aspect ratio.'
   }
   if (lower.includes('econnrefused') || lower.includes('unreachable') || lower.includes('fetch failed')) {
@@ -15,14 +15,30 @@ function classifyError(message: string) {
   return message
 }
 
-export async function watchComfyJob(job: Job) {
+function sleep(ms: number) {
+  return new Promise(resolve => setTimeout(resolve, ms))
+}
+
+export function watchComfyJob(job: Job): Promise<void> {
   const ws = new WebSocket(comfyWsUrl(job.clientId))
   let settled = false
+  let pollTimer: ReturnType<typeof setInterval> | null = null
+  let timeout: ReturnType<typeof setTimeout> | null = null
+
+  return new Promise((resolve) => {
+    const cleanup = () => {
+      if (timeout) clearTimeout(timeout)
+      if (pollTimer) clearInterval(pollTimer)
+      timeout = null
+      pollTimer = null
+      try { ws.close() } catch { /* ignore */ }
+    }
 
     const finish = async (error?: string) => {
       if (settled) return
       settled = true
-    try { ws.close() } catch { /* ignore */ }
+      cleanup()
+      try {
         if (error) {
           job.status = job.status === 'cancelled' ? 'cancelled' : 'error'
           job.error = classifyError(error)
@@ -59,8 +75,8 @@ export async function watchComfyJob(job: Job) {
             job.clipId = clip.id
             job.hideThumbnail = clip.hideThumbnail
             job.library.thumb = undefined
-      } catch (error) {
+          } catch (saveError) {
-        const message = error instanceof Error ? error.message : String(error)
+            const message = saveError instanceof Error ? saveError.message : String(saveError)
             job.status = 'error'
             job.error = `Video generated but library save failed: ${message}`
             emitJob(job, { type: 'error', error: job.error, message: job.error })
@@ -78,32 +94,49 @@ export async function watchComfyJob(job: Job) {
           clipId: job.clipId,
           hideThumbnail: job.hideThumbnail
         })
+      } finally {
+        resolve()
+      }
     }
 
-  const timeout = setTimeout(() => {
+    const pollHistory = async () => {
-    finish('Timed out waiting for ComfyUI (15 minutes).')
+      if (settled || !job.promptId) return
+      try {
+        const inspected = inspectHistory(await fetchHistory(job.promptId), job.promptId)
+        if (inspected.error) {
+          await finish(inspected.error)
+          return
+        }
+        if (inspected.video || inspected.completed) {
+          await finish()
+        }
+      } catch {
+        // Transient ComfyUI history misses are expected while the graph is still running.
+      }
+    }
+
+    timeout = setTimeout(() => {
+      void finish('Timed out waiting for ComfyUI (15 minutes).')
     }, 15 * 60 * 1000)
 
+    pollTimer = setInterval(() => {
+      void pollHistory()
+    }, 2000)
+
     ws.addEventListener('open', () => {
-    emitJob(job, { type: 'status', message: 'Connected to ComfyUI', progress: 8 })
+      job.socketReady = true
+      emitJob(job, { type: 'status', message: 'Connected to ComfyUI', progress: Math.max(job.progress, 4) })
     })
 
     ws.addEventListener('error', () => {
-    if (!settled) finish('ComfyUI WebSocket connection dropped.')
-  })
-
-  ws.addEventListener('close', () => {
-    if (!settled) {
-      // Fall back to history polling in case completion arrived as HTTP only.
-      setTimeout(async () => {
       if (settled) return
-        const history = await fetchHistory(job.promptId || '')
+      emitJob(job, {
-        if (extractVideo(history, job.promptId || '')) {
+        type: 'status',
-          clearTimeout(timeout)
+        message: job.promptId
-          await finish()
+          ? 'Live socket unavailable, polling ComfyUI history...'
-        }
+          : 'Waiting for ComfyUI socket...',
-      }, 1500)
+        progress: Math.max(job.progress, 4)
-    }
+      })
     })
 
     ws.addEventListener('message', async (event) => {
@@ -134,7 +167,6 @@ export async function watchComfyJob(job: Job) {
       if (payload.type === 'executing') {
         const node = data.node === null || data.node === undefined ? null : String(data.node)
         if (node === null) {
-        clearTimeout(timeout)
           await finish()
           return
         }
@@ -149,20 +181,22 @@ export async function watchComfyJob(job: Job) {
       }
 
       if (payload.type === 'execution_error') {
-      clearTimeout(timeout)
         const message = String(data.exception_message || data.message || 'ComfyUI node execution failed')
         await finish(message)
       }
 
       if (payload.type === 'execution_interrupted') {
-      clearTimeout(timeout)
         job.status = 'cancelled'
         await finish('Job interrupted.')
       }
     })
+  })
+}
 
-  return () => {
+export async function waitForComfySocket(job: Job, ms = 4000) {
-    clearTimeout(timeout)
+  const started = Date.now()
-    try { ws.close() } catch { /* ignore */ }
+  while (Date.now() - started < ms) {
+    if (job.socketReady) return
+    await sleep(100)
   }
 }