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

Author: huchenlei

browser_tests/ComfyPage.ts

@@ -37,6 +37,42 @@ class ComfyNodeSearchBox {
   }
 }
 
+class NodeLibrarySideBarTab {
+  public readonly tabId: string = 'node-library'
+  constructor(public readonly page: Page) {}
+
+  get tabButton() {
+    return this.page.locator(`.${this.tabId}-tab-button`)
+  }
+
+  get selectedTabButton() {
+    return this.page.locator(
+      `.${this.tabId}-tab-button.side-bar-button-selected`
+    )
+  }
+
+  get nodeLibraryTree() {
+    return this.page.locator('.node-lib-tree')
+  }
+
+  get nodePreview() {
+    return this.page.locator('.node-lib-node-preview')
+  }
+
+  async open() {
+    if (await this.selectedTabButton.isVisible()) {
+      return
+    }
+
+    await this.tabButton.click()
+    await this.nodeLibraryTree.waitFor({ state: 'visible' })
+  }
+
+  async toggleFirstFolder() {
+    await this.page.locator('.p-tree-node-toggle-button').nth(0).click()
+  }
+}
+
 class ComfyMenu {
   public readonly sideToolBar: Locator
   public readonly themeToggleButton: Locator
@@ -46,6 +82,10 @@ class ComfyMenu {
     this.themeToggleButton = page.locator('.comfy-vue-theme-toggle')
   }
 
+  get nodeLibraryTab() {
+    return new NodeLibrarySideBarTab(this.page)
+  }
+
   async toggleTheme() {
     await this.themeToggleButton.click()
     await this.page.evaluate(() => {
@@ -96,6 +136,12 @@ export class ComfyPage {
     this.menu = new ComfyMenu(page)
   }
 
+  async getGraphNodesCount(): Promise<number> {
+    return await this.page.evaluate(() => {
+      return window['app']?.graph?._nodes?.length || 0
+    })
+  }
+
   async setup() {
     await this.goto()
     // Unify font for consistent screenshots.

browser_tests/menu.spec.ts

@@ -67,4 +67,29 @@ test.describe('Menu', () => {
     )
     expect(newChildrenCount).toBe(initialChildrenCount + 1)
   })
+
+  test('Sidebar node preview and drag to canvas', async ({ comfyPage }) => {
+    // Open the sidebar
+    const tab = comfyPage.menu.nodeLibraryTab
+    await tab.open()
+    await tab.toggleFirstFolder()
+
+    // Hover over a node to display the preview
+    const nodeSelector = '.p-tree-node-leaf'
+    await comfyPage.page.hover(nodeSelector)
+
+    // Verify the preview is displayed
+    const previewVisible = await comfyPage.page.isVisible(
+      '.node-lib-node-preview'
+    )
+    expect(previewVisible).toBe(true)
+
+    const count = await comfyPage.getGraphNodesCount()
+    // Drag the node onto the canvas
+    const canvasSelector = '#graph-canvas'
+    await comfyPage.page.dragAndDrop(nodeSelector, canvasSelector)
+
+    // Verify the node is added to the canvas
+    expect(await comfyPage.getGraphNodesCount()).toBe(count + 1)
+  })
 })

package-lock.json

@@ -46,6 +46,7 @@
     "zip-dir": "^2.0.0"
   },
   "dependencies": {
+    "@atlaskit/pragmatic-drag-and-drop": "^1.2.1",
     "@comfyorg/litegraph": "^0.7.29",
     "@primevue/themes": "^4.0.0-rc.2",
     "@vitejs/plugin-vue": "^5.0.5",

package.json

@@ -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/App.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/primevueOverride/TreePlus.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">