Keep queued extension chains running after refresh instead of saving the first clip as final.

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

server/utils/watch.ts

@@ -1,6 +1,8 @@
 import { existsSync } from 'node:fs'
 import type { Job, JobEvent } from '~/server/utils/jobs'
 import { getJob, restoreJob } from '~/server/utils/jobs'
+import type { PendingJob } from '~/server/utils/pending'
+import { isLastChainShot, libraryFromPending, pendingFromJob, writePendingJob, deletePendingJob } from '~/server/utils/pending'
 
 function classifyError(message: string) {
   const lower = message.toLowerCase()
@@ -219,12 +221,16 @@ export function watchComfyJob(job: Job, options: { persist?: boolean } = {}): Pr
             })
             if (!persist) {
               job.segmentBuffer = buffer
-              deletePendingJob(job.id)
+              job.status = 'running'
+              writePendingJob(pendingFromJob(job, {
+                promptId: '',
+                currentClipId: clip.id
+              }))
               job.saving = false
               emitLocal({
                 type: 'checkpoint',
                 message: stitching ? 'Extension checkpoint saved' : 'Initial segment saved',
-                progress: 99,
+                progress: chainMeta(job).chained ? mapChainProgress(job, 99) : 99,
                 clipId: clip.id,
                 hideThumbnail: job.hideThumbnail,
                 folderLocked: job.library.folderLocked
@@ -385,70 +391,42 @@ export async function waitForComfySocket(job: Job, ms = 4000) {
 
 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
-  referenceImageNames?: string[]
-}) {
+export function ensurePendingWatch(pending: PendingJob) {
   const existing = getJob(pending.jobId)
   if (existing) return existing
+  const lastShot = isLastChainShot(pending)
   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,
-      referenceImageNames: pending.referenceImageNames
-    }
+    library: libraryFromPending(pending)
   })
+  if (pending.currentClipId) job.clipId = pending.currentClipId
   if (!pendingWatches.has(job.id)) {
     pendingWatches.add(job.id)
-    void watchComfyJob(job, { persist: true }).finally(() => pendingWatches.delete(job.id))
+    const resume = async () => {
+      try {
+        if (pending.promptId) {
+          await watchComfyJob(job, { persist: lastShot })
+        }
+        if (job.status === 'error' || job.status === 'cancelled' || job.status === 'complete') return
+        if (!lastShot) {
+          const { continueQueuedExtensionsIfNeeded } = await import('~/server/utils/videoChain')
+          await continueQueuedExtensionsIfNeeded(job)
+        }
+      } catch (error) {
+        const message = error instanceof Error ? error.message : String(error)
+        if (job.status === 'error' || job.status === 'cancelled') return
+        job.status = 'error'
+        job.error = message
+        emitChainJob(job, { type: 'error', error: message, message })
+      } finally {
+        pendingWatches.delete(job.id)
+      }
+    }
+    void resume()
   }
   return job
 }