diff --git a/src/components/graph/GraphCanvas.vue b/src/components/graph/GraphCanvas.vue index 67ef64d744..50a17c155a 100644 --- a/src/components/graph/GraphCanvas.vue +++ b/src/components/graph/GraphCanvas.vue @@ -146,6 +146,7 @@ import type { VueNodeData } from '@/composables/graph/useGraphNodeManager' import { useVueNodeLifecycle } from '@/composables/graph/useVueNodeLifecycle' import { useNodeBadge } from '@/composables/node/useNodeBadge' import { useCanvasDrop } from '@/composables/useCanvasDrop' +import { useCanvasSearchBoxMenu } from '@/composables/useCanvasSearchBoxMenu' import { useContextMenuTranslation } from '@/composables/useContextMenuTranslation' import { useCopy } from '@/composables/useCopy' import { useGlobalLitegraph } from '@/composables/useGlobalLitegraph' @@ -459,6 +460,7 @@ useLitegraphSettings() useNodeBadge() useGlobalLitegraph() +useCanvasSearchBoxMenu() useContextMenuTranslation() useCopy() usePaste() diff --git a/src/composables/useCanvasSearchBoxMenu.test.ts b/src/composables/useCanvasSearchBoxMenu.test.ts new file mode 100644 index 0000000000..6e02f1bb28 --- /dev/null +++ b/src/composables/useCanvasSearchBoxMenu.test.ts @@ -0,0 +1,88 @@ +import { createPinia, setActivePinia } from 'pinia' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +import { useCanvasSearchBoxMenu } from '@/composables/useCanvasSearchBoxMenu' +import type { IContextMenuValue } from '@/lib/litegraph/src/litegraph' +import { LGraphCanvas } from '@/lib/litegraph/src/litegraph' +import { useSettingStore } from '@/platform/settings/settingStore' +import { useSearchBoxStore } from '@/stores/workspace/searchBoxStore' +import { createMockCanvas } from '@/utils/__tests__/litegraphTestUtils' + +describe('useCanvasSearchBoxMenu', () => { + let originalGetCanvasMenuOptions: typeof LGraphCanvas.prototype.getCanvasMenuOptions + let mockCanvas: LGraphCanvas + + beforeEach(() => { + setActivePinia(createPinia()) + originalGetCanvasMenuOptions = LGraphCanvas.prototype.getCanvasMenuOptions + + LGraphCanvas.prototype.getCanvasMenuOptions = + function (): (IContextMenuValue | null)[] { + const items: (IContextMenuValue | null)[] = [ + { + content: 'Add Node', + has_submenu: true, + callback: LGraphCanvas.onMenuAdd + }, + { content: 'Add Group', callback: vi.fn() } + ] + return items as (IContextMenuValue | null)[] + } + + mockCanvas = createMockCanvas({ + constructor: { prototype: LGraphCanvas.prototype } as typeof LGraphCanvas + } as Partial) + }) + + afterEach(() => { + LGraphCanvas.prototype.getCanvasMenuOptions = originalGetCanvasMenuOptions + vi.restoreAllMocks() + }) + + it('leaves the default Add Node entry untouched when the setting is off', () => { + vi.spyOn(useSettingStore(), 'get').mockImplementation((id) => + id === 'Comfy.NodeSearchBox.ReplaceCanvasMenu' ? false : undefined + ) + + useCanvasSearchBoxMenu() + const items = LGraphCanvas.prototype.getCanvasMenuOptions.call(mockCanvas) + + const addNode = items.find((i) => i?.content === 'Add Node') + expect(addNode?.callback).toBe(LGraphCanvas.onMenuAdd) + expect(addNode?.has_submenu).toBe(true) + }) + + it('replaces the Add Node callback with the search box trigger when the setting is on', () => { + vi.spyOn(useSettingStore(), 'get').mockImplementation((id) => + id === 'Comfy.NodeSearchBox.ReplaceCanvasMenu' ? true : undefined + ) + const toggleVisible = vi.spyOn(useSearchBoxStore(), 'toggleVisible') + + useCanvasSearchBoxMenu() + const items = LGraphCanvas.prototype.getCanvasMenuOptions.call(mockCanvas) + const addNode = items.find((i) => i?.content === 'Add Node') + + expect(addNode).toBeTruthy() + expect(addNode?.has_submenu).toBe(false) + expect(addNode?.callback).not.toBe(LGraphCanvas.onMenuAdd) + + void addNode?.callback?.call( + addNode as never, + undefined, + undefined, + undefined as never, + undefined + ) + expect(toggleVisible).toHaveBeenCalledTimes(1) + }) + + it('preserves other canvas menu entries', () => { + vi.spyOn(useSettingStore(), 'get').mockReturnValue(true) + + useCanvasSearchBoxMenu() + const items = LGraphCanvas.prototype.getCanvasMenuOptions.call(mockCanvas) + + const contents = items.map((i) => i?.content) + expect(contents).toEqual(['Add Node', 'Add Group']) + }) +}) diff --git a/src/composables/useCanvasSearchBoxMenu.ts b/src/composables/useCanvasSearchBoxMenu.ts new file mode 100644 index 0000000000..a0da34f323 --- /dev/null +++ b/src/composables/useCanvasSearchBoxMenu.ts @@ -0,0 +1,67 @@ +import { legacyMenuCompat } from '@/lib/litegraph/src/contextMenuCompat' +import type { IContextMenuValue } from '@/lib/litegraph/src/litegraph' +import { LGraphCanvas } from '@/lib/litegraph/src/litegraph' +import { useSettingStore } from '@/platform/settings/settingStore' +import { useSearchBoxStore } from '@/stores/workspace/searchBoxStore' + +const REPLACE_SETTING_ID = 'Comfy.NodeSearchBox.ReplaceCanvasMenu' + +/** + * When the experimental "replace canvas menu" setting is enabled, the + * right-click canvas menu's "Add Node" entry opens the Vue node search box + * (which already includes blueprints, partner nodes, core nodes, and + * extensions) instead of the legacy LiteGraph category submenu. + * + * The replacement is identified by callback identity against + * {@link LGraphCanvas.onMenuAdd} so it remains stable across the translation + * wrapper installed by {@link useContextMenuTranslation}. + */ +export const useCanvasSearchBoxMenu = () => { + legacyMenuCompat.install(LGraphCanvas.prototype, 'getCanvasMenuOptions') + + const originalGetCanvasMenuOptions = + LGraphCanvas.prototype.getCanvasMenuOptions + + const wrapped = function ( + this: LGraphCanvas, + ...args: Parameters + ) { + const items = originalGetCanvasMenuOptions.apply(this, args) + + const settingStore = useSettingStore() + if (!settingStore.get(REPLACE_SETTING_ID)) return items + + return items.map((item) => + isLegacyAddNode(item) ? buildSearchBoxAddNode(item) : item + ) + } + + LGraphCanvas.prototype.getCanvasMenuOptions = wrapped + legacyMenuCompat.registerWrapper( + 'getCanvasMenuOptions', + wrapped, + originalGetCanvasMenuOptions, + LGraphCanvas.prototype + ) +} + +function isLegacyAddNode( + item: IContextMenuValue | null +): item is IContextMenuValue { + return ( + !!item && + typeof item === 'object' && + item.callback === LGraphCanvas.onMenuAdd + ) +} + +function buildSearchBoxAddNode(original: IContextMenuValue): IContextMenuValue { + return { + ...original, + has_submenu: false, + submenu: undefined, + callback: () => { + useSearchBoxStore().toggleVisible() + } + } +} diff --git a/src/locales/en/settings.json b/src/locales/en/settings.json index 398f7b6888..144c1e57bf 100644 --- a/src/locales/en/settings.json +++ b/src/locales/en/settings.json @@ -328,6 +328,10 @@ "name": "Show node frequency in search results", "tooltip": "Only applies to v1 (legacy)" }, + "Comfy_NodeSearchBox_ReplaceCanvasMenu": { + "name": "Replace canvas right-click \"Add Node\" with search box", + "tooltip": "When enabled, the right-click canvas menu opens the node search box instead of the LiteGraph category submenu. The search box includes blueprints, partner nodes, core nodes, and extensions." + }, "Comfy_NodeSuggestions_number": { "name": "Number of nodes suggestions", "tooltip": "Only for litegraph searchbox/context menu" diff --git a/src/platform/settings/constants/coreSettings.ts b/src/platform/settings/constants/coreSettings.ts index 4a6a6f79fa..1d87e21f11 100644 --- a/src/platform/settings/constants/coreSettings.ts +++ b/src/platform/settings/constants/coreSettings.ts @@ -38,6 +38,17 @@ export const CORE_SETTINGS: SettingParams[] = [ options: ['default', 'v1 (legacy)', 'litegraph (legacy)'], defaultValue: 'default' }, + { + id: 'Comfy.NodeSearchBox.ReplaceCanvasMenu', + category: ['Comfy', 'Node Search Box', 'ReplaceCanvasMenu'], + experimental: true, + name: 'Replace canvas right-click "Add Node" with search box', + tooltip: + 'When enabled, the right-click canvas menu opens the node search box instead of the LiteGraph category submenu. The search box includes blueprints, partner nodes, core nodes, and extensions.', + type: 'boolean', + defaultValue: false, + versionAdded: '1.46.0' + }, { id: 'Comfy.LinkRelease.Action', category: ['LiteGraph', 'LinkRelease', 'Action'], diff --git a/src/schemas/apiSchema.ts b/src/schemas/apiSchema.ts index e7e40abdf7..6bfa0ac570 100644 --- a/src/schemas/apiSchema.ts +++ b/src/schemas/apiSchema.ts @@ -344,6 +344,7 @@ const zSettings = z.object({ 'Comfy.NodeSearchBoxImpl.ShowCategory': z.boolean(), 'Comfy.NodeSearchBoxImpl.ShowIdName': z.boolean(), 'Comfy.NodeSearchBoxImpl.ShowNodeFrequency': z.boolean(), + 'Comfy.NodeSearchBox.ReplaceCanvasMenu': z.boolean(), 'Comfy.NodeSuggestions.number': z.number(), 'Comfy.Node.BypassAllLinksOnDelete': z.boolean(), 'Comfy.Node.Opacity': z.number(),