"""Budget policy and accounting tests.""" import os import sys import json ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) sys.path.insert(0, os.path.join(ROOT, "pilot")) from review.budget import Budget, BudgetState # noqa: E402 from review import opencode_runtime # noqa: E402 from review.opencode import parse_opencode_events # noqa: E402 def test_budget_reads_config_over_environment(monkeypatch): monkeypatch.setenv("PRAGENT_MAX_REVIEW_STEPS", "3") budget = Budget.from_config({"budget": {"max_steps": 7}}) assert budget.max_steps == 7 def test_budget_state_stops_at_token_limit(): state = BudgetState(Budget(max_steps=20, max_total_tokens=100)) assert state.record({"steps": 1, "total": 60, "output": 10}) == "" assert state.record({"steps": 1, "total": 40, "output": 10}) == "max_total_tokens" assert state.snapshot()["cap_hit"] is True def test_budget_state_tracks_cost_cap(): state = BudgetState(Budget(max_equivalent_cost_usd=1.0)) assert state.record({"steps": 1, "total": 1}, 0.75) == "" assert state.record({"steps": 1, "total": 1}, 0.25) == "max_equivalent_cost_usd" def test_process_terminates_after_step_budget(): code = ( "import json,time; " "print(json.dumps({'type':'step_finish','part':{'tokens':{" "'input':1,'output':1,'total':2}}}), flush=True); " "time.sleep(30)" ) budget = Budget(max_steps=1, max_duration_seconds=10) state = BudgetState(budget) proc = opencode_runtime._run_process( [sys.executable, "-u", "-c", code], cwd=".", env=os.environ.copy(), timeout=10, parse_events=parse_opencode_events, budget=budget, budget_state=state, model="glm-5.2:cloud", runner=None, ) assert state.snapshot()["cap_reason"] == "max_steps" assert proc.stdout.count("step_finish") == 1