Reattach Comfy watchers after refresh so pending jobs finish and save.

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

server/utils/jobs.ts

@@ -118,6 +118,36 @@ export function getJob(id: string) {
   return jobs.get(id)
 }
 
+export function restoreJob(params: {
+  id: string
+  clientId: string
+  promptId: string
+  startedAt: number
+  hideThumbnail?: boolean
+  library: Job['library']
+}) {
+  const existing = jobs.get(params.id)
+  if (existing) return existing
+  const job: Job = {
+    id: params.id,
+    clientId: params.clientId,
+    kind: 'video',
+    promptId: params.promptId,
+    status: 'running',
+    message: 'Waiting for ComfyUI to finish this job...',
+    progress: 8,
+    step: 0,
+    maxStep: params.library?.steps || 0,
+    startedAt: params.startedAt,
+    hideThumbnail: params.hideThumbnail === true,
+    library: params.library,
+    events: [],
+    listeners: new Set()
+  }
+  jobs.set(job.id, job)
+  return job
+}
+
 export function listJobs() {
   return [...jobs.values()]
 }

server/utils/watch.ts

@@ -1,5 +1,6 @@
 import { existsSync } from 'node:fs'
 import type { Job, JobEvent } from '~/server/utils/jobs'
+import { getJob, restoreJob } from '~/server/utils/jobs'
 
 function classifyError(message: string) {
   const lower = message.toLowerCase()
@@ -375,3 +376,71 @@ export async function waitForComfySocket(job: Job, ms = 4000) {
     await sleep(100)
   }
 }
+
+const pendingWatches = new Set<string>()
+
+export function ensurePendingWatch(pending: {
+  jobId: string
+  promptId: string
+  clientId: string
+  ownerKey: string
+  folderId: string
+  hideThumbnail: boolean
+  folderLocked?: boolean
+  name?: string
+  prompt: string
+  aspect: string
+  width: number
+  height: number
+  steps: number
+  turbo: boolean
+  seed: number
+  startedAt: number
+  imageName?: string
+  imageSubfolder?: string
+  extendTmpDir?: string
+  extendPart1Path?: string
+  familyId?: string
+  parentClipId?: string
+  chainIndex?: number
+  stillId?: string
+  sound?: boolean
+}) {
+  const existing = getJob(pending.jobId)
+  if (existing) return existing
+  const job = restoreJob({
+    id: pending.jobId,
+    clientId: pending.clientId,
+    promptId: pending.promptId,
+    startedAt: pending.startedAt,
+    hideThumbnail: pending.hideThumbnail,
+    library: {
+      ownerKey: pending.ownerKey,
+      folderId: pending.folderId,
+      hideThumbnail: pending.hideThumbnail,
+      folderLocked: pending.folderLocked,
+      name: pending.name,
+      prompt: pending.prompt,
+      aspect: pending.aspect,
+      width: pending.width,
+      height: pending.height,
+      steps: pending.steps,
+      turbo: pending.turbo,
+      seed: pending.seed,
+      imageName: pending.imageName,
+      imageSubfolder: pending.imageSubfolder,
+      extendTmpDir: pending.extendTmpDir,
+      extendPart1Path: pending.extendPart1Path,
+      familyId: pending.familyId,
+      parentClipId: pending.parentClipId,
+      chainIndex: pending.chainIndex,
+      stillId: pending.stillId,
+      sound: pending.sound
+    }
+  })
+  if (!pendingWatches.has(job.id)) {
+    pendingWatches.add(job.id)
+    void watchComfyJob(job, { persist: true }).finally(() => pendingWatches.delete(job.id))
+  }
+  return job
+}