From ebdd7b8e403212995a27aa0b5e7adf5c5134eddf Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Wed, 24 Jul 2024 12:00:23 -0400
Subject: [PATCH] Use store to manage nodeSearchService (#214)

---
 browser_tests/ComfyPage.ts              |  4 +++-
 src/App.vue                             | 15 +++++----------
 src/components/NodeSearchBox.vue        | 20 +++++++++-----------
 src/components/NodeSearchBoxPopover.vue | 17 +++++++----------
 src/components/NodeSearchFilter.vue     | 12 ++++--------
 src/main.ts                             |  2 --
 src/stores/nodeDefStore.ts              |  6 +++++-
 7 files changed, 33 insertions(+), 43 deletions(-)

diff --git a/browser_tests/ComfyPage.ts b/browser_tests/ComfyPage.ts
index 154165f8a..835e82914 100644
--- a/browser_tests/ComfyPage.ts
+++ b/browser_tests/ComfyPage.ts
@@ -110,7 +110,9 @@ export class ComfyPage {
 			}`,
     });
     await this.page.waitForFunction(() => document.fonts.ready);
-    await this.page.waitForFunction(() => window["app"] !== undefined);
+    await this.page.waitForFunction(
+      () => window["app"] !== undefined && window["app"].vueAppReady
+    );
     await this.page.evaluate(() => {
       window["app"]["canvas"].show_info = false;
     });
diff --git a/src/App.vue b/src/App.vue
index f39fb3770..55f881180 100644
--- a/src/App.vue
+++ b/src/App.vue
@@ -13,20 +13,16 @@
 </template>
 
 <script setup lang="ts">
-import { computed, onMounted, onUnmounted, provide, ref, watch } from "vue";
+import { computed, onMounted, ref, watch } from "vue";
 import NodeSearchboxPopover from "@/components/NodeSearchBoxPopover.vue";
 import SideToolBar from "@/components/sidebar/SideToolBar.vue";
 import LiteGraphCanvasSplitterOverlay from "@/components/LiteGraphCanvasSplitterOverlay.vue";
 import ProgressSpinner from "primevue/progressspinner";
-import { NodeSearchService } from "./services/nodeSearchService";
 import { app } from "./scripts/app";
 import { useSettingStore } from "./stores/settingStore";
 import { useNodeDefStore } from "./stores/nodeDefStore";
 
 const isLoading = ref(true);
-const nodeSearchService = computed(
-  () => new NodeSearchService(useNodeDefStore().nodeDefs)
-);
 const nodeSearchEnabled = computed<boolean>(
   () => useSettingStore().get("Comfy.NodeSearchBoxImpl") === "default"
 );
@@ -47,22 +43,21 @@ watch(
   { immediate: true }
 );
 
-const init = async () => {
+const init = () => {
   useNodeDefStore().addNodeDefs(Object.values(app.nodeDefs));
   useSettingStore().addSettings(app.ui.settings);
+  app.vueAppReady = true;
 };
 
-onMounted(async () => {
+onMounted(() => {
   try {
-    await init();
+    init();
   } catch (e) {
     console.error("Failed to init Vue app", e);
   } finally {
     isLoading.value = false;
   }
 });
-
-provide("nodeSearchService", nodeSearchService);
 </script>
 
 <style scoped>
diff --git a/src/components/NodeSearchBox.vue b/src/components/NodeSearchBox.vue
index e8d5b9f8f..abd4bd9ef 100644
--- a/src/components/NodeSearchBox.vue
+++ b/src/components/NodeSearchBox.vue
@@ -60,11 +60,9 @@ import Badge from "primevue/badge";
 import NodeSearchFilter from "@/components/NodeSearchFilter.vue";
 import NodeSourceChip from "@/components/NodeSourceChip.vue";
 import { ComfyNodeDef } from "@/types/apiTypes";
-import {
-  NodeSearchService,
-  type FilterAndValue,
-} from "@/services/nodeSearchService";
+import { type FilterAndValue } from "@/services/nodeSearchService";
 import NodePreview from "./NodePreview.vue";
+import { useNodeDefStore } from "@/stores/nodeDefStore";
 
 const props = defineProps({
   filters: {
@@ -76,10 +74,6 @@ const props = defineProps({
   },
 });
 
-const nodeSearchService = (
-  inject("nodeSearchService") as Ref<NodeSearchService>
-).value;
-
 const inputId = `comfy-vue-node-search-box-input-${Math.random()}`;
 const suggestions = ref<ComfyNodeDef[]>([]);
 const hoveredSuggestion = ref<ComfyNodeDef | null>(null);
@@ -88,9 +82,13 @@ const placeholder = computed(() => {
 });
 
 const search = (query: string) => {
-  suggestions.value = nodeSearchService.searchNode(query, props.filters, {
-    limit: props.searchLimit,
-  });
+  suggestions.value = useNodeDefStore().nodeSearchService.searchNode(
+    query,
+    props.filters,
+    {
+      limit: props.searchLimit,
+    }
+  );
 };
 
 const emit = defineEmits(["addFilter", "removeFilter", "addNode"]);
diff --git a/src/components/NodeSearchBoxPopover.vue b/src/components/NodeSearchBoxPopover.vue
index e3cf8c9a1..532e12143 100644
--- a/src/components/NodeSearchBoxPopover.vue
+++ b/src/components/NodeSearchBoxPopover.vue
@@ -32,11 +32,9 @@ import {
   LGraphNode,
   LinkReleaseContext,
 } from "@comfyorg/litegraph";
-import {
-  FilterAndValue,
-  NodeSearchService,
-} from "@/services/nodeSearchService";
+import { FilterAndValue } from "@/services/nodeSearchService";
 import { ComfyNodeDef } from "@/types/apiTypes";
+import { useNodeDefStore } from "@/stores/nodeDefStore";
 
 interface LiteGraphPointerEvent extends Event {
   canvasX: number;
@@ -114,9 +112,6 @@ const addNode = (nodeDef: ComfyNodeDef) => {
     }
   }
 };
-const nodeSearchService = (
-  inject("nodeSearchService") as Ref<NodeSearchService>
-).value;
 
 const canvasEventHandler = (e: LiteGraphCanvasEvent) => {
   const shiftPressed = (e.detail.originalEvent as KeyboardEvent).shiftKey;
@@ -128,9 +123,10 @@ const canvasEventHandler = (e: LiteGraphCanvasEvent) => {
 
   if (e.detail.subType === "empty-release") {
     const destIsInput = e.detail.linkReleaseContext.node_from !== undefined;
-    const filter = destIsInput
-      ? nodeSearchService.getFilterById("input")
-      : nodeSearchService.getFilterById("output");
+    const filter = useNodeDefStore().nodeSearchService.getFilterById(
+      destIsInput ? "input" : "output"
+    );
+
     const value = destIsInput
       ? e.detail.linkReleaseContext.type_filter_in
       : e.detail.linkReleaseContext.type_filter_out;
@@ -172,6 +168,7 @@ onUnmounted(() => {
   background-color: transparent !important;
   margin-top: 25vh;
 }
+
 .node-search-box-dialog-mask {
   align-items: flex-start !important;
 }
diff --git a/src/components/NodeSearchFilter.vue b/src/components/NodeSearchFilter.vue
index 42b2e05e4..1c45a1a3e 100644
--- a/src/components/NodeSearchFilter.vue
+++ b/src/components/NodeSearchFilter.vue
@@ -34,26 +34,22 @@
 </template>
 
 <script setup lang="ts">
-import {
-  NodeFilter,
-  NodeSearchService,
-  type FilterAndValue,
-} from "@/services/nodeSearchService";
+import { NodeFilter, type FilterAndValue } from "@/services/nodeSearchService";
 import Button from "primevue/button";
 import Dialog from "primevue/dialog";
 import SelectButton from "primevue/selectbutton";
 import AutoComplete from "primevue/autocomplete";
-import { inject, ref, onMounted } from "vue";
+import { ref, onMounted } from "vue";
+import { useNodeDefStore } from "@/stores/nodeDefStore";
 
 const visible = ref<boolean>(false);
-const nodeSearchService: NodeSearchService = inject("nodeSearchService").value;
-
 const filters = ref<NodeFilter[]>([]);
 const selectedFilter = ref<NodeFilter>();
 const filterValues = ref<string[]>([]);
 const selectedFilterValue = ref<string>("");
 
 onMounted(() => {
+  const nodeSearchService = useNodeDefStore().nodeSearchService;
   filters.value = nodeSearchService.nodeFilters;
   selectedFilter.value = nodeSearchService.nodeFilters[0];
 });
diff --git a/src/main.ts b/src/main.ts
index f63b00219..0cde8caa8 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -45,6 +45,4 @@ comfyApp.setup().then(() => {
     .use(pinia)
     .use(i18n)
     .mount("#vue-app");
-
-  comfyApp.vueAppReady = true;
 });
diff --git a/src/stores/nodeDefStore.ts b/src/stores/nodeDefStore.ts
index e5539678d..8d2724c13 100644
--- a/src/stores/nodeDefStore.ts
+++ b/src/stores/nodeDefStore.ts
@@ -1,3 +1,4 @@
+import { NodeSearchService } from "@/services/nodeSearchService";
 import { ComfyNodeDef } from "@/types/apiTypes";
 import { defineStore } from "pinia";
 
@@ -54,6 +55,9 @@ export const useNodeDefStore = defineStore("nodeDef", {
     nodeDefs(state) {
       return Object.values(state.nodeDefsByName);
     },
+    nodeSearchService(state) {
+      return new NodeSearchService(Object.values(state.nodeDefsByName));
+    },
   },
   actions: {
     addNodeDef(nodeDef: ComfyNodeDef) {
@@ -61,7 +65,7 @@ export const useNodeDefStore = defineStore("nodeDef", {
     },
     addNodeDefs(nodeDefs: ComfyNodeDef[]) {
       for (const nodeDef of nodeDefs) {
-        this.addNodeDef(nodeDef);
+        this.nodeDefsByName[nodeDef.name] = nodeDef;
       }
     },
   },