fix(hardening): address all PR #157 review findings

Critical:
  - _load_module(): wrap exec_module() in try/except, pop from sys.modules
    on failure so broken modules aren't cached
  - get_comfyui_path(): add _is_comfyui_directory() check to third (CWD walk)
    strategy so a non-ComfyUI directory named "ComfyUI" isn't accepted
  - get_node_class_mappings(): return cached sys.modules["nodes"] if already
    loaded to avoid re-executing nodes.py and losing init_extra_nodes() state

Suggestions:
  - import_custom_nodes(): wrap sys.path filtering in try/finally so path is
    always restored even if _load_module("server") raises mid-execution
  - render.py: auto-discover helpers from generated_helpers.__all__ instead of
    a manually maintained list
  - pyproject.toml: remove duplicate ruff dev dependency (kept standard
    [project.optional-dependencies] form)

Nitpicks:
  - Remove dead imports (types, StringIO) from test_import_path_resolution.py
  - Trim trailing blank lines in pyproject.toml
  - Tone down _load_module docstring: "eliminates ALL" -> "significantly reduces"

Tests (7 new cases):
  - TestLoadModuleFailureCleanup: broken module removed on exec failure, retry
    starts fresh
  - TestLoadModuleTemp: temp modules purged from sys.modules (success + failure)
  - TestRepeatedNodeLoad: repeated _load_module("nodes") returns cached copy
  - TestGetComfyuiPathThirdStrategy: fake "ComfyUI" dir rejected
  - TestImportCustomNodesSysPathRestoration: sys.path restored after crash

comfyui_to_python/generator/render.py

@@ -5,19 +5,7 @@ from typing import Any
 import black
 
 from ..node_runtime import import_custom_nodes
-from .generated_helpers import (
-    _find_from_extension_location,
-    _is_comfyui_directory,
-    _load_module,
-    _load_module_temp,
-    add_comfyui_directory_to_sys_path,
-    add_extra_model_paths,
-    bootstrap_comfyui_runtime,
-    cleanup_comfyui_runtime,
-    find_path,
-    get_comfyui_path,
-    get_value_at_index,
-)
+from . import generated_helpers
 from .model import GenerationPlan
 
 
@@ -33,20 +21,11 @@ class WorkflowRenderer:
                 {"workflow": plan.metadata_workflow_data}
             )
 
+        # Auto-discover helpers from generated_helpers.__all__ so new helpers
+        # are picked up without manually updating this list.
         func_strings = []
