diff --git a/strix/core/execution.py b/strix/core/execution.py index e3e1c551..3bb9a5fe 100644 --- a/strix/core/execution.py +++ b/strix/core/execution.py @@ -599,7 +599,7 @@ async def _run_until_lifecycle( hooks=hooks, ) except BudgetPausedError as exc: - if coordinator.budget_policy != "pause" and not coordinator.operator_paused: + if coordinator.budget_policy != "pause" and not exc.operator_pause: raise # The agent parked right before an LLM call; everything up to that # point is already in its session. Once resumed, the same call goes @@ -612,6 +612,7 @@ async def _run_until_lifecycle( # Stopped while parked (operator stop or a parent's stop_agent). await coordinator.reset_recovery(agent_id) return result + await coordinator.consume_pending(agent_id) if session is not None: input_data = [] continue @@ -831,7 +832,7 @@ async def _run_cycle( # noqa: PLR0912, PLR0915 await coordinator.detach_stream(agent_id, stream) except BudgetPausedError as exc: logger.info("agent %s paused at the scan budget limit: %s", agent_id, exc) - if coordinator.budget_policy == "pause" or coordinator.operator_paused: + if coordinator.budget_policy == "pause" or exc.operator_pause: await coordinator.park_for_budget(agent_id) else: await coordinator.pause_for_budget(agent_id) diff --git a/strix/core/hooks.py b/strix/core/hooks.py index 55ce5e44..13fc98f9 100644 --- a/strix/core/hooks.py +++ b/strix/core/hooks.py @@ -65,9 +65,12 @@ class BudgetPausedError(RuntimeError): decided to park; the agent waits for a resume newer than that. """ - def __init__(self, message: str, *, resume_epoch: int = 0) -> None: + def __init__( + self, message: str, *, resume_epoch: int = 0, operator_pause: bool = False + ) -> None: super().__init__(message) self.resume_epoch = resume_epoch + self.operator_pause = operator_pause def _validate_budget(max_budget_usd: float | None) -> None: @@ -214,6 +217,7 @@ class ReportUsageHooks(RunHooks[dict[str, Any]]): raise BudgetPausedError( "scan paused; waiting for the operator to resume", resume_epoch=coordinator.resume_epoch, + operator_pause=True, ) def _pause_if_limited(self, context: RunContextWrapper[dict[str, Any]]) -> None: diff --git a/tests/test_budget_pause_policy.py b/tests/test_budget_pause_policy.py index 03cd18b2..c5ddab8a 100644 --- a/tests/test_budget_pause_policy.py +++ b/tests/test_budget_pause_policy.py @@ -21,7 +21,7 @@ from strix.core.sessions import open_agent_session if TYPE_CHECKING: - from collections.abc import AsyncIterator, Callable + from collections.abc import AsyncIterator, Awaitable, Callable from pathlib import Path @@ -31,6 +31,7 @@ _CALL_LATENCY_S = 0.005 class _FakeLedger: def __init__(self) -> None: + self.on_pause: Callable[[], Awaitable[None]] | None = None self.cost = 0.0 self.calls: list[str] = [] self.remaining: dict[str, int] = {} @@ -80,6 +81,8 @@ class _FakeStream: await self._hooks.on_llm_start(ctx_wrapper, self._agent, None, input_items) except BudgetPausedError as exc: self.run_loop_exception = exc + if self._ledger.on_pause is not None: + await self._ledger.on_pause() return self._ledger.warned_inputs.append(input_items) if self._ledger.gate is not None: @@ -583,6 +586,39 @@ async def test_operator_pause_holds_under_the_stop_policy( "child-a": "completed", "child-b": "completed", } + root_items = await scan.sessions[0].get_items() + assert any("keep going" in str(item) for item in root_items) + + for session in scan.sessions: + session.close() + + +@pytest.mark.asyncio +async def test_a_resume_that_lands_before_the_park_still_continues( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + scan = _Scan(tmp_path, monkeypatch, max_budget_usd=None, budget_policy="stop") + resumed = 0 + + async def resume_at_once() -> None: + nonlocal resumed + resumed += 1 + await scan.coordinator.resume_budget() + + scan.ledger.on_pause = resume_at_once + with patch("strix.core.hooks.get_global_report_state", return_value=scan.ledger): + await scan.start_root(calls=6) + await scan.start_child("child-a", calls=6) + await _wait_until(lambda: scan.ledger.cost >= 2.0) + await scan.coordinator.pause_budget() + await asyncio.wait_for(asyncio.gather(*scan.tasks(), return_exceptions=True), timeout=5.0) + assert resumed >= 1 + assert _pause_flags(scan.coordinator) == (False, False) + assert scan.ledger.cost == pytest.approx(12.0) + assert {aid: str(s) for aid, s in scan.coordinator.statuses.items()} == { + "root": "completed", + "child-a": "completed", + } for session in scan.sessions: session.close()