From c83ce863d79c4274aa391179abc9566350330379 Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Sun, 6 Oct 2024 23:31:57 -0400 Subject: [PATCH] Rework command menu extension API (#1144) * Rework command menu API * Update test * Update README * Prevent register other extension's command --- README.md | 19 +++++++- browser_tests/extensionAPI.spec.ts | 40 ++++++++++++++--- src/stores/extensionStore.ts | 2 + src/stores/menuItemStore.ts | 70 +++++++++++++++--------------- src/stores/workspaceStateStore.ts | 6 --- src/types/comfy.d.ts | 16 +++++++ 6 files changed, 104 insertions(+), 49 deletions(-) diff --git a/README.md b/README.md index f7a37a9cf..749af6207 100644 --- a/README.md +++ b/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) diff --git a/browser_tests/extensionAPI.spec.ts b/browser_tests/extensionAPI.spec.ts index 9902d668e..4c028ab60 100644 --- a/browser_tests/extensionAPI.spec.ts +++ b/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'] diff --git a/src/stores/extensionStore.ts b/src/stores/extensionStore.ts index 675f41cdb..e2d71382a 100644 --- a/src/stores/extensionStore.ts +++ b/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. diff --git a/src/stores/menuItemStore.ts b/src/stores/menuItemStore.ts index 21b26fd8c..50a0d6cb4 100644 --- a/src/stores/menuItemStore.ts +++ b/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 } }) diff --git a/src/stores/workspaceStateStore.ts b/src/stores/workspaceStateStore.ts index 747a7dacc..c685bea15 100644 --- a/src/stores/workspaceStateStore.ts +++ b/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 diff --git a/src/types/comfy.d.ts b/src/types/comfy.d.ts index aaeb8f39c..7704e3d60 100644 --- a/src/types/comfy.d.ts +++ b/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