Feat: Persist all unsaved workflow tabs (#6050)

## Summary

- Keep all drafts in localStorage, mirroring the logic from VSCode.

- Fix a bug where newly created blank workflow tabs would incorrectly
restore as defaultGraph instead of blankGraph after page refresh.

Resolves https://github.com/Comfy-Org/desktop/issues/910, Resolves
https://github.com/Comfy-Org/ComfyUI_frontend/issues/4057, Fixes
https://github.com/Comfy-Org/ComfyUI_frontend/issues/3665

## Changes

### What
- Fix `restoreWorkflowTabsState` to parse and pass workflow data from
drafts when recreating temporary workflows
- Add error handling for invalid draft data with fallback to default
workflow
- Fix E2E test `should not serialize color adjustments in workflow` to
wait for workflow persistence before assertions
- Add proper validation for workflow nodes array in test assertions

### Breaking
- None

### Dependencies
- No new dependencies added

## Review Focus

1. **Workflow restoration**: Verify that blank workflows correctly
restore as blankGraph after page refresh
2. **Error handling**: Check that invalid draft data gracefully falls
back to default workflow
3. **Test coverage**: Ensure E2E test correctly waits for workflow
persistence before checking node properties
4. **Edge cases**: Test with multiple tabs, switching between tabs, and
rapid refresh scenarios

---------

Co-authored-by: Yourz <crazilou@vip.qq.com>

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

@@ -11,6 +11,7 @@ import {
   useWorkflowBookmarkStore,
   useWorkflowStore
 } from '@/platform/workflow/management/stores/workflowStore'
+import { useWorkflowDraftStore } from '@/platform/workflow/persistence/stores/workflowDraftStore'
 import { api } from '@/scripts/api'
 import { app as comfyApp } from '@/scripts/app'
 import { defaultGraph, defaultGraphJSON } from '@/scripts/defaultGraph'
@@ -66,6 +67,9 @@ describe('useWorkflowStore', () => {
     store = useWorkflowStore()
     bookmarkStore = useWorkflowBookmarkStore()
     vi.clearAllMocks()
+    localStorage.clear()
+    sessionStorage.clear()
+    useWorkflowDraftStore().reset()
 
     // Add default mock implementations
     vi.mocked(api.getUserData).mockResolvedValue({
@@ -235,6 +239,60 @@ describe('useWorkflowStore', () => {
       expect(workflow.isModified).toBe(false)
     })
 
+    it('prefers local draft snapshots when available', async () => {
+      localStorage.clear()
+      await syncRemoteWorkflows(['a.json'])
+      const workflow = store.getWorkflowByPath('workflows/a.json')!
+
+      const draftGraph = {
+        ...defaultGraph,
+        nodes: [...defaultGraph.nodes]
+      }
+
+      useWorkflowDraftStore().saveDraft(workflow.path, {
+        data: JSON.stringify(draftGraph),
+        updatedAt: Date.now(),
+        name: workflow.key,
+        isTemporary: workflow.isTemporary
+      })
+
+      vi.mocked(api.getUserData).mockResolvedValue({
+        status: 200,
+        text: () => Promise.resolve(defaultGraphJSON)
+      } as Response)
+
+      await workflow.load()
+
+      expect(workflow.isModified).toBe(true)
+      expect(workflow.changeTracker?.activeState).toEqual(draftGraph)
+    })
+
+    it('ignores stale drafts when server version is newer', async () => {
+      await syncRemoteWorkflows(['a.json'])
+      const workflow = store.getWorkflowByPath('workflows/a.json')!
+      const draftStore = useWorkflowDraftStore()
+
+      const draftSnapshot = {
+        data: JSON.stringify(defaultGraph),
+        updatedAt: Date.now(),
+        name: workflow.key,
+        isTemporary: workflow.isTemporary
+      }
+
+      draftStore.saveDraft(workflow.path, draftSnapshot)
+      workflow.lastModified = draftSnapshot.updatedAt + 1000
+
+      vi.mocked(api.getUserData).mockResolvedValue({
+        status: 200,
+        text: () => Promise.resolve(defaultGraphJSON)
+      } as Response)
+
+      await workflow.load()
+
+      expect(workflow.isModified).toBe(false)
+      expect(draftStore.getDraft(workflow.path)).toBeUndefined()
+    })
+
     it('should load and open a remote workflow', async () => {
       await syncRemoteWorkflows(['a.json', 'b.json'])
 
@@ -413,6 +471,20 @@ describe('useWorkflowStore', () => {
       expect(store.isOpen(workflow)).toBe(false)
       expect(store.getWorkflowByPath(workflow.path)).toBeNull()
     })
+
+    it('should remove draft when closing temporary workflow', async () => {
+      const workflow = store.createTemporary('test.json')
+      const draftStore = useWorkflowDraftStore()
+      draftStore.saveDraft(workflow.path, {
+        data: defaultGraphJSON,
+        updatedAt: Date.now(),
+        name: workflow.key,
+        isTemporary: true
+      })
+      expect(draftStore.getDraft(workflow.path)).toBeDefined()
+      await store.closeWorkflow(workflow)
+      expect(draftStore.getDraft(workflow.path)).toBeUndefined()
+    })
   })
 
   describe('deleteWorkflow', () => {

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

@@ -13,6 +13,7 @@ import type {
   ComfyWorkflowJSON,
   NodeId
 } from '@/platform/workflow/validation/schemas/workflowSchema'
+import { useWorkflowDraftStore } from '@/platform/workflow/persistence/stores/workflowDraftStore'
 import { useWorkflowThumbnail } from '@/renderer/core/thumbnail/useWorkflowThumbnail'
 import { api } from '@/scripts/api'
 import { app as comfyApp } from '@/scripts/app'
@@ -86,6 +87,28 @@ export class ComfyWorkflow extends UserFile {
   override async load({ force = false }: { force?: boolean } = {}): Promise<
     this & LoadedComfyWorkflow
   > {
+    const draftStore = useWorkflowDraftStore()
+    let draft = !force ? draftStore.getDraft(this.path) : undefined
+    let draftState: ComfyWorkflowJSON | null = null
+    let draftContent: string | null = null
+
+    if (draft) {
+      if (draft.updatedAt < this.lastModified) {
+        draftStore.removeDraft(this.path)
+        draft = undefined
+      }
+    }
+
+    if (draft) {
+      try {
+        draftState = JSON.parse(draft.data)
+        draftContent = draft.data
+      } catch (err) {
+        console.warn('Failed to parse workflow draft, clearing it', err)
+        draftStore.removeDraft(this.path)
+      }
+    }
+
     await super.load({ force })
     if (!force && this.isLoaded) return this as this & LoadedComfyWorkflow
 
@@ -93,13 +116,14 @@ export class ComfyWorkflow extends UserFile {
       throw new Error('[ASSERT] Workflow content should be loaded')
     }
 
-    // Note: originalContent is populated by super.load()
-    this.changeTracker = markRaw(
-      new ChangeTracker(
-        this,
-        /* initialState= */ JSON.parse(this.originalContent)
-      )
-    )
+    const initialState = JSON.parse(this.originalContent)
+    this.changeTracker = markRaw(new ChangeTracker(this, initialState))
+    if (draftState && draftContent) {
+      this.changeTracker.activeState = draftState
+      this.content = draftContent
+      this._isModified = true
+      draftStore.markDraftUsed(this.path)
+    }
     return this as this & LoadedComfyWorkflow
   }
 
@@ -109,12 +133,14 @@ export class ComfyWorkflow extends UserFile {
   }
 
   override async save() {
+    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 })
     this.changeTracker?.reset()
     this.isModified = false
+    draftStore.removeDraft(this.path)
     return ret
   }
 
@@ -124,8 +150,11 @@ export class ComfyWorkflow extends UserFile {
    * @returns this
    */
   override async saveAs(path: string) {
+    const draftStore = useWorkflowDraftStore()
     this.content = JSON.stringify(this.activeState)
-    return await super.saveAs(path)
+    const result = await super.saveAs(path)
+    draftStore.removeDraft(path)
+    return result
   }
 
   async promptSave(): Promise<string | null> {
@@ -436,6 +465,8 @@ export const useWorkflowStore = defineStore('workflow', () => {
     if (workflow.isTemporary) {
       // Clear thumbnail when temporary workflow is closed
       clearThumbnail(workflow.key)
+      // Clear draft when unsaved workflow tab is closed
+      useWorkflowDraftStore().removeDraft(workflow.path)
       delete workflowLookup.value[workflow.path]
     } else {
       workflow.unload()
@@ -565,6 +596,7 @@ export const useWorkflowStore = defineStore('workflow', () => {
       const oldPath = workflow.path
       const oldKey = workflow.key
       const wasBookmarked = bookmarkStore.isBookmarked(oldPath)
+      const draftStore = useWorkflowDraftStore()
 
       const openIndex = detachWorkflow(workflow)
       // Perform the actual rename operation first
@@ -574,6 +606,8 @@ export const useWorkflowStore = defineStore('workflow', () => {
         attachWorkflow(workflow, openIndex)
       }
 
+      draftStore.moveDraft(oldPath, newPath, workflow.key)
+
       // Move thumbnail from old key to new key (using workflow keys, not full paths)
       const newKey = workflow.key
       moveWorkflowThumbnail(oldKey, newKey)
@@ -591,6 +625,7 @@ export const useWorkflowStore = defineStore('workflow', () => {
     isBusy.value = true
     try {
       await workflow.delete()
+      useWorkflowDraftStore().removeDraft(workflow.path)
       if (bookmarkStore.isBookmarked(workflow.path)) {
         await bookmarkStore.setBookmarked(workflow.path, false)
       }