[Refactor] Support handleAddFolder in TreeExplorer (#3101)

browser_tests/menu.spec.ts

@@ -119,7 +119,10 @@ test.describe('Menu', () => {
     test('Can add new bookmark folder', async ({ comfyPage }) => {
       const tab = comfyPage.menu.nodeLibraryTab
       await tab.newFolderButton.click()
-      await comfyPage.page.keyboard.press('Enter')
+      const textInput = comfyPage.page.locator('.editable-text input')
+      await textInput.waitFor({ state: 'visible' })
+      await textInput.fill('New Folder')
+      await textInput.press('Enter')
       expect(await tab.getFolder('New Folder').count()).toBe(1)
       expect(
         await comfyPage.getSetting('Comfy.NodeLibrary.Bookmarks.V2')
@@ -132,8 +135,10 @@ test.describe('Menu', () => {
 
       await tab.getFolder('foo').click({ button: 'right' })
       await comfyPage.page.getByLabel('New Folder').click()
-      await comfyPage.page.keyboard.type('bar')
-      await comfyPage.page.keyboard.press('Enter')
+      const textInput = comfyPage.page.locator('.editable-text input')
+      await textInput.waitFor({ state: 'visible' })
+      await textInput.fill('bar')
+      await textInput.press('Enter')
 
       expect(await tab.getFolder('bar').count()).toBe(1)
       expect(

src/components/common/TreeExplorer.vue

@@ -43,14 +43,18 @@ import { computed, provide, ref } from 'vue'
 import { useI18n } from 'vue-i18n'
 
 import TreeExplorerTreeNode from '@/components/common/TreeExplorerTreeNode.vue'
+import { useTreeFolderOperations } from '@/composables/tree/useTreeFolderOperations'
 import { useErrorHandling } from '@/composables/useErrorHandling'
 import {
+  InjectKeyExpandedKeys,
   InjectKeyHandleEditLabelFunction,
   type RenderedTreeExplorerNode,
   type TreeExplorerNode
 } from '@/types/treeExplorerTypes'
+import { combineTrees, findNodeByKey } from '@/utils/treeUtil'
 
 const expandedKeys = defineModel<Record<string, boolean>>('expandedKeys')
+provide(InjectKeyExpandedKeys, expandedKeys)
 const selectionKeys = defineModel<Record<string, boolean>>('selectionKeys')
 // Tracks whether the caller has set the selectionKeys model.
 const storeSelectionKeys = selectionKeys.value !== undefined
@@ -64,8 +68,23 @@ const emit = defineEmits<{
   (e: 'nodeDelete', node: RenderedTreeExplorerNode): void
   (e: 'contextMenu', node: RenderedTreeExplorerNode, event: MouseEvent): void
 }>()
+
+const {
+  newFolderNode,
+  getAddFolderMenuItem,
+  handleFolderCreation,
+  addFolderCommand
+} = useTreeFolderOperations(
+  /* expandNode */ (node: TreeExplorerNode) => {
+    expandedKeys.value[node.key] = true
+  }
+)
+
 const renderedRoot = computed<RenderedTreeExplorerNode>(() => {
-  return fillNodeInfo(props.root)
+  const renderedRoot = fillNodeInfo(props.root)
+  return newFolderNode.value
+    ? combineTrees(renderedRoot, newFolderNode.value)
+    : renderedRoot
 })
 const getTreeNodeIcon = (node: TreeExplorerNode) => {
   if (node.getIcon) {
@@ -127,7 +146,11 @@ const handleNodeLabelEdit = async (
 ) => {
   await errorHandling.wrapWithErrorHandlingAsync(
     async () => {
-      await node.handleRename(newName)
+      if (node.key === newFolderNode.value?.key) {
+        await handleFolderCreation(newName)
+      } else {
+        await node.handleRename(newName)
+      }
     },
     node.handleError,
     () => {
@@ -147,6 +170,7 @@ const deleteCommand = async (node: RenderedTreeExplorerNode) => {
 }
 const menuItems = computed<MenuItem[]>(() =>
   [
+    getAddFolderMenuItem(menuTargetNode.value),
     {
       label: t('g.rename'),
       icon: 'pi pi-file-edit',
@@ -194,7 +218,14 @@ const wrapCommandWithErrorHandler = (
 
 defineExpose({
   renameCommand,
-  deleteCommand
+  deleteCommand,
+  /**
+   * The command to add a folder to a node via the context menu
+   * @param targetNodeKey - The key of the node where the folder will be added under
+   */
+  addFolderCommand: (targetNodeKey: string) => {
+    addFolderCommand(findNodeByKey(renderedRoot.value, targetNodeKey))
+  }
 })
 </script>
 

src/components/sidebar/tabs/nodeLibrary/NodeBookmarkTreeExplorer.vue

@@ -40,7 +40,6 @@ import type {
   TreeExplorerDragAndDropData,
   TreeExplorerNode
 } from '@/types/treeExplorerTypes'
-import { findNodeByKey } from '@/utils/treeUtil'
 
 const props = defineProps<{
   filteredNodeDefs: ComfyNodeDefImpl[]
@@ -94,14 +93,6 @@ const { t } = useI18n()
 const extraMenuItems = (
   menuTargetNode: RenderedTreeExplorerNode<ComfyNodeDefImpl>
 ) => [
-  {
-    label: t('g.newFolder'),
-    icon: 'pi pi-folder-plus',
-    command: () => {
-      addNewBookmarkFolder(menuTargetNode)
-    },
-    visible: !menuTargetNode?.leaf
-  },
   {
     label: t('g.customize'),
     icon: 'pi pi-palette',
@@ -152,6 +143,11 @@ const renderedBookmarkedRoot = computed<TreeExplorerNode<ComfyNodeDefImpl>>(
         },
         children: sortedChildren,
         draggable: node.leaf,
+        handleAddFolder(newName: string) {
+          if (newName !== '') {
+            nodeBookmarkStore.addNewBookmarkFolder(this.data, newName)
+          }
+        },
         renderDragPreview(container) {
           const vnode = h(NodePreview, { nodeDef: node.data })
           render(vnode, container)
@@ -197,25 +193,8 @@ const renderedBookmarkedRoot = computed<TreeExplorerNode<ComfyNodeDefImpl>>(
 )
 
 const treeExplorerRef = ref<InstanceType<typeof TreeExplorer> | null>(null)
-const addNewBookmarkFolder = (
-  parent?: RenderedTreeExplorerNode<ComfyNodeDefImpl>
-) => {
-  const newFolderKey =
-    'root/' + nodeBookmarkStore.addNewBookmarkFolder(parent?.data).slice(0, -1)
-  nextTick(() => {
-    treeExplorerRef.value?.renameCommand(
-      findNodeByKey(
-        renderedBookmarkedRoot.value,
-        newFolderKey
-      ) as RenderedTreeExplorerNode
-    )
-    if (parent) {
-      expandedKeys.value[parent.key] = true
-    }
-  })
-}
 defineExpose({
-  addNewBookmarkFolder
+  addNewBookmarkFolder: () => treeExplorerRef.value?.addFolderCommand('root')
 })
 
 const showCustomizationDialog = ref(false)

src/components/sidebar/tabs/nodeLibrary/NodeTreeFolder.vue

@@ -8,21 +8,27 @@
 </template>
 
 <script setup lang="ts">
-import { Ref, computed, inject, onMounted, onUnmounted, ref, watch } from 'vue'
+import { computed, inject, onMounted, onUnmounted, ref, watch } from 'vue'
 
 import TreeExplorerTreeNode from '@/components/common/TreeExplorerTreeNode.vue'
 import type { BookmarkCustomization } from '@/schemas/apiSchema'
 import { useNodeBookmarkStore } from '@/stores/nodeBookmarkStore'
 import { ComfyNodeDefImpl } from '@/stores/nodeDefStore'
-import { RenderedTreeExplorerNode } from '@/types/treeExplorerTypes'
+import {
+  InjectKeyExpandedKeys,
+  type RenderedTreeExplorerNode
+} from '@/types/treeExplorerTypes'
 
-const props = defineProps<{
+const { node } = defineProps<{
   node: RenderedTreeExplorerNode<ComfyNodeDefImpl>
 }>()
 
 const nodeBookmarkStore = useNodeBookmarkStore()
 const customization = computed<BookmarkCustomization | undefined>(() => {
-  return nodeBookmarkStore.bookmarksCustomization[props.node.data.nodePath]
+  const nodeDef = node.data
+  return nodeDef
+    ? nodeBookmarkStore.bookmarksCustomization[nodeDef.nodePath]
+    : undefined
 })
 
 const treeNodeElement = ref<HTMLElement | null>(null)
@@ -56,7 +62,7 @@ onUnmounted(() => {
   }
 })
 
-const expandedKeys = inject<Ref<Record<string, boolean>>>('expandedKeys')
+const expandedKeys = inject(InjectKeyExpandedKeys)
 const handleItemDrop = (node: RenderedTreeExplorerNode) => {
   expandedKeys.value[node.key] = true
 }

src/composables/tree/useTreeFolderOperations.ts

@@ -0,0 +1,76 @@
+import { ref } from 'vue'
+import { useI18n } from 'vue-i18n'
+
+import type { RenderedTreeExplorerNode } from '@/types/treeExplorerTypes'
+
+/**
+ * Use this to handle folder operations in a tree.
+ * @param expandNode - The function to expand a node.
+ */
+export function useTreeFolderOperations(
+  expandNode: (node: RenderedTreeExplorerNode) => void
+) {
+  const { t } = useI18n()
+  const newFolderNode = ref<RenderedTreeExplorerNode | null>(null)
+  const addFolderTargetNode = ref<RenderedTreeExplorerNode | null>(null)
+
+  // Generate a unique temporary key for the new folder
+  const generateTempKey = (parentKey: string) => {
+    return `${parentKey}/new_folder_${Date.now()}`
+  }
+
+  // Handle folder creation after name is confirmed
+  const handleFolderCreation = async (newName: string) => {
+    if (!newFolderNode.value || !addFolderTargetNode.value) return
+
+    try {
+      // Call the handleAddFolder method with the new folder name
+      await addFolderTargetNode.value?.handleAddFolder?.(newName)
+    } finally {
+      newFolderNode.value = null
+      addFolderTargetNode.value = null
+    }
+  }
+
+  /**
+   * The command to add a folder to a node via the context menu
+   * @param targetNode - The node where the folder will be added under
+   */
+  const addFolderCommand = (targetNode: RenderedTreeExplorerNode) => {
+    expandNode(targetNode)
+    newFolderNode.value = {
+      key: generateTempKey(targetNode.key),
+      label: '',
+      leaf: false,
+      children: [],
+      icon: 'pi pi-folder',
+      type: 'folder',
+      totalLeaves: 0,
+      badgeText: '',
+      isEditingLabel: true
+    }
+    addFolderTargetNode.value = targetNode
+  }
+
+  // Generate the "Add Folder" menu item
+  const getAddFolderMenuItem = (
+    targetNode: RenderedTreeExplorerNode | null
+  ) => {
+    return {
+      label: t('g.newFolder'),
+      icon: 'pi pi-folder-plus',
+      command: () => {
+        if (targetNode) addFolderCommand(targetNode)
+      },
+      visible: targetNode && !targetNode.leaf && !!targetNode.handleAddFolder,
+      isAsync: false
+    }
+  }
+
+  return {
+    newFolderNode,
+    addFolderCommand,
+    getAddFolderMenuItem,
+    handleFolderCreation
+  }
+}

src/composables/useTreeExpansion.ts

@@ -23,30 +23,20 @@ export function useTreeExpansion(expandedKeys: Ref<Record<string, boolean>>) {
   }
 
   const expandNode = (node: TreeNode) => {
-    if (
-      node.key &&
-      typeof node.key === 'string' &&
-      node.children &&
-      node.children.length
-    ) {
+    if (node.key && typeof node.key === 'string' && !node.leaf) {
       expandedKeys.value[node.key] = true
 
-      for (const child of node.children) {
+      for (const child of node.children ?? []) {
         expandNode(child)
       }
     }
   }
 
   const collapseNode = (node: TreeNode) => {
-    if (
-      node.key &&
-      typeof node.key === 'string' &&
-      node.children &&
-      node.children.length
-    ) {
+    if (node.key && typeof node.key === 'string' && !node.leaf) {
       delete expandedKeys.value[node.key]
 
-      for (const child of node.children) {
+      for (const child of node.children ?? []) {
         collapseNode(child)
       }
     }

src/stores/nodeBookmarkStore.ts

@@ -73,14 +73,12 @@ export const useNodeBookmarkStore = defineStore('nodeBookmark', () => {
     )
   }
 
-  const addNewBookmarkFolder = (parent?: ComfyNodeDefImpl) => {
+  const addNewBookmarkFolder = (
+    parent: ComfyNodeDefImpl | undefined,
+    folderName: string
+  ) => {
     const parentPath = parent ? parent.nodePath : ''
-    let newFolderPath = parentPath + 'New Folder/'
-    let suffix = 1
-    while (bookmarks.value.some((b: string) => b.startsWith(newFolderPath))) {
-      newFolderPath = parentPath + `New Folder ${suffix}/`
-      suffix++
-    }
+    const newFolderPath = parentPath + folderName + '/'
     addBookmark(newFolderPath)
     return newFolderPath
   }

src/types/treeExplorerTypes.ts

@@ -1,5 +1,5 @@
 import type { MenuItem } from 'primevue/menuitem'
-import type { InjectionKey } from 'vue'
+import type { InjectionKey, Ref } from 'vue'
 
 export interface TreeExplorerNode<T = any> {
   key: string
@@ -19,10 +19,10 @@ export interface TreeExplorerNode<T = any> {
   ) => void | Promise<void>
   /** Function to handle deleting the node */
   handleDelete?: (this: TreeExplorerNode<T>) => void | Promise<void>
-  /** Function to handle adding a child node */
-  handleAddChild?: (
+  /** Function to handle adding a folder */
+  handleAddFolder?: (
     this: TreeExplorerNode<T>,
-    child: TreeExplorerNode<T>
+    folderName: string
   ) => void | Promise<void>
   /** Whether the node is draggable */
   draggable?: boolean
@@ -71,3 +71,6 @@ export type TreeExplorerDragAndDropData<T = any> = {
 export const InjectKeyHandleEditLabelFunction: InjectionKey<
   (node: RenderedTreeExplorerNode, newName: string) => void
 > = Symbol()
+
+export const InjectKeyExpandedKeys: InjectionKey<Ref<Record<string, boolean>>> =
+  Symbol()

src/utils/treeUtil.ts

@@ -105,7 +105,10 @@ export function sortedTree(
   return newNode
 }
 
-export const findNodeByKey = (root: TreeNode, key: string): TreeNode | null => {
+export const findNodeByKey = <T extends TreeNode>(
+  root: T,
+  key: string
+): T | null => {
   if (root.key === key) {
     return root
   }
@@ -113,10 +116,46 @@ export const findNodeByKey = (root: TreeNode, key: string): TreeNode | null => {
     return null
   }
   for (const child of root.children) {
-    const result = findNodeByKey(child, key)
+    const result = findNodeByKey(child as T, key)
     if (result) {
       return result
     }
   }
   return null
 }
+
+/**
+ * Deep clone a tree node and its children.
+ * @param node - The node to clone.
+ * @returns A deep clone of the node.
+ */
+export function cloneTree<T extends TreeNode>(node: T): T {
+  const clone: T = { ...node } as T
+
+  // Clone children recursively
+  if (node.children && node.children.length > 0) {
+    clone.children = node.children.map((child) => cloneTree(child as T))
+  }
+
+  return clone
+}
+
+/**
+ * Merge a subtree into the tree.
+ * @param root - The root of the tree.
+ * @param subtree - The subtree to merge.
+ * @returns A new tree with the subtree merged.
+ */
+export const combineTrees = <T extends TreeNode>(root: T, subtree: T): T => {
+  const newRoot = cloneTree(root)
+
+  const parentKey = subtree.key.slice(0, subtree.key.lastIndexOf('/'))
+  const parent = findNodeByKey(newRoot, parentKey)
+
+  if (parent) {
+    parent.children ??= []
+    parent.children.push(cloneTree(subtree))
+  }
+
+  return newRoot
+}