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>
70 lines
2.2 KiB
Python
70 lines
2.2 KiB
Python
"""SDK run hooks used by Strix orchestration."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from agents.lifecycle import RunHooks
|
|
|
|
from strix.report.state import get_global_report_state
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
from agents import RunContextWrapper
|
|
from agents.agent import Agent
|
|
from agents.items import ModelResponse
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class BudgetExceededError(RuntimeError):
|
|
"""Raised when the accumulated LLM cost reaches the configured budget."""
|
|
|
|
|
|
class ReportUsageHooks(RunHooks[dict[str, Any]]):
|
|
"""Persist SDK-native usage after every model response."""
|
|
|
|
def __init__(self, *, model: str, max_budget_usd: float | None = None) -> None:
|
|
import math
|
|
if max_budget_usd is not None and (not math.isfinite(max_budget_usd) or max_budget_usd <= 0):
|
|
raise ValueError("max_budget_usd must be a finite number greater than 0")
|
|
self._model = model
|
|
self._max_budget_usd = max_budget_usd
|
|
|
|
async def on_llm_end(
|
|
self,
|
|
context: RunContextWrapper[dict[str, Any]],
|
|
agent: Agent[dict[str, Any]],
|
|
response: ModelResponse,
|
|
) -> None:
|
|
report_state = get_global_report_state()
|
|
if report_state is None:
|
|
return
|
|
|
|
ctx = context.context if isinstance(context.context, dict) else {}
|
|
agent_name = getattr(agent, "name", None)
|
|
if not isinstance(agent_name, str):
|
|
agent_name = None
|
|
agent_id = ctx.get("agent_id")
|
|
if not isinstance(agent_id, str) or not agent_id:
|
|
agent_id = agent_name or "unknown"
|
|
|
|
try:
|
|
report_state.record_sdk_usage(
|
|
agent_id=agent_id,
|
|
agent_name=agent_name,
|
|
model=self._model,
|
|
usage=response.usage,
|
|
)
|
|
except Exception:
|
|
logger.exception("failed to record SDK usage for agent %s", agent_id)
|
|
|
|
if self._max_budget_usd is not None:
|
|
cost = report_state.get_total_llm_cost()
|
|
if cost >= self._max_budget_usd:
|
|
raise BudgetExceededError(
|
|
f"Token budget of ${self._max_budget_usd:.2f} exceeded (spent ${cost:.4f})"
|
|
)
|