Node library side bar tab (#237)

* Basic tree

* Add node filter

* Fix key issue

* Add icons

* Node count

* nit

* Add node preview basics

* Node preview

* Make comfy node in node lib draggable

* Set drop target

* Add node on drop

* Drop on dynamic location

* nit

* nit

* Fix hover preview issue

* nit

* More visual diff between node and folder

* Add playwright test

* Add dep

* Get rid of screenshot test

src/App.vue

@@ -13,7 +13,7 @@
 </template>
 
 <script setup lang="ts">
-import { computed, markRaw, onMounted, ref, watch } from 'vue'
+import { computed, markRaw, onMounted, onUnmounted, ref, watch } from 'vue'
 import NodeSearchboxPopover from '@/components/NodeSearchBoxPopover.vue'
 import SideToolBar from '@/components/sidebar/SideToolBar.vue'
 import LiteGraphCanvasSplitterOverlay from '@/components/LiteGraphCanvasSplitterOverlay.vue'
@@ -23,6 +23,9 @@ import { app } from './scripts/app'
 import { useSettingStore } from './stores/settingStore'
 import { useI18n } from 'vue-i18n'
 import { useWorkspaceStore } from './stores/workspaceStateStore'
+import NodeLibrarySideBarTab from './components/sidebar/tabs/NodeLibrarySideBarTab.vue'
+import { dropTargetForElements } from '@atlaskit/pragmatic-drag-and-drop/element/adapter'
+import { useNodeDefStore } from './stores/nodeDefStore'
 
 const isLoading = ref(true)
 const nodeSearchEnabled = computed<boolean>(
@@ -49,6 +52,7 @@ const betaMenuEnabled = computed(
 )
 
 const { t } = useI18n()
+let dropTargetCleanup = () => {}
 const init = () => {
   useSettingStore().addSettings(app.ui.settings)
   app.vueAppReady = true
@@ -61,6 +65,29 @@ const init = () => {
     component: markRaw(QueueSideBarTab),
     type: 'vue'
   })
+  app.extensionManager.registerSidebarTab({
+    id: 'node-library',
+    icon: 'pi pi-book',
+    title: t('sideToolBar.nodeLibrary'),
+    tooltip: t('sideToolBar.nodeLibrary'),
+    component: markRaw(NodeLibrarySideBarTab),
+    type: 'vue'
+  })
+
+  dropTargetCleanup = dropTargetForElements({
+    element: document.querySelector('.graph-canvas-container'),
+    onDrop: (event) => {
+      const loc = event.location.current.input
+      // Add an offset on x to make sure after adding the node, the cursor
+      // is on the node (top left corner)
+      const pos = app.clientPosToCanvasPos([loc.clientX - 20, loc.clientY])
+      const comfyNodeName = event.source.element.getAttribute(
+        'data-comfy-node-name'
+      )
+      const nodeDef = useNodeDefStore().nodeDefsByName[comfyNodeName]
+      app.addNodeOnGraph(nodeDef, { pos })
+    }
+  })
 }
 
 onMounted(() => {
@@ -72,6 +99,10 @@ onMounted(() => {
     isLoading.value = false
   }
 })
+
+onUnmounted(() => {
+  dropTargetCleanup()
+})
 </script>
 
 <style scoped>

src/components/primevueOverride/TreePlus.vue

@@ -0,0 +1,106 @@
+<!-- Tree with all leaf nodes draggable -->
+<script>
+import Tree from 'primevue/tree'
+import { draggable } from '@atlaskit/pragmatic-drag-and-drop/element/adapter'
+import { h, onMounted, onBeforeUnmount, computed } from 'vue'
+
+export default {
+  name: 'TreePlus',
+  extends: Tree,
+  props: {
+    dragSelector: {
+      type: String,
+      default: '.p-tree-node'
+    },
+    // Explicitly declare all v-model props
+    expandedKeys: {
+      type: Object,
+      default: () => ({})
+    },
+    selectionKeys: {
+      type: Object,
+      default: () => ({})
+    }
+  },
+  emits: ['update:expandedKeys', 'update:selectionKeys'],
+  setup(props, context) {
+    // Create computed properties for each v-model prop
+    const computedExpandedKeys = computed({
+      get: () => props.expandedKeys,
+      set: (value) => context.emit('update:expandedKeys', value)
+    })
+
+    const computedSelectionKeys = computed({
+      get: () => props.selectionKeys,
+      set: (value) => context.emit('update:selectionKeys', value)
+    })
+
+    let observer = null
+
+    const makeDraggable = (element) => {
+      if (!element._draggableCleanup) {
+        element._draggableCleanup = draggable({
+          element
+        })
+      }
+    }
+
+    const observeTreeChanges = (treeElement) => {
+      observer = new MutationObserver((mutations) => {
+        mutations.forEach((mutation) => {
+          if (mutation.type === 'childList') {
+            mutation.addedNodes.forEach((node) => {
+              if (node.nodeType === Node.ELEMENT_NODE) {
+                node.querySelectorAll(props.dragSelector).forEach(makeDraggable)
+              }
+            })
+          }
+        })
+      })
+
+      observer.observe(treeElement, { childList: true, subtree: true })
+
+      // Make existing nodes draggable
+      treeElement.querySelectorAll(props.dragSelector).forEach(makeDraggable)
+    }
+
+    onMounted(() => {
+      const treeElement = document.querySelector('.p-tree')
+      if (treeElement) {
+        observeTreeChanges(treeElement)
+      }
+    })
+
+    onBeforeUnmount(() => {
+      if (observer) {
+        observer.disconnect()
+      }
+      // Clean up draggable instances if necessary
+      const treeElement = document.querySelector('.p-tree')
+      if (treeElement) {
+        treeElement.querySelectorAll(props.dragSelector).forEach((node) => {
+          if (node._draggableCleanup) {
+            node._draggableCleanup()
+          }
+        })
+      }
+    })
+
+    return () =>
+      h(
+        Tree,
+        {
+          ...context.attrs,
+          ...props,
+          expandedKeys: computedExpandedKeys.value,
+          selectionKeys: computedSelectionKeys.value,
+          'onUpdate:expandedKeys': (value) =>
+            (computedExpandedKeys.value = value),
+          'onUpdate:selectionKeys': (value) =>
+            (computedSelectionKeys.value = value)
+        },
+        context.slots
+      )
+  }
+}
+</script>

src/components/sidebar/SideToolBar.vue

@@ -7,6 +7,7 @@
         :icon="tab.icon"
         :tooltip="tab.tooltip"
         :selected="tab === selectedTab"
+        :class="tab.id + '-tab-button'"
         @click="onTabClick(tab)"
       />
       <div class="side-tool-bar-end">

src/components/sidebar/tabs/NodeLibrarySideBarTab.vue

@@ -0,0 +1,99 @@
+<template>
+  <TreePlus
+    class="node-lib-tree"
+    v-model:expandedKeys="expandedKeys"
+    selectionMode="single"
+    :value="renderedRoot.children"
+    :filter="true"
+    filterMode="lenient"
+    dragSelector=".p-tree-node-leaf"
+    :pt="{
+      nodeLabel: 'node-lib-tree-node-label',
+      nodeChildren: ({ props }) => ({
+        'data-comfy-node-name': props.node?.data?.name,
+        onMouseenter: (event: MouseEvent) => {
+          hoveredComfyNodeName = props.node?.data?.name
+
+          const hoverTarget = event.target as HTMLElement
+          const targetRect = hoverTarget.getBoundingClientRect()
+          nodePreviewStyle.top = `${targetRect.top - 40}px`
+          nodePreviewStyle.left = `${targetRect.right}px`
+        },
+        onMouseleave: () => {
+          hoveredComfyNodeName = null
+        }
+      })
+    }"
+  >
+    <template #folder="{ node }">
+      <span class="folder-label">{{ node.label }}</span>
+      <Badge
+        :value="node.totalNodes"
+        severity="secondary"
+        :style="{ marginLeft: '0.5rem' }"
+      ></Badge>
+    </template>
+    <template #node="{ node }">
+      <span class="node-label">{{ node.label }}</span>
+    </template>
+  </TreePlus>
+  <div
+    v-if="hoveredComfyNode"
+    class="node-lib-node-preview"
+    :style="nodePreviewStyle"
+  >
+    <NodePreview
+      :key="hoveredComfyNode.name"
+      :nodeDef="hoveredComfyNode"
+    ></NodePreview>
+  </div>
+</template>
+
+<script setup lang="ts">
+import Badge from 'primevue/badge'
+import { ComfyNodeDefImpl, useNodeDefStore } from '@/stores/nodeDefStore'
+import { computed, ref } from 'vue'
+import { TreeNode } from 'primevue/treenode'
+import TreePlus from '@/components/primevueOverride/TreePlus.vue'
+import NodePreview from '@/components/NodePreview.vue'
+
+const nodeDefStore = useNodeDefStore()
+const expandedKeys = ref({})
+const hoveredComfyNodeName = ref<string | null>(null)
+const hoveredComfyNode = computed<ComfyNodeDefImpl | null>(() => {
+  if (!hoveredComfyNodeName.value) {
+    return null
+  }
+  return nodeDefStore.nodeDefsByName[hoveredComfyNodeName.value] || null
+})
+const nodePreviewStyle = ref<Record<string, string>>({
+  position: 'absolute',
+  top: '0px',
+  left: '0px'
+})
+
+const root = computed(() => nodeDefStore.nodeTree)
+const renderedRoot = computed(() => {
+  return fillNodeInfo(root.value)
+})
+const fillNodeInfo = (node: TreeNode): TreeNode => {
+  const isLeaf = node.children === undefined || node.children.length === 0
+  const isExpanded = expandedKeys.value[node.key]
+  const icon = isLeaf
+    ? 'pi pi-circle-fill'
+    : isExpanded
+      ? 'pi pi-folder-open'
+      : 'pi pi-folder'
+  const children = node.children?.map(fillNodeInfo)
+
+  return {
+    ...node,
+    icon,
+    children,
+    type: isLeaf ? 'node' : 'folder',
+    totalNodes: isLeaf
+      ? 1
+      : children.reduce((acc, child) => acc + child.totalNodes, 0)
+  }
+}
+</script>

src/i18n.ts

@@ -5,14 +5,16 @@ const messages = {
     sideToolBar: {
       settings: 'Settings',
       themeToggle: 'Toggle Theme',
-      queue: 'Queue'
+      queue: 'Queue',
+      nodeLibrary: 'Node Library'
     }
   },
   zh: {
     sideToolBar: {
       settings: '设置',
       themeToggle: '主题切换',
-      queue: '队列'
+      queue: '队列',
+      nodeLibrary: '节点库'
     }
   }
   // TODO: Add more languages

src/scripts/app.ts

@@ -42,6 +42,7 @@ import {
   useNodeDefStore
 } from '@/stores/nodeDefStore'
 import { Vector2 } from '@comfyorg/litegraph'
+import _ from 'lodash'
 
 export const ANIM_PREVIEW_WIDGET = '$$comfy_animation_preview'
 
@@ -2951,6 +2952,14 @@ export class ComfyApp {
     this.graph.add(node)
     return node
   }
+
+  clientPosToCanvasPos(pos: Vector2): Vector2 {
+    const rect = this.canvasContainer.getBoundingClientRect()
+    const containerOffsets = [rect.left, rect.top]
+    return _.zip(pos, this.canvas.ds.offset, containerOffsets).map(
+      ([p, o1, o2]) => p / this.canvas.ds.scale - o1 - o2
+    ) as Vector2
+  }
 }
 
 export const app = new ComfyApp()

