From fef780a72f8a3463b172686caae5b38d48f29b66 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Sat, 14 Sep 2024 11:27:38 +0900
Subject: [PATCH] Make useTreeExpansion hook accept expandedKeys as param
 (#826)

---
 src/components/sidebar/tabs/NodeLibrarySidebarTab.vue      | 3 ++-
 .../sidebar/tabs/nodeLibrary/NodeBookmarkTreeExplorer.vue  | 3 ++-
 src/hooks/treeHooks.ts                                     | 7 ++-----
 3 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue b/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
index e9d56d48f..d264e2b5a 100644
--- a/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
+++ b/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue
@@ -92,7 +92,8 @@ import { useNodeBookmarkStore } from '@/stores/nodeBookmarkStore'
 
 const nodeDefStore = useNodeDefStore()
 const nodeBookmarkStore = useNodeBookmarkStore()
-const { expandedKeys, expandNode, toggleNodeOnEvent } = useTreeExpansion()
+const expandedKeys = ref<Record<string, boolean>>({})
+const { expandNode, toggleNodeOnEvent } = useTreeExpansion(expandedKeys)
 
 const nodeBookmarkTreeExplorerRef = ref<InstanceType<
   typeof NodeBookmarkTreeExplorer
diff --git a/src/components/sidebar/tabs/nodeLibrary/NodeBookmarkTreeExplorer.vue b/src/components/sidebar/tabs/nodeLibrary/NodeBookmarkTreeExplorer.vue
index bbbf32fb8..608fa1aa9 100644
--- a/src/components/sidebar/tabs/nodeLibrary/NodeBookmarkTreeExplorer.vue
+++ b/src/components/sidebar/tabs/nodeLibrary/NodeBookmarkTreeExplorer.vue
@@ -47,7 +47,8 @@ const props = defineProps<{
   filteredNodeDefs: ComfyNodeDefImpl[]
 }>()
 
-const { expandedKeys, expandNode, toggleNodeOnEvent } = useTreeExpansion()
+const expandedKeys = ref<Record<string, boolean>>({})
+const { expandNode, toggleNodeOnEvent } = useTreeExpansion(expandedKeys)
 
 const handleNodeClick = (
   node: RenderedTreeExplorerNode<ComfyNodeDefImpl>,
diff --git a/src/hooks/treeHooks.ts b/src/hooks/treeHooks.ts
index 421e626e6..b9b235bf6 100644
--- a/src/hooks/treeHooks.ts
+++ b/src/hooks/treeHooks.ts
@@ -1,9 +1,7 @@
-import { ref } from 'vue'
+import { Ref } from 'vue'
 import type { TreeNode } from 'primevue/treenode'
 
-export function useTreeExpansion() {
-  const expandedKeys = ref<Record<string, boolean>>({})
-
+export function useTreeExpansion(expandedKeys: Ref<Record<string, boolean>>) {
   const toggleNode = (node: TreeNode) => {
     if (node.key && typeof node.key === 'string') {
       if (node.key in expandedKeys.value) {
@@ -63,7 +61,6 @@ export function useTreeExpansion() {
   }
 
   return {
-    expandedKeys,
     toggleNode,
     toggleNodeRecursive,
     expandNode,