"""Budget: Reservierung vor, Abrechnung nach der Ausführung; persistent; nebenläufig.""" import subprocess import sys import tempfile import threading import unittest from pathlib import Path from aria_router.budget import BudgetLedger, BudgetLimit from aria_router.gateway import Handling from tests.helpers import FakeClock, GatewayTestCase PACKAGE_ROOT = str(Path(__file__).resolve().parents[1]) class BudgetLedgerTests(unittest.TestCase): def setUp(self): self._tmp = tempfile.TemporaryDirectory() self.path = Path(self._tmp.name) / "budget.json" self.clock = FakeClock() def tearDown(self): self._tmp.cleanup() def ledger(self, jobs=2, tokens=None): return BudgetLedger(self.path, {"codex": BudgetLimit(jobs, tokens)}, clock=self.clock) def test_reserve_and_settle_counts(self): ledger = self.ledger(2) for job in ("a", "b"): self.assertTrue(ledger.reserve("codex", job, 1)) ledger.settle(job, consumed=True) self.assertFalse(ledger.reserve("codex", "c", 1)) self.assertEqual(ledger.usage_today("codex")["jobs"], 2) def test_release_when_not_consumed(self): ledger = self.ledger(1) self.assertTrue(ledger.reserve("codex", "a", 1)) ledger.settle("a", consumed=False) self.assertEqual(ledger.usage_today("codex")["jobs"], 0) self.assertTrue(ledger.reserve("codex", "b", 1)) def test_pending_reservation_counts_against_limit(self): ledger = self.ledger(1) self.assertTrue(ledger.reserve("codex", "a", 1)) self.assertFalse(ledger.available("codex")) self.assertFalse(ledger.reserve("codex", "b", 1)) def test_persistence_across_instances(self): first = self.ledger(5) first.reserve("codex", "a", 1) first.settle("a", consumed=True) self.assertEqual(self.ledger(5).usage_today("codex")["jobs"], 1) def test_token_estimate_limit(self): ledger = self.ledger(10, tokens=10) self.assertTrue(ledger.reserve("codex", "a", 8)) self.assertFalse(ledger.reserve("codex", "b", 5)) self.assertTrue(ledger.reserve("codex", "c", 2)) def test_day_rollover_utc(self): ledger = self.ledger(1) ledger.reserve("codex", "a", 1) ledger.settle("a", consumed=True) self.assertFalse(ledger.available("codex")) self.clock.advance(86400) self.assertTrue(ledger.reserve("codex", "b", 1)) def test_unconfigured_backend_denied(self): ledger = self.ledger(5) self.assertFalse(ledger.available("claude")) self.assertFalse(ledger.reserve("claude", "a", 1)) def test_stale_reservation_counted_as_used(self): ledger = self.ledger(1) ledger.reserve("codex", "crashed-job", 1) self.clock.advance(3601) self.assertFalse(ledger.reserve("codex", "b", 1)) self.assertEqual(ledger.usage_today("codex")["jobs"], 1) def test_concurrent_reservations_threads(self): ledger = self.ledger(5) barrier = threading.Barrier(24) results = [] lock = threading.Lock() def worker(index): barrier.wait() granted = BudgetLedger(self.path, {"codex": BudgetLimit(5)}, clock=self.clock).reserve( "codex", f"job-{index}", 1) with lock: results.append(granted) threads = [threading.Thread(target=worker, args=(i,)) for i in range(24)] for thread in threads: thread.start() for thread in threads: thread.join(30) self.assertEqual(sum(results), 5) self.assertEqual(ledger.usage_today("codex")["jobs"], 5) def test_concurrent_reservations_processes(self): code = ( "import sys; sys.path.insert(0, sys.argv[1])\n" "from pathlib import Path\n" "from aria_router.budget import BudgetLedger, BudgetLimit\n" "ledger = BudgetLedger(Path(sys.argv[2]), {'codex': BudgetLimit(5)}, clock=lambda: 1800000000.0)\n" "print(sum(ledger.reserve('codex', f'{sys.argv[3]}-{i}', 1) for i in range(4)))\n" ) procs = [subprocess.Popen([sys.executable, "-I", "-c", code, PACKAGE_ROOT, str(self.path), f"p{i}"], stdout=subprocess.PIPE, text=True) for i in range(4)] granted = sum(int(p.communicate(timeout=60)[0].strip()) for p in procs) self.assertEqual(granted, 5) class GatewayBudgetTests(GatewayTestCase): def test_budget_exhausted_blocks_manual_codex(self): gateway = self.gateway({"budget": {"codex": {"daily_jobs": 1}}}) self.assertEqual(gateway.handle(self.request(1, "/codex #public a")).backend, "codex") result = gateway.handle(self.request(2, "/codex #public b")) self.assertEqual((result.handling, result.code), (Handling.BLOCKED, "MANUAL_CHOICE_NOT_PERMITTED")) self.assertIn("Tagesbudget", result.reply_text) self.assertEqual(len(gateway.adapters["codex"].runner.calls), 1) def test_budget_exhausted_auto_routes_to_eligible_claude(self): gateway = self.gateway({"budget": {"codex": {"daily_jobs": 1}}}) self.assertEqual(gateway.handle(self.request(3, "#public a")).backend, "codex") result = gateway.handle(self.request(4, "#public b")) self.assertEqual(result.backend, "claude") record = self.decision_records(gateway)[-1] codex = next(c for c in record["routing"]["candidates"] if c["backend"] == "codex") self.assertIn("BUDGET_EXHAUSTED", codex["reasons"]) def test_budget_released_when_not_executed(self): gateway = self.gateway({"budget": {"codex": {"daily_jobs": 1}}}) (self.root / "keys" / "clearance.key").unlink() self.assertEqual(gateway.handle(self.request(5, "#public a")).code, "KEY_UNAVAILABLE") self.assertEqual(gateway.status()["backends"]["codex"]["budget"]["jobs"], 0) def test_shadow_does_not_consume_budget(self): gateway = self.gateway({"budget": {"codex": {"daily_jobs": 1}}}) for message_id in (6, 7, 8): self.assertEqual(gateway.shadow(self.request(message_id, "#public a")).backend, "codex") self.assertEqual(gateway.handle(self.request(9, "#public a")).backend, "codex") def test_failed_run_consumes_budget(self): from aria_router.runner import MockRunner, RunResult gateway = self.gateway({"budget": {"codex": {"daily_jobs": 5}}}, codex_runner=MockRunner(lambda spec: RunResult(False, "NONZERO_EXIT", exit_code=1))) self.assertEqual(gateway.handle(self.request(10, "#public a")).handling, Handling.FAILED) self.assertEqual(gateway.status()["backends"]["codex"]["budget"]["jobs"], 1) if __name__ == "__main__": unittest.main()