Only persist workflow on workflow change/switch (#1576)

* Only persist workflow on workflow change/switch

* nit

* Add playwright test

* Add modify test

* nit

* Fix playwright tests

* Update test expectations [skip ci]

---------

Co-authored-by: github-actions <github-actions@github.com>

browser_tests/fixtures/ComfyPage.ts

@@ -204,13 +204,15 @@ export class ComfyPage {
     }
   }
 
-  async setup() {
+  async setup({ clearStorage = true }: { clearStorage?: boolean } = {}) {
     await this.goto()
-    await this.page.evaluate((id) => {
-      localStorage.clear()
-      sessionStorage.clear()
-      localStorage.setItem('Comfy.userId', id)
-    }, this.id)
+    if (clearStorage) {
+      await this.page.evaluate((id) => {
+        localStorage.clear()
+        sessionStorage.clear()
+        localStorage.setItem('Comfy.userId', id)
+      }, this.id)
+    }
     await this.goto()
 
     // Unify font for consistent screenshots.
@@ -314,9 +316,9 @@ export class ComfyPage {
     }, settingId)
   }
 
-  async reload() {
+  async reload({ clearStorage = true }: { clearStorage?: boolean } = {}) {
     await this.page.reload({ timeout: 15000 })
-    await this.setup()
+    await this.setup({ clearStorage })
   }
 
   async goto() {

browser_tests/interaction.spec.ts

@@ -537,6 +537,34 @@ test.describe('Load workflow', () => {
     await comfyPage.loadWorkflow('string_input')
     await expect(comfyPage.canvas).toHaveScreenshot('string_input.png')
   })
+
+  test('Restore workflow on reload (switch workflow)', async ({
+    comfyPage
+  }) => {
+    await comfyPage.loadWorkflow('single_ksampler')
+    await expect(comfyPage.canvas).toHaveScreenshot('single_ksampler.png')
+    await comfyPage.reload({ clearStorage: false })
+    await expect(comfyPage.canvas).toHaveScreenshot('single_ksampler.png')
+  })
+
+  test('Restore workflow on reload (modify workflow)', async ({
+    comfyPage
+  }) => {
+    await comfyPage.loadWorkflow('single_ksampler')
+    const node = (await comfyPage.getFirstNodeRef())!
+    await node.click('collapse')
+    // Wait 300ms between 2 clicks so that it is not treated as a double click
+    // by litegraph.
+    await comfyPage.page.waitForTimeout(300)
+    await comfyPage.clickEmptySpace()
+    await expect(comfyPage.canvas).toHaveScreenshot(
+      'single_ksampler_modified.png'
+    )
+    await comfyPage.reload({ clearStorage: false })
+    await expect(comfyPage.canvas).toHaveScreenshot(
+      'single_ksampler_modified.png'
+    )
+  })
 })
 
 test.describe('Load duplicate workflow', () => {

src/components/graph/GraphCanvas.vue

@@ -61,6 +61,7 @@ import { usePragmaticDroppable } from '@/hooks/dndHooks'
 import { useWorkflowStore } from '@/stores/workflowStore'
 import { setStorageValue } from '@/scripts/utils'
 import { ChangeTracker } from '@/scripts/changeTracker'
+import { api } from '@/scripts/api'
 
 const emit = defineEmits(['ready'])
 const canvasRef = ref<HTMLCanvasElement | null>(null)
@@ -178,13 +179,26 @@ watchEffect(() => {
 })
 
 const workflowStore = useWorkflowStore()
+const persistCurrentWorkflow = () => {
+  const workflow = JSON.stringify(comfyApp.serializeGraph())
+  localStorage.setItem('workflow', workflow)
+  if (api.clientId) {
+    sessionStorage.setItem(`workflow:${api.clientId}`, workflow)
+  }
+}
+
 watchEffect(() => {
   if (workflowStore.activeWorkflow) {
     const workflow = workflowStore.activeWorkflow
     setStorageValue('Comfy.PreviousWorkflow', workflow.key)
+    // When the activeWorkflow changes, the graph has already been loaded.
+    // Saving the current state of the graph to the localStorage.
+    persistCurrentWorkflow()
   }
 })
 
+api.addEventListener('graphChanged', persistCurrentWorkflow)
+
 usePragmaticDroppable(() => canvasRef.value, {
   onDrop: (event) => {
     const loc = event.location.current.input

src/scripts/app.ts

@@ -1865,15 +1865,6 @@ export class ComfyApp {
       await this.loadGraphData()
     }
 
-    // Save current workflow automatically
-    setInterval(() => {
-      const workflow = JSON.stringify(this.serializeGraph())
-      localStorage.setItem('workflow', workflow)
-      if (api.clientId) {
-        sessionStorage.setItem(`workflow:${api.clientId}`, workflow)
-      }
-    }, 1000)
-
     this.#addDrawNodeHandler()
     this.#addDrawGroupsHandler()
     this.#addDropHandler()

src/stores/coreSettings.ts

@@ -1,9 +1,6 @@
 import type { Keybinding } from '@/types/keyBindingTypes'
 import { NodeBadgeMode } from '@/types/nodeSource'
-import {
-  LinkReleaseTriggerAction,
-  LinkReleaseTriggerMode
-} from '@/types/searchBoxTypes'
+import { LinkReleaseTriggerAction } from '@/types/searchBoxTypes'
 import type { SettingParams } from '@/types/settingTypes'
 import { LinkMarkerShape } from '@comfyorg/litegraph'
 import { LiteGraph } from '@comfyorg/litegraph'