source/study/budget_proxy.py

"""Local-only Reflection gateway with atomic, restart-safe token reservations.

Pi never receives the provider credential. Every request, including summaries,
passes through this gateway. A failed request with unknown usage keeps its full
reservation as a conservative charge. No prompts or credentials enter logs.
"""
from __future__ import annotations

import argparse
from contextlib import contextmanager
from datetime import datetime, timezone
from email.utils import parsedate_to_datetime
import hashlib
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import json
import math
import os
from pathlib import Path
import re
import secrets
import select
import socket
import sqlite3
import threading
import time
import urllib.error
import urllib.request
import uuid

ENDPOINT = "https://api.reflection.ai/openai/v1/chat/completions"
CONTEXT_LIMIT = 262_144
MAX_OUTPUT = 16_000
MAX_PROVIDER_ATTEMPTS = 6
MAX_RETRY_WAIT_SECONDS = 120
ROUTE = re.compile(r"^/runs/([A-Za-z0-9_-]+)/agents/([A-Za-z0-9_-]+)/v1/chat/completions$")


def remaining_tokens(headers):
    try:
        value = int(headers.get("x-ratelimit-remaining-tokens-day", ""))
        return value if value >= 0 else None
    except (TypeError, ValueError):
        return None


def rate_limit_delay(headers, attempt, elapsed):
    """Bounded backoff; never retry earlier than a valid provider Retry-After."""
    if headers.get("x-should-retry", "").lower() == "false":
        return None
    delay = min(2 ** (attempt + 1), 32) + secrets.randbelow(1000) / 1000
    retry_after = headers.get("Retry-After")
    if retry_after:
        try:
            seconds = float(retry_after)
        except (ValueError, TypeError):
            try:
                seconds = parsedate_to_datetime(retry_after).timestamp() - time.time()
            except (ValueError, TypeError, OverflowError):
                seconds = 0
        if math.isfinite(seconds):
            delay = max(delay, seconds)
    return delay if elapsed + delay <= MAX_RETRY_WAIT_SECONDS else None


def utcnow():
    return datetime.now(timezone.utc).isoformat()


class BudgetExceeded(Exception):
    pass


