From 5b37fc59e7f13fb060aff82777c6955efafd0cc8 Mon Sep 17 00:00:00 2001 From: Johnpaul Date: Thu, 9 Oct 2025 01:33:29 +0100 Subject: [PATCH] feat(contextMenu): add extension API for context menu items Introduces a new extension API that allows extensions to provide context menu items directly, without monkey-patching. This provides a clean, type-safe way for extensions to add menu items. **New API methods:** - `getCanvasMenuItems(canvas)`: Add items to canvas right-click menus - `getNodeMenuItems(node)`: Add items to node right-click menus **Implementation:** - Added TypeScript interfaces in `src/types/comfy.ts` - Added collection methods in `ComfyApp` class - Comprehensive test coverage for the new API --- src/scripts/app.ts | 51 +++ src/types/comfy.ts | 21 +- .../extensions/contextMenuExtension.test.ts | 303 ++++++++++++++++++ .../contextMenuExtensionName.test.ts | 65 ++++ 4 files changed, 438 insertions(+), 2 deletions(-) create mode 100644 tests-ui/tests/extensions/contextMenuExtension.test.ts create mode 100644 tests-ui/tests/extensions/contextMenuExtensionName.test.ts diff --git a/src/scripts/app.ts b/src/scripts/app.ts index 8522aeb2c..6582480e4 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,56 @@ 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[] { + const items: IContextMenuValue[] = [] + + for (const ext of this.extensions) { + if (ext.getCanvasMenuItems) { + try { + const extItems = ext.getCanvasMenuItems(canvas) + items.push(...extItems) + } catch (error) { + console.error( + `[Context Menu] Extension "${ext.name}" failed to provide canvas menu items:`, + error + ) + } + } + } + + return items + } + + /** + * 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[] { + const items: IContextMenuValue[] = [] + + for (const ext of this.extensions) { + if (ext.getNodeMenuItems) { + try { + const extItems = ext.getNodeMenuItems(node) + items.push(...extItems) + } catch (error) { + console.error( + `[Context Menu] Extension "${ext.name}" failed to provide node menu items:`, + error + ) + } + } + } + + return items + } + /** * Refresh combo list on whole nodes */ diff --git a/src/types/comfy.ts b/src/types/comfy.ts index 0ce8d89dc..184aa7a4e 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 000000000..65e279be8 --- /dev/null +++ b/tests-ui/tests/extensions/contextMenuExtension.test.ts @@ -0,0 +1,303 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +import type { IContextMenuValue } from '@/lib/litegraph/src/interfaces' +import type { LGraphCanvas, LGraphNode } from '@/lib/litegraph/src/litegraph' +import { app } from '@/scripts/app' +import type { ComfyExtension } from '@/types/comfy' + +describe('Context Menu Extension API', () => { + let mockCanvas: LGraphCanvas + let mockNode: LGraphNode + + beforeEach(() => { + // Create mock canvas + mockCanvas = { + graph_mouse: [100, 100], + selectedItems: new Set() + } as unknown as LGraphCanvas + + // Create mock node + mockNode = { + id: 1, + type: 'TestNode', + pos: [0, 0] + } as unknown as LGraphNode + + // Clear console + vi.spyOn(console, 'error').mockImplementation(() => {}) + }) + + afterEach(() => { + vi.restoreAllMocks() + }) + + describe('collectCanvasMenuItems', () => { + it('should collect items from extensions with getCanvasMenuItems', () => { + const mockExtension1: ComfyExtension = { + name: 'Test Extension 1', + getCanvasMenuItems: (_canvas: LGraphCanvas) => { + return [ + { content: 'Extension 1 Item', callback: () => {} } + ] as IContextMenuValue[] + } + } + + const mockExtension2: ComfyExtension = { + name: 'Test Extension 2', + getCanvasMenuItems: (_canvas: LGraphCanvas) => { + return [ + { content: 'Extension 2 Item', callback: () => {} } + ] as IContextMenuValue[] + } + } + + // Mock extensions array + vi.spyOn(app, 'extensions', 'get').mockReturnValue([ + mockExtension1, + mockExtension2 + ]) + + const items = app.collectCanvasMenuItems(mockCanvas) + + expect(items).toHaveLength(2) + expect(items[0]).toMatchObject({ content: 'Extension 1 Item' }) + expect(items[1]).toMatchObject({ content: 'Extension 2 Item' }) + }) + + it('should skip extensions without getCanvasMenuItems', () => { + const mockExtension1: ComfyExtension = { + name: 'Test Extension 1', + getCanvasMenuItems: () => { + return [{ content: 'Item 1', callback: () => {} }] + } + } + + const mockExtension2: ComfyExtension = { + name: 'Test Extension 2' + // No getCanvasMenuItems + } + + vi.spyOn(app, 'extensions', 'get').mockReturnValue([ + mockExtension1, + mockExtension2 + ]) + + const items = app.collectCanvasMenuItems(mockCanvas) + + expect(items).toHaveLength(1) + expect(items[0]).toMatchObject({ content: 'Item 1' }) + }) + + it('should handle errors in extension gracefully', () => { + const errorSpy = vi.spyOn(console, 'error').mockImplementation(() => {}) + + const mockExtension1: ComfyExtension = { + name: 'Failing Extension', + getCanvasMenuItems: () => { + throw new Error('Extension error') + } + } + + const mockExtension2: ComfyExtension = { + name: 'Working Extension', + getCanvasMenuItems: () => { + return [{ content: 'Working Item', callback: () => {} }] + } + } + + vi.spyOn(app, 'extensions', 'get').mockReturnValue([ + mockExtension1, + mockExtension2 + ]) + + const items = app.collectCanvasMenuItems(mockCanvas) + + // Should have logged error with extension name + expect(errorSpy).toHaveBeenCalledWith( + expect.stringContaining('[Context Menu]'), + expect.any(Error) + ) + expect(errorSpy).toHaveBeenCalledWith( + expect.stringContaining('"Failing Extension"'), + expect.any(Error) + ) + + // Should still return items from working extension + expect(items).toHaveLength(1) + expect(items[0]).toMatchObject({ content: 'Working Item' }) + }) + + it('should pass canvas to extension method', () => { + const canvasCapture = vi.fn() + + const mockExtension: ComfyExtension = { + name: 'Test Extension', + getCanvasMenuItems: (canvas: LGraphCanvas) => { + canvasCapture(canvas) + return [] + } + } + + vi.spyOn(app, 'extensions', 'get').mockReturnValue([mockExtension]) + + app.collectCanvasMenuItems(mockCanvas) + + expect(canvasCapture).toHaveBeenCalledWith(mockCanvas) + }) + + it('should return empty array when no extensions', () => { + vi.spyOn(app, 'extensions', 'get').mockReturnValue([]) + + const items = app.collectCanvasMenuItems(mockCanvas) + + expect(items).toHaveLength(0) + }) + }) + + describe('collectNodeMenuItems', () => { + it('should collect items from extensions with getNodeMenuItems', () => { + const mockExtension1: ComfyExtension = { + name: 'Test Extension 1', + getNodeMenuItems: (_node: LGraphNode) => { + return [ + { content: 'Node Item 1', callback: () => {} } + ] as IContextMenuValue[] + } + } + + const mockExtension2: ComfyExtension = { + name: 'Test Extension 2', + getNodeMenuItems: (_node: LGraphNode) => { + return [ + { content: 'Node Item 2', callback: () => {} } + ] as IContextMenuValue[] + } + } + + vi.spyOn(app, 'extensions', 'get').mockReturnValue([ + mockExtension1, + mockExtension2 + ]) + + const items = app.collectNodeMenuItems(mockNode) + + expect(items).toHaveLength(2) + expect(items[0]).toMatchObject({ content: 'Node Item 1' }) + expect(items[1]).toMatchObject({ content: 'Node Item 2' }) + }) + + it('should skip extensions without getNodeMenuItems', () => { + const mockExtension1: ComfyExtension = { + name: 'Test Extension 1', + getNodeMenuItems: () => { + return [{ content: 'Node Item', callback: () => {} }] + } + } + + const mockExtension2: ComfyExtension = { + name: 'Test Extension 2' + // No getNodeMenuItems + } + + vi.spyOn(app, 'extensions', 'get').mockReturnValue([ + mockExtension1, + mockExtension2 + ]) + + const items = app.collectNodeMenuItems(mockNode) + + expect(items).toHaveLength(1) + expect(items[0]).toMatchObject({ content: 'Node Item' }) + }) + + it('should handle errors in extension gracefully', () => { + const errorSpy = vi.spyOn(console, 'error').mockImplementation(() => {}) + + const mockExtension1: ComfyExtension = { + name: 'Failing Extension', + getNodeMenuItems: () => { + throw new Error('Extension error') + } + } + + const mockExtension2: ComfyExtension = { + name: 'Working Extension', + getNodeMenuItems: () => { + return [{ content: 'Working Node Item', callback: () => {} }] + } + } + + vi.spyOn(app, 'extensions', 'get').mockReturnValue([ + mockExtension1, + mockExtension2 + ]) + + const items = app.collectNodeMenuItems(mockNode) + + // Should have logged error with extension name + expect(errorSpy).toHaveBeenCalledWith( + expect.stringContaining('[Context Menu]'), + expect.any(Error) + ) + expect(errorSpy).toHaveBeenCalledWith( + expect.stringContaining('"Failing Extension"'), + expect.any(Error) + ) + + // Should still return items from working extension + expect(items).toHaveLength(1) + expect(items[0]).toMatchObject({ content: 'Working Node Item' }) + }) + + it('should pass node to extension method', () => { + const nodeCapture = vi.fn() + + const mockExtension: ComfyExtension = { + name: 'Test Extension', + getNodeMenuItems: (node: LGraphNode) => { + nodeCapture(node) + return [] + } + } + + vi.spyOn(app, 'extensions', 'get').mockReturnValue([mockExtension]) + + app.collectNodeMenuItems(mockNode) + + expect(nodeCapture).toHaveBeenCalledWith(mockNode) + }) + + it('should return empty array when no extensions', () => { + vi.spyOn(app, 'extensions', 'get').mockReturnValue([]) + + const items = app.collectNodeMenuItems(mockNode) + + expect(items).toHaveLength(0) + }) + }) + + describe('integration with both methods', () => { + it('should work when extension provides both methods', () => { + const mockExtension: ComfyExtension = { + name: 'Full Extension', + getCanvasMenuItems: () => { + return [{ content: 'Canvas Item', callback: () => {} }] + }, + getNodeMenuItems: () => { + return [{ content: 'Node Item', callback: () => {} }] + } + } + + vi.spyOn(app, 'extensions', 'get').mockReturnValue([mockExtension]) + + const canvasItems = app.collectCanvasMenuItems(mockCanvas) + const nodeItems = app.collectNodeMenuItems(mockNode) + + expect(canvasItems).toHaveLength(1) + expect(canvasItems[0]).toMatchObject({ content: 'Canvas Item' }) + + expect(nodeItems).toHaveLength(1) + expect(nodeItems[0]).toMatchObject({ content: 'Node Item' }) + }) + }) +}) diff --git a/tests-ui/tests/extensions/contextMenuExtensionName.test.ts b/tests-ui/tests/extensions/contextMenuExtensionName.test.ts new file mode 100644 index 000000000..485df79de --- /dev/null +++ b/tests-ui/tests/extensions/contextMenuExtensionName.test.ts @@ -0,0 +1,65 @@ +import { describe, expect, it, vi } from 'vitest' + +import { legacyMenuCompat } from '@/lib/litegraph/src/contextMenuCompat' +import { LGraphCanvas } from '@/lib/litegraph/src/litegraph' + +/** + * Test that demonstrates the extension name appearing in deprecation warnings + */ +describe('Context Menu Extension Name in Warnings', () => { + it('should include extension name in deprecation warning', () => { + const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {}) + + // Install compatibility layer + legacyMenuCompat.install(LGraphCanvas.prototype, 'getCanvasMenuOptions') + + // Simulate what happens during extension setup + legacyMenuCompat.setCurrentExtension('MyCustomExtension') + + // Extension monkey-patches the method + const original = LGraphCanvas.prototype.getCanvasMenuOptions + LGraphCanvas.prototype.getCanvasMenuOptions = function (...args: any[]) { + const items = (original as any).apply(this, args) + items.push({ content: 'My Custom Menu Item', callback: () => {} }) + return items + } + + // Clear extension (happens after setup completes) + legacyMenuCompat.setCurrentExtension(null) + + // Verify the warning includes the extension name + expect(warnSpy).toHaveBeenCalled() + const warningMessage = warnSpy.mock.calls[0][0] + + expect(warningMessage).toContain('[DEPRECATED]') + expect(warningMessage).toContain('getCanvasMenuOptions') + expect(warningMessage).toContain('"MyCustomExtension"') + + vi.restoreAllMocks() + }) + + it('should not include extension name if not set', () => { + const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {}) + + // Install compatibility layer + legacyMenuCompat.install(LGraphCanvas.prototype, 'getNodeMenuOptions') + + // Extension monkey-patches without setting current extension + const original = LGraphCanvas.prototype.getNodeMenuOptions + LGraphCanvas.prototype.getNodeMenuOptions = function (...args: any[]) { + const items = (original as any).apply(this, args) + items.push({ content: 'My Node Menu Item', callback: () => {} }) + return items + } + + // Verify the warning does NOT include extension info + expect(warnSpy).toHaveBeenCalled() + const warningMessage = warnSpy.mock.calls[0][0] + + expect(warningMessage).toContain('[DEPRECATED]') + expect(warningMessage).toContain('getNodeMenuOptions') + expect(warningMessage).not.toContain('Extension:') + + vi.restoreAllMocks() + }) +})