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
This commit is contained in:
Peyton DeNiro
2026-05-10 10:43:12 -05:00
parent 7ebd45350d
commit ae8d5a2f2a
5 changed files with 216 additions and 51 deletions
+5 -26
View File
@@ -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 = [