fix: Implement atomic workflow overwrite via persistence-level flag (#9355)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

src/platform/workflow/core/services/workflowService.test.ts

@@ -99,7 +99,9 @@ vi.mock('@/renderer/core/canvas/canvasStore', () => ({
 vi.mock('@/renderer/core/thumbnail/useWorkflowThumbnail', () => ({
   useWorkflowThumbnail: () => ({
     storeThumbnail: vi.fn(),
-    getThumbnail: vi.fn()
+    getThumbnail: vi.fn(),
+    clearThumbnail: vi.fn(),
+    moveWorkflowThumbnail: vi.fn()
   })
 }))
 
@@ -853,12 +855,13 @@ describe('useWorkflowService', () => {
 
       const existing = createSaveableWorkflow('workflows/test.app.json')
       vi.spyOn(workflowStore, 'getWorkflowByPath').mockReturnValue(existing)
-      vi.spyOn(workflowStore, 'deleteWorkflow').mockResolvedValue()
+      vi.spyOn(workflowStore, 'removeWorkflowEntry').mockResolvedValue()
       mockConfirm.mockResolvedValue(true)
 
       await service.saveWorkflow(workflow)
 
       expect(mockConfirm).toHaveBeenCalled()
+      expect(workflowStore.removeWorkflowEntry).toHaveBeenCalledWith(existing)
       expect(workflowStore.renameWorkflow).toHaveBeenCalledWith(
         workflow,
         'workflows/test.app.json'

src/platform/workflow/core/services/workflowService.ts

@@ -130,11 +130,20 @@ export const useWorkflowService = () => {
 
     if (existingWorkflow && !existingWorkflow.isTemporary) {
       if ((await confirmOverwrite(newPath)) !== true) return false
+    }
 
-      if (!isSelfOverwrite) {
-        const deleted = await deleteWorkflow(existingWorkflow, true)
-        if (!deleted) return false
+    const needsOverwrite =
+      !!existingWorkflow && !existingWorkflow.isTemporary && !isSelfOverwrite
+
+    // Close and remove the old workflow entry before saving the new content.
+    // The file on disk is intentionally kept so that a save failure does not
+    // cause data loss.  The subsequent save with overwrite: true will
+    // atomically replace it.
+    if (needsOverwrite) {
+      if (workflowStore.isOpen(existingWorkflow)) {
+        await closeWorkflow(existingWorkflow, { warnIfUnsaved: false })
       }
+      await workflowStore.removeWorkflowEntry(existingWorkflow)
     }
 
     workflow.changeTracker?.checkState()
@@ -143,11 +152,19 @@ export const useWorkflowService = () => {
       await saveWorkflow(workflow)
     } else if (workflow.isTemporary) {
       await renameWorkflow(workflow, newPath)
-      await workflowStore.saveWorkflow(workflow)
+      if (needsOverwrite) {
+        await workflowStore.saveWorkflow(workflow, { overwrite: true })
+      } else {
+        await workflowStore.saveWorkflow(workflow)
+      }
     } else {
       const tempWorkflow = workflowStore.saveAs(workflow, newPath)
       await openWorkflow(tempWorkflow)
-      await workflowStore.saveWorkflow(tempWorkflow)
+      if (needsOverwrite) {
+        await workflowStore.saveWorkflow(tempWorkflow, { overwrite: true })
+      } else {
+        await workflowStore.saveWorkflow(tempWorkflow)
+      }
     }
     return true
   }
@@ -174,7 +191,12 @@ export const useWorkflowService = () => {
             await workflowStore.saveWorkflow(workflow)
             return
           }
-          await deleteWorkflow(existing, true)
+          // Remove the old entry without deleting the file; the rename
+          // will atomically replace it, preventing data loss on failure.
+          if (workflowStore.isOpen(existing)) {
+            await closeWorkflow(existing, { warnIfUnsaved: false })
+          }
+          await workflowStore.removeWorkflowEntry(existing)
         }
         await renameWorkflow(workflow, expectedPath)
         toastStore.add({

src/platform/workflow/management/stores/comfyWorkflow.ts

@@ -151,14 +151,16 @@ export class ComfyWorkflow extends UserFile {
     super.unload()
   }
 
-  override async save() {
+  override async save({
+    overwrite
+  }: { force?: boolean; overwrite?: boolean } = {}) {
     const { useWorkflowDraftStore } =
       await import('@/platform/workflow/persistence/stores/workflowDraftStore')
     const draftStore = useWorkflowDraftStore()
     this.content = JSON.stringify(this.activeState)
     // Force save to ensure the content is updated in remote storage incase
     // the isModified state is screwed by changeTracker.
-    const ret = await super.save({ force: true })
+    const ret = await super.save({ force: true, overwrite })
     this.changeTracker?.reset()
     this.isModified = false
     draftStore.removeDraft(this.path)

src/platform/workflow/management/stores/workflowStore.ts

@@ -63,7 +63,11 @@ interface WorkflowStore {
   ) => ComfyWorkflow
   renameWorkflow: (workflow: ComfyWorkflow, newPath: string) => Promise<void>
   deleteWorkflow: (workflow: ComfyWorkflow) => Promise<void>
-  saveWorkflow: (workflow: ComfyWorkflow) => Promise<void>
+  removeWorkflowEntry: (workflow: ComfyWorkflow) => Promise<void>
+  saveWorkflow: (
+    workflow: ComfyWorkflow,
+    options?: { overwrite?: boolean }
+  ) => Promise<void>
 
   workflows: ComfyWorkflow[]
   bookmarkedWorkflows: ComfyWorkflow[]
@@ -539,14 +543,32 @@ export const useWorkflowStore = defineStore('workflow', () => {
     }
   }
 
+  /**
+   * Remove a workflow entry from the store without deleting the file on disk.
+   * Used during atomic overwrite to clear the old entry before saving the new
+   * content, so that the save can use overwrite: true to replace the file.
+   */
+  const removeWorkflowEntry = async (workflow: ComfyWorkflow) => {
+    useWorkflowDraftStore().removeDraft(workflow.path)
+    if (bookmarkStore.isBookmarked(workflow.path)) {
+      await bookmarkStore.setBookmarked(workflow.path, false)
+    }
+    clearThumbnail(workflow.key)
+    delete workflowLookup.value[workflow.path]
+  }
+
   /**
    * Save a workflow.
    * @param workflow The workflow to save.
+   * @param options.overwrite Force overwrite of existing file at the path.
    */
-  const saveWorkflow = async (workflow: ComfyWorkflow) => {
+  const saveWorkflow = async (
+    workflow: ComfyWorkflow,
+    options?: { overwrite?: boolean }
+  ) => {
     isBusy.value = true
     try {
-      await workflow.save()
+      await workflow.save({ overwrite: options?.overwrite })
       // Synchronously detach and re-attach to force refresh the tree objects
       // without an async gap that would cause the tab to disappear.
       const openIndex = detachWorkflow(workflow)
@@ -774,6 +796,7 @@ export const useWorkflowStore = defineStore('workflow', () => {
     createNewTemporary,
     renameWorkflow,
     deleteWorkflow,
+    removeWorkflowEntry,
     saveAs,
     saveWorkflow,
     reorderWorkflows,

src/stores/nodeOutputStore.test.ts

@@ -7,11 +7,18 @@ import { LiteGraph } from '@/lib/litegraph/src/litegraph'
 import type { ExecutedWsMessage } from '@/schemas/apiSchema'
 import { app } from '@/scripts/app'
 import { useNodeOutputStore } from '@/stores/nodeOutputStore'
+import * as litegraphUtil from '@/utils/litegraphUtil'
+
+vi.mock('@/utils/litegraphUtil', () => ({
+  isAnimatedOutput: vi.fn(),
+  isVideoNode: vi.fn()
+}))
 
 const mockGetNodeById = vi.fn()
 
 vi.mock('@/scripts/app', () => ({
   app: {
+    getPreviewFormatParam: vi.fn(() => '&format=test_webp'),
     rootGraph: {
       getNodeById: (...args: unknown[]) => mockGetNodeById(...args)
     },
@@ -144,6 +151,76 @@ describe('nodeOutputStore restoreOutputs', () => {
   })
 })
 
+describe('nodeOutputStore getPreviewParam', () => {
+  beforeEach(() => {
+    setActivePinia(createTestingPinia({ stubActions: false }))
+    vi.clearAllMocks()
+    vi.mocked(litegraphUtil.isAnimatedOutput).mockReturnValue(false)
+    vi.mocked(litegraphUtil.isVideoNode).mockReturnValue(false)
+  })
+
+  it('should return empty string if output is animated', () => {
+    const store = useNodeOutputStore()
+    vi.mocked(litegraphUtil.isAnimatedOutput).mockReturnValue(true)
+    const node = createMockNode()
+    const outputs = createMockOutputs([{ filename: 'img.png' }])
+    expect(store.getPreviewParam(node, outputs)).toBe('')
+    expect(vi.mocked(app).getPreviewFormatParam).not.toHaveBeenCalled()
+  })
+
+  it('should return empty string if isVideoNode returns true', () => {
+    const store = useNodeOutputStore()
+    vi.mocked(litegraphUtil.isVideoNode).mockReturnValue(true)
+    const node = createMockNode()
+    const outputs = createMockOutputs([{ filename: 'img.png' }])
+    expect(store.getPreviewParam(node, outputs)).toBe('')
+    expect(vi.mocked(app).getPreviewFormatParam).not.toHaveBeenCalled()
+  })
+
+  it('should return empty string if outputs.images is undefined', () => {
+    const store = useNodeOutputStore()
+    const node = createMockNode()
+    const outputs: ExecutedWsMessage['output'] = {}
+    expect(store.getPreviewParam(node, outputs)).toBe('')
+    expect(vi.mocked(app).getPreviewFormatParam).not.toHaveBeenCalled()
+  })
+
+  it('should return empty string if outputs.images is empty', () => {
+    const store = useNodeOutputStore()
+    const node = createMockNode()
+    const outputs = createMockOutputs([])
+    expect(store.getPreviewParam(node, outputs)).toBe('')
+    expect(vi.mocked(app).getPreviewFormatParam).not.toHaveBeenCalled()
+  })
+
+  it('should return empty string if outputs.images contains SVG images', () => {
+    const store = useNodeOutputStore()
+    const node = createMockNode()
+    const outputs = createMockOutputs([{ filename: 'img.svg' }])
+    expect(store.getPreviewParam(node, outputs)).toBe('')
+    expect(vi.mocked(app).getPreviewFormatParam).not.toHaveBeenCalled()
+  })
+
+  it('should return format param for standard image outputs', () => {
+    const store = useNodeOutputStore()
+    const node = createMockNode()
+    const outputs = createMockOutputs([{ filename: 'img.png' }])
+    expect(store.getPreviewParam(node, outputs)).toBe('&format=test_webp')
+    expect(vi.mocked(app).getPreviewFormatParam).toHaveBeenCalledTimes(1)
+  })
+
+  it('should return format param for multiple standard images', () => {
+    const store = useNodeOutputStore()
+    const node = createMockNode()
+    const outputs = createMockOutputs([
+      { filename: 'img1.png' },
+      { filename: 'img2.jpg' }
+    ])
+    expect(store.getPreviewParam(node, outputs)).toBe('&format=test_webp')
+    expect(vi.mocked(app).getPreviewFormatParam).toHaveBeenCalledTimes(1)
+  })
+})
+
 describe('nodeOutputStore syncLegacyNodeImgs', () => {
   beforeEach(() => {
     setActivePinia(createTestingPinia({ stubActions: false }))

src/stores/nodeOutputStore.ts

@@ -11,8 +11,6 @@ import type {
   ResultItemType
 } from '@/schemas/apiSchema'
 import { appendCloudResParam } from '@/platform/distribution/cloudPreviewUtil'
-import { isCloud } from '@/platform/distribution/types'
-import { useSettingStore } from '@/platform/settings/settingStore'
 import { api } from '@/scripts/api'
 import { app } from '@/scripts/app'
 import { clone } from '@/scripts/utils'
@@ -99,6 +97,20 @@ export const useNodeOutputStore = defineStore('nodeOutput', () => {
     return true
   }
 
+  /**
+   * Get the preview param for the node's outputs.
+   *
+   * If the output is an image, use the user's preferred format (from settings).
+   * For non-image outputs, return an empty string, as including the preview param
+   * will force the server to load the output file as an image.
+   */
+  function getPreviewParam(
+    node: LGraphNode,
+    outputs: ExecutedWsMessage['output']
+  ): string {
+    return isImageOutputs(node, outputs) ? app.getPreviewFormatParam() : ''
+  }
+
   function getNodeImageUrls(node: LGraphNode): string[] | undefined {
     const previews = getNodePreviews(node)
     if (previews?.length) return previews
@@ -106,17 +118,14 @@ export const useNodeOutputStore = defineStore('nodeOutput', () => {
     const outputs = getNodeOutputs(node)
     if (!outputs?.images?.length) return
 
+    const rand = app.getRandParam()
+    const previewParam = getPreviewParam(node, outputs)
     const isImage = isImageOutputs(node, outputs)
 
     return outputs.images.map((image) => {
       const params = new URLSearchParams(image)
-      if (isImage) {
-        appendCloudResParam(params, image.filename)
-        const previewFormat = useSettingStore().get('Comfy.PreviewFormat')
-        if (previewFormat) params.set('preview', previewFormat)
-      }
-      if (!isCloud) params.set('rand', String(Math.random()))
-      return api.apiURL(`/view?${params}`)
+      if (isImage) appendCloudResParam(params, image.filename)
+      return api.apiURL(`/view?${params}${previewParam}${rand}`)
     })
   }
 
@@ -434,6 +443,8 @@ export const useNodeOutputStore = defineStore('nodeOutput', () => {
     getNodeOutputs,
     getNodeImageUrls,
     getNodePreviews,
+    getPreviewParam,
+
     // Setters
     setNodeOutputs,
     setNodeOutputsByExecutionId,

src/stores/userFileStore.ts

@@ -140,11 +140,14 @@ export class UserFile {
    * Saves the file to the remote storage.
    * @param force Whether to force the save even if the file is not modified.
    */
-  async save({ force = false }: { force?: boolean } = {}): Promise<UserFile> {
+  async save({
+    force = false,
+    overwrite
+  }: { force?: boolean; overwrite?: boolean } = {}): Promise<UserFile> {
     if (this.isPersisted && !this.isModified && !force) return this
 
     const resp = await api.storeUserData(this.path, this.content, {
-      overwrite: this.isPersisted,
+      overwrite: overwrite ?? this.isPersisted,
       throwOnError: true,
       full_info: true
     })