Fix registration issues (#15890)

main.py

@@ -196,9 +196,9 @@ def execute_prestartup_script():
         return False
 
     node_paths = folder_paths.get_folder_paths("custom_nodes")
+    node_prestartup_times = []
     for custom_node_path in node_paths:
         possible_modules = os.listdir(custom_node_path)
-        node_prestartup_times = []
 
         for possible_module in possible_modules:
             module_path = os.path.join(custom_node_path, possible_module)

nodes.py

@@ -2298,7 +2298,9 @@ async def load_custom_node(module_path: str, ignore=set(), module_parent="custom
                     NODE_CLASS_MAPPINGS[name] = node_cls
                     node_cls.RELATIVE_PYTHON_MODULE = "{}.{}".format(module_parent, get_module_name(module_path))
             if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS") and getattr(module, "NODE_DISPLAY_NAME_MAPPINGS") is not None:
-                NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS)
+                for name, display_name in module.NODE_DISPLAY_NAME_MAPPINGS.items():
+                    if name not in ignore:
+                        NODE_DISPLAY_NAME_MAPPINGS[name] = display_name
             return True
         # V3 Extension Definition
         elif hasattr(module, "comfy_entrypoint"):
@@ -2325,8 +2327,8 @@ async def load_custom_node(module_path: str, ignore=set(), module_parent="custom
                     if schema.node_id not in ignore:
                         NODE_CLASS_MAPPINGS[schema.node_id] = node_cls
                         node_cls.RELATIVE_PYTHON_MODULE = "{}.{}".format(module_parent, get_module_name(module_path))
-                    if schema.display_name is not None:
-                        NODE_DISPLAY_NAME_MAPPINGS[schema.node_id] = schema.display_name
+                        if schema.display_name is not None:
+                            NODE_DISPLAY_NAME_MAPPINGS[schema.node_id] = schema.display_name
                 return True
             except Exception as e:
                 logging.warning(f"Error while calling comfy_entrypoint in {module_path}: {e}")

tests-unit/main_prestartup_test.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import ast
+import importlib
+import logging
+import os
+from pathlib import Path
+from types import SimpleNamespace
+import time
+
+import folder_paths
+
+
+def _load_execute_prestartup_script():
+    main_path = Path(__file__).resolve().parents[1] / "main.py"
+    module = ast.parse(main_path.read_text(), filename=str(main_path))
+    function = next(node for node in module.body if isinstance(node, ast.FunctionDef) and node.name == "execute_prestartup_script")
+    compiled = compile(ast.Module(body=[function], type_ignores=[]), filename=str(main_path), mode="exec")
+    namespace = {
+        "args": SimpleNamespace(disable_all_custom_nodes=False, whitelist_custom_nodes=[], enable_manager=False),
+        "folder_paths": folder_paths,
+        "importlib": importlib,
+        "logging": logging,
+        "os": os,
+        "time": time,
+    }
+    exec(compiled, namespace)  # noqa: S102 - trusted AST extracted from main.py itself, not external input
+    return namespace["execute_prestartup_script"]
+
+
+def _load_prestartup_script_for_paths(monkeypatch, custom_nodes_paths: list[str]):
+    monkeypatch.setattr(
+        folder_paths,
+        "get_folder_paths",
+        lambda name: list(custom_nodes_paths) if name == "custom_nodes" else [],
+    )
+    return _load_execute_prestartup_script()
+
+
+def _make_pack(root: Path, name: str) -> Path:
+    pack = root / name
+    pack.mkdir(parents=True)
+    (pack / "prestartup_script.py").write_text("VALUE = 1\n")
+    return pack
+
+
+def test_execute_prestartup_script_handles_empty_custom_nodes_paths(monkeypatch):
+    execute_prestartup_script = _load_prestartup_script_for_paths(monkeypatch, [])
+
+    execute_prestartup_script()
+
+
+def test_execute_prestartup_script_keeps_all_timing_entries(monkeypatch, tmp_path):
+    first_custom_nodes = tmp_path / "custom_nodes_1"
+    second_custom_nodes = tmp_path / "custom_nodes_2"
+    pack_one = _make_pack(first_custom_nodes, "pack_one")
+    pack_two = _make_pack(second_custom_nodes, "pack_two")
+
+    execute_prestartup_script = _load_prestartup_script_for_paths(monkeypatch, [str(first_custom_nodes), str(second_custom_nodes)])
+
+    messages: list[str] = []
+    monkeypatch.setattr(logging, "info", lambda message, *args, **kwargs: messages.append(message))
+
+    execute_prestartup_script()
+
+    joined = "\n".join(messages)
+    assert str(pack_one) in joined
+    assert str(pack_two) in joined

tests-unit/nodes_test/test_ignore_display_name.py

@@ -0,0 +1,62 @@
+import sys
+
+import pytest
+import torch
+
+from comfy.cli_args import args
+
+if not torch.cuda.is_available():
+    args.cpu = True
+
+import nodes
+
+
+pytestmark = pytest.mark.asyncio
+
+
+@pytest.fixture(autouse=True)
+def _restore_node_mappings():
+    class_mappings = dict(nodes.NODE_CLASS_MAPPINGS)
+    display_name_mappings = dict(nodes.NODE_DISPLAY_NAME_MAPPINGS)
+    try:
+        yield
+    finally:
+        nodes.NODE_CLASS_MAPPINGS.clear()
+        nodes.NODE_CLASS_MAPPINGS.update(class_mappings)
+        nodes.NODE_DISPLAY_NAME_MAPPINGS.clear()
+        nodes.NODE_DISPLAY_NAME_MAPPINGS.update(display_name_mappings)
+        sys.modules.pop("test_v1_custom_node", None)
+        sys.modules.pop("test_v3_custom_node", None)
+
+
+async def test_load_custom_node_skips_display_names_for_ignored_nodes(tmp_path, monkeypatch):
+    v1_module = tmp_path / "test_v1_custom_node.py"
+    v1_module.write_text(
+        "NODE_CLASS_MAPPINGS = {\"LeakTest\": object}\n"
+        "NODE_DISPLAY_NAME_MAPPINGS = {\"LeakTest\": \"Leak Test\"}\n",
+    )
+
+    v3_module = tmp_path / "test_v3_custom_node.py"
+    v3_module.write_text(
+        "from comfy_api.latest import ComfyExtension\n\n"
+        "class LeakTestV3Node:\n"
+        "    @classmethod\n"
+        "    def GET_SCHEMA(cls):\n"
+        "        class Schema:\n"
+        "            node_id = \"LeakTestV3\"\n"
+        "            display_name = \"Leak Test V3\"\n\n"
+        "        return Schema()\n\n\n"
+        "class TestExtension(ComfyExtension):\n"
+        "    async def get_node_list(self):\n"
+        "        return [LeakTestV3Node]\n\n\n"
+        "async def comfy_entrypoint():\n"
+        "    return TestExtension()\n",
+    )
+
+    monkeypatch.syspath_prepend(str(tmp_path))
+
+    assert await nodes.load_custom_node(str(v1_module), ignore={"LeakTest"})
+    assert await nodes.load_custom_node(str(v3_module), ignore={"LeakTestV3"})
+
+    assert "LeakTest" not in nodes.NODE_DISPLAY_NAME_MAPPINGS
+    assert "LeakTestV3" not in nodes.NODE_DISPLAY_NAME_MAPPINGS