diff --git a/strix/llm/warmup.py b/strix/llm/warmup.py index 98da959d..0459469d 100644 --- a/strix/llm/warmup.py +++ b/strix/llm/warmup.py @@ -14,6 +14,7 @@ from __future__ import annotations import importlib import logging +import sys import threading @@ -30,12 +31,38 @@ _lock = threading.Lock() _thread: threading.Thread | None = None +def _purge_orphaned_modules(before: frozenset[str]) -> None: + """Remove submodules stranded by an import attempt that just failed. + + When a package import fails partway (for example CPython's import-lock + deadlock avoidance breaking a cross-thread cycle), the failed package is + removed from ``sys.modules`` but submodules it already finished stay + behind. A later import of one of those submodules then short-circuits on + the cached entry without re-importing its parent, and re-entering the + parent from inside a submodule crashes with "partially initialized + module". Dropping the orphans (cached submodules whose ancestor package is + gone) restores a clean slate, and touches nothing another thread imported + successfully. + """ + added = set(sys.modules) - before + for name in added: + parent = name.rpartition(".")[0] + while parent: + if parent not in sys.modules: + sys.modules.pop(name, None) + logger.debug("Import warm-up purged orphaned module %r", name) + break + parent = parent.rpartition(".")[0] + + def _warm(modules: tuple[str, ...]) -> None: for name in modules: + before = frozenset(sys.modules) try: importlib.import_module(name) except Exception: # noqa: BLE001 - a failed warm-up must never fail the run. logger.debug("Import warm-up for %r failed", name, exc_info=True) + _purge_orphaned_modules(before) def start_import_warmup(modules: tuple[str, ...] = WARMUP_MODULES) -> threading.Thread: diff --git a/strix/report/__init__.py b/strix/report/__init__.py index c8c80eec..dca4114b 100644 --- a/strix/report/__init__.py +++ b/strix/report/__init__.py @@ -1,12 +1,26 @@ """Report/finding helpers.""" -from strix.report.dedupe import check_duplicate +from importlib import import_module +from typing import TYPE_CHECKING, Any + from strix.report.state import ReportState, get_global_report_state, set_global_report_state +if TYPE_CHECKING: + from strix.report.dedupe import check_duplicate + __all__ = [ "ReportState", "check_duplicate", "get_global_report_state", "set_global_report_state", ] + + +def __getattr__(name: str) -> Any: + # check_duplicate pulls in the agents SDK import graph, so it resolves + # lazily: importing this package must stay lightweight and never enter + # that graph (the import warm-up thread may be walking it concurrently). + if name == "check_duplicate": + return import_module("strix.report.dedupe").check_duplicate + raise AttributeError(name) diff --git a/strix/report/state.py b/strix/report/state.py index 623cb9e9..a665adea 100644 --- a/strix/report/state.py +++ b/strix/report/state.py @@ -6,18 +6,15 @@ from collections.abc import Callable from datetime import UTC, datetime from importlib.metadata import PackageNotFoundError, version from pathlib import Path -from typing import Any, Optional, cast +from typing import TYPE_CHECKING, Any, Optional, cast from uuid import uuid4 -from agents.usage import Usage - from strix.config import codex from strix.config.loader import load_settings from strix.core.paths import run_dir_for, runtime_state_dir from strix.report.coverage import write_coverage from strix.report.pricing import resolve_litellm_model from strix.report.sarif import write_sarif -from strix.report.usage import LLMUsageLedger from strix.report.writer import ( read_run_record, write_executive_report, @@ -27,6 +24,10 @@ from strix.report.writer import ( from strix.telemetry import posthog, scarf +if TYPE_CHECKING: + from agents.usage import Usage + + logger = logging.getLogger(__name__) _global_report_state: Optional["ReportState"] = None @@ -131,6 +132,10 @@ class ReportState: self.scan_results: dict[str, Any] | None = None self.scan_config: dict[str, Any] | None = None + # Imported here so importing this module never enters the agents SDK + # package (which the warm-up thread may be initializing concurrently). + from strix.report.usage import LLMUsageLedger + self._llm_usage = LLMUsageLedger() self._telemetry_llm_usage_baseline: dict[str, Any] = {} auth_mode = codex.auth_mode(load_settings().llm.model) @@ -338,7 +343,7 @@ class ReportState: self, *, agent_id: str, - usage: Usage | None, + usage: "Usage | None", agent_name: str | None = None, model: str | None = None, ) -> None: diff --git a/tests/test_import_warmup.py b/tests/test_import_warmup.py new file mode 100644 index 00000000..64d3bfba --- /dev/null +++ b/tests/test_import_warmup.py @@ -0,0 +1,99 @@ +"""The import warm-up thread must never leave the import system poisoned. + +Field failure: the warm-up thread's ``strix.core.runner`` import and the main +thread's ``strix.report`` import both walked the agents SDK graph, and the two +held each other's import locks (report -> dedupe -> agents while runner -> +hooks -> report.state). CPython's deadlock avoidance breaks such a cycle by +failing one import, which strands finished submodules in ``sys.modules`` with +their parent package gone — and the next import of one of those submodules +crashes with "partially initialized module". +""" + +from __future__ import annotations + +import subprocess +import sys +import textwrap + +from strix.llm import warmup + + +def _run(code: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( # noqa: S603 + [sys.executable, "-c", textwrap.dedent(code)], + capture_output=True, + text=True, + check=False, + timeout=300, + ) + + +def test_strix_report_does_not_import_the_agents_graph() -> None: + result = _run( + """ + import sys + + import strix.report + + agents_modules = [m for m in sys.modules if m == "agents" or m.startswith("agents.")] + assert not agents_modules, agents_modules + assert "strix.report.dedupe" not in sys.modules + """ + ) + assert result.returncode == 0, result.stderr + + +def test_check_duplicate_resolves_lazily() -> None: + result = _run( + """ + import strix.report + from strix.report import check_duplicate + from strix.report.dedupe import check_duplicate as direct + + assert strix.report.check_duplicate is direct is check_duplicate + """ + ) + assert result.returncode == 0, result.stderr + + +def test_failed_warm_import_purges_orphaned_submodules() -> None: + result = _run( + """ + import sys + + from strix.llm.warmup import _warm + + # A package whose import fails after a submodule already completed: + # CPython removes the package but leaves the submodule stranded. + import pathlib + import tempfile + + root = pathlib.Path(tempfile.mkdtemp()) + pkg = root / "stranded_pkg" + pkg.mkdir() + (pkg / "ok.py").write_text("VALUE = 1") + (pkg / "__init__.py").write_text("from . import ok\\nraise RuntimeError('boom')") + sys.path.insert(0, str(root)) + + _warm(("stranded_pkg",)) + + assert "stranded_pkg" not in sys.modules + assert "stranded_pkg.ok" not in sys.modules, "orphan survived the purge" + + # And the subtree imports cleanly afterwards up to the real error. + try: + import stranded_pkg # noqa: F401 + except RuntimeError: + pass + else: + raise AssertionError("expected the package's own error") + """ + ) + assert result.returncode == 0, result.stderr + + +def test_purge_does_not_touch_preexisting_or_healthy_modules() -> None: + before = frozenset(sys.modules) - {"strix.llm.warmup"} + warmup._purge_orphaned_modules(before) + assert "strix.llm.warmup" in sys.modules # parent chain intact -> kept + assert "strix" in sys.modules