From 77aaa38a92b9912c2a0022e4cc9afeba43025072 Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Wed, 2 Oct 2024 21:38:04 -0400 Subject: [PATCH] [Extension API] Custom commands and keybindings (#1075) * Add keybinding schema * nit * Keybinding store * nit * wip * Bind condition on ComfyCommand * Add settings * nit * Revamp keybinding store * Add tests * Add load keybinding * load extension keybindings * Load extension commands * Handle keybindings * test * Keybinding playwright test * Update README * nit * Remove log * Remove system stats fromt logging.ts --- README.md | 29 +++ browser_tests/extensionAPI.spec.ts | 28 +++ src/extensions/core/keybinds.ts | 19 +- src/scripts/app.ts | 6 + src/scripts/logging.ts | 2 - src/scripts/ui/menu/index.ts | 11 +- src/stores/commandStore.ts | 12 +- src/stores/coreKeybindings.ts | 3 + src/stores/coreSettings.ts | 15 ++ src/stores/keybindingStore.ts | 209 +++++++++++++++++++ src/types/apiTypes.ts | 5 +- src/types/comfy.d.ts | 10 + src/types/keyBindingTypes.ts | 20 ++ tests-ui/tests/store/keybindingStore.test.ts | 122 +++++++++++ 14 files changed, 478 insertions(+), 13 deletions(-) create mode 100644 src/stores/coreKeybindings.ts create mode 100644 src/stores/keybindingStore.ts create mode 100644 src/types/keyBindingTypes.ts create mode 100644 tests-ui/tests/store/keybindingStore.test.ts diff --git a/README.md b/README.md index 0724d2148..f7a37a9cf 100644 --- a/README.md +++ b/README.md @@ -145,6 +145,35 @@ https://github.com/user-attachments/assets/c142c43f-2fe9-4030-8196-b3bfd4c6977d ### Node developers API +
+ v1.3.7: Register commands and keybindings + + Extensions can call the following API to register commands and keybindings. Do + note that keybindings defined in core cannot be overwritten, and some keybindings + are reserved by the browser. + +```js + app.registerExtension({ + name: 'TestExtension1', + commands: [ + { + id: 'TestCommand', + function: () => { + alert('TestCommand') + } + } + ], + keybindings: [ + { + combo: { key: 'k' }, + commandId: 'TestCommand' + } + ] + }) +``` + +
+
v1.3.1: Extension API to register custom topbar menu items diff --git a/browser_tests/extensionAPI.spec.ts b/browser_tests/extensionAPI.spec.ts index 8002bda53..7bbbbac3d 100644 --- a/browser_tests/extensionAPI.spec.ts +++ b/browser_tests/extensionAPI.spec.ts @@ -29,4 +29,32 @@ test.describe('Topbar commands', () => { await comfyPage.menu.topbar.triggerTopbarCommand(['ext', 'foo']) expect(await comfyPage.page.evaluate(() => window['foo'])).toBe(true) }) + + test('Should allow registering keybindings', async ({ comfyPage }) => { + await comfyPage.page.evaluate(() => { + const app = window['app'] + app.registerExtension({ + name: 'TestExtension1', + commands: [ + { + id: 'TestCommand', + function: () => { + window['TestCommand'] = true + } + } + ], + keybindings: [ + { + combo: { key: 'k' }, + commandId: 'TestCommand' + } + ] + }) + }) + + await comfyPage.page.keyboard.press('k') + expect(await comfyPage.page.evaluate(() => window['TestCommand'])).toBe( + true + ) + }) }) diff --git a/src/extensions/core/keybinds.ts b/src/extensions/core/keybinds.ts index d1f4ffd3c..98bf8cbe9 100644 --- a/src/extensions/core/keybinds.ts +++ b/src/extensions/core/keybinds.ts @@ -1,11 +1,26 @@ import { app } from '../../scripts/app' import { api } from '../../scripts/api' import { useToastStore } from '@/stores/toastStore' +import { KeyComboImpl, useKeybindingStore } from '@/stores/keybindingStore' +import { useCommandStore } from '@/stores/commandStore' app.registerExtension({ name: 'Comfy.Keybinds', init() { - const keybindListener = async function (event) { + const keybindListener = async function (event: KeyboardEvent) { + // Ignore keybindings for legacy jest tests as jest tests don't have + // a Vue app instance or pinia stores. + if (!app.vueAppReady) return + + const keyCombo = KeyComboImpl.fromEvent(event) + const keybindingStore = useKeybindingStore() + const commandStore = useCommandStore() + const keybinding = keybindingStore.getKeybinding(keyCombo) + if (keybinding) { + await commandStore.getCommandFunction(keybinding.commandId)() + return + } + const modifierPressed = event.ctrlKey || event.metaKey // Queue prompt using (ctrl or command) + enter @@ -26,7 +41,7 @@ app.registerExtension({ return } - const target = event.composedPath()[0] + const target = event.composedPath()[0] as HTMLElement if ( target.tagName === 'TEXTAREA' || target.tagName === 'INPUT' || diff --git a/src/scripts/app.ts b/src/scripts/app.ts index 2ff895b25..7b8a014c9 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -53,6 +53,8 @@ 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' export const ANIM_PREVIEW_WIDGET = '$$comfy_animation_preview' @@ -2951,6 +2953,10 @@ export class ComfyApp { 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) + } this.extensions.push(extension) } diff --git a/src/scripts/logging.ts b/src/scripts/logging.ts index dd6c3ffc1..b25d65881 100644 --- a/src/scripts/logging.ts +++ b/src/scripts/logging.ts @@ -377,7 +377,5 @@ export class ComfyLogging { if (!this.enabled) return const source = 'ComfyUI.Logging' this.addEntry(source, 'debug', { UserAgent: navigator.userAgent }) - const systemStats = await api.getSystemStats() - this.addEntry(source, 'debug', systemStats) } } diff --git a/src/scripts/ui/menu/index.ts b/src/scripts/ui/menu/index.ts index 66c068189..0e708e443 100644 --- a/src/scripts/ui/menu/index.ts +++ b/src/scripts/ui/menu/index.ts @@ -4,13 +4,10 @@ import { downloadBlob } from '../../utils' import { ComfyButtonGroup } from '../components/buttonGroup' import './menu.css' -// Import ComfyButton to make sure it's shimmed and exported by vite -import { ComfyButton } from '../components/button' -import { ComfySplitButton } from '../components/splitButton' -import { ComfyPopup } from '../components/popup' -console.debug( - `Keep following definitions ${ComfyButton} ${ComfySplitButton} ${ComfyPopup}` -) +// Export to make sure following components are shimmed and exported by vite +export { ComfyButton } from '../components/button' +export { ComfySplitButton } from '../components/splitButton' +export { ComfyPopup } from '../components/popup' export class ComfyAppMenu { app: ComfyApp diff --git a/src/stores/commandStore.ts b/src/stores/commandStore.ts index 2e9e39a23..7cdbc3923 100644 --- a/src/stores/commandStore.ts +++ b/src/stores/commandStore.ts @@ -8,6 +8,7 @@ import { useToastStore } from '@/stores/toastStore' import { showTemplateWorkflowsDialog } from '@/services/dialogService' import { useQueueStore } from './queueStore' import { LiteGraph } from '@comfyorg/litegraph' +import { ComfyExtension } from '@/types/comfy' export interface ComfyCommand { id: string @@ -246,10 +247,19 @@ export const useCommandStore = defineStore('command', () => { return !!commands.value[command] } + const loadExtensionCommands = (extension: ComfyExtension) => { + if (extension.commands) { + for (const command of extension.commands) { + registerCommand(command) + } + } + } + return { getCommand, getCommandFunction, registerCommand, - isRegistered + isRegistered, + loadExtensionCommands } }) diff --git a/src/stores/coreKeybindings.ts b/src/stores/coreKeybindings.ts new file mode 100644 index 000000000..94bb1865a --- /dev/null +++ b/src/stores/coreKeybindings.ts @@ -0,0 +1,3 @@ +import type { Keybinding } from '@/types/keyBindingTypes' + +export const CORE_KEYBINDINGS: Keybinding[] = [] diff --git a/src/stores/coreSettings.ts b/src/stores/coreSettings.ts index fd9f7e8ee..c13689ada 100644 --- a/src/stores/coreSettings.ts +++ b/src/stores/coreSettings.ts @@ -1,3 +1,4 @@ +import type { Keybinding } from '@/types/keyBindingTypes' import { NodeBadgeMode } from '@/types/nodeSource' import { LinkReleaseTriggerAction, @@ -404,5 +405,19 @@ export const CORE_SETTINGS: SettingParams[] = [ type: 'number', defaultValue: 100, versionAdded: '1.3.5' + }, + { + id: 'Comfy.Keybinding.UnsetBindings', + name: 'Keybindings unset by the user', + type: 'hidden', + defaultValue: [] as Keybinding[], + versionAdded: '1.3.7' + }, + { + id: 'Comfy.Keybinding.NewBindings', + name: 'Keybindings set by the user', + type: 'hidden', + defaultValue: [] as Keybinding[], + versionAdded: '1.3.7' } ] diff --git a/src/stores/keybindingStore.ts b/src/stores/keybindingStore.ts new file mode 100644 index 000000000..7c42305ea --- /dev/null +++ b/src/stores/keybindingStore.ts @@ -0,0 +1,209 @@ +import { defineStore } from 'pinia' +import { computed, Ref, ref, toRaw } from 'vue' +import { Keybinding, KeyCombo } from '@/types/keyBindingTypes' +import { useSettingStore } from './settingStore' +import { CORE_KEYBINDINGS } from './coreKeybindings' +import type { ComfyExtension } from '@/types/comfy' + +export class KeybindingImpl implements Keybinding { + commandId: string + combo: KeyComboImpl + + constructor(obj: Keybinding) { + this.commandId = obj.commandId + this.combo = new KeyComboImpl(obj.combo) + } + + equals(other: any): boolean { + if (toRaw(other) instanceof KeybindingImpl) { + return ( + this.commandId === other.commandId && this.combo.equals(other.combo) + ) + } + return false + } +} + +export class KeyComboImpl implements KeyCombo { + key: string + // ctrl or meta(cmd on mac) + ctrl: boolean + alt: boolean + shift: boolean + + constructor(obj: KeyCombo) { + this.key = obj.key + this.ctrl = obj.ctrl ?? false + this.alt = obj.alt ?? false + this.shift = obj.shift ?? false + } + + static fromEvent(event: KeyboardEvent) { + return new KeyComboImpl({ + key: event.key, + ctrl: event.ctrlKey, + alt: event.altKey, + shift: event.shiftKey + }) + } + + equals(other: any): boolean { + if (toRaw(other) instanceof KeyComboImpl) { + return ( + this.key === other.key && + this.ctrl === other.ctrl && + this.alt === other.alt && + this.shift === other.shift + ) + } + return false + } + + serialize(): string { + return `${this.key}:${this.ctrl}:${this.alt}:${this.shift}` + } + + deserialize(serialized: string): KeyComboImpl { + const [key, ctrl, alt, shift] = serialized.split(':') + return new KeyComboImpl({ + key, + ctrl: ctrl === 'true', + alt: alt === 'true', + shift: shift === 'true' + }) + } + + toString(): string { + return `${this.key} + ${this.ctrl ? 'Ctrl' : ''}${this.alt ? 'Alt' : ''}${this.shift ? 'Shift' : ''}` + } +} + +export const useKeybindingStore = defineStore('keybinding', () => { + /** + * Default keybindings provided by core and extensions. + */ + const defaultKeybindings = ref>({}) + /** + * User-defined keybindings. + */ + const userKeybindings = ref>({}) + /** + * User-defined keybindings that unset default keybindings. + */ + const userUnsetKeybindings = ref>({}) + + const keybindingByKeyCombo = computed>(() => { + const result: Record = { + ...defaultKeybindings.value, + ...userKeybindings.value + } + + for (const keybinding of Object.values(userUnsetKeybindings.value)) { + const serializedCombo = keybinding.combo.serialize() + if (result[serializedCombo]?.equals(keybinding)) { + delete result[serializedCombo] + } + } + return result + }) + + const keybindings = computed(() => + Object.values(keybindingByKeyCombo.value) + ) + + function getKeybinding(combo: KeyComboImpl) { + return keybindingByKeyCombo.value[combo.serialize()] + } + + function addKeybinding( + target: Ref>, + keybinding: KeybindingImpl, + { existOk = false }: { existOk: boolean } + ) { + if (!existOk && keybinding.combo.serialize() in target.value) { + throw new Error( + `Keybinding on ${keybinding.combo} already exists on ${ + target.value[keybinding.combo.serialize()].commandId + }` + ) + } + target.value[keybinding.combo.serialize()] = keybinding + } + + function addDefaultKeybinding(keybinding: KeybindingImpl) { + addKeybinding(defaultKeybindings, keybinding, { existOk: false }) + } + + function addUserKeybinding(keybinding: KeybindingImpl) { + const defaultKeybinding = + defaultKeybindings.value[keybinding.combo.serialize()] + if (defaultKeybinding) { + unsetKeybinding(defaultKeybinding) + } + addKeybinding(userKeybindings, keybinding, { existOk: true }) + } + + function unsetKeybinding(keybinding: KeybindingImpl) { + const serializedCombo = keybinding.combo.serialize() + if (!(serializedCombo in keybindingByKeyCombo.value)) { + throw new Error(`Keybinding on ${keybinding.combo} does not exist`) + } + + if (userKeybindings.value[serializedCombo]?.equals(keybinding)) { + delete userKeybindings.value[serializedCombo] + return + } + + if (defaultKeybindings.value[serializedCombo]?.equals(keybinding)) { + addKeybinding(userUnsetKeybindings, keybinding, { existOk: false }) + return + } + + throw new Error(`NOT_REACHED`) + } + + function loadUserKeybindings() { + const settingStore = useSettingStore() + // Unset bindings first as new bindings might conflict with default bindings. + const unsetBindings = settingStore.get('Comfy.Keybinding.UnsetBindings') + for (const keybinding of unsetBindings) { + unsetKeybinding(new KeybindingImpl(keybinding)) + } + const newBindings = settingStore.get('Comfy.Keybinding.NewBindings') + for (const keybinding of newBindings) { + addUserKeybinding(new KeybindingImpl(keybinding)) + } + } + + function loadCoreKeybindings() { + for (const keybinding of CORE_KEYBINDINGS) { + addDefaultKeybinding(new KeybindingImpl(keybinding)) + } + } + + function loadExtensionKeybindings(extension: ComfyExtension) { + if (extension.keybindings) { + for (const keybinding of extension.keybindings) { + try { + addDefaultKeybinding(new KeybindingImpl(keybinding)) + } catch (error) { + console.warn( + `Failed to load keybinding for extension ${extension.name}`, + error + ) + } + } + } + } + + return { + keybindings, + getKeybinding, + addDefaultKeybinding, + addUserKeybinding, + unsetKeybinding, + loadUserKeybindings, + loadCoreKeybindings, + loadExtensionKeybindings + } +}) diff --git a/src/types/apiTypes.ts b/src/types/apiTypes.ts index 08116ac97..4a460f452 100644 --- a/src/types/apiTypes.ts +++ b/src/types/apiTypes.ts @@ -4,6 +4,7 @@ import { fromZodError } from 'zod-validation-error' import { colorPalettesSchema } from './colorPalette' import { LinkReleaseTriggerAction } from './searchBoxTypes' import { NodeBadgeMode } from './nodeSource' +import { zKeybinding } from './keyBindingTypes' const zNodeType = z.string() const zQueueIndex = z.number() @@ -504,7 +505,9 @@ const zSettings = z.record(z.any()).and( 'Comfy.NodeBadge.NodeSourceBadgeMode': zNodeBadgeMode, 'Comfy.NodeBadge.NodeIdBadgeMode': zNodeBadgeMode, 'Comfy.NodeBadge.NodeLifeCycleBadgeMode': zNodeBadgeMode, - 'Comfy.QueueButton.BatchCountLimit': z.number() + 'Comfy.QueueButton.BatchCountLimit': z.number(), + 'Comfy.Keybinding.UnsetBindings': z.array(zKeybinding), + 'Comfy.Keybinding.NewBindings': z.array(zKeybinding) }) .optional() ) diff --git a/src/types/comfy.d.ts b/src/types/comfy.d.ts index 0771ebf73..aaeb8f39c 100644 --- a/src/types/comfy.d.ts +++ b/src/types/comfy.d.ts @@ -1,6 +1,8 @@ import { LGraphNode, IWidget } from './litegraph' import { ComfyApp } from '../scripts/app' import type { ComfyNodeDef } from '@/types/apiTypes' +import type { Keybinding } from '@/types/keyBindingTypes' +import type { ComfyCommand } from '@/stores/commandStore' export type Widgets = Record< string, @@ -17,6 +19,14 @@ export interface ComfyExtension { * The name of the extension */ name: string + /** + * The commands defined by the extension + */ + commands?: ComfyCommand[] + /** + * The keybindings defined by the extension + */ + keybindings?: Keybinding[] /** * Allows any initialisation, e.g. loading resources. Called after the canvas is created but before nodes are added * @param app The ComfyUI app instance diff --git a/src/types/keyBindingTypes.ts b/src/types/keyBindingTypes.ts new file mode 100644 index 000000000..c637e2084 --- /dev/null +++ b/src/types/keyBindingTypes.ts @@ -0,0 +1,20 @@ +import { z } from 'zod' + +// KeyCombo schema +export const zKeyCombo = z.object({ + key: z.string(), + ctrl: z.boolean().optional(), + alt: z.boolean().optional(), + shift: z.boolean().optional(), + meta: z.boolean().optional() +}) + +// Keybinding schema +export const zKeybinding = z.object({ + commandId: z.string(), + combo: zKeyCombo +}) + +// Infer types from schemas +export type KeyCombo = z.infer +export type Keybinding = z.infer diff --git a/tests-ui/tests/store/keybindingStore.test.ts b/tests-ui/tests/store/keybindingStore.test.ts new file mode 100644 index 000000000..97e64a059 --- /dev/null +++ b/tests-ui/tests/store/keybindingStore.test.ts @@ -0,0 +1,122 @@ +import { setActivePinia, createPinia } from 'pinia' +import { + useKeybindingStore, + KeybindingImpl +} from '../../../src/stores/keybindingStore' + +describe('useKeybindingStore', () => { + beforeEach(() => { + setActivePinia(createPinia()) + }) + + it('should add and retrieve default keybindings', () => { + const store = useKeybindingStore() + const keybinding = new KeybindingImpl({ + commandId: 'test.command', + combo: { key: 'A', ctrl: true } + }) + + store.addDefaultKeybinding(keybinding) + + expect(store.keybindings).toHaveLength(1) + expect(store.getKeybinding(keybinding.combo)).toEqual(keybinding) + }) + + it('should add and retrieve user keybindings', () => { + const store = useKeybindingStore() + const keybinding = new KeybindingImpl({ + commandId: 'test.command', + combo: { key: 'B', alt: true } + }) + + store.addUserKeybinding(keybinding) + + expect(store.keybindings).toHaveLength(1) + expect(store.getKeybinding(keybinding.combo)).toEqual(keybinding) + }) + + it('should override default keybindings with user keybindings', () => { + const store = useKeybindingStore() + const defaultKeybinding = new KeybindingImpl({ + commandId: 'test.command1', + combo: { key: 'C', ctrl: true } + }) + const userKeybinding = new KeybindingImpl({ + commandId: 'test.command2', + combo: { key: 'C', ctrl: true } + }) + + store.addDefaultKeybinding(defaultKeybinding) + store.addUserKeybinding(userKeybinding) + + expect(store.keybindings).toHaveLength(1) + expect(store.getKeybinding(userKeybinding.combo)).toEqual(userKeybinding) + }) + + it('should unset user keybindings', () => { + const store = useKeybindingStore() + const keybinding = new KeybindingImpl({ + commandId: 'test.command', + combo: { key: 'D', meta: true } + }) + + store.addUserKeybinding(keybinding) + expect(store.keybindings).toHaveLength(1) + + store.unsetKeybinding(keybinding) + expect(store.keybindings).toHaveLength(0) + }) + + it('should unset default keybindings', () => { + const store = useKeybindingStore() + const keybinding = new KeybindingImpl({ + commandId: 'test.command', + combo: { key: 'E', ctrl: true, alt: true } + }) + + store.addDefaultKeybinding(keybinding) + expect(store.keybindings).toHaveLength(1) + + store.unsetKeybinding(keybinding) + expect(store.keybindings).toHaveLength(0) + }) + + it('should throw an error when adding duplicate default keybindings', () => { + const store = useKeybindingStore() + const keybinding = new KeybindingImpl({ + commandId: 'test.command', + combo: { key: 'F', shift: true } + }) + + store.addDefaultKeybinding(keybinding) + expect(() => store.addDefaultKeybinding(keybinding)).toThrow() + }) + + it('should allow adding duplicate user keybindings', () => { + const store = useKeybindingStore() + const keybinding1 = new KeybindingImpl({ + commandId: 'test.command1', + combo: { key: 'G', ctrl: true } + }) + const keybinding2 = new KeybindingImpl({ + commandId: 'test.command2', + combo: { key: 'G', ctrl: true } + }) + + store.addUserKeybinding(keybinding1) + store.addUserKeybinding(keybinding2) + + expect(store.keybindings).toHaveLength(1) + expect(store.getKeybinding(keybinding2.combo)).toEqual(keybinding2) + }) + + it('should throw an error when unsetting non-existent keybindings', () => { + const store = useKeybindingStore() + const keybinding = new KeybindingImpl({ + commandId: 'test.command', + combo: { key: 'H', alt: true, shift: true } + }) + + expect(() => store.unsetKeybinding(keybinding)).toThrow() + }) +})