From 73bdbddf90af6db19ce1f6c6d8b14335f98e978a Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Mon, 11 Nov 2024 11:06:41 -0500
Subject: [PATCH] Fix rename open/bookmark workflow (#1507)

* Fix rename open/bookmark workflow

* nit

* Fix save as

* Add browser test
---
 .../fixtures/components/SidebarTab.ts         | 29 +++++++--
 browser_tests/menu.spec.ts                    | 28 ++++++++-
 .../sidebar/tabs/WorkflowsSidebarTab.vue      | 63 ++++++++++++++++---
 src/services/workflowService.ts               | 15 ++---
 src/stores/workflowStore.ts                   | 42 +++----------
 .../tests/fast/store/workflowStore.test.ts    |  3 +-
 6 files changed, 121 insertions(+), 59 deletions(-)

diff --git a/browser_tests/fixtures/components/SidebarTab.ts b/browser_tests/fixtures/components/SidebarTab.ts
index 6e21e49e0..76b74c6dc 100644
--- a/browser_tests/fixtures/components/SidebarTab.ts
+++ b/browser_tests/fixtures/components/SidebarTab.ts
@@ -1,4 +1,4 @@
-import { Page } from '@playwright/test'
+import { Locator, Page } from '@playwright/test'
 
 class SidebarTab {
   constructor(
@@ -110,11 +110,30 @@ export class WorkflowsSidebarTab extends SidebarTab {
   }
 
   async switchToWorkflow(workflowName: string) {
-    const workflowLocator = this.page.locator(
-      '.comfyui-workflows-open .node-label',
-      { hasText: workflowName }
-    )
+    const workflowLocator = this.getOpenedItem(workflowName)
     await workflowLocator.click()
     await this.page.waitForTimeout(300)
   }
+
+  getOpenedItem(name: string) {
+    return this.page.locator('.comfyui-workflows-open .node-label', {
+      hasText: name
+    })
+  }
+
+  getPersistedItem(name: string) {
+    return this.page.locator('.comfyui-workflows-browse .node-label', {
+      hasText: name
+    })
+  }
+
+  async renameWorkflow(locator: Locator, newName: string) {
+    await locator.click({ button: 'right' })
+    await this.page
+      .locator('.p-contextmenu-item-content', { hasText: 'Rename' })
+      .click()
+    await this.page.keyboard.type(newName)
+    await this.page.keyboard.press('Enter')
+    await this.page.waitForTimeout(300)
+  }
 }
diff --git a/browser_tests/menu.spec.ts b/browser_tests/menu.spec.ts
index adf509238..083a4a4a5 100644
--- a/browser_tests/menu.spec.ts
+++ b/browser_tests/menu.spec.ts
@@ -401,7 +401,6 @@ test.describe('Menu', () => {
         'workflow1.json': 'default.json',
         'workflow2.json': 'default.json'
       })
-      // Avoid reset view as the button is not visible in BetaMenu UI.
       await comfyPage.setup()
 
       const tab = comfyPage.menu.workflowsTab
@@ -411,6 +410,33 @@ test.describe('Menu', () => {
       )
     })
 
+    test('Can rename nested workflow from opened workflow item', async ({
+      comfyPage
+    }) => {
+      await comfyPage.setupWorkflowsDirectory({
+        foo: {
+          'bar.json': 'default.json'
+        }
+      })
+      await comfyPage.setup()
+
+      const tab = comfyPage.menu.workflowsTab
+      await tab.open()
+      // Switch to the parent folder
+      await tab.getPersistedItem('foo').click()
+      await comfyPage.page.waitForTimeout(300)
+      // Switch to the nested workflow
+      await tab.getPersistedItem('bar').click()
+      await comfyPage.page.waitForTimeout(300)
+
+      const openedWorkflow = tab.getOpenedItem('foo/bar')
+      await tab.renameWorkflow(openedWorkflow, 'foo/baz')
+      expect(await tab.getOpenedWorkflowNames()).toEqual([
+        '*Unsaved Workflow.json',
+        'foo/baz.json'
+      ])
+    })
+
     test('Can save workflow as', async ({ comfyPage }) => {
       await comfyPage.menu.workflowsTab.newBlankWorkflowButton.click()
       await comfyPage.menu.topbar.saveWorkflowAs('workflow3.json')
diff --git a/src/components/sidebar/tabs/WorkflowsSidebarTab.vue b/src/components/sidebar/tabs/WorkflowsSidebarTab.vue
index c9ab9088f..c2dea138c 100644
--- a/src/components/sidebar/tabs/WorkflowsSidebarTab.vue
+++ b/src/components/sidebar/tabs/WorkflowsSidebarTab.vue
@@ -42,7 +42,9 @@
         >
           <TextDivider text="Open" type="dashed" class="ml-2" />
           <TreeExplorer
-            :roots="renderTreeNode(workflowStore.openWorkflowsTree).children"
+            :roots="
+              renderTreeNode(openWorkflowsTree, WorkflowTreeType.Open).children
+            "
             :selectionKeys="selectionKeys"
           >
             <template #node="{ node }">
@@ -72,7 +74,10 @@
           <TextDivider text="Bookmarks" type="dashed" class="ml-2" />
           <TreeExplorer
             :roots="
-              renderTreeNode(workflowStore.bookmarkedWorkflowsTree).children
+              renderTreeNode(
+                bookmarkedWorkflowsTree,
+                WorkflowTreeType.Bookmarks
+              ).children
             "
           >
             <template #node="{ node }">
@@ -83,7 +88,9 @@
         <div class="comfyui-workflows-browse">
           <TextDivider text="Browse" type="dashed" class="ml-2" />
           <TreeExplorer
-            :roots="renderTreeNode(workflowStore.workflowsTree).children"
+            :roots="
+              renderTreeNode(workflowsTree, WorkflowTreeType.Browse).children
+            "
             v-model:expandedKeys="expandedKeys"
           >
             <template #node="{ node }">
@@ -94,7 +101,9 @@
       </div>
       <div class="comfyui-workflows-search-panel" v-else>
         <TreeExplorer
-          :roots="renderTreeNode(filteredRoot).children"
+          :roots="
+            renderTreeNode(filteredRoot, WorkflowTreeType.Browse).children
+          "
           v-model:expandedKeys="expandedKeys"
         >
           <template #node="{ node }">
@@ -128,6 +137,8 @@ import { useTreeExpansion } from '@/hooks/treeHooks'
 import { useSettingStore } from '@/stores/settingStore'
 import { workflowService } from '@/services/workflowService'
 import { useWorkspaceStore } from '@/stores/workspaceStore'
+import { appendJsonExt } from '@/utils/formatUtil'
+import { buildTree, sortedTree } from '@/utils/treeUtil'
 
 const settingStore = useSettingStore()
 const workflowTabsPosition = computed(() =>
@@ -138,9 +149,7 @@ const searchQuery = ref('')
 const isSearching = computed(() => searchQuery.value.length > 0)
 const filteredWorkflows = ref<ComfyWorkflow[]>([])
 const filteredRoot = computed<TreeNode>(() => {
-  return workflowStore.buildWorkflowTree(
-    filteredWorkflows.value as ComfyWorkflow[]
-  )
+  return buildWorkflowTree(filteredWorkflows.value as ComfyWorkflow[])
 })
 const handleSearch = (query: string) => {
   if (query.length === 0) {
@@ -172,8 +181,37 @@ const handleCloseWorkflow = (workflow?: ComfyWorkflow) => {
   }
 }
 
-const renderTreeNode = (node: TreeNode): TreeExplorerNode<ComfyWorkflow> => {
-  const children = node.children?.map(renderTreeNode)
+enum WorkflowTreeType {
+  Open = 'Open',
+  Bookmarks = 'Bookmarks',
+  Browse = 'Browse'
+}
+
+const buildWorkflowTree = (workflows: ComfyWorkflow[]) => {
+  return buildTree(workflows, (workflow: ComfyWorkflow) =>
+    workflow.key.split('/')
+  )
+}
+
+const workflowsTree = computed(() =>
+  sortedTree(buildWorkflowTree(workflowStore.persistedWorkflows), {
+    groupLeaf: true
+  })
+)
+// Bookmarked workflows tree is flat.
+const bookmarkedWorkflowsTree = computed(() =>
+  buildTree(workflowStore.bookmarkedWorkflows, (workflow) => [workflow.key])
+)
+// Open workflows tree is flat.
+const openWorkflowsTree = computed(() =>
+  buildTree(workflowStore.openWorkflows, (workflow) => [workflow.key])
+)
+
+const renderTreeNode = (
+  node: TreeNode,
+  type: WorkflowTreeType
+): TreeExplorerNode<ComfyWorkflow> => {
+  const children = node.children?.map((child) => renderTreeNode(child, type))
 
   const workflow: ComfyWorkflow = node.data
 
@@ -194,7 +232,12 @@ const renderTreeNode = (node: TreeNode): TreeExplorerNode<ComfyWorkflow> => {
           node: TreeExplorerNode<ComfyWorkflow>,
           newName: string
         ) => {
-          await workflowService.renameWorkflow(workflow, newName)
+          const newPath =
+            type === WorkflowTreeType.Browse
+              ? workflow.directory + '/' + appendJsonExt(newName)
+              : ComfyWorkflow.basePath + appendJsonExt(newName)
+
+          await workflowService.renameWorkflow(workflow, newPath)
         },
         handleDelete: workflow.isTemporary
           ? undefined
diff --git a/src/services/workflowService.ts b/src/services/workflowService.ts
index a412f10bf..bd6355ec3 100644
--- a/src/services/workflowService.ts
+++ b/src/services/workflowService.ts
@@ -61,14 +61,15 @@ export const workflowService = {
     })
     if (!newFilename) return
 
+    const newPath = workflow.directory + '/' + appendJsonExt(newFilename)
+    const newKey = newPath.substring(ComfyWorkflow.basePath.length)
+
     if (workflow.isTemporary) {
-      await this.renameWorkflow(workflow, newFilename)
+      await this.renameWorkflow(workflow, newPath)
       await useWorkflowStore().saveWorkflow(workflow)
     } else {
       const tempWorkflow = useWorkflowStore().createTemporary(
-        (workflow.directory + '/' + appendJsonExt(newFilename)).substring(
-          'workflows/'.length
-        ),
+        newKey,
         workflow.activeState as ComfyWorkflowJSON
       )
       await this.openWorkflow(tempWorkflow)
@@ -162,8 +163,8 @@ export const workflowService = {
     await workflowStore.closeWorkflow(workflow)
   },
 
-  async renameWorkflow(workflow: ComfyWorkflow, newName: string) {
-    await useWorkflowStore().renameWorkflow(workflow, newName)
+  async renameWorkflow(workflow: ComfyWorkflow, newPath: string) {
+    await useWorkflowStore().renameWorkflow(workflow, newPath)
   },
 
   async deleteWorkflow(workflow: ComfyWorkflow) {
@@ -212,7 +213,7 @@ export const workflowService = {
     const workflowStore = useWorkspaceStore().workflow
     if (typeof value === 'string') {
       const workflow = workflowStore.getWorkflowByPath(
-        'workflows/' + appendJsonExt(value)
+        ComfyWorkflow.basePath + appendJsonExt(value)
       )
       if (workflow?.isPersisted) {
         const loadedWorkflow = await workflowStore.openWorkflow(workflow)
diff --git a/src/stores/workflowStore.ts b/src/stores/workflowStore.ts
index bb393741b..b757dd37f 100644
--- a/src/stores/workflowStore.ts
+++ b/src/stores/workflowStore.ts
@@ -1,15 +1,16 @@
 import { defineStore } from 'pinia'
 import { computed, markRaw, ref } from 'vue'
-import { buildTree, sortedTree } from '@/utils/treeUtil'
 import { api } from '@/scripts/api'
 import { UserFile } from './userFileStore'
 import { ChangeTracker } from '@/scripts/changeTracker'
 import { ComfyWorkflowJSON } from '@/types/comfyWorkflow'
-import { appendJsonExt, getPathDetails } from '@/utils/formatUtil'
+import { getPathDetails } from '@/utils/formatUtil'
 import { defaultGraphJSON } from '@/scripts/defaultGraph'
 import { syncEntities } from '@/utils/syncUtil'
 
 export class ComfyWorkflow extends UserFile {
+  static readonly basePath = 'workflows/'
+
   /**
    * The change tracker for the workflow. Non-reactive raw object.
    */
@@ -28,7 +29,7 @@ export class ComfyWorkflow extends UserFile {
   }
 
   get key() {
-    return this.path.substring('workflows/'.length)
+    return this.path.substring(ComfyWorkflow.basePath.length)
   }
 
   get activeState(): ComfyWorkflowJSON | null {
@@ -104,12 +105,6 @@ export class ComfyWorkflow extends UserFile {
     this.content = JSON.stringify(this.activeState)
     return await super.saveAs(path)
   }
-
-  async rename(newName: string) {
-    const newPath = this.directory + '/' + appendJsonExt(newName)
-    await super.rename(newPath)
-    return this
-  }
 }
 
 export interface LoadedComfyWorkflow extends ComfyWorkflow {
@@ -210,7 +205,7 @@ export const useWorkflowStore = defineStore('workflow', () => {
 
   const createTemporary = (path?: string, workflowData?: ComfyWorkflowJSON) => {
     const fullPath = getUnconflictedPath(
-      'workflows/' + (path ?? 'Unsaved Workflow.json')
+      ComfyWorkflow.basePath + (path ?? 'Unsaved Workflow.json')
     )
     const workflow = new ComfyWorkflow({
       path: fullPath,
@@ -289,33 +284,15 @@ export const useWorkflowStore = defineStore('workflow', () => {
     workflows.value.filter((workflow) => workflow.isModified)
   )
 
-  const buildWorkflowTree = (workflows: ComfyWorkflow[]) => {
-    return buildTree(workflows, (workflow: ComfyWorkflow) =>
-      workflow.key.split('/')
-    )
-  }
-  const workflowsTree = computed(() =>
-    sortedTree(buildWorkflowTree(persistedWorkflows.value), { groupLeaf: true })
-  )
-  // Bookmarked workflows tree is flat.
-  const bookmarkedWorkflowsTree = computed(() =>
-    buildTree(bookmarkedWorkflows.value, (workflow) => [workflow.key])
-  )
-  // Open workflows tree is flat.
-  const openWorkflowsTree = computed(() =>
-    buildTree(openWorkflows.value, (workflow) => [workflow.key])
-  )
-
-  const renameWorkflow = async (workflow: ComfyWorkflow, newName: string) => {
+  const renameWorkflow = async (workflow: ComfyWorkflow, newPath: string) => {
     // Capture all needed values upfront
     const oldPath = workflow.path
-    const newPath = workflow.directory + '/' + appendJsonExt(newName)
     const wasBookmarked = bookmarkStore.isBookmarked(oldPath)
 
     const openIndex = detachWorkflow(workflow)
     // Perform the actual rename operation first
     try {
-      await workflow.rename(newName)
+      await workflow.rename(newPath)
     } finally {
       attachWorkflow(workflow, openIndex)
     }
@@ -353,7 +330,6 @@ export const useWorkflowStore = defineStore('workflow', () => {
     activeWorkflow,
     isActive,
     openWorkflows,
-    openWorkflowsTree,
     openedWorkflowIndexShift,
     openWorkflow,
     isOpen,
@@ -365,11 +341,9 @@ export const useWorkflowStore = defineStore('workflow', () => {
 
     workflows,
     bookmarkedWorkflows,
+    persistedWorkflows,
     modifiedWorkflows,
     getWorkflowByPath,
-    workflowsTree,
-    bookmarkedWorkflowsTree,
-    buildWorkflowTree,
     syncWorkflows
   }
 })
diff --git a/tests-ui/tests/fast/store/workflowStore.test.ts b/tests-ui/tests/fast/store/workflowStore.test.ts
index 4b647f4ca..ec30456b8 100644
--- a/tests-ui/tests/fast/store/workflowStore.test.ts
+++ b/tests-ui/tests/fast/store/workflowStore.test.ts
@@ -178,9 +178,8 @@ describe('useWorkflowStore', () => {
         } as any)
 
       // Perform rename
-      const newName = 'renamed.json'
       const newPath = 'workflows/dir/renamed.json'
-      await store.renameWorkflow(workflow, newName)
+      await store.renameWorkflow(workflow, newPath)
 
       // Check that bookmark was transferred
       expect(bookmarkStore.isBookmarked(newPath)).toBe(true)