source/study/tests/test_budget_proxy.py

"""Budget gateway checks. All HTTP requests target loopback fake servers only."""

import concurrent.futures
import contextlib
import copy
import http.client
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import importlib.util
import io
import json
from pathlib import Path
import tempfile
import threading
import time
import unittest
from unittest import mock


spec = importlib.util.spec_from_file_location("budget_proxy", Path(__file__).parents[1] / "budget_proxy.py")
proxy = importlib.util.module_from_spec(spec)
spec.loader.exec_module(proxy)


def configuration():
    return {
        "global_api_token_cap": 20_000_000,
        "category_caps": {"benchmark": 18_000_000, "validation": 2_000_000},
        "daily_headroom_tokens": 20_000_000,
        "max_inflight_requests": 2,
        "runs": [
            {"run_id": "run-a", "api_token_cap": 8_000_000, "category": "benchmark", "task": "jq", "condition": "single"},
            {"run_id": "run-b", "api_token_cap": 8_000_000, "category": "benchmark", "task": "xsv", "condition": "peers"},
            {"run_id": "validation-a", "api_token_cap": 2_000_000, "category": "validation"},
        ],
    }


class LedgerTests(unittest.TestCase):
    def setUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.path = Path(self.temp.name) / "ledger.sqlite"
        self.config = configuration()

    def tearDown(self):
        self.temp.cleanup()

    def test_concurrent_reservations_cannot_overspend(self):
        unit = proxy.CONTEXT_LIMIT + 100
        self.config["global_api_token_cap"] = 3 * unit
        ledger = proxy.Ledger(self.path, self.config)
        barrier = threading.Barrier(12)

        def reserve(index):
            barrier.wait()
            try:
                return ledger.reserve("run-a", f"agent-{index}", "hash", 100)
            except proxy.BudgetExceeded:
                return None

        with concurrent.futures.ThreadPoolExecutor(max_workers=12) as pool:
            results = list(pool.map(reserve, range(12)))
        admitted = [item for item in results if item is not None]
        self.assertEqual(len(admitted), 3)
        self.assertEqual(ledger.totals()["api_tokens_reserved"], 3 * unit)

    def test_per_run_category_and_global_caps_are_independent(self):
        unit = proxy.CONTEXT_LIMIT + 100
        for field, expected in [("run", "episode_api_budget"), ("category", "category_api_budget"), ("global", "global_api_budget")]:
            with self.subTest(field=field):
                config = configuration()
                if field == "run":
                    config["runs"][0]["api_token_cap"] = unit
                elif field == "category":
                    config["category_caps"]["benchmark"] = unit
                else:
                    config["global_api_token_cap"] = unit
                ledger = proxy.Ledger(Path(self.temp.name) / f"{field}.sqlite", config)
                ledger.reserve("run-a", "agent-0", "hash", 100)
                next_run = "run-a" if field == "run" else "run-b"
                with self.assertRaisesRegex(proxy.BudgetExceeded, expected):
                    ledger.reserve(next_run, "agent-1", "hash", 100)
                if field != "global":
                    ledger.reserve("validation-a", "agent-0", "hash", 100)

    def test_unknown_failure_keeps_full_charge_and_restart_never_refunds(self):
        ledger = proxy.Ledger(self.path, self.config)
        first, amount = ledger.reserve("run-a", "agent-0", "one", 100)
        ledger.settle(first, None, 1, 502)
        second, second_amount = ledger.reserve("run-b", "agent-1", "two", 100)
        restored = proxy.Ledger(self.path, self.config)
        totals = restored.totals()
        self.assertEqual(totals["api_tokens_charged"], amount + second_amount)
        self.assertEqual(totals["api_tokens_reserved"], 0)
        self.assertEqual(totals["unknown_usage_requests"], 2)
        with restored.connect() as db:
            self.assertEqual(db.execute("SELECT status FROM requests WHERE id=?", (second,)).fetchone()[0], "unknown_after_restart")
        self.assertEqual(proxy.Ledger(self.path, self.config).totals(), totals)

    def test_usage_refunds_unused_reservation_without_double_counting_reasoning(self):
        ledger = proxy.Ledger(self.path, self.config)
        ident, _ = ledger.reserve("run-a", "agent-0", "hash", 100)
        row = ledger.settle(ident, {"prompt_tokens": 1000, "completion_tokens": 90,
                                   "total_tokens": 1090, "completion_tokens_details": {"reasoning_tokens": 70}}, 0.2, 200)
        self.assertEqual(row["charged"], 1090)
        self.assertEqual(row["reserved"], 0)
        self.assertEqual(ledger.totals()["api_output_tokens"], 90)
        with self.assertRaises(ValueError):
            ledger.settle(ident, {"prompt_tokens": 1000, "completion_tokens": 90}, 0.2, 200)

    def test_known_rejection_is_zero_charge_and_daily_headroom_stops_new_calls(self):
        ledger = proxy.Ledger(self.path, self.config)
        ident, _ = ledger.reserve("run-a", "agent-0", "hash", 100)
        row = ledger.settle(ident, None, 0.2, 429, remaining_daily=19_000_000, rejected=True)
        self.assertEqual(row["charged"], 0)
        with self.assertRaisesRegex(proxy.BudgetExceeded, "daily_quota_headroom"):
            ledger.reserve("run-b", "agent-0", "hash", 100)

    def test_mutating_budget_on_restart_is_rejected(self):
        proxy.Ledger(self.path, self.config)
        changed = copy.deepcopy(self.config)
        changed["global_api_token_cap"] += 1
        with self.assertRaises(ValueError):
            proxy.Ledger(self.path, changed)

    def test_repeated_history_is_not_recounted_and_compaction_starts_new_window(self):
        self.config["runs"][0].update(context_token_cap=1_000_000, agent_context_token_cap=1_000_000)
        ledger = proxy.Ledger(self.path, self.config)
        body = {"messages": [{"role": "user", "content": "Implement this task"}]}
        ident, _ = ledger.reserve("run-a", "agent-0", "one", 100, body)
        ledger.settle(ident, {"prompt_tokens": 1000, "completion_tokens": 90}, 0.2, 200)
        first = ledger.totals()["context_tokens_estimated"]
        self.assertEqual(first, 1090)
        body["messages"].extend([{"role": "assistant", "content": "already counted output"},
                                 {"role": "tool", "content": "new tool result"}])
        ident, _ = ledger.reserve("run-a", "agent-0", "two", 100, body)
        ledger.settle(ident, {"prompt_tokens": 10_000, "completion_tokens": 90}, 0.2, 200)
        second = ledger.totals()["context_tokens_estimated"]
        self.assertGreater(second, first + 90)
        self.assertLess(second, first + 300)
        compacted = {"messages": [{"role": "user", "content": "Summary of the preceding conversation"}]}
        ident, _ = ledger.reserve("run-a", "agent-0", "three", 100, compacted)
        ledger.settle(ident, {"prompt_tokens": 500, "completion_tokens": 90}, 0.2, 200)
        self.assertEqual(ledger.totals()["context_tokens_estimated"] - second, 590)

    def test_agent_and_shared_context_limits_block_before_dispatch(self):
        for scope in ("agent", "episode"):
            with self.subTest(scope=scope):
                config = configuration()
                config["runs"][0].update(context_token_cap=2000 if scope == "episode" else 1_000_000,
                                         agent_context_token_cap=1500 if scope == "agent" else 1_000_000)
                ledger = proxy.Ledger(Path(self.temp.name) / f"context-{scope}.sqlite", config)
                body = {"messages": [{"role": "user", "content": "Implement this task"}]}
                ident, _ = ledger.reserve("run-a", "agent-0", "one", 100, body)
                ledger.settle(ident, {"prompt_tokens": 1000, "completion_tokens": 100}, 0.2, 200)
                with self.assertRaisesRegex(proxy.BudgetExceeded, f"{scope}_context_budget_estimated"):
                    ledger.reserve("run-a", "agent-0" if scope == "agent" else "agent-1", "two", 500, body)


