Contextmenu Monkeypatch Standardization

src/scripts/app.ts

@@ -7,6 +7,7 @@ import { shallowRef } from 'vue'
 import { useCanvasPositionConversion } from '@/composables/element/useCanvasPositionConversion'
 import { registerProxyWidgets } from '@/core/graph/subgraph/proxyWidget'
 import { st, t } from '@/i18n'
+import type { IContextMenuValue } from '@/lib/litegraph/src/interfaces'
 import {
   LGraph,
   LGraphCanvas,
@@ -1667,6 +1668,28 @@ export class ComfyApp {
     useExtensionService().registerExtension(extension)
   }
 
+  /**
+   * Collects context menu items from all extensions for canvas menus
+   * @param canvas The canvas instance
+   * @returns Array of context menu items from all extensions
+   */
+  collectCanvasMenuItems(canvas: LGraphCanvas): IContextMenuValue[] {
+    return useExtensionService()
+      .invokeExtensions('getCanvasMenuItems', canvas)
+      .flat() as IContextMenuValue[]
+  }
+
+  /**
+   * Collects context menu items from all extensions for node menus
+   * @param node The node being right-clicked
+   * @returns Array of context menu items from all extensions
+   */
+  collectNodeMenuItems(node: LGraphNode): IContextMenuValue[] {
+    return useExtensionService()
+      .invokeExtensions('getNodeMenuItems', node)
+      .flat() as IContextMenuValue[]
+  }
+
   /**
    * Refresh combo list on whole nodes
    */

src/types/comfy.ts

@@ -1,5 +1,8 @@
-import type { Positionable } from '@/lib/litegraph/src/interfaces'
-import type { LGraphNode } from '@/lib/litegraph/src/litegraph'
+import type {
+  IContextMenuValue,
+  Positionable
+} from '@/lib/litegraph/src/interfaces'
+import type { LGraphCanvas, LGraphNode } from '@/lib/litegraph/src/litegraph'
 import type { SettingParams } from '@/platform/settings/types'
 import type { ComfyWorkflowJSON } from '@/platform/workflow/validation/schemas/workflowSchema'
 import type { Keybinding } from '@/schemas/keyBindingSchema'
@@ -106,6 +109,20 @@ export interface ComfyExtension {
    */
   getSelectionToolboxCommands?(selectedItem: Positionable): string[]
 
+  /**
+   * Allows the extension to add context menu items to canvas right-click menus
+   * @param canvas The canvas instance
+   * @returns An array of context menu items to add
+   */
+  getCanvasMenuItems?(canvas: LGraphCanvas): IContextMenuValue[]
+
+  /**
+   * Allows the extension to add context menu items to node right-click menus
+   * @param node The node being right-clicked
+   * @returns An array of context menu items to add
+   */
+  getNodeMenuItems?(node: LGraphNode): IContextMenuValue[]
+
   /**
    * Allows the extension to add additional handling to the node before it is registered with **LGraph**
    * @param nodeType The node class (not an instance)

tests-ui/tests/extensions/contextMenuExtension.test.ts

@@ -0,0 +1,200 @@
+import { createTestingPinia } from '@pinia/testing'
+import { setActivePinia } from 'pinia'
+import { beforeEach, describe, expect, it } from 'vitest'
+
+import type { IContextMenuValue } from '@/lib/litegraph/src/interfaces'
+import type { LGraphCanvas, LGraphNode } from '@/lib/litegraph/src/litegraph'
+import { useExtensionService } from '@/services/extensionService'
+import { useExtensionStore } from '@/stores/extensionStore'
+import type { ComfyExtension } from '@/types/comfy'
+
+describe('Context Menu Extension API', () => {
+  let mockCanvas: LGraphCanvas
+  let mockNode: LGraphNode
+  let extensionStore: ReturnType<typeof useExtensionStore>
+  let extensionService: ReturnType<typeof useExtensionService>
+
+  // Mock menu items
+  const canvasMenuItem1: IContextMenuValue = {
+    content: 'Canvas Item 1',
+    callback: () => {}
+  }
+  const canvasMenuItem2: IContextMenuValue = {
+    content: 'Canvas Item 2',
+    callback: () => {}
+  }
+  const nodeMenuItem1: IContextMenuValue = {
+    content: 'Node Item 1',
+    callback: () => {}
+  }
+  const nodeMenuItem2: IContextMenuValue = {
+    content: 'Node Item 2',
+    callback: () => {}
+  }
+
+  // Mock extensions
+  const createCanvasMenuExtension = (
+    name: string,
+    items: IContextMenuValue[]
+  ): ComfyExtension => ({
+    name,
+    getCanvasMenuItems: () => items
+  })
+
+  const createNodeMenuExtension = (
+    name: string,
+    items: IContextMenuValue[]
+  ): ComfyExtension => ({
+    name,
+    getNodeMenuItems: () => items
+  })
+
+  beforeEach(() => {
+    setActivePinia(createTestingPinia({ stubActions: false }))
+    extensionStore = useExtensionStore()
+    extensionService = useExtensionService()
+
+    mockCanvas = {
+      graph_mouse: [100, 100],
+      selectedItems: new Set()
+    } as unknown as LGraphCanvas
+
+    mockNode = {
+      id: 1,
+      type: 'TestNode',
+      pos: [0, 0]
+    } as unknown as LGraphNode
+  })
+
+  describe('collectCanvasMenuItems', () => {
+    it('should call getCanvasMenuItems and collect into flat array', () => {
+      const ext1 = createCanvasMenuExtension('Extension 1', [canvasMenuItem1])
+      const ext2 = createCanvasMenuExtension('Extension 2', [
+        canvasMenuItem2,
+        { content: 'Item 3', callback: () => {} }
+      ])
+
+      extensionStore.registerExtension(ext1)
+      extensionStore.registerExtension(ext2)
+
+      const items = extensionService
+        .invokeExtensions('getCanvasMenuItems', mockCanvas)
+        .flat() as IContextMenuValue[]
+
+      expect(items).toHaveLength(3)
+      expect(items[0]).toMatchObject({ content: 'Canvas Item 1' })
+      expect(items[1]).toMatchObject({ content: 'Canvas Item 2' })
+      expect(items[2]).toMatchObject({ content: 'Item 3' })
+    })
+
+    it('should support submenus and separators', () => {
+      const extension = createCanvasMenuExtension('Test Extension', [
+        {
+          content: 'Menu with Submenu',
+          has_submenu: true,
+          submenu: {
+            options: [
+              { content: 'Submenu Item 1', callback: () => {} },
+              { content: 'Submenu Item 2', callback: () => {} }
+            ]
+          }
+        },
+        null as unknown as IContextMenuValue,
+        { content: 'After Separator', callback: () => {} }
+      ])
+
+      extensionStore.registerExtension(extension)
+
+      const items = extensionService
+        .invokeExtensions('getCanvasMenuItems', mockCanvas)
+        .flat() as IContextMenuValue[]
+
+      expect(items).toHaveLength(3)
+      expect(items[0].content).toBe('Menu with Submenu')
+      expect(items[0].submenu?.options).toHaveLength(2)
+      expect(items[1]).toBeNull()
+      expect(items[2].content).toBe('After Separator')
+    })
+
+    it('should skip extensions without getCanvasMenuItems', () => {
+      const canvasExtension = createCanvasMenuExtension('Canvas Ext', [
+        canvasMenuItem1
+      ])
+      const extensionWithoutCanvasMenu: ComfyExtension = {
+        name: 'No Canvas Menu'
+      }
+
+      extensionStore.registerExtension(canvasExtension)
+      extensionStore.registerExtension(extensionWithoutCanvasMenu)
+
+      const items = extensionService
+        .invokeExtensions('getCanvasMenuItems', mockCanvas)
+        .flat() as IContextMenuValue[]
+
+      expect(items).toHaveLength(1)
+      expect(items[0].content).toBe('Canvas Item 1')
+    })
+  })
+
+  describe('collectNodeMenuItems', () => {
+    it('should call getNodeMenuItems and collect into flat array', () => {
+      const ext1 = createNodeMenuExtension('Extension 1', [nodeMenuItem1])
+      const ext2 = createNodeMenuExtension('Extension 2', [
+        nodeMenuItem2,
+        { content: 'Item 3', callback: () => {} }
+      ])
+
+      extensionStore.registerExtension(ext1)
+      extensionStore.registerExtension(ext2)
+
+      const items = extensionService
+        .invokeExtensions('getNodeMenuItems', mockNode)
+        .flat() as IContextMenuValue[]
+
+      expect(items).toHaveLength(3)
+      expect(items[0]).toMatchObject({ content: 'Node Item 1' })
+      expect(items[1]).toMatchObject({ content: 'Node Item 2' })
+    })
+
+    it('should support submenus', () => {
+      const extension = createNodeMenuExtension('Submenu Extension', [
+        {
+          content: 'Node Menu with Submenu',
+          has_submenu: true,
+          submenu: {
+            options: [
+              { content: 'Node Submenu 1', callback: () => {} },
+              { content: 'Node Submenu 2', callback: () => {} }
+            ]
+          }
+        }
+      ])
+
+      extensionStore.registerExtension(extension)
+
+      const items = extensionService
+        .invokeExtensions('getNodeMenuItems', mockNode)
+        .flat() as IContextMenuValue[]
+
+      expect(items[0].content).toBe('Node Menu with Submenu')
+      expect(items[0].submenu?.options).toHaveLength(2)
+    })
+
+    it('should skip extensions without getNodeMenuItems', () => {
+      const nodeExtension = createNodeMenuExtension('Node Ext', [nodeMenuItem1])
+      const extensionWithoutNodeMenu: ComfyExtension = {
+        name: 'No Node Menu'
+      }
+
+      extensionStore.registerExtension(nodeExtension)
+      extensionStore.registerExtension(extensionWithoutNodeMenu)
+
+      const items = extensionService
+        .invokeExtensions('getNodeMenuItems', mockNode)
+        .flat() as IContextMenuValue[]
+
+      expect(items).toHaveLength(1)
+      expect(items[0].content).toBe('Node Item 1')
+    })
+  })
+})