Add frontend extension management panel (#1141)

* Manage register of extension in pinia

* Add disabled extensions setting

* nit

* Disable extension

* Add virtual divider

* Basic extension panel

* Style cell

* nit

* Fix loading

* inactive rules

* nit

* Calculate changes

* nit

* Experimental setting guard

src/components/dialog/content/SettingDialogContent.vue

@@ -57,6 +57,9 @@
           <TabPanel key="keybinding" value="Keybinding">
             <KeybindingPanel />
           </TabPanel>
+          <TabPanel key="extension" value="Extension">
+            <ExtensionPanel />
+          </TabPanel>
         </TabPanels>
       </Tabs>
     </div>
@@ -78,6 +81,7 @@ import NoResultsPlaceholder from '@/components/common/NoResultsPlaceholder.vue'
 import { flattenTree } from '@/utils/treeUtil'
 import AboutPanel from './setting/AboutPanel.vue'
 import KeybindingPanel from './setting/KeybindingPanel.vue'
+import ExtensionPanel from './setting/ExtensionPanel.vue'
 
 interface ISettingGroup {
   label: string
@@ -96,11 +100,24 @@ const keybindingPanelNode: SettingTreeNode = {
   children: []
 }
 
+const extensionPanelNode: SettingTreeNode = {
+  key: 'extension',
+  label: 'Extension',
+  children: []
+}
+
+const extensionPanelNodeList = computed<SettingTreeNode[]>(() => {
+  const settingStore = useSettingStore()
+  const showExtensionPanel = settingStore.get('Comfy.Settings.ExtensionPanel')
+  return showExtensionPanel ? [extensionPanelNode] : []
+})
+
 const settingStore = useSettingStore()
 const settingRoot = computed<SettingTreeNode>(() => settingStore.settingTree)
 const categories = computed<SettingTreeNode[]>(() => [
   ...(settingRoot.value.children || []),
   keybindingPanelNode,
+  ...extensionPanelNodeList.value,
   aboutPanelNode
 ])
 const activeCategory = ref<SettingTreeNode | null>(null)
@@ -226,4 +243,15 @@ const tabValue = computed(() =>
     width: 100%;
   }
 }
+
+/* Show a separator line above the Keybinding tab */
+/* This indicates the start of custom setting panels */
+.settings-sidebar :deep(.p-listbox-option[aria-label='Keybinding']) {
+  position: relative;
+}
+
+.settings-sidebar :deep(.p-listbox-option[aria-label='Keybinding'])::before {
+  @apply content-[''] top-0 left-0 absolute w-full;
+  border-top: 1px solid var(--p-divider-border-color);
+}
 </style>

src/components/dialog/content/setting/ExtensionPanel.vue

@@ -0,0 +1,93 @@
+<template>
+  <div class="extension-panel">
+    <DataTable :value="extensionStore.extensions" stripedRows size="small">
+      <Column field="name" :header="$t('extensionName')" sortable></Column>
+      <Column
+        :pt="{
+          bodyCell: 'flex items-center justify-end'
+        }"
+      >
+        <template #body="slotProps">
+          <ToggleSwitch
+            v-model="editingEnabledExtensions[slotProps.data.name]"
+            @change="updateExtensionStatus"
+          />
+        </template>
+      </Column>
+    </DataTable>
+    <div class="mt-4">
+      <Message v-if="hasChanges" severity="info">
+        <ul>
+          <li v-for="ext in changedExtensions" :key="ext.name">
+            <span>
+              {{ extensionStore.isExtensionEnabled(ext.name) ? '[-]' : '[+]' }}
+            </span>
+            {{ ext.name }}
+          </li>
+        </ul>
+      </Message>
+      <Button
+        :label="$t('reloadToApplyChanges')"
+        icon="pi pi-refresh"
+        @click="applyChanges"
+        :disabled="!hasChanges"
+        text
+        fluid
+        severity="danger"
+      />
+    </div>
+  </div>
+</template>
+
+<script setup lang="ts">
+import { ref, computed, onMounted } from 'vue'
+import { useExtensionStore } from '@/stores/extensionStore'
+import { useSettingStore } from '@/stores/settingStore'
+import DataTable from 'primevue/datatable'
+import Column from 'primevue/column'
+import ToggleSwitch from 'primevue/toggleswitch'
+import Button from 'primevue/button'
+import Message from 'primevue/message'
+
+const extensionStore = useExtensionStore()
+const settingStore = useSettingStore()
+
+const editingEnabledExtensions = ref<Record<string, boolean>>({})
+
+onMounted(() => {
+  extensionStore.extensions.forEach((ext) => {
+    editingEnabledExtensions.value[ext.name] =
+      extensionStore.isExtensionEnabled(ext.name)
+  })
+})
+
+const changedExtensions = computed(() => {
+  return extensionStore.extensions.filter(
+    (ext) =>
+      editingEnabledExtensions.value[ext.name] !==
+      extensionStore.isExtensionEnabled(ext.name)
+  )
+})
+
+const hasChanges = computed(() => {
+  return changedExtensions.value.length > 0
+})
+
+const updateExtensionStatus = () => {
+  const editingDisabledExtensionNames = Object.entries(
+    editingEnabledExtensions.value
+  )
+    .filter(([_, enabled]) => !enabled)
+    .map(([name]) => name)
+
+  settingStore.set('Comfy.Extension.Disabled', [
+    ...extensionStore.inactiveDisabledExtensionNames,
+    ...editingDisabledExtensionNames
+  ])
+}
+
+const applyChanges = () => {
+  // Refresh the page to apply changes
+  window.location.reload()
+}
+</script>

