refactor workflow node traversal

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(

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
 

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,
+            )

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()

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(),