class Ledger:
    def __init__(self, path, config):
        self.path = str(path)
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        self.config = config
        self.lock = threading.RLock()
        with self.connect() as db:
            db.executescript("""
              CREATE TABLE IF NOT EXISTS requests (
                id TEXT PRIMARY KEY, run_id TEXT NOT NULL, agent_id TEXT NOT NULL,
                category TEXT NOT NULL, reserved INTEGER NOT NULL,
                charged INTEGER NOT NULL DEFAULT 0, status TEXT NOT NULL,
                prompt_tokens INTEGER, completion_tokens INTEGER,
                started_at TEXT NOT NULL, ended_at TEXT, elapsed REAL,
                request_sha256 TEXT NOT NULL, http_status INTEGER,
                remaining_daily_tokens INTEGER,
                context_reserved INTEGER NOT NULL DEFAULT 0,
                context_charged INTEGER NOT NULL DEFAULT 0,
                fresh_input_estimate INTEGER NOT NULL DEFAULT 0,
                new_window INTEGER NOT NULL DEFAULT 0,
                history_hashes TEXT, client_request_id TEXT, purpose TEXT,
                context_window TEXT
              );
              CREATE TABLE IF NOT EXISTS state (key TEXT PRIMARY KEY, value TEXT);
            """)
            encoded = json.dumps(config, sort_keys=True)
            previous = db.execute("SELECT value FROM state WHERE key='config'").fetchone()
            if previous and previous[0] != encoded:
                raise ValueError("Refusing to change budgets or run identities in an existing ledger")
            db.execute("INSERT OR IGNORE INTO state VALUES ('config', ?)", (encoded,))
            # After a crash, dispatched calls may have completed remotely. Keep
            # their entire reservation charged; never silently refund them.
            db.execute("UPDATE requests SET status='unknown_after_restart', charged=reserved, reserved=0, context_charged=context_reserved, context_reserved=0 WHERE status='inflight'")

    @contextmanager
    def connect(self):
        db = sqlite3.connect(self.path, timeout=30)
        db.row_factory = sqlite3.Row
        db.execute("PRAGMA journal_mode=WAL")
        try:
            with db:
                yield db
        finally:
            db.close()

    def run_config(self, run_id):
        for row in self.config["runs"]:
            if row["run_id"] == run_id:
                return row
        raise BudgetExceeded("unregistered_run")

    def reserve(self, run_id, agent_id, request_hash, output_limit, body=None, metadata=None):
        run = self.run_config(run_id)
        # Reflection's model context bound is a conservative prompt bound. It
        # covers serialization/tool-schema overhead without guessing char/4.
        reservation = CONTEXT_LIMIT + output_limit
        with self.lock, self.connect() as db:
            db.execute("BEGIN IMMEDIATE")
            rows = db.execute("SELECT run_id,category,charged+reserved AS used FROM requests").fetchall()
            total = sum(row["used"] for row in rows)
            per_run = sum(row["used"] for row in rows if row["run_id"] == run_id)
            category = run.get("category", "benchmark")
            category_used = sum(row["used"] for row in rows if row["category"] == category)
            caps = [(total, self.config["global_api_token_cap"], "global_api_budget"),
                    (per_run, run["api_token_cap"], "episode_api_budget"),
                    (category_used, self.config["category_caps"][category], "category_api_budget")]
            for used, cap, reason in caps:
                if used + reservation > cap:
                    raise BudgetExceeded(reason)
            quota = db.execute("SELECT value FROM state WHERE key='quota_blocked'").fetchone()
            if quota and quota[0] == "true":
                raise BudgetExceeded("daily_quota_headroom")
            hashes = []
            fresh_input = 0
            new_window = False
            context_reservation = 0
            if body is not None and "context_token_cap" in run:
                incoming = [m for m in body.get("messages", []) if m.get("role") != "assistant"]
                encoded_messages = [json.dumps(m, sort_keys=True, separators=(",", ":")).encode() for m in incoming]
                hashes = [hashlib.sha256(m).hexdigest() for m in encoded_messages]
                key = f"history:{run_id}:{agent_id}"
                previous = db.execute("SELECT value FROM state WHERE key=?", (key,)).fetchone()
                previous = json.loads(previous[0]) if previous else []
                new_window = not previous or hashes[:len(previous)] != previous or (metadata or {}).get("purpose") == "compaction"
                if new_window:
                    # Approximation for admission only; reconciled to the actual
                    # full prompt count on a context-window transition.
                    fresh_input = math.ceil(len(json.dumps(body,ensure_ascii=False).encode()) / 4) + 1024
                else:
                    fresh_input = sum(math.ceil(len(m) / 4) + 8 for m in encoded_messages[len(previous):])
                context_reservation = fresh_input + output_limit
                agent_used, episode_used = 0, 0
                for context_row in db.execute("SELECT agent_id,context_charged+context_reserved AS n FROM requests WHERE run_id=?", (run_id,)):
                    episode_used += context_row["n"]
                    if context_row["agent_id"] == agent_id:
                        agent_used += context_row["n"]
                if episode_used + context_reservation > run["context_token_cap"]:
                    raise BudgetExceeded("episode_context_budget_estimated")
                if agent_used + context_reservation > run["agent_context_token_cap"]:
                    raise BudgetExceeded("agent_context_budget_estimated")
            ident = uuid.uuid4().hex
            db.execute("INSERT INTO requests(id,run_id,agent_id,category,reserved,status,started_at,request_sha256,context_reserved,fresh_input_estimate,new_window,history_hashes) VALUES(?,?,?,?,?,'inflight',?,?,?,?,?,?)",
                       (ident,run_id,agent_id,category,reservation,utcnow(),request_hash,context_reservation,fresh_input,int(new_window),json.dumps(hashes)))
            meta = metadata or {}
            db.execute("UPDATE requests SET client_request_id=?,purpose=?,context_window=? WHERE id=?",(meta.get("request_id"),meta.get("purpose"),meta.get("context_window"),ident))
            return ident, reservation

    def settle(self, ident, usage, elapsed, http_status, remaining_daily=None, rejected=False):
        with self.lock, self.connect() as db:
            db.execute("BEGIN IMMEDIATE")
            row = db.execute("SELECT * FROM requests WHERE id=?", (ident,)).fetchone()
            if row["status"] != "inflight":
                raise ValueError("Request already settled")
            if usage is not None:
                prompt = int(usage["prompt_tokens"])
                output = int(usage["completion_tokens"])
                if min(prompt, output) < 0:
                    raise ValueError("Negative provider usage")
                charge = prompt + output
                status = "complete"
                if charge > row["reserved"]:
                    db.execute("INSERT OR REPLACE INTO state VALUES('quota_blocked','true')")
                    status = "reservation_bound_violated"
            elif rejected:
                prompt = output = None
                charge = 0
                status = "rejected"
            else:
                prompt = output = None
                charge = row["reserved"]
                status = "unknown_usage_conservatively_charged"
            context_charge = 0
            if row["context_reserved"]:
                context_charge = (prompt if row["new_window"] else row["fresh_input_estimate"]) + output if usage else (0 if rejected else row["context_reserved"])
            if usage and row["history_hashes"]:
                db.execute("INSERT OR REPLACE INTO state VALUES(?,?)", (f"history:{row['run_id']}:{row['agent_id']}",row["history_hashes"]))
            db.execute("UPDATE requests SET charged=?,reserved=0,status=?,prompt_tokens=?,completion_tokens=?,ended_at=?,elapsed=?,http_status=?,remaining_daily_tokens=? WHERE id=?",
                       (charge,status,prompt,output,utcnow(),elapsed,http_status,remaining_daily,ident))
            db.execute("UPDATE requests SET context_charged=?,context_reserved=0 WHERE id=?", (context_charge,ident))
            # Headers count estimated input at admission, not finalized usage.
            # Preserve enough space for this completion and all other calls.
            if remaining_daily is not None:
                outstanding = db.execute("SELECT COALESCE(SUM(reserved),0) FROM requests").fetchone()[0]
                if remaining_daily < self.config.get("daily_headroom_tokens", 20_000_000) + outstanding + MAX_OUTPUT:
                    db.execute("INSERT OR REPLACE INTO state VALUES('quota_blocked','true')")
            return dict(db.execute("SELECT * FROM requests WHERE id=?", (ident,)).fetchone())

    def totals(self, run_id=None):
        with self.connect() as db:
            rows = db.execute("SELECT * FROM requests" + (" WHERE run_id=?" if run_id else ""), (run_id,) if run_id else ()).fetchall()
        return {
            "api_tokens_charged": sum(r["charged"] for r in rows),
            "api_tokens_reserved": sum(r["reserved"] for r in rows),
            "api_input_tokens": sum(r["prompt_tokens"] or 0 for r in rows),
            "api_output_tokens": sum(r["completion_tokens"] or 0 for r in rows),
            "api_requests": len(rows),
            "unknown_usage_requests": sum(r["status"].startswith("unknown") for r in rows),
            "context_tokens_estimated": sum(r["context_charged"] for r in rows),
        }


