mirror of
https://github.com/usestrix/strix.git
synced 2026-10-24 05:43:56 +00:00
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.
This commit is contained in:
parent
6bd1cead1a
commit
a0f8ea2cd3
3 changed files with 45 additions and 4 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue