Add configurable token / cost usage limits (#576)

* 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>
This commit is contained in:
Mads Hvelplund
2026-06-22 17:17:08 +02:00
committed by GitHub
parent 11e5d1c2b3
commit 962d4459d9
17 changed files with 351 additions and 23 deletions
+18
View File
@@ -43,6 +43,24 @@ strix --target <target> [options]
Path to a custom config file (JSON) to use instead of `~/.strix/cli-config.json`.
</ParamField>
<ParamField path="--max-budget-usd" type="number">
Maximum LLM spend in USD for the whole scan, counted cumulatively across the
root agent and every child agent. The budget is checked after each model
response; once the running cost reaches the threshold, the scan stops cleanly
with a `stopped` status (not a failure) and the sandbox is torn down.
Must be greater than `0`. Omit the flag for no limit.
**Limitations**
- The check fires *after* a response is returned, so the final spend can
slightly overshoot the limit by any calls already in flight when the
threshold is crossed (most relevant with several child agents running
concurrently).
- Cost is a best-effort estimate derived from token usage and model pricing;
providers that do not expose priced usage may under-count.
</ParamField>
## Examples
```bash