src/extensions/core/nodeBadge.ts

@@ -9,7 +9,6 @@ import _ from 'lodash'
 import { getColorPalette, defaultColorPalette } from './colorPalette'
 import { BadgePosition } from '@comfyorg/litegraph'
 import type { Palette } from '@/types/colorPalette'
-import type { ComfyNodeDef } from '@/types/apiTypes'
 import { useNodeDefStore } from '@/stores/nodeDefStore'
 
 function getNodeSource(node: LGraphNode): NodeSource | null {

src/i18n.ts

@@ -2,6 +2,8 @@ import { createI18n } from 'vue-i18n'
 
 const messages = {
   en: {
+    extensionName: 'Extension Name',
+    reloadToApplyChanges: 'Reload to apply changes',
     insert: 'Insert',
     systemInfo: 'System Info',
     devices: 'Devices',
@@ -108,6 +110,8 @@ const messages = {
     }
   },
   zh: {
+    extensionName: '扩展名称',
+    reloadToApplyChanges: '重新加载以应用更改',
     insert: '插入',
     systemInfo: '系统信息',
     devices: '设备',

src/scripts/app.ts

@@ -52,8 +52,7 @@ import type { ToastMessageOptions } from 'primevue/toast'
 import { useWorkspaceStore } from '@/stores/workspaceStateStore'
 import { useExecutionStore } from '@/stores/executionStore'
 import { IWidget } from '@comfyorg/litegraph'
-import { useKeybindingStore } from '@/stores/keybindingStore'
-import { useCommandStore } from '@/stores/commandStore'
+import { useExtensionStore } from '@/stores/extensionStore'
 
 export const ANIM_PREVIEW_WIDGET = '$$comfy_animation_preview'
 
@@ -357,6 +356,13 @@ export class ComfyApp {
     }
   }
 
+  get enabledExtensions() {
+    if (!this.vueAppReady) {
+      return this.extensions
+    }
+    return useExtensionStore().enabledExtensions
+  }
+
   /**
    * Invoke an extension callback
    * @param {keyof ComfyExtension} method The extension callback to execute
@@ -365,7 +371,7 @@ export class ComfyApp {
    */
   #invokeExtensions(method, ...args) {
     let results = []
-    for (const ext of this.extensions) {
+    for (const ext of this.enabledExtensions) {
       if (method in ext) {
         try {
           results.push(ext[method](...args, this))
@@ -391,7 +397,7 @@ export class ComfyApp {
    */
   async #invokeExtensionsAsync(method, ...args) {
     return await Promise.all(
-      this.extensions.map(async (ext) => {
+      this.enabledExtensions.map(async (ext) => {
         if (method in ext) {
           try {
             return await ext[method](...args, this)
@@ -1773,6 +1779,8 @@ export class ComfyApp {
    * Loads all extensions from the API into the window in parallel
    */
   async #loadExtensions() {
+    useExtensionStore().loadDisabledExtensionNames()
+
     const extensions = await api.getExtensions()
     this.logging.addEntry('Comfy.App', 'debug', { Extensions: extensions })
 
@@ -2943,22 +2951,12 @@ export class ComfyApp {
    * @param {ComfyExtension} extension
    */
   registerExtension(extension: ComfyExtension) {
-    if (!extension.name) {
-      throw new Error("Extensions must have a 'name' property.")
-    }
-    // https://github.com/Comfy-Org/litegraph.js/pull/117
-    if (extension.name === 'pysssss.Locking') {
-      console.log('pysssss.Locking is replaced by pin/unpin in ComfyUI core.')
-      return
-    }
-    if (this.extensions.find((ext) => ext.name === extension.name)) {
-      throw new Error(`Extension named '${extension.name}' already registered.`)
-    }
     if (this.vueAppReady) {
-      useKeybindingStore().loadExtensionKeybindings(extension)
-      useCommandStore().loadExtensionCommands(extension)
+      useExtensionStore().registerExtension(extension)
+    } else {
+      // For jest testing.
+      this.extensions.push(extension)
     }
-    this.extensions.push(extension)
   }
 
   /**

src/stores/coreSettings.ts

@@ -420,5 +420,20 @@ export const CORE_SETTINGS: SettingParams[] = [
     type: 'hidden',
     defaultValue: [] as Keybinding[],
     versionAdded: '1.3.7'
+  },
+  {
+    id: 'Comfy.Extension.Disabled',
+    name: 'Disabled extension names',
+    type: 'hidden',
+    defaultValue: [] as string[],
+    versionAdded: '1.3.11'
+  },
+  {
+    id: 'Comfy.Settings.ExtensionPanel',
+    name: 'Show extension panel in settings dialog',
+    type: 'boolean',
+    defaultValue: false,
+    experimental: true,
+    versionAdded: '1.3.11'
   }
 ]

src/stores/extensionStore.ts

@@ -0,0 +1,81 @@
+import { ref, computed, markRaw } from 'vue'
+import { defineStore } from 'pinia'
+import type { ComfyExtension } from '@/types/comfy'
+import { useKeybindingStore } from './keybindingStore'
+import { useCommandStore } from './commandStore'
+import { useSettingStore } from './settingStore'
+import { app } from '@/scripts/app'
+
+export const useExtensionStore = defineStore('extension', () => {
+  // For legacy reasons, the name uniquely identifies an extension
+  const extensionByName = ref<Record<string, ComfyExtension>>({})
+  const extensions = computed(() => Object.values(extensionByName.value))
+  // Not using computed because disable extension requires reloading of the page.
+  // Dynamically update this list won't affect extensions that are already loaded.
+  const disabledExtensionNames = ref<Set<string>>(new Set())
+
+  // Disabled extension names that are currently not in the extension list.
+  // If a node pack is disabled in the backend, we shouldn't remove the configuration
+  // of the frontend extension disable list, in case the node pack is re-enabled.
+  const inactiveDisabledExtensionNames = computed(() => {
+    return Array.from(disabledExtensionNames.value).filter(
+      (name) => !(name in extensionByName.value)
+    )
+  })
+
+  const isExtensionEnabled = (name: string) =>
+    !disabledExtensionNames.value.has(name)
+  const enabledExtensions = computed(() => {
+    return extensions.value.filter((ext) => isExtensionEnabled(ext.name))
+  })
+
+  function registerExtension(extension: ComfyExtension) {
+    if (!extension.name) {
+      throw new Error("Extensions must have a 'name' property.")
+    }
+
+    if (extensionByName.value[extension.name]) {
+      throw new Error(`Extension named '${extension.name}' already registered.`)
+    }
+
+    if (disabledExtensionNames.value.has(extension.name)) {
+      console.log(`Extension ${extension.name} is disabled.`)
+    }
+
+    extensionByName.value[extension.name] = markRaw(extension)
+    useKeybindingStore().loadExtensionKeybindings(extension)
+    useCommandStore().loadExtensionCommands(extension)
+
+    /*
+     * Extensions are currently stored in both extensionStore and app.extensions.
+     * Legacy jest tests still depend on app.extensions being populated.
+     */
+    app.extensions.push(extension)
+  }
+
+  function loadDisabledExtensionNames() {
+    disabledExtensionNames.value = new Set(
+      useSettingStore().get('Comfy.Extension.Disabled')
+    )
+    // pysssss.Locking is replaced by pin/unpin in ComfyUI core.
+    // https://github.com/Comfy-Org/litegraph.js/pull/117
+    disabledExtensionNames.value.add('pysssss.Locking')
+  }
+
+  // Some core extensions are registered before the store is initialized, e.g.
+  // colorPalette.
+  // Register them manually here so the state of app.extensions and
+  // extensionByName are in sync.
+  for (const ext of app.extensions) {
+    extensionByName.value[ext.name] = markRaw(ext)
+  }
+
+  return {
+    extensions,
+    enabledExtensions,
+    inactiveDisabledExtensionNames,
+    isExtensionEnabled,
+    registerExtension,
+    loadDisabledExtensionNames
+  }
+})

src/types/apiTypes.ts

@@ -507,7 +507,9 @@ const zSettings = z.record(z.any()).and(
       'Comfy.NodeBadge.NodeLifeCycleBadgeMode': zNodeBadgeMode,
       'Comfy.QueueButton.BatchCountLimit': z.number(),
       'Comfy.Keybinding.UnsetBindings': z.array(zKeybinding),
-      'Comfy.Keybinding.NewBindings': z.array(zKeybinding)
+      'Comfy.Keybinding.NewBindings': z.array(zKeybinding),
+      'Comfy.Extension.Disabled': z.array(z.string()),
+      'Comfy.Settings.ExtensionPanel': z.boolean()
     })
     .optional()
 )

tests-ui/globalSetup.ts

@@ -36,6 +36,15 @@ module.exports = async function () {
     }
   })
 
+  jest.mock('@/stores/extensionStore', () => {
+    return {
+      useExtensionStore: () => ({
+        registerExtension: jest.fn(),
+        loadDisabledExtensionNames: jest.fn()
+      })
+    }
+  })
+
   jest.mock('vue-i18n', () => {
     return {
       useI18n: jest.fn()