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", {})