diff --git a/core/services/loaded_models_service.py b/core/services/loaded_models_service.py index a35ace5..dd74ed3 100644 --- a/core/services/loaded_models_service.py +++ b/core/services/loaded_models_service.py @@ -1,6 +1,7 @@ """Loaded model inspection used by the HTTP route adapter.""" from ..routes.context import RouteContext +from ..workflow.traversal import iter_workflow_nodes_with_scope class LoadedModelsService: @@ -40,25 +41,9 @@ class LoadedModelsService: self.extension._update_loaded_progress(loaded_id, *args, **kwargs) def get_workflow_node_count(): - node_count = 0 - nodes = workflow_json.get("nodes", []) - if isinstance(nodes, list): - node_count += len(nodes) - - definitions = workflow_json.get("definitions", {}) - subgraphs = ( - definitions.get("subgraphs", []) - if isinstance(definitions, dict) - else [] + return sum( + 1 for _ in iter_workflow_nodes_with_scope(workflow_json) ) - if isinstance(subgraphs, list): - for subgraph in subgraphs: - if not isinstance(subgraph, dict): - continue - subgraph_nodes = subgraph.get("nodes", []) - if isinstance(subgraph_nodes, list): - node_count += len(subgraph_nodes) - return node_count def interpolate_percent(start, end, current, total): if not total: @@ -183,40 +168,19 @@ class LoadedModelsService: pass # Also extract from node.properties.models - workflow_nodes = workflow_json.get("nodes", []) - nodes = list(workflow_nodes) if isinstance(workflow_nodes, list) else [] - node_scopes = { - id(node): { - "subgraph_id": "", - "subgraph_name": "", - "is_top_level": True, - } - for node in nodes - if isinstance(node, dict) - } - definitions = workflow_json.get("definitions", {}) - subgraphs = ( - definitions.get("subgraphs", []) - if isinstance(definitions, dict) - else [] + node_contexts = list( + iter_workflow_nodes_with_scope(workflow_json) ) - for subgraph in subgraphs: - if isinstance(subgraph, dict): - subgraph_nodes = subgraph.get("nodes", []) - if isinstance(subgraph_nodes, list): - nodes.extend(subgraph_nodes) - for node in subgraph_nodes: - if not isinstance(node, dict): - continue - node_scopes[id(node)] = { - "subgraph_id": subgraph.get("id") or "", - "subgraph_name": ( - subgraph.get("name") - or subgraph.get("id") - or "" - ), - "is_top_level": False, - } + nodes = [context.node for context in node_contexts] + node_scopes = { + id(context.node): { + "subgraph_id": context.subgraph_id, + "subgraph_name": context.subgraph_name, + "is_top_level": context.is_top_level, + } + for context in node_contexts + if isinstance(context.node, dict) + } def get_node_scope(node): return node_scopes.get( diff --git a/core/workflow/analysis.py b/core/workflow/analysis.py index c25c997..c51d04d 100644 --- a/core/workflow/analysis.py +++ b/core/workflow/analysis.py @@ -7,6 +7,7 @@ from typing import Any, Callable, Dict, List, Optional from ..log_system import create_module_logger from . import references from . import subgraphs as subgraph_utils +from .traversal import iter_workflow_nodes_with_scope log = create_module_logger(__name__) @@ -60,35 +61,15 @@ def _get_workflow_node_fingerprints( workflow_json: Dict[str, Any], ) -> Dict[tuple[str, str, str], str]: fingerprints = {} - nodes = workflow_json.get("nodes", []) - if isinstance(nodes, list): - for node in nodes: - if not isinstance(node, dict): - continue - key = _get_workflow_node_cache_key(node, is_top_level=True) - fingerprints[key] = _get_workflow_node_fingerprint(node) - - definitions = workflow_json.get("definitions", {}) - subgraphs = definitions.get("subgraphs", []) if isinstance(definitions, dict) else [] - if not isinstance(subgraphs, list): - return fingerprints - - for subgraph in subgraphs: - if not isinstance(subgraph, dict): + for context in iter_workflow_nodes_with_scope(workflow_json): + if not isinstance(context.node, dict): continue - subgraph_id = subgraph.get("id") - subgraph_nodes = subgraph.get("nodes", []) - if not isinstance(subgraph_nodes, list): - continue - for node in subgraph_nodes: - if not isinstance(node, dict): - continue - key = _get_workflow_node_cache_key( - node, - subgraph_id=subgraph_id, - is_top_level=False, - ) - fingerprints[key] = _get_workflow_node_fingerprint(node) + key = _get_workflow_node_cache_key( + context.node, + subgraph_id=context.subgraph_id, + is_top_level=context.is_top_level, + ) + fingerprints[key] = _get_workflow_node_fingerprint(context.node) return fingerprints @@ -262,12 +243,8 @@ def analyze_workflow_models( previous_node_cache, is_top_level=True, ) - total_nodes = len(nodes) + sum( - len(subgraph_nodes) - for subgraph in subgraphs - if isinstance(subgraph, dict) - for subgraph_nodes in [subgraph.get("nodes", [])] - if isinstance(subgraph_nodes, list) + total_nodes = sum( + 1 for _ in iter_workflow_nodes_with_scope(workflow_json) ) processed_nodes = 0 diff --git a/core/workflow/traversal.py b/core/workflow/traversal.py new file mode 100644 index 0000000..5337a91 --- /dev/null +++ b/core/workflow/traversal.py @@ -0,0 +1,50 @@ +"""Shared workflow node traversal helpers.""" + +from dataclasses import dataclass +from typing import Any, Dict, Iterator, Optional + + +@dataclass(frozen=True) +class WorkflowNodeContext: + """A workflow node together with the scope where it is defined.""" + + node: Any + subgraph_id: Optional[Any] = "" + subgraph_name: Optional[Any] = "" + is_top_level: bool = True + + +def iter_workflow_nodes_with_scope( + workflow_json: Dict[str, Any], +) -> Iterator[WorkflowNodeContext]: + """Yield top-level and subgraph nodes in serialized workflow order.""" + nodes = workflow_json.get("nodes", []) + if isinstance(nodes, list): + for node in nodes: + yield WorkflowNodeContext(node=node) + + definitions = workflow_json.get("definitions", {}) + subgraphs = ( + definitions.get("subgraphs", []) + if isinstance(definitions, dict) + else [] + ) + if not isinstance(subgraphs, list): + return + + for subgraph in subgraphs: + if not isinstance(subgraph, dict): + continue + subgraph_nodes = subgraph.get("nodes", []) + if not isinstance(subgraph_nodes, list): + continue + + subgraph_id = subgraph.get("id") or "" + subgraph_name = subgraph.get("name") or subgraph_id or "" + for node in subgraph_nodes: + yield WorkflowNodeContext( + node=node, + subgraph_id=subgraph_id, + subgraph_name=subgraph_name, + is_top_level=False, + ) diff --git a/tests/test_loaded_models_routes.py b/tests/test_loaded_models_routes.py index eca93de..95eace3 100644 --- a/tests/test_loaded_models_routes.py +++ b/tests/test_loaded_models_routes.py @@ -235,6 +235,43 @@ async def test_loaded_models_route_handles_invalid_node_collections_without_fold assert json.loads(response.text) == {"loaded_models": [], "total": 0} +@pytest.mark.asyncio +async def test_loaded_models_progress_counts_top_level_and_subgraph_nodes(): + handlers, values = _build_routes() + + def inventory_with_progress(_workflow, progress_callback): + progress_callback( + { + "stage": "analyzing", + "current": 1, + "total": 5, + } + ) + return {"available_models": [], "model_refs": []} + + values["get_workflow_model_inventory"].side_effect = inventory_with_progress + workflow = { + "nodes": [{"id": 1}, {"id": 2}], + "definitions": { + "subgraphs": [ + { + "id": "group-1", + "nodes": [{"id": 3}, {"id": 4}, {"id": 5}], + } + ] + }, + } + + with patch.dict(sys.modules, {"folder_paths": None}): + response = await handlers[("POST", "/model_resolver/loaded")]( + _request({"workflow": workflow}) + ) + + assert response.status == 200 + values["self"]._update_workflow_analysis_progress.assert_called_once() + assert values["self"]._update_workflow_analysis_progress.call_args.args[1] == 5 + + @pytest.mark.asyncio async def test_loaded_models_route_reads_folder_paths_and_ignores_category_errors(): folder_paths = MagicMock() diff --git a/tests/test_workflow_components.py b/tests/test_workflow_components.py index eb0c7e7..8da0eef 100644 --- a/tests/test_workflow_components.py +++ b/tests/test_workflow_components.py @@ -1,7 +1,7 @@ from types import SimpleNamespace from unittest.mock import patch -from core.workflow import analysis, references, subgraphs +from core.workflow import analysis, references, subgraphs, traversal from core.workflow_updater import update_workflow_nodes @@ -22,6 +22,40 @@ def _loader_node(node_id=10, value="model.safetensors", link=None): } +def test_workflow_node_traversal_preserves_node_order_and_scope(): + top_level_node = {"id": 1, "type": "TopLevel"} + subgraph_node = {"id": 2, "type": "Inner"} + workflow = { + "nodes": [top_level_node, "invalid-top-level-node"], + "definitions": { + "subgraphs": [ + None, + { + "id": "group-1", + "name": "Group one", + "nodes": [subgraph_node, "invalid-subgraph-node"], + }, + ] + }, + } + + contexts = list(traversal.iter_workflow_nodes_with_scope(workflow)) + + assert [context.node for context in contexts] == [ + top_level_node, + "invalid-top-level-node", + subgraph_node, + "invalid-subgraph-node", + ] + assert contexts[0].is_top_level is True + assert contexts[0].subgraph_id == "" + assert contexts[1].is_top_level is True + assert contexts[2].is_top_level is False + assert contexts[2].subgraph_id == "group-1" + assert contexts[2].subgraph_name == "Group one" + assert contexts[3].subgraph_id == "group-1" + + def test_subgraph_widget_indexes_support_named_inputs_and_links(): node = { **_loader_node(),