class BackoffTests(unittest.TestCase):
    def test_exponential_jitter_and_provider_delay_are_bounded(self):
        with mock.patch.object(proxy.secrets, "randbelow", return_value=500):
            self.assertEqual(proxy.rate_limit_delay({}, 0, 0), 2.5)
            self.assertEqual(proxy.rate_limit_delay({}, 4, 0), 32.5)
            self.assertEqual(proxy.rate_limit_delay({"Retry-After": "60"}, 0, 0), 60)
            self.assertEqual(proxy.rate_limit_delay({"Retry-After": "nonsense"}, 0, 0), 2.5)
            self.assertIsNone(proxy.rate_limit_delay({"Retry-After": "60"}, 0, 61))
            self.assertIsNone(proxy.rate_limit_delay({"x-should-retry": "false"}, 0, 0))
            with mock.patch.object(proxy.time, "time", return_value=0):
                self.assertEqual(proxy.rate_limit_delay({"Retry-After": "Thu, 01 Jan 1970 00:01:00 GMT"}, 0, 0), 60)


class FakeProvider(BaseHTTPRequestHandler):
    def log_message(self, *_):
        pass

    def do_POST(self):
        body = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
        with self.server.lock:
            self.server.calls.append({"body": body, "authorization": self.headers.get("Authorization")})
            mode = self.server.modes.pop(0) if self.server.modes else self.server.mode
            self.server.inflight += 1
            self.server.peak = max(self.server.peak, self.server.inflight)
        try:
            if mode in ("http_500", "http_429"):
                self.send_response(429 if mode == "http_429" else 500)
                if mode == "http_429":
                    self.send_header("Retry-After", self.server.retry_after)
                    self.send_header("x-ratelimit-remaining-tokens-day", self.server.remaining_tokens)
                    self.send_header("x-should-retry", self.server.should_retry)
                details = b"private-provider-error-detail"
                self.send_header("Content-Length", str(len(details)))
                self.end_headers()
                self.wfile.write(details)
                return
            self.send_response(200)
            self.send_header("Content-Type", "text/event-stream")
            self.send_header("x-ratelimit-remaining-tokens-day", "149000000")
            self.end_headers()
            time.sleep(0.02)
            events = [{"choices": [{"delta": {"reasoning_content": "fake reasoning"}}]},
                      {"choices": [{"delta": {"content": "fake answer"}}]}]
            if mode != "missing_usage":
                events.append({"choices": [], "usage": {"prompt_tokens": 1000, "completion_tokens": 90,
                                                       "total_tokens": 1090, "completion_tokens_details": {"reasoning_tokens": 70}}})
            for event in events:
                self.wfile.write(b"data: " + json.dumps(event).encode() + b"\n\n")
            self.wfile.write(b"data: [DONE]\n\n")
        finally:
            with self.server.lock:
                self.server.inflight -= 1


