Bookmark nodes in node library (#612)

* Basic bookmark

* Extract node leaf as component

* bigger hitbox

src/components/sidebar/tabs/NodeLibrarySidebarTab.vue

@@ -63,17 +63,11 @@
           ></Badge>
         </template>
         <template #node="{ node }">
-          <Tag
-            v-if="node.data.experimental"
-            :value="$t('experimental')"
-            severity="primary"
+          <NodeTreeLeaf
+            :node="node.data"
+            :isBookmarked="isBookmarked(node.data)"
+            @toggleBookmark="toggleBookmark(node.data.display_name)"
           />
-          <Tag
-            v-if="node.data.deprecated"
-            :value="$t('deprecated')"
-            severity="danger"
-          />
-          <span class="node-label">{{ node.label }}</span>
         </template>
       </TreePlus>
       <div
@@ -93,18 +87,23 @@
 
 <script setup lang="ts">
 import Badge from 'primevue/badge'
-import Tag from 'primevue/tag'
 import ToggleButton from 'primevue/togglebutton'
-import { ComfyNodeDefImpl, useNodeDefStore } from '@/stores/nodeDefStore'
+import {
+  buildNodeDefTree,
+  ComfyNodeDefImpl,
+  useNodeDefStore
+} from '@/stores/nodeDefStore'
 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 NodePreview from '@/components/node/NodePreview.vue'
 import SearchBox from '@/components/common/SearchBox.vue'
 import SidebarTabTemplate from '@/components/sidebar/tabs/SidebarTabTemplate.vue'
 import { useSettingStore } from '@/stores/settingStore'
 import { app } from '@/scripts/app'
-import { buildTree, sortedTree } from '@/utils/treeUtil'
+import { sortedTree } from '@/utils/treeUtil'
+import _ from 'lodash'
 
 const nodeDefStore = useNodeDefStore()
 const alphabeticalSort = ref(false)
@@ -130,13 +129,63 @@ const nodePreviewStyle = ref<Record<string, string>>({
   left: '0px'
 })
 
+// Bookmarks are in format of category/display_name. e.g. "comfy/conditioning/CLIPTextEncode"
+const bookmarks = computed(() =>
+  settingStore.get('Comfy.NodeLibrary.Bookmarks')
+)
+const bookmarkedNodes = computed(
+  () =>
+    new Set(
+      bookmarks.value.map((bookmark: string) => bookmark.split('/').pop())
+    )
+)
+const isBookmarked = (node: ComfyNodeDefImpl) =>
+  bookmarkedNodes.value.has(node.display_name)
+const toggleBookmark = (bookmark: string) => {
+  if (bookmarks.value.includes(bookmark)) {
+    settingStore.set(
+      'Comfy.NodeLibrary.Bookmarks',
+      bookmarks.value.filter((b: string) => b !== bookmark)
+    )
+  } else {
+    settingStore.set('Comfy.NodeLibrary.Bookmarks', [
+      ...bookmarks.value,
+      bookmark
+    ])
+  }
+}
+const bookmarkedRoot = computed<TreeNode>(() => {
+  const bookmarkNodes = bookmarks.value.map((bookmark: string) => {
+    const parts = bookmark.split('/')
+    const nodeName = parts.pop()
+    const category = parts.join('/')
+    const nodeDef = _.clone(nodeDefStore.nodeDefsByDisplayName[nodeName])
+    nodeDef.category = category
+    return nodeDef
+  })
+  return buildNodeDefTree(bookmarkNodes)
+})
+
+const allNodesRoot = computed<TreeNode>(() => {
+  return {
+    key: 'all-nodes',
+    label: 'All Nodes',
+    children: [
+      ...(bookmarkedRoot.value?.children ?? []),
+      ...nodeDefStore.nodeTree.children
+    ]
+  }
+})
+
 const root = computed(() => {
-  const root = filteredRoot.value || nodeDefStore.nodeTree
+  const root = filteredRoot.value || allNodesRoot.value
   return alphabeticalSort.value ? sortedTree(root) : root
 })
+
 const renderedRoot = computed(() => {
   return fillNodeInfo(root.value)
 })
+
 const fillNodeInfo = (node: TreeNode): TreeNode => {
   const isExpanded = expandedKeys.value[node.key]
   const icon = node.leaf
@@ -208,10 +257,7 @@ const handleSearch = (query: string) => {
     limit: 64
   })
 
-  filteredRoot.value = buildTree(matchedNodes, (nodeDef: ComfyNodeDefImpl) => [
-    ...nodeDef.category.split('/'),
-    nodeDef.display_name
-  ])
+  filteredRoot.value = buildNodeDefTree(matchedNodes)
   expandNode(filteredRoot.value)
 }
 
@@ -231,6 +277,7 @@ const expandNode = (node: TreeNode) => {
   display: flex;
   align-items: center;
   margin-left: var(--p-tree-node-gap);
+  flex-grow: 1;
 }
 </style>
 

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

