from __future__ import annotations import asyncio import hashlib import json import math import time import uuid from pathlib import Path from urllib.parse import urlparse import httpx class GatewayUnavailable(RuntimeError): pass class ProviderGateway: DEFAULT_LIMITS = {"startlux": 5, "jev": 3} def __init__(self, db): self.db = db self.preflight_locks = {} self.active = set() self.health_cache = {} self.health_locks = {} def configure(self, path): if not Path(path).exists(): return configs = json.loads(Path(path).read_text()) with self.db.connect(transaction=True) as c: for name, cfg in configs.items(): if name not in self.DEFAULT_LIMITS: raise ValueError("unknown provider") endpoint = cfg["endpoint"] u = urlparse(endpoint) if u.scheme not in ("http", "https") or not u.hostname or u.username: raise ValueError("invalid provider endpoint") rate_rpm = cfg.get("rate_rpm") limit_count = self.DEFAULT_LIMITS[name] if name == "jev": limit_count = cfg.get("max_inflight", limit_count) if isinstance(limit_count, bool) or not isinstance(limit_count, int) or limit_count < 1: raise ValueError("Jev max_inflight must be a positive integer") if rate_rpm is not None and ( isinstance(rate_rpm, bool) or not isinstance(rate_rpm, (int, float)) or not math.isfinite(rate_rpm) or rate_rpm <= 0 ): raise ValueError("Jev rate_rpm must be finite and positive") elif rate_rpm is not None or "max_inflight" in cfg: raise ValueError("rate_rpm/max_inflight are supported only for Jev") c.execute( "INSERT INTO providers(id,limit_count,endpoint,timeout,enabled,health_url,credential_path,deployment_fingerprint) VALUES(?,?,?,?,?,?,?,?) ON CONFLICT(id) DO UPDATE SET endpoint=excluded.endpoint,timeout=excluded.timeout,enabled=excluded.enabled,health_url=excluded.health_url,credential_path=excluded.credential_path,limit_count=excluded.limit_count,deployment_fingerprint=excluded.deployment_fingerprint", ( name, limit_count, endpoint, cfg.get("timeout", 180), int(cfg.get("enabled", True)), cfg.get("health_url"), cfg.get("credential_path"), cfg.get("deployment_fingerprint"), ), ) c.execute( "UPDATE providers SET parameters_json=?,rate_rpm=? WHERE id=?", (self.db.dumps(cfg.get("parameters", {})), rate_rpm, name), ) def mark_drift(self, run_id, reason): with self.db.connect(transaction=True) as c: if ( c.execute("SELECT 1 FROM runs WHERE id=?", (run_id,)).fetchone() and not c.execute( "SELECT 1 FROM events WHERE run_id=? AND type='provider_identity_drift'", (run_id,), ).fetchone() ): c.execute( "INSERT INTO events(run_id,at,type,payload) VALUES(?,?,?,?)", ( run_id, time.time(), "provider_identity_drift", self.db.dumps({"error": reason}), ), ) def _identity(self, run_id): with self.db.connect() as c: run = c.execute( "SELECT config,provider FROM runs WHERE id=?", (run_id,) ).fetchone() if not run or not run["provider"]: return None config = json.loads(run["config"]) expected = config.get("provider_snapshot") current = c.execute( "SELECT * FROM providers WHERE id=?", (run["provider"],) ).fetchone() return config, expected, dict(current) if current else None async def check_identity(self, run_id): identity = await asyncio.to_thread(self._identity, run_id) if identity is None: return config, expected, current = identity if not expected: raise GatewayUnavailable("missing frozen provider identity") changed = [k for k, v in expected.items() if current is None or current[k] != v] if changed: reason = "frozen Provider configuration changed: " + ",".join(changed) await asyncio.to_thread(self.mark_drift, run_id, reason) raise GatewayUnavailable(reason) url = expected.get("health_url") if not url: return # Jev has no deployment-health API; native receipts attest the served model. lock = self.health_locks.setdefault(url, asyncio.Lock()) async with lock: cached = self.health_cache.get(url) if not cached or time.monotonic() - cached[0] > 5: try: async with httpx.AsyncClient(timeout=5, trust_env=False) as client: response = await client.get(url) response.raise_for_status() health = response.json() self.health_cache[url] = (time.monotonic(), health) except Exception as exc: raise GatewayUnavailable( "Provider health unavailable: " + type(exc).__name__ ) from exc else: health = cached[1] mismatches = [] fingerprint = expected.get("deployment_fingerprint") parameters = json.loads(expected.get("parameters_json", "{}")) if fingerprint and health.get("deployment_fingerprint") != fingerprint: mismatches.append("deployment_fingerprint") if health.get("model") != config["prompt_config"].get("model"): mismatches.append("model") if any(health.get("parameters", {}).get(k) != v for k, v in parameters.items()): mismatches.append("parameters") if mismatches: reason = "Provider health identity drift: " + ",".join(mismatches) await asyncio.to_thread(self.mark_drift, run_id, reason) raise GatewayUnavailable(reason) if health.get("status") != "ok" or not health.get("ready_gpu_count"): raise GatewayUnavailable("Provider workers are not ready") def _acquire(self, provider, run_id, call_id, received_at, body_sha256, queue_deadline=None): # The lease transaction is unchanged, but never runs on the event loop. with self.db.connect(transaction=True) as c: if queue_deadline is not None and time.monotonic() >= queue_deadline: return None row = c.execute( "SELECT * FROM providers WHERE id=? AND enabled=1", (provider,) ).fetchone() if not row: raise GatewayUnavailable("provider not configured or disabled") p = dict(row) now = time.time() paced = provider == "jev" and p.get("rate_rpm") is not None if paced: # A lease not yet dispatched excludes other dispatch admissions. # This closes the acquire -> identity check -> send gap across # Gateway instances sharing this database, without a new queue. pending = c.execute( "SELECT 1 FROM provider_calls WHERE provider=? AND ended_at IS NULL AND expires_at>? AND sent_at IS NULL LIMIT 1", (provider, now), ).fetchone() last_sent = c.execute( "SELECT MAX(sent_at) FROM provider_calls WHERE provider=?", (provider,) ).fetchone()[0] if pending or (last_sent is not None and now < last_sent + 60.0 / p["rate_rpm"]): return None active = c.execute( "SELECT COUNT(*) FROM provider_calls WHERE provider=? AND ended_at IS NULL AND expires_at>?", (provider, time.time()), ).fetchone()[0] if active < p["limit_count"]: # A lost upstream connection retains its lease to the configured upstream deadline. c.execute( "INSERT INTO provider_calls(id,provider,run_id,started_at,expires_at,received_at,body_sha256) VALUES(?,?,?,?,?,?,?)", ( call_id, provider, run_id, time.time(), time.time() + p["timeout"] + 60, received_at, body_sha256, ), ) return p return None def _queue_timeout(self, provider, run_id, call_id, received_at, body_sha256): now = time.time() with self.db.connect(transaction=True) as c: c.execute( "INSERT INTO events(run_id,at,type,payload) VALUES(?,?,?,?)", (run_id, now, "provider_queue_timeout", self.db.dumps({ "provider": provider, "call_id": call_id, "received_at": received_at, "queue_seconds": max(0, now - received_at), "body_sha256": body_sha256, })), ) def _mark_sent(self, call_id): with self.db.connect(transaction=True) as c: c.execute("UPDATE provider_calls SET sent_at=? WHERE id=?", (time.time(), call_id)) def _finish(self, call_id, completed, status, error, response_limits, upstream_error): with self.db.connect(transaction=True) as c: c.execute( "UPDATE provider_calls SET ended_at=?,status=?,error=?,completed_at=?,response_limits=?,upstream_error=? WHERE id=?", (time.time() if completed else None, status, error, time.time() if completed else None, self.db.dumps(response_limits), self.db.dumps(upstream_error), call_id), ) async def request(self, provider, run_id, body, headers, call=None): received_at = time.time() await self.check_identity(run_id) deadline = time.monotonic() + 120 call_id = uuid.uuid4().hex body_sha256 = hashlib.sha256(body).hexdigest() while True: if provider == "jev" and time.monotonic() >= deadline: await asyncio.to_thread(self._queue_timeout, provider, run_id, call_id, received_at, body_sha256) raise GatewayUnavailable("provider request queue timeout") p = await asyncio.to_thread( self._acquire, provider, run_id, call_id, received_at, body_sha256, deadline if provider == "jev" else None, ) if p is not None: break if time.monotonic() > deadline: if provider == "jev": await asyncio.to_thread(self._queue_timeout, provider, run_id, call_id, received_at, body_sha256) raise GatewayUnavailable("provider request queue timeout") await asyncio.sleep(0.05) async def send(): request_headers = { "content-type": headers.get("content-type", "application/json") } credential = p.get("credential_path") if credential: request_headers["Authorization"] = ( "Bearer " + Path(credential).read_text().strip() ) elif headers.get("authorization"): request_headers["Authorization"] = headers["authorization"] request_headers["X-Experiment-Run-ID"] = run_id if headers.get("x-startlux-request-class"): request_headers["X-StartLux-Request-Class"] = headers[ "x-startlux-request-class" ] status = None error = None completed = False response_limits = {} upstream_error = None try: await self.check_identity(run_id) paced = provider == "jev" and p.get("rate_rpm") is not None dispatched = asyncio.Event() initiated = False async def invoke(): nonlocal initiated try: if call is not None: initiated = True dispatched.set() return await call() async with httpx.AsyncClient(timeout=p["timeout"], trust_env=False) as client: initiated = True dispatched.set() return await client.post(p["endpoint"], content=body, headers=request_headers) finally: # Wake the sender even if client construction fails. dispatched.set() if paced: # Keep the unsent lease until the transport coroutine has # entered call/post. Recording afterwards is conservative: # later dispatches wait a full interval from this observed # initiation, including arbitrary pre-dispatch delays. upstream = asyncio.create_task(invoke()) try: await dispatched.wait() if initiated: await asyncio.to_thread(self._mark_sent, call_id) except BaseException: # A dispatched request still owns its concurrency lease. await asyncio.gather(upstream, return_exceptions=True) raise response = await upstream else: await asyncio.to_thread(self._mark_sent, call_id) response = await invoke() if call is not None: status = response[0] if isinstance(response, tuple) and response else 200 completed = True return response status = response.status_code response_limits = { key: response.headers[key][:256] for key in ( "retry-after", "x-request-id", "x-ratelimit-limit-requests", "x-ratelimit-remaining-requests", "x-ratelimit-reset-requests", "x-ratelimit-limit-tokens", "x-ratelimit-remaining-tokens", "x-ratelimit-reset-tokens", ) if key in response.headers } if status >= 400: try: detail = response.json().get("error", {}) if isinstance(detail, dict): upstream_error = { key: str(detail[key])[:128] for key in ("code", "type") if key in detail } except (ValueError, AttributeError): pass completed = True return ( status, response.content, response.headers.get("content-type", "application/json"), ) except Exception as exc: error = type(exc).__name__ raise GatewayUnavailable(error) from exc finally: await asyncio.to_thread(self._finish, call_id, completed, status, error, response_limits, upstream_error) task = asyncio.create_task(send()) self.active.add(task) task.add_done_callback(self.active.discard) return await asyncio.shield(task) def _preflight_cached(self, provider, key): with self.db.connect() as c: row = c.execute( "SELECT * FROM preflights WHERE key=? AND provider=?", (key, provider) ).fetchone() return dict(row) if row else None def _preflight_store(self, provider, key, result): with self.db.connect(transaction=True) as c: c.execute( "INSERT OR REPLACE INTO preflights VALUES(?,?,?,?)", (key, provider, time.time(), json.dumps(result)), ) async def preflight(self, provider, key, call): lock = self.preflight_locks.setdefault(provider, asyncio.Lock()) async with lock: row = await asyncio.to_thread(self._preflight_cached, provider, key) if row and time.time() - row["checked_at"] < 120: return {**json.loads(row["result"]), "cached": True} result = await call() if result.get("ok"): await asyncio.to_thread(self._preflight_store, provider, key, result) return {**result, "cached": False} def overview(self): with self.db.connect() as c: rows = c.execute( "SELECT p.*, (SELECT COUNT(*) FROM provider_calls q WHERE q.provider=p.id AND q.ended_at IS NULL AND q.expires_at>?) AS active_requests FROM providers p ORDER BY p.id", (time.time(),), ).fetchall() return [ {k: v for k, v in dict(r).items() if k != "credential_path"} for r in rows ]