src/stores/nodeDefStore.ts

@@ -3,6 +3,7 @@ import { ComfyNodeDef } from '@/types/apiTypes'
 import { defineStore } from 'pinia'
 import { Type, Transform, plainToClass } from 'class-transformer'
 import { ComfyWidgetConstructor } from '@/scripts/widgets'
+import { TreeNode } from 'primevue/treenode'
 
 export class BaseInputSpec<T = any> {
   name: string
@@ -232,6 +233,33 @@ export const useNodeDefStore = defineStore('nodeDef', {
     },
     nodeSearchService(state) {
       return new NodeSearchService(Object.values(state.nodeDefsByName))
+    },
+    nodeTree(state): TreeNode {
+      const root: TreeNode = {
+        key: 'root',
+        label: 'Nodes',
+        children: []
+      }
+      for (const nodeDef of Object.values(state.nodeDefsByName)) {
+        const path = nodeDef.category.split('/')
+        let current = root
+        let key = 'root'
+        for (const part of path) {
+          key += `/${part}`
+          let next = current.children.find((child) => child.label === part)
+          if (!next) {
+            next = { key, label: part, children: [] }
+            current.children.push(next)
+          }
+          current = next
+        }
+        current.children.push({
+          label: nodeDef.display_name,
+          data: nodeDef,
+          key: `${key}/${nodeDef.name}`
+        })
+      }
+      return root
     }
   },
   actions: {