diff --git a/src/composables/graph/contextMenuConverter.test.ts b/src/composables/graph/contextMenuConverter.test.ts new file mode 100644 index 0000000000..cc157f220a --- /dev/null +++ b/src/composables/graph/contextMenuConverter.test.ts @@ -0,0 +1,190 @@ +import { describe, it, expect } from 'vitest' + +import type { MenuOption } from './useMoreOptionsMenu' +import { + buildStructuredMenu, + convertContextMenuToOptions +} from './contextMenuConverter' + +describe('contextMenuConverter', () => { + describe('buildStructuredMenu', () => { + it('should order core items before extension items', () => { + const options: MenuOption[] = [ + { label: 'Custom Extension Item', source: 'litegraph' }, + { label: 'Copy', source: 'vue' }, + { label: 'Rename', source: 'vue' } + ] + + const result = buildStructuredMenu(options) + + // Core items (Rename, Copy) should come before extension items + const renameIndex = result.findIndex((opt) => opt.label === 'Rename') + const copyIndex = result.findIndex((opt) => opt.label === 'Copy') + const extensionIndex = result.findIndex( + (opt) => opt.label === 'Custom Extension Item' + ) + + expect(renameIndex).toBeLessThan(extensionIndex) + expect(copyIndex).toBeLessThan(extensionIndex) + }) + + it('should add Extensions category label before extension items', () => { + const options: MenuOption[] = [ + { label: 'Copy', source: 'vue' }, + { label: 'My Custom Extension', source: 'litegraph' } + ] + + const result = buildStructuredMenu(options) + + const extensionsLabel = result.find( + (opt) => opt.label === 'Extensions' && opt.type === 'category' + ) + expect(extensionsLabel).toBeDefined() + expect(extensionsLabel?.disabled).toBe(true) + }) + + it('should place Delete at the very end', () => { + const options: MenuOption[] = [ + { label: 'Delete', action: () => {}, source: 'vue' }, + { label: 'Copy', source: 'vue' }, + { label: 'Rename', source: 'vue' } + ] + + const result = buildStructuredMenu(options) + + const lastNonDivider = [...result] + .reverse() + .find((opt) => opt.type !== 'divider') + expect(lastNonDivider?.label).toBe('Delete') + }) + + it('should deduplicate items with same label, preferring vue source', () => { + const options: MenuOption[] = [ + { label: 'Copy', action: () => {}, source: 'litegraph' }, + { label: 'Copy', action: () => {}, source: 'vue' } + ] + + const result = buildStructuredMenu(options) + + const copyItems = result.filter((opt) => opt.label === 'Copy') + expect(copyItems).toHaveLength(1) + expect(copyItems[0].source).toBe('vue') + }) + + it('should preserve dividers between sections', () => { + const options: MenuOption[] = [ + { label: 'Rename', source: 'vue' }, + { label: 'Copy', source: 'vue' }, + { label: 'Pin', source: 'vue' } + ] + + const result = buildStructuredMenu(options) + + const dividers = result.filter((opt) => opt.type === 'divider') + expect(dividers.length).toBeGreaterThan(0) + }) + + it('should handle empty input', () => { + const result = buildStructuredMenu([]) + expect(result).toEqual([]) + }) + + it('should handle only dividers', () => { + const options: MenuOption[] = [{ type: 'divider' }, { type: 'divider' }] + + const result = buildStructuredMenu(options) + + // Should be empty since dividers are filtered initially + expect(result).toEqual([]) + }) + + it('should recognize Remove as equivalent to Delete', () => { + const options: MenuOption[] = [ + { label: 'Remove', action: () => {}, source: 'vue' }, + { label: 'Copy', source: 'vue' } + ] + + const result = buildStructuredMenu(options) + + // Remove should be placed at the end like Delete + const lastNonDivider = [...result] + .reverse() + .find((opt) => opt.type !== 'divider') + expect(lastNonDivider?.label).toBe('Remove') + }) + + it('should group core items in correct section order', () => { + const options: MenuOption[] = [ + { label: 'Color', source: 'vue' }, + { label: 'Node Info', source: 'vue' }, + { label: 'Pin', source: 'vue' }, + { label: 'Rename', source: 'vue' } + ] + + const result = buildStructuredMenu(options) + + // Get indices of items (excluding dividers and categories) + const getIndex = (label: string) => + result.findIndex((opt) => opt.label === label) + + // Rename (section 1) should come before Pin (section 2) + expect(getIndex('Rename')).toBeLessThan(getIndex('Pin')) + // Pin (section 2) should come before Node Info (section 4) + expect(getIndex('Pin')).toBeLessThan(getIndex('Node Info')) + // Node Info (section 4) should come before or with Color (section 4) + expect(getIndex('Node Info')).toBeLessThanOrEqual(getIndex('Color')) + }) + }) + + describe('convertContextMenuToOptions', () => { + it('should convert empty array to empty result', () => { + const result = convertContextMenuToOptions([]) + expect(result).toEqual([]) + }) + + it('should convert null items to dividers', () => { + const result = convertContextMenuToOptions([null], undefined, false) + expect(result).toHaveLength(1) + expect(result[0].type).toBe('divider') + }) + + it('should skip blacklisted items like Properties', () => { + const items = [{ content: 'Properties', callback: () => {} }] + const result = convertContextMenuToOptions(items, undefined, false) + expect(result.find((opt) => opt.label === 'Properties')).toBeUndefined() + }) + + it('should convert basic menu items with content', () => { + const items = [{ content: 'Test Item', callback: () => {} }] + const result = convertContextMenuToOptions(items, undefined, false) + expect(result).toHaveLength(1) + expect(result[0].label).toBe('Test Item') + }) + + it('should mark items as litegraph source', () => { + const items = [{ content: 'Test Item', callback: () => {} }] + const result = convertContextMenuToOptions(items, undefined, false) + expect(result[0].source).toBe('litegraph') + }) + + it('should pass through disabled state', () => { + const items = [{ content: 'Disabled Item', disabled: true }] + const result = convertContextMenuToOptions(items, undefined, false) + expect(result[0].disabled).toBe(true) + }) + + it('should apply structuring by default', () => { + const items = [ + { content: 'Copy', callback: () => {} }, + { content: 'Custom Extension', callback: () => {} } + ] + const result = convertContextMenuToOptions(items) + + // With structuring, there should be Extensions category + const hasExtensionsCategory = result.some( + (opt) => opt.label === 'Extensions' && opt.type === 'category' + ) + expect(hasExtensionsCategory).toBe(true) + }) + }) +}) diff --git a/src/composables/graph/contextMenuConverter.ts b/src/composables/graph/contextMenuConverter.ts new file mode 100644 index 0000000000..3395e18be8 --- /dev/null +++ b/src/composables/graph/contextMenuConverter.ts @@ -0,0 +1,620 @@ +import { default as DOMPurify } from 'dompurify' + +import { LiteGraph } from '@/lib/litegraph/src/litegraph' +import type { + IContextMenuValue, + LGraphNode, + IContextMenuOptions, + ContextMenu +} from '@/lib/litegraph/src/litegraph' + +import type { MenuOption, SubMenuOption } from './useMoreOptionsMenu' +import type { ContextMenuDivElement } from '@/lib/litegraph/src/interfaces' + +/** + * Hard blacklist - items that should NEVER be included + */ +const HARD_BLACKLIST = new Set([ + 'Properties', // Never include Properties submenu + 'Colors', // Use singular "Color" instead + 'Shapes', // Use singular "Shape" instead + 'Title', + 'Mode', + 'Properties Panel', + 'Copy (Clipspace)' +]) + +/** + * Core menu items - items that should appear in the main menu, not under Extensions + * Includes both LiteGraph base menu items and ComfyUI built-in functionality + */ +const CORE_MENU_ITEMS = new Set([ + // Basic operations + 'Rename', + 'Copy', + 'Duplicate', + 'Clone', + // Node state operations + 'Run Branch', + 'Pin', + 'Unpin', + 'Bypass', + 'Remove Bypass', + 'Mute', + // Structure operations + 'Convert to Subgraph', + 'Frame selection', + 'Minimize Node', + 'Expand', + 'Collapse', + // Info and adjustments + 'Node Info', + 'Resize', + 'Title', + 'Properties Panel', + 'Adjust Size', + // Visual + 'Color', + 'Colors', + 'Shape', + 'Shapes', + 'Mode', + // Built-in node operations (node-specific) + 'Open Image', + 'Copy Image', + 'Save Image', + 'Open in Mask Editor', + 'Edit Subgraph Widgets', + 'Unpack Subgraph', + 'Copy (Clipspace)', + 'Paste (Clipspace)', + // Selection and alignment + 'Align Selected To', + 'Distribute Nodes', + // Deletion + 'Delete', + 'Remove', + // LiteGraph base items + 'Show Advanced', + 'Hide Advanced' +]) + +/** + * Normalize menu item label for duplicate detection + * Handles variations like Colors/Color, Shapes/Shape, Pin/Unpin, Remove/Delete + */ +function normalizeLabel(label: string): string { + return label + .toLowerCase() + .replace(/^un/, '') // Remove 'un' prefix (Unpin -> Pin) + .trim() +} + +/** + * Check if a similar menu item already exists in the results + * Returns true if an item with the same normalized label exists + */ +function isDuplicateItem(label: string, existingItems: MenuOption[]): boolean { + const normalizedLabel = normalizeLabel(label) + + // Map of equivalent items + const equivalents: Record = { + color: ['color', 'colors'], + shape: ['shape', 'shapes'], + pin: ['pin', 'unpin'], + delete: ['remove', 'delete'], + duplicate: ['clone', 'duplicate'] + } + + return existingItems.some((item) => { + if (!item.label) return false + + const existingNormalized = normalizeLabel(item.label) + + // Check direct match + if (existingNormalized === normalizedLabel) return true + + // Check if they're in the same equivalence group + for (const values of Object.values(equivalents)) { + if ( + values.includes(normalizedLabel) && + values.includes(existingNormalized) + ) { + return true + } + } + + return false + }) +} + +/** + * Check if a menu item is a core menu item (not an extension) + * Core items include LiteGraph base items and ComfyUI built-in functionality + */ +function isCoreMenuItem(label: string): boolean { + return CORE_MENU_ITEMS.has(label) +} + +/** + * Filter out duplicate menu items based on label + * Gives precedence to Vue hardcoded options over LiteGraph options + */ +function removeDuplicateMenuOptions(options: MenuOption[]): MenuOption[] { + // Group items by label + const itemsByLabel = new Map() + const itemsWithoutLabel: MenuOption[] = [] + + for (const opt of options) { + // Always keep dividers and category items + if (opt.type === 'divider' || opt.type === 'category') { + itemsWithoutLabel.push(opt) + continue + } + + // Items without labels are kept as-is + if (!opt.label) { + itemsWithoutLabel.push(opt) + continue + } + + // Group by label + if (!itemsByLabel.has(opt.label)) { + itemsByLabel.set(opt.label, []) + } + itemsByLabel.get(opt.label)!.push(opt) + } + + // Select best item for each label (prefer vue over litegraph) + const result: MenuOption[] = [] + const seenLabels = new Set() + + for (const opt of options) { + // Add non-labeled items in original order + if (opt.type === 'divider' || opt.type === 'category' || !opt.label) { + if (itemsWithoutLabel.includes(opt)) { + result.push(opt) + const idx = itemsWithoutLabel.indexOf(opt) + itemsWithoutLabel.splice(idx, 1) + } + continue + } + + // Skip if we already processed this label + if (seenLabels.has(opt.label)) { + continue + } + seenLabels.add(opt.label) + + // Get all items with this label + const duplicates = itemsByLabel.get(opt.label)! + + // If only one item, add it + if (duplicates.length === 1) { + result.push(duplicates[0]) + continue + } + + // Multiple items: prefer vue source over litegraph + const vueItem = duplicates.find((item) => item.source === 'vue') + if (vueItem) { + result.push(vueItem) + } else { + // No vue item, just take the first one + result.push(duplicates[0]) + } + } + + return result +} + +/** + * Order groups for menu items - defines the display order of sections + */ +const MENU_ORDER: string[] = [ + // Section 1: Basic operations + 'Rename', + 'Copy', + 'Duplicate', + // Section 2: Node actions + 'Run Branch', + 'Pin', + 'Unpin', + 'Bypass', + 'Remove Bypass', + 'Mute', + // Section 3: Structure operations + 'Convert to Subgraph', + 'Frame selection', + 'Minimize Node', + 'Expand', + 'Collapse', + 'Resize', + 'Clone', + // Section 4: Node properties + 'Node Info', + 'Color', + // Section 5: Node-specific operations + 'Open in Mask Editor', + 'Open Image', + 'Copy Image', + 'Save Image', + 'Copy (Clipspace)', + 'Paste (Clipspace)', + // Fallback for other core items + 'Convert to Group Node (Deprecated)' +] + +/** + * Get the order index for a menu item (lower = earlier in menu) + */ +function getMenuItemOrder(label: string): number { + const index = MENU_ORDER.indexOf(label) + return index === -1 ? 999 : index +} + +/** + * Build structured menu with core items first, then extensions under a labeled section + * Ensures Delete always appears at the bottom + */ +export function buildStructuredMenu(options: MenuOption[]): MenuOption[] { + // First, remove duplicates (giving precedence to Vue hardcoded options) + const deduplicated = removeDuplicateMenuOptions(options) + const coreItemsMap = new Map() + const extensionItems: MenuOption[] = [] + let deleteItem: MenuOption | undefined + + // Separate items into core and extension categories + for (const option of deduplicated) { + // Skip dividers for now - we'll add them between sections later + if (option.type === 'divider') { + continue + } + + // Skip category labels (they'll be added separately) + if (option.type === 'category') { + continue + } + + // Check if this is the Delete/Remove item - save it for the end + const isDeleteItem = option.label === 'Delete' || option.label === 'Remove' + if (isDeleteItem && !option.hasSubmenu) { + deleteItem = option + continue + } + + // Categorize based on label + if (option.label && isCoreMenuItem(option.label)) { + coreItemsMap.set(option.label, option) + } else { + extensionItems.push(option) + } + } + // Build ordered core items based on MENU_ORDER + const orderedCoreItems: MenuOption[] = [] + const coreLabels = Array.from(coreItemsMap.keys()) + coreLabels.sort((a, b) => getMenuItemOrder(a) - getMenuItemOrder(b)) + + // Section boundaries based on MENU_ORDER indices + // Section 1: 0-2 (Rename, Copy, Duplicate) + // Section 2: 3-8 (Run Branch, Pin, Unpin, Bypass, Remove Bypass, Mute) + // Section 3: 9-15 (Convert to Subgraph, Frame selection, Minimize Node, Expand, Collapse, Resize, Clone) + // Section 4: 16-17 (Node Info, Color) + // Section 5: 18+ (Image operations and fallback items) + const getSectionNumber = (index: number): number => { + if (index <= 2) return 1 + if (index <= 8) return 2 + if (index <= 15) return 3 + if (index <= 17) return 4 + return 5 + } + + let lastSection = 0 + for (const label of coreLabels) { + const item = coreItemsMap.get(label)! + const itemIndex = getMenuItemOrder(label) + const currentSection = getSectionNumber(itemIndex) + + // Add divider when moving to a new section + if (lastSection > 0 && currentSection !== lastSection) { + orderedCoreItems.push({ type: 'divider' }) + } + + orderedCoreItems.push(item) + lastSection = currentSection + } + + // Build the final menu structure + const result: MenuOption[] = [] + + // Add ordered core items with their dividers + result.push(...orderedCoreItems) + + // Add extensions section if there are extension items + if (extensionItems.length > 0) { + // Add divider before Extensions section + result.push({ type: 'divider' }) + + // Add non-clickable Extensions label + result.push({ + label: 'Extensions', + type: 'category', + disabled: true + }) + + // Add extension items + result.push(...extensionItems) + } + + // Add Delete at the bottom if it exists + if (deleteItem) { + result.push({ type: 'divider' }) + result.push(deleteItem) + } + + return result +} + +/** + * Convert LiteGraph IContextMenuValue items to Vue MenuOption format + * Used to bridge LiteGraph context menus into Vue node menus + * @param items - The LiteGraph menu items to convert + * @param node - The node context (optional) + * @param applyStructuring - Whether to apply menu structuring (core/extensions separation). Defaults to true. + */ +export function convertContextMenuToOptions( + items: (IContextMenuValue | null)[], + node?: LGraphNode, + applyStructuring: boolean = true +): MenuOption[] { + const result: MenuOption[] = [] + + for (const item of items) { + // Null items are separators in LiteGraph + if (item === null) { + result.push({ type: 'divider' }) + continue + } + + // Skip items without content (shouldn't happen, but be safe) + if (!item.content) { + continue + } + + // Skip hard blacklisted items + if (HARD_BLACKLIST.has(item.content)) { + continue + } + + // Skip if a similar item already exists in results + if (isDuplicateItem(item.content, result)) { + continue + } + + const option: MenuOption = { + label: item.content, + source: 'litegraph' + } + + // Pass through disabled state + if (item.disabled) { + option.disabled = true + } + + // Handle submenus + if (item.has_submenu) { + // Static submenu with pre-defined options + if (item.submenu?.options) { + option.hasSubmenu = true + option.submenu = convertSubmenuToOptions(item.submenu.options) + } + // Dynamic submenu - callback creates it on-demand + else if (item.callback && !item.disabled) { + option.hasSubmenu = true + // Intercept the callback to capture dynamic submenu items + const capturedSubmenu = captureDynamicSubmenu(item, node) + if (capturedSubmenu) { + option.submenu = capturedSubmenu + } else { + console.warn( + '[ContextMenuConverter] Failed to capture submenu for:', + item.content + ) + } + } + } + // Handle callback (only if not disabled and not a submenu) + else if (item.callback && !item.disabled) { + // Wrap the callback to match the () => void signature + option.action = () => { + try { + void item.callback?.call( + item as unknown as ContextMenuDivElement, + item.value, + {}, + undefined, + undefined, + item + ) + } catch (error) { + console.error('Error executing context menu callback:', error) + } + } + } + + result.push(option) + } + + // Apply structured menu with core items and extensions section (if requested) + if (applyStructuring) { + return buildStructuredMenu(result) + } + + return result +} + +/** + * Capture submenu items from a dynamic submenu callback + * Intercepts ContextMenu constructor to extract items without creating HTML menu + */ +function captureDynamicSubmenu( + item: IContextMenuValue, + node?: LGraphNode +): SubMenuOption[] | undefined { + let capturedItems: readonly (IContextMenuValue | string | null)[] | undefined + let capturedOptions: IContextMenuOptions | undefined + + // Store original ContextMenu constructor + const OriginalContextMenu = LiteGraph.ContextMenu + + try { + // Mock ContextMenu constructor to capture submenu items and options + LiteGraph.ContextMenu = function ( + items: readonly (IContextMenuValue | string | null)[], + options?: IContextMenuOptions + ) { + // Capture both items and options + capturedItems = items + capturedOptions = options + // Return a minimal mock object to prevent errors + return { + close: () => {}, + root: document.createElement('div') + } as unknown as ContextMenu + } as unknown as typeof ContextMenu + + // Execute the callback to trigger submenu creation + try { + // Create a mock MouseEvent for the callback + const mockEvent = new MouseEvent('click', { + bubbles: true, + cancelable: true, + clientX: 0, + clientY: 0 + }) + + // Create a mock parent menu + const mockMenu = { + close: () => {}, + root: document.createElement('div') + } as unknown as ContextMenu + + // Call the callback which should trigger ContextMenu constructor + // Callback signature varies, but typically: (value, options, event, menu, node) + void item.callback?.call( + item as unknown as ContextMenuDivElement, + item.value, + {}, + mockEvent, + mockMenu, + node // Pass the node context for callbacks that need it + ) + } catch (error) { + console.warn( + '[ContextMenuConverter] Error executing callback for:', + item.content, + error + ) + } + } finally { + // Always restore original constructor + LiteGraph.ContextMenu = OriginalContextMenu + } + + // Convert captured items to Vue submenu format + if (capturedItems) { + const converted = convertSubmenuToOptions(capturedItems, capturedOptions) + return converted + } + + console.warn('[ContextMenuConverter] No items captured for:', item.content) + return undefined +} + +/** + * Convert LiteGraph submenu items to Vue SubMenuOption format + */ +function convertSubmenuToOptions( + items: readonly (IContextMenuValue | string | null)[], + options?: IContextMenuOptions +): SubMenuOption[] { + const result: SubMenuOption[] = [] + + for (const item of items) { + // Skip null separators + if (item === null) { + continue + } + + // Handle string items (simple labels like in Mode/Shapes menus) + if (typeof item === 'string') { + const subOption: SubMenuOption = { + label: item, + action: () => { + try { + // Call the options callback with the string value + if (options?.callback) { + void options.callback.call( + null, + item, + options, + undefined, + undefined, + options.extra + ) + } + } catch (error) { + console.error('Error executing string item callback:', error) + } + } + } + result.push(subOption) + continue + } + + // Handle object items + if (!item.content) { + continue + } + + // Extract text content from HTML if present + const content = stripHtmlTags(item.content) + + const subOption: SubMenuOption = { + label: content, + action: () => { + try { + void item.callback?.call( + item as unknown as ContextMenuDivElement, + item.value, + {}, + undefined, + undefined, + item + ) + } catch (error) { + console.error('Error executing submenu callback:', error) + } + } + } + + // Pass through disabled state + if (item.disabled) { + subOption.disabled = true + } + + result.push(subOption) + } + return result +} + +/** + * Strip HTML tags from content string safely + * LiteGraph menu items often include HTML for styling + */ +function stripHtmlTags(html: string): string { + // Use DOMPurify to sanitize and strip all HTML tags + const sanitized = DOMPurify.sanitize(html, { ALLOWED_TAGS: [] }) + const result = sanitized.trim() + return result || html.replace(/<[^>]*>/g, '').trim() || html +} diff --git a/src/composables/graph/useMoreOptionsMenu.ts b/src/composables/graph/useMoreOptionsMenu.ts index 45ea6eb1f2..29d6fec919 100644 --- a/src/composables/graph/useMoreOptionsMenu.ts +++ b/src/composables/graph/useMoreOptionsMenu.ts @@ -15,10 +15,13 @@ export interface MenuOption { icon?: string shortcut?: string hasSubmenu?: boolean - type?: 'divider' + type?: 'divider' | 'category' action?: () => void submenu?: SubMenuOption[] badge?: BadgeVariant + disabled?: boolean + source?: 'litegraph' | 'vue' + isColorPicker?: boolean } export interface SubMenuOption { @@ -26,6 +29,7 @@ export interface SubMenuOption { icon?: string action: () => void color?: string + disabled?: boolean } export enum BadgeVariant { diff --git a/src/composables/useContextMenuTranslation.ts b/src/composables/useContextMenuTranslation.ts index cfcf5c809f..a85cccdc07 100644 --- a/src/composables/useContextMenuTranslation.ts +++ b/src/composables/useContextMenuTranslation.ts @@ -22,7 +22,10 @@ export const useContextMenuTranslation = () => { this: LGraphCanvas, ...args: Parameters ) { - const res: IContextMenuValue[] = getCanvasMenuOptions.apply(this, args) + const res: (IContextMenuValue | null)[] = getCanvasMenuOptions.apply( + this, + args + ) // Add items from new extension API const newApiItems = app.collectCanvasMenuItems(this) @@ -58,13 +61,16 @@ export const useContextMenuTranslation = () => { LGraphCanvas.prototype ) + // Install compatibility layer for getNodeMenuOptions + legacyMenuCompat.install(LGraphCanvas.prototype, 'getNodeMenuOptions') + // Wrap getNodeMenuOptions to add new API items const nodeMenuFn = LGraphCanvas.prototype.getNodeMenuOptions const getNodeMenuOptionsWithExtensions = function ( this: LGraphCanvas, ...args: Parameters ) { - const res = nodeMenuFn.apply(this, args) + const res = nodeMenuFn.apply(this, args) as (IContextMenuValue | null)[] // Add items from new extension API const node = args[0] @@ -73,11 +79,28 @@ export const useContextMenuTranslation = () => { res.push(item) } + // Add legacy monkey-patched items + const legacyItems = legacyMenuCompat.extractLegacyItems( + 'getNodeMenuOptions', + this, + ...args + ) + for (const item of legacyItems) { + res.push(item) + } + return res } LGraphCanvas.prototype.getNodeMenuOptions = getNodeMenuOptionsWithExtensions + legacyMenuCompat.registerWrapper( + 'getNodeMenuOptions', + getNodeMenuOptionsWithExtensions, + nodeMenuFn, + LGraphCanvas.prototype + ) + function translateMenus( values: readonly (IContextMenuValue | string | null)[] | undefined, options: IContextMenuOptions diff --git a/src/lib/litegraph/src/LGraphCanvas.ts b/src/lib/litegraph/src/LGraphCanvas.ts index 0f2fae1da4..0573bc3682 100644 --- a/src/lib/litegraph/src/LGraphCanvas.ts +++ b/src/lib/litegraph/src/LGraphCanvas.ts @@ -712,8 +712,8 @@ export class LGraphCanvas getMenuOptions?(): IContextMenuValue[] getExtraMenuOptions?( canvas: LGraphCanvas, - options: IContextMenuValue[] - ): IContextMenuValue[] + options: (IContextMenuValue | null)[] + ): (IContextMenuValue | null)[] static active_node: LGraphNode /** called before modifying the graph */ onBeforeChange?(graph: LGraph): void @@ -8020,8 +8020,8 @@ export class LGraphCanvas } } - getCanvasMenuOptions(): IContextMenuValue[] { - let options: IContextMenuValue[] + getCanvasMenuOptions(): (IContextMenuValue | null)[] { + let options: (IContextMenuValue | null)[] if (this.getMenuOptions) { options = this.getMenuOptions() } else { diff --git a/src/lib/litegraph/src/contextMenuCompat.ts b/src/lib/litegraph/src/contextMenuCompat.ts index 5d2cd09ad5..38865bf206 100644 --- a/src/lib/litegraph/src/contextMenuCompat.ts +++ b/src/lib/litegraph/src/contextMenuCompat.ts @@ -7,7 +7,9 @@ import type { IContextMenuValue } from './interfaces' */ const ENABLE_LEGACY_SUPPORT = true -type ContextMenuValueProvider = (...args: unknown[]) => IContextMenuValue[] +type ContextMenuValueProvider = ( + ...args: unknown[] +) => (IContextMenuValue | null)[] class LegacyMenuCompat { private originalMethods = new Map() @@ -37,16 +39,22 @@ class LegacyMenuCompat { * @param preWrapperFn The method that existed before the wrapper * @param prototype The prototype to verify wrapper installation */ - registerWrapper( - methodName: keyof LGraphCanvas, - wrapperFn: ContextMenuValueProvider, - preWrapperFn: ContextMenuValueProvider, + registerWrapper( + methodName: K, + wrapperFn: LGraphCanvas[K], + preWrapperFn: LGraphCanvas[K], prototype?: LGraphCanvas ) { - this.wrapperMethods.set(methodName, wrapperFn) - this.preWrapperMethods.set(methodName, preWrapperFn) + this.wrapperMethods.set( + methodName as string, + wrapperFn as unknown as ContextMenuValueProvider + ) + this.preWrapperMethods.set( + methodName as string, + preWrapperFn as unknown as ContextMenuValueProvider + ) const isInstalled = prototype && prototype[methodName] === wrapperFn - this.wrapperInstalled.set(methodName, !!isInstalled) + this.wrapperInstalled.set(methodName as string, !!isInstalled) } /** @@ -54,11 +62,17 @@ class LegacyMenuCompat { * @param prototype The prototype to install on * @param methodName The method name to track */ - install(prototype: LGraphCanvas, methodName: keyof LGraphCanvas) { + install( + prototype: LGraphCanvas, + methodName: K + ) { if (!ENABLE_LEGACY_SUPPORT) return const originalMethod = prototype[methodName] - this.originalMethods.set(methodName, originalMethod) + this.originalMethods.set( + methodName as string, + originalMethod as unknown as ContextMenuValueProvider + ) let currentImpl = originalMethod @@ -66,13 +80,13 @@ class LegacyMenuCompat { get() { return currentImpl }, - set: (newImpl: ContextMenuValueProvider) => { - const fnKey = `${methodName}:${newImpl.toString().slice(0, 100)}` + set: (newImpl: LGraphCanvas[K]) => { + const fnKey = `${methodName as string}:${newImpl.toString().slice(0, 100)}` if (!this.hasWarned.has(fnKey) && this.currentExtension) { this.hasWarned.add(fnKey) console.warn( - `%c[DEPRECATED]%c Monkey-patching ${methodName} is deprecated. (Extension: "${this.currentExtension}")\n` + + `%c[DEPRECATED]%c Monkey-patching ${methodName as string} is deprecated. (Extension: "${this.currentExtension}")\n` + `Please use the new context menu API instead.\n\n` + `See: https://docs.comfy.org/custom-nodes/js/context-menu-migration`, 'color: orange; font-weight: bold', @@ -85,7 +99,15 @@ class LegacyMenuCompat { } /** - * Extract items that were added by legacy monkey patches + * Extract items that were added by legacy monkey patches. + * + * Uses set-based diffing by reference to reliably detect additions regardless + * of item reordering or replacement. Items present in patchedItems but not in + * originalItems (by reference equality) are considered additions. + * + * Note: If a monkey patch removes items (patchedItems has fewer unique items + * than originalItems), a warning is logged but we still return any new items. + * * @param methodName The method name that was monkey-patched * @param context The context to call methods with * @param args Arguments to pass to the methods @@ -95,7 +117,7 @@ class LegacyMenuCompat { methodName: keyof LGraphCanvas, context: LGraphCanvas, ...args: unknown[] - ): IContextMenuValue[] { + ): (IContextMenuValue | null)[] { if (!ENABLE_LEGACY_SUPPORT) return [] if (this.isExtracting) return [] @@ -106,7 +128,7 @@ class LegacyMenuCompat { this.isExtracting = true const originalItems = originalMethod.apply(context, args) as - | IContextMenuValue[] + | (IContextMenuValue | null)[] | undefined if (!originalItems) return [] @@ -127,15 +149,26 @@ class LegacyMenuCompat { const methodToCall = shouldSkipWrapper ? preWrapperMethod : currentMethod const patchedItems = methodToCall.apply(context, args) as - | IContextMenuValue[] + | (IContextMenuValue | null)[] | undefined if (!patchedItems) return [] - if (patchedItems.length > originalItems.length) { - return patchedItems.slice(originalItems.length) as IContextMenuValue[] + // Use set-based diff to detect additions by reference + const originalSet = new Set(originalItems) + const addedItems = patchedItems.filter((item) => !originalSet.has(item)) + + // Warn if items were removed (patched has fewer original items than expected) + const retainedOriginalCount = patchedItems.filter((item) => + originalSet.has(item) + ).length + if (retainedOriginalCount < originalItems.length) { + console.warn( + `[Context Menu Compat] Monkey patch for ${methodName} removed ${originalItems.length - retainedOriginalCount} original menu item(s). ` + + `This may cause unexpected behavior.` + ) } - return [] + return addedItems } catch (e) { console.error('[Context Menu Compat] Failed to extract legacy items:', e) return [] diff --git a/tests-ui/tests/litegraph/core/contextMenuCompat.test.ts b/tests-ui/tests/litegraph/core/contextMenuCompat.test.ts index 623082c68b..fc4c7d156b 100644 --- a/tests-ui/tests/litegraph/core/contextMenuCompat.test.ts +++ b/tests-ui/tests/litegraph/core/contextMenuCompat.test.ts @@ -1,6 +1,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { legacyMenuCompat } from '@/lib/litegraph/src/contextMenuCompat' +import type { IContextMenuValue } from '@/lib/litegraph/src/litegraph' import { LGraphCanvas } from '@/lib/litegraph/src/litegraph' describe('contextMenuCompat', () => { @@ -98,13 +99,15 @@ describe('contextMenuCompat', () => { }) describe('extractLegacyItems', () => { + // Cache base items to ensure reference equality for set-based diffing + const baseItem1 = { content: 'Item 1', callback: () => {} } + const baseItem2 = { content: 'Item 2', callback: () => {} } + beforeEach(() => { - // Setup a mock original method + // Setup a mock original method that returns cached items + // This ensures reference equality when set-based diffing compares items LGraphCanvas.prototype.getCanvasMenuOptions = function () { - return [ - { content: 'Item 1', callback: () => {} }, - { content: 'Item 2', callback: () => {} } - ] + return [baseItem1, baseItem2] } // Install compatibility layer @@ -114,12 +117,13 @@ describe('contextMenuCompat', () => { it('should extract items added by monkey patches', () => { // Monkey-patch to add items const original = LGraphCanvas.prototype.getCanvasMenuOptions - LGraphCanvas.prototype.getCanvasMenuOptions = function (...args: any[]) { - const items = (original as any).apply(this, args) - items.push({ content: 'Custom Item 1', callback: () => {} }) - items.push({ content: 'Custom Item 2', callback: () => {} }) - return items - } + LGraphCanvas.prototype.getCanvasMenuOptions = + function (): (IContextMenuValue | null)[] { + const items = original.apply(this) + items.push({ content: 'Custom Item 1', callback: () => {} }) + items.push({ content: 'Custom Item 2', callback: () => {} }) + return items + } // Extract legacy items const legacyItems = legacyMenuCompat.extractLegacyItems( @@ -142,8 +146,11 @@ describe('contextMenuCompat', () => { expect(legacyItems).toHaveLength(0) }) - it('should return empty array when patched method returns same count', () => { - // Monkey-patch that replaces items but keeps same count + it('should detect replaced items as additions and warn about removed items', () => { + const warnSpy = vi.spyOn(console, 'warn') + + // Monkey-patch that replaces items with different ones (same count) + // With set-based diffing, these are detected as new items since they're different references LGraphCanvas.prototype.getCanvasMenuOptions = function () { return [ { content: 'Replaced 1', callback: () => {} }, @@ -156,7 +163,13 @@ describe('contextMenuCompat', () => { mockCanvas ) - expect(legacyItems).toHaveLength(0) + // Set-based diffing detects the replaced items as additions + expect(legacyItems).toHaveLength(2) + expect(legacyItems[0]).toMatchObject({ content: 'Replaced 1' }) + expect(legacyItems[1]).toMatchObject({ content: 'Replaced 2' }) + + // Should warn about removed original items + expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining('removed')) }) it('should handle errors gracefully', () => { @@ -181,29 +194,36 @@ describe('contextMenuCompat', () => { }) describe('integration', () => { + // Cache base items to ensure reference equality for set-based diffing + const integrationBaseItem = { content: 'Base Item', callback: () => {} } + const integrationBaseItem1 = { content: 'Base Item 1', callback: () => {} } + const integrationBaseItem2 = { content: 'Base Item 2', callback: () => {} } + it('should work with multiple extensions patching', () => { - // Setup base method + // Setup base method with cached item LGraphCanvas.prototype.getCanvasMenuOptions = function () { - return [{ content: 'Base Item', callback: () => {} }] + return [integrationBaseItem] } legacyMenuCompat.install(LGraphCanvas.prototype, 'getCanvasMenuOptions') // First extension patches const original1 = LGraphCanvas.prototype.getCanvasMenuOptions - LGraphCanvas.prototype.getCanvasMenuOptions = function (...args: any[]) { - const items = (original1 as any).apply(this, args) - items.push({ content: 'Extension 1 Item', callback: () => {} }) - return items - } + LGraphCanvas.prototype.getCanvasMenuOptions = + function (): (IContextMenuValue | null)[] { + const items = original1.apply(this) + items.push({ content: 'Extension 1 Item', callback: () => {} }) + return items + } // Second extension patches const original2 = LGraphCanvas.prototype.getCanvasMenuOptions - LGraphCanvas.prototype.getCanvasMenuOptions = function (...args: any[]) { - const items = (original2 as any).apply(this, args) - items.push({ content: 'Extension 2 Item', callback: () => {} }) - return items - } + LGraphCanvas.prototype.getCanvasMenuOptions = + function (): (IContextMenuValue | null)[] { + const items = original2.apply(this) + items.push({ content: 'Extension 2 Item', callback: () => {} }) + return items + } // Extract legacy items const legacyItems = legacyMenuCompat.extractLegacyItems( @@ -218,24 +238,22 @@ describe('contextMenuCompat', () => { }) it('should extract legacy items only once even when called multiple times', () => { - // Setup base method + // Setup base method with cached items LGraphCanvas.prototype.getCanvasMenuOptions = function () { - return [ - { content: 'Base Item 1', callback: () => {} }, - { content: 'Base Item 2', callback: () => {} } - ] + return [integrationBaseItem1, integrationBaseItem2] } legacyMenuCompat.install(LGraphCanvas.prototype, 'getCanvasMenuOptions') // Simulate legacy extension monkey-patching the prototype const original = LGraphCanvas.prototype.getCanvasMenuOptions - LGraphCanvas.prototype.getCanvasMenuOptions = function (...args: any[]) { - const items = (original as any).apply(this, args) - items.push({ content: 'Legacy Item 1', callback: () => {} }) - items.push({ content: 'Legacy Item 2', callback: () => {} }) - return items - } + LGraphCanvas.prototype.getCanvasMenuOptions = + function (): (IContextMenuValue | null)[] { + const items = original.apply(this) + items.push({ content: 'Legacy Item 1', callback: () => {} }) + items.push({ content: 'Legacy Item 2', callback: () => {} }) + return items + } // Extract legacy items multiple times (simulating repeated menu opens) const legacyItems1 = legacyMenuCompat.extractLegacyItems( @@ -268,17 +286,19 @@ describe('contextMenuCompat', () => { }) it('should not extract items from registered wrapper methods', () => { - // Setup base method + // Setup base method with cached item LGraphCanvas.prototype.getCanvasMenuOptions = function () { - return [{ content: 'Base Item', callback: () => {} }] + return [integrationBaseItem] } legacyMenuCompat.install(LGraphCanvas.prototype, 'getCanvasMenuOptions') // Create a wrapper that adds new API items (simulating useContextMenuTranslation) const originalMethod = LGraphCanvas.prototype.getCanvasMenuOptions - const wrapperMethod = function (this: LGraphCanvas) { - const items = (originalMethod as any).apply(this, []) + const wrapperMethod = function ( + this: LGraphCanvas + ): (IContextMenuValue | null)[] { + const items = originalMethod.apply(this) // Add new API items items.push({ content: 'New API Item 1', callback: () => {} }) items.push({ content: 'New API Item 2', callback: () => {} }) @@ -306,16 +326,16 @@ describe('contextMenuCompat', () => { }) it('should extract legacy items even when a wrapper is registered but not active', () => { - // Setup base method + // Setup base method with cached item LGraphCanvas.prototype.getCanvasMenuOptions = function () { - return [{ content: 'Base Item', callback: () => {} }] + return [integrationBaseItem] } legacyMenuCompat.install(LGraphCanvas.prototype, 'getCanvasMenuOptions') // Register a wrapper (but don't set it as the current method) const originalMethod = LGraphCanvas.prototype.getCanvasMenuOptions - const wrapperMethod = function () { + const wrapperMethod = function (): (IContextMenuValue | null)[] { return [{ content: 'Wrapper Item', callback: () => {} }] } legacyMenuCompat.registerWrapper( @@ -327,11 +347,12 @@ describe('contextMenuCompat', () => { // Monkey-patch with a different function (legacy extension) const original = LGraphCanvas.prototype.getCanvasMenuOptions - LGraphCanvas.prototype.getCanvasMenuOptions = function (...args: any[]) { - const items = (original as any).apply(this, args) - items.push({ content: 'Legacy Item', callback: () => {} }) - return items - } + LGraphCanvas.prototype.getCanvasMenuOptions = + function (): (IContextMenuValue | null)[] { + const items = original.apply(this) + items.push({ content: 'Legacy Item', callback: () => {} }) + return items + } // Extract legacy items - should return the legacy item because current method is NOT the wrapper const legacyItems = legacyMenuCompat.extractLegacyItems(