From 04e13446767261f6b8853692b671b1e582940471 Mon Sep 17 00:00:00 2001
From: bymyself <abolkonsky.rem@gmail.com>
Date: Mon, 30 Sep 2024 12:09:24 -0700
Subject: [PATCH] Fix closing saved workflows (#1049)

* Fix workflow save. Resolves #996 Resolves #1048

* Add test on closing tabs

* Add test on closing open workflows in sidebar
---
 browser_tests/ComfyPage.ts | 28 ++++++++++++++++++++++++++++
 browser_tests/menu.spec.ts | 32 ++++++++++++++++++++++++++++++++
 src/scripts/workflows.ts   |  6 ++++++
 3 files changed, 66 insertions(+)

diff --git a/browser_tests/ComfyPage.ts b/browser_tests/ComfyPage.ts
index 7426b082e1..cde31a0dd3 100644
--- a/browser_tests/ComfyPage.ts
+++ b/browser_tests/ComfyPage.ts
@@ -230,6 +230,34 @@ class Topbar {
       .allInnerTexts()
   }
 
+  async getMenuItem(itemLabel: string): Promise<Locator> {
+    return this.page.locator(`.p-menubar-item-label:text-is("${itemLabel}")`)
+  }
+
+  async getWorkflowTab(tabName: string): Promise<Locator> {
+    return this.page
+      .locator(`.workflow-tabs .workflow-label:has-text("${tabName}")`)
+      .locator('..')
+  }
+
+  async closeWorkflowTab(tabName: string) {
+    const tab = await this.getWorkflowTab(tabName)
+    await tab.locator('.close-button').click({ force: true })
+  }
+
+  async saveWorkflow(workflowName: string) {
+    this.page.on('dialog', async (dialog) => {
+      await dialog.accept(workflowName)
+    })
+    const workflowMenuItem = await this.getMenuItem('Workflow')
+    workflowMenuItem.click()
+    await this.page.evaluate(() => {
+      return new Promise<number>(requestAnimationFrame)
+    })
+    const saveButton = await this.getMenuItem('Save')
+    await saveButton.click()
+  }
+
   async triggerTopbarCommand(path: string[]) {
     if (path.length < 2) {
       throw new Error('Path is too short')
diff --git a/browser_tests/menu.spec.ts b/browser_tests/menu.spec.ts
index eab8042553..241742eb4a 100644
--- a/browser_tests/menu.spec.ts
+++ b/browser_tests/menu.spec.ts
@@ -437,6 +437,19 @@ test.describe('Menu', () => {
         comfyPage.page.locator('.comfy-missing-nodes')
       ).not.toBeVisible()
     })
+
+    test('Can close saved-workflows from the open workflows section', async ({
+      comfyPage
+    }) => {
+      await comfyPage.menu.topbar.saveWorkflow('deault')
+      const closeButton = comfyPage.page.locator(
+        '.comfyui-workflows-open .p-button-icon.pi-times'
+      )
+      await closeButton.click()
+      expect(
+        await comfyPage.menu.workflowsTab.getOpenedWorkflowNames()
+      ).toEqual(['*Unsaved Workflow (2).json'])
+    })
   })
 
   test.describe('Workflows topbar tabs', () => {
@@ -447,11 +460,30 @@ test.describe('Menu', () => {
       )
     })
 
+    test.afterEach(async ({ comfyPage }) => {
+      // Delete the saved workflow for cleanup.
+      await comfyPage.page.evaluate(async () => {
+        window['app'].workflowManager.activeWorkflow.delete()
+      })
+    })
+
     test('Can show opened workflows', async ({ comfyPage }) => {
       expect(await comfyPage.menu.topbar.getTabNames()).toEqual([
         'Unsaved Workflow'
       ])
     })
+
+    test('Can close saved-workflow tabs', async ({ comfyPage }) => {
+      const savedWorkflowName = 'default'
+      await comfyPage.menu.topbar.saveWorkflow(savedWorkflowName)
+      expect(await comfyPage.menu.topbar.getTabNames()).toEqual([
+        savedWorkflowName
+      ])
+      await comfyPage.menu.topbar.closeWorkflowTab('default')
+      expect(await comfyPage.menu.topbar.getTabNames()).toEqual([
+        'Unsaved Workflow (2)'
+      ])
+    })
   })
 
   // Only test 'Top' to reduce test time.
diff --git a/src/scripts/workflows.ts b/src/scripts/workflows.ts
index 165762a742..8ad88c807f 100644
--- a/src/scripts/workflows.ts
+++ b/src/scripts/workflows.ts
@@ -405,7 +405,13 @@ export class ComfyWorkflow {
 
     if (!this.path) {
       // Saved new workflow, patch this instance
+      const oldKey = this.key
       this.updatePath(path, null)
+
+      // Update workflowLookup: change the key from the old unsaved path to the new saved path
+      delete this.manager.workflowStore.workflowLookup[oldKey]
+      this.manager.workflowStore.workflowLookup[this.key] = this
+
       await this.manager.loadWorkflows()
       this.unsaved = false
       this.manager.dispatchEvent(new CustomEvent('rename', { detail: this }))