obey-cli 0.1.0__py3-none-any.whl

obey_cli/runner/gateway.py

@@ -0,0 +1,529 @@
+"""Model gateway: holds the API key outside the agent's container (standard library only).
+
+Agents get a dummy token and a base URL of the form http://<gateway>/r/<run_id>/<api>, where <api> is
+`anthropic` (Messages API), `openai` (Chat Completions / Responses) or `gemini`. The gateway checks the dummy token,
+swaps in the real key, pins the model, streams the response back, records token usage and cost per run in the
+ledger file, and refuses calls once a run or the whole batch reaches its cap. Only the model endpoints of each
+API are reachable, and only for run ids the host registered (when it registered any).
+
+    OBEY_UPSTREAM_KEY=... python3 gateway.py --listen 0.0.0.0:8787 --model vendor/model --ledger ledger.jsonl \
+        --price-prompt 3e-7 --price-completion 1.2e-6 --run-cap 0.5 --budget 20 --run-id r1 --run-id r2
+"""
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import hmac
+import http.client
+import json
+import os
+import re
+import sys
+import threading
+import time
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from typing import Any, Dict, Iterable, List, Optional, Set, Tuple
+from urllib.parse import parse_qs, quote, unquote, urlsplit
+
+UPSTREAMS = {
+    "anthropic": "https://openrouter.ai/api",
+    "openai": "https://openrouter.ai/api/v1",
+    "gemini": "https://generativelanguage.googleapis.com",
+}
+DUMMY_TOKEN = "obey-dummy-token"  # noqa: S105 - what agents send; same value as runner/agents/base.py (stand-alone file)
+MAX_BODY = 64 * 1024 * 1024
+_PATH_RE = re.compile(r"/r/([A-Za-z0-9._-]{1,80})/(anthropic|openai|gemini)(/.*)?")
+_CHUNK_SIZE_RE = re.compile(rb"[0-9A-Fa-f]{1,8}")
+# What an agent may call, per API, relative to the upstream base URL (the openai base already ends in /v1).
+_ROUTES: Dict[str, List[Tuple[str, "re.Pattern[str]"]]] = {
+    "anthropic": [("POST", re.compile(r"/v1/messages(?:/count_tokens)?")), ("GET", re.compile(r"/v1/models"))],
+    "openai": [("POST", re.compile(r"/(?:chat/completions|responses)")), ("GET", re.compile(r"/models"))],
+    "gemini": [
+        (
+            "POST",
+            re.compile(
+                r"/v1beta/models/(?P<model>[A-Za-z0-9._-]+):(?P<verb>generateContent|streamGenerateContent|countTokens)"
+            ),
+        )
+    ],
+}
+_HOP_HEADERS = {
+    "connection",
+    "keep-alive",
+    "proxy-authenticate",
+    "proxy-authorization",
+    "te",
+    "trailers",
+    "transfer-encoding",
+    "upgrade",
+    "host",
+    "content-length",
+    "authorization",
+    "x-api-key",
+    "x-goog-api-key",
+    "accept-encoding",
+}
+_RESPONSE_SKIP = _HOP_HEADERS | {"server", "date"}  # send_response adds our own
+
+
+class Config:
+    def __init__(
+        self,
+        key: str,
+        model: Optional[str],
+        ledger: Optional[str],
+        price: Dict[str, float],
+        run_cap: Optional[float],
+        budget: Optional[float],
+        upstreams: Dict[str, str],
+        allowed_hosts: Optional[List[str]] = None,
+        token: str = DUMMY_TOKEN,
+        run_ids: Optional[Iterable[str]] = None,
+    ):
+        self.key = key
+        self.model = model
+        self.ledger = ledger
+        self.price = price
+        self.run_cap = run_cap
+        self.budget = budget
+        self.upstreams = upstreams
+        self.allowed_hosts = allowed_hosts
+        self.token = token
+        # Run ids the host registered for this batch; None accepts any id (and register_run() switches the check on).
+        self.run_ids: Optional[Set[str]] = set(run_ids) if run_ids is not None else None
+        self.lock = threading.Lock()
+        self.spent: Dict[str, float] = {}
+        self.total = 0.0
+        if ledger and os.path.exists(ledger):
+            with open(ledger, encoding="utf-8") as fh:
+                for line in fh:
+                    if line.strip():
+                        row = json.loads(line)
+                        self._add(str(row.get("run_id")), float(row.get("cost_usd", 0.0)))
+
+    def _add(self, run_id: str, cost: float) -> None:
+        self.spent[run_id] = self.spent.get(run_id, 0.0) + cost
+        self.total += cost
+
+    def register_run(self, run_id: str) -> None:
+        with self.lock:
+            if self.run_ids is None:
+                self.run_ids = set()
+            self.run_ids.add(run_id)
+
+    def known_run(self, run_id: str) -> bool:
+        with self.lock:
+            return self.run_ids is None or run_id in self.run_ids
+
+    def over_cap(self, run_id: str) -> Optional[str]:
+        with self.lock:
+            if self.budget is not None and self.total >= self.budget:
+                return f"batch budget ${self.budget:.2f} reached"
+            if self.run_cap is not None and self.spent.get(run_id, 0.0) >= self.run_cap:
+                return f"per-run cap ${self.run_cap:.2f} reached"
+        return None
+
+    def record(self, run_id: str, requested: str, usage: Dict[str, int], provider_cost: Optional[float]) -> None:
+        if provider_cost is not None:
+            cost, source = provider_cost, "provider"
+        else:
+            cost, source = price_cost(self.price, usage), "price_table"
+        entry = {
+            "ts": round(time.time(), 3),
+            "run_id": run_id,
+            "kind": "agent",
+            "model": self.model or requested,
+            "input_tokens": usage.get("input_tokens", 0),
+            "output_tokens": usage.get("output_tokens", 0),
+            "cache_read_tokens": usage.get("cache_read_tokens", 0),
+            "cache_write_tokens": usage.get("cache_write_tokens", 0),
+            "cost_usd": cost,
+            "cost_source": source,
+        }
+        with self.lock:
+            self._add(run_id, cost)
+            if self.ledger:
+                with open(self.ledger, "a", encoding="utf-8") as fh:
+                    fh.write(json.dumps(entry, sort_keys=True) + "\n")
+
+
+def price_cost(price: Dict[str, float], usage: Dict[str, int]) -> float:
+    cached = usage.get("cache_read_tokens", 0)
+    written = usage.get("cache_write_tokens", 0)
+    fresh = max(0, usage.get("input_tokens", 0) - cached - written)
+    prompt = price.get("prompt", 0.0)
+    return (
+        fresh * prompt
+        + usage.get("output_tokens", 0) * price.get("completion", 0.0)
+        + cached * price.get("cache_read", prompt)
+        + written * price.get("cache_write", prompt)
+    )
+
+
+class UsageMeter:
+    """Accumulates usage from a JSON body or from server-sent events, for both API styles."""
+
+    def __init__(self) -> None:
+        self.usage = {"input_tokens": 0, "output_tokens": 0, "cache_read_tokens": 0, "cache_write_tokens": 0}
+        self.cost: Optional[float] = None
+        self._buffer = b""
+
+    def feed_json(self, body: Any) -> None:
+        if not isinstance(body, dict):
+            return
+        if body.get("type") == "message_start" and isinstance(body.get("message"), dict):
+            body = body["message"]
+        if body.get("type") == "response.completed" and isinstance(body.get("response"), dict):
+            body = body["response"]
+        meta = body.get("usageMetadata")
+        if isinstance(meta, dict):  # Gemini API
+            self.usage["input_tokens"] = max(self.usage["input_tokens"], int(meta.get("promptTokenCount") or 0))
+            self.usage["output_tokens"] = max(
+                self.usage["output_tokens"],
+                int(meta.get("candidatesTokenCount") or 0) + int(meta.get("thoughtsTokenCount") or 0),
+            )
+            self.usage["cache_read_tokens"] = max(
+                self.usage["cache_read_tokens"], int(meta.get("cachedContentTokenCount") or 0)
+            )
+            return
+        usage = body.get("usage")
+        if not isinstance(usage, dict):
+            return
+        if isinstance(usage.get("cost"), (int, float)):
+            self.cost = (self.cost or 0.0) + float(usage["cost"])
+        if "prompt_tokens" in usage or "completion_tokens" in usage:
+            details = usage.get("prompt_tokens_details") or {}
+            self.usage["input_tokens"] = max(self.usage["input_tokens"], int(usage.get("prompt_tokens") or 0))
+            self.usage["output_tokens"] = max(self.usage["output_tokens"], int(usage.get("completion_tokens") or 0))
+            self.usage["cache_read_tokens"] = max(
+                self.usage["cache_read_tokens"], int(details.get("cached_tokens") or 0)
+            )
+            return
+        if "input_tokens" in usage and "output_tokens" in usage and "input_tokens_details" in usage:
+            details = usage.get("input_tokens_details") or {}  # Responses API
+            self.usage["input_tokens"] = max(self.usage["input_tokens"], int(usage.get("input_tokens") or 0))
+            self.usage["output_tokens"] = max(self.usage["output_tokens"], int(usage.get("output_tokens") or 0))
+            self.usage["cache_read_tokens"] = max(
+                self.usage["cache_read_tokens"], int(details.get("cached_tokens") or 0)
+            )
+            return
+        read = int(usage.get("cache_read_input_tokens") or 0)
+        write = int(usage.get("cache_creation_input_tokens") or 0)
+        if "input_tokens" in usage:
+            self.usage["input_tokens"] = max(
+                self.usage["input_tokens"], int(usage.get("input_tokens") or 0) + read + write
+            )
+            self.usage["cache_read_tokens"] = max(self.usage["cache_read_tokens"], read)
+            self.usage["cache_write_tokens"] = max(self.usage["cache_write_tokens"], write)
+        if "output_tokens" in usage:
+            self.usage["output_tokens"] = max(self.usage["output_tokens"], int(usage.get("output_tokens") or 0))
+
+    def feed_sse(self, chunk: bytes) -> None:
+        self._buffer += chunk
+        while b"\n" in self._buffer:
+            line, self._buffer = self._buffer.split(b"\n", 1)
+            line = line.strip()
+            if line.startswith(b"data:"):
+                data = line[5:].strip()
+                if data and data != b"[DONE]":
+                    with contextlib.suppress(ValueError):
+                        self.feed_json(json.loads(data))
+
+
+def rewrite_body(api: str, raw: bytes, model: Optional[str]) -> Tuple[bytes, str]:
+    """Pin the model and ask the provider to report usage. Returns (body, model the agent asked for)."""
+    try:
+        body = json.loads(raw) if raw else None
+    except ValueError:
+        return raw, ""
+    if not isinstance(body, dict):
+        return raw, ""
+    requested = str(body.get("model", ""))
+    if model and "model" in body:
+        body["model"] = model
+    if api == "openai":
+        body["usage"] = {"include": True}
+        if body.get("stream") and "messages" in body:
+            options = body.get("stream_options")
+            merged: Dict[str, Any] = dict(options) if isinstance(options, dict) else {}
+            merged["include_usage"] = True
+            body["stream_options"] = merged
+    return json.dumps(body).encode(), requested
+
+
+def clean_path(rest: str) -> Optional[str]:
+    """Percent-decode a request path once; None if it could climb out of its route (dot segments, encoded slashes)."""
+    low = rest.lower()
+    if "%2f" in low or "%5c" in low or "%25" in low or "\\" in rest:
+        return None
+    path = unquote(rest)
+    if any(ord(ch) < 32 or ord(ch) == 127 for ch in path) or "//" in path:
+        return None
+    if any(segment in (".", "..") for segment in path.split("/")):
+        return None
+    return path
+
+
+def route(api: str, method: str, rest: str, pinned: Optional[str]) -> Optional[Tuple[str, str]]:
+    """Check a path against the API's allow-list. Returns (upstream path, model named in the path), or None to refuse.
+
+    Gemini names its model in the path, so the pin is applied here rather than in the body."""
+    path = clean_path(rest)
+    if path is None:
+        return None
+    for verb, pattern in _ROUTES[api]:
+        found = pattern.fullmatch(path)
+        if verb == method and found:
+            requested = found.groupdict().get("model") or ""
+            if requested and pinned:
+                path = f"/v1beta/models/{quote(pinned, safe='')}:{found.group('verb')}"
+            return path, requested
+    return None
+
+
+class BadBody(Exception):
+    def __init__(self, status: int, message: str):
+        super().__init__(message)
+        self.status = status
+
+
+class Handler(BaseHTTPRequestHandler):
+    server_version = "obey-gateway"
+    protocol_version = "HTTP/1.1"
+    config: Config
+
+    def log_message(self, format: str, *args: Any) -> None:
+        pass
+
+    def _reply(self, status: int, payload: Dict[str, Any], close: bool = False) -> None:
+        data = json.dumps(payload).encode()
+        self.send_response(status)
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Length", str(len(data)))
+        if close:  # the request body was not read, so the connection cannot be reused
+            self.send_header("Connection", "close")
+        self.end_headers()
+        self.wfile.write(data)
+
+    def _guard(self) -> bool:
+        cfg = self.config
+        if self.headers.get("Origin"):
+            self._reply(403, {"error": {"message": "browser requests are not accepted"}}, close=True)
+            return False
+        host = (self.headers.get("Host") or "").lower()
+        if cfg.allowed_hosts is not None and host not in cfg.allowed_hosts:
+            self._reply(403, {"error": {"message": "unexpected Host header"}}, close=True)
+            return False
+        return True
+
+    def _authorised(self, api: str, query: str) -> bool:
+        """The agent must present the batch's dummy token (bearer, x-api-key, x-goog-api-key or Gemini's ?key=)."""
+        supplied = [self.headers.get("x-api-key") or "", self.headers.get("x-goog-api-key") or ""]
+        auth = self.headers.get("Authorization") or ""
+        if auth[:7].lower() == "bearer ":
+            supplied.append(auth[7:].strip())
+        if api == "gemini":
+            supplied += parse_qs(query).get("key", [])
+        token = self.config.token.encode()
+        return any(value and hmac.compare_digest(value.encode("utf-8", "replace"), token) for value in supplied)
+
+    def _read_body(self) -> bytes:
+        if "chunked" in (self.headers.get("Transfer-Encoding") or "").lower():
+            return self._read_chunked()
+        text = (self.headers.get("Content-Length") or "0").strip()
+        if not (text.isascii() and text.isdigit()):
+            raise BadBody(400, "bad Content-Length")
+        length = int(text)
+        if length > MAX_BODY:
+            raise BadBody(413, "request body too large")
+        data = self.rfile.read(length) if length else b""
+        if len(data) < length:
+            raise BadBody(400, "request body ended early")
+        return data
+
+    def _read_chunked(self) -> bytes:
+        body = bytearray()
+        while True:
+            size_text = self.rfile.readline(1024).split(b";", 1)[0].strip()
+            if not _CHUNK_SIZE_RE.fullmatch(size_text):
+                raise BadBody(400, "bad chunk size")
+            size = int(size_text, 16)
+            if size == 0:
+                break
+            if len(body) + size > MAX_BODY:
+                raise BadBody(413, "request body too large")
+            chunk = self.rfile.read(size)
+            if len(chunk) < size or self.rfile.read(2) != b"\r\n":
+                raise BadBody(400, "bad chunked body")
+            body += chunk
+        while self.rfile.readline(1024).strip():  # trailers
+            pass
+        return bytes(body)
+
+    def do_GET(self) -> None:
+        self._forward("GET")
+
+    def do_POST(self) -> None:
+        self._forward("POST")
+
+    def _forward(self, method: str) -> None:
+        if not self._guard():
+            return
+        raw_path, _, query = self.path.partition("?")
+        match = _PATH_RE.fullmatch(raw_path)
+        if not match:
+            self._reply(404, {"error": {"message": "use /r/<run_id>/<anthropic|openai|gemini>/..."}}, close=True)
+            return
+        run_id, api, rest = match.group(1), match.group(2), match.group(3) or "/"
+        if not self._authorised(api, query):
+            self._reply(401, {"error": {"message": "missing or wrong gateway token"}}, close=True)
+            return
+        try:
+            raw = self._read_body()
+        except BadBody as exc:
+            self._reply(exc.status, {"error": {"message": str(exc)}}, close=True)
+            return
+        if not self.config.known_run(run_id):
+            self._reply(403, {"error": {"message": "unknown run id"}})
+            return
+        routed = route(api, method, rest, self.config.model)
+        if routed is None:
+            self._reply(403, {"error": {"message": "path not allowed"}})
+            return
+        upath, path_model = routed
+        reason = self.config.over_cap(run_id)
+        if reason:
+            self._reply(402, {"type": "error", "error": {"type": "budget_exceeded", "message": reason}})
+            return
+        body, requested = rewrite_body(api, raw, self.config.model) if method == "POST" else (raw, "")
+        requested = requested or path_model
+        base = urlsplit(self.config.upstreams[api])
+        headers = {k: v for k, v in self.headers.items() if k.lower() not in _HOP_HEADERS}
+        if api == "gemini":
+            headers["x-goog-api-key"] = self.config.key
+            query = "&".join(q for q in query.split("&") if not q.startswith("key="))
+        else:
+            headers["Authorization"] = f"Bearer {self.config.key}"
+        path = base.path.rstrip("/") + upath + (f"?{query}" if query else "")
+        if method == "POST":
+            headers["Content-Length"] = str(len(body))
+        headers["Accept-Encoding"] = "identity"
+        conn_cls = http.client.HTTPSConnection if base.scheme == "https" else http.client.HTTPConnection
+        conn = conn_cls(base.netloc, timeout=600)
+        meter = UsageMeter()
+        started = False  # response headers are on the wire, so a failure can no longer become a 502
+        streaming = False
+        collected = b""
+        settled = False
+
+        def settle() -> None:
+            # Record the call before the response ends, so the client's next request sees this cost in the caps.
+            nonlocal settled
+            if settled:
+                return
+            settled = True
+            if not streaming and collected:
+                with contextlib.suppress(ValueError):
+                    meter.feed_json(json.loads(collected))
+            if method == "POST" and (any(meter.usage.values()) or meter.cost):
+                self.config.record(run_id, requested, meter.usage, meter.cost)
+
+        try:
+            conn.request(method, path, body=body if method == "POST" else None, headers=headers)
+            resp = conn.getresponse()
+            streaming = "text/event-stream" in (resp.getheader("Content-Type") or "")
+            self.send_response(resp.status)
+            for key, value in resp.getheaders():
+                if key.lower() not in _RESPONSE_SKIP:
+                    self.send_header(key, value)
+            if resp.status in (204, 304):  # no body, so no chunked framing either
+                self.end_headers()
+                return
+            self.send_header("Transfer-Encoding", "chunked")
+            self.end_headers()
+            started = True
+            while True:
+                chunk = resp.read1(65536) if hasattr(resp, "read1") else resp.read(65536)
+                if not chunk:
+                    break
+                if streaming:
+                    meter.feed_sse(chunk)
+                elif len(collected) < 8_000_000:
+                    collected += chunk
+                self.wfile.write(f"{len(chunk):x}\r\n".encode() + chunk + b"\r\n")
+                self.wfile.flush()
+            if resp.length:  # read1 does not notice an upstream that closed before its Content-Length
+                raise http.client.IncompleteRead(b"", resp.length)
+            settle()
+            self.wfile.write(b"0\r\n\r\n")
+        except (OSError, http.client.HTTPException) as exc:
+            self.close_connection = True
+            settle()
+            with contextlib.suppress(OSError):
+                if started:  # end the chunked body here instead of writing a second response into it
+                    self.wfile.write(b"0\r\n\r\n")
+                else:
+                    self._reply(502, {"error": {"message": f"upstream error: {type(exc).__name__}"}})
+        finally:
+            conn.close()
+            settle()
+
+
+def make_server(listen: str, config: Config) -> ThreadingHTTPServer:
+    host, _, port = listen.rpartition(":")
+    handler = type("BoundHandler", (Handler,), {"config": config})
+    server = ThreadingHTTPServer((host or "127.0.0.1", int(port)), handler)
+    server.daemon_threads = True
+    return server
+
+
+def main(argv: Optional[List[str]] = None) -> int:
+    parser = argparse.ArgumentParser(description="obey model gateway")
+    parser.add_argument("--listen", default="127.0.0.1:8787")
+    parser.add_argument("--model", default=None, help="pin every request to this model")
+    parser.add_argument("--ledger", default=None)
+    parser.add_argument("--price-prompt", type=float, default=0.0)
+    parser.add_argument("--price-completion", type=float, default=0.0)
+    parser.add_argument("--price-cache-read", type=float, default=None)
+    parser.add_argument("--price-cache-write", type=float, default=None)
+    parser.add_argument("--run-cap", type=float, default=None, help="USD per run")
+    parser.add_argument("--budget", type=float, default=None, help="USD for everything in the ledger")
+    parser.add_argument("--upstream", action="append", default=[], help="api=https://base (repeatable)")
+    parser.add_argument("--allowed-host", action="append", default=None, help="accepted Host header (repeatable)")
+    parser.add_argument("--token", default=DUMMY_TOKEN, help="token agents must present instead of a real key")
+    parser.add_argument("--run-id", action="append", default=None, help="accepted run id (repeatable; default: any)")
+    args = parser.parse_args(argv)
+    key = os.environ.get("OBEY_UPSTREAM_KEY", "")
+    if not key:
+        print("OBEY_UPSTREAM_KEY is not set", file=sys.stderr)
+        return 2
+    upstreams = dict(UPSTREAMS)
+    for item in args.upstream:
+        api, _, url = item.partition("=")
+        upstreams[api] = url
+    price = {"prompt": args.price_prompt, "completion": args.price_completion}
+    if args.price_cache_read is not None:
+        price["cache_read"] = args.price_cache_read
+    if args.price_cache_write is not None:
+        price["cache_write"] = args.price_cache_write
+    config = Config(
+        key,
+        args.model,
+        args.ledger,
+        price,
+        args.run_cap,
+        args.budget,
+        upstreams,
+        [h.lower() for h in args.allowed_host] if args.allowed_host else None,
+        args.token,
+        args.run_id,
+    )
+    server = make_server(args.listen, config)
+    print(f"gateway listening on {args.listen}", flush=True)
+    server.serve_forever()
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

