fix: harden workflow draft v2 shadow writes

browser_tests/tests/workflowPersistence.spec.ts

@@ -112,17 +112,21 @@ test.describe('Workflow Persistence', () => {
       await workflowsTab.open()
       await workflowsTab.getPersistedItem('restore-a').dblclick()
       await comfyPage.workflow.waitForWorkflowIdle()
+      const restoreANodeCount = await comfyPage.nodeOps.getNodeCount()
       await workflowsTab.getPersistedItem('restore-b').dblclick()
       await comfyPage.workflow.waitForWorkflowIdle()
 
       await expect.poll(() => comfyPage.nodeOps.getNodeCount()).toBe(1)
+      await expect
+        .poll(() => comfyPage.menu.topbar.getActiveTabName())
+        .toBe('restore-b')
       const tabNamesBeforeReload = await comfyPage.menu.topbar.getTabNames()
       expect(tabNamesBeforeReload).toEqual(
         expect.arrayContaining(['restore-a', 'restore-b'])
       )
+      expect(tabNamesBeforeReload.slice(-2)).toEqual(['restore-a', 'restore-b'])
       const activeTabBeforeReload =
         await comfyPage.menu.topbar.getActiveTabName()
-      expect(activeTabBeforeReload).toBe('restore-b')
 
       await comfyPage.workflow.reloadAndWaitForApp()
 
@@ -133,6 +137,12 @@ test.describe('Workflow Persistence', () => {
         .poll(() => comfyPage.menu.topbar.getActiveTabName())
         .toBe(activeTabBeforeReload)
       await expect.poll(() => comfyPage.nodeOps.getNodeCount()).toBe(1)
+
+      await comfyPage.menu.topbar.getWorkflowTab('restore-a').click()
+      await comfyPage.workflow.waitForWorkflowIdle()
+      await expect
+        .poll(() => comfyPage.nodeOps.getNodeCount())
+        .toBe(restoreANodeCount)
     }
   )
 

src/platform/workflow/persistence/composables/useWorkflowPersistenceV2.test.ts

@@ -1,7 +1,7 @@
 import { createTestingPinia } from '@pinia/testing'
 import { setActivePinia } from 'pinia'
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
-import { createApp, defineComponent } from 'vue'
+import { createApp, defineComponent, nextTick } from 'vue'
 import { createI18n } from 'vue-i18n'
 
 import { useWorkflowStore } from '@/platform/workflow/management/stores/workflowStore'
@@ -391,6 +391,30 @@ describe('useWorkflowPersistenceV2', () => {
       expect(mocks.apiMock.listUserDataFullInfo).not.toHaveBeenCalled()
     })
 
+    it('does not overwrite stored tab state before tab restore runs', async () => {
+      const workflowStore = useWorkflowStore()
+      writeTabState(['workflows/A.json'], 0)
+
+      const { initializeWorkflow } = mountWorkflowPersistence()
+      await initializeWorkflow()
+
+      const workflow = await workflowStore
+        .createTemporary('Current.json')
+        .load()
+      workflowStore.attachWorkflow(workflow, 0)
+      workflowStore.activeWorkflow = workflow
+      await nextTick()
+
+      expect(
+        JSON.parse(
+          sessionStorage.getItem('Comfy.Workflow.OpenPaths:test-client')!
+        )
+      ).toMatchObject({
+        paths: ['workflows/A.json'],
+        activeIndex: 0
+      })
+    })
+
     it('loads the onboarding blank workflow when persistence is disabled before tutorial completion', async () => {
       settingMocks.persistRef!.value = false
       settingMocks.tutorialCompletedRef!.value = false

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

@@ -281,8 +281,9 @@ export function useWorkflowPersistenceV2() {
     }
   )
 
-  // Track whether tab state has been properly restored to avoid
-  // overwriting with stale data during initialization
+  // initializeWorkflow() may intentionally return early when a valid tab
+  // pointer exists. Keep tab-state writes disabled until
+  // restoreWorkflowTabsState() has consumed that pointer.
   let tabStateRestored = false
 
   watch(restoreState, ({ paths, activeIndex }) => {

src/platform/workflow/persistence/stores/workflowDraftStoreV2.test.ts

@@ -20,6 +20,19 @@ vi.mock('@/scripts/app', () => ({
 }))
 
 describe('workflowDraftStoreV2', () => {
+  function failV2IndexWrites() {
+    const originalSetItem = localStorage.setItem.bind(localStorage)
+    return vi
+      .spyOn(localStorage, 'setItem')
+      .mockImplementation((key: string, value: string) => {
+        if (key.startsWith('Comfy.Workflow.DraftIndex.v2:')) {
+          throw new DOMException('Quota exceeded', 'QuotaExceededError')
+        }
+
+        return originalSetItem(key, value)
+      })
+  }
+
   beforeEach(() => {
     setActivePinia(createTestingPinia({ stubActions: false }))
     localStorage.clear()
@@ -30,6 +43,8 @@ describe('workflowDraftStoreV2', () => {
   afterEach(() => {
     localStorage.clear()
     sessionStorage.clear()
+    vi.useRealTimers()
+    vi.restoreAllMocks()
   })
 
   describe('saveDraft', () => {
@@ -76,6 +91,32 @@ describe('workflowDraftStoreV2', () => {
       })
     })
 
+    it('uses the same timestamp for V2 and legacy shadow-writes', () => {
+      vi.useFakeTimers()
+      vi.setSystemTime(new Date('2025-01-01T00:00:00Z'))
+      const store = useWorkflowDraftStoreV2()
+      const path = 'workflows/test.json'
+
+      store.saveDraft(path, '{"source":"v2"}', {
+        name: 'test',
+        isTemporary: true
+      })
+
+      const v2Draft = store.getDraft(path)
+      const legacyDraft = useWorkflowDraftStore().getDraft(path)
+      expect(v2Draft?.updatedAt).toBe(Date.now())
+      expect(legacyDraft?.updatedAt).toBe(Date.now())
+
+      useWorkflowDraftStore().saveDraft(path, {
+        data: '{"source":"legacy"}',
+        updatedAt: Date.now(),
+        name: 'test',
+        isTemporary: true
+      })
+
+      expect(store.getDraft(path)?.data).toBe('{"source":"v2"}')
+    })
+
     it('updates existing draft', () => {
       const store = useWorkflowDraftStoreV2()
 
@@ -112,6 +153,40 @@ describe('workflowDraftStoreV2', () => {
       expect(store.getDraft('workflows/test.json')?.data).toBe('{}')
     })
 
+    it('returns false without shadow-writing legacy when V2 save fails', () => {
+      const store = useWorkflowDraftStoreV2()
+      const path = 'workflows/test.json'
+
+      const setItemSpy = failV2IndexWrites()
+      const result = store.saveDraft(path, '{}', {
+        name: 'test',
+        isTemporary: true
+      })
+      setItemSpy.mockRestore()
+
+      expect(result).toBe(false)
+      expect(store.getDraft(path)).toBeNull()
+      expect(useWorkflowDraftStore().getDraft(path)).toBeUndefined()
+    })
+
+    it('does not cache a draft index entry when storage write fails', () => {
+      const store = useWorkflowDraftStoreV2()
+      store.saveDraft('workflows/old.json', '{"id":"old"}', {
+        name: 'old',
+        isTemporary: true
+      })
+
+      const setItemSpy = failV2IndexWrites()
+      const result = store.saveDraft('workflows/new.json', '{"id":"new"}', {
+        name: 'new',
+        isTemporary: true
+      })
+      setItemSpy.mockRestore()
+
+      expect(result).toBe(false)
+      expect(store.getMostRecentPath()).toBe('workflows/old.json')
+    })
+
     it('evicts oldest when over limit', () => {
       const store = useWorkflowDraftStoreV2()
 
@@ -183,10 +258,38 @@ describe('workflowDraftStoreV2', () => {
         data: '{"data":"test"}'
       })
     })
+
+    it('does not move legacy shadow data when V2 move fails', () => {
+      const store = useWorkflowDraftStoreV2()
+      const legacyStore = useWorkflowDraftStore()
+
+      store.saveDraft('workflows/old.json', '{"data":"test"}', {
+        name: 'old',
+        isTemporary: true
+      })
+
+      const setItemSpy = failV2IndexWrites()
+      const moved = store.moveDraft(
+        'workflows/old.json',
+        'workflows/new.json',
+        'new'
+      )
+      setItemSpy.mockRestore()
+
+      expect(moved).toBe(false)
+      expect(store.getDraft('workflows/old.json')).toMatchObject({
+        data: '{"data":"test"}'
+      })
+      expect(store.getDraft('workflows/new.json')).toBeNull()
+      expect(legacyStore.getDraft('workflows/old.json')).toMatchObject({
+        data: '{"data":"test"}'
+      })
+      expect(legacyStore.getDraft('workflows/new.json')).toBeUndefined()
+    })
   })
 
   describe('getDraft', () => {
-    it('prefers legacy draft content when it is newer than V2', () => {
+    it('uses legacy fallback content only when it is strictly newer than V2', () => {
       const store = useWorkflowDraftStoreV2()
       const path = 'workflows/test.json'
 
@@ -204,6 +307,39 @@ describe('workflowDraftStoreV2', () => {
       expect(store.getDraft(path)?.data).toBe('{"source":"legacy"}')
     })
 
+    it('does not read or write legacy drafts outside personal workspace', () => {
+      sessionStorage.setItem(
+        'Comfy.Workspace.Current',
+        JSON.stringify({ id: 'workspace-1', type: 'org' })
+      )
+      const store = useWorkflowDraftStoreV2()
+      const legacyStore = useWorkflowDraftStore()
+      const path = 'workflows/test.json'
+
+      store.saveDraft(path, '{"source":"v2"}', {
+        name: 'test',
+        isTemporary: true
+      })
+      legacyStore.saveDraft(path, {
+        data: '{"source":"legacy"}',
+        updatedAt: Date.now() + 1000,
+        name: 'test',
+        isTemporary: true
+      })
+      legacyStore.saveDraft('workflows/legacy-only.json', {
+        data: '{"source":"legacy-only"}',
+        updatedAt: Date.now() + 1000,
+        name: 'legacy-only',
+        isTemporary: true
+      })
+
+      expect(legacyStore.getDraft(path)).toMatchObject({
+        data: '{"source":"legacy"}'
+      })
+      expect(store.getDraft(path)?.data).toBe('{"source":"v2"}')
+      expect(store.getDraft('workflows/legacy-only.json')).toBeNull()
+    })
+
     it('returns null when legacy fallback read fails', () => {
       const store = useWorkflowDraftStoreV2()
       const legacyStore = useWorkflowDraftStore()

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

@@ -49,6 +49,8 @@ interface LoadPersistedWorkflowOptions {
   fallbackToLatestDraft?: boolean
 }
 
+type DraftMoveStatus = 'moved' | 'missing' | 'failed'
+
 export const useWorkflowDraftStoreV2 = defineStore('workflowDraftV2', () => {
   // In-memory cache of the index per workspace (synced with localStorage)
   // Key is workspaceId, value is the cached index
@@ -96,18 +98,23 @@ export const useWorkflowDraftStoreV2 = defineStore('workflowDraftV2', () => {
    */
   function persistIndex(index: DraftIndexV2): boolean {
     const workspaceId = currentWorkspaceId()
+    const written = writeIndex(workspaceId, index)
+    if (!written) return false
+
     indexCacheByWorkspace.value[workspaceId] = index
-    return writeIndex(workspaceId, index)
+    return true
   }
 
   /**
    * Saves a draft to V2 and shadow-writes the legacy store for rollback.
    */
   function saveDraft(path: string, data: string, meta: DraftMeta): boolean {
-    const savedV2 = saveDraftV2(path, data, meta)
-    const savedLegacy = saveLegacyDraft(path, data, meta)
+    const now = Date.now()
+    const savedV2 = saveDraftV2(path, data, meta, now)
+    if (!savedV2) return false
 
-    return savedV2 || savedLegacy
+    saveLegacyDraft(path, data, meta, now)
+    return true
   }
 
   /**
@@ -121,13 +128,19 @@ export const useWorkflowDraftStoreV2 = defineStore('workflowDraftV2', () => {
   /**
    * Moves a draft from one path to another in V2 and legacy rollback storage.
    */
-  function moveDraft(oldPath: string, newPath: string, name: string): void {
-    moveDraftV2(oldPath, newPath, name)
-    moveLegacyDraft(oldPath, newPath, name)
+  function moveDraft(oldPath: string, newPath: string, name: string): boolean {
+    const status = moveDraftV2(oldPath, newPath, name)
+    if (status !== 'failed') {
+      moveLegacyDraft(oldPath, newPath, name)
+    }
+    return status !== 'failed'
   }
 
   /**
-   * Gets draft data by path, preferring whichever store has newer content.
+   * Gets draft data by path, using legacy only as a rollback fallback.
+   *
+   * V2 wins equal timestamps because normal shadow-writes use the same
+   * timestamp in both stores.
    */
   function getDraft(path: string): WorkflowDraftSnapshot | null {
     const v2Draft = getDraftV2(path)
@@ -204,14 +217,18 @@ export const useWorkflowDraftStoreV2 = defineStore('workflowDraftV2', () => {
 
   /**
    * Saves a draft (data + metadata) to V2.
-   * Primes index cache, writes payload, then persists updated index.
+   * Loads the index before writing payload, then persists the updated index.
    */
-  function saveDraftV2(path: string, data: string, meta: DraftMeta): boolean {
+  function saveDraftV2(
+    path: string,
+    data: string,
+    meta: DraftMeta,
+    now: number
+  ): boolean {
     if (!isStorageAvailable()) return false
 
     const workspaceId = currentWorkspaceId()
     const draftKey = hashPath(path)
-    const now = Date.now()
 
     // Prime the index cache before writing payload.
     // loadIndex() runs orphan cleanup on cache miss, which would
@@ -226,7 +243,7 @@ export const useWorkflowDraftStoreV2 = defineStore('workflowDraftV2', () => {
 
     if (!payloadWritten) {
       // Quota exceeded - try eviction loop
-      return handleQuotaExceeded(path, data, meta)
+      return handleQuotaExceeded(path, data, meta, now)
     }
     const { index: newIndex, evicted } = upsertEntry(
       index,
@@ -254,7 +271,8 @@ export const useWorkflowDraftStoreV2 = defineStore('workflowDraftV2', () => {
   function handleQuotaExceeded(
     path: string,
     data: string,
-    meta: DraftMeta
+    meta: DraftMeta,
+    now: number
   ): boolean {
     const workspaceId = currentWorkspaceId()
     const index = loadIndex()
@@ -281,7 +299,7 @@ export const useWorkflowDraftStoreV2 = defineStore('workflowDraftV2', () => {
       // Try writing again
       const success = writePayload(workspaceId, draftKey, {
         data,
-        updatedAt: Date.now()
+        updatedAt: now
       })
 
       if (success) {
@@ -289,7 +307,7 @@ export const useWorkflowDraftStoreV2 = defineStore('workflowDraftV2', () => {
         const { index: finalIndex } = upsertEntry(
           currentIndex,
           path,
-          { ...meta, updatedAt: Date.now() },
+          { ...meta, updatedAt: now },
           MAX_DRAFTS
         )
         if (!persistIndex(finalIndex)) {
@@ -316,27 +334,45 @@ export const useWorkflowDraftStoreV2 = defineStore('workflowDraftV2', () => {
     }
   }
 
-  function moveDraftV2(oldPath: string, newPath: string, name: string): void {
+  function moveDraftV2(
+    oldPath: string,
+    newPath: string,
+    name: string
+  ): DraftMoveStatus {
     const workspaceId = currentWorkspaceId()
     const index = loadIndex()
+    const oldEntry = getEntryByPath(index, oldPath)
+    if (!oldEntry) return 'missing'
+
+    const oldKey = hashPath(oldPath)
+    const newKey = hashPath(newPath)
+    const newEntry = getEntryByPath(index, newPath)
+    if (oldKey !== newKey && newEntry) return 'failed'
+
     const result = moveEntry(index, oldPath, newPath, name)
+    if (!result) return 'failed'
 
-    if (result) {
-      const oldPayload = readPayload(workspaceId, result.oldKey)
-      if (oldPayload) {
-        const written = writePayload(workspaceId, result.newKey, {
-          data: oldPayload.data,
-          updatedAt: Date.now()
-        })
-        if (!written) return
-
-        if (!persistIndex(result.index)) {
-          deletePayload(workspaceId, result.newKey)
-          return
-        }
-        deletePayload(workspaceId, result.oldKey)
-      }
+    const oldPayload = readPayload(workspaceId, result.oldKey)
+    if (!oldPayload) {
+      removeDraftV2(oldPath)
+      return 'failed'
     }
+
+    const written = writePayload(workspaceId, result.newKey, {
+      data: oldPayload.data,
+      updatedAt: Date.now()
+    })
+    if (!written) return 'failed'
+
+    if (!persistIndex(result.index)) {
+      deletePayload(workspaceId, result.newKey)
+      return 'failed'
+    }
+
+    if (result.oldKey !== result.newKey) {
+      deletePayload(workspaceId, result.oldKey)
+    }
+    return 'moved'
   }
 
   function getDraftV2(path: string): WorkflowDraftSnapshot | null {
@@ -397,11 +433,14 @@ export const useWorkflowDraftStoreV2 = defineStore('workflowDraftV2', () => {
   function saveLegacyDraft(
     path: string,
     data: string,
-    meta: DraftMeta
+    meta: DraftMeta,
+    updatedAt: number
   ): boolean {
+    if (!canUseLegacyDraftStore()) return false
+
     const snapshot: WorkflowDraftSnapshot = {
       data,
-      updatedAt: Date.now(),
+      updatedAt,
       name: meta.name,
       isTemporary: meta.isTemporary
     }
@@ -417,6 +456,8 @@ export const useWorkflowDraftStoreV2 = defineStore('workflowDraftV2', () => {
   }
 
   function getLegacyDraft(path: string): WorkflowDraftSnapshot | null {
+    if (!canUseLegacyDraftStore()) return null
+
     try {
       return useWorkflowDraftStore().getDraft(path) ?? null
     } catch {
@@ -426,6 +467,8 @@ export const useWorkflowDraftStoreV2 = defineStore('workflowDraftV2', () => {
   }
 
   function removeLegacyDraft(path: string): void {
+    if (!canUseLegacyDraftStore()) return
+
     try {
       useWorkflowDraftStore().removeDraft(path)
     } catch {
@@ -435,6 +478,8 @@ export const useWorkflowDraftStoreV2 = defineStore('workflowDraftV2', () => {
   }
 
   function moveLegacyDraft(oldPath: string, newPath: string, name: string) {
+    if (!canUseLegacyDraftStore()) return
+
     try {
       useWorkflowDraftStore().moveDraft(oldPath, newPath, name)
     } catch {
@@ -444,6 +489,8 @@ export const useWorkflowDraftStoreV2 = defineStore('workflowDraftV2', () => {
   }
 
   function markLegacyDraftUsed(path: string) {
+    if (!canUseLegacyDraftStore()) return
+
     try {
       useWorkflowDraftStore().markDraftUsed(path)
     } catch {
@@ -452,6 +499,10 @@ export const useWorkflowDraftStoreV2 = defineStore('workflowDraftV2', () => {
     }
   }
 
+  function canUseLegacyDraftStore(): boolean {
+    return currentWorkspaceId() === 'personal'
+  }
+
   async function loadLegacyPersistedWorkflow(
     options: LoadPersistedWorkflowOptions
   ): Promise<boolean> {