"""The request gate — authenticate, balance-check, and burn. Every /v1/* request flows through here: 1. resolve the Bearer token → (user_id | admin | None) 2. if a user: check balance > 0, else 402 3. after the upstream responds: compute cost from usage → spend from cache + record usage + ledger entry (async, best-effort, never blocks the response) The master key bypasses billing entirely (admin tier). """ from __future__ import annotations import hashlib import logging import re from typing import Optional from . import cache, pricing, registry, user_store logger = logging.getLogger("router.gate") # a valid 20-byte ethereum address: 0x + 40 hex chars (case-insensitive) _ADDR_RE = re.compile(r"^0x[0-9a-f]{40}$") def is_enabled() -> bool: """True if the billing system is fully configured (gist store + wallet). Used throughout to gracefully degrade to unmetered mode when billing secrets aren't set (e.g. local dev, or a Space without GIST_TOKEN). """ from .. import gist_client, wallet # local import avoids a cycle at module load return gist_client.is_configured() and wallet.is_configured() def _sha256_hex(s: str) -> str: """sha256 hex digest of a string (how api keys are stored/compared).""" return hashlib.sha256(s.encode()).hexdigest() def _normalize_addr(token: str) -> Optional[str]: """If `token` is an ETH-address key, return the lowercased address, else None. Accepts either `sk-0x…` (our key format) or a bare `0x…` address. """ t = token.strip() if t[:3].lower() == "sk-": # strip the sk- key prefix if present t = t[3:] t = t.strip().lower() return t if _ADDR_RE.match(t) else None # valid address or None def resolve(authorization: Optional[str], master_key: str) -> tuple[Optional[str], bool]: """Resolve a Bearer header to (user_id | None, is_admin). Returns (None, True) for the master key, (user_id, False) for a valid user credential, (None, False) for an invalid/missing token. A user credential is their ETH deposit address (as `sk-0x…` or bare `0x…`). Legacy `dr_u_…` keys still resolve via the sha256 index (backward compat). """ if not authorization or not authorization.startswith("Bearer "): return None, False token = authorization[7:].strip() if not token: return None, False if master_key and token == master_key: # admin tier — unmetered return None, True addr = _normalize_addr(token) # address-key path (the primary credential) if addr: uid = registry.find_user_by_address(addr) # O(1) address -> user lookup if uid: return uid, False key_hash = _sha256_hex(token) # legacy dr_u_ path (sha256 -> user) user_id = registry.find_user_by_key_hash(key_hash) return user_id, False async def check_balance(user_id: Optional[str]) -> bool: """True if the caller may proceed: admin (None) or a user with credits.""" if user_id is None: # admin return True return await cache.get_balance(user_id) > 0 async def burn(user_id: str, models: dict, model_name: str, tok_in: int, tok_out: int) -> int: """Charge a user for a request: spend credits + log usage + ledger. Returns credits charged. `free:` models cost 0 credits (pricing returns 0), but we STILL record the usage event + ledger entry so the dashboard reflects free-tier activity. """ cost = pricing.cost_credits(models, model_name, tok_in, tok_out) _touch_active(user_id) # stamp last-active (so the sweeper keeps this account) if cost > 0: # paid model — spend credits now cache.spend(user_id, cost) # decrement the in-memory balance instantly gid = cache.gist_id_for(user_id) # fire-and-forget persistence (best-effort) if gid: try: await user_store.record_usage(gid, model_name, tok_in, tok_out, cost) # counters + recent ring await user_store.append_ledger(gid, { "id": f"b_{tok_out}_{cost}", # best-effort unique id "type": "burn", "model": model_name, "tok_in": tok_in, "tok_out": tok_out, "cost": cost, # 0 for free models, >0 for paid "ts": registry._now(), }) except Exception as e: # logging failures must never break the request logger.warning("usage log failed for %s: %s", user_id, e) return cost def _touch_active(user_id: str) -> None: """Stamp the current epoch on the user's in-memory registry record (for the sweeper).""" import time # stdlib epoch rec = registry.find_user(user_id) # the cached record if rec: # found rec["last_active_epoch"] = time.time() # update in-memory (persisted by the sweeper) def extract_usage(payload: dict) -> tuple[int, int]: """Pull (prompt_tokens, completion_tokens) from an OpenAI-shaped response.""" u = payload.get("usage") or {} return int(u.get("prompt_tokens", 0)), int(u.get("completion_tokens", 0))