From 19c70d95d34d4e9a34c71860b784921ad980e1e4 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Wed, 24 Jul 2024 21:31:59 -0400
Subject: [PATCH] Sidebar tab API for extensions (#215)

* Add extensionManager to manage tabs

* Fix null bug

* nit
---
 src/App.vue                                   | 20 ++++-
 .../LiteGraphCanvasSplitterOverlay.vue        | 12 +--
 src/components/sidebar/SideToolBar.vue        | 83 ++++++++++++-------
 .../sidebar/items/NodeDetailSideBarItem.vue   | 20 -----
 .../QueueSideBarTab.vue}                      |  0
 src/scripts/app.ts                            |  2 +
 src/scripts/extensionManager.ts               | 43 ++++++++++
 src/stores/workspaceStateStore.ts             | 32 +++++++
 src/types/extensionTypes.ts                   | 30 +++++++
 9 files changed, 181 insertions(+), 61 deletions(-)
 delete mode 100644 src/components/sidebar/items/NodeDetailSideBarItem.vue
 rename src/components/sidebar/{items/QueueSideBarItem.vue => tabs/QueueSideBarTab.vue} (100%)
 create mode 100644 src/scripts/extensionManager.ts
 create mode 100644 src/stores/workspaceStateStore.ts
 create mode 100644 src/types/extensionTypes.ts

diff --git a/src/App.vue b/src/App.vue
index 55f8811809..7a4d69de80 100644
--- a/src/App.vue
+++ b/src/App.vue
@@ -4,8 +4,8 @@
     <NodeSearchboxPopover v-if="nodeSearchEnabled" />
     <teleport to="#graph-canvas-container">
       <LiteGraphCanvasSplitterOverlay>
-        <template #side-bar-panel="{ setPanelVisible }">
-          <SideToolBar @change="setPanelVisible($event)" />
+        <template #side-bar-panel>
+          <SideToolBar />
         </template>
       </LiteGraphCanvasSplitterOverlay>
     </teleport>
@@ -13,14 +13,17 @@
 </template>
 
 <script setup lang="ts">
