Restore all open workflows on load (#2238)

browser_tests/fixtures/components/Topbar.ts

@@ -9,6 +9,12 @@ export class Topbar {
       .allInnerTexts()
   }
 
+  async getActiveTabName(): Promise<string> {
+    return this.page
+      .locator('.workflow-tabs .p-togglebutton-checked')
+      .innerText()
+  }
+
   async openSubmenuMobile() {
     await this.page.locator('.p-menubar-mobile .p-menubar-button').click()
   }

browser_tests/interaction.spec.ts

@@ -615,6 +615,67 @@ test.describe('Load workflow', () => {
       'single_ksampler_modified.png'
     )
   })
+
+  test.describe('Restore all open workflows on reload', () => {
+    let workflowA: string
+    let workflowB: string
+
+    const generateUniqueFilename = (extension = '') =>
+      `${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 8)}${extension}`
+
+    test.beforeEach(async ({ comfyPage }) => {
+      await comfyPage.setSetting('Comfy.UseNewMenu', 'Top')
+
+      workflowA = generateUniqueFilename()
+      await comfyPage.menu.topbar.saveWorkflow(workflowA)
+      workflowB = generateUniqueFilename()
+      await comfyPage.menu.topbar.triggerTopbarCommand(['Workflow', 'New'])
+      await comfyPage.menu.topbar.saveWorkflow(workflowB)
+
+      // Wait for localStorage to persist the workflow paths before reloading
+      await comfyPage.page.waitForFunction(
+        () => !!window.localStorage.getItem('Comfy.OpenWorkflowsPaths')
+      )
+      await comfyPage.setup({ clearStorage: false })
+    })
+
+    test('Restores topbar workflow tabs after reload', async ({
+      comfyPage
+    }) => {
+      await comfyPage.setSetting(
+        'Comfy.Workflow.WorkflowTabsPosition',
+        'Topbar'
+      )
+      const tabs = await comfyPage.menu.topbar.getTabNames()
+      const activeWorkflowName = await comfyPage.menu.topbar.getActiveTabName()
+
+      expect(tabs).toEqual(expect.arrayContaining([workflowA, workflowB]))
+      expect(tabs.indexOf(workflowA)).toBeLessThan(tabs.indexOf(workflowB))
+      expect(activeWorkflowName).toEqual(workflowB)
+    })
+
+    test('Restores sidebar workflows after reload', async ({ comfyPage }) => {
+      await comfyPage.setSetting(
+        'Comfy.Workflow.WorkflowTabsPosition',
+        'Sidebar'
+      )
+      await comfyPage.menu.workflowsTab.open()
+      const openWorkflows =
+        await comfyPage.menu.workflowsTab.getOpenedWorkflowNames()
+      const activeWorkflowName =
+        await comfyPage.menu.workflowsTab.getActiveWorkflowName()
+      const workflowPathA = `${workflowA}.json`
+      const workflowPathB = `${workflowB}.json`
+
+      expect(openWorkflows).toEqual(
+        expect.arrayContaining([workflowPathA, workflowPathB])
+      )
+      expect(openWorkflows.indexOf(workflowPathA)).toBeLessThan(
+        openWorkflows.indexOf(workflowPathB)
+      )
+      expect(activeWorkflowName).toEqual(workflowPathB)
+    })
+  })
 })
 
 test.describe('Load duplicate workflow', () => {

src/components/graph/GraphCanvas.vue

@@ -57,7 +57,7 @@ import { usePragmaticDroppable } from '@/hooks/dndHooks'
 import { api } from '@/scripts/api'
 import { app as comfyApp } from '@/scripts/app'
 import { ChangeTracker } from '@/scripts/changeTracker'
-import { setStorageValue } from '@/scripts/utils'
+import { getStorageValue, setStorageValue } from '@/scripts/utils'
 import { IS_CONTROL_WIDGET, updateControlWidgetLabel } from '@/scripts/widgets'
 import { useColorPaletteService } from '@/services/colorPaletteService'
 import { useLitegraphService } from '@/services/litegraphService'
@@ -95,6 +95,29 @@ const canvasMenuEnabled = computed(() =>
 )
 const tooltipEnabled = computed(() => settingStore.get('Comfy.EnableTooltips'))
 
+const storedWorkflows = JSON.parse(
+  getStorageValue('Comfy.OpenWorkflowsPaths') || '[]'
+)
+const storedActiveIndex = JSON.parse(
+  getStorageValue('Comfy.ActiveWorkflowIndex') || '-1'
+)
+const openWorkflows = computed(() => workspaceStore?.workflow?.openWorkflows)
+const activeWorkflow = computed(() => workspaceStore?.workflow?.activeWorkflow)
+const restoreState = computed<{ paths: string[]; activeIndex: number }>(() => {
+  if (!openWorkflows.value || !activeWorkflow.value) {
+    return { paths: [], activeIndex: -1 }
+  }
+
+  const paths = openWorkflows.value
+    .filter((workflow) => workflow?.isPersisted && !workflow.isModified)
+    .map((workflow) => workflow.path)
+  const activeIndex = openWorkflows.value.findIndex(
+    (workflow) => workflow.path === activeWorkflow.value?.path
+  )
+
+  return { paths, activeIndex }
+})
+
 watchEffect(() => {
   const canvasInfoEnabled = settingStore.get('Comfy.Graph.CanvasInfo')
   if (canvasStore.canvas) {
@@ -364,6 +387,18 @@ onMounted(async () => {
     'Comfy.CustomColorPalettes'
   )
 
+  const isRestorable = storedWorkflows?.length > 0 && storedActiveIndex >= 0
+  if (isRestorable)
+    workflowStore.openWorkflowsInBackground({
+      left: storedWorkflows.slice(0, storedActiveIndex),
+      right: storedWorkflows.slice(storedActiveIndex)
+    })
+
+  watch(restoreState, ({ paths, activeIndex }) => {
+    setStorageValue('Comfy.OpenWorkflowsPaths', JSON.stringify(paths))
+    setStorageValue('Comfy.ActiveWorkflowIndex', JSON.stringify(activeIndex))
+  })
+
   // Start watching for locale change after the initial value is loaded.
   watch(
     () => settingStore.get('Comfy.Locale'),

src/services/workflowService.ts

@@ -176,10 +176,6 @@ export const useWorkflowService = () => {
     workflow: ComfyWorkflow,
     options: { warnIfUnsaved: boolean } = { warnIfUnsaved: true }
   ): Promise<boolean> => {
-    if (!workflow.isLoaded) {
-      return true
-    }
-
     if (workflow.isModified && options.warnIfUnsaved) {
       const confirmed = await dialogService.confirm({
         title: t('sideToolbar.workflowTab.dirtyCloseTitle'),

src/stores/workflowStore.ts

@@ -1,3 +1,4 @@
+import _ from 'lodash'
 import { defineStore } from 'pinia'
 import { computed, markRaw, ref } from 'vue'
 
@@ -130,6 +131,10 @@ export interface WorkflowStore {
   openWorkflows: LoadedComfyWorkflow[]
   openedWorkflowIndexShift: (shift: number) => LoadedComfyWorkflow | null
   openWorkflow: (workflow: ComfyWorkflow) => Promise<LoadedComfyWorkflow>
+  openWorkflowsInBackground: (paths: {
+    left?: string[]
+    right?: string[]
+  }) => void
   isOpen: (workflow: ComfyWorkflow) => boolean
   isBusy: boolean
   closeWorkflow: (workflow: ComfyWorkflow) => Promise<void>
@@ -213,6 +218,36 @@ export const useWorkflowStore = defineStore('workflow', () => {
   const isOpen = (workflow: ComfyWorkflow) =>
     openWorkflowPathSet.value.has(workflow.path)
 
+  /**
+   * Add paths to the list of open workflow paths without loading the files
+   * or changing the active workflow.
+   *
+   * @param paths - The workflows to open, specified as:
+   *   - `left`: Workflows to be added to the left.
+   *   - `right`: Workflows to be added to the right.
+   *
+   * Invalid paths (non-strings or paths not found in `workflowLookup.value`)
+   * will be ignored. Duplicate paths are automatically removed.
+   */
+  const openWorkflowsInBackground = (paths: {
+    left?: string[]
+    right?: string[]
+  }) => {
+    const { left = [], right = [] } = paths
+    if (!left.length && !right.length) return
+
+    const isValidPath = (
+      path: unknown
+    ): path is keyof typeof workflowLookup.value =>
+      typeof path === 'string' && path in workflowLookup.value
+
+    openWorkflowPaths.value = _.union(
+      left,
+      openWorkflowPaths.value,
+      right
+    ).filter(isValidPath)
+  }
+
   /**
    * Set the workflow as the active workflow.
    * @param workflow The workflow to open.
@@ -389,6 +424,7 @@ export const useWorkflowStore = defineStore('workflow', () => {
     openWorkflows,
     openedWorkflowIndexShift,
     openWorkflow,
+    openWorkflowsInBackground,
     isOpen,
     isBusy,
     closeWorkflow,

tests-ui/tests/store/workflowStore.test.ts

@@ -3,6 +3,7 @@ import { createPinia, setActivePinia } from 'pinia'
 import { api } from '@/scripts/api'
 import { defaultGraph, defaultGraphJSON } from '@/scripts/defaultGraph'
 import {
+  ComfyWorkflow,
   LoadedComfyWorkflow,
   useWorkflowBookmarkStore,
   useWorkflowStore
@@ -161,6 +162,97 @@ describe('useWorkflowStore', () => {
     })
   })
 
+  describe('openWorkflowsInBackground', () => {
+    let workflowA: ComfyWorkflow
+    let workflowB: ComfyWorkflow
+    let workflowC: ComfyWorkflow
+
+    const openWorkflowPaths = () =>
+      store.openWorkflows.filter((w) => store.isOpen(w)).map((w) => w.path)
+
+    beforeEach(async () => {
+      await syncRemoteWorkflows(['a.json', 'b.json', 'c.json'])
+      workflowA = store.getWorkflowByPath('workflows/a.json')!
+      workflowB = store.getWorkflowByPath('workflows/b.json')!
+      workflowC = store.getWorkflowByPath('workflows/c.json')!
+      ;(api.getUserData as jest.Mock).mockResolvedValue({
+        status: 200,
+        text: () => Promise.resolve(defaultGraphJSON)
+      })
+    })
+
+    it('should open workflows adjacent to the active workflow', async () => {
+      await store.openWorkflow(workflowA)
+      store.openWorkflowsInBackground({
+        left: [workflowB.path],
+        right: [workflowC.path]
+      })
+      expect(openWorkflowPaths()).toEqual([
+        workflowB.path,
+        workflowA.path,
+        workflowC.path
+      ])
+    })
+
+    it('should not change the active workflow', async () => {
+      await store.openWorkflow(workflowA)
+      store.openWorkflowsInBackground({
+        left: [workflowC.path],
+        right: [workflowB.path]
+      })
+      expect(store.activeWorkflow).not.toBeUndefined()
+      expect(store.activeWorkflow!.path).toBe(workflowA.path)
+    })
+
+    it('should open workflows when none are active', async () => {
+      expect(store.openWorkflows.length).toBe(0)
+      store.openWorkflowsInBackground({
+        left: [workflowA.path],
+        right: [workflowB.path]
+      })
+      expect(openWorkflowPaths()).toEqual([workflowA.path, workflowB.path])
+    })
+
+    it('should not open duplicate workflows', async () => {
+      store.openWorkflowsInBackground({
+        left: [workflowA.path, workflowB.path, workflowA.path],
+        right: [workflowB.path, workflowA.path, workflowB.path]
+      })
+      expect(openWorkflowPaths()).toEqual([workflowA.path, workflowB.path])
+    })
+
+    it('should not open workflow that is already open', async () => {
+      await store.openWorkflow(workflowA)
+      store.openWorkflowsInBackground({
+        left: [workflowA.path]
+      })
+      expect(openWorkflowPaths()).toEqual([workflowA.path])
+      expect(store.activeWorkflow?.path).toBe(workflowA.path)
+    })
+
+    it('should ignore invalid or deleted workflow paths', async () => {
+      await store.openWorkflow(workflowA)
+      store.openWorkflowsInBackground({
+        left: ['workflows/invalid::$-path.json'],
+        right: ['workflows/deleted-since-last-session.json']
+      })
+      expect(openWorkflowPaths()).toEqual([workflowA.path])
+      expect(store.activeWorkflow?.path).toBe(workflowA.path)
+    })
+
+    it('should do nothing when given an empty argument', async () => {
+      await store.openWorkflow(workflowA)
+
+      store.openWorkflowsInBackground({})
+      expect(openWorkflowPaths()).toEqual([workflowA.path])
+
+      store.openWorkflowsInBackground({ left: [], right: [] })
+      expect(openWorkflowPaths()).toEqual([workflowA.path])
+
+      expect(store.activeWorkflow?.path).toBe(workflowA.path)
+    })
+  })
+
   describe('renameWorkflow', () => {
     it('should rename workflow and update bookmarks', async () => {
       const workflow = store.createTemporary('dir/test.json')