Mads Hvelplund
962d4459d9
* fix: resolve pre-commit check failures - Change RuntimeError to TypeError for type validation in report/writer.py - Update pyupgrade to v3.21.2 for Python 3.14 compatibility * feat(cli): add --max-budget-usd flag Raises BudgetExceededError in ReportUsageHooks after each LLM call when accumulated cost reaches the limit, with clean "stopped" status and child-agent cancellation in non-interactive mode. * test: add budget enforcement unit tests 7 tests covering no-budget, under-budget, at-limit, over-limit, error message content, None report state, and exception hierarchy. Also adds pytest/pytest-asyncio to dev deps and a mypy override for tests. * fix(budget): validate positive budget and check the live cost ledger Two hardening fixes for --max-budget-usd enforcement: - Reject non-positive budgets. ReportUsageHooks now raises ValueError for max_budget_usd <= 0, and the CLI validates the flag via a custom argparse type so '--max-budget-usd 0' fails fast with a friendly message instead of silently killing the scan on the first model response. - Read the live cost. The budget check now reads ReportState.get_total_llm_cost() (the live ledger) instead of the persisted run-record snapshot, so it stays accurate even when a usage save fails after a model call. * fix(budget): stop the entire scan deterministically when the limit is hit Previously a BudgetExceededError was handled per-agent: it was swallowed in interactive mode (the loop kept waiting), a child's error escaped its detached task as an unretrieved-exception warning, the parent was never released from wait_for_message, and the stop was logged at ERROR with a traceback as if the agent had failed. Replace that with a single scan-wide signal on the coordinator: - AgentCoordinator.trigger_budget_stop() sets a flag and wakes every parked agent; wait_for_message returns as soon as the flag is set. - The run loops check coordinator.budget_stopped and raise to exit cleanly, marking themselves 'stopped'. The root's exception reaches run_strix_scan's handler, which cancels descendants and tears the scan down once; child exceptions are swallowed in their detached task. - The budget stop is logged at INFO, not as a failure. This is deterministic regardless of tree depth or which agent first sees the limit, fixing the interactive/TUI hang where a deep agent's stop never reached a parked root. Also re-raises BudgetExceededError explicitly in the stream handler so it can't be mistaken for the LiteLLM 'after shutdown' race. * fix(budget): treat a budget stop as a clean stop in the TUI Add an explicit BudgetExceededError handler in the TUI scan thread so that, if the error ever reaches it, the budget stop is logged as a graceful stop rather than surfaced as a red scan error by the broad 'except Exception'. The runner normally absorbs the error and returns cleanly, so this is defensive depth for a money-spending feature. * docs(cli): document --max-budget-usd behavior and limitations Clarify that the budget is cumulative across all agents, checked after each model response, that the scan stops cleanly (not as a failure), that the value must be > 0, and that spend can slightly overshoot due to in-flight calls and best-effort cost estimation. * Apply suggestions from code review Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
109 lines
3.7 KiB
Python
109 lines
3.7 KiB
Python
"""Tests for budget enforcement in ReportUsageHooks."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from strix.core.hooks import BudgetExceededError, ReportUsageHooks
|
|
|
|
|
|
def _make_hooks(max_budget: float | None) -> ReportUsageHooks:
|
|
return ReportUsageHooks(model="test-model", max_budget_usd=max_budget)
|
|
|
|
|
|
def _make_report_state(cost: float) -> MagicMock:
|
|
state = MagicMock()
|
|
state.get_total_llm_cost.return_value = cost
|
|
state.record_sdk_usage = MagicMock()
|
|
return state
|
|
|
|
|
|
def _make_context(agent_id: str = "test-agent") -> MagicMock:
|
|
ctx: MagicMock = MagicMock()
|
|
ctx.context = {"agent_id": agent_id}
|
|
return ctx
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_budget_never_raises() -> None:
|
|
hooks = _make_hooks(None)
|
|
state = _make_report_state(9999.0)
|
|
with patch("strix.core.hooks.get_global_report_state", return_value=state):
|
|
await hooks.on_llm_end(_make_context(), MagicMock(), MagicMock())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_under_budget_does_not_raise() -> None:
|
|
hooks = _make_hooks(10.0)
|
|
state = _make_report_state(9.99)
|
|
with patch("strix.core.hooks.get_global_report_state", return_value=state):
|
|
await hooks.on_llm_end(_make_context(), MagicMock(), MagicMock())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_at_budget_raises() -> 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(), MagicMock(), MagicMock())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_over_budget_raises() -> None:
|
|
hooks = _make_hooks(10.0)
|
|
state = _make_report_state(10.01)
|
|
with (
|
|
patch("strix.core.hooks.get_global_report_state", return_value=state),
|
|
pytest.raises(BudgetExceededError),
|
|
):
|
|
await hooks.on_llm_end(_make_context(), MagicMock(), MagicMock())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_budget_check_uses_live_cost_accessor() -> None:
|
|
# The check must read the live ledger, not the persisted run-record snapshot,
|
|
# so it stays accurate even when a save fails after a usage record.
|
|
hooks = _make_hooks(5.0)
|
|
state = _make_report_state(6.0)
|
|
with (
|
|
patch("strix.core.hooks.get_global_report_state", return_value=state),
|
|
pytest.raises(BudgetExceededError),
|
|
):
|
|
await hooks.on_llm_end(_make_context(), MagicMock(), MagicMock())
|
|
state.get_total_llm_cost.assert_called_once()
|
|
state.get_total_llm_usage.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_error_message_includes_amounts() -> None:
|
|
hooks = _make_hooks(5.0)
|
|
state = _make_report_state(7.1234)
|
|
with patch("strix.core.hooks.get_global_report_state", return_value=state):
|
|
with pytest.raises(BudgetExceededError, match=r"\$5\.00") as exc_info:
|
|
await hooks.on_llm_end(_make_context(), MagicMock(), MagicMock())
|
|
assert "7.1234" in str(exc_info.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_raise_when_report_state_none() -> None:
|
|
hooks = _make_hooks(1.0)
|
|
with patch("strix.core.hooks.get_global_report_state", return_value=None):
|
|
# Should return early without raising, even with budget set
|
|
await hooks.on_llm_end(_make_context(), MagicMock(), MagicMock())
|
|
|
|
|
|
@pytest.mark.parametrize("bad_budget", [0.0, -0.01, -5.0])
|
|
def test_non_positive_budget_rejected(bad_budget: float) -> None:
|
|
with pytest.raises(ValueError, match="greater than 0"):
|
|
ReportUsageHooks(model="test-model", max_budget_usd=bad_budget)
|
|
|
|
|
|
def test_budget_exceeded_error_is_runtime_error() -> None:
|
|
err = BudgetExceededError("test")
|
|
assert isinstance(err, RuntimeError)
|