Feat: Persist all unsaved workflow tabs

browser_tests/tests/colorPalette.spec.ts

@@ -232,11 +232,25 @@ test.describe('Node Color Adjustments', () => {
   }) => {
     await comfyPage.setSetting('Comfy.Node.Opacity', 0.5)
     await comfyPage.setSetting('Comfy.ColorPalette', 'light')
-    const saveWorkflowInterval = 1000
-    const workflow = await comfyPage.page.evaluate(() => {
-      return localStorage.getItem('workflow')
-    })
-    for (const node of JSON.parse(workflow ?? '{}').nodes) {
+    await comfyPage.nextFrame()
+    const parsed = await (
+      await comfyPage.page.waitForFunction(
+        () => {
+          const workflow = localStorage.getItem('workflow')
+          if (!workflow) return null
+          try {
+            const data = JSON.parse(workflow)
+            return Array.isArray(data?.nodes) ? data : null
+          } catch {
+            return null
+          }
+        },
+        { timeout: 3000 }
+      )
+    ).jsonValue()
+    expect(parsed.nodes).toBeDefined()
+    expect(Array.isArray(parsed.nodes)).toBe(true)
+    for (const node of parsed.nodes) {
       if (node.bgcolor) expect(node.bgcolor).not.toMatch(/hsla/)
       if (node.color) expect(node.color).not.toMatch(/hsla/)
     }

src/locales/en/main.json

@@ -1749,6 +1749,7 @@
     "nothingToQueue": "Nothing to queue",
     "pleaseSelectOutputNodes": "Please select output nodes",
     "failedToQueue": "Failed to queue",
+    "failedToSaveDraft": "Failed to save workflow draft",
     "failedExecutionPathResolution": "Could not resolve path to selected nodes",
     "no3dScene": "No 3D scene to apply texture",
     "failedToApplyTexture": "Failed to apply texture",

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

@@ -6,6 +6,7 @@ import { LGraph, LGraphCanvas } from '@/lib/litegraph/src/litegraph'
 import type { Point, SerialisableGraph } from '@/lib/litegraph/src/litegraph'
 import { useSettingStore } from '@/platform/settings/settingStore'
 import { useToastStore } from '@/platform/updates/common/toastStore'
+import { useWorkflowDraftStore } from '@/platform/workflow/persistence/stores/workflowDraftStore'
 import {
   ComfyWorkflow,
   useWorkflowStore
@@ -28,6 +29,7 @@ export const useWorkflowService = () => {
   const dialogService = useDialogService()
   const workflowThumbnail = useWorkflowThumbnail()
   const domWidgetStore = useDomWidgetStore()
+  const workflowDraftStore = useWorkflowDraftStore()
 
   async function getFilename(defaultName: string): Promise<string | null> {
     if (settingStore.get('Comfy.PromptFilename')) {
@@ -291,6 +293,27 @@ export const useWorkflowService = () => {
     const activeWorkflow = workflowStore.activeWorkflow
     if (activeWorkflow) {
       activeWorkflow.changeTracker.store()
+      if (settingStore.get('Comfy.Workflow.Persist') && activeWorkflow.path) {
+        const activeState = activeWorkflow.activeState
+        if (activeState) {
+          try {
+            const workflowJson = JSON.stringify(activeState)
+            workflowDraftStore.saveDraft(activeWorkflow.path, {
+              data: workflowJson,
+              updatedAt: Date.now(),
+              name: activeWorkflow.key,
+              isTemporary: activeWorkflow.isTemporary
+            })
+          } catch {
+            toastStore.add({
+              severity: 'error',
+              summary: t('g.error'),
+              detail: t('toastMessages.failedToSaveDraft'),
+              life: 3000
+            })
+          }
+        }
+      }
       // Capture thumbnail before loading new graph
       void workflowThumbnail.storeThumbnail(activeWorkflow)
       domWidgetStore.clear()

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,20 +87,43 @@ 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
 
     if (!this.originalContent) {
       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)
       }

src/platform/workflow/persistence/base/draftCache.test.ts

@@ -0,0 +1,79 @@
+import { describe, expect, it } from 'vitest'
+
+import {
+  MAX_DRAFTS,
+  createDraftCacheState,
+  mostRecentDraftPath,
+  moveDraft,
+  removeDraft,
+  touchEntry,
+  upsertDraft
+} from './draftCache'
+import type { WorkflowDraftSnapshot } from './draftCache'
+
+function createSnapshot(name: string): WorkflowDraftSnapshot {
+  return {
+    data: JSON.stringify({ name }),
+    updatedAt: Date.now(),
+    name,
+    isTemporary: true
+  }
+}
+
+describe('draftCache helpers', () => {
+  it('touchEntry moves path to end', () => {
+    expect(touchEntry(['a', 'b'], 'a')).toEqual(['b', 'a'])
+    expect(touchEntry(['a', 'b'], 'c')).toEqual(['a', 'b', 'c'])
+  })
+
+  it('upsertDraft stores snapshot and applies LRU', () => {
+    let state = createDraftCacheState()
+    for (let i = 0; i < MAX_DRAFTS; i++) {
+      const path = `workflows/Draft${i}.json`
+      state = upsertDraft(state, path, createSnapshot(String(i)))
+    }
+
+    expect(Object.keys(state.drafts).length).toBe(MAX_DRAFTS)
+
+    state = upsertDraft(state, 'workflows/New.json', createSnapshot('new'))
+    expect(Object.keys(state.drafts).length).toBe(MAX_DRAFTS)
+    expect(state.drafts).not.toHaveProperty('workflows/Draft0.json')
+    expect(state.order[state.order.length - 1]).toBe('workflows/New.json')
+  })
+
+  it('removeDraft clears entry and order', () => {
+    const state = upsertDraft(
+      createDraftCacheState(),
+      'workflows/test.json',
+      createSnapshot('test')
+    )
+
+    const nextState = removeDraft(state, 'workflows/test.json')
+    expect(nextState.drafts).toEqual({})
+    expect(nextState.order).toEqual([])
+  })
+
+  it('moveDraft renames entry and updates order', () => {
+    const state = upsertDraft(
+      createDraftCacheState(),
+      'workflows/old.json',
+      createSnapshot('old')
+    )
+
+    const nextState = moveDraft(
+      state,
+      'workflows/old.json',
+      'workflows/new.json',
+      'new'
+    )
+    expect(nextState.drafts).not.toHaveProperty('workflows/old.json')
+    expect(nextState.drafts['workflows/new.json']?.name).toBe('new')
+    expect(nextState.order).toEqual(['workflows/new.json'])
+  })
+
+  it('mostRecentDraftPath returns last entry', () => {
+    const state = createDraftCacheState({}, ['a', 'b', 'c'])
+    expect(mostRecentDraftPath(state.order)).toBe('c')
+    expect(mostRecentDraftPath([])).toBeNull()
+  })
+})