workflow caching cascade persistence

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'])
 

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'
@@ -93,6 +94,28 @@ export class ComfyWorkflow extends UserFile {
       throw new Error('[ASSERT] Workflow content should be loaded')
     }
 
+    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)
+      }
+    }
+
     // Note: originalContent is populated by super.load()
     this.changeTracker = markRaw(
       new ChangeTracker(
@@ -100,6 +123,12 @@ export class ComfyWorkflow extends UserFile {
         /* initialState= */ JSON.parse(this.originalContent)
       )
     )
+    if (draftState && draftContent) {
+      this.changeTracker.activeState = draftState
+      this.content = draftContent
+      this._isModified = true
+      draftStore.markDraftUsed(this.path)
+    }
     return this as this & LoadedComfyWorkflow
   }
 
@@ -109,12 +138,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 +155,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> {
@@ -565,6 +599,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 +609,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 +628,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.ts

@@ -0,0 +1,77 @@
+export interface WorkflowDraftSnapshot {
+  data: string
+  updatedAt: number
+  name: string
+  isTemporary: boolean
+}
+
+export interface DraftCacheState {
+  drafts: Record<string, WorkflowDraftSnapshot>
+  order: string[]
+}
+
+export const MAX_DRAFTS = 32
+
+export const createDraftCacheState = (
+  drafts: Record<string, WorkflowDraftSnapshot> = {},
+  order: string[] = []
+): DraftCacheState => ({ drafts, order })
+
+export const touchEntry = (order: string[], path: string): string[] => {
+  const next = order.filter((entry) => entry !== path)
+  next.push(path)
+  return next
+}
+
+export const upsertDraft = (
+  state: DraftCacheState,
+  path: string,
+  snapshot: WorkflowDraftSnapshot,
+  limit: number = MAX_DRAFTS
+): DraftCacheState => {
+  const drafts = { ...state.drafts, [path]: snapshot }
+  const order = touchEntry(state.order, path)
+
+  while (order.length > limit) {
+    const oldest = order.shift()
+    if (!oldest) continue
+    if (oldest !== path) {
+      delete drafts[oldest]
+    }
+  }
+
+  return createDraftCacheState(drafts, order)
+}
+
+export const removeDraft = (
+  state: DraftCacheState,
+  path: string
+): DraftCacheState => {
+  if (!(path in state.drafts)) return state
+  const drafts = { ...state.drafts }
+  delete drafts[path]
+  const order = state.order.filter((entry) => entry !== path)
+  return createDraftCacheState(drafts, order)
+}
+
+export const moveDraft = (
+  state: DraftCacheState,
+  oldPath: string,
+  newPath: string,
+  name: string
+): DraftCacheState => {
+  const draft = state.drafts[oldPath]
+  if (!draft) return state
+  const updatedDraft = { ...draft, name }
+  const drafts = { ...state.drafts }
+  delete drafts[oldPath]
+  drafts[newPath] = updatedDraft
+  const order = touchEntry(
+    state.order.filter((entry) => entry !== oldPath && entry !== newPath),
+    newPath
+  )
+  return createDraftCacheState(drafts, order)
+}
+
+export const mostRecentDraftPath = (order: string[]): string | null =>
+  order.length ? order[order.length - 1] : null

src/platform/workflow/persistence/composables/useWorkflowPersistence.ts

@@ -9,7 +9,11 @@ import {
 import { PRESERVED_QUERY_NAMESPACES } from '@/platform/navigation/preservedQueryNamespaces'
 import { useSettingStore } from '@/platform/settings/settingStore'
 import { useWorkflowService } from '@/platform/workflow/core/services/workflowService'
-import { useWorkflowStore } from '@/platform/workflow/management/stores/workflowStore'
+import {
+  ComfyWorkflow,
+  useWorkflowStore
+} from '@/platform/workflow/management/stores/workflowStore'
+import { useWorkflowDraftStore } from '@/platform/workflow/persistence/stores/workflowDraftStore'
 import { useTemplateUrlLoader } from '@/platform/workflow/templates/composables/useTemplateUrlLoader'
 import { api } from '@/scripts/api'
 import { app as comfyApp } from '@/scripts/app'
@@ -23,6 +27,8 @@ export function useWorkflowPersistence() {
   const router = useRouter()
   const templateUrlLoader = useTemplateUrlLoader()
   const TEMPLATE_NAMESPACE = PRESERVED_QUERY_NAMESPACES.TEMPLATE
+  const workflowDraftStore = useWorkflowDraftStore()
+
 
   const ensureTemplateQueryFromIntent = async () => {
     hydratePreservedQuery(TEMPLATE_NAMESPACE)
@@ -44,6 +50,9 @@ export function useWorkflowPersistence() {
 
   const persistCurrentWorkflow = () => {
     if (!workflowPersistenceEnabled.value) return
+    const activeWorkflow = workflowStore.activeWorkflow
+    if (!activeWorkflow) return
+
     const workflow = JSON.stringify(comfyApp.rootGraph.serialize())
 
     try {
@@ -68,33 +77,31 @@ export function useWorkflowPersistence() {
       })
       throw error
     }
-  }
 
-  const loadWorkflowFromStorage = async (
-    json: string | null,
-    workflowName: string | null
-  ) => {
-    if (!json) return false
-    const workflow = JSON.parse(json)
-    await comfyApp.loadGraphData(workflow, true, true, workflowName)
-    return true
+    if (!activeWorkflow.isTemporary && !activeWorkflow.isModified) {
+      workflowDraftStore.removeDraft(activeWorkflow.path)
+      return
+    }
+
+    workflowDraftStore.saveDraft(activeWorkflow.path, {
+      data: workflow,
+      updatedAt: Date.now(),
+      name: activeWorkflow.key,
+      isTemporary: activeWorkflow.isTemporary
+    })
   }
 
   const loadPreviousWorkflowFromStorage = async () => {
     const workflowName = getStorageValue('Comfy.PreviousWorkflow')
-    const clientId = api.initialClientId ?? api.clientId
+    const preferredPath = workflowName
+      ? `${ComfyWorkflow.basePath}${workflowName}`
+      : null
 
-    // Try loading from session storage first
-    if (clientId) {
-      const sessionWorkflow = sessionStorage.getItem(`workflow:${clientId}`)
-      if (await loadWorkflowFromStorage(sessionWorkflow, workflowName)) {
-        return true
-      }
-    }
-
-    // Fall back to local storage
-    const localWorkflow = localStorage.getItem('workflow')
-    return await loadWorkflowFromStorage(localWorkflow, workflowName)
+    return await workflowDraftStore.loadPersistedWorkflow({
+      workflowName,
+      preferredPath,
+      fallbackToLatestDraft: !workflowName
+    })
   }
 
   const loadDefaultWorkflow = async () => {
@@ -158,11 +165,12 @@ export function useWorkflowPersistence() {
       }
 
       const paths = openWorkflows.value
-        .filter((workflow) => workflow?.isPersisted)
-        .map((workflow) => workflow.path)
-      const activeIndex = openWorkflows.value.findIndex(
-        (workflow) => workflow.path === activeWorkflow.value?.path
-      )
+        .map((workflow) => workflow?.path)
+        .filter(
+          (path): path is string =>
+            typeof path === 'string' && path.startsWith(ComfyWorkflow.basePath)
+        )
+      const activeIndex = paths.indexOf(activeWorkflow.value.path)
 
       return { paths, activeIndex }
     }
@@ -171,10 +179,10 @@ export function useWorkflowPersistence() {
   // Get storage values before setting watchers
   const storedWorkflows = JSON.parse(
     getStorageValue('Comfy.OpenWorkflowsPaths') || '[]'
-  )
+  ) as string[]
   const storedActiveIndex = JSON.parse(
     getStorageValue('Comfy.ActiveWorkflowIndex') || '-1'
-  )
+  ) as number
 
   watch(restoreState, ({ paths, activeIndex }) => {
     if (workflowPersistenceEnabled.value) {
@@ -186,12 +194,19 @@ export function useWorkflowPersistence() {
   const restoreWorkflowTabsState = () => {
     if (!workflowPersistenceEnabled.value) return
     const isRestorable = storedWorkflows?.length > 0 && storedActiveIndex >= 0
-    if (isRestorable) {
-      workflowStore.openWorkflowsInBackground({
-        left: storedWorkflows.slice(0, storedActiveIndex),
-        right: storedWorkflows.slice(storedActiveIndex)
-      })
-    }
+    if (!isRestorable) return
+
+    storedWorkflows.forEach((path: string) => {
+      if (workflowStore.getWorkflowByPath(path)) return
+      const draft = workflowDraftStore.getDraft(path)
+      if (!draft?.isTemporary) return
+      workflowStore.createTemporary(draft.name)
+    })
+
+    workflowStore.openWorkflowsInBackground({
+      left: storedWorkflows.slice(0, storedActiveIndex),
+      right: storedWorkflows.slice(storedActiveIndex)
+    })
   }
 
   return {

src/platform/workflow/persistence/stores/workflowDraftStore.ts

@@ -0,0 +1,138 @@
+import { useStorage } from '@vueuse/core'
+import { defineStore } from 'pinia'
+import { computed } from 'vue'
+
+import {
+  type DraftCacheState,
+  MAX_DRAFTS,
+  type WorkflowDraftSnapshot,
+  createDraftCacheState,
+  mostRecentDraftPath,
+  moveDraft as moveDraftEntry,
+  removeDraft as removeDraftEntry,
+  touchEntry,
+  upsertDraft
+} from '@/platform/workflow/persistence/base/draftCache'
+import { api } from '@/scripts/api'
+import { app as comfyApp } from '@/scripts/app'
+
+const DRAFTS_STORAGE_KEY = 'Comfy.Workflow.Drafts'
+const ORDER_STORAGE_KEY = 'Comfy.Workflow.DraftOrder'
+
+interface LoadPersistedWorkflowOptions {
+  workflowName: string | null
+  preferredPath?: string | null
+  fallbackToLatestDraft?: boolean
+}
+
+export const useWorkflowDraftStore = defineStore('workflowDraft', () => {
+  const storedDrafts = useStorage<Record<string, WorkflowDraftSnapshot>>(
+    DRAFTS_STORAGE_KEY,
+    {}
+  )
+  const storedOrder = useStorage<string[]>(ORDER_STORAGE_KEY, [])
+
+  const mostRecentDraft = computed(() => mostRecentDraftPath(storedOrder.value))
+
+  const currentState = (): DraftCacheState =>
+    createDraftCacheState(storedDrafts.value, storedOrder.value)
+
+  const updateState = (state: DraftCacheState) => {
+    storedDrafts.value = state.drafts
+    storedOrder.value = state.order
+  }
+
+  const saveDraft = (path: string, snapshot: WorkflowDraftSnapshot) => {
+    updateState(upsertDraft(currentState(), path, snapshot, MAX_DRAFTS))
+  }
+
+  const removeDraft = (path: string) => {
+    updateState(removeDraftEntry(currentState(), path))
+  }
+
+  const moveDraft = (oldPath: string, newPath: string, name: string) => {
+    updateState(moveDraftEntry(currentState(), oldPath, newPath, name))
+  }
+
+  const markDraftUsed = (path: string) => {
+    if (!(path in storedDrafts.value)) return
+    storedOrder.value = touchEntry(storedOrder.value, path)
+  }
+
+  const getDraft = (path: string) => storedDrafts.value[path]
+
+  const tryLoadGraph = async (
+    payload: string | null,
+    workflowName: string | null,
+    onFailure?: () => void
+  ) => {
+    if (!payload) return false
+    try {
+      const workflow = JSON.parse(payload)
+      await comfyApp.loadGraphData(workflow, true, true, workflowName)
+      return true
+    } catch (err) {
+      console.error('Failed to load persisted workflow', err)
+      onFailure?.()
+      return false
+    }
+  }
+
+  const loadDraft = async (path: string) => {
+    const draft = getDraft(path)
+    if (!draft) return false
+    const loaded = await tryLoadGraph(draft.data, draft.name, () => {
+      removeDraft(path)
+    })
+    if (loaded) {
+      markDraftUsed(path)
+    }
+    return loaded
+  }
+
+  const loadPersistedWorkflow = async (
+    options: LoadPersistedWorkflowOptions
+  ): Promise<boolean> => {
+    const {
+      workflowName,
+      preferredPath,
+      fallbackToLatestDraft = false
+    } = options
+
+    if (preferredPath && (await loadDraft(preferredPath))) {
+      return true
+    }
+
+    if (!preferredPath && fallbackToLatestDraft) {
+      const fallbackPath = mostRecentDraft.value
+      if (fallbackPath && (await loadDraft(fallbackPath))) {
+        return true
+      }
+    }
+
+    const clientId = api.initialClientId ?? api.clientId
+    if (clientId) {
+      const sessionPayload = sessionStorage.getItem(`workflow:${clientId}`)
+      if (await tryLoadGraph(sessionPayload, workflowName)) {
+        return true
+      }
+    }
+
+    const localPayload = localStorage.getItem('workflow')
+    return await tryLoadGraph(localPayload, workflowName)
+  }
+
+  return {
+    saveDraft,
+    removeDraft,
+    moveDraft,
+    markDraftUsed,
+    getDraft,
+    loadPersistedWorkflow,
+    reset: () => {
+      updateState(createDraftCacheState())
+    }
+  }
+})
+
+export type { WorkflowDraftSnapshot } from '@/platform/workflow/persistence/base/draftCache'

tests-ui/tests/platform/workflow/persistence/draftCache.test.ts

@@ -0,0 +1,77 @@
+import { describe, expect, it } from 'vitest'
+
+import {
+  MAX_DRAFTS,
+  type WorkflowDraftSnapshot,
+  createDraftCacheState,
+  mostRecentDraftPath,
+  moveDraft,
+  removeDraft,
+  touchEntry,
+  upsertDraft
+} from '@/platform/workflow/persistence/base/draftCache'
+
+const createSnapshot = (name: string): WorkflowDraftSnapshot => ({
+  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()
+  })
+})

tests-ui/tests/platform/workflow/persistence/useWorkflowPersistence.test.ts

@@ -0,0 +1,197 @@
+import { createPinia, setActivePinia } from 'pinia'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+
+import { useWorkflowStore } from '@/platform/workflow/management/stores/workflowStore'
+import { useWorkflowPersistence } from '@/platform/workflow/persistence/composables/useWorkflowPersistence'
+import { useWorkflowDraftStore } from '@/platform/workflow/persistence/stores/workflowDraftStore'
+import { defaultGraphJSON } from '@/scripts/defaultGraph'
+import { setStorageValue } from '@/scripts/utils'
+
+vi.mock('@/platform/settings/settingStore', () => ({
+  useSettingStore: vi.fn(() => ({
+    get: vi.fn((key: string) =>
+      key === 'Comfy.Workflow.Persist' ? true : undefined
+    ),
+    set: vi.fn()
+  }))
+}))
+
+const loadBlankWorkflow = vi.fn()
+vi.mock('@/platform/workflow/core/services/workflowService', () => ({
+  useWorkflowService: () => ({
+    loadBlankWorkflow
+  })
+}))
+
+const executeCommand = vi.fn()
+vi.mock('@/stores/commandStore', () => ({
+  useCommandStore: () => ({
+    execute: executeCommand
+  })
+}))
+
+type GraphChangedHandler = (() => void) | null
+
+const mocks = vi.hoisted(() => {
+  const state = {
+    graphChangedHandler: null as GraphChangedHandler,
+    currentGraph: {} as Record<string, unknown>
+  }
+  const serializeMock = vi.fn(() => state.currentGraph)
+  const loadGraphDataMock = vi.fn()
+  const apiMock = {
+    clientId: 'test-client',
+    initialClientId: 'test-client',
+    addEventListener: vi.fn((event: string, handler: () => void) => {
+      if (event === 'graphChanged') {
+        state.graphChangedHandler = handler
+      }
+    }),
+    removeEventListener: vi.fn(),
+    getUserData: vi.fn(),
+    storeUserData: vi.fn(),
+    listUserDataFullInfo: vi.fn(),
+    storeSetting: vi.fn(),
+    getSettings: vi.fn(),
+    deleteUserData: vi.fn(),
+    moveUserData: vi.fn(),
+    apiURL: vi.fn((path: string) => path)
+  }
+  return { state, serializeMock, loadGraphDataMock, apiMock }
+})
+
+vi.mock('@/scripts/app', () => ({
+  app: {
+    graph: {
+      serialize: () => mocks.serializeMock()
+    },
+    loadGraphData: (...args: unknown[]) => mocks.loadGraphDataMock(...args),
+    canvas: {}
+  }
+}))
+
+vi.mock('@/scripts/api', () => ({
+  api: mocks.apiMock
+}))
+
+describe('useWorkflowPersistence', () => {
+  beforeEach(() => {
+    vi.useFakeTimers()
+    vi.setSystemTime(new Date('2025-01-01T00:00:00Z'))
+    setActivePinia(createPinia())
+    localStorage.clear()
+    sessionStorage.clear()
+    vi.clearAllMocks()
+    useWorkflowDraftStore().reset()
+    mocks.state.graphChangedHandler = null
+    mocks.state.currentGraph = { initial: true }
+    mocks.serializeMock.mockImplementation(() => mocks.state.currentGraph)
+    mocks.loadGraphDataMock.mockReset()
+    mocks.apiMock.clientId = 'test-client'
+    mocks.apiMock.initialClientId = 'test-client'
+    mocks.apiMock.addEventListener.mockImplementation(
+      (event: string, handler: () => void) => {
+        if (event === 'graphChanged') {
+          mocks.state.graphChangedHandler = handler
+        }
+      }
+    )
+    mocks.apiMock.removeEventListener.mockImplementation(() => {})
+    mocks.apiMock.listUserDataFullInfo.mockResolvedValue([])
+    mocks.apiMock.getUserData.mockResolvedValue({
+      status: 200,
+      text: () => Promise.resolve(defaultGraphJSON)
+    } as Response)
+    mocks.apiMock.apiURL.mockImplementation((path: string) => path)
+  })
+
+  afterEach(() => {
+    vi.useRealTimers()
+  })
+
+  it('persists snapshots for multiple workflows', async () => {
+    const workflowStore = useWorkflowStore()
+    const workflowA = workflowStore.createTemporary('DraftA.json')
+    await workflowStore.openWorkflow(workflowA)
+
+    const persistence = useWorkflowPersistence()
+    expect(persistence).toBeDefined()
+    expect(mocks.state.graphChangedHandler).toBeTypeOf('function')
+
+    const graphA = { title: 'A' }
+    mocks.state.currentGraph = graphA
+    mocks.state.graphChangedHandler!()
+    await vi.advanceTimersByTimeAsync(800)
+
+    const workflowB = workflowStore.createTemporary('DraftB.json')
+    await workflowStore.openWorkflow(workflowB)
+    const graphB = { title: 'B' }
+    mocks.state.currentGraph = graphB
+    mocks.state.graphChangedHandler!()
+    await vi.advanceTimersByTimeAsync(800)
+
+    const drafts = JSON.parse(
+      localStorage.getItem('Comfy.Workflow.Drafts') ?? '{}'
+    ) as Record<string, any>
+
+    expect(Object.keys(drafts)).toEqual(
+      expect.arrayContaining(['workflows/DraftA.json', 'workflows/DraftB.json'])
+    )
+    expect(JSON.parse(drafts['workflows/DraftA.json'].data)).toEqual(graphA)
+    expect(JSON.parse(drafts['workflows/DraftB.json'].data)).toEqual(graphB)
+    expect(drafts['workflows/DraftA.json'].isTemporary).toBe(true)
+    expect(drafts['workflows/DraftB.json'].isTemporary).toBe(true)
+  })
+
+  it('evicts least recently used drafts beyond the limit', async () => {
+    const workflowStore = useWorkflowStore()
+    useWorkflowPersistence()
+    expect(mocks.state.graphChangedHandler).toBeTypeOf('function')
+
+    for (let i = 0; i < 33; i++) {
+      const workflow = workflowStore.createTemporary(`Draft${i}.json`)
+      await workflowStore.openWorkflow(workflow)
+      mocks.state.currentGraph = { index: i }
+      mocks.state.graphChangedHandler!()
+      await vi.advanceTimersByTimeAsync(800)
+      vi.setSystemTime(new Date(Date.now() + 60000))
+    }
+
+    const drafts = JSON.parse(
+      localStorage.getItem('Comfy.Workflow.Drafts') ?? '{}'
+    ) as Record<string, any>
+
+    expect(Object.keys(drafts).length).toBe(32)
+    expect(drafts['workflows/Draft0.json']).toBeUndefined()
+    expect(drafts['workflows/Draft32.json']).toBeDefined()
+  })
+
+  it('restores temporary tabs from cached drafts', async () => {
+    const workflowStore = useWorkflowStore()
+    const draftStore = useWorkflowDraftStore()
+    const draftData = JSON.parse(defaultGraphJSON)
+    draftStore.saveDraft('workflows/Unsaved Workflow.json', {
+      data: JSON.stringify(draftData),
+      updatedAt: Date.now(),
+      name: 'Unsaved Workflow.json',
+      isTemporary: true
+    })
+    setStorageValue(
+      'Comfy.OpenWorkflowsPaths',
+      JSON.stringify(['workflows/Unsaved Workflow.json'])
+    )
+    setStorageValue('Comfy.ActiveWorkflowIndex', JSON.stringify(0))
+
+    const { restoreWorkflowTabsState } = useWorkflowPersistence()
+    restoreWorkflowTabsState()
+
+    const restored = workflowStore.getWorkflowByPath(
+      'workflows/Unsaved Workflow.json'
+    )
+    expect(restored).toBeTruthy()
+    expect(restored?.isTemporary).toBe(true)
+    expect(
+      workflowStore.openWorkflows.map((workflow) => workflow?.path)
+    ).toContain('workflows/Unsaved Workflow.json')
+  })
+})