diff --git a/src/scripts/app.ts b/src/scripts/app.ts index 8522aeb2c3..c962fd1eb1 100644 --- a/src/scripts/app.ts +++ b/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 */ diff --git a/src/types/comfy.ts b/src/types/comfy.ts index 0ce8d89dc3..184aa7a4e6 100644 --- a/src/types/comfy.ts +++ b/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) diff --git a/tests-ui/tests/extensions/contextMenuExtension.test.ts b/tests-ui/tests/extensions/contextMenuExtension.test.ts new file mode 100644 index 0000000000..e6766b0da3 --- /dev/null +++ b/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 + let extensionService: ReturnType + + // 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') + }) + }) +})