Rework command menu extension API (#1144)

* Rework command menu API

* Update test

* Update README

* Prevent register other extension's command

README.md

@@ -180,7 +180,24 @@ https://github.com/user-attachments/assets/c142c43f-2fe9-4030-8196-b3bfd4c6977d
   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)}])
+  app.registerExtension({
+    name: 'TestExtension1',
+    commands: [
+      {
+        id: 'foo-id',
+        label: 'foo',
+        function: () => {
+          alert(1)
+        }
+      }
+    ],
+    menuCommands: [
+      {
+        path: ['ext', 'ext2'],
+        commands: ['foo-id']
+      }
+    ]
+  })
 ```
 
 ![image](https://github.com/user-attachments/assets/ae7b082f-7ce9-4549-a446-4563567102fe)

browser_tests/extensionAPI.spec.ts

@@ -1,4 +1,4 @@
-import { expect } from '@playwright/test'
+import { expect, Locator } from '@playwright/test'
 import { comfyPageFixture as test } from './ComfyPage'
 
 test.describe('Topbar commands', () => {
@@ -12,24 +12,50 @@ test.describe('Topbar commands', () => {
 
   test('Should allow registering topbar commands', async ({ comfyPage }) => {
     await comfyPage.page.evaluate(() => {
-      window['app'].extensionManager.menu.registerTopbarCommands(
-        ['ext'],
-        [
+      window['app'].registerExtension({
+        name: 'TestExtension1',
+        commands: [
           {
             id: 'foo',
-            label: 'foo',
+            label: 'foo-command',
             function: () => {
               window['foo'] = true
             }
           }
+        ],
+        menuCommands: [
+          {
+            path: ['ext'],
+            commands: ['foo']
+          }
         ]
-      )
+      })
     })
 
-    await comfyPage.menu.topbar.triggerTopbarCommand(['ext', 'foo'])
+    await comfyPage.menu.topbar.triggerTopbarCommand(['ext', 'foo-command'])
     expect(await comfyPage.page.evaluate(() => window['foo'])).toBe(true)
   })
 
+  test('Should not allow register command defined in other extension', async ({
+    comfyPage
+  }) => {
+    await comfyPage.registerCommand('foo', () => alert(1))
+    await comfyPage.page.evaluate(() => {
+      window['app'].registerExtension({
+        name: 'TestExtension1',
+        menuCommands: [
+          {
+            path: ['ext'],
+            commands: ['foo']
+          }
+        ]
+      })
+    })
+
+    const menuItem: Locator = await comfyPage.menu.topbar.getMenuItem('ext')
+    expect(await menuItem.count()).toBe(0)
+  })
+
   test('Should allow registering keybindings', async ({ comfyPage }) => {
     await comfyPage.page.evaluate(() => {
       const app = window['app']

src/stores/extensionStore.ts

@@ -5,6 +5,7 @@ import { useKeybindingStore } from './keybindingStore'
 import { useCommandStore } from './commandStore'
 import { useSettingStore } from './settingStore'
 import { app } from '@/scripts/app'
+import { useMenuItemStore } from './menuItemStore'
 
 export const useExtensionStore = defineStore('extension', () => {
   // For legacy reasons, the name uniquely identifies an extension
@@ -45,6 +46,7 @@ export const useExtensionStore = defineStore('extension', () => {
     extensionByName.value[extension.name] = markRaw(extension)
     useKeybindingStore().loadExtensionKeybindings(extension)
     useCommandStore().loadExtensionCommands(extension)
+    useMenuItemStore().loadExtensionMenuCommands(extension)
 
     /*
      * Extensions are currently stored in both extensionStore and app.extensions.

src/stores/menuItemStore.ts

@@ -1,7 +1,8 @@
 import { defineStore } from 'pinia'
 import type { MenuItem } from 'primevue/menuitem'
 import { ref } from 'vue'
-import { type ComfyCommand, useCommandStore } from './commandStore'
+import { useCommandStore } from './commandStore'
+import { ComfyExtension } from '@/types/comfy'
 
 export const useMenuItemStore = defineStore('menuItem', () => {
   const commandStore = useCommandStore()
@@ -42,18 +43,9 @@ export const useMenuItemStore = defineStore('menuItem', () => {
     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
-      // Convert command to commandImpl
-      .map((command) => commandStore.getCommand(command.id))
+  const registerCommands = (path: string[], commandIds: string[]) => {
+    const items = commandIds
+      .map((commandId) => commandStore.getCommand(commandId))
       .map(
         (command) =>
           ({
@@ -67,41 +59,49 @@ export const useMenuItemStore = defineStore('menuItem', () => {
     registerMenuGroup(path, items)
   }
 
-  registerCommands(
-    ['Workflow'],
-    [commandStore.getCommand('Comfy.NewBlankWorkflow')]
+  const loadExtensionMenuCommands = (extension: ComfyExtension) => {
+    if (!extension.menuCommands) {
+      return
+    }
+
+    const extensionCommandIds = new Set(
+      extension.commands?.map((command) => command.id) ?? []
     )
+    extension.menuCommands.forEach((menuCommand) => {
+      const commands = menuCommand.commands.filter((command) =>
+        extensionCommandIds.has(command)
+      )
+      if (commands.length) {
+        registerCommands(menuCommand.path, commands)
+      }
+    })
+  }
+
+  // Core menu commands
+  registerCommands(['Workflow'], ['Comfy.NewBlankWorkflow'])
 
   registerCommands(
     ['Workflow'],
-    [
-      commandStore.getCommand('Comfy.OpenWorkflow'),
-      commandStore.getCommand('Comfy.BrowseTemplates')
-    ]
+    ['Comfy.OpenWorkflow', 'Comfy.BrowseTemplates']
   )
   registerCommands(
     ['Workflow'],
     [
-      commandStore.getCommand('Comfy.SaveWorkflow'),
-      commandStore.getCommand('Comfy.SaveWorkflowAs'),
-      commandStore.getCommand('Comfy.ExportWorkflow'),
-      commandStore.getCommand('Comfy.ExportWorkflowAPI')
+      'Comfy.SaveWorkflow',
+      'Comfy.SaveWorkflowAs',
+      'Comfy.ExportWorkflow',
+      'Comfy.ExportWorkflowAPI'
     ]
   )
 
-  registerCommands(
-    ['Edit'],
-    [
-      commandStore.getCommand('Comfy.Undo'),
-      commandStore.getCommand('Comfy.Redo')
-    ]
-  )
-  registerCommands(['Edit'], [commandStore.getCommand('Comfy.ClearWorkflow')])
-  registerCommands(['Edit'], [commandStore.getCommand('Comfy.OpenClipspace')])
+  registerCommands(['Edit'], ['Comfy.Undo', 'Comfy.Redo'])
+  registerCommands(['Edit'], ['Comfy.ClearWorkflow'])
+  registerCommands(['Edit'], ['Comfy.OpenClipspace'])
 
   return {
     menuItems,
     registerMenuGroup,
-    registerCommands
+    registerCommands,
+    loadExtensionMenuCommands
   }
 })

src/stores/workspaceStateStore.ts

@@ -2,7 +2,6 @@ import type { SidebarTabExtension, ToastManager } from '@/types/extensionTypes'
 import { defineStore } from 'pinia'
 import { useToastStore } from './toastStore'
 import { useQueueSettingsStore } from './queueStore'
-import { useMenuItemStore } from './menuItemStore'
 import { useCommandStore } from './commandStore'
 import { useSidebarTabStore } from './workspace/sidebarTabStore'
 
@@ -21,11 +20,6 @@ export const useWorkspaceStore = defineStore('workspace', {
     queueSettings() {
       return useQueueSettingsStore()
     },
-    menu() {
-      return {
-        registerTopbarCommands: useMenuItemStore().registerCommands
-      }
-    },
     command() {
       return {
         execute: useCommandStore().execute

src/types/comfy.d.ts

@@ -14,6 +14,18 @@ export type Widgets = Record<
   ) => { widget?: IWidget; minWidth?: number; minHeight?: number }
 >
 
+export type MenuCommandGroup = {
+  /**
+   * The path to the menu group.
+   */
+  path: string[]
+  /**
+   * Command ids.
+   * Note: Commands must be defined in `commands` array in the extension.
+   */
+  commands: string[]
+}
+
 export interface ComfyExtension {
   /**
    * The name of the extension
@@ -27,6 +39,10 @@ export interface ComfyExtension {
    * The keybindings defined by the extension
    */
   keybindings?: Keybinding[]
+  /**
+   * Menu commands to add to the menu bar
+   */
+  menuCommands?: MenuCommandGroup[]
   /**
    * Allows any initialisation, e.g. loading resources. Called after the canvas is created but before nodes are added
    * @param app The ComfyUI app instance