@@ -0,0 +1,64 @@
+<template>
+  <div class="node-tree-leaf">
+    <div class="node-content">
+      <Tag
+        v-if="node.experimental"
+        :value="$t('experimental')"
+        severity="primary"
+      />
+      <Tag v-if="node.deprecated" :value="$t('deprecated')" severity="danger" />
+      <span class="node-label">{{ node.display_name }}</span>
+    </div>
+    <Button
+      class="bookmark-button"
+      size="small"
+      :icon="isBookmarked ? 'pi pi-bookmark-fill' : 'pi pi-bookmark'"
+      text
+      severity="secondary"
+      @click.stop="toggleBookmark"
+    />
+  </div>
+</template>
+
+<script setup lang="ts">
+import Button from 'primevue/button'
+import Tag from 'primevue/tag'
+import { ComfyNodeDefImpl } from '@/stores/nodeDefStore'
+
+const props = defineProps<{
+  node: ComfyNodeDefImpl
+  isBookmarked: boolean
+}>()
+
+const emit = defineEmits<{
+  (e: 'toggle-bookmark', value: ComfyNodeDefImpl): void
+}>()
+
+const toggleBookmark = () => {
+  emit('toggle-bookmark', props.node)
+}
+</script>
+
+<style scoped>
+.node-tree-leaf {
+  width: 100%;
+  display: flex;
+  align-items: center;
+  justify-content: space-between;
+}
+
+.node-content {
+  display: flex;
+  align-items: center;
+  flex-grow: 1;
+}
+
+.node-label {
+  margin-left: 0.5rem;
+}
+
+.bookmark-button {
+  width: unset;
+  padding: 0.25rem;
+}
+</style>

src/stores/nodeDefStore.ts

@@ -243,8 +243,16 @@ export const SYSTEM_NODE_DEFS: Record<string, ComfyNodeDef> = {
   }
 }
 
+export function buildNodeDefTree(nodeDefs: ComfyNodeDefImpl[]): TreeNode {
+  return buildTree(nodeDefs, (nodeDef: ComfyNodeDefImpl) => [
+    ...nodeDef.category.split('/').filter((s) => s !== ''),
+    nodeDef.display_name
+  ])
+}
+
 interface State {
   nodeDefsByName: Record<string, ComfyNodeDefImpl>
+  nodeDefsByDisplayName: Record<string, ComfyNodeDefImpl>
   widgets: Record<string, ComfyWidgetConstructor>
   showDeprecated: boolean
   showExperimental: boolean
@@ -253,6 +261,7 @@ interface State {
 export const useNodeDefStore = defineStore('nodeDef', {
   state: (): State => ({
     nodeDefsByName: {},
+    nodeDefsByDisplayName: {},
     widgets: {},
     showDeprecated: false,
     showExperimental: false
@@ -262,7 +271,7 @@ export const useNodeDefStore = defineStore('nodeDef', {
       return Object.values(state.nodeDefsByName)
     },
     // Node defs that are not deprecated
-    visibleNodeDefs(state) {
+    visibleNodeDefs(state): ComfyNodeDefImpl[] {
       return this.nodeDefs.filter(
         (nodeDef: ComfyNodeDefImpl) =>
           (state.showDeprecated || !nodeDef.deprecated) &&
@@ -273,22 +282,20 @@ export const useNodeDefStore = defineStore('nodeDef', {
       return new NodeSearchService(this.visibleNodeDefs)
     },
     nodeTree(): TreeNode {
-      return buildTree(this.visibleNodeDefs, (nodeDef: ComfyNodeDefImpl) => [
-        ...nodeDef.category.split('/'),
-        nodeDef.display_name
-      ])
+      return buildNodeDefTree(this.visibleNodeDefs)
     }
   },
   actions: {
     updateNodeDefs(nodeDefs: ComfyNodeDef[]) {
       const newNodeDefsByName: { [key: string]: ComfyNodeDefImpl } = {}
+      const nodeDefsByDisplayName: { [key: string]: ComfyNodeDefImpl } = {}
       for (const nodeDef of nodeDefs) {
-        newNodeDefsByName[nodeDef.name] = plainToClass(
-          ComfyNodeDefImpl,
-          nodeDef
-        )
+        const nodeDefImpl = plainToClass(ComfyNodeDefImpl, nodeDef)
+        newNodeDefsByName[nodeDef.name] = nodeDefImpl
+        nodeDefsByDisplayName[nodeDef.display_name] = nodeDefImpl
       }
       this.nodeDefsByName = newNodeDefsByName
+      this.nodeDefsByDisplayName = nodeDefsByDisplayName
     },
     updateWidgets(widgets: Record<string, ComfyWidgetConstructor>) {
       this.widgets = widgets

src/stores/settingStore.ts

@@ -191,6 +191,15 @@ export const useSettingStore = defineStore('setting', {
           step: 0.01
         }
       })
+
+      // Bookmarks are stored in the settings store.
+      // Bookmarks are in format of category/display_name. e.g. "conditioning/CLIPTextEncode"
+      app.ui.settings.addSetting({
+        id: 'Comfy.NodeLibrary.Bookmarks',
+        name: 'Node library bookmarks',
+        type: 'hidden',
+        defaultValue: []
+      })
     },
 
     set<K extends keyof Settings>(key: K, value: Settings[K]) {

src/types/apiTypes.ts

@@ -77,7 +77,7 @@ const zDownloadModelStatus = z.object({
   status: z.string(),
   progress_percentage: z.number(),
   message: z.string(),
-  already_existed: z.boolean(),
+  already_existed: z.boolean()
 })
 
 export type StatusWsMessageStatus = z.infer<typeof zStatusWsMessageStatus>
@@ -423,6 +423,7 @@ const zSettings = z.record(z.any()).and(
       'Comfy.Graph.ZoomSpeed': z.number(),
       'Comfy.InvertMenuScrolling': z.boolean(),
       'Comfy.Logging.Enabled': z.boolean(),
+      'Comfy.NodeLibrary.Bookmarks': z.array(z.string()),
       'Comfy.NodeInputConversionSubmenus': z.boolean(),
       'Comfy.NodeSearchBoxImpl.LinkReleaseTrigger': z.enum([
         'always',