From a0f8ea2cd379c4bf56e665d3db2ba8a4f7c0683c Mon Sep 17 00:00:00 2001 From: Ahmed Allam Date: Sun, 18 Oct 2026 11:47:16 +0000 Subject: [PATCH] fix(core): a parked agent keeps its pause reason and drains queued messages on resume BudgetPausedError carries whether the operator pause caused it, so a resume that lands between the hook's refusal and the loop's handler still parks and continues instead of exiting under the stop policy. After the resume the agent drains its mailbox into the session before retrying the call, so a message queued while it was held is not skipped. --- strix/core/execution.py | 5 ++-- strix/core/hooks.py | 6 ++++- tests/test_budget_pause_policy.py | 38 ++++++++++++++++++++++++++++++- 3 files changed, 45 insertions(+), 4 deletions(-) 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()