[Refactor] Support handleAddFolder in TreeExplorer (#3101)

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
 }