obey_cli/runner/ledger.py

@@ -0,0 +1,132 @@
+"""Spend ledger: every model call is recorded, and runs stop before the budget is exhausted.
+
+The ledger is an append-only JSON-lines file guarded by an advisory lock, so the model gateway (which records
+usage as responses stream through) and the runner (which checks the budget before starting a run) can share it.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import fcntl
+import json
+import time
+from collections.abc import Iterator, Mapping
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import IO, Any
+
+
+class BudgetExceeded(RuntimeError):
+    pass
+
+
+@dataclass
+class Price:
+    """USD per token."""
+
+    prompt: float
+    completion: float
+    cache_read: float | None = None
+    cache_write: float | None = None
+
+    def cost(self, usage: Mapping[str, int]) -> float:
+        cached = int(usage.get("cache_read_tokens", 0))
+        written = int(usage.get("cache_write_tokens", 0))
+        fresh = max(0, int(usage.get("input_tokens", 0)) - cached - written)
+        total = fresh * self.prompt + int(usage.get("output_tokens", 0)) * self.completion
+        total += cached * (self.cache_read if self.cache_read is not None else self.prompt)
+        total += written * (self.cache_write if self.cache_write is not None else self.prompt)
+        return total
+
+
+@dataclass
+class Entry:
+    ts: float
+    run_id: str
+    kind: str  # agent | llm | estimate
+    model: str
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cache_read_tokens: int = 0
+    cache_write_tokens: int = 0
+    cost_usd: float = 0.0
+    cost_source: str = "price_table"  # provider | price_table
+
+
+@contextlib.contextmanager
+def _locked(path: Path, mode: str) -> Iterator[IO[str]]:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open(mode, encoding="utf-8") as fh:
+        fcntl.flock(fh.fileno(), fcntl.LOCK_EX if "a" in mode or "w" in mode else fcntl.LOCK_SH)
+        try:
+            yield fh
+        finally:
+            fcntl.flock(fh.fileno(), fcntl.LOCK_UN)
+
+
+class Ledger:
+    def __init__(self, path: Path, budget_usd: float | None = None):
+        self.path = path
+        self.budget_usd = budget_usd
+
+    def record(self, entry: Entry) -> None:
+        with _locked(self.path, "a") as fh:
+            fh.write(json.dumps(asdict(entry), sort_keys=True) + "\n")
+
+    def entries(self) -> list[Entry]:
+        if not self.path.exists():
+            return []
+        with _locked(self.path, "r") as fh:
+            rows = [json.loads(line) for line in fh if line.strip()]
+        return [Entry(**row) for row in rows]
+
+    def spent(self, run_id: str | None = None) -> float:
+        return sum(e.cost_usd for e in self.entries() if run_id is None or e.run_id == run_id)
+
+    def remaining(self) -> float | None:
+        return None if self.budget_usd is None else self.budget_usd - self.spent()
+
+    def check(self, next_cost_usd: float = 0.0) -> None:
+        """Raise BudgetExceeded when spending `next_cost_usd` more would cross the budget."""
+        if self.budget_usd is None:
+            return
+        spent = self.spent()
+        if spent + next_cost_usd > self.budget_usd:
+            raise BudgetExceeded(
+                f"budget ${self.budget_usd:.2f}: spent ${spent:.4f}, next run estimated at ${next_cost_usd:.4f}"
+            )
+
+
+def usage_from_response(body: Mapping[str, Any]) -> dict[str, int]:
+    """Normalise OpenAI- and Anthropic-style usage blocks to input/output/cache token counts."""
+    usage = body.get("usage") or {}
+    if not isinstance(usage, Mapping):
+        return {}
+    if "prompt_tokens" in usage or "completion_tokens" in usage:
+        details = usage.get("prompt_tokens_details") or {}
+        return {
+            "input_tokens": int(usage.get("prompt_tokens") or 0),
+            "output_tokens": int(usage.get("completion_tokens") or 0),
+            "cache_read_tokens": int(details.get("cached_tokens") or 0),
+            "cache_write_tokens": int(details.get("cache_write_tokens") or 0),
+        }
+    read = int(usage.get("cache_read_input_tokens") or 0)
+    write = int(usage.get("cache_creation_input_tokens") or 0)
+    return {
+        # Anthropic counts cached tokens separately from input_tokens.
+        "input_tokens": int(usage.get("input_tokens") or 0) + read + write,
+        "output_tokens": int(usage.get("output_tokens") or 0),
+        "cache_read_tokens": read,
+        "cache_write_tokens": write,
+    }
+
+
+def provider_cost(body: Mapping[str, Any]) -> float | None:
+    usage = body.get("usage") or {}
+    if isinstance(usage, Mapping) and isinstance(usage.get("cost"), int | float):
+        return float(usage["cost"])
+    return None
+
+
+def now() -> float:
+    return time.time()