class Gateway(ThreadingHTTPServer):
    daemon_threads = False

    def __init__(self, address, config, ledger_path, event_path, api_key, local_key, endpoint=ENDPOINT):
        super().__init__(address, Handler)
        self.config = config
        self.ledger = Ledger(ledger_path, config)
        self.events = Path(event_path)
        self.events.parent.mkdir(parents=True, exist_ok=True)
        self.event_lock = threading.Lock()
        self.api_key, self.local_key, self.endpoint = api_key, local_key, endpoint
        self.slots = threading.BoundedSemaphore(config.get("max_inflight_requests", 5))

    def emit(self, run_id, kind, metrics=None, message=None):
        run = self.ledger.run_config(run_id)
        event = dict(timestamp=utcnow(),run_id=run_id,task=run.get("task"),condition=run.get("condition"),kind=kind)
        if metrics is not None:
            event["metrics"] = metrics
        if message is not None:
            event["message"] = message
        with self.event_lock:
            with self.events.open("a") as stream:
                stream.write(json.dumps(event) + "\n")


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def log_message(self, *_):
        pass  # The URL and raw request/response body are not application logs.

    def send_json(self, status, payload):
        data = json.dumps(payload).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(data)))
        self.send_header("Connection", "close")
        self.end_headers()
        self.wfile.write(data)
        self.close_connection = True

    def do_GET(self):
        if self.path == "/health":
            self.send_json(200, {"status":"ready", "model":"Beam-501B-A23B"})
        else:
            self.send_json(404, {"error":"not_found"})

    def do_POST(self):
        match = ROUTE.fullmatch(self.path)
        auth = self.headers.get("Authorization", "")
        if not secrets.compare_digest(auth, "Bearer " + self.server.local_key):
            return self.send_json(401, {"error":{"message":"Invalid local gateway credential"}})
        if not match:
            return self.send_json(404, {"error":{"message":"Unregistered gateway route"}})
        run_id, agent_id = match.groups()
        try:
            size = int(self.headers.get("Content-Length", "0"))
        except ValueError:
            return self.send_json(400, {"error":{"message":"Invalid request size"}})
        if not 0 < size <= 8_000_000:
            return self.send_json(400, {"error":{"message":"Invalid request size"}})
        try:
            body = json.loads(self.rfile.read(size))
            if not isinstance(body, dict):
                raise ValueError("Request must be a JSON object")
            if body.get("model") != "Beam-501B-A23B":
                raise ValueError("Only the pinned Beam model is allowed")
            if not body.get("stream"):
                raise ValueError("Streaming with final usage is required")
            body.pop("max_completion_tokens", None)
            body.update(temperature=0.7, top_p=0.9, reasoning_effort="medium")
            body["max_tokens"] = min(int(body.get("max_tokens") or MAX_OUTPUT), MAX_OUTPUT)
            if body["max_tokens"] <= 0:
                raise ValueError("Positive output cap required")
            body["stream_options"] = {"include_usage": True}
            data = json.dumps(body, separators=(",", ":")).encode()
        except (ValueError, TypeError, KeyError) as error:
            return self.send_json(400, {"error":{"message": str(error)}})
        if not self.server.slots.acquire(timeout=660):
            return self.send_json(429, {"error":{"message":"Local concurrency limit; no provider request sent"}})
        try:
            self.proxy_request(run_id, agent_id, body, data)
        finally:
            self.close_connection = True
            self.server.slots.release()

    def client_disconnected(self):
        try:
            return bool(select.select([self.connection], [], [], 0)[0]) and not self.connection.recv(1, socket.MSG_PEEK)
        except (ConnectionResetError, OSError):
            return True

    def wait_retry(self, delay):
        # An aborted solver may have hit its episode deadline during backoff.
        # Check once per second so it never purchases a later retry.
        deadline = time.monotonic() + delay
        while not self.client_disconnected():
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                return True
            time.sleep(min(remaining, 1))
        return False

    def proxy_request(self, run_id, agent_id, body, data):
        metadata = {"request_id": self.headers.get("X-Study-Request-ID"),
                    "purpose": self.headers.get("X-Study-Purpose"),
                    "context_window": self.headers.get("X-Study-Context-Window")}
        metadata = {k: v for k, v in metadata.items() if v and re.fullmatch(r"[A-Za-z0-9_-]{1,100}", v)}
        request_started = time.monotonic()
        for attempt in range(MAX_PROVIDER_ATTEMPTS):
            if self.client_disconnected():
                return
            try:
                # Every dispatch, including a retry, must pass the original
                # global, category, episode, context and daily quota checks.
                ident, _ = self.server.ledger.reserve(run_id, agent_id, hashlib.sha256(data).hexdigest(), body["max_tokens"], body, metadata)
            except BudgetExceeded as error:
                if str(error) != "unregistered_run":
                    self.server.emit(run_id, "budget_stop", message=str(error))
                return self.send_json(402, {"error": {"message": "Study budget stop: " + str(error), "code": "study_budget_exhausted"}})
            except Exception:
                return self.send_json(500, {"error": {"message": "Budget ledger unavailable; no provider request sent"}})
            started = time.monotonic()
            usage = status = remaining = retry_delay = None
            rejected = headers_sent = False
            client_alive = True
            try:
                request = urllib.request.Request(self.server.endpoint, data, headers={"Content-Type": "application/json", "Authorization": "Bearer " + self.server.api_key})
                with urllib.request.urlopen(request, timeout=180) as response:
                    status = response.status
                    remaining = remaining_tokens(response.headers)
                    self.send_response(status)
                    self.send_header("Content-Type", "text/event-stream")
                    self.send_header("Cache-Control", "no-cache")
                    self.send_header("Connection", "close")
                    self.end_headers()
                    headers_sent = True
                    for line in response:
                        if time.monotonic() - started > 600:
                            raise TimeoutError("Provider request exceeded the ten-minute transport ceiling")
                        if line.startswith(b"data:") and line[5:].strip() != b"[DONE]":
                            event = json.loads(line[5:])
                            if event.get("usage"):
                                candidate = event["usage"]
                                if isinstance(candidate, dict) and all(isinstance(candidate.get(k), int) and candidate[k] >= 0 for k in ("prompt_tokens", "completion_tokens")):
                                    usage = candidate
                        if client_alive:
                            try:
                                self.wfile.write(line)
                                self.wfile.flush()
                            except (BrokenPipeError, ConnectionResetError):
                                client_alive = False
                    # Keep reading even if Pi exits: usage must still be settled.
            except urllib.error.HTTPError as error:
                status = error.code
                remaining = remaining_tokens(error.headers)
                # Only explicit pre-stream rejections can refund a reservation.
                # Retry 429 only: 5xx and transport failures may have been billed.
                rejected = not headers_sent and status in (400, 401, 403, 404, 413, 422, 429)
                if status == 429 and not headers_sent and attempt + 1 < MAX_PROVIDER_ATTEMPTS:
                    retry_delay = rate_limit_delay(error.headers, attempt, time.monotonic() - request_started)
                error.close()
                if not headers_sent and retry_delay is None:
                    try:
                        self.send_json(status, {"error": {"message": f"Reflection HTTP {status}; details withheld from solver logs"}})
                    except (BrokenPipeError, ConnectionResetError):
                        pass
            except Exception as error:
                if not headers_sent:
                    try:
                        self.send_json(502, {"error": {"message": "Reflection transport error; reservation retained"}})
                    except (BrokenPipeError, ConnectionResetError):
                        pass
                self.server.emit(run_id, "error", message=f"Provider transport failure: {type(error).__name__}")
            finally:
                record = self.server.ledger.settle(ident, usage, time.monotonic() - started, status, remaining, rejected)
                totals = self.server.ledger.totals(run_id)
                totals["last_request_seconds"] = record["elapsed"]
                if remaining is not None:
                    totals["provider_daily_remaining_at_admission"] = remaining
                self.server.emit(run_id, "progress", metrics=totals)
                if record["status"] != "complete" and retry_delay is None:
                    self.server.emit(run_id, "error", message=record["status"])
            if retry_delay is None:
                return
            self.server.emit(run_id, "progress", metrics={"provider_retry": 1, "http_status": status,
                             "attempt": attempt + 1, "retry_wait_seconds": retry_delay})
            if not self.wait_retry(retry_delay):
                return


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument("--ledger", type=Path, required=True)
    parser.add_argument("--events", type=Path, required=True)
    parser.add_argument("--port", type=int, default=18093)
    args = parser.parse_args()
    config = json.loads(args.config.read_text())
    server = Gateway(("127.0.0.1",args.port),config,args.ledger,args.events,
                     os.environ["REFLECTION_API_KEY"],os.environ["PI_STUDY_API_KEY"])
    print(json.dumps({"status":"ready","port":server.server_port,"global_api_token_cap":config["global_api_token_cap"]}),flush=True)
    try:
        server.serve_forever(poll_interval=0.25)
    except KeyboardInterrupt:
        pass
    finally:
        server.server_close()


if __name__ == "__main__":
    main()