[Major Refactor] Use pinia store to manage setting & nodeDef (#202)

* Node def store and settings tore

* Fix initial values

* Remove legacy setting listen

* Fix searchbox test

browser_tests/ComfyPage.ts

@@ -30,6 +30,9 @@ class ComfyNodeSearchBox {
     await this.input.waitFor({ state: "visible" });
     await this.input.fill(nodeName);
     await this.dropdown.waitFor({ state: "visible" });
+    // Wait for some time for the auto complete list to update.
+    // The auto complete list is debounced and may take some time to update.
+    await this.page.waitForTimeout(500);
     await this.dropdown.locator("li").nth(0).click();
   }
 }

src/App.vue

@@ -13,51 +13,43 @@
 </template>
 
 <script setup lang="ts">
-import { onMounted, onUnmounted, provide, ref } from "vue";
+import { computed, onMounted, onUnmounted, provide, 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,
-  SYSTEM_NODE_DEFS,
-} from "./services/nodeSearchService";
+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 nodeSearchEnabled = ref(false);
-const nodeSearchService = ref<NodeSearchService>();
-
-const updateTheme = (e) => {
-  const DARK_THEME_CLASS = "dark-theme";
-  const isDarkTheme = e.detail.value !== "light";
-
-  if (isDarkTheme) {
-    document.body.classList.add(DARK_THEME_CLASS);
-  } else {
-    document.body.classList.remove(DARK_THEME_CLASS);
-  }
-};
-
-const updateNodeSearchSetting = (e) => {
-  const settingValue = e.detail.value || "default";
-  nodeSearchEnabled.value = settingValue === "default";
-};
+const nodeSearchService = computed(
+  () => new NodeSearchService(useNodeDefStore().nodeDefs)
+);
+const nodeSearchEnabled = computed<boolean>(
+  () => useSettingStore().get("Comfy.NodeSearchBoxImpl") === "default"
+);
+const theme = computed<string>(() =>
+  useSettingStore().get("Comfy.ColorPalette")
+);
+watch(
+  theme,
+  (newTheme) => {
+    const DARK_THEME_CLASS = "dark-theme";
+    const isDarkTheme = newTheme !== "light";
+    if (isDarkTheme) {
+      document.body.classList.add(DARK_THEME_CLASS);
+    } else {
+      document.body.classList.remove(DARK_THEME_CLASS);
+    }
+  },
+  { immediate: true }
+);
 
 const init = async () => {
-  const nodeDefs = Object.values(app.nodeDefs);
-  nodeSearchService.value = new NodeSearchService([
-    ...nodeDefs,
-    ...SYSTEM_NODE_DEFS,
-  ]);
-
-  app.ui.settings.addEventListener("Comfy.ColorPalette.change", updateTheme);
-  app.ui.settings.addEventListener(
-    "Comfy.NodeSearchBoxImpl.change",
-    updateNodeSearchSetting
-  );
-  app.ui.settings.refreshSetting("Comfy.NodeSearchBoxImpl");
-  app.ui.settings.refreshSetting("Comfy.ColorPalette");
+  useNodeDefStore().addNodeDefs(Object.values(app.nodeDefs));
+  useSettingStore().addSettings(app.ui.settings);
 };
 
 onMounted(async () => {
@@ -70,14 +62,6 @@ onMounted(async () => {
   }
 });
 
-onUnmounted(() => {
-  app.ui.settings.removeEventListener("Comfy.ColorPalette.change", updateTheme);
-  app.ui.settings.removeEventListener(
-    "Comfy.NodeSearchBoxImpl.change",
-    updateNodeSearchSetting
-  );
-});
-
 provide("nodeSearchService", nodeSearchService);
 </script>
 

src/components/sidebar/SideBarThemeToggleIcon.vue

@@ -8,31 +8,23 @@
 </template>
 
 <script setup lang="ts">
-import { computed, onMounted, onUnmounted, ref, watch } from "vue";
+import { computed, ref } from "vue";
 import SideBarIcon from "./SideBarIcon.vue";
-import { app } from "@/scripts/app";
+import { useSettingStore } from "@/stores/settingStore";
 
-const isDarkMode = ref(false);
+const previousDarkTheme = ref("dark");
+const currentTheme = computed(() =>
+  useSettingStore().get("Comfy.ColorPalette", "dark")
+);
+const isDarkMode = computed(() => currentTheme.value !== "light");
 const icon = computed(() => (isDarkMode.value ? "pi pi-moon" : "pi pi-sun"));
-const themeId = computed(() => (isDarkMode.value ? "dark" : "light"));
+
 const toggleTheme = () => {
-  isDarkMode.value = !isDarkMode.value;
+  if (isDarkMode.value) {
+    previousDarkTheme.value = currentTheme.value;
+    useSettingStore().set("Comfy.ColorPalette", "light");
+  } else {
+    useSettingStore().set("Comfy.ColorPalette", previousDarkTheme.value);
+  }
 };
-
-watch(themeId, (newThemeId) => {
-  app.ui.settings.setSettingValue("Comfy.ColorPalette", newThemeId);
-});
-
-const updateTheme = (e) => {
-  isDarkMode.value = e.detail.value !== "light";
-};
-
-onMounted(() => {
-  app.ui.settings.addEventListener("Comfy.ColorPalette.change", updateTheme);
-  app.ui.settings.refreshSetting("Comfy.ColorPalette");
-});
-
-onUnmounted(() => {
-  app.ui.settings.removeEventListener("Comfy.ColorPalette.change", updateTheme);
-});
 </script>

src/components/sidebar/SideToolBar.vue

@@ -23,8 +23,9 @@ import SideBarThemeToggleIcon from "./SideBarThemeToggleIcon.vue";
 import SideBarSettingsToggleIcon from "./SideBarSettingsToggleIcon.vue";
 import NodeDetailSideBarItem from "./items/NodeDetailSideBarItem.vue";
 import QueueSideBarItem from "./items/QueueSideBarItem.vue";
-import { markRaw, onMounted, onUnmounted, ref, watch } from "vue";
+import { computed, markRaw, ref, watch } from "vue";
 import { useI18n } from "vue-i18n";
+import { useSettingStore } from "@/stores/settingStore";
 
 const { t } = useI18n();
 const items = ref([
@@ -49,22 +50,13 @@ watch(selectedItem, (newVal) => {
   emit("change", newVal !== null);
 });
 
-const onBetaMenuDisabled = () => {
-  selectedItem.value = null;
-};
-
-onMounted(() => {
-  document.addEventListener(
-    "comfy:setting:beta-menu-disabled",
-    onBetaMenuDisabled
-  );
-});
-
-onUnmounted(() => {
-  document.removeEventListener(
-    "comfy:setting:beta-menu-disabled",
-    onBetaMenuDisabled
-  );
+const betaMenuEnabled = computed(
+  () => useSettingStore().get("Comfy.UseNewMenu") !== "Disabled"
+);
+watch(betaMenuEnabled, (newValue) => {
+  if (!newValue) {
+    selectedItem.value = null;
+  }
 });
 </script>
 

src/main.ts

@@ -45,4 +45,6 @@ comfyApp.setup().then(() => {
     .use(pinia)
     .use(i18n)
     .mount("#vue-app");
+
+  comfyApp.vueAppReady = true;
 });

src/scripts/app.ts

@@ -86,6 +86,7 @@ export class ComfyApp {
     DraggableList,
   };
 
+  vueAppReady: boolean;
   ui: ComfyUI;
   logging: ComfyLogging;
   extensions: ComfyExtension[];
@@ -118,6 +119,7 @@ export class ComfyApp {
   nodeDefs: Record<string, ComfyNodeDef>;
 
   constructor() {
+    this.vueAppReady = false;
     this.ui = new ComfyUI(this);
     this.logging = new ComfyLogging(this);
     this.workflowManager = new ComfyWorkflowManager(this);
@@ -1877,6 +1879,7 @@ export class ComfyApp {
     this.#addAfterConfigureHandler();
 
     this.canvas = new LGraphCanvas(canvasEl, this.graph);
+    this.ui.settings.refreshSetting("Comfy.NodeSearchBoxImpl");
     this.ctx = canvasEl.getContext("2d");
 
     LiteGraph.release_link_on_empty_shows_menu = true;

src/scripts/ui/menu/index.ts

@@ -182,7 +182,6 @@ export class ComfyAppMenu {
           app.ui.menuContainer.style.removeProperty("display");
           this.element.style.display = "none";
           app.ui.restoreMenuPosition();
-          document.dispatchEvent(new Event("comfy:setting:beta-menu-disabled"));
         }
         window.dispatchEvent(new Event("resize"));
       },

src/scripts/ui/settings.ts

@@ -3,11 +3,13 @@ import { api } from "../api";
 import { ComfyDialog } from "./dialog";
 import type { ComfyApp } from "../app";
 import type { Setting, SettingParams } from "@/types/settingTypes";
+import { useSettingStore } from "@/stores/settingStore";
 
 export class ComfySettingsDialog extends ComfyDialog<HTMLDialogElement> {
   app: ComfyApp;
   settingsValues: any;
   settingsLookup: Record<string, Setting>;
+  settingsParamLookup: Record<string, SettingParams>;
 
   constructor(app: ComfyApp) {
     super();
@@ -15,6 +17,7 @@ export class ComfySettingsDialog extends ComfyDialog<HTMLDialogElement> {
     this.app = app;
     this.settingsValues = {};
     this.settingsLookup = {};
+    this.settingsParamLookup = {};
     this.element = $el(
       "dialog",
       {
@@ -55,6 +58,14 @@ export class ComfySettingsDialog extends ComfyDialog<HTMLDialogElement> {
   }
 
   #dispatchChange<T>(id: string, value: T, oldValue?: T) {
+    // Keep the settingStore updated. Not using `store.set` as it would trigger
+    // setSettingValue again.
+    // `load` re-dispatch the change for any settings added before load so
+    // settingStore is always up to date.
+    if (this.app.vueAppReady) {
+      useSettingStore().settingValues[id] = value;
+    }
+
     this.dispatchEvent(
       new CustomEvent(id + ".change", {
         detail: {
@@ -99,6 +110,11 @@ export class ComfySettingsDialog extends ComfyDialog<HTMLDialogElement> {
     return value ?? defaultValue;
   }
 
+  getSettingDefaultValue(id: string) {
+    const param = this.settingsParamLookup[id];
+    return param?.defaultValue;
+  }
+
   async setSettingValueAsync(id: string, value: any) {
     const json = JSON.stringify(value);
     localStorage["Comfy.Settings." + id] = json; // backwards compatibility for extensions keep setting in storage
@@ -165,8 +181,10 @@ export class ComfySettingsDialog extends ComfyDialog<HTMLDialogElement> {
     // Trigger initial setting of value
     if (!skipOnChange) {
       onChange?.(value, undefined);
+      this.#dispatchChange(id, value);
     }
 
+    this.settingsParamLookup[id] = params;
     this.settingsLookup[id] = {
       id,
       onChange,

src/stores/nodeDefStore.ts

@@ -0,0 +1,68 @@
+import { ComfyNodeDef } from "@/types/apiTypes";
+import { defineStore } from "pinia";
+
+export const SYSTEM_NODE_DEFS: ComfyNodeDef[] = [
+  {
+    name: "PrimitiveNode",
+    display_name: "Primitive",
+    category: "utils",
+    input: { required: {}, optional: {} },
+    output: ["*"],
+    output_name: ["connect to widget input"],
+    output_is_list: [false],
+    python_module: "nodes",
+    description: "Primitive values like numbers, strings, and booleans.",
+  },
+  {
+    name: "Reroute",
+    display_name: "Reroute",
+    category: "utils",
+    input: { required: { "": ["*"] }, optional: {} },
+    output: ["*"],
+    output_name: [""],
+    output_is_list: [false],
+    python_module: "nodes",
+    description: "Reroute the connection to another node.",
+  },
+  {
+    name: "Note",
+    display_name: "Note",
+    category: "utils",
+    input: { required: {}, optional: {} },
+    output: [],
+    output_name: [],
+    output_is_list: [],
+    python_module: "nodes",
+    description: "Node that add notes to your project",
+  },
+];
+
+const SYSTEM_NODE_DEFS_BY_NAME = SYSTEM_NODE_DEFS.reduce((acc, nodeDef) => {
+  acc[nodeDef.name] = nodeDef;
+  return acc;
+}, {}) as Record<string, ComfyNodeDef>;
+
+interface State {
+  nodeDefsByName: Record<string, ComfyNodeDef>;
+}
+
+export const useNodeDefStore = defineStore("nodeDef", {
+  state: (): State => ({
+    nodeDefsByName: SYSTEM_NODE_DEFS_BY_NAME,
+  }),
+  getters: {
+    nodeDefs(state) {
+      return Object.values(state.nodeDefsByName);
+    },
+  },
+  actions: {
+    addNodeDef(nodeDef: ComfyNodeDef) {
+      this.nodeDefsByName[nodeDef.name] = nodeDef;
+    },
+    addNodeDefs(nodeDefs: ComfyNodeDef[]) {
+      for (const nodeDef of nodeDefs) {
+        this.addNodeDef(nodeDef);
+      }
+    },
+  },
+});

src/stores/settingStore.ts

@@ -0,0 +1,41 @@
+/**
+ * TODO: Migrate scripts/ui/settings.ts here
+ *
+ * Currently the reactive settings act as a proxy of the legacy settings.
+ * Every time a setting is changed, the settingStore dispatch the change to the
+ * legacy settings. Every time the legacy settings are changed, the legacy
+ * settings directly updates the settingStore.settingValues.
+ */
+
+import { app } from "@/scripts/app";
+import { ComfySettingsDialog } from "@/scripts/ui/settings";
+import { defineStore } from "pinia";
+
+interface State {
+  settingValues: Record<string, any>;
+}
+
+export const useSettingStore = defineStore("setting", {
+  state: (): State => ({
+    settingValues: {},
+  }),
+  actions: {
+    addSettings(settings: ComfySettingsDialog) {
+      for (const id in settings.settingsLookup) {
+        const value = settings.getSettingValue(id);
+        this.settingValues[id] = value;
+      }
+    },
+
+    set(key: string, value: any) {
+      this.settingValues[key] = value;
+      app.ui.settings.setSettingValue(key, value);
+    },
+
+    get(key: string) {
+      return (
+        this.settingValues[key] ?? app.ui.settings.getSettingDefaultValue(key)
+      );
+    },
+  },
+});