diff --git a/strix/core/runner.py b/strix/core/runner.py index ee996cd3..cc09a083 100644 --- a/strix/core/runner.py +++ b/strix/core/runner.py @@ -18,7 +18,7 @@ from openai import RateLimitError from strix.agents.factory import build_strix_agent, make_child_factory from strix.agents.prompt import render_system_prompt -from strix.config import load_settings +from strix.config import codex, load_settings from strix.config.models import ( StrixProvider, configure_sdk_model_defaults, @@ -521,24 +521,36 @@ async def run_strix_scan( with contextlib.suppress(Exception): await coordinator.set_status(root_id, "stopped") return None - except RateLimitError as exc: - logger.warning( - "Scan %s stopped: persistent rate limit from the LLM provider (%s). " - "Resume with 'strix --resume %s' once the limit clears.", - scan_id, - exc, - scan_id, - ) - if root_id is not None: - with contextlib.suppress(Exception): - await coordinator.set_status(root_id, "stopped") - return None except (asyncio.CancelledError, KeyboardInterrupt): logger.info("Scan %s interrupted by the user", scan_id) if root_id is not None: with contextlib.suppress(Exception): await coordinator.set_status(root_id, "running") raise + except Exception as exc: + # A usage-limit rejection is resumable regardless of the provider's + # exception class: OpenAI raises RateLimitError, but a LiteLLM-routed + # provider may surface the same exhausted-window error as a different + # type. Both retry layers already stop retrying it (see + # codex.is_usage_limit_error), so route it to the same resumable stop + # instead of the generic failure path. + if isinstance(exc, RateLimitError) or codex.is_usage_limit_error(exc): + logger.warning( + "Scan %s stopped: persistent rate limit from the LLM provider (%s). " + "Resume with 'strix --resume %s' once the limit clears.", + scan_id, + exc, + scan_id, + ) + if root_id is not None: + with contextlib.suppress(Exception): + await coordinator.set_status(root_id, "stopped") + return None + logger.exception("Strix scan %s failed", scan_id) + if root_id is not None: + with contextlib.suppress(Exception): + await coordinator.set_status(root_id, "failed") + raise except BaseException: logger.exception("Strix scan %s failed", scan_id) if root_id is not None: diff --git a/tests/test_runner_rate_limit.py b/tests/test_runner_rate_limit.py index 3110ae2c..b9eb504d 100644 --- a/tests/test_runner_rate_limit.py +++ b/tests/test_runner_rate_limit.py @@ -1,4 +1,4 @@ -"""Tests for graceful handling of persistent RateLimitError in run_strix_scan.""" +"""Tests for graceful handling of persistent usage/rate-limit errors in run_strix_scan.""" from __future__ import annotations @@ -24,11 +24,22 @@ def _make_rate_limit_error() -> RateLimitError: return RateLimitError("rate limited", response=response, body=None) -@pytest.mark.asyncio -async def test_persistent_rate_limit_stops_gracefully( - monkeypatch: pytest.MonkeyPatch, tmp_path: Any, caplog: pytest.LogCaptureFixture -) -> None: - """A persistent RateLimitError stops the scan (root -> 'stopped') without raising.""" +class _ProviderUsageLimitError(Exception): + """A non-OpenAI (e.g. LiteLLM-routed) provider exception whose text carries the + usage-limit marker but which is *not* an ``openai.RateLimitError``.""" + + +def _make_litellm_usage_limit_error() -> Exception: + return _ProviderUsageLimitError( + "litellm.RateLimitError: 429 - {'error': {'type': 'usage_limit_reached', " + "'message': 'The usage limit has been reached'}}" + ) + + +async def _run_scan_raising( + monkeypatch: pytest.MonkeyPatch, tmp_path: Any, error: BaseException +) -> tuple[Any, AgentCoordinator]: + """Drive run_strix_scan with a fully mocked scan whose agent loop raises ``error``.""" monkeypatch.setattr(runner, "run_dir_for", lambda _scan_id: tmp_path) monkeypatch.setattr(runner, "runtime_state_dir", lambda _run_dir: tmp_path) monkeypatch.setattr(runner, "setup_scan_logging", lambda _run_dir: lambda: None) @@ -70,21 +81,24 @@ async def test_persistent_rate_limit_stops_gracefully( monkeypatch.setattr(runner, "make_child_factory", lambda **_kwargs: lambda **_k: object()) monkeypatch.setattr(runner, "open_agent_session", lambda _root_id, _db: object()) - async def _raise_rate_limit(*_args: Any, **_kwargs: Any) -> None: - raise _make_rate_limit_error() + async def _raise(*_args: Any, **_kwargs: Any) -> None: + raise error - monkeypatch.setattr(runner, "run_agent_loop", _raise_rate_limit) + monkeypatch.setattr(runner, "run_agent_loop", _raise) coordinator = AgentCoordinator() + result = await runner.run_strix_scan( + scan_config={"targets": [], "scan_mode": "deep"}, + scan_id="scan-test", + image="img", + coordinator=coordinator, + ) + return result, coordinator - with caplog.at_level(logging.WARNING): - result = await runner.run_strix_scan( - scan_config={"targets": [], "scan_mode": "deep"}, - scan_id="scan-test", - image="img", - coordinator=coordinator, - ) +def _assert_stopped_resumably( + result: Any, coordinator: AgentCoordinator, caplog: pytest.LogCaptureFixture +) -> None: assert result is None root_ids = [aid for aid, parent in coordinator.parent_of.items() if parent is None] assert len(root_ids) == 1 @@ -92,3 +106,29 @@ async def test_persistent_rate_limit_stops_gracefully( # the resume hint must carry the real scan id, not a literal placeholder assert "strix --resume scan-test" in caplog.text assert "" not in caplog.text + + +@pytest.mark.asyncio +async def test_persistent_rate_limit_stops_gracefully( + monkeypatch: pytest.MonkeyPatch, tmp_path: Any, caplog: pytest.LogCaptureFixture +) -> None: + """A persistent OpenAI RateLimitError stops the scan (root -> 'stopped') without raising.""" + with caplog.at_level(logging.WARNING): + result, coordinator = await _run_scan_raising( + monkeypatch, tmp_path, _make_rate_limit_error() + ) + _assert_stopped_resumably(result, coordinator, caplog) + + +@pytest.mark.asyncio +async def test_persistent_usage_limit_non_openai_stops_gracefully( + monkeypatch: pytest.MonkeyPatch, tmp_path: Any, caplog: pytest.LogCaptureFixture +) -> None: + """A usage-limit error from a non-OpenAI (LiteLLM-routed) provider — not an + ``openai.RateLimitError`` — must also land on the resumable stop path, not the + generic failure path.""" + with caplog.at_level(logging.WARNING): + result, coordinator = await _run_scan_raising( + monkeypatch, tmp_path, _make_litellm_usage_limit_error() + ) + _assert_stopped_resumably(result, coordinator, caplog)