From a69858c87a74ad364b92832d974835753f1c9453 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Mon, 26 Aug 2024 11:04:08 -0400
Subject: [PATCH] Fix unbookmark node in node library (#639)

* Add unbookmark test

* Fix unbookmark
---
 browser_tests/menu.spec.ts      | 20 +++++++++++++++++++-
 src/stores/nodeBookmarkStore.ts |  4 ++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/browser_tests/menu.spec.ts b/browser_tests/menu.spec.ts
index 3048694f70..fa83f4684d 100644
--- a/browser_tests/menu.spec.ts
+++ b/browser_tests/menu.spec.ts
@@ -206,7 +206,7 @@ test.describe('Menu', () => {
       )
     })
 
-    test('Can unbookmark node', async ({ comfyPage }) => {
+    test('Can unbookmark node (Top level bookmark)', async ({ comfyPage }) => {
       await comfyPage.setSetting('Comfy.NodeLibrary.Bookmarks', [
         'KSampler (Advanced)'
       ])
@@ -219,6 +219,24 @@ test.describe('Menu', () => {
         []
       )
     })
+
+    test('Can unbookmark node (Library node bookmark)', async ({
+      comfyPage
+    }) => {
+      await comfyPage.setSetting('Comfy.NodeLibrary.Bookmarks', [
+        'KSampler (Advanced)'
+      ])
+      const tab = comfyPage.menu.nodeLibraryTab
+      await tab.getFolder('sampling').click()
+      await comfyPage.page
+        .locator(tab.nodeSelector('KSampler (Advanced)'))
+        .nth(1)
+        .locator('.bookmark-button')
+        .click()
+      expect(await comfyPage.getSetting('Comfy.NodeLibrary.Bookmarks')).toEqual(
+        []
+      )
+    })
   })
 
   test('Can change canvas zoom speed setting', async ({ comfyPage }) => {
diff --git a/src/stores/nodeBookmarkStore.ts b/src/stores/nodeBookmarkStore.ts
index 58d9a61452..a550d2209a 100644
--- a/src/stores/nodeBookmarkStore.ts
+++ b/src/stores/nodeBookmarkStore.ts
@@ -30,6 +30,10 @@ export const useNodeBookmarkStore = defineStore('nodeBookmark', () => {
   const toggleBookmark = (node: ComfyNodeDefImpl) => {
     if (isBookmarked(node)) {
       deleteBookmark(node.nodePath)
+      // Delete the bookmark at the top level if it exists
+      // This is used for clicking the bookmark button in the node library, i.e.
+      // the node is inside original/standard node library tree node
+      deleteBookmark(node.display_name)
     } else {
       addBookmark(node.display_name)
     }