Refactor node library drag and drop

src/components/graph/GraphCanvas.vue

@@ -20,7 +20,7 @@ import { ref, computed, onUnmounted, watch, onMounted, watchEffect } from 'vue'
 import { app as comfyApp } from '@/scripts/app'
 import { useSettingStore } from '@/stores/settingStore'
 import { dropTargetForElements } from '@atlaskit/pragmatic-drag-and-drop/element/adapter'
-import { useNodeDefStore } from '@/stores/nodeDefStore'
+import { ComfyNodeDefImpl, useNodeDefStore } from '@/stores/nodeDefStore'
 import { useWorkspaceStore } from '@/stores/workspaceStateStore'
 import {
  LiteGraph,
@@ -113,14 +113,18 @@ onMounted(async () => {
  element: canvasRef.value,
  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 = comfyApp.clientPosToCanvasPos([loc.clientX - 20, loc.clientY])
- const comfyNodeName = event.source.element.getAttribute(
- 'data-comfy-node-name'
- )
- const nodeDef = nodeDefStore.nodeDefsByName[comfyNodeName]
- comfyApp.addNodeOnGraph(nodeDef, { pos })
+ const dndData = event.source.data
+
+ if (dndData.type === 'add-node') {
+ const nodeDef = dndData.data as ComfyNodeDefImpl
+ // Add an offset on x to make sure after adding the node, the cursor
+ // is on the node (top left corner)
+ const pos = comfyApp.clientPosToCanvasPos([
+ loc.clientX - 20,
+ loc.clientY
+ ])
+ comfyApp.addNodeOnGraph(nodeDef, { pos })
+ }
  }
  })
  emit('ready')

src/components/sidebar/tabs/NodeLibrarySidebarTab.vue

@@ -20,7 +20,7 @@
  @search="handleSearch"
  :placeholder="$t('searchNodes') + '...'"
  />
- <TreePlus
+ <Tree
  class="node-lib-tree"
  v-model:expandedKeys="expandedKeys"
  selectionMode="single"
@@ -62,7 +62,7 @@
  @toggleBookmark="toggleBookmark(node.data.display_name)"
  />
  </template>
- </TreePlus>
+ </Tree>
  <div
  v-if="hoveredComfyNode"
  class="node-lib-node-preview"
@@ -89,7 +89,7 @@ import {
 import { computed, ref, nextTick } from 'vue'
 import type { TreeNode } from 'primevue/treenode'
 import NodeTreeLeaf from './nodeLibrary/NodeTreeLeaf.vue'
-import TreePlus from '@/components/primevueOverride/TreePlus.vue'
+import Tree from 'primevue/tree'
 import NodePreview from '@/components/node/NodePreview.vue'
 import SearchBox from '@/components/common/SearchBox.vue'
 import SidebarTabTemplate from '@/components/sidebar/tabs/SidebarTabTemplate.vue'

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

@@ -1,5 +1,5 @@
 <template>
- <div class="node-tree-leaf">
+ <div class="node-tree-leaf" ref="container">
  <div class="node-content">
  <Tag
  v-if="node.experimental"
@@ -24,6 +24,9 @@
 import Button from 'primevue/button'
 import Tag from 'primevue/tag'
 import { ComfyNodeDefImpl } from '@/stores/nodeDefStore'
+import { onMounted, onUnmounted, ref } from 'vue'
+import { draggable } from '@atlaskit/pragmatic-drag-and-drop/element/adapter'
+import { CanvasDragAndDropData } from '@/types/litegraphTypes'
 
 const props = defineProps<{
  node: ComfyNodeDefImpl
@@ -37,6 +40,26 @@ const emit = defineEmits<{
 const toggleBookmark = () => {
  emit('toggle-bookmark', props.node)
 }
+
+const container = ref<HTMLElement | null>(null)
+let draggableCleanup: () => void
+onMounted(() => {
+ const treeNodeElement = container.value?.closest(
+ '.p-tree-node'
+ ) as HTMLElement
+ draggableCleanup = draggable({
+ element: treeNodeElement,
+ getInitialData() {
+ return {
+ type: 'add-node',
+ data: props.node
+ } as CanvasDragAndDropData<ComfyNodeDefImpl>
+ }
+ })
+})
+onUnmounted(() => {
+ draggableCleanup()
+})
 </script>
 
 <style scoped>

src/types/litegraphTypes.ts

@@ -73,3 +73,8 @@ export class ConnectingLinkImpl implements ConnectingLink {
  }
  }
 }
+
+export type CanvasDragAndDropData<T = any> = {
+ type: 'add-node'
+ data: T
+}