class GatewayTests(unittest.TestCase):
    def setUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.provider = ThreadingHTTPServer(("127.0.0.1", 0), FakeProvider)
        self.provider.mode = "success"
        self.provider.modes = []
        self.provider.retry_after = "0"
        self.provider.remaining_tokens = "149000000"
        self.provider.should_retry = "true"
        self.provider.calls = []
        self.provider.lock = threading.Lock()
        self.provider.inflight = self.provider.peak = 0
        self.provider_thread = threading.Thread(target=self.provider.serve_forever, daemon=True)
        self.provider_thread.start()
        config = configuration()
        config["runs"][0].update(context_token_cap=1_000_000, agent_context_token_cap=1_000_000)
        self.gateway = proxy.Gateway(("127.0.0.1", 0), config,
                                     Path(self.temp.name) / "ledger.sqlite", Path(self.temp.name) / "events.jsonl",
                                     "synthetic-provider-secret", "synthetic-local-secret",
                                     endpoint=f"http://127.0.0.1:{self.provider.server_port}/chat/completions")
        self.gateway_thread = threading.Thread(target=self.gateway.serve_forever, daemon=True)
        self.gateway_thread.start()

    def tearDown(self):
        self.gateway.shutdown()
        self.gateway.server_close()
        self.provider.shutdown()
        self.provider.server_close()
        self.gateway_thread.join(timeout=2)
        self.provider_thread.join(timeout=2)
        self.temp.cleanup()

    def call(self, message="test prompt", authorization="Bearer synthetic-local-secret", path=None, headers=None):
        conn = http.client.HTTPConnection("127.0.0.1", self.gateway.server_port, timeout=5)
        body = {"model": "Beam-501B-A23B", "messages": [{"role": "user", "content": message}],
                "stream": True, "max_tokens": 100, "temperature": 0, "reasoning_effort": "low"}
        conn.request("POST", path or "/runs/run-a/agents/agent-0/v1/chat/completions",
                     json.dumps(body), {"Authorization": authorization, "Content-Type": "application/json", **(headers or {})})
        response = conn.getresponse()
        status, data = response.status, response.read()
        conn.close()
        return status, data

    def wait_settled(self):
        deadline = time.monotonic() + 3
        while self.gateway.ledger.totals()["api_tokens_reserved"] and time.monotonic() < deadline:
            time.sleep(0.01)

    def test_stream_usage_auth_payload_and_redaction(self):
        output = io.StringIO()
        with contextlib.redirect_stdout(output), contextlib.redirect_stderr(output):
            status, data = self.call("private-test-prompt")
            self.wait_settled()
        self.assertEqual(status, 200)
        self.assertIn(b"[DONE]", data)
        self.assertEqual(self.gateway.ledger.totals()["api_tokens_charged"], 1090)
        sent = self.provider.calls[0]
        self.assertEqual(sent["authorization"], "Bearer synthetic-provider-secret")
        self.assertEqual(sent["body"]["stream_options"], {"include_usage": True})
        self.assertEqual(sent["body"]["reasoning_effort"], "medium")
        self.assertEqual(sent["body"]["temperature"], 0.7)
        logs = output.getvalue() + (Path(self.temp.name) / "events.jsonl").read_text()
        with self.gateway.ledger.connect() as db:
            logs += str([dict(row) for row in db.execute("SELECT * FROM requests")])
        for forbidden in ["synthetic-provider-secret", "synthetic-local-secret", "private-test-prompt", "fake reasoning"]:
            self.assertNotIn(forbidden, logs)

    def test_rate_limit_rejection_is_retried_and_every_attempt_is_accounted(self):
        self.provider.modes = ["http_429", "success"]
        with mock.patch.object(proxy, "rate_limit_delay", return_value=0):
            status, data = self.call(headers={"X-Study-Request-ID": "model-call-1", "X-Study-Purpose": "generation"})
        self.wait_settled()
        self.assertEqual(status, 200)
        self.assertIn(b"[DONE]", data)
        self.assertEqual(len(self.provider.calls), 2)
        with self.gateway.ledger.connect() as db:
            rows = list(db.execute("SELECT * FROM requests ORDER BY started_at"))
        self.assertEqual([r["status"] for r in rows], ["rejected", "complete"])
        self.assertEqual([r["charged"] for r in rows], [0, 1090])
        self.assertEqual([r["context_charged"] for r in rows], [0, 1090])
        self.assertEqual({r["run_id"] for r in rows}, {"run-a"})
        self.assertEqual({r["client_request_id"] for r in rows}, {"model-call-1"})
        self.assertEqual({r["purpose"] for r in rows}, {"generation"})
        self.assertEqual({r["request_sha256"] for r in rows}, {rows[0]["request_sha256"]})

    def test_retry_events_pass_real_tracking_sanitizer_and_advance_outbox_cursor(self):
        tracking_spec = importlib.util.spec_from_file_location("gateway_test_tracking", Path(__file__).parents[1] / "tracking.py")
        tracking = importlib.util.module_from_spec(tracking_spec)
        tracking_spec.loader.exec_module(tracking)
        self.provider.modes = ["http_429", "success"]
        with mock.patch.object(proxy, "rate_limit_delay", return_value=0):
            status, _ = self.call()
        self.assertEqual(status, 200)
        events = [json.loads(line) for line in self.gateway.events.read_text().splitlines()]
        sanitized = [tracking.sanitize(event) for event in events]
        self.assertEqual(sum(event["metrics"].get("provider_retry", 0) for event in sanitized), 1)
        tracker = tracking.Tracker(Path(self.temp.name) / "tracking")
        try:
            tracker.ingest(self.gateway.events)
            self.assertEqual(tracker.pending(), len(events))
            self.assertEqual(tracker.db.execute("SELECT offset FROM cursor").fetchone()[0], self.gateway.events.stat().st_size)
            tracker.ingest(self.gateway.events)
            self.assertEqual(tracker.pending(), len(events))
        finally:
            tracker.close()

    def test_concurrent_agents_recover_without_bypassing_concurrency_or_accounting(self):
        self.provider.modes = ["http_429"] * 4
        with mock.patch.object(proxy, "rate_limit_delay", return_value=0):
            with concurrent.futures.ThreadPoolExecutor(max_workers=5) as pool:
                results = list(pool.map(lambda index: self.call(path=f"/runs/run-a/agents/agent-{index}/v1/chat/completions"), range(5)))
        self.wait_settled()
        self.assertTrue(all(status == 200 for status, _ in results))
        self.assertEqual(len(self.provider.calls), 9)
        self.assertLessEqual(self.provider.peak, 2)
        self.assertEqual(self.gateway.ledger.totals()["api_tokens_charged"], 5 * 1090)
        self.assertEqual(self.gateway.ledger.totals()["context_tokens_estimated"], 5 * 1090)

    def test_persistent_rate_limits_stop_after_bounded_attempts_without_charge(self):
        self.provider.mode = "http_429"
        with mock.patch.object(proxy, "rate_limit_delay", return_value=0):
            status, data = self.call()
        self.wait_settled()
        self.assertEqual(status, 429)
        self.assertIn(b"Reflection HTTP 429", data)
        self.assertNotIn(b"private-provider-error-detail", data)
        self.assertEqual(len(self.provider.calls), proxy.MAX_PROVIDER_ATTEMPTS)
        self.assertEqual(self.gateway.ledger.totals()["api_tokens_charged"], 0)
        events = [json.loads(line) for line in self.gateway.events.read_text().splitlines()]
        self.assertEqual(len([e for e in events if e.get("metrics", {}).get("provider_retry") == 1]), proxy.MAX_PROVIDER_ATTEMPTS - 1)
        for secret in ("synthetic-provider-secret", "synthetic-local-secret", "test prompt", "private-provider-error-detail"):
            self.assertNotIn(secret, self.gateway.events.read_text())

    def test_provider_no_retry_header_stops_retries(self):
        self.provider.mode = "http_429"
        self.provider.should_retry = "false"
        status, _ = self.call()
        self.assertEqual(status, 429)
        self.assertEqual(len(self.provider.calls), 1)

    def test_rate_limit_daily_quota_header_blocks_retry(self):
        self.provider.mode = "http_429"
        self.provider.remaining_tokens = "19000000"
        with mock.patch.object(proxy, "rate_limit_delay", return_value=0):
            status, data = self.call()
        self.assertEqual(status, 402)
        self.assertIn(b"daily_quota_headroom", data)
        self.assertEqual(len(self.provider.calls), 1)
        self.assertEqual(self.gateway.ledger.totals()["api_tokens_charged"], 0)

    def test_every_retry_rechecks_original_episode_budget(self):
        self.provider.modes = ["http_429", "success"]
        reserve = self.gateway.ledger.reserve
        with mock.patch.object(self.gateway.ledger, "reserve", side_effect=[reserve("run-a", "agent-0", "hash", 100), proxy.BudgetExceeded("episode_api_budget")]) as calls:
            with mock.patch.object(proxy, "rate_limit_delay", return_value=0):
                status, data = self.call()
        self.assertEqual(status, 402)
        self.assertIn(b"episode_api_budget", data)
        self.assertEqual(calls.call_count, 2)
        self.assertEqual(len(self.provider.calls), 1)
        self.assertEqual(self.gateway.ledger.totals()["api_tokens_charged"], 0)

    def test_provider_retry_after_beyond_deadline_does_not_retry_early(self):
        self.provider.mode = "http_429"
        self.provider.retry_after = "600"
        status, _ = self.call()
        self.assertEqual(status, 429)
        self.assertEqual(len(self.provider.calls), 1)

    def test_disconnected_client_does_not_dispatch_after_backoff(self):
        self.provider.mode = "http_429"
        self.provider.retry_after = "5"
        conn = http.client.HTTPConnection("127.0.0.1", self.gateway.server_port, timeout=5)
        body = {"model": "Beam-501B-A23B", "messages": [], "stream": True, "max_tokens": 100}
        conn.request("POST", "/runs/run-a/agents/agent-0/v1/chat/completions", json.dumps(body),
                     {"Authorization": "Bearer synthetic-local-secret", "Content-Type": "application/json"})
        deadline = time.monotonic() + 3
        while not self.provider.calls and time.monotonic() < deadline:
            time.sleep(0.01)
        self.assertEqual(len(self.provider.calls), 1)
        conn.close()
        self.wait_settled()
        # Acquiring every slot proves the waiting request has released its slot.
        for _ in range(2):
            self.assertTrue(self.gateway.slots.acquire(timeout=2))
        self.gateway.slots.release()
        self.gateway.slots.release()
        self.assertEqual(len(self.provider.calls), 1)

    def test_unauthorized_request_never_reaches_provider_or_reserves_budget(self):
        status, _ = self.call(authorization="Bearer wrong-key")
        self.assertEqual(status, 401)
        self.assertEqual(self.provider.calls, [])
        self.assertEqual(self.gateway.ledger.totals()["api_requests"], 0)

    def test_unregistered_run_is_rejected_without_dispatch(self):
        status, data = self.call(path="/runs/unknown/agents/agent-0/v1/chat/completions")
        self.assertEqual(status, 402)
        self.assertIn(b"unregistered_run", data)
        self.assertEqual(self.provider.calls, [])

    def test_ledger_failure_releases_concurrency_slot_without_dispatch(self):
        with mock.patch.object(self.gateway.ledger, "reserve", side_effect=RuntimeError("sensitive detail")):
            status, data = self.call()
        self.assertEqual(status, 500)
        self.assertNotIn(b"sensitive detail", data)
        self.assertEqual(self.provider.calls, [])
        status, _ = self.call()
        self.assertEqual(status, 200)
        self.wait_settled()

    def test_missing_usage_and_transport_failure_retain_full_reservation(self):
        for mode, expected_status in [("missing_usage", 200), ("http_500", 500)]:
            with self.subTest(mode=mode):
                self.provider.mode = mode
                before = self.gateway.ledger.totals()["api_tokens_charged"]
                status, _ = self.call()
                self.wait_settled()
                self.assertEqual(status, expected_status)
                self.assertEqual(self.gateway.ledger.totals()["api_tokens_charged"] - before, proxy.CONTEXT_LIMIT + 100)

    def test_global_concurrency_limit_covers_all_agents_and_summary_requests(self):
        def request(index):
            # Compaction is another completion on this same registered provider route.
            prompt = "Summarize the prior conversation" if index % 2 else "Continue implementation"
            return self.call(prompt, path=f"/runs/run-a/agents/agent-{index}/v1/chat/completions")
        with concurrent.futures.ThreadPoolExecutor(max_workers=6) as pool:
            results = list(pool.map(request, range(6)))
        self.wait_settled()
        self.assertTrue(all(status == 200 for status, _ in results))
        self.assertEqual(len(self.provider.calls), 6)
        self.assertLessEqual(self.provider.peak, 2)
        self.assertEqual(self.gateway.ledger.totals()["api_tokens_charged"], 6 * 1090)


if __name__ == "__main__":
    unittest.main()