-        for func in [
-            _load_module,
-            _load_module_temp,
-            _is_comfyui_directory,
-            _find_from_extension_location,
-            get_value_at_index,
-            get_comfyui_path,
-            find_path,
-            add_comfyui_directory_to_sys_path,
-            add_extra_model_paths,
-            bootstrap_comfyui_runtime,
-            cleanup_comfyui_runtime,
-        ]:
+        for name in generated_helpers.__all__:
+            func = getattr(generated_helpers, name)
             func_strings.append(f"\n{inspect.getsource(func)}")
 
         static_imports = [

comfyui_to_python/node_runtime.py

@@ -18,9 +18,15 @@ def _is_comfyui_directory(path: str) -> bool:
 def _load_module(module_name: str, filepath: str) -> Any:
     """Load a Python module from an explicit file path, bypassing sys.path.
 
-    This eliminates ALL bare import shadowing attacks by loading modules
+    Significantly reduces bare import shadowing risk by loading modules
     from verified file paths instead of relying on sys.path resolution.
+    If exec_module() raises, the partially-loaded module is removed from
+    sys.modules so subsequent calls start fresh.
     """
+    # Return cached module if already loaded — prevents re-execution
+    # that would reset state (e.g. NODE_CLASS_MAPPINGS after init_extra_nodes).
+    if module_name in sys.modules:
+        return sys.modules[module_name]
     if not os.path.isfile(filepath):
         log.debug("Module file not found: %s (%s)", module_name, filepath)
         return None
@@ -31,7 +37,11 @@ def _load_module(module_name: str, filepath: str) -> Any:
             return None
         mod = importlib.util.module_from_spec(spec)
         sys.modules[module_name] = mod
-        spec.loader.exec_module(mod)
+        try:
+            spec.loader.exec_module(mod)
+        except BaseException:
+            sys.modules.pop(module_name, None)
+            raise
         return mod
     except Exception as e:
         log.debug("Failed to load %s from %s: %s", module_name, filepath, e)
@@ -84,7 +94,10 @@ def get_comfyui_path() -> str | None:
     p = _find_from_extension_location()
     if p:
         return p
-    return find_path("ComfyUI", max_depth=20)
+    p = find_path("ComfyUI", max_depth=20)
+    if p and _is_comfyui_directory(p):
+        return p
+    return None
 
 
 def add_comfyui_directory_to_sys_path() -> None:
@@ -267,10 +280,12 @@ def import_custom_nodes() -> None:
     original_sys_path = list(sys.path)
     sys.path[:] = [p for p in sys.path if p != comfy_subdir]
 
-    server_mod = _load_module("server", os.path.join(comfyui_path, "server.py"))
-
-    # Restore sys.path so nodes and other modules that need comfy/ still work
-    sys.path[:] = original_sys_path
+    try:
+        server_mod = _load_module("server", os.path.join(comfyui_path, "server.py"))
+    finally:
+        # Restore sys.path so nodes and other modules that need comfy/ still work.
+        # Guaranteed even if _load_module raises mid-execution.
+        sys.path[:] = original_sys_path
 
     if execution_mod is None or server_mod is None:
         log.debug(
@@ -292,14 +307,21 @@ def import_custom_nodes() -> None:
 
 
 def get_node_class_mappings() -> dict:
-    """Load ComfyUI node mappings on demand via _load_module()."""
+    """Load ComfyUI node mappings on demand via _load_module().
+
+    Reuses the cached "nodes" module from sys.modules if already loaded
+    (e.g. by import_custom_nodes) to avoid resetting NODE_CLASS_MAPPINGS.
+    """
     add_comfyui_directory_to_sys_path()
     comfyui_path = get_comfyui_path()
     if comfyui_path is None:
         log.debug("get_node_class_mappings: ComfyUI path not found")
         return {}
 
-    nodes_mod = _load_module("nodes", os.path.join(comfyui_path, "nodes.py"))
+    # Reuse cached module if already loaded to avoid resetting mappings.
+    nodes_mod = sys.modules.get("nodes")
+    if nodes_mod is None:
+        nodes_mod = _load_module("nodes", os.path.join(comfyui_path, "nodes.py"))
     if nodes_mod is None:
         return {}
     return getattr(nodes_mod, "NODE_CLASS_MAPPINGS", {})

pyproject.toml

@@ -9,9 +9,6 @@ dependencies = ["black"]
 [project.optional-dependencies]
 dev = ["ruff"]
 
-[tool.uv]
-dev-dependencies = ["ruff"]
-
 [project.urls]
 Repository = "https://github.com/pydn/ComfyUI-to-Python-Extension"
 #  Used by Comfy Registry https://comfyregistry.org
@@ -20,5 +17,3 @@ Repository = "https://github.com/pydn/ComfyUI-to-Python-Extension"
 PublisherId = "pydn"
 DisplayName = "ComfyUI-to-Python-Extension"
 Icon = ""
-
-

tests/test_import_path_resolution.py

@@ -12,9 +12,7 @@ Covers Approach B from docs/specs/harden-import-path-resolution.html:
 import os
 import sys
 import tempfile
-import types
 import unittest
-from io import StringIO
 from unittest.mock import patch
 
 # Real ComfyUI checkout for integration-style unit tests (parent of this repo)
@@ -192,7 +190,9 @@ class TestGeneratedScriptIsolation(unittest.TestCase):
         renderer = WorkflowRenderer()
         generated = renderer.render(plan)
 
-        self.assertIn("_load_module", generated, "Generated script must embed _load_module()")
+        self.assertIn(
+            "_load_module", generated, "Generated script must embed _load_module()"
+        )
         self.assertIn(
             "spec_from_file_location",
             generated,
@@ -283,5 +283,182 @@ class TestIsComfyuiDirectory(unittest.TestCase):
         self.assertFalse(_is_comfyui_directory("/nonexistent/path/to/comfyui"))
 
 
+class TestLoadModuleFailureCleanup(unittest.TestCase):
+    """Tests for _load_module() behavior when exec_module() raises."""
+
+    def tearDown(self):
+        mod_names_to_clean = [
+            name for name in sys.modules if name.startswith("_test_fail_")
+        ]
+        for name in mod_names_to_clean:
+            del sys.modules[name]
+
+    def test_broken_module_removed_from_sys_modules_on_failure(self):
+        """Given a module that raises during exec, it is removed from sys.modules."""
+        from comfyui_to_python.node_runtime import _load_module
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            mod_file = os.path.join(tmpdir, "badmodule.py")
+            with open(mod_file, "w") as f:
+                f.write("raise RuntimeError('mid-execution failure')\n")
+
+            result = _load_module("_test_fail_bad", mod_file)
+        self.assertIsNone(result)
+        self.assertNotIn(
+            "_test_fail_bad",
+            sys.modules,
+            "Failed module must not remain cached in sys.modules",
+        )
+
+    def test_broken_module_retried_cleanly(self):
+        """Given a failed load, a subsequent load attempt starts fresh (not from cache)."""
+        from comfyui_to_python.node_runtime import _load_module
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            bad_file = os.path.join(tmpdir, "bad.py")
+            good_file = os.path.join(tmpdir, "good.py")
+
+            # First load: bad module that raises
+            with open(bad_file, "w") as f:
+                f.write("raise RuntimeError('fail')\n")
+            result_bad = _load_module("_test_fail_retry", bad_file)
+            self.assertIsNone(result_bad)
+
+            # Overwrite with good module
+            os.remove(bad_file)
+            with open(good_file, "w") as f:
+                f.write("VALUE = 99\n")
+
+            # Reload same name from a different (working) file
+            result_good = _load_module("_test_fail_retry", good_file)
+        self.assertIsNotNone(result_good)
+        self.assertEqual(result_good.VALUE, 99)
+
+
+class TestLoadModuleTemp(unittest.TestCase):
+    """Tests for _load_module_temp() removal from sys.modules."""
+
+    def tearDown(self):
+        mod_names_to_clean = [
+            name for name in sys.modules if name.startswith("_test_temp_")
+        ]
+        for name in mod_names_to_clean:
+            del sys.modules[name]
+
+    def test_module_removed_from_sys_modules_after_temp_load(self):
+        """After _load_module_temp(), the module key is absent from sys.modules."""
+        from comfyui_to_python.node_runtime import _load_module_temp
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            mod_file = os.path.join(tmpdir, "tempmod.py")
+            with open(mod_file, "w") as f:
+                f.write("VALUE = 123\n")
+
+            mod = _load_module_temp("_test_temp_mod", mod_file)
+        self.assertIsNotNone(mod)
+        self.assertEqual(mod.VALUE, 123)
+        self.assertNotIn(
+            "_test_temp_mod",
+            sys.modules,
+            "Temp module must be removed from sys.modules after load",
+        )
+
+    def test_module_removed_on_temp_load_failure(self):
+        """After _load_module_temp() fails, the module key is absent from sys.modules."""
+        from comfyui_to_python.node_runtime import _load_module_temp
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            mod_file = os.path.join(tmpdir, "badtemp.py")
+            with open(mod_file, "w") as f:
+                f.write("raise RuntimeError('fail')\n")
+
+            result = _load_module_temp("_test_temp_bad", mod_file)
+        self.assertIsNone(result)
+        self.assertNotIn(
+            "_test_temp_bad",
+            sys.modules,
+            "Failed temp module must be removed from sys.modules",
+        )
+
+
+class TestRepeatedNodeLoad(unittest.TestCase):
+    """Tests that repeated _load_module('nodes', ...) calls respect the cache."""
+
+    def tearDown(self):
+        if "_test_nodes" in sys.modules:
+            del sys.modules["_test_nodes"]
+
+    def test_repeated_load_preserves_cached_attributes(self):
+        """Second _load_module() for same name returns cached module, not a fresh copy."""
+        from comfyui_to_python.node_runtime import _load_module
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            mod_file = os.path.join(tmpdir, "mynodes.py")
+            with open(mod_file, "w") as f:
+                f.write("NODE_CLASS_MAPPINGS = {'A': 1}\n")
+
+            first = _load_module("_test_nodes", mod_file)
+            self.assertIsNotNone(first)
+
+            # Mutate the cached module to simulate init_extra_nodes()
+            first.NODE_CLASS_MAPPINGS["B"] = 2
+
+            # Second load with same name must return cached copy, preserving mutation
+            second = _load_module("_test_nodes", mod_file)
+        self.assertIs(first, second)
+        self.assertEqual(second.NODE_CLASS_MAPPINGS.get("B"), 2)
+
+
+class TestGetComfyuiPathThirdStrategy(unittest.TestCase):
+    """Tests for get_comfyui_path() third-strategy verification."""
+
+    def test_fallback_rejects_non_comfyui_directory(self):
+        """Given a directory named ComfyUI without nodes.py, third strategy rejects it."""
+        from comfyui_to_python.node_runtime import get_comfyui_path
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Create a fake "ComfyUI" directory (no nodes.py)
+            fake_dir = os.path.join(tmpdir, "ComfyUI")
+            os.makedirs(fake_dir)
+
+            original_cwd = os.getcwd()
+            try:
+                os.chdir(tmpdir)
+                with patch.dict("os.environ", {}, clear=False):
+                    # Strip COMFYUI_PATH so fallback strategies activate
+                    env_copy = dict(os.environ)
+                    env_copy.pop("COMFYUI_PATH", None)
+                    with patch.dict("os.environ", env_copy, clear=True):
+                        result = get_comfyui_path()
+            finally:
+                os.chdir(original_cwd)
+
+        # Must NOT return the fake directory
+        self.assertNotEqual(result, fake_dir)
+
+
+class TestImportCustomNodesSysPathRestoration(unittest.TestCase):
+    """Tests for sys.path restoration in import_custom_nodes()."""
+
+    def test_sys_path_restored_on_crash(self):
+        """Given a crash during import_custom_nodes(), sys.path is still restored."""
+        from comfyui_to_python.node_runtime import import_custom_nodes
+
+        original_path = sys.path[:]
+
+        with patch(
+            "comfyui_to_python.node_runtime._load_module",
+            side_effect=RuntimeError("simulated crash during server load"),
+        ):
+            try:
+                import_custom_nodes()
+            except RuntimeError:
+                pass  # Expected crash
+
+        # sys.path should still contain all original entries (no dangling state)
+        for entry in original_path:
+            self.assertIn(entry, sys.path)
+
+
 if __name__ == "__main__":
     unittest.main()

uv.lock

@@ -68,11 +68,6 @@ dev = [
     { name = "ruff" },
 ]
 
-[package.dev-dependencies]
-dev = [
-    { name = "ruff" },
-]
-
 [package.metadata]
 requires-dist = [
     { name = "black" },
@@ -80,9 +75,6 @@ requires-dist = [
 ]
 provides-extras = ["dev"]
 
-[package.metadata.requires-dev]
-dev = [{ name = "ruff" }]
-
 [[package]]
 name = "mypy-extensions"
 version = "1.1.0"