From 2f6461ddf832191d49ee060fbda3cc7831012463 Mon Sep 17 00:00:00 2001 From: Connor Byrne Date: Wed, 13 May 2026 13:49:25 -0700 Subject: [PATCH] fix: defer removeDraft until after switch succeeds Moves removeDraft() call to after the switch/close logic completes, preserving the draft if closeWorkflow returns false due to a failed switch. Adds test coverage for draft preservation on switch failure. Addresses review feedback: https://github.com/Comfy-Org/ComfyUI_frontend/pull/11389#pullrequestreview-2960339131 Co-Authored-By: Claude Opus 4.5 --- .../core/services/workflowService.test.ts | 16 ++++++++++++++++ .../workflow/core/services/workflowService.ts | 6 ++++-- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/src/platform/workflow/core/services/workflowService.test.ts b/src/platform/workflow/core/services/workflowService.test.ts index 7583d90086..bd8ded2e7a 100644 --- a/src/platform/workflow/core/services/workflowService.test.ts +++ b/src/platform/workflow/core/services/workflowService.test.ts @@ -11,6 +11,7 @@ import { useSettingStore } from '@/platform/settings/settingStore' import { useToastStore } from '@/platform/updates/common/toastStore' import type { ComfyWorkflow } from '@/platform/workflow/management/stores/workflowStore' import { useWorkflowStore } from '@/platform/workflow/management/stores/workflowStore' +import { useWorkflowDraftStore } from '@/platform/workflow/persistence/stores/workflowDraftStore' import { useWorkflowService } from '@/platform/workflow/core/services/workflowService' import { useMissingNodesErrorStore } from '@/platform/nodeReplacement/missingNodesErrorStore' import { useExecutionErrorStore } from '@/stores/executionErrorStore' @@ -1292,6 +1293,21 @@ describe('useWorkflowService', () => { expect(workflowStore.isOpen(active)).toBe(true) }) + it('does not remove draft when switch fails', async () => { + const draftStore = useWorkflowDraftStore() + const active = createAndRegister('workflows/active.json', 0) + createAndRegister('workflows/other.json', 1) + workflowStore.activeWorkflow = active as LoadedComfyWorkflow + + vi.mocked(app.loadGraphData).mockRejectedValue( + new Error('configure failed') + ) + + await service.closeWorkflow(active, { warnIfUnsaved: false }) + + expect(draftStore.removeDraft).not.toHaveBeenCalled() + }) + it('falls back to default workflow when replacement throws', async () => { const active = createAndRegister('workflows/active.json', 0) createAndRegister('workflows/other.json', 1) diff --git a/src/platform/workflow/core/services/workflowService.ts b/src/platform/workflow/core/services/workflowService.ts index d4d9ed35d2..071d535e22 100644 --- a/src/platform/workflow/core/services/workflowService.ts +++ b/src/platform/workflow/core/services/workflowService.ts @@ -324,8 +324,6 @@ export const useWorkflowService = () => { } } - workflowDraftStore.removeDraft(workflow.path) - // If this is the last workflow, create a new default temporary workflow. // Route through trySwitch so a rejection from loadGraphData // (validation / extension hooks / node-replacement loading) keeps the tab @@ -340,6 +338,10 @@ export const useWorkflowService = () => { if (!didSwitch) return false } + // Remove draft only after switch succeeds — keeps draft intact if close is + // aborted due to switch failure, so the tab retains its persisted state. + workflowDraftStore.removeDraft(workflow.path) + await workflowStore.closeWorkflow(workflow) return true }