Allow extension register custom topbar menu command (#982)

* Refactor command store

* Rename coreMenuStore to menuStore

* Extension API to register command

* Update README

* Add playwright test

README.md

@@ -137,6 +137,37 @@ https://github.com/user-attachments/assets/c142c43f-2fe9-4030-8196-b3bfd4c6977d
 </details>
 
 ### Node developers API
+
+<details>
+  <summary>v1.3.1: Extension API to register custom topbar menu items</summary>
+
+  Extensions can call the following API to register custom topbar menu items.
+
+```js
+  app.extensionManager.menu.registerTopbarCommands(["ext", "ext2"], [{id:"foo", label: "foo", function: () => alert(1)}])
+```
+
+![image](https://github.com/user-attachments/assets/ae7b082f-7ce9-4549-a446-4563567102fe)
+</details>
+
+<details>
+  <summary>v1.2.27: Extension API to add toast message</summary>i
+
+  Extensions can call the following API to add toast messages.
+
+```js
+  app.extensionManager.toast.add({
+    severity: 'info',
+    summary: 'Loaded!',
+    detail: 'Extension loaded!',
+    life: 3000
+  })
+```
+Documentation of all supported options can be found here: <https://primevue.org/toast/#api.toast.interfaces.ToastMessageOptions>
+
+![image](https://github.com/user-attachments/assets/de02cd7e-cd81-43d1-a0b0-bccef92ff487)
+</details>
+
 <details>
   <summary>v1.2.4: Extension API to register custom sidebar tab</summary>
 
@@ -162,24 +193,6 @@ We will support custom icons later.
 ![image](https://github.com/user-attachments/assets/7bff028a-bf91-4cab-bf97-55c243b3f5e0)
 </details>
 
-<details>
-  <summary>v1.2.27: Extension API to add toast message</summary>
-
-  Extensions can call the following API to add toast messages.
-
-```js
-  app.extensionManager.toast.add({
-    severity: 'info',
-    summary: 'Loaded!',
-    detail: 'Extension loaded!',
-    life: 3000
-  })
-```
-Documentation of all supported options can be found here: <https://primevue.org/toast/#api.toast.interfaces.ToastMessageOptions>
-
-![image](https://github.com/user-attachments/assets/de02cd7e-cd81-43d1-a0b0-bccef92ff487)
-</details>
-
 ## Road Map
 
 ### What has been done

browser_tests/ComfyPage.ts

@@ -229,6 +229,32 @@ class Topbar {
       .locator('.workflow-tabs .workflow-label')
       .allInnerTexts()
   }
+
+  async triggerTopbarCommand(path: string[]) {
+    if (path.length < 2) {
+      throw new Error('Path is too short')
+    }
+
+    const tabName = path[0]
+    const topLevelMenu = this.page.locator(
+      `.top-menubar .p-menubar-item:has-text("${tabName}")`
+    )
+    await topLevelMenu.waitFor({ state: 'visible' })
+    await topLevelMenu.click()
+
+    for (let i = 1; i < path.length; i++) {
+      const commandName = path[i]
+      const menuItem = this.page.locator(
+        `.top-menubar .p-menubar-submenu .p-menubar-item:has-text("${commandName}")`
+      )
+      await menuItem.waitFor({ state: 'visible' })
+      await menuItem.hover()
+
+      if (i === path.length - 1) {
+        await menuItem.click()
+      }
+    }
+  }
 }
 
 class ComfyMenu {

browser_tests/extensionAPI.spec.ts

@@ -0,0 +1,32 @@
+import { expect } from '@playwright/test'
+import { comfyPageFixture as test } from './ComfyPage'
+
+test.describe('Topbar commands', () => {
+  test.beforeEach(async ({ comfyPage }) => {
+    await comfyPage.setSetting('Comfy.UseNewMenu', 'Floating')
+  })
+
+  test.afterEach(async ({ comfyPage }) => {
+    await comfyPage.setSetting('Comfy.UseNewMenu', 'Disabled')
+  })
+
+  test('Should allow registering topbar commands', async ({ comfyPage }) => {
+    await comfyPage.page.evaluate(() => {
+      window['app'].extensionManager.menu.registerTopbarCommands(
+        ['ext'],
+        [
+          {
+            id: 'foo',
+            label: 'foo',
+            function: () => {
+              window['foo'] = true
+            }
+          }
+        ]
+      )
+    })
+
+    await comfyPage.menu.topbar.triggerTopbarCommand(['ext', 'foo'])
+    expect(await comfyPage.page.evaluate(() => window['foo'])).toBe(true)
+  })
+})

src/components/appMenu/AppMenu.vue

@@ -29,7 +29,7 @@
             icon="pi pi-times"
             :severity="executingPrompt ? 'danger' : 'secondary'"
             :disabled="!executingPrompt"
-            @click="() => commandStore.getCommand('Comfy.Interrupt')()"
+            @click="() => commandStore.getCommandFunction('Comfy.Interrupt')()"
           >
           </Button>
           <Button
@@ -37,7 +37,9 @@
             icon="pi pi-stop"
             :severity="hasPendingTasks ? 'danger' : 'secondary'"
             :disabled="!hasPendingTasks"
-            @click="() => commandStore.getCommand('Comfy.ClearPendingTasks')()"
+            @click="
+              () => commandStore.getCommandFunction('Comfy.ClearPendingTasks')()
+            "
           />
         </ButtonGroup>
       </div>
@@ -48,14 +50,15 @@
           icon="pi pi-refresh"
           severity="secondary"
           @click="
-            () => commandStore.getCommand('Comfy.RefreshNodeDefinitions')()
+            () =>
+              commandStore.getCommandFunction('Comfy.RefreshNodeDefinitions')()
           "
         />
         <Button
           v-tooltip.bottom="$t('menu.resetView')"
           icon="pi pi-expand"
           severity="secondary"
-          @click="() => commandStore.getCommand('Comfy.ResetView')()"
+          @click="() => commandStore.getCommandFunction('Comfy.ResetView')()"
         />
       </ButtonGroup>
     </div>

src/components/sidebar/tabs/QueueSidebarTab.vue

@@ -36,7 +36,9 @@
           icon="pi pi-stop"
           severity="danger"
           text
-          @click="() => commandStore.getCommand('Comfy.ClearPendingTasks')()"
+          @click="
+            () => commandStore.getCommandFunction('Comfy.ClearPendingTasks')()
+          "
           v-tooltip.bottom="$t('sideToolbar.queueTab.clearPendingTasks')"
         />
         <Button

src/components/sidebar/tabs/WorkflowsSidebarTab.vue

@@ -6,27 +6,33 @@
         icon="pi pi-th-large"
         v-tooltip="$t('sideToolbar.browseTemplates')"
         text
-        @click="() => commandStore.getCommand('Comfy.BrowseTemplates')()"
+        @click="
+          () => commandStore.getCommandFunction('Comfy.BrowseTemplates')()
+        "
       />
       <Button
         class="browse-workflows-button"
         icon="pi pi-folder-open"
         v-tooltip="'Browse for an image or exported workflow'"
         text
-        @click="() => commandStore.getCommand('Comfy.OpenWorkflow')()"
+        @click="() => commandStore.getCommandFunction('Comfy.OpenWorkflow')()"
       />
       <Button
         class="new-default-workflow-button"
         icon="pi pi-code"
         v-tooltip="'Load default workflow'"
         text
-        @click="() => commandStore.getCommand('Comfy.LoadDefaultWorkflow')()"
+        @click="
+          () => commandStore.getCommandFunction('Comfy.LoadDefaultWorkflow')()
+        "
       />
       <Button
         class="new-blank-workflow-button"
         icon="pi pi-plus"
         v-tooltip="'Create a new blank workflow'"
-        @click="() => commandStore.getCommand('Comfy.NewBlankWorkflow')()"
+        @click="
+          () => commandStore.getCommandFunction('Comfy.NewBlankWorkflow')()
+        "
         text
       />
     </template>

src/components/topbar/TopMenubar.vue

@@ -1,13 +1,10 @@
 <template>
   <teleport to=".comfyui-body-top">
-    <div
-      class="top-menubar comfyui-menu flex items-center"
-      v-show="betaMenuEnabled"
-    >
+    <div class="comfyui-menu flex items-center" v-show="betaMenuEnabled">
       <h1 class="comfyui-logo mx-2">ComfyUI</h1>
       <Menubar
         :model="items"
-        class="border-none p-0 bg-transparent"
+        class="top-menubar border-none p-0 bg-transparent"
         :pt="{
           rootList: 'gap-0 flex-nowrap'
         }"
@@ -26,7 +23,7 @@
 import Menubar from 'primevue/menubar'
 import Divider from 'primevue/divider'
 import WorkflowTabs from '@/components/topbar/WorkflowTabs.vue'
-import { useCoreMenuItemStore } from '@/stores/coreMenuItemStore'
+import { useMenuItemStore } from '@/stores/menuItemStore'
 import { computed, onMounted, ref } from 'vue'
 import { useSettingStore } from '@/stores/settingStore'
 import { app } from '@/scripts/app'
@@ -38,8 +35,8 @@ const workflowTabsPosition = computed(() =>
 const betaMenuEnabled = computed(
   () => settingStore.get('Comfy.UseNewMenu') !== 'Disabled'
 )
-const coreMenuItemsStore = useCoreMenuItemStore()
-const items = coreMenuItemsStore.menuItems
+const menuItemsStore = useMenuItemStore()
+const items = menuItemsStore.menuItems
 
 const menuRight = ref<HTMLDivElement | null>(null)
 // Menu-right holds legacy topbar elements attached by custom scripts

src/stores/commandStore.ts

@@ -8,90 +8,198 @@ import { useToastStore } from '@/stores/toastStore'
 import { showTemplateWorkflowsDialog } from '@/services/dialogService'
 import { useQueueStore } from './queueStore'
 
-type Command = () => void | Promise<void>
+export interface ComfyCommand {
+  id: string
+  function: () => void | Promise<void>
+
+  label?: string | (() => string)
+  icon?: string | (() => string)
+  tooltip?: string | (() => string)
+  shortcut?: string
+}
 
 const getTracker = () =>
   app.workflowManager.activeWorkflow?.changeTracker ?? globalTracker
 
 export const useCommandStore = defineStore('command', () => {
   const settingStore = useSettingStore()
-  const commands = ref<Record<string, Command>>({
-    'Comfy.NewBlankWorkflow': () => {
-      app.workflowManager.setWorkflow(null)
-      app.clean()
-      app.graph.clear()
-      app.workflowManager.activeWorkflow.track()
-    },
-    'Comfy.OpenWorkflow': () => {
-      app.ui.loadFile()
-    },
-    'Comfy.LoadDefaultWorkflow': async () => {
-      await app.loadGraphData()
-    },
-    'Comfy.SaveWorkflow': () => {
-      app.workflowManager.activeWorkflow.save()
-    },
-    'Comfy.SaveWorkflowAs': () => {
-      app.workflowManager.activeWorkflow.save(true)
-    },
-    'Comfy.ExportWorkflow': () => {
-      app.menu.exportWorkflow('workflow', 'workflow')
-    },
-    'Comfy.ExportWorkflowAPI': () => {
-      app.menu.exportWorkflow('workflow_api', 'output')
-    },
-    'Comfy.Undo': async () => {
-      await getTracker().undo()
-    },
-    'Comfy.Redo': async () => {
-      await getTracker().redo()
-    },
-    'Comfy.ClearWorkflow': () => {
-      if (
-        !settingStore.get('Comfy.ComfirmClear') ||
-        confirm('Clear workflow?')
-      ) {
+
+  const commands = ref<Record<string, ComfyCommand>>({})
+  const registerCommand = (command: ComfyCommand) => {
+    if (commands.value[command.id]) {
+      console.warn(`Command ${command.id} already registered`)
+    }
+    commands.value[command.id] = command
+  }
+
+  const commandDefinitions: ComfyCommand[] = [
+    {
+      id: 'Comfy.NewBlankWorkflow',
+      icon: 'pi pi-plus',
+      label: 'New Blank Workflow',
+      function: () => {
+        app.workflowManager.setWorkflow(null)
         app.clean()
         app.graph.clear()
-        api.dispatchEvent(new CustomEvent('graphCleared'))
+        app.workflowManager.activeWorkflow.track()
       }
     },
-    'Comfy.ResetView': () => {
-      app.resetView()
+    {
+      id: 'Comfy.OpenWorkflow',
+      icon: 'pi pi-folder-open',
+      label: 'Open Workflow',
+      function: () => {
+        app.ui.loadFile()
+      }
     },
-    'Comfy.OpenClipspace': () => {
-      app['openClipspace']?.()
+    {
+      id: 'Comfy.LoadDefaultWorkflow',
+      icon: 'pi pi-code',
+      label: 'Load Default Workflow',
+      function: async () => {
+        await app.loadGraphData()
+      }
     },
-    'Comfy.RefreshNodeDefinitions': async () => {
-      await app.refreshComboInNodes()
+    {
+      id: 'Comfy.SaveWorkflow',
+      icon: 'pi pi-save',
+      label: 'Save Workflow',
+      function: () => {
+        app.workflowManager.activeWorkflow.save()
+      }
     },
-    'Comfy.Interrupt': async () => {
-      await api.interrupt()
-      useToastStore().add({
-        severity: 'info',
-        summary: 'Interrupted',
-        detail: 'Execution has been interrupted',
-        life: 1000
-      })
+    {
+      id: 'Comfy.SaveWorkflowAs',
+      icon: 'pi pi-save',
+      label: 'Save Workflow As',
+      function: () => {
+        app.workflowManager.activeWorkflow.save(true)
+      }
     },
-    'Comfy.ClearPendingTasks': async () => {
-      await useQueueStore().clear(['queue'])
-      useToastStore().add({
-        severity: 'info',
-        summary: 'Confirmed',
-        detail: 'Pending tasks deleted',
-        life: 3000
-      })
+    {
+      id: 'Comfy.ExportWorkflow',
+      icon: 'pi pi-download',
+      label: 'Export Workflow',
+      function: () => {
+        app.menu.exportWorkflow('workflow', 'workflow')
+      }
     },
-    'Comfy.BrowseTemplates': showTemplateWorkflowsDialog
-  })
+    {
+      id: 'Comfy.ExportWorkflowAPI',
+      icon: 'pi pi-download',
+      label: 'Export Workflow (API Format)',
+      function: () => {
+        app.menu.exportWorkflow('workflow_api', 'output')
+      }
+    },
+    {
+      id: 'Comfy.Undo',
+      icon: 'pi pi-undo',
+      label: 'Undo',
+      function: async () => {
+        await getTracker().undo()
+      }
+    },
+    {
+      id: 'Comfy.Redo',
+      icon: 'pi pi-refresh',
+      label: 'Redo',
+      function: async () => {
+        await getTracker().redo()
+      }
+    },
+    {
+      id: 'Comfy.ClearWorkflow',
+      icon: 'pi pi-trash',
+      label: 'Clear Workflow',
+      function: () => {
+        if (
+          !settingStore.get('Comfy.ComfirmClear') ||
+          confirm('Clear workflow?')
+        ) {
+          app.clean()
+          app.graph.clear()
+          api.dispatchEvent(new CustomEvent('graphCleared'))
+        }
+      }
+    },
+    {
+      id: 'Comfy.ResetView',
+      icon: 'pi pi-expand',
+      label: 'Reset View',
+      function: () => {
+        app.resetView()
+      }
+    },
+    {
+      id: 'Comfy.OpenClipspace',
+      icon: 'pi pi-clipboard',
+      label: 'Clipspace',
+      function: () => {
+        app['openClipspace']?.()
+      }
+    },
+    {
+      id: 'Comfy.RefreshNodeDefinitions',
+      icon: 'pi pi-refresh',
+      label: 'Refresh Node Definitions',
+      function: async () => {
+        await app.refreshComboInNodes()
+      }
+    },
+    {
+      id: 'Comfy.Interrupt',
+      icon: 'pi pi-stop',
+      label: 'Interrupt',
+      function: async () => {
+        await api.interrupt()
+        useToastStore().add({
+          severity: 'info',
+          summary: 'Interrupted',
+          detail: 'Execution has been interrupted',
+          life: 1000
+        })
+      }
+    },
+    {
+      id: 'Comfy.ClearPendingTasks',
+      icon: 'pi pi-stop',
+      label: 'Clear Pending Tasks',
+      function: async () => {
+        await useQueueStore().clear(['queue'])
+        useToastStore().add({
+          severity: 'info',
+          summary: 'Confirmed',
+          detail: 'Pending tasks deleted',
+          life: 3000
+        })
+      }
+    },
+    {
+      id: 'Comfy.BrowseTemplates',
+      icon: 'pi pi-folder-open',
+      label: 'Browse Templates',
+      function: showTemplateWorkflowsDialog
+    }
+  ]
+
+  commandDefinitions.forEach(registerCommand)
+  const getCommandFunction = (command: string) => {
+    return commands.value[command]?.function ?? (() => {})
+  }
 
   const getCommand = (command: string) => {
-    return commands.value[command] ?? (() => {})
+    return commands.value[command]
+  }
+
+  const isRegistered = (command: string) => {
+    return !!commands.value[command]
   }
 
   return {
-    commands,
-    getCommand
+    getCommand,
+    getCommandFunction,
+    registerCommand,
+    isRegistered
   }
 })

src/stores/coreMenuItemStore.ts

@@ -1,93 +0,0 @@
-import { defineStore } from 'pinia'
-import type { MenuItem } from 'primevue/menuitem'
-import { computed } from 'vue'
-import { useCommandStore } from './commandStore'
-
-export const useCoreMenuItemStore = defineStore('coreMenuItem', () => {
-  const commandStore = useCommandStore()
-  const menuItems = computed<MenuItem[]>(() => {
-    return [
-      {
-        label: 'Workflow',
-        items: [
-          {
-            label: 'New',
-            icon: 'pi pi-plus',
-            command: commandStore.commands['Comfy.NewBlankWorkflow']
-          },
-          {
-            separator: true
-          },
-          {
-            label: 'Open',
-            icon: 'pi pi-folder-open',
-            command: commandStore.commands['Comfy.OpenWorkflow']
-          },
-          {
-            label: 'Browse Templates',
-            icon: 'pi pi-th-large',
-            command: commandStore.commands['Comfy.BrowseTemplates']
-          },
-          {
-            separator: true
-          },
-          {
-            label: 'Save',
-            icon: 'pi pi-save',
-            command: commandStore.commands['Comfy.SaveWorkflow']
-          },
-          {
-            label: 'Save As',
-            icon: 'pi pi-save',
-            command: commandStore.commands['Comfy.SaveWorkflowAs']
-          },
-          {
-            label: 'Export',
-            icon: 'pi pi-download',
-            command: commandStore.commands['Comfy.ExportWorkflow']
-          },
-          {
-            label: 'Export (API Format)',
-            icon: 'pi pi-download',
-            command: commandStore.commands['Comfy.ExportWorkflowAPI']
-          }
-        ]
-      },
-      {
-        label: 'Edit',
-        items: [
-          {
-            label: 'Undo',
-            icon: 'pi pi-undo',
-            command: commandStore.commands['Comfy.Undo']
-          },
-          {
-            label: 'Redo',
-            icon: 'pi pi-refresh',
-            command: commandStore.commands['Comfy.Redo']
-          },
-          {
-            separator: true
-          },
-          {
-            label: 'Clear Workflow',
-            icon: 'pi pi-trash',
-            command: commandStore.commands['Comfy.ClearWorkflow']
-          },
-          {
-            separator: true
-          },
-          {
-            label: 'Clipspace',
-            icon: 'pi pi-clipboard',
-            command: commandStore.commands['Comfy.OpenClipspace']
-          }
-        ]
-      }
-    ]
-  })
-
-  return {
-    menuItems
-  }
-})

src/stores/menuItemStore.ts

@@ -0,0 +1,124 @@
+import { defineStore } from 'pinia'
+import type { MenuItem } from 'primevue/menuitem'
+import { ref } from 'vue'
+import { type ComfyCommand, useCommandStore } from './commandStore'
+
+export const useMenuItemStore = defineStore('menuItem', () => {
+  const commandStore = useCommandStore()
+  const menuItems = ref<MenuItem[]>([])
+
+  const registerMenuGroup = (path: string[], items: MenuItem[]) => {
+    let currentLevel = menuItems.value
+
+    // Traverse the path, creating nodes if necessary
+    for (let i = 0; i < path.length; i++) {
+      const segment = path[i]
+      let found = currentLevel.find((item) => item.label === segment)
+
+      if (!found) {
+        // Create a new node if it doesn't exist
+        found = {
+          label: segment,
+          items: []
+        }
+        currentLevel.push(found)
+      }
+
+      // Ensure the found item has an 'items' array
+      if (!found.items) {
+        found.items = []
+      }
+
+      // Move to the next level
+      currentLevel = found.items
+    }
+
+    if (currentLevel.length > 0) {
+      currentLevel.push({
+        separator: true
+      })
+    }
+    // Add the new items to the last level
+    currentLevel.push(...items)
+  }
+
+  const registerCommands = (path: string[], commands: ComfyCommand[]) => {
+    // Register commands that are not already registered
+    for (const command of commands) {
+      if (commandStore.isRegistered(command.id)) {
+        continue
+      }
+      commandStore.registerCommand(command)
+    }
+
+    const items = commands.map(
+      (command) =>
+        ({
+          ...command,
+          command: command.function
+        }) as MenuItem
+    )
+    registerMenuGroup(path, items)
+  }
+
+  const workflowMenuGroup: MenuItem[] = [
+    {
+      label: 'New',
+      icon: 'pi pi-plus',
+      command: commandStore.getCommandFunction('Comfy.NewBlankWorkflow')
+    },
+    {
+      separator: true
+    },
+    {
+      label: 'Open',
+      icon: 'pi pi-folder-open',
+      command: commandStore.getCommandFunction('Comfy.OpenWorkflow')
+    },
+    {
+      label: 'Browse Templates',
+      icon: 'pi pi-th-large',
+      command: commandStore.getCommandFunction('Comfy.BrowseTemplates')
+    },
+    {
+      separator: true
+    },
+    {
+      label: 'Save',
+      icon: 'pi pi-save',
+      command: commandStore.getCommandFunction('Comfy.SaveWorkflow')
+    },
+    {
+      label: 'Save As',
+      icon: 'pi pi-save',
+      command: commandStore.getCommandFunction('Comfy.SaveWorkflowAs')
+    },
+    {
+      label: 'Export',
+      icon: 'pi pi-download',
+      command: commandStore.getCommandFunction('Comfy.ExportWorkflow')
+    },
+    {
+      label: 'Export (API Format)',
+      icon: 'pi pi-download',
+      command: commandStore.getCommandFunction('Comfy.ExportWorkflowAPI')
+    }
+  ]
+
+  registerMenuGroup(['Workflow'], workflowMenuGroup)
+  registerCommands(
+    ['Edit'],
+    [
+      commandStore.getCommand('Comfy.Undo'),
+      commandStore.getCommand('Comfy.Redo')
+    ]
+  )
+  registerCommands(['Edit'], [commandStore.getCommand('Comfy.ClearWorkflow')])
+  registerCommands(['Edit'], [commandStore.getCommand('Comfy.OpenClipspace')])
+
+  return {
+    menuItems,
+    registerMenuGroup,
+    registerCommands
+  }
+})

src/stores/workspaceStateStore.ts

@@ -2,6 +2,7 @@ import type { SidebarTabExtension, ToastManager } from '@/types/extensionTypes'
 import { defineStore } from 'pinia'
 import { useToastStore } from './toastStore'
 import { useQueueSettingsStore } from './queueStore'
+import { useMenuItemStore } from './menuItemStore'
 
 interface WorkspaceState {
   spinner: boolean
@@ -21,6 +22,11 @@ export const useWorkspaceStore = defineStore('workspace', {
     },
     queueSettings() {
       return useQueueSettingsStore()
+    },
+    menu() {
+      return {
+        registerTopbarCommands: useMenuItemStore().registerCommands
+      }
     }
   },
   actions: {