From 9a3bbcc1b7284e4019d0b4cf4acbd233b3e83eea Mon Sep 17 00:00:00 2001 From: jaeone94 <89377375+jaeone94@users.noreply.github.com> Date: Tue, 12 May 2026 15:25:50 +0900 Subject: [PATCH] fix: harden workflow draft v2 shadow writes --- .../tests/workflowPersistence.spec.ts | 12 +- .../useWorkflowPersistenceV2.test.ts | 26 +++- .../composables/useWorkflowPersistenceV2.ts | 5 +- .../stores/workflowDraftStoreV2.test.ts | 138 +++++++++++++++++- .../stores/workflowDraftStoreV2.ts | 117 ++++++++++----- 5 files changed, 260 insertions(+), 38 deletions(-) diff --git a/browser_tests/tests/workflowPersistence.spec.ts b/browser_tests/tests/workflowPersistence.spec.ts index da1858a05b..b112983ea3 100644 --- a/browser_tests/tests/workflowPersistence.spec.ts +++ b/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) } ) diff --git a/src/platform/workflow/persistence/composables/useWorkflowPersistenceV2.test.ts b/src/platform/workflow/persistence/composables/useWorkflowPersistenceV2.test.ts index 1716650dd6..a7217c6987 100644 --- a/src/platform/workflow/persistence/composables/useWorkflowPersistenceV2.test.ts +++ b/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 diff --git a/src/platform/workflow/persistence/composables/useWorkflowPersistenceV2.ts b/src/platform/workflow/persistence/composables/useWorkflowPersistenceV2.ts index c42bc8ee21..21bda3030b 100644 --- a/src/platform/workflow/persistence/composables/useWorkflowPersistenceV2.ts +++ b/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 }) => { diff --git a/src/platform/workflow/persistence/stores/workflowDraftStoreV2.test.ts b/src/platform/workflow/persistence/stores/workflowDraftStoreV2.test.ts index dc414cf93a..f0e38ae73d 100644 --- a/src/platform/workflow/persistence/stores/workflowDraftStoreV2.test.ts +++ b/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() diff --git a/src/platform/workflow/persistence/stores/workflowDraftStoreV2.ts b/src/platform/workflow/persistence/stores/workflowDraftStoreV2.ts index 1dc9c24e70..bd93a2493b 100644 --- a/src/platform/workflow/persistence/stores/workflowDraftStoreV2.ts +++ b/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 {