diff --git a/tests/test_downloads_tab_workflow_route.mjs b/tests/test_downloads_tab_workflow_route.mjs index 7ff0346..dd0c2b4 100644 --- a/tests/test_downloads_tab_workflow_route.mjs +++ b/tests/test_downloads_tab_workflow_route.mjs @@ -1195,6 +1195,107 @@ test('node context integration preserves existing menu hooks and refreshes after ]); }); +test('node context integration injects the menu when a node overrides getMenuOptions', () => { + const configureNodeContextMenu = eval(`(${extractMethod(modelResolverSource, 'configureNodeContextMenu')})`); + class NodeType {} + NodeType.prototype.getMenuOptions = function() { + return [{ content: 'Native subgraph action' }]; + }; + + const resolver = { + getResolvedModelsForNodeContextMenu: () => [{ + node_id: 30, + widget_index: 0, + original_path: 'model.safetensors', + full_path: 'E:\\models\\model.safetensors', + category: 'checkpoints', + exists: true, + }], + showResolvedNodeModelInResolver: () => {}, + resolveNodeContextMenuModel: model => model, + dialog: { + getCategoryDisplayName: () => 'checkpoint', + showModelInfo: () => {}, + openContainingFolder: () => {}, + }, + }; + + configureNodeContextMenu.call(resolver, NodeType); + const options = new NodeType().getMenuOptions(null); + + assert.equal(options[0].title, 'Model Resolver'); + assert.equal(options[1].content, 'Native subgraph action'); +}); + +test('node context integration patches instance menu hooks installed by subgraph handlers', () => { + const configureNodeContextMenu = eval(`(${extractMethod(modelResolverSource, 'configureNodeContextMenu')})`); + class NodeType {} + const node = new NodeType(); + node.getExtraMenuOptions = function(_canvas, options) { + options.push({ content: 'Native subgraph action' }); + return []; + }; + + const resolver = { + getResolvedModelsForNodeContextMenu: () => [{ + node_id: 31, + widget_index: 0, + original_path: 'model.safetensors', + full_path: 'E:\\models\\model.safetensors', + category: 'checkpoints', + exists: true, + }], + showResolvedNodeModelInResolver: () => {}, + resolveNodeContextMenuModel: model => model, + dialog: { + getCategoryDisplayName: () => 'checkpoint', + showModelInfo: () => {}, + openContainingFolder: () => {}, + }, + }; + + configureNodeContextMenu.call(resolver, NodeType); + configureNodeContextMenu.call(resolver, node); + const options = []; + node.getExtraMenuOptions(null, options); + + assert.equal(options[0].title, 'Model Resolver'); + assert.equal(options[1].content, 'Native subgraph action'); +}); + +test('node context scope identifies an outer subgraph instance', () => { + const getSubgraphDefinitionIdForNode = eval( + `(${extractMethod(modelResolverSource, 'getSubgraphDefinitionIdForNode')})` + ); + const getNodeContextScope = eval( + `(${extractMethod(modelResolverSource, 'getNodeContextScope')})` + ); + const rootGraph = {}; + const previousApp = globalThis.app; + globalThis.app = { graph: rootGraph }; + + try { + const scope = getNodeContextScope.call( + { getSubgraphDefinitionIdForNode }, + { type: 'subgraph-a', graph: rootGraph }, + { definitions: { subgraphs: [{ id: 'subgraph-a' }] } }, + ); + + assert.deepEqual(scope, { + is_top_level: true, + subgraph_id: 'subgraph-a', + subgraph_instance_id: 'subgraph-a', + is_subgraph_instance: true, + }); + } finally { + if (previousApp === undefined) { + delete globalThis.app; + } else { + globalThis.app = previousApp; + } + } +}); + test('custom node registry exposes normalized LoRA Manager entries', () => { const node = { comfyClass: 'Lora Loader (LoraManager)', diff --git a/tests/test_node_context_menu.mjs b/tests/test_node_context_menu.mjs index 57a7398..c13dfb8 100644 --- a/tests/test_node_context_menu.mjs +++ b/tests/test_node_context_menu.mjs @@ -4,6 +4,7 @@ import test from 'node:test'; import { buildModelResolverNodeMenu, getImmediateModelsForNode, + getModelsForSubgraphInstance, getResolvedModelsForNode, isExistingResolvedModel, matchesWorkflowModelReference, @@ -46,6 +47,61 @@ test('node context menu filters existing models by path and workflow scope', () assert.deepEqual(models.map(item => item.widget_index), [1, 2]); }); +test('node context menu aggregates inner models for an outer subgraph instance', () => { + const models = getModelsForSubgraphInstance({ + resolved_models: [ + { + node_id: 7, + widget_index: 0, + original_path: 'promoted.safetensors', + full_path: 'C:\\promoted', + is_top_level: true, + subgraph_id: 'subgraph-a', + promoted_inner_node_id: 10, + promoted_inner_widget_index: 0, + exists: true, + }, + { + node_id: 10, + widget_index: 0, + original_path: 'promoted.safetensors', + full_path: 'C:\\promoted', + is_top_level: false, + subgraph_id: 'subgraph-a', + exists: true, + }, + { + node_id: 11, + widget_index: 1, + original_path: 'hidden.safetensors', + full_path: 'C:\\hidden', + is_top_level: false, + subgraph_id: 'subgraph-a', + exists: true, + }, + { + node_id: 12, + widget_index: 0, + original_path: 'other.safetensors', + full_path: 'C:\\other', + is_top_level: false, + subgraph_id: 'subgraph-b', + exists: true, + }, + ], + }, 7, { + is_top_level: true, + subgraph_id: 'subgraph-a', + subgraph_instance_id: 'subgraph-a', + is_subgraph_instance: true, + }); + + assert.deepEqual( + models.map(model => model.original_path), + ['promoted.safetensors', 'hidden.safetensors'], + ); +}); + test('node context menu resolves immediate widget selections and ignores invalid references', () => { const node = { id: 7, diff --git a/web/resolver.js b/web/resolver.js index 41f2108..6f6e661 100644 --- a/web/resolver.js +++ b/web/resolver.js @@ -35,12 +35,16 @@ app.registerExtension({ }, nodeCreated(node) { modelResolver.configureNodeContextMenu(node?.constructor); + modelResolver.configureNodeContextMenu(node); + queueMicrotask(() => modelResolver.configureNodeContextMenu(node)); modelResolver.configureWorkflowDependencyMarkerNode(node); modelResolver.configureCustomNodeModelAdapter(node); modelResolver.scheduleNodeContextMenuAnalysis(); }, loadedGraphNode(node) { modelResolver.configureNodeContextMenu(node?.constructor); + modelResolver.configureNodeContextMenu(node); + queueMicrotask(() => modelResolver.configureNodeContextMenu(node)); modelResolver.configureWorkflowDependencyMarkerNode(node); modelResolver.configureCustomNodeModelAdapter(node); modelResolver.scheduleNodeContextMenuAnalysis(); diff --git a/web/resolver/model_resolver.js b/web/resolver/model_resolver.js index cf1783e..822a4df 100644 --- a/web/resolver/model_resolver.js +++ b/web/resolver/model_resolver.js @@ -304,11 +304,33 @@ export class ModelResolver { return { workflow, signature }; } - getNodeContextScope(node) { + getSubgraphDefinitionIdForNode(node, workflow = null) { + const runtimeSubgraphId = node?.subgraph?.id + || node?._subgraph?.id + || node?.subgraph_id + || ''; + const nodeType = String(node?.type || '').trim(); + const definitions = workflow?.definitions?.subgraphs; + if (Array.isArray(definitions)) { + const definition = definitions.find(item => ( + String(item?.id || '') === nodeType + )); + if (definition?.id) return String(definition.id); + } + return runtimeSubgraphId ? String(runtimeSubgraphId) : ''; + } + + getNodeContextScope(node, workflow = null) { const graph = node?.graph; const isTopLevel = !graph || graph === app?.graph; + const subgraphInstanceId = this.getSubgraphDefinitionIdForNode(node, workflow); if (isTopLevel) { - return { is_top_level: true }; + return { + is_top_level: true, + subgraph_id: subgraphInstanceId || undefined, + subgraph_instance_id: subgraphInstanceId || undefined, + is_subgraph_instance: Boolean(subgraphInstanceId), + }; } const subgraphId = graph?._subgraph?.id @@ -319,6 +341,8 @@ export class ModelResolver { return { is_top_level: false, subgraph_id: subgraphId ? String(subgraphId) : undefined, + subgraph_instance_id: subgraphInstanceId || undefined, + is_subgraph_instance: Boolean(subgraphInstanceId), }; } @@ -402,7 +426,8 @@ export class ModelResolver { } getResolvedModelsForNodeContextMenu(node) { - const { signature } = this.getNodeContextWorkflowState(); + const { workflow, signature } = this.getNodeContextWorkflowState(); + const scope = this.getNodeContextScope(node, workflow); const data = this.getCurrentNodeContextAnalysis(signature); if (!data) { this.scheduleNodeContextMenuAnalysis(0); @@ -411,14 +436,14 @@ export class ModelResolver { return getImmediateModelsForNode( previousData || {}, node, - this.getNodeContextScope(node) + scope ); } return getResolvedModelsForNode( data, node?.id, - this.getNodeContextScope(node) + scope ); } @@ -456,16 +481,46 @@ export class ModelResolver { } configureNodeContextMenu(nodeType) { - if (!nodeType?.prototype || nodeType.prototype.__modelResolverContextMenuPatched) { + const isNodeInstance = Boolean( + nodeType + && !nodeType.prototype + && nodeType.constructor?.prototype + ); + const target = isNodeInstance ? nodeType : nodeType?.prototype; + if (!target) { + return; + } + + if (isNodeInstance) { + const hasOwnMenuHook = Object.prototype.hasOwnProperty.call(target, 'getExtraMenuOptions') + || Object.prototype.hasOwnProperty.call(target, 'getMenuOptions'); + if (!hasOwnMenuHook || target.__modelResolverContextMenuInstancePatched) { + return; + } + } else if (target.__modelResolverContextMenuPatched) { return; } const owner = this; - const originalGetExtraMenuOptions = nodeType.prototype.getExtraMenuOptions; - const originalOnWidgetChanged = nodeType.prototype.onWidgetChanged; - nodeType.prototype.getExtraMenuOptions = function(_, options) { - const result = originalGetExtraMenuOptions?.apply(this, arguments); - const models = owner.getResolvedModelsForNodeContextMenu(this); + const hasOwnExtraMenuOptions = Object.prototype.hasOwnProperty.call( + target, + 'getExtraMenuOptions' + ); + const hasOwnMenuOptions = Object.prototype.hasOwnProperty.call( + target, + 'getMenuOptions' + ); + const originalGetExtraMenuOptions = target.getExtraMenuOptions; + const originalGetMenuOptions = target.getMenuOptions; + const addModelResolverMenu = (node, targetOptions) => { + if (!Array.isArray(targetOptions)) return; + if (targetOptions.some(option => ( + String(option?.className || '').includes('mr-model-resolver-node-menu') + ))) { + return; + } + + const models = owner.getResolvedModelsForNodeContextMenu(node); const menu = buildModelResolverNodeMenu(models, { showInResolver: async model => { const resolvedModel = await owner.resolveNodeContextMenuModel(model); @@ -489,32 +544,54 @@ export class ModelResolver { formatCategory: category => owner.dialog?.getCategoryDisplayName?.(category), }); - const targetOptions = Array.isArray(options) - ? options - : (Array.isArray(result) ? result : null); - if (menu && targetOptions) { - targetOptions.unshift(menu); - } - return result; + if (menu) targetOptions.unshift(menu); }; - nodeType.prototype.onWidgetChanged = function() { - const result = originalOnWidgetChanged?.apply(this, arguments); - const widgetName = arguments[0]; - const handledByCustomNodeAdapter = owner.isCustomNodeModelWidget?.( - this, - widgetName - ); - if (handledByCustomNodeAdapter) { - owner.customNodeModelAdapterStates?.get(this)?.notify?.(); - } else if (!owner.dialog?.isWorkflowRefreshSuppressed?.()) { - if (!owner.dialog?.isWorkflowStrengthWidgetName?.(widgetName)) { - owner.scheduleNodeContextMenuAnalysis(); + const originalOnWidgetChanged = target.onWidgetChanged; + if (!isNodeInstance || hasOwnExtraMenuOptions) { + target.getExtraMenuOptions = function(_, options) { + const result = originalGetExtraMenuOptions?.apply(this, arguments); + const targetOptions = Array.isArray(options) + ? options + : (Array.isArray(result) ? result : null); + addModelResolverMenu(this, targetOptions); + return result; + }; + } + // ComfyUI treats the presence of getMenuOptions as an alternate menu + // provider. Only wrap node types that already implement it; adding an + // undefined method makes the native context menu return no options. + if (typeof originalGetMenuOptions === 'function' && (!isNodeInstance || hasOwnMenuOptions)) { + target.getMenuOptions = function() { + const result = originalGetMenuOptions.apply(this, arguments); + const targetOptions = Array.isArray(result) + ? result + : (Array.isArray(arguments[1]) ? arguments[1] : null); + addModelResolverMenu(this, targetOptions); + return result; + }; + } + if (!isNodeInstance) { + target.onWidgetChanged = function() { + const result = originalOnWidgetChanged?.apply(this, arguments); + const widgetName = arguments[0]; + const handledByCustomNodeAdapter = owner.isCustomNodeModelWidget?.( + this, + widgetName + ); + if (handledByCustomNodeAdapter) { + owner.customNodeModelAdapterStates?.get(this)?.notify?.(); + } else if (!owner.dialog?.isWorkflowRefreshSuppressed?.()) { + if (!owner.dialog?.isWorkflowStrengthWidgetName?.(widgetName)) { + owner.scheduleNodeContextMenuAnalysis(); + } + owner.dialog?.scheduleActiveWorkflowRefresh?.('node-widget-change'); } - owner.dialog?.scheduleActiveWorkflowRefresh?.('node-widget-change'); - } - return result; - }; - nodeType.prototype.__modelResolverContextMenuPatched = true; + return result; + }; + target.__modelResolverContextMenuPatched = true; + } else { + target.__modelResolverContextMenuInstancePatched = true; + } } isCustomNodeModelWidget(node, widgetName = '') { diff --git a/web/resolver/node_context_menu.js b/web/resolver/node_context_menu.js index 89427ac..d3ac3ec 100644 --- a/web/resolver/node_context_menu.js +++ b/web/resolver/node_context_menu.js @@ -119,17 +119,25 @@ export function getResolvedModelsForNode(data = {}, nodeId, scope = {}) { subgraph_id: scope.subgraph_id, }; - return (data.resolved_models || []) + const directModels = (data.resolved_models || []) .filter(isExistingResolvedModel) .filter(model => matchesWorkflowModelReference(model, nodeReference)) - .sort((left, right) => { - const widgetOrder = Number(left.widget_index ?? 0) - Number(right.widget_index ?? 0); - if (widgetOrder) return widgetOrder; - return getResolvedModelMenuLabel(left).localeCompare(getResolvedModelMenuLabel(right)); - }); + if (!scope.is_subgraph_instance) { + return sortResolvedModels(directModels); + } + + return getModelsForSubgraphInstance(data, nodeId, scope, { + directModels, + }); } export function getImmediateModelsForNode(data = {}, node = {}, scope = {}) { + if (scope.is_subgraph_instance) { + return getModelsForSubgraphInstance(data, node?.id, scope, { + includeMissing: true, + }); + } + const references = [ ...(data.resolved_models || []), ...(data.missing_models || []), @@ -177,6 +185,111 @@ export function getImmediateModelsForNode(data = {}, node = {}, scope = {}) { )); } +function sortResolvedModels(models = []) { + return [...models].sort((left, right) => { + const widgetOrder = Number(left.widget_index ?? 0) - Number(right.widget_index ?? 0); + if (widgetOrder) return widgetOrder; + return getResolvedModelMenuLabel(left).localeCompare(getResolvedModelMenuLabel(right)); + }); +} + +function getSubgraphModelReferenceKey(model = {}) { + return [ + model.is_top_level !== false ? 'top' : 'nested', + model.subgraph_id || '', + model.node_id ?? '', + model.widget_index ?? '', + normalizeIdentity(model.original_path || model.name || ''), + ].join('|'); +} + +function getMissingModelReferences(data = {}) { + const references = []; + for (const missing of data.missing_models || []) { + const nodeReferences = Array.isArray(missing?.all_node_refs) + && missing.all_node_refs.length + ? missing.all_node_refs + : [missing]; + for (const reference of nodeReferences) { + references.push({ + ...missing, + ...reference, + full_path: '', + path: '', + resolved_path: '', + exists: false, + resolution_pending: true, + }); + } + } + return references; +} + +export function getModelsForSubgraphInstance( + data = {}, + nodeId, + scope = {}, + { directModels = null, includeMissing = false } = {}, +) { + const instanceSubgraphId = String(scope.subgraph_instance_id || '').trim(); + if (!instanceSubgraphId) return []; + + const direct = directModels || (data.resolved_models || []) + .filter(isExistingResolvedModel) + .filter(model => matchesWorkflowModelReference(model, { + node_id: nodeId, + is_top_level: scope.is_top_level, + subgraph_id: scope.subgraph_id, + })); + const directPromotedInnerKeys = new Set( + direct + .filter(model => model.promoted_inner_node_id !== undefined) + .map(model => [ + instanceSubgraphId, + model.promoted_inner_node_id, + model.promoted_inner_widget_index, + ].join('|')) + ); + const nested = (data.resolved_models || []) + .filter(isExistingResolvedModel) + .filter(model => ( + model.is_top_level === false + && String(model.subgraph_id || '') === instanceSubgraphId + && !directPromotedInnerKeys.has([ + instanceSubgraphId, + model.node_id, + model.widget_index, + ].join('|')) + )); + + const models = [...direct, ...nested]; + if (includeMissing) { + const missing = getMissingModelReferences(data).filter(model => { + const isDirect = ( + String(model.node_id ?? '') === String(nodeId ?? '') + && Boolean(model.is_top_level !== false) === Boolean(scope.is_top_level) + && String(model.subgraph_id || '') === String(scope.subgraph_id || '') + ); + const isNested = ( + model.is_top_level === false + && String(model.subgraph_id || '') === instanceSubgraphId + ); + return isDirect || isNested; + }); + models.push(...missing); + } + + const uniqueModels = []; + const seen = new Set(); + for (const model of models) { + const key = getSubgraphModelReferenceKey(model); + if (seen.has(key)) continue; + seen.add(key); + uniqueModels.push(model); + } + return sortResolvedModels(uniqueModels); +} + export function buildModelResolverNodeMenu(models = [], handlers = {}, options = {}) { const resolvedModels = models.filter(model => ( isExistingResolvedModel(model) || model.resolution_pending === true