-import { computed, onMounted, ref, watch } from "vue";
+import { computed, markRaw, 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 QueueSideBarTab from "@/components/sidebar/tabs/QueueSideBarTab.vue";
 import ProgressSpinner from "primevue/progressspinner";
 import { app } from "./scripts/app";
 import { useSettingStore } from "./stores/settingStore";
 import { useNodeDefStore } from "./stores/nodeDefStore";
+import { ExtensionManagerImpl } from "./scripts/extensionManager";
+import { useI18n } from "vue-i18n";
 
 const isLoading = ref(true);
 const nodeSearchEnabled = computed<boolean>(
@@ -43,10 +46,21 @@ watch(
   { immediate: true }
 );
 
+const { t } = useI18n();
 const init = () => {
   useNodeDefStore().addNodeDefs(Object.values(app.nodeDefs));
   useSettingStore().addSettings(app.ui.settings);
   app.vueAppReady = true;
+  // Late init as extension manager needs to access pinia store.
+  app.extensionManager = new ExtensionManagerImpl();
+  app.extensionManager.registerSidebarTab({
+    id: "queue",
+    icon: "pi pi-history",
+    title: t("sideToolBar.queue"),
+    tooltip: t("sideToolBar.queue"),
+    component: markRaw(QueueSideBarTab),
+    type: "vue",
+  });
 };
 
 onMounted(() => {
diff --git a/src/components/LiteGraphCanvasSplitterOverlay.vue b/src/components/LiteGraphCanvasSplitterOverlay.vue
index bacd64ff7c..836df58651 100644
--- a/src/components/LiteGraphCanvasSplitterOverlay.vue
+++ b/src/components/LiteGraphCanvasSplitterOverlay.vue
@@ -6,7 +6,7 @@
       :size="20"
       v-show="sideBarPanelVisible"
     >
-      <slot name="side-bar-panel" :setPanelVisible="setPanelVisible"></slot>
+      <slot name="side-bar-panel"></slot>
     </SplitterPanel>
     <SplitterPanel class="graph-canvas-panel" :size="100">
       <div></div>
@@ -15,14 +15,14 @@
 </template>
 
 <script setup lang="ts">
+import { useWorkspaceStore } from "@/stores/workspaceStateStore";
 import Splitter from "primevue/splitter";
 import SplitterPanel from "primevue/splitterpanel";
-import { computed, ref } from "vue";
+import { computed } from "vue";
 
-const sideBarPanelVisible = ref(false);
-const setPanelVisible = (visible: boolean) => {
-  sideBarPanelVisible.value = visible;
-};
+const sideBarPanelVisible = computed(
+  () => useWorkspaceStore().activeSidebarTab !== null
+);
 const gutterClass = computed(() => {
   return sideBarPanelVisible.value ? "" : "gutter-hidden";
 });
diff --git a/src/components/sidebar/SideToolBar.vue b/src/components/sidebar/SideToolBar.vue
index 88d437e7e6..b1e97771c5 100644
--- a/src/components/sidebar/SideToolBar.vue
+++ b/src/components/sidebar/SideToolBar.vue
@@ -2,11 +2,12 @@
   <teleport to=".comfyui-body-left">
     <nav class="side-tool-bar-container">
       <SideBarIcon
-        v-for="item in items"
-        :icon="item.icon"
-        :tooltip="item.tooltip"
-        :selected="item === selectedItem"
-        @click="onItemClick(item)"
+        v-for="tab in tabs"
+        :key="tab.id"
+        :icon="tab.icon"
+        :tooltip="tab.tooltip"
+        :selected="tab === selectedTab"
+        @click="onTabClick(tab)"
       />
       <div class="side-tool-bar-end">
         <SideBarThemeToggleIcon />
@@ -14,50 +15,68 @@
       </div>
     </nav>
   </teleport>
-  <component :is="selectedItem?.component" />
+  <div v-if="!selectedTab"></div>
+  <component
+    v-else-if="selectedTab.type === 'vue'"
+    :is="selectedTab.component"
+  />
+  <div
+    v-else
+    :ref="
+      (el) => {
+        if (el)
+          mountCustomTab(
+            selectedTab as CustomSidebarTabExtension,
+            el as HTMLElement
+          );
+      }
+    "
+  ></div>
 </template>
 
 <script setup lang="ts">
 import SideBarIcon from "./SideBarIcon.vue";
 import SideBarThemeToggleIcon from "./SideBarThemeToggleIcon.vue";
 import SideBarSettingsToggleIcon from "./SideBarSettingsToggleIcon.vue";
-import NodeDetailSideBarItem from "./items/NodeDetailSideBarItem.vue";
-import QueueSideBarItem from "./items/QueueSideBarItem.vue";
-import { computed, markRaw, ref, watch } from "vue";
-import { useI18n } from "vue-i18n";
+import { computed, onBeforeUnmount, watch } from "vue";
 import { useSettingStore } from "@/stores/settingStore";
+import { app } from "@/scripts/app";
+import { useWorkspaceStore } from "@/stores/workspaceStateStore";
+import {
+  CustomSidebarTabExtension,
+  SidebarTabExtension,
+} from "@/types/extensionTypes";
 
-const { t } = useI18n();
-const items = ref([
-  // { icon: "pi pi-map", component: markRaw(NodeDetailSideBarItem) },
-  {
-    icon: "pi pi-history",
-    tooltip: t("sideToolBar.queue"),
-    component: markRaw(QueueSideBarItem),
-  },
-]);
-const selectedItem = ref(null);
-const onItemClick = (item) => {
-  if (selectedItem.value === item) {
-    selectedItem.value = null;
-    return;
-  }
-  selectedItem.value = item;
-};
-
-const emit = defineEmits(["change"]);
-watch(selectedItem, (newVal) => {
-  emit("change", newVal !== null);
+const workspaceStateStore = useWorkspaceStore();
+const tabs = computed(() => app.extensionManager.getSidebarTabs());
+const selectedTab = computed<SidebarTabExtension | null>(() => {
+  const tabId = workspaceStateStore.activeSidebarTab;
+  return tabs.value.find((tab) => tab.id === tabId) || null;
 });
+const mountCustomTab = (tab: CustomSidebarTabExtension, el: HTMLElement) => {
+  tab.render(el);
+};
+const onTabClick = (item: SidebarTabExtension) => {
+  workspaceStateStore.updateActiveSidebarTab(
+    workspaceStateStore.activeSidebarTab === item.id ? null : item.id
+  );
+};
 
 const betaMenuEnabled = computed(
   () => useSettingStore().get("Comfy.UseNewMenu") !== "Disabled"
 );
 watch(betaMenuEnabled, (newValue) => {
   if (!newValue) {
-    selectedItem.value = null;
+    workspaceStateStore.updateActiveSidebarTab(null);
   }
 });
+onBeforeUnmount(() => {
+  tabs.value.forEach((tab) => {
+    if (tab.type === "custom" && tab.destroy) {
+      tab.destroy();
+    }
+  });
+});
 </script>
 
 <style>
diff --git a/src/components/sidebar/items/NodeDetailSideBarItem.vue b/src/components/sidebar/items/NodeDetailSideBarItem.vue
deleted file mode 100644
index ad19897db8..0000000000
--- a/src/components/sidebar/items/NodeDetailSideBarItem.vue
+++ /dev/null
@@ -1,20 +0,0 @@
-<template>
-  <Accordion>
-    <AccordionPanel v-for="node in nodes" :key="node.id">
-      <AccordionHeader>
-        <div>{{ node.id }} - {{ node.type }}</div>
-      </AccordionHeader>
-    </AccordionPanel>
-  </Accordion>
-</template>
-
-<script setup lang="ts">
-import { defaultGraph } from "@/scripts/defaultGraph";
-import Accordion from "primevue/accordion";
-import AccordionHeader from "primevue/accordionheader";
-import AccordionPanel from "primevue/accordionpanel";
-
-/* Placeholder workflow display */
-const workflow = defaultGraph;
-const nodes = workflow.nodes.sort((a, b) => a.id - b.id);
-</script>
diff --git a/src/components/sidebar/items/QueueSideBarItem.vue b/src/components/sidebar/tabs/QueueSideBarTab.vue
similarity index 100%
rename from src/components/sidebar/items/QueueSideBarItem.vue
rename to src/components/sidebar/tabs/QueueSideBarTab.vue
diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index 4b8c2dcbc0..e56e516b6d 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -35,6 +35,7 @@ import { StorageLocation } from "@/types/settingTypes";
 // CSS imports. style.css must be imported later as it overwrites some litegraph styles.
 import "@comfyorg/litegraph/css/litegraph.css";
 import "../assets/css/style.css";
+import { ExtensionManager } from "@/types/extensionTypes";
 
 export const ANIM_PREVIEW_WIDGET = "$$comfy_animation_preview";
 
@@ -90,6 +91,7 @@ export class ComfyApp {
   ui: ComfyUI;
   logging: ComfyLogging;
   extensions: ComfyExtension[];
+  extensionManager: ExtensionManager;
   _nodeOutputs: Record<string, any>;
   nodePreviewImages: Record<string, typeof Image>;
   shiftDown: boolean;
diff --git a/src/scripts/extensionManager.ts b/src/scripts/extensionManager.ts
new file mode 100644
index 0000000000..16e95435c4
--- /dev/null
+++ b/src/scripts/extensionManager.ts
@@ -0,0 +1,43 @@
+import { useWorkspaceStore } from "@/stores/workspaceStateStore";
+import { ExtensionManager, SidebarTabExtension } from "@/types/extensionTypes";
+
+export class ExtensionManagerImpl implements ExtensionManager {
+  private sidebarTabs: SidebarTabExtension[] = [];
+  private workspaceStore = useWorkspaceStore();
+
+  registerSidebarTab(tab: SidebarTabExtension) {
+    this.sidebarTabs.push(tab);
+    this.updateSidebarOrder();
+  }
+
+  unregisterSidebarTab(id: string) {
+    const index = this.sidebarTabs.findIndex((tab) => tab.id === id);
+    if (index !== -1) {
+      const tab = this.sidebarTabs[index];
+      if (tab.type === "custom" && tab.destroy) {
+        tab.destroy();
+      }
+      this.sidebarTabs.splice(index, 1);
+      this.updateSidebarOrder();
+    }
+  }
+
+  getSidebarTabs() {
+    return this.sidebarTabs.sort((a, b) => {
+      const orderA = this.workspaceStore.sidebarTabsOrder.indexOf(a.id);
+      const orderB = this.workspaceStore.sidebarTabsOrder.indexOf(b.id);
+      return orderA - orderB;
+    });
+  }
+
+  private updateSidebarOrder() {
+    const currentOrder = this.workspaceStore.sidebarTabsOrder;
+    const newTabs = this.sidebarTabs.filter(
+      (tab) => !currentOrder.includes(tab.id)
+    );
+    this.workspaceStore.updateSidebarOrder([
+      ...currentOrder,
+      ...newTabs.map((tab) => tab.id),
+    ]);
+  }
+}
diff --git a/src/stores/workspaceStateStore.ts b/src/stores/workspaceStateStore.ts
new file mode 100644
index 0000000000..865956813c
--- /dev/null
+++ b/src/stores/workspaceStateStore.ts
@@ -0,0 +1,32 @@
+import { defineStore } from "pinia";
+
+interface WorkspaceState {
+  activeSidebarTab: string | null;
+  sidebarTabsOrder: string[]; // Array of tab IDs in order
+}
+
+export const useWorkspaceStore = defineStore("workspace", {
+  state: (): WorkspaceState => ({
+    activeSidebarTab: null,
+    sidebarTabsOrder: [],
+  }),
+  actions: {
+    updateActiveSidebarTab(tabId: string) {
+      this.activeSidebarTab = tabId;
+    },
+    updateSidebarOrder(newOrder: string[]) {
+      this.sidebarTabsOrder = newOrder;
+    },
+    serialize() {
+      return JSON.stringify({
+        activeSidebarTab: this.activeSidebarTab,
+        sidebarTabsOrder: this.sidebarTabsOrder,
+      });
+    },
+    deserialize(state: string) {
+      const parsedState = JSON.parse(state);
+      this.sidebarTabsOrder = parsedState.sidebarTabsOrder;
+      this.activeSidebarTab = parsedState.activeSidebarTab;
+    },
+  },
+});
diff --git a/src/types/extensionTypes.ts b/src/types/extensionTypes.ts
new file mode 100644
index 0000000000..8d000585d2
--- /dev/null
+++ b/src/types/extensionTypes.ts
@@ -0,0 +1,30 @@
+import { Component } from "vue";
+
+export interface BaseSidebarTabExtension {
+  id: string;
+  title: string;
+  icon?: string;
+  order?: number;
+  tooltip?: string;
+}
+
+export interface VueSidebarTabExtension extends BaseSidebarTabExtension {
+  type: "vue";
+  component: Component;
+}
+
+export interface CustomSidebarTabExtension extends BaseSidebarTabExtension {
+  type: "custom";
+  render: (container: HTMLElement) => void;
+  destroy?: () => void;
+}
+
+export type SidebarTabExtension =
+  | VueSidebarTabExtension
+  | CustomSidebarTabExtension;
+
+export interface ExtensionManager {
+  registerSidebarTab(tab: SidebarTabExtension): void;
+  unregisterSidebarTab(id: string): void;
+  getSidebarTabs(): SidebarTabExtension[];
+}