diff --git a/src/platform/workflow/management/stores/workflowStore.ts b/src/platform/workflow/management/stores/workflowStore.ts index f3c8c96515..5d439f771e 100644 --- a/src/platform/workflow/management/stores/workflowStore.ts +++ b/src/platform/workflow/management/stores/workflowStore.ts @@ -8,6 +8,7 @@ import type { LGraphNode, Subgraph } from '@/lib/litegraph/src/litegraph' +import { useWorkflowDraftStore } from '@/platform/workflow/persistence/stores/workflowDraftStore' import type { ComfyWorkflowJSON } from '@/platform/workflow/validation/schemas/workflowSchema' import type { NodeId } from '@/platform/workflow/validation/schemas/workflowSchema' import { useWorkflowThumbnail } from '@/renderer/core/thumbnail/useWorkflowThumbnail' @@ -83,6 +84,28 @@ export class ComfyWorkflow extends UserFile { override async load({ force = false }: { force?: boolean } = {}): Promise { + 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 LoadedComfyWorkflow @@ -90,13 +113,16 @@ 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 LoadedComfyWorkflow } @@ -106,12 +132,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 } @@ -121,8 +149,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 { @@ -448,6 +479,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 @@ -457,6 +489,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) @@ -474,6 +508,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) } diff --git a/src/platform/workflow/persistence/base/draftCache.ts b/src/platform/workflow/persistence/base/draftCache.ts new file mode 100644 index 0000000000..2ac15e4b96 --- /dev/null +++ b/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 + order: string[] +} + +export const MAX_DRAFTS = 32 + +export const createDraftCacheState = ( + drafts: Record = {}, + 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 diff --git a/src/platform/workflow/persistence/composables/useWorkflowPersistence.ts b/src/platform/workflow/persistence/composables/useWorkflowPersistence.ts index f6f5f18280..eac9424306 100644 --- a/src/platform/workflow/persistence/composables/useWorkflowPersistence.ts +++ b/src/platform/workflow/persistence/composables/useWorkflowPersistence.ts @@ -3,7 +3,11 @@ import { computed, watch } from 'vue' 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 { api } from '@/scripts/api' import { app as comfyApp } from '@/scripts/app' import { getStorageValue, setStorageValue } from '@/scripts/utils' @@ -12,6 +16,7 @@ import { useCommandStore } from '@/stores/commandStore' export function useWorkflowPersistence() { const workflowStore = useWorkflowStore() const settingStore = useSettingStore() + const workflowDraftStore = useWorkflowDraftStore() const workflowPersistenceEnabled = computed(() => settingStore.get('Comfy.Workflow.Persist') @@ -19,38 +24,40 @@ export function useWorkflowPersistence() { const persistCurrentWorkflow = () => { if (!workflowPersistenceEnabled.value) return - const workflow = JSON.stringify(comfyApp.graph.serialize()) - localStorage.setItem('workflow', workflow) + const activeWorkflow = workflowStore.activeWorkflow + if (!activeWorkflow) return + + const graphData = comfyApp.graph.serialize() + const workflowJson = JSON.stringify(graphData) + localStorage.setItem('workflow', workflowJson) if (api.clientId) { - sessionStorage.setItem(`workflow:${api.clientId}`, workflow) + sessionStorage.setItem(`workflow:${api.clientId}`, workflowJson) + } + + if (!activeWorkflow.isTemporary && !activeWorkflow.isModified) { + workflowDraftStore.removeDraft(activeWorkflow.path) + return } - } - 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 + workflowDraftStore.saveDraft(activeWorkflow.path, { + data: workflowJson, + updatedAt: Date.now(), + name: activeWorkflow.key, + isTemporary: activeWorkflow.isTemporary + }) } const loadPreviousWorkflowFromStorage = async () => { const workflowName = getStorageValue('Comfy.PreviousWorkflow') - const clientId = api.initialClientId ?? api.clientId - - // 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) + const preferredPath = workflowName + ? `${ComfyWorkflow.basePath}${workflowName}` + : null + + return await workflowDraftStore.loadPersistedWorkflow({ + workflowName, + preferredPath, + fallbackToLatestDraft: !workflowName + }) } const loadDefaultWorkflow = async () => { @@ -104,11 +111,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 } } @@ -117,10 +125,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) { @@ -132,12 +140,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 { diff --git a/src/platform/workflow/persistence/stores/workflowDraftStore.ts b/src/platform/workflow/persistence/stores/workflowDraftStore.ts new file mode 100644 index 0000000000..a53d669a37 --- /dev/null +++ b/src/platform/workflow/persistence/stores/workflowDraftStore.ts @@ -0,0 +1,136 @@ +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>( + DRAFTS_STORAGE_KEY, + {} + ) + const storedOrder = useStorage(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 => { + 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()) + } + } +}) diff --git a/tests-ui/tests/platform/workflow/persistence/draftCache.test.ts b/tests-ui/tests/platform/workflow/persistence/draftCache.test.ts new file mode 100644 index 0000000000..ca10f3ab0d --- /dev/null +++ b/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() + }) +}) diff --git a/tests-ui/tests/platform/workflow/persistence/useWorkflowPersistence.test.ts b/tests-ui/tests/platform/workflow/persistence/useWorkflowPersistence.test.ts new file mode 100644 index 0000000000..1922897bee --- /dev/null +++ b/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 + } + 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 + + 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 + + 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') + }) +}) diff --git a/tests-ui/tests/store/workflowStore.test.ts b/tests-ui/tests/store/workflowStore.test.ts index 648d00fa88..75afb12b86 100644 --- a/tests-ui/tests/store/workflowStore.test.ts +++ b/tests-ui/tests/store/workflowStore.test.ts @@ -9,6 +9,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' @@ -57,6 +58,9 @@ describe('useWorkflowStore', () => { store = useWorkflowStore() bookmarkStore = useWorkflowBookmarkStore() vi.clearAllMocks() + localStorage.clear() + sessionStorage.clear() + useWorkflowDraftStore().reset() // Add default mock implementations vi.mocked(api.getUserData).mockResolvedValue({ @@ -154,6 +158,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'])