mirror of
https://github.com/usestrix/strix.git
synced 2026-08-20 10:33:34 +02:00
feat(runtime): graduated wrap-up warnings, budget reserve, and interactive budget pause/continue (#893)
Co-authored-by: Ahmed Allam <ahmed39652003@gmail.com>
This commit is contained in:
devin-ai-integration[bot]
co-authored by
Ahmed Allam
Ahmed Allam
parent
47617969d3
commit
c55a8fa4ba
@@ -0,0 +1,357 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from strix.core import execution
|
||||
from strix.core.agents import AgentCoordinator
|
||||
from strix.core.execution import _start_child_runner, run_agent_loop
|
||||
from strix.core.hooks import BudgetExceededError, ReportUsageHooks
|
||||
from strix.core.sessions import open_agent_session
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
MAX_BUDGET = 10.0
|
||||
COST_PER_CALL = 1.0
|
||||
|
||||
|
||||
class _FakeLedger:
|
||||
def __init__(self) -> None:
|
||||
self.cost = 0.0
|
||||
self.calls: list[str] = []
|
||||
|
||||
def record_sdk_usage(self, **_kwargs: Any) -> None:
|
||||
return
|
||||
|
||||
def get_total_llm_cost(self) -> float:
|
||||
return self.cost
|
||||
|
||||
|
||||
class _FakeStream:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
ledger: _FakeLedger,
|
||||
hooks: ReportUsageHooks,
|
||||
context: dict[str, Any],
|
||||
agent: Any,
|
||||
) -> None:
|
||||
self._ledger = ledger
|
||||
self._hooks = hooks
|
||||
self._context = context
|
||||
self._agent = agent
|
||||
self.run_loop_exception: BaseException | None = None
|
||||
self.final_output = None
|
||||
|
||||
async def stream_events(self) -> AsyncIterator[Any]:
|
||||
self._ledger.cost += COST_PER_CALL
|
||||
self._ledger.calls.append(str(self._context.get("agent_id")))
|
||||
ctx_wrapper = MagicMock()
|
||||
ctx_wrapper.context = self._context
|
||||
try:
|
||||
await self._hooks.on_llm_end(ctx_wrapper, self._agent, MagicMock())
|
||||
except Exception as exc: # noqa: BLE001
|
||||
self.run_loop_exception = exc
|
||||
items: tuple[Any, ...] = ()
|
||||
for item in items:
|
||||
yield item
|
||||
|
||||
def cancel(self, mode: str = "immediate") -> None: # noqa: ARG002
|
||||
return
|
||||
|
||||
|
||||
def _fake_runner(ledger: _FakeLedger) -> Any:
|
||||
class _FakeRunner:
|
||||
@staticmethod
|
||||
def run_streamed(
|
||||
agent: Any,
|
||||
input: Any, # noqa: A002, ARG004
|
||||
*,
|
||||
run_config: Any, # noqa: ARG004
|
||||
context: dict[str, Any],
|
||||
max_turns: int, # noqa: ARG004
|
||||
session: Any, # noqa: ARG004
|
||||
hooks: ReportUsageHooks,
|
||||
) -> _FakeStream:
|
||||
return _FakeStream(ledger=ledger, hooks=hooks, context=context, agent=agent)
|
||||
|
||||
return _FakeRunner
|
||||
|
||||
|
||||
async def _noop_compact(*_args: Any, **_kwargs: Any) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
async def _wait_until(predicate: Callable[[], bool], *, timeout: float = 5.0) -> None:
|
||||
async def _poll() -> None:
|
||||
while not predicate():
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
await asyncio.wait_for(_poll(), timeout=timeout)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_budget_lifecycle_reserve_then_cap( # noqa: PLR0915
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
ledger = _FakeLedger()
|
||||
hooks = ReportUsageHooks(model="test-model", max_budget_usd=MAX_BUDGET)
|
||||
monkeypatch.setattr(execution, "Runner", _fake_runner(ledger))
|
||||
monkeypatch.setattr(execution, "_compact_session", _noop_compact)
|
||||
|
||||
coordinator = AgentCoordinator()
|
||||
db_path = tmp_path / "agents.sqlite"
|
||||
sessions: list[Any] = []
|
||||
run_config = MagicMock()
|
||||
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
root_session = open_agent_session("root", db_path)
|
||||
sessions.append(root_session)
|
||||
|
||||
root_exc: list[BaseException] = []
|
||||
|
||||
async def _root_loop() -> None:
|
||||
try:
|
||||
await run_agent_loop(
|
||||
agent=MagicMock(),
|
||||
initial_input=[],
|
||||
run_config=run_config,
|
||||
context={"agent_id": "root", "parent_id": None},
|
||||
max_turns=500,
|
||||
coordinator=coordinator,
|
||||
agent_id="root",
|
||||
interactive=True,
|
||||
session=root_session,
|
||||
start_parked=True,
|
||||
hooks=hooks,
|
||||
)
|
||||
except BaseException as exc:
|
||||
root_exc.append(exc)
|
||||
raise
|
||||
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=ledger):
|
||||
root_task = asyncio.create_task(_root_loop())
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
for child_id in ("child-a", "child-b"):
|
||||
await coordinator.register(child_id, "recon", parent_id="root")
|
||||
await _start_child_runner(
|
||||
parent_ctx={"agent_id": "root", "parent_id": None},
|
||||
coordinator=coordinator,
|
||||
agents_db_path=db_path,
|
||||
sessions_to_close=sessions,
|
||||
run_config=run_config,
|
||||
max_turns=500,
|
||||
interactive=True,
|
||||
child_agent=MagicMock(),
|
||||
child_id=child_id,
|
||||
name=f"recon-{child_id}",
|
||||
parent_id="root",
|
||||
task="probe things",
|
||||
initial_input=[],
|
||||
hooks=hooks,
|
||||
)
|
||||
await _wait_until(lambda: ledger.cost >= 2.0)
|
||||
reserve_before = coordinator.reserve_stopped
|
||||
assert reserve_before is False
|
||||
|
||||
async def _wait_spend_above(amount: float) -> None:
|
||||
await _wait_until(lambda: ledger.cost > amount)
|
||||
|
||||
turn = 0
|
||||
while ledger.cost < MAX_BUDGET * 0.90 - 1e-9:
|
||||
target = ("child-a", "child-b")[turn % 2]
|
||||
spent_before = ledger.cost
|
||||
assert await coordinator.send(target, {"from": "user", "content": "keep going"})
|
||||
await _wait_spend_above(spent_before)
|
||||
turn += 1
|
||||
|
||||
await _wait_until(lambda: coordinator.reserve_stopped)
|
||||
|
||||
await _wait_until(
|
||||
lambda: (
|
||||
coordinator.statuses["child-a"] == "stopped"
|
||||
and coordinator.statuses["child-b"] == "stopped"
|
||||
)
|
||||
)
|
||||
|
||||
assert coordinator.reserve_stopped is True
|
||||
|
||||
await _wait_until(lambda: coordinator.budget_stopped)
|
||||
assert ledger.cost == pytest.approx(MAX_BUDGET)
|
||||
|
||||
assert len(ledger.calls) == 10
|
||||
assert set(ledger.calls[:9]) == {"child-a", "child-b"}
|
||||
assert ledger.calls[9] == "root"
|
||||
|
||||
root_items = await root_session.get_items()
|
||||
notices = [item for item in root_items if "Budget reserve" in str(item)]
|
||||
assert len(notices) == 1
|
||||
|
||||
with pytest.raises(BudgetExceededError):
|
||||
await root_task
|
||||
assert root_exc and isinstance(root_exc[0], BudgetExceededError)
|
||||
|
||||
assert {aid: str(status) for aid, status in coordinator.statuses.items()} == {
|
||||
"root": "stopped",
|
||||
"child-a": "stopped",
|
||||
"child-b": "stopped",
|
||||
}
|
||||
assert coordinator.budget_stopped is True
|
||||
assert coordinator.reserve_stopped is True
|
||||
|
||||
for session in sessions:
|
||||
session.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_respawned_children_after_reserve_never_spend(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
ledger = _FakeLedger()
|
||||
ledger.cost = 9.5
|
||||
hooks = ReportUsageHooks(model="test-model", max_budget_usd=MAX_BUDGET)
|
||||
monkeypatch.setattr(execution, "Runner", _fake_runner(ledger))
|
||||
monkeypatch.setattr(execution, "_compact_session", _noop_compact)
|
||||
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.register("child-a", "recon", parent_id="root")
|
||||
snap = await coordinator.snapshot()
|
||||
snap["reserve_stopped"] = True
|
||||
|
||||
restored = AgentCoordinator()
|
||||
await restored.restore(snap)
|
||||
assert restored.reserve_stopped is True
|
||||
|
||||
sessions: list[Any] = []
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=ledger):
|
||||
await _start_child_runner(
|
||||
parent_ctx={"agent_id": "root", "parent_id": None},
|
||||
coordinator=restored,
|
||||
agents_db_path=tmp_path / "agents.sqlite",
|
||||
sessions_to_close=sessions,
|
||||
run_config=MagicMock(),
|
||||
max_turns=500,
|
||||
interactive=True,
|
||||
child_agent=MagicMock(),
|
||||
child_id="child-a",
|
||||
name="recon-child-a",
|
||||
parent_id="root",
|
||||
task="probe things",
|
||||
initial_input=[],
|
||||
hooks=hooks,
|
||||
)
|
||||
await _wait_until(lambda: restored.statuses["child-a"] == "stopped")
|
||||
|
||||
assert ledger.cost == pytest.approx(9.5)
|
||||
assert ledger.calls == []
|
||||
for session in sessions:
|
||||
session.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resumed_parked_root_after_reserve_is_renotified_and_finalizes(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
ledger = _FakeLedger()
|
||||
ledger.cost = 9.0
|
||||
hooks = ReportUsageHooks(model="test-model", max_budget_usd=MAX_BUDGET)
|
||||
monkeypatch.setattr(execution, "Runner", _fake_runner(ledger))
|
||||
monkeypatch.setattr(execution, "_compact_session", _noop_compact)
|
||||
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.set_status("root", "waiting")
|
||||
snap = await coordinator.snapshot()
|
||||
snap["reserve_stopped"] = True
|
||||
|
||||
restored = AgentCoordinator()
|
||||
await restored.restore(snap)
|
||||
assert restored.reserve_stopped is True
|
||||
|
||||
root_session = open_agent_session("root", tmp_path / "agents.sqlite")
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=ledger):
|
||||
root_task = asyncio.create_task(
|
||||
run_agent_loop(
|
||||
agent=MagicMock(),
|
||||
initial_input=[],
|
||||
run_config=MagicMock(),
|
||||
context={"agent_id": "root", "parent_id": None},
|
||||
max_turns=500,
|
||||
coordinator=restored,
|
||||
agent_id="root",
|
||||
interactive=True,
|
||||
session=root_session,
|
||||
start_parked=True,
|
||||
hooks=hooks,
|
||||
)
|
||||
)
|
||||
with pytest.raises(BudgetExceededError):
|
||||
await asyncio.wait_for(root_task, timeout=5.0)
|
||||
|
||||
assert ledger.calls == ["root"]
|
||||
assert ledger.cost == pytest.approx(MAX_BUDGET)
|
||||
root_items = await root_session.get_items()
|
||||
notices = [item for item in root_items if "Budget reserve" in str(item)]
|
||||
assert len(notices) == 1
|
||||
root_session.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_interactive_budget_pause_then_user_message_extends_and_resumes(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
ledger = _FakeLedger()
|
||||
ledger.cost = 9.0
|
||||
hooks = ReportUsageHooks(model="test-model", max_budget_usd=MAX_BUDGET, interactive=True)
|
||||
monkeypatch.setattr(execution, "Runner", _fake_runner(ledger))
|
||||
monkeypatch.setattr(execution, "_compact_session", _noop_compact)
|
||||
|
||||
coordinator = AgentCoordinator()
|
||||
coordinator.set_budget_extender(hooks.extend_budget)
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
root_session = open_agent_session("root", tmp_path / "agents.sqlite")
|
||||
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=ledger):
|
||||
root_task = asyncio.create_task(
|
||||
run_agent_loop(
|
||||
agent=MagicMock(),
|
||||
initial_input=[],
|
||||
run_config=MagicMock(),
|
||||
context={"agent_id": "root", "parent_id": None},
|
||||
max_turns=500,
|
||||
coordinator=coordinator,
|
||||
agent_id="root",
|
||||
interactive=True,
|
||||
session=root_session,
|
||||
start_parked=True,
|
||||
hooks=hooks,
|
||||
)
|
||||
)
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
assert await coordinator.send("root", {"from": "user", "content": "go"})
|
||||
await _wait_until(lambda: coordinator.budget_paused)
|
||||
assert coordinator.statuses["root"] == "budget_paused"
|
||||
assert ledger.cost == pytest.approx(MAX_BUDGET)
|
||||
assert not root_task.done()
|
||||
assert coordinator.budget_stopped is False
|
||||
|
||||
assert await coordinator.send("root", {"from": "user", "content": "keep going"})
|
||||
await _wait_until(lambda: not coordinator.budget_paused)
|
||||
await _wait_until(lambda: ledger.cost > MAX_BUDGET)
|
||||
await _wait_until(lambda: coordinator.statuses["root"] == "waiting")
|
||||
assert not root_task.done()
|
||||
|
||||
root_task.cancel()
|
||||
await root_task
|
||||
|
||||
root_session.close()
|
||||
@@ -3,10 +3,265 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from agents.memory import SQLiteSession
|
||||
from agents.tool_context import ToolContext
|
||||
|
||||
from strix.core.agents import AgentCoordinator
|
||||
from strix.core.execution import _notify_root_on_budget_reserve
|
||||
from strix.tools.finish.tool import finish_scan
|
||||
|
||||
|
||||
async def _call_finish_scan(
|
||||
coordinator: AgentCoordinator, agent_id: str, parent_id: str | None
|
||||
) -> dict[str, Any]:
|
||||
ctx = ToolContext(
|
||||
context={"coordinator": coordinator, "agent_id": agent_id, "parent_id": parent_id},
|
||||
tool_name="finish_scan",
|
||||
tool_call_id="call-1",
|
||||
tool_arguments="{}",
|
||||
)
|
||||
fields = ("executive_summary", "methodology", "technical_analysis", "recommendations")
|
||||
result: str = await finish_scan.on_invoke_tool(ctx, json.dumps(dict.fromkeys(fields, "x")))
|
||||
parsed: dict[str, Any] = json.loads(result)
|
||||
return parsed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reserve_stop_notifies_root_once(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.register("child-a", "recon", parent_id="root")
|
||||
await coordinator.register("child-b", "recon", parent_id="root")
|
||||
|
||||
sent: list[tuple[str, dict[str, Any]]] = []
|
||||
|
||||
async def _record(target_agent_id: str, message: dict[str, Any]) -> bool:
|
||||
sent.append((target_agent_id, message))
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(coordinator, "send", _record)
|
||||
|
||||
await _notify_root_on_budget_reserve(coordinator)
|
||||
await _notify_root_on_budget_reserve(coordinator)
|
||||
|
||||
assert len(sent) == 1
|
||||
target, message = sent[0]
|
||||
assert target == "root"
|
||||
assert message["type"] == "budget_reserve_stop"
|
||||
assert "finish_scan" in str(message["content"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_reserve_claims_yield_single_root() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
for i in range(12):
|
||||
await coordinator.register(f"child-{i}", "recon", parent_id="root")
|
||||
|
||||
results = await asyncio.gather(*(coordinator.claim_reserve_notification() for _ in range(12)))
|
||||
|
||||
assert results.count("root") == 1
|
||||
assert all(r is None for r in results if r != "root")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_claim_reserve_sets_flag_and_wakes_parked_agents() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.register("child", "recon", parent_id="root")
|
||||
|
||||
flag_before = coordinator.reserve_stopped
|
||||
assert flag_before is False
|
||||
waiter = asyncio.create_task(coordinator.wait_for_message("child"))
|
||||
await asyncio.sleep(0)
|
||||
assert not waiter.done()
|
||||
|
||||
await coordinator.claim_reserve_notification()
|
||||
|
||||
flag_after = coordinator.reserve_stopped
|
||||
assert flag_after is True
|
||||
await asyncio.wait_for(waiter, timeout=1.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finish_scan_bypasses_active_agent_guard_after_reserve() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.register("child", "recon", parent_id="root")
|
||||
await coordinator.set_status("child", "running")
|
||||
|
||||
blocked = await _call_finish_scan(coordinator, "root", None)
|
||||
assert blocked["scan_completed"] is False
|
||||
assert blocked["active_agents"]
|
||||
|
||||
await coordinator.claim_reserve_notification()
|
||||
|
||||
finished = await _call_finish_scan(coordinator, "root", None)
|
||||
assert finished["scan_completed"] is True
|
||||
assert coordinator.statuses["root"] == "completed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finish_scan_gate_ignores_sub_agent_caller() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.register("child", "recon", parent_id="root")
|
||||
await coordinator.set_status("child", "running")
|
||||
|
||||
result = await _call_finish_scan(coordinator, "child", "root")
|
||||
assert "active_agents" not in result
|
||||
assert result["success"] is False
|
||||
assert "root" in result["error"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reserve_stop_notify_noop_without_root(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("child", "recon", parent_id="missing")
|
||||
|
||||
sent: list[tuple[str, dict[str, Any]]] = []
|
||||
|
||||
async def _record(target_agent_id: str, message: dict[str, Any]) -> bool:
|
||||
sent.append((target_agent_id, message))
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(coordinator, "send", _record)
|
||||
await _notify_root_on_budget_reserve(coordinator)
|
||||
|
||||
assert sent == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_snapshot_round_trip_preserves_stop_flags() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.trigger_budget_stop()
|
||||
await coordinator.claim_reserve_notification()
|
||||
|
||||
snap = await coordinator.snapshot()
|
||||
assert snap["budget_stopped"] is True
|
||||
assert snap["reserve_stopped"] is True
|
||||
|
||||
restored = AgentCoordinator()
|
||||
await restored.restore(snap)
|
||||
assert restored.budget_stopped is True
|
||||
assert restored.reserve_stopped is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_legacy_snapshot_without_stop_flags_defaults_to_false() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
snap = await coordinator.snapshot()
|
||||
del snap["budget_stopped"]
|
||||
del snap["reserve_stopped"]
|
||||
|
||||
restored = AgentCoordinator()
|
||||
await restored.restore(snap)
|
||||
assert restored.budget_stopped is False
|
||||
assert restored.reserve_stopped is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_randomized_reserve_claim_race_many_interleavings() -> None:
|
||||
for seed in range(25):
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
child_ids = [f"child-{i}" for i in range(8)]
|
||||
for child_id in child_ids:
|
||||
await coordinator.register(child_id, "recon", parent_id="root")
|
||||
|
||||
waiters = [asyncio.create_task(coordinator.wait_for_message(cid)) for cid in child_ids]
|
||||
await asyncio.sleep(0)
|
||||
|
||||
async def _claim(delay: float, coord: AgentCoordinator = coordinator) -> str | None:
|
||||
await asyncio.sleep(delay)
|
||||
return await coord.claim_reserve_notification()
|
||||
|
||||
delays = [((seed * 31 + i * 17) % 50) / 10_000 for i in range(len(child_ids))]
|
||||
results = await asyncio.gather(*(_claim(delay) for delay in delays))
|
||||
|
||||
assert results.count("root") == 1, f"seed {seed}: expected exactly one winner"
|
||||
await asyncio.wait_for(asyncio.gather(*waiters), timeout=1.0)
|
||||
assert coordinator.reserve_stopped is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reserve_claim_never_loses_root_wake() -> None:
|
||||
for _ in range(10):
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.register("child", "recon", parent_id="root")
|
||||
|
||||
root_waiter = asyncio.create_task(coordinator.wait_for_message("root"))
|
||||
await asyncio.sleep(0)
|
||||
assert not root_waiter.done()
|
||||
|
||||
await coordinator.claim_reserve_notification()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
async with coordinator._lock:
|
||||
coordinator.pending_counts["root"] = 1
|
||||
coordinator.runtimes["root"].wake.set()
|
||||
|
||||
await asyncio.wait_for(root_waiter, timeout=1.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_stop_takes_precedence_over_reserve_for_all_roles() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.register("child", "recon", parent_id="root")
|
||||
await coordinator.claim_reserve_notification()
|
||||
await coordinator.trigger_budget_stop()
|
||||
|
||||
await asyncio.wait_for(coordinator.wait_for_message("root"), timeout=1.0)
|
||||
await asyncio.wait_for(coordinator.wait_for_message("child"), timeout=1.0)
|
||||
assert coordinator.budget_stopped is True
|
||||
assert coordinator.reserve_stopped is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_root_not_released_by_reserve_alone() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.register("child", "recon", parent_id="root")
|
||||
|
||||
await coordinator.claim_reserve_notification()
|
||||
|
||||
root_waiter = asyncio.create_task(coordinator.wait_for_message("root"))
|
||||
await asyncio.sleep(0.02)
|
||||
assert not root_waiter.done()
|
||||
|
||||
root_waiter.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await root_waiter
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_snapshot_during_concurrent_claims_is_consistent() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
for i in range(6):
|
||||
await coordinator.register(f"child-{i}", "recon", parent_id="root")
|
||||
|
||||
claims = [asyncio.create_task(coordinator.claim_reserve_notification()) for _ in range(6)]
|
||||
snap = await coordinator.snapshot()
|
||||
await asyncio.gather(*claims)
|
||||
|
||||
assert isinstance(snap["reserve_stopped"], bool)
|
||||
final_snap = await coordinator.snapshot()
|
||||
assert final_snap["reserve_stopped"] is True
|
||||
|
||||
restored = AgentCoordinator()
|
||||
await restored.restore(final_snap)
|
||||
assert restored.reserve_stopped is True
|
||||
assert await restored.claim_reserve_notification() is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -42,3 +297,133 @@ async def test_wait_for_message_returns_immediately_after_budget_stop() -> None:
|
||||
|
||||
# No pending messages, but the stop flag short-circuits the wait.
|
||||
await asyncio.wait_for(coordinator.wait_for_message("agent"), timeout=1.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pause_for_budget_sets_flag_and_status() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
|
||||
await coordinator.pause_for_budget("root")
|
||||
assert coordinator.budget_paused is True
|
||||
assert coordinator.statuses["root"] == "budget_paused"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resume_from_budget_pause_extends_and_nudges(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.register("child-a", "recon", parent_id="root")
|
||||
await coordinator.register("child-b", "recon", parent_id="root")
|
||||
await coordinator.pause_for_budget("root")
|
||||
await coordinator.pause_for_budget("child-a")
|
||||
await coordinator.pause_for_budget("child-b")
|
||||
|
||||
extensions: list[int] = []
|
||||
coordinator.set_budget_extender(lambda: extensions.append(1))
|
||||
|
||||
sent: list[tuple[str, dict[str, Any]]] = []
|
||||
|
||||
async def _record(target_agent_id: str, message: dict[str, Any]) -> bool:
|
||||
sent.append((target_agent_id, message))
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(coordinator, "send", _record)
|
||||
|
||||
await coordinator.resume_from_budget_pause(exclude="root")
|
||||
|
||||
assert coordinator.budget_paused is False
|
||||
assert len(extensions) == 1
|
||||
assert all(coordinator.statuses[aid] == "waiting" for aid in ("root", "child-a", "child-b"))
|
||||
assert sorted(target for target, _ in sent) == ["child-a", "child-b"]
|
||||
assert all(message["type"] == "budget_extended" for _, message in sent)
|
||||
|
||||
await coordinator.resume_from_budget_pause(exclude="root")
|
||||
assert len(extensions) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_send_resumes_budget_pause(tmp_path: Any) -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
session = SQLiteSession("root", tmp_path / "agents.db")
|
||||
await coordinator.attach_runtime("root", session=session)
|
||||
await coordinator.pause_for_budget("root")
|
||||
|
||||
extensions: list[int] = []
|
||||
coordinator.set_budget_extender(lambda: extensions.append(1))
|
||||
|
||||
delivered = await coordinator.send("root", {"from": "user", "content": "keep going"})
|
||||
|
||||
assert delivered is True
|
||||
assert coordinator.budget_paused is False
|
||||
assert len(extensions) == 1
|
||||
assert coordinator.statuses["root"] == "waiting"
|
||||
assert coordinator.pending_counts["root"] == 1
|
||||
session.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_user_send_does_not_resume_budget_pause(tmp_path: Any) -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
session = SQLiteSession("root", tmp_path / "agents.db")
|
||||
await coordinator.attach_runtime("root", session=session)
|
||||
await coordinator.pause_for_budget("root")
|
||||
|
||||
extensions: list[int] = []
|
||||
coordinator.set_budget_extender(lambda: extensions.append(1))
|
||||
|
||||
await coordinator.send("root", {"from": "system", "content": "status"})
|
||||
|
||||
assert coordinator.budget_paused is True
|
||||
assert extensions == []
|
||||
assert coordinator.statuses["root"] == "budget_paused"
|
||||
session.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_budget_stops_clears_pause_and_normalizes_statuses() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.trigger_budget_stop()
|
||||
await coordinator.claim_reserve_notification()
|
||||
await coordinator.pause_for_budget("root")
|
||||
|
||||
await coordinator.reset_budget_stops(budget_stopped=False, reserve_stopped=False)
|
||||
|
||||
assert coordinator.budget_stopped is False
|
||||
assert coordinator.reserve_stopped is False
|
||||
assert coordinator.budget_paused is False
|
||||
assert coordinator.statuses["root"] == "waiting"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_budget_stops_can_preserve_pause() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.pause_for_budget("root")
|
||||
|
||||
await coordinator.reset_budget_stops(
|
||||
budget_stopped=False, reserve_stopped=False, budget_paused=True
|
||||
)
|
||||
|
||||
assert coordinator.budget_paused is True
|
||||
assert coordinator.statuses["root"] == "budget_paused"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_snapshot_round_trip_preserves_budget_pause() -> None:
|
||||
coordinator = AgentCoordinator()
|
||||
await coordinator.register("root", "strix", parent_id=None)
|
||||
await coordinator.pause_for_budget("root")
|
||||
|
||||
snap = await coordinator.snapshot()
|
||||
assert snap["budget_paused"] is True
|
||||
|
||||
restored = AgentCoordinator()
|
||||
await restored.restore(snap)
|
||||
assert restored.budget_paused is True
|
||||
assert restored.statuses["root"] == "budget_paused"
|
||||
|
||||
+388
-3
@@ -2,11 +2,18 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from strix.core.hooks import BudgetExceededError, ReportUsageHooks
|
||||
from strix.core.hooks import (
|
||||
BudgetExceededError,
|
||||
BudgetPausedError,
|
||||
ReportUsageHooks,
|
||||
SubagentBudgetReservedError,
|
||||
recomputed_budget_flags,
|
||||
)
|
||||
|
||||
|
||||
def _make_hooks(max_budget: float | None) -> ReportUsageHooks:
|
||||
@@ -20,9 +27,22 @@ def _make_report_state(cost: float) -> MagicMock:
|
||||
return state
|
||||
|
||||
|
||||
def _make_context(agent_id: str = "test-agent") -> MagicMock:
|
||||
def _make_context(agent_id: str = "test-agent", parent_id: str | None = None) -> MagicMock:
|
||||
ctx: MagicMock = MagicMock()
|
||||
ctx.context = {"agent_id": agent_id}
|
||||
ctx.context = {"agent_id": agent_id, "parent_id": parent_id}
|
||||
return ctx
|
||||
|
||||
|
||||
def _make_warn_context(
|
||||
*,
|
||||
requests: int,
|
||||
parent_id: str | None = None,
|
||||
agent_id: str = "test-agent",
|
||||
) -> MagicMock:
|
||||
ctx: MagicMock = MagicMock()
|
||||
ctx.context = {"agent_id": agent_id, "parent_id": parent_id}
|
||||
ctx.usage = MagicMock()
|
||||
ctx.usage.requests = requests
|
||||
return ctx
|
||||
|
||||
|
||||
@@ -89,6 +109,127 @@ async def test_error_message_includes_amounts() -> None:
|
||||
assert "7.1234" in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subagent_stops_at_reserve() -> None:
|
||||
hooks = _make_hooks(10.0)
|
||||
state = _make_report_state(9.0)
|
||||
with (
|
||||
patch("strix.core.hooks.get_global_report_state", return_value=state),
|
||||
pytest.raises(SubagentBudgetReservedError),
|
||||
):
|
||||
await hooks.on_llm_end(_make_context(parent_id="root-1"), MagicMock(), MagicMock())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subagent_below_reserve_does_not_raise() -> None:
|
||||
hooks = _make_hooks(10.0)
|
||||
state = _make_report_state(8.99)
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=state):
|
||||
await hooks.on_llm_end(_make_context(parent_id="root-1"), MagicMock(), MagicMock())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subagent_overshoot_to_full_budget_triggers_scan_wide_stop() -> None:
|
||||
hooks = _make_hooks(10.0)
|
||||
state = _make_report_state(10.5)
|
||||
with (
|
||||
patch("strix.core.hooks.get_global_report_state", return_value=state),
|
||||
pytest.raises(BudgetExceededError),
|
||||
):
|
||||
await hooks.on_llm_end(_make_context(parent_id="root-1"), MagicMock(), MagicMock())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_root_keeps_running_inside_reserve() -> None:
|
||||
hooks = _make_hooks(10.0)
|
||||
state = _make_report_state(9.5)
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=state):
|
||||
await hooks.on_llm_end(_make_context(parent_id=None), MagicMock(), MagicMock())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_root_hard_stop_stays_at_full_budget() -> None:
|
||||
hooks = _make_hooks(10.0)
|
||||
state = _make_report_state(10.0)
|
||||
with (
|
||||
patch("strix.core.hooks.get_global_report_state", return_value=state),
|
||||
pytest.raises(BudgetExceededError),
|
||||
):
|
||||
await hooks.on_llm_end(_make_context(parent_id=None), MagicMock(), MagicMock())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_warning_mentions_reserve() -> None:
|
||||
hooks = ReportUsageHooks(model="test-model", max_budget_usd=10.0)
|
||||
state = _make_report_state(7.5)
|
||||
root_items: list[Any] = []
|
||||
sub_items: list[Any] = []
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=state):
|
||||
await hooks.on_llm_start(
|
||||
_make_warn_context(requests=0, parent_id=None), MagicMock(), None, root_items
|
||||
)
|
||||
await hooks.on_llm_start(
|
||||
_make_warn_context(requests=0, parent_id="root-1"), MagicMock(), None, sub_items
|
||||
)
|
||||
assert "stopped at 90%" in root_items[0]["content"]
|
||||
assert "stopped at 90%" in sub_items[0]["content"]
|
||||
assert "root agent's final report" in sub_items[0]["content"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subagent_critical_budget_warning_reachable_before_reserve() -> None:
|
||||
hooks = ReportUsageHooks(model="test-model", max_budget_usd=10.0)
|
||||
state = _make_report_state(8.6)
|
||||
sub_items: list[Any] = []
|
||||
root_items: list[Any] = []
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=state):
|
||||
await hooks.on_llm_start(
|
||||
_make_warn_context(requests=0, parent_id="root-1"), MagicMock(), None, sub_items
|
||||
)
|
||||
await hooks.on_llm_start(
|
||||
_make_warn_context(requests=0, parent_id=None), MagicMock(), None, root_items
|
||||
)
|
||||
assert "[CRITICAL]" in sub_items[0]["content"]
|
||||
assert "[URGENT]" in root_items[0]["content"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("parent_id", "cost", "expected"),
|
||||
[
|
||||
("root-1", 0.0, None),
|
||||
("root-1", 8.9999, None),
|
||||
("root-1", 9.0, SubagentBudgetReservedError),
|
||||
("root-1", 9.0001, SubagentBudgetReservedError),
|
||||
("root-1", 9.5, SubagentBudgetReservedError),
|
||||
("root-1", 9.9999, SubagentBudgetReservedError),
|
||||
("root-1", 10.0, BudgetExceededError),
|
||||
("root-1", 10.0001, BudgetExceededError),
|
||||
("root-1", 25.0, BudgetExceededError),
|
||||
(None, 0.0, None),
|
||||
(None, 8.9999, None),
|
||||
(None, 9.0, None),
|
||||
(None, 9.5, None),
|
||||
(None, 9.9999, None),
|
||||
(None, 10.0, BudgetExceededError),
|
||||
(None, 10.0001, BudgetExceededError),
|
||||
(None, 25.0, BudgetExceededError),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_enforcement_decision_table(
|
||||
parent_id: str | None, cost: float, expected: type[Exception] | None
|
||||
) -> None:
|
||||
hooks = _make_hooks(10.0)
|
||||
state = _make_report_state(cost)
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=state):
|
||||
if expected is None:
|
||||
await hooks.on_llm_end(_make_context(parent_id=parent_id), MagicMock(), MagicMock())
|
||||
else:
|
||||
with pytest.raises(expected):
|
||||
await hooks.on_llm_end(_make_context(parent_id=parent_id), MagicMock(), MagicMock())
|
||||
state.record_sdk_usage.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_raise_when_report_state_none() -> None:
|
||||
hooks = _make_hooks(1.0)
|
||||
@@ -106,3 +247,247 @@ def test_non_positive_budget_rejected(bad_budget: float) -> None:
|
||||
def test_budget_exceeded_error_is_runtime_error() -> None:
|
||||
err = BudgetExceededError("test")
|
||||
assert isinstance(err, RuntimeError)
|
||||
|
||||
|
||||
def test_non_positive_max_turns_rejected() -> None:
|
||||
with pytest.raises(ValueError, match="positive integer"):
|
||||
ReportUsageHooks(model="test-model", max_turns=0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_turn_warning_below_first_band() -> None:
|
||||
hooks = ReportUsageHooks(model="test-model", max_turns=100)
|
||||
items: list[Any] = []
|
||||
await hooks.on_llm_start(_make_warn_context(requests=68), MagicMock(), None, items)
|
||||
assert items == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_turn_warning_notice_band() -> None:
|
||||
hooks = ReportUsageHooks(model="test-model", max_turns=100)
|
||||
items: list[Any] = []
|
||||
await hooks.on_llm_start(_make_warn_context(requests=69), MagicMock(), None, items)
|
||||
assert len(items) == 1
|
||||
content = items[0]["content"]
|
||||
assert "[NOTICE]" in content
|
||||
assert "finish_scan" in content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_turn_warning_escalates_and_names_subagent_tool() -> None:
|
||||
hooks = ReportUsageHooks(model="test-model", max_turns=100)
|
||||
items: list[Any] = []
|
||||
await hooks.on_llm_start(
|
||||
_make_warn_context(requests=95, parent_id="root-1"), MagicMock(), None, items
|
||||
)
|
||||
assert len(items) == 1
|
||||
content = items[0]["content"]
|
||||
assert "[CRITICAL]" in content
|
||||
assert "agent_finish" in content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_turn_warning_root_directive_distinct_from_subagent() -> None:
|
||||
hooks = ReportUsageHooks(model="test-model", max_turns=100)
|
||||
|
||||
root_items: list[Any] = []
|
||||
await hooks.on_llm_start(
|
||||
_make_warn_context(requests=85, parent_id=None), MagicMock(), None, root_items
|
||||
)
|
||||
root = root_items[0]["content"]
|
||||
|
||||
sub_items: list[Any] = []
|
||||
await hooks.on_llm_start(
|
||||
_make_warn_context(requests=85, parent_id="root-1"), MagicMock(), None, sub_items
|
||||
)
|
||||
sub = sub_items[0]["content"]
|
||||
|
||||
assert root != sub
|
||||
assert "root agent" in root
|
||||
assert "finish_scan" in root
|
||||
assert "agent_finish" not in root
|
||||
assert "whole scan" in root
|
||||
assert "sub-agent" in sub
|
||||
assert "agent_finish" in sub
|
||||
assert "finish_scan" not in sub
|
||||
assert "confirmed" in sub
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_warning_root_directive_distinct_from_subagent() -> None:
|
||||
hooks = ReportUsageHooks(model="test-model", max_budget_usd=10.0)
|
||||
state = _make_report_state(8.6)
|
||||
|
||||
root_items: list[Any] = []
|
||||
sub_items: list[Any] = []
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=state):
|
||||
await hooks.on_llm_start(
|
||||
_make_warn_context(requests=0, parent_id=None), MagicMock(), None, root_items
|
||||
)
|
||||
await hooks.on_llm_start(
|
||||
_make_warn_context(requests=0, parent_id="root-1"), MagicMock(), None, sub_items
|
||||
)
|
||||
|
||||
root = root_items[0]["content"]
|
||||
sub = sub_items[0]["content"]
|
||||
assert "finish_scan" in root and "agent_finish" not in root
|
||||
assert "agent_finish" in sub and "finish_scan" not in sub
|
||||
assert "confirmed" in sub
|
||||
|
||||
|
||||
@pytest.mark.parametrize("parent_id", [None, "root-1"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_turn_warning_directive_escalates_per_stage(parent_id: str | None) -> None:
|
||||
hooks = ReportUsageHooks(model="test-model", max_turns=100)
|
||||
contents: dict[str, str] = {}
|
||||
for label, requests in (("notice", 69), ("urgent", 85), ("critical", 95)):
|
||||
items: list[Any] = []
|
||||
await hooks.on_llm_start(
|
||||
_make_warn_context(requests=requests, parent_id=parent_id), MagicMock(), None, items
|
||||
)
|
||||
contents[label] = items[0]["content"]
|
||||
|
||||
assert len({contents["notice"], contents["urgent"], contents["critical"]}) == 3
|
||||
assert "[NOTICE]" in contents["notice"] and "begin planning" in contents["notice"]
|
||||
assert "[URGENT]" in contents["urgent"] and "prioritize" in contents["urgent"]
|
||||
assert "[CRITICAL]" in contents["critical"] and "STOP" in contents["critical"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_turn_warning_when_max_turns_unset() -> None:
|
||||
hooks = ReportUsageHooks(model="test-model")
|
||||
items: list[Any] = []
|
||||
await hooks.on_llm_start(_make_warn_context(requests=999), MagicMock(), None, items)
|
||||
assert items == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_budget_warning_below_first_band() -> None:
|
||||
hooks = ReportUsageHooks(model="test-model", max_budget_usd=10.0)
|
||||
state = _make_report_state(6.9)
|
||||
items: list[Any] = []
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=state):
|
||||
await hooks.on_llm_start(_make_warn_context(requests=0), MagicMock(), None, items)
|
||||
assert items == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_warning_broadcast_content() -> None:
|
||||
hooks = ReportUsageHooks(model="test-model", max_budget_usd=10.0)
|
||||
state = _make_report_state(9.6)
|
||||
items: list[Any] = []
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=state):
|
||||
await hooks.on_llm_start(_make_warn_context(requests=0), MagicMock(), None, items)
|
||||
assert len(items) == 1
|
||||
content = items[0]["content"]
|
||||
assert "[CRITICAL]" in content
|
||||
assert "shared across every agent" in content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_turn_and_budget_warnings_stack() -> None:
|
||||
hooks = ReportUsageHooks(model="test-model", max_budget_usd=10.0, max_turns=100)
|
||||
state = _make_report_state(8.6)
|
||||
items: list[Any] = []
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=state):
|
||||
await hooks.on_llm_start(_make_warn_context(requests=89), MagicMock(), None, items)
|
||||
assert len(items) == 2
|
||||
joined = " ".join(i["content"] for i in items)
|
||||
assert "Turn budget" in joined
|
||||
assert "cost budget" in joined
|
||||
|
||||
|
||||
def _make_interactive_hooks(max_budget: float | None) -> ReportUsageHooks:
|
||||
return ReportUsageHooks(model="test-model", max_budget_usd=max_budget, interactive=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_interactive_at_budget_pauses_instead_of_stopping() -> None:
|
||||
hooks = _make_interactive_hooks(10.0)
|
||||
state = _make_report_state(10.0)
|
||||
with (
|
||||
patch("strix.core.hooks.get_global_report_state", return_value=state),
|
||||
pytest.raises(BudgetPausedError),
|
||||
):
|
||||
await hooks.on_llm_end(_make_context(parent_id=None), MagicMock(), MagicMock())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_interactive_subagent_has_no_reserve() -> None:
|
||||
hooks = _make_interactive_hooks(10.0)
|
||||
state = _make_report_state(9.5)
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=state):
|
||||
await hooks.on_llm_end(_make_context(parent_id="root-1"), MagicMock(), MagicMock())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_interactive_subagent_pauses_at_full_budget() -> None:
|
||||
hooks = _make_interactive_hooks(10.0)
|
||||
state = _make_report_state(10.5)
|
||||
with (
|
||||
patch("strix.core.hooks.get_global_report_state", return_value=state),
|
||||
pytest.raises(BudgetPausedError),
|
||||
):
|
||||
await hooks.on_llm_end(_make_context(parent_id="root-1"), MagicMock(), MagicMock())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extend_budget_lifts_the_pause() -> None:
|
||||
hooks = _make_interactive_hooks(10.0)
|
||||
state = _make_report_state(10.5)
|
||||
hooks.extend_budget()
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=state):
|
||||
await hooks.on_llm_end(_make_context(parent_id=None), MagicMock(), MagicMock())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extend_budget_adds_original_amount_each_time() -> None:
|
||||
hooks = _make_interactive_hooks(10.0)
|
||||
hooks.extend_budget()
|
||||
hooks.extend_budget()
|
||||
state = _make_report_state(29.9)
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=state):
|
||||
await hooks.on_llm_end(_make_context(parent_id=None), MagicMock(), MagicMock())
|
||||
state = _make_report_state(30.0)
|
||||
with (
|
||||
patch("strix.core.hooks.get_global_report_state", return_value=state),
|
||||
pytest.raises(BudgetPausedError),
|
||||
):
|
||||
await hooks.on_llm_end(_make_context(parent_id=None), MagicMock(), MagicMock())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_interactive_subagent_uses_root_warning_bands() -> None:
|
||||
hooks = _make_interactive_hooks(10.0)
|
||||
state = _make_report_state(7.4)
|
||||
items: list[Any] = []
|
||||
with patch("strix.core.hooks.get_global_report_state", return_value=state):
|
||||
await hooks.on_llm_start(
|
||||
_make_warn_context(requests=0, parent_id="root-1"), MagicMock(), None, items
|
||||
)
|
||||
assert len(items) == 1
|
||||
content = items[0]["content"]
|
||||
assert "[NOTICE]" in content
|
||||
assert "paused until the user chooses to continue" in content
|
||||
assert "reserve" not in content.lower()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("cost", "max_budget", "interactive", "expected"),
|
||||
[
|
||||
(0.0, None, False, (False, False)),
|
||||
(100.0, None, False, (False, False)),
|
||||
(5.0, 10.0, False, (False, False)),
|
||||
(9.0, 10.0, False, (False, True)),
|
||||
(10.0, 10.0, False, (True, True)),
|
||||
(10.0, 20.0, False, (False, False)),
|
||||
(10.0, 10.0, True, (False, False)),
|
||||
],
|
||||
)
|
||||
def test_recomputed_budget_flags(
|
||||
cost: float,
|
||||
max_budget: float | None,
|
||||
interactive: bool,
|
||||
expected: tuple[bool, bool],
|
||||
) -> None:
|
||||
assert recomputed_budget_flags(cost, max_budget, interactive=interactive) == expected
|
||||
|
||||
Reference in New Issue
Block a user