fix: restore correct workflow on page reload (#9318)

## Summary

- Fixes incorrect workflow loading after page refresh when the active
workflow is saved and unmodified
- Adds saved-workflow fallback in `loadPreviousWorkflowFromStorage()`
before falling back to latest draft
- Calls `openWorkflow()` in `restoreWorkflowTabsState()` to activate the
correct tab after restoration

- Fixes #9317

## Test plan

- [x] Unit tests for saved-workflow fallback (draft missing, saved
workflow exists)
- [x] Unit tests for correct tab activation after restoration
- [x] Unit tests for existing behavior preservation (draft preferred
over saved, no-session-path fallback)
- [x] `pnpm typecheck` passes
- [x] `pnpm lint` passes
- [x] All 117 persistence tests pass

🤖 Generated with [Claude Code](https://claude.com/claude-code)

┆Issue is synchronized with this [Notion
page](https://www.notion.so/PR-9318-fix-restore-correct-workflow-on-page-reload-3166d73d365081ba9139f7c23c917aa4)
by [Unito](https://www.unito.io)

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Alexander Brown <drjkl@comfy.org>
Co-authored-by: GitHub Action <action@github.com>

.gitignore

@@ -26,6 +26,7 @@ dist-ssr
 .claude/*.local.json
 .claude/*.local.md
 .claude/*.local.txt
+.claude/worktrees
 CLAUDE.local.md
 
 # Editor directories and files

src/components/graph/GraphCanvas.vue

@@ -538,7 +538,7 @@ onMounted(async () => {
 
   // Restore saved workflow and workflow tabs state
   await workflowPersistence.initializeWorkflow()
-  workflowPersistence.restoreWorkflowTabsState()
+  await workflowPersistence.restoreWorkflowTabsState()
 
   const sharedWorkflowLoadStatus =
     await workflowPersistence.loadSharedWorkflowFromUrlIfPresent()

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

@@ -0,0 +1,360 @@
+import { createTestingPinia } from '@pinia/testing'
+import { setActivePinia } from 'pinia'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+import type * as I18n from 'vue-i18n'
+
+import { useWorkflowStore } from '@/platform/workflow/management/stores/workflowStore'
+import { useWorkflowDraftStoreV2 } from '../stores/workflowDraftStoreV2'
+import { useWorkflowPersistenceV2 } from './useWorkflowPersistenceV2'
+
+const settingMocks = vi.hoisted(() => ({
+  persistRef: null as { value: boolean } | null
+}))
+
+vi.mock('@/platform/settings/settingStore', async () => {
+  const { ref } = await import('vue')
+  settingMocks.persistRef = ref(true)
+  return {
+    useSettingStore: vi.fn(() => ({
+      get: vi.fn((key: string) => {
+        if (key === 'Comfy.Workflow.Persist')
+          return settingMocks.persistRef!.value
+        return undefined
+      }),
+      set: vi.fn()
+    }))
+  }
+})
+
+const mockToastAdd = vi.fn()
+vi.mock('primevue', () => ({
+  useToast: () => ({
+    add: mockToastAdd
+  })
+}))
+
+vi.mock('primevue/usetoast', () => ({
+  useToast: () => ({
+    add: mockToastAdd
+  })
+}))
+
+vi.mock(
+  '@/platform/workflow/sharing/composables/useSharedWorkflowUrlLoader',
+  () => ({
+    useSharedWorkflowUrlLoader: () => ({
+      loadSharedWorkflowFromUrl: vi.fn().mockResolvedValue('not-present')
+    })
+  })
+)
+
+vi.mock('vue-i18n', async (importOriginal) => {
+  const actual = await importOriginal<typeof I18n>()
+  return {
+    ...actual,
+    useI18n: () => ({
+      t: (key: string) => key
+    })
+  }
+})
+
+const openWorkflowMock = vi.fn()
+const loadBlankWorkflowMock = vi.fn()
+vi.mock('@/platform/workflow/core/services/workflowService', () => ({
+  useWorkflowService: () => ({
+    openWorkflow: openWorkflowMock,
+    loadBlankWorkflow: loadBlankWorkflowMock
+  })
+}))
+
+vi.mock(
+  '@/platform/workflow/templates/composables/useTemplateUrlLoader',
+  () => ({
+    useTemplateUrlLoader: () => ({
+      loadTemplateFromUrl: vi.fn()
+    })
+  })
+)
+
+vi.mock('@/stores/commandStore', () => ({
+  useCommandStore: () => ({
+    execute: vi.fn()
+  })
+}))
+
+vi.mock('vue-router', () => ({
+  useRoute: () => ({
+    query: {}
+  }),
+  useRouter: () => ({
+    replace: vi.fn()
+  })
+}))
+
+vi.mock('@/composables/auth/useCurrentUser', () => ({
+  useCurrentUser: () => ({
+    onUserLogout: vi.fn()
+  })
+}))
+
+vi.mock('@/platform/navigation/preservedQueryManager', () => ({
+  hydratePreservedQuery: vi.fn(),
+  mergePreservedQueryIntoQuery: vi.fn(() => null)
+}))
+
+vi.mock('@/platform/navigation/preservedQueryNamespaces', () => ({
+  PRESERVED_QUERY_NAMESPACES: { TEMPLATE: 'template' }
+}))
+
+vi.mock('@/platform/distribution/types', () => ({
+  isCloud: false
+}))
+
+vi.mock('../migration/migrateV1toV2', () => ({
+  migrateV1toV2: vi.fn()
+}))
+
+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()
+  }
+  return { state, serializeMock, loadGraphDataMock, apiMock }
+})
+
+vi.mock('@/scripts/app', () => ({
+  app: {
+    graph: {
+      serialize: () => mocks.serializeMock()
+    },
+    rootGraph: {
+      serialize: () => mocks.serializeMock()
+    },
+    loadGraphData: (...args: unknown[]) => mocks.loadGraphDataMock(...args),
+    canvas: {}
+  }
+}))
+
+vi.mock('@/scripts/api', () => ({
+  api: mocks.apiMock
+}))
+
+describe('useWorkflowPersistenceV2', () => {
+  beforeEach(() => {
+    vi.useFakeTimers()
+    vi.setSystemTime(new Date('2025-01-01T00:00:00Z'))
+    setActivePinia(createTestingPinia({ stubActions: false }))
+    localStorage.clear()
+    sessionStorage.clear()
+    vi.clearAllMocks()
+    settingMocks.persistRef!.value = true
+    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(() => {})
+    openWorkflowMock.mockReset()
+    loadBlankWorkflowMock.mockReset()
+  })
+
+  afterEach(() => {
+    vi.useRealTimers()
+  })
+
+  function writeTabState(paths: string[], activeIndex: number) {
+    const pointer = {
+      workspaceId: 'personal',
+      paths,
+      activeIndex
+    }
+    sessionStorage.setItem(
+      `Comfy.Workflow.OpenPaths:test-client`,
+      JSON.stringify(pointer)
+    )
+  }
+
+  function writeActivePath(path: string) {
+    const pointer = {
+      workspaceId: 'personal',
+      path
+    }
+    sessionStorage.setItem(
+      `Comfy.Workflow.ActivePath:test-client`,
+      JSON.stringify(pointer)
+    )
+  }
+
+  describe('loadPreviousWorkflowFromStorage', () => {
+    it('loads saved workflow when draft is missing for session path', async () => {
+      const workflowStore = useWorkflowStore()
+      const savedWorkflow = workflowStore.createTemporary('SavedWorkflow.json')
+
+      // Set session path to the saved workflow but do NOT create a draft
+      writeActivePath(savedWorkflow.path)
+
+      const { initializeWorkflow } = useWorkflowPersistenceV2()
+      await initializeWorkflow()
+
+      // Should call workflowService.openWorkflow with the saved workflow
+      expect(openWorkflowMock).toHaveBeenCalledWith(savedWorkflow)
+      // Should NOT fall through to loadGraphData (fallbackToLatestDraft)
+      expect(mocks.loadGraphDataMock).not.toHaveBeenCalled()
+    })
+
+    it('prefers draft over saved workflow when draft exists', async () => {
+      const workflowStore = useWorkflowStore()
+      const draftStore = useWorkflowDraftStoreV2()
+
+      const workflow = workflowStore.createTemporary('DraftWorkflow.json')
+      const draftData = JSON.stringify({ nodes: [], title: 'draft' })
+      draftStore.saveDraft(workflow.path, draftData, {
+        name: 'DraftWorkflow.json',
+        isTemporary: true
+      })
+      writeActivePath(workflow.path)
+
+      mocks.loadGraphDataMock.mockResolvedValue(undefined)
+
+      const { initializeWorkflow } = useWorkflowPersistenceV2()
+      await initializeWorkflow()
+
+      // Should load draft via loadGraphData, not via workflowService.openWorkflow
+      expect(mocks.loadGraphDataMock).toHaveBeenCalled()
+      expect(openWorkflowMock).not.toHaveBeenCalled()
+    })
+
+    it('falls back to latest draft only when no session path exists', async () => {
+      const draftStore = useWorkflowDraftStoreV2()
+
+      // No session path set, but a draft exists
+      const draftData = JSON.stringify({ nodes: [], title: 'latest' })
+      draftStore.saveDraft('workflows/Other.json', draftData, {
+        name: 'Other.json',
+        isTemporary: true
+      })
+
+      mocks.loadGraphDataMock.mockResolvedValue(undefined)
+
+      const { initializeWorkflow } = useWorkflowPersistenceV2()
+      await initializeWorkflow()
+
+      // Should load via fallbackToLatestDraft
+      expect(mocks.loadGraphDataMock).toHaveBeenCalled()
+      expect(openWorkflowMock).not.toHaveBeenCalled()
+    })
+  })
+
+  describe('restoreWorkflowTabsState', () => {
+    it('activates the correct workflow at storedActiveIndex', async () => {
+      const workflowStore = useWorkflowStore()
+      const draftStore = useWorkflowDraftStoreV2()
+
+      // Create two temporary workflows with drafts
+      const workflowA = workflowStore.createTemporary('WorkflowA.json')
+      const workflowB = workflowStore.createTemporary('WorkflowB.json')
+
+      draftStore.saveDraft(workflowA.path, JSON.stringify({ title: 'A' }), {
+        name: 'WorkflowA.json',
+        isTemporary: true
+      })
+      draftStore.saveDraft(workflowB.path, JSON.stringify({ title: 'B' }), {
+        name: 'WorkflowB.json',
+        isTemporary: true
+      })
+
+      // storedActiveIndex = 1 → WorkflowB should be activated
+      writeTabState([workflowA.path, workflowB.path], 1)
+
+      const { restoreWorkflowTabsState } = useWorkflowPersistenceV2()
+      await restoreWorkflowTabsState()
+
+      expect(openWorkflowMock).toHaveBeenCalledWith(workflowB)
+    })
+
+    it('activates first tab when storedActiveIndex is 0', async () => {
+      const workflowStore = useWorkflowStore()
+      const draftStore = useWorkflowDraftStoreV2()
+
+      const workflowA = workflowStore.createTemporary('WorkflowA.json')
+      const workflowB = workflowStore.createTemporary('WorkflowB.json')
+
+      draftStore.saveDraft(workflowA.path, JSON.stringify({ title: 'A' }), {
+        name: 'WorkflowA.json',
+        isTemporary: true
+      })
+      draftStore.saveDraft(workflowB.path, JSON.stringify({ title: 'B' }), {
+        name: 'WorkflowB.json',
+        isTemporary: true
+      })
+
+      writeTabState([workflowA.path, workflowB.path], 0)
+
+      const { restoreWorkflowTabsState } = useWorkflowPersistenceV2()
+      await restoreWorkflowTabsState()
+
+      expect(openWorkflowMock).toHaveBeenCalledWith(workflowA)
+    })
+
+    it('does not call openWorkflow when no restorable state', async () => {
+      // No tab state written to sessionStorage
+      const { restoreWorkflowTabsState } = useWorkflowPersistenceV2()
+      await restoreWorkflowTabsState()
+
+      expect(openWorkflowMock).not.toHaveBeenCalled()
+    })
+
+    it('restores temporary workflows and adds them to tabs', async () => {
+      const workflowStore = useWorkflowStore()
+      const draftStore = useWorkflowDraftStoreV2()
+
+      // Save a draft for a workflow that doesn't exist in the store yet
+      const path = 'workflows/Unsaved.json'
+      draftStore.saveDraft(path, JSON.stringify({ title: 'Unsaved' }), {
+        name: 'Unsaved.json',
+        isTemporary: true
+      })
+
+      writeTabState([path], 0)
+
+      const { restoreWorkflowTabsState } = useWorkflowPersistenceV2()
+      await restoreWorkflowTabsState()
+
+      const restored = workflowStore.getWorkflowByPath(path)
+      expect(restored).toBeTruthy()
+      expect(restored?.isTemporary).toBe(true)
+      expect(workflowStore.openWorkflows.map((w) => w?.path)).toContain(path)
+    })
+
+    it('skips activation when persistence is disabled', async () => {
+      settingMocks.persistRef!.value = false
+
+      const { restoreWorkflowTabsState } = useWorkflowPersistenceV2()
+      await restoreWorkflowTabsState()
+
+      expect(openWorkflowMock).not.toHaveBeenCalled()
+    })
+  })
+})

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

@@ -132,19 +132,28 @@ export function useWorkflowPersistenceV2() {
   const debouncedPersist = debounce(persistCurrentWorkflow, PERSIST_DEBOUNCE_MS)
 
   const loadPreviousWorkflowFromStorage = async () => {
-    // 1. Try session pointer (for tab restoration)
     const sessionPath = tabState.getActivePath()
+
+    // 1. Try draft for session path
     if (
       sessionPath &&
       (await draftStore.loadPersistedWorkflow({
         workflowName: null,
         preferredPath: sessionPath
       }))
-    ) {
+    )
       return true
+
+    // 2. Try saved workflow by path (draft may not exist for saved+unmodified workflows)
+    if (sessionPath) {
+      const saved = workflowStore.getWorkflowByPath(sessionPath)
+      if (saved) {
+        await useWorkflowService().openWorkflow(saved)
+        return true
+      }
     }
 
-    // 2. Fall back to most recent draft
+    // 3. Fall back to most recent draft
     return await draftStore.loadPersistedWorkflow({
       workflowName: null,
       fallbackToLatestDraft: true
@@ -242,7 +251,7 @@ export function useWorkflowPersistenceV2() {
     }
   })
 
-  const restoreWorkflowTabsState = () => {
+  const restoreWorkflowTabsState = async () => {
     if (!workflowPersistenceEnabled.value) {
       tabStateRestored = true
       return
@@ -254,10 +263,11 @@ export function useWorkflowPersistenceV2() {
     const storedWorkflows = storedTabState?.paths ?? []
     const storedActiveIndex = storedTabState?.activeIndex ?? -1
 
-    tabStateRestored = true
-
     const isRestorable = storedWorkflows.length > 0 && storedActiveIndex >= 0
-    if (!isRestorable) return
+    if (!isRestorable) {
+      tabStateRestored = true
+      return
+    }
 
     storedWorkflows.forEach((path: string) => {
       if (workflowStore.getWorkflowByPath(path)) return
@@ -280,6 +290,17 @@ export function useWorkflowPersistenceV2() {
       left: storedWorkflows.slice(0, storedActiveIndex),
       right: storedWorkflows.slice(storedActiveIndex)
     })
+
+    tabStateRestored = true
+
+    // Activate the correct workflow at storedActiveIndex
+    const activePath = storedWorkflows[storedActiveIndex]
+    const workflow = activePath
+      ? workflowStore.getWorkflowByPath(activePath)
+      : null
+    if (workflow) {
+      await useWorkflowService().openWorkflow(workflow)
+    }
   }
 
   return {