[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

@@ -145,6 +145,35 @@ https://github.com/user-attachments/assets/c142c43f-2fe9-4030-8196-b3bfd4c6977d
 
 ### Node developers API
 
+<details>
+  <summary>v1.3.7: Register commands and keybindings</summary>
+
+  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'
+      }
+    ]
+  })
+```
+
+</details>
+
 <details>
   <summary>v1.3.1: Extension API to register custom topbar menu items</summary>
 

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
+    )
+  })
 })

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' ||

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)
   }
 

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)
   }
 }

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

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
   }
 })

src/stores/coreKeybindings.ts

@@ -0,0 +1,3 @@
+import type { Keybinding } from '@/types/keyBindingTypes'
+
+export const CORE_KEYBINDINGS: Keybinding[] = []

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'
   }
 ]

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<Record<string, KeybindingImpl>>({})
+  /**
+   * User-defined keybindings.
+   */
+  const userKeybindings = ref<Record<string, KeybindingImpl>>({})
+  /**
+   * User-defined keybindings that unset default keybindings.
+   */
+  const userUnsetKeybindings = ref<Record<string, KeybindingImpl>>({})
+
+  const keybindingByKeyCombo = computed<Record<string, KeybindingImpl>>(() => {
+    const result: Record<string, KeybindingImpl> = {
+      ...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<KeybindingImpl[]>(() =>
+    Object.values(keybindingByKeyCombo.value)
+  )
+
+  function getKeybinding(combo: KeyComboImpl) {
+    return keybindingByKeyCombo.value[combo.serialize()]
+  }
+
+  function addKeybinding(
+    target: Ref<Record<string, KeybindingImpl>>,
+    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
+  }
+})

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()
 )

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

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<typeof zKeyCombo>
+export type Keybinding = z.infer<typeof zKeybinding>

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()
+  })
+})