Bookmark folder icon customization (#647)

* Add bookmark customization support

* WIP

* Fix bugs

* Fix color update

* Handle rename and delete of customization

* nit

* Add custom color picker

* Computed final color

* i18n

* Remove cancel button as dialog already has it

* Add playwright test

browser_tests/menu.spec.ts

@@ -62,6 +62,7 @@ test.describe('Menu', () => {
   test.describe('Node library sidebar', () => {
     test.beforeEach(async ({ comfyPage }) => {
       await comfyPage.setSetting('Comfy.NodeLibrary.Bookmarks', [])
+      await comfyPage.setSetting('Comfy.NodeLibrary.BookmarksCustomization', {})
       // Open the sidebar
       const tab = comfyPage.menu.nodeLibraryTab
       await tab.open()
@@ -237,6 +238,93 @@ test.describe('Menu', () => {
         []
       )
     })
+    test('Can customize icon', async ({ comfyPage }) => {
+      await comfyPage.setSetting('Comfy.NodeLibrary.Bookmarks', ['foo/'])
+      const tab = comfyPage.menu.nodeLibraryTab
+      await tab.getFolder('foo').click({ button: 'right' })
+      await comfyPage.page.getByLabel('Customize').click()
+      await comfyPage.page
+        .locator('.icon-field .p-selectbutton > *:nth-child(2)')
+        .click()
+      await comfyPage.page
+        .locator('.color-field .p-selectbutton > *:nth-child(2)')
+        .click()
+      await comfyPage.page.getByLabel('Confirm').click()
+      await comfyPage.nextFrame()
+      expect(
+        await comfyPage.getSetting('Comfy.NodeLibrary.BookmarksCustomization')
+      ).toEqual({
+        'foo/': {
+          icon: 'pi-folder',
+          color: '#007bff'
+        }
+      })
+    })
+    // If color is left as default, it should not be saved
+    test('Can customize icon (default field)', async ({ comfyPage }) => {
+      await comfyPage.setSetting('Comfy.NodeLibrary.Bookmarks', ['foo/'])
+      const tab = comfyPage.menu.nodeLibraryTab
+      await tab.getFolder('foo').click({ button: 'right' })
+      await comfyPage.page.getByLabel('Customize').click()
+      await comfyPage.page
+        .locator('.icon-field .p-selectbutton > *:nth-child(2)')
+        .click()
+      await comfyPage.page.getByLabel('Confirm').click()
+      await comfyPage.nextFrame()
+      expect(
+        await comfyPage.getSetting('Comfy.NodeLibrary.BookmarksCustomization')
+      ).toEqual({
+        'foo/': {
+          icon: 'pi-folder'
+        }
+      })
+    })
+    test('Can rename customized bookmark folder', async ({ comfyPage }) => {
+      await comfyPage.setSetting('Comfy.NodeLibrary.Bookmarks', ['foo/'])
+      await comfyPage.setSetting('Comfy.NodeLibrary.BookmarksCustomization', {
+        'foo/': {
+          icon: 'pi-folder',
+          color: '#007bff'
+        }
+      })
+      const tab = comfyPage.menu.nodeLibraryTab
+      await tab.getFolder('foo').click({ button: 'right' })
+      await comfyPage.page.getByLabel('Rename').click()
+      await comfyPage.page.keyboard.insertText('bar')
+      await comfyPage.page.keyboard.press('Enter')
+      await comfyPage.nextFrame()
+      expect(await comfyPage.getSetting('Comfy.NodeLibrary.Bookmarks')).toEqual(
+        ['bar/']
+      )
+      expect(
+        await comfyPage.getSetting('Comfy.NodeLibrary.BookmarksCustomization')
+      ).toEqual({
+        'bar/': {
+          icon: 'pi-folder',
+          color: '#007bff'
+        }
+      })
+    })
+
+    test('Can delete customized bookmark folder', async ({ comfyPage }) => {
+      await comfyPage.setSetting('Comfy.NodeLibrary.Bookmarks', ['foo/'])
+      await comfyPage.setSetting('Comfy.NodeLibrary.BookmarksCustomization', {
+        'foo/': {
+          icon: 'pi-folder',
+          color: '#007bff'
+        }
+      })
+      const tab = comfyPage.menu.nodeLibraryTab
+      await tab.getFolder('foo').click({ button: 'right' })
+      await comfyPage.page.getByLabel('Delete').click()
+      await comfyPage.nextFrame()
+      expect(await comfyPage.getSetting('Comfy.NodeLibrary.Bookmarks')).toEqual(
+        []
+      )
+      expect(
+        await comfyPage.getSetting('Comfy.NodeLibrary.BookmarksCustomization')
+      ).toEqual({})
+    })
   })
 
   test('Can change canvas zoom speed setting', async ({ comfyPage }) => {

src/components/common/CustomizationDialog.vue

@@ -0,0 +1,197 @@
+<template>
+  <Dialog v-model:visible="visible" :header="$t('customizeFolder')">
+    <div class="p-fluid">
+      <div class="field icon-field">
+        <label for="icon">{{ $t('icon') }}</label>
+        <SelectButton
+          v-model="selectedIcon"
+          :options="iconOptions"
+          optionLabel="name"
+          dataKey="value"
+        >
+          <template #option="slotProps">
+            <i
+              :class="['pi', slotProps.option.value, 'mr-2']"
+              :style="{ color: finalColor }"
+            ></i>
+          </template>
+        </SelectButton>
+      </div>
+      <Divider />
+      <div class="field color-field">
+        <label for="color">{{ $t('color') }}</label>
+        <div class="color-picker-container">
+          <SelectButton
+            v-model="selectedColor"
+            :options="colorOptions"
+            optionLabel="name"
+            dataKey="value"
+          >
+            <template #option="slotProps">
+              <div
+                v-if="slotProps.option.value !== 'custom'"
+                :style="{
+                  width: '20px',
+                  height: '20px',
+                  backgroundColor: slotProps.option.value,
+                  borderRadius: '50%'
+                }"
+              ></div>
+              <i
+                v-else
+                class="pi pi-palette"
+                :style="{ fontSize: '1.2rem' }"
+                v-tooltip="$t('customColor')"
+              ></i>
+            </template>
+          </SelectButton>
+          <ColorPicker
+            v-if="selectedColor.value === 'custom'"
+            v-model="customColor"
+          />
+        </div>
+      </div>
+    </div>
+    <template #footer>
+      <Button
+        :label="$t('reset')"
+        icon="pi pi-refresh"
+        @click="resetCustomization"
+        class="p-button-text"
+      />
+      <Button
+        :label="$t('confirm')"
+        icon="pi pi-check"
+        @click="confirmCustomization"
+        autofocus
+      />
+    </template>
+  </Dialog>
+</template>
+
+<script setup lang="ts">
+import { ref, computed, watch } from 'vue'
+import Dialog from 'primevue/dialog'
+import SelectButton from 'primevue/selectbutton'
+import Button from 'primevue/button'
+import Divider from 'primevue/divider'
+import ColorPicker from 'primevue/colorpicker'
+import { useNodeBookmarkStore } from '@/stores/nodeBookmarkStore'
+import { useI18n } from 'vue-i18n'
+
+const { t } = useI18n()
+
+const props = defineProps<{
+  modelValue: boolean
+  initialIcon?: string
+  initialColor?: string
+}>()
+
+const emit = defineEmits<{
+  (e: 'update:modelValue', value: boolean): void
+  (e: 'confirm', icon: string, color: string): void
+}>()
+
+const visible = computed({
+  get: () => props.modelValue,
+  set: (value) => emit('update:modelValue', value)
+})
+
+const nodeBookmarkStore = useNodeBookmarkStore()
+
+const iconOptions = [
+  { name: t('bookmark'), value: nodeBookmarkStore.defaultBookmarkIcon },
+  { name: t('folder'), value: 'pi-folder' },
+  { name: t('star'), value: 'pi-star' },
+  { name: t('heart'), value: 'pi-heart' },
+  { name: t('file'), value: 'pi-file' },
+  { name: t('inbox'), value: 'pi-inbox' },
+  { name: t('box'), value: 'pi-box' },
+  { name: t('briefcase'), value: 'pi-briefcase' }
+]
+
+const colorOptions = [
+  { name: t('default'), value: nodeBookmarkStore.defaultBookmarkColor },
+  { name: t('blue'), value: '#007bff' },
+  { name: t('green'), value: '#28a745' },
+  { name: t('red'), value: '#dc3545' },
+  { name: t('pink'), value: '#e83e8c' },
+  { name: t('yellow'), value: '#ffc107' },
+  { name: t('custom'), value: 'custom' }
+]
+
+const defaultIcon = iconOptions.find(
+  (option) => option.value === nodeBookmarkStore.defaultBookmarkIcon
+)
+const defaultColor = colorOptions.find(
+  (option) => option.value === nodeBookmarkStore.defaultBookmarkColor
+)
+
+const selectedIcon = ref<{ name: string; value: string }>(defaultIcon)
+const selectedColor = ref<{ name: string; value: string }>(defaultColor)
+const finalColor = computed(() =>
+  selectedColor.value.value === 'custom'
+    ? `#${customColor.value}`
+    : selectedColor.value.value
+)
+
+const customColor = ref('000000')
+
+const closeDialog = () => {
+  visible.value = false
+}
+
+const confirmCustomization = () => {
+  emit('confirm', selectedIcon.value.value, finalColor.value)
+  closeDialog()
+}
+
+const resetCustomization = () => {
+  selectedIcon.value =
+    iconOptions.find((option) => option.value === props.initialIcon) ||
+    defaultIcon
+  const colorOption = colorOptions.find(
+    (option) => option.value === props.initialColor
+  )
+  if (!props.initialColor) {
+    selectedColor.value = defaultColor
+  } else if (!colorOption) {
+    customColor.value = props.initialColor.replace('#', '')
+    selectedColor.value = { name: 'Custom', value: 'custom' }
+  } else {
+    selectedColor.value = colorOption
+  }
+}
+
+watch(
+  () => props.modelValue,
+  (newValue: boolean) => {
+    if (newValue) {
+      resetCustomization()
+    }
+  },
+  { immediate: true }
+)
+</script>
+
+<style scoped>
+.p-selectbutton .p-button {
+  padding: 0.5rem;
+}
+
+.p-selectbutton .p-button .pi {
+  font-size: 1.5rem;
+}
+
+.field {
+  display: flex;
+  flex-direction: column;
+  gap: 0.5rem;
+}
+
+.color-picker-container {
+  display: flex;
+  align-items: center;
+  gap: 0.5rem;
+}
+</style>

src/components/sidebar/tabs/NodeLibrarySidebarTab.vue

@@ -86,6 +86,12 @@
     </template>
   </SidebarTabTemplate>
   <ContextMenu ref="menu" :model="menuItems" />
+  <FolderCustomizationDialog
+    v-model="showCustomizationDialog"
+    @confirm="updateCustomization"
+    :initialIcon="initialIcon"
+    :initialColor="initialColor"
+  />
 </template>
 
 <script setup lang="ts">
@@ -105,6 +111,7 @@ import ContextMenu from 'primevue/contextmenu'
 import EditableText from '@/components/common/EditableText.vue'
 import NodePreview from '@/components/node/NodePreview.vue'
 import SearchBox from '@/components/common/SearchBox.vue'
+import FolderCustomizationDialog from '@/components/common/CustomizationDialog.vue'
 import SidebarTabTemplate from '@/components/sidebar/tabs/SidebarTabTemplate.vue'
 import { useSettingStore } from '@/stores/settingStore'
 import { app } from '@/scripts/app'
@@ -171,6 +178,11 @@ const getTreeNodeIcon = (node: TreeNode) => {
 
   // If the node is a bookmark folder, show a bookmark icon
   if (node.data && node.data.isDummyFolder) {
+    const customization =
+      nodeBookmarkStore.bookmarksCustomization[node.data.nodePath]
+    if (customization?.icon) {
+      return 'pi ' + customization.icon
+    }
     return 'pi pi-bookmark-fill'
   }
 
@@ -280,13 +292,22 @@ const menuItems = computed<MenuItem[]>(() => [
     command: () => {
       renameEditingNode.value = menuTargetNode.value
     }
+  },
+  {
+    label: t('customize'),
+    icon: 'pi pi-palette',
+    command: () => {
+      initialIcon.value =
+        nodeBookmarkStore.bookmarksCustomization[
+          menuTargetNode.value.data.nodePath
+        ]?.icon || nodeBookmarkStore.defaultBookmarkIcon
+      initialColor.value =
+        nodeBookmarkStore.bookmarksCustomization[
+          menuTargetNode.value.data.nodePath
+        ]?.color || nodeBookmarkStore.defaultBookmarkColor
+      showCustomizationDialog.value = true
+    }
   }
-  // TODO: Add customize color and icon feature.
-  // {
-  //   label: t('customize'),
-  //   icon: 'pi pi-palette',
-  //   command: () => console.log('customize')
-  // }
 ])
 
 const handleContextMenu = (node: TreeNode, e: MouseEvent) => {
@@ -320,6 +341,18 @@ const addNewBookmarkFolder = (parent?: ComfyNodeDefImpl) => {
     renameEditingNode.value = findNodeByKey(renderedRoot.value, newFolderKey)
   })
 }
+
+const showCustomizationDialog = ref(false)
+const initialIcon = ref(nodeBookmarkStore.defaultBookmarkIcon)
+const initialColor = ref(nodeBookmarkStore.defaultBookmarkColor)
+const updateCustomization = (icon: string, color: string) => {
+  if (menuTargetNode.value?.data) {
+    nodeBookmarkStore.updateBookmarkCustomization(
+      menuTargetNode.value.data.nodePath,
+      { icon, color }
+    )
+  }
+}
 </script>
 
 <style>

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

@@ -26,7 +26,8 @@ import type { CanvasDragAndDropData } from '@/types/litegraphTypes'
 import { dropTargetForElements } from '@atlaskit/pragmatic-drag-and-drop/element/adapter'
 import Badge from 'primevue/badge'
 import type { TreeNode } from 'primevue/treenode'
-import { onMounted, onUnmounted, ref } from 'vue'
+import { computed, onMounted, onUnmounted, ref, watch } from 'vue'
+import type { BookmarkCustomization } from '@/types/apiTypes'
 
 const props = defineProps<{
   node: TreeNode
@@ -39,6 +40,10 @@ const emit = defineEmits<{
 
 const nodeBookmarkStore = useNodeBookmarkStore()
 
+const customization = computed<BookmarkCustomization | undefined>(() => {
+  return nodeBookmarkStore.bookmarksCustomization[props.node.data.nodePath]
+})
+
 const addNodeToBookmarkFolder = (node: ComfyNodeDefImpl) => {
   if (!props.node.data) {
     console.error('Bookmark folder does not have data!')
@@ -56,15 +61,20 @@ const addNodeToBookmarkFolder = (node: ComfyNodeDefImpl) => {
 const container = ref<HTMLElement | null>(null)
 const canDrop = ref(false)
 
+const treeNodeElement = ref<HTMLElement | null>(null)
+const iconElement = ref<HTMLElement | null>(null)
+
 let dropTargetCleanup = () => {}
+let stopWatchCustomization: (() => void) | null = null
+
 onMounted(() => {
   if (!props.isBookmarkFolder) return
 
-  const treeNodeElement = container.value?.closest(
+  treeNodeElement.value = container.value?.closest(
     '.p-tree-node-content'
   ) as HTMLElement
   dropTargetCleanup = dropTargetForElements({
-    element: treeNodeElement,
+    element: treeNodeElement.value,
     onDrop: (event) => {
       const dndData = event.source.data as CanvasDragAndDropData
       if (dndData.type === 'add-node') {
@@ -83,8 +93,34 @@ onMounted(() => {
       canDrop.value = false
     }
   })
+
+  iconElement.value = treeNodeElement.value.querySelector(
+    ':scope > .p-tree-node-icon'
+  ) as HTMLElement
+
+  updateIconColor()
+
+  // Start watching after the component is mounted
+  stopWatchCustomization = watch(customization, updateIconColor, { deep: true })
 })
+
+const updateIconColor = () => {
+  if (iconElement.value && customization.value) {
+    iconElement.value.style.color = customization.value.color
+  }
+}
+
 onUnmounted(() => {
   dropTargetCleanup()
+  if (stopWatchCustomization) {
+    stopWatchCustomization()
+  }
 })
 </script>
+
+<style scoped>
+.node-tree-folder {
+  display: flex;
+  align-items: center;
+}
+</style>

src/i18n.ts

@@ -2,6 +2,17 @@ import { createI18n } from 'vue-i18n'
 
 const messages = {
   en: {
+    customizeFolder: 'Customize Folder',
+    icon: 'Icon',
+    color: 'Color',
+    bookmark: 'Bookmark',
+    folder: 'Folder',
+    star: 'Star',
+    heart: 'Heart',
+    file: 'File',
+    inbox: 'Inbox',
+    box: 'Box',
+    briefcase: 'Briefcase',
     error: 'Error',
     findIssues: 'Find Issues',
     copyToClipboard: 'Copy to Clipboard',
@@ -39,6 +50,17 @@ const messages = {
     }
   },
   zh: {
+    customizeFolder: '定制文件夹',
+    icon: '图标',
+    color: '颜色',
+    bookmark: '书签',
+    folder: '文件夹',
+    star: '星星',
+    heart: '心',
+    file: '文件',
+    inbox: '收件箱',
+    box: '盒子',
+    briefcase: '公文包',
     error: '错误',
     showReport: '显示报告',
     imageFailedToLoad: '图像加载失败',

src/stores/nodeBookmarkStore.ts

@@ -6,6 +6,7 @@ import { ComfyNodeDefImpl, createDummyFolderNodeDef } from './nodeDefStore'
 import { buildNodeDefTree } from './nodeDefStore'
 import type { TreeNode } from 'primevue/treenode'
 import _ from 'lodash'
+import type { BookmarkCustomization } from '@/types/apiTypes'
 
 export const useNodeBookmarkStore = defineStore('nodeBookmark', () => {
   const settingStore = useSettingStore()
@@ -113,6 +114,7 @@ export const useNodeBookmarkStore = defineStore('nodeBookmark', () => {
           : b
       )
     )
+    renameBookmarkCustomization(folderNode.nodePath, newNodePath)
   }
 
   const deleteBookmarkFolder = (folderNode: ComfyNodeDefImpl) => {
@@ -126,8 +128,64 @@ export const useNodeBookmarkStore = defineStore('nodeBookmark', () => {
           b !== folderNode.nodePath && !b.startsWith(folderNode.nodePath)
       )
     )
+    deleteBookmarkCustomization(folderNode.nodePath)
   }
 
+  const bookmarksCustomization = computed<
+    Record<string, BookmarkCustomization>
+  >(() => settingStore.get('Comfy.NodeLibrary.BookmarksCustomization'))
+
+  const updateBookmarkCustomization = (
+    nodePath: string,
+    customization: BookmarkCustomization
+  ) => {
+    const currentCustomization = bookmarksCustomization.value[nodePath] || {}
+    const newCustomization = { ...currentCustomization, ...customization }
+
+    // Remove attributes that are set to default values
+    if (newCustomization.icon === defaultBookmarkIcon) {
+      delete newCustomization.icon
+    }
+    if (newCustomization.color === defaultBookmarkColor) {
+      delete newCustomization.color
+    }
+
+    // If the customization is empty, remove it entirely
+    if (Object.keys(newCustomization).length === 0) {
+      deleteBookmarkCustomization(nodePath)
+    } else {
+      settingStore.set('Comfy.NodeLibrary.BookmarksCustomization', {
+        ...bookmarksCustomization.value,
+        [nodePath]: newCustomization
+      })
+    }
+  }
+
+  const deleteBookmarkCustomization = (nodePath: string) => {
+    settingStore.set('Comfy.NodeLibrary.BookmarksCustomization', {
+      ...bookmarksCustomization.value,
+      [nodePath]: undefined
+    })
+  }
+
+  const renameBookmarkCustomization = (
+    oldNodePath: string,
+    newNodePath: string
+  ) => {
+    const updatedCustomization = { ...bookmarksCustomization.value }
+    if (updatedCustomization[oldNodePath]) {
+      updatedCustomization[newNodePath] = updatedCustomization[oldNodePath]
+      delete updatedCustomization[oldNodePath]
+    }
+    settingStore.set(
+      'Comfy.NodeLibrary.BookmarksCustomization',
+      updatedCustomization
+    )
+  }
+
+  const defaultBookmarkIcon = 'pi-bookmark-fill'
+  const defaultBookmarkColor = '#a1a1aa'
+
   return {
     bookmarks,
     bookmarkedRoot,
@@ -136,6 +194,13 @@ export const useNodeBookmarkStore = defineStore('nodeBookmark', () => {
     addBookmark,
     addNewBookmarkFolder,
     renameBookmarkFolder,
-    deleteBookmarkFolder
+    deleteBookmarkFolder,
+
+    bookmarksCustomization,
+    updateBookmarkCustomization,
+    deleteBookmarkCustomization,
+    renameBookmarkCustomization,
+    defaultBookmarkIcon,
+    defaultBookmarkColor
   }
 })

src/stores/settingStore.ts

@@ -211,6 +211,14 @@ export const useSettingStore = defineStore('setting', {
         type: 'hidden',
         defaultValue: []
       })
+
+      // Stores mapping from bookmark folder name to its customization.
+      app.ui.settings.addSetting({
+        id: 'Comfy.NodeLibrary.BookmarksCustomization',
+        name: 'Node library bookmarks customization',
+        type: 'hidden',
+        defaultValue: {}
+      })
     },
 
     set<K extends keyof Settings>(key: K, value: Settings[K]) {

src/types/apiTypes.ts

@@ -408,6 +408,13 @@ const zUser = z.object({
   users: z.record(z.string(), z.unknown())
 })
 const zUserData = z.array(z.array(z.string(), z.string()))
+
+const zBookmarkCustomization = z.object({
+  icon: z.string().optional(),
+  color: z.string().optional()
+})
+export type BookmarkCustomization = z.infer<typeof zBookmarkCustomization>
+
 const zSettings = z.record(z.any()).and(
   z
     .object({
@@ -428,6 +435,10 @@ const zSettings = z.record(z.any()).and(
       'Comfy.InvertMenuScrolling': z.boolean(),
       'Comfy.Logging.Enabled': z.boolean(),
       'Comfy.NodeLibrary.Bookmarks': z.array(z.string()),
+      'Comfy.NodeLibrary.BookmarksCustomization': z.record(
+        z.string(),
+        zBookmarkCustomization
+      ),
       'Comfy.NodeInputConversionSubmenus': z.boolean(),
       'Comfy.NodeSearchBoxImpl.LinkReleaseTrigger': z.enum([
         'always',