1from __future__ import annotations2import asyncio3import hashlib4import json5import math6import time7import uuid8from pathlib import Path9from urllib.parse import urlparse10import httpx111213class GatewayUnavailable(RuntimeError):14 pass151617class ProviderGateway:18 DEFAULT_LIMITS = {"startlux": 5, "jev": 3}1920 def __init__(self, db):21 self.db = db22 self.preflight_locks = {}23 self.active = set()24 self.health_cache = {}25 self.health_locks = {}2627 def configure(self, path):28 if not Path(path).exists():29 return30 configs = json.loads(Path(path).read_text())31 with self.db.connect(transaction=True) as c:32 for name, cfg in configs.items():33 if name not in self.DEFAULT_LIMITS:34 raise ValueError("unknown provider")35 endpoint = cfg["endpoint"]36 u = urlparse(endpoint)37 if u.scheme not in ("http", "https") or not u.hostname or u.username:38 raise ValueError("invalid provider endpoint")39 rate_rpm = cfg.get("rate_rpm")40 limit_count = self.DEFAULT_LIMITS[name]41 if name == "jev":42 limit_count = cfg.get("max_inflight", limit_count)43 if isinstance(limit_count, bool) or not isinstance(limit_count, int) or limit_count < 1:44 raise ValueError("Jev max_inflight must be a positive integer")45 if rate_rpm is not None and (46 isinstance(rate_rpm, bool) or not isinstance(rate_rpm, (int, float))47 or not math.isfinite(rate_rpm) or rate_rpm <= 048 ):49 raise ValueError("Jev rate_rpm must be finite and positive")50 elif rate_rpm is not None or "max_inflight" in cfg:51 raise ValueError("rate_rpm/max_inflight are supported only for Jev")52 c.execute(53 "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",54 (55 name,56 limit_count,57 endpoint,58 cfg.get("timeout", 180),59 int(cfg.get("enabled", True)),60 cfg.get("health_url"),61 cfg.get("credential_path"),62 cfg.get("deployment_fingerprint"),63 ),64 )65 c.execute(66 "UPDATE providers SET parameters_json=?,rate_rpm=? WHERE id=?",67 (self.db.dumps(cfg.get("parameters", {})), rate_rpm, name),68 )6970 def mark_drift(self, run_id, reason):71 with self.db.connect(transaction=True) as c:72 if (73 c.execute("SELECT 1 FROM runs WHERE id=?", (run_id,)).fetchone()74 and not c.execute(75 "SELECT 1 FROM events WHERE run_id=? AND type='provider_identity_drift'",76 (run_id,),77 ).fetchone()78 ):79 c.execute(80 "INSERT INTO events(run_id,at,type,payload) VALUES(?,?,?,?)",81 (82 run_id,83 time.time(),84 "provider_identity_drift",85 self.db.dumps({"error": reason}),86 ),87 )8889 def _identity(self, run_id):90 with self.db.connect() as c:91 run = c.execute(92 "SELECT config,provider FROM runs WHERE id=?", (run_id,)93 ).fetchone()94 if not run or not run["provider"]:95 return None96 config = json.loads(run["config"])97 expected = config.get("provider_snapshot")98 current = c.execute(99 "SELECT * FROM providers WHERE id=?", (run["provider"],)100 ).fetchone()101 return config, expected, dict(current) if current else None102103 async def check_identity(self, run_id):104 identity = await asyncio.to_thread(self._identity, run_id)105 if identity is None:106 return107 config, expected, current = identity108 if not expected:109 raise GatewayUnavailable("missing frozen provider identity")110 changed = [k for k, v in expected.items() if current is None or current[k] != v]111 if changed:112 reason = "frozen Provider configuration changed: " + ",".join(changed)113 await asyncio.to_thread(self.mark_drift, run_id, reason)114 raise GatewayUnavailable(reason)115 url = expected.get("health_url")116 if not url:117 return # Jev has no deployment-health API; native receipts attest the served model.118 lock = self.health_locks.setdefault(url, asyncio.Lock())119 async with lock:120 cached = self.health_cache.get(url)121 if not cached or time.monotonic() - cached[0] > 5:122 try:123 async with httpx.AsyncClient(timeout=5, trust_env=False) as client:124 response = await client.get(url)125 response.raise_for_status()126 health = response.json()127 self.health_cache[url] = (time.monotonic(), health)128 except Exception as exc:129 raise GatewayUnavailable(130 "Provider health unavailable: " + type(exc).__name__131 ) from exc132 else:133 health = cached[1]134 mismatches = []135 fingerprint = expected.get("deployment_fingerprint")136 parameters = json.loads(expected.get("parameters_json", "{}"))137 if fingerprint and health.get("deployment_fingerprint") != fingerprint:138 mismatches.append("deployment_fingerprint")139 if health.get("model") != config["prompt_config"].get("model"):140 mismatches.append("model")141 if any(health.get("parameters", {}).get(k) != v for k, v in parameters.items()):142 mismatches.append("parameters")143 if mismatches:144 reason = "Provider health identity drift: " + ",".join(mismatches)145 await asyncio.to_thread(self.mark_drift, run_id, reason)146 raise GatewayUnavailable(reason)147 if health.get("status") != "ok" or not health.get("ready_gpu_count"):148 raise GatewayUnavailable("Provider workers are not ready")149150 def _acquire(self, provider, run_id, call_id, received_at, body_sha256, queue_deadline=None):151 # The lease transaction is unchanged, but never runs on the event loop.152 with self.db.connect(transaction=True) as c:153 if queue_deadline is not None and time.monotonic() >= queue_deadline:154 return None155 row = c.execute(156 "SELECT * FROM providers WHERE id=? AND enabled=1", (provider,)157 ).fetchone()158 if not row:159 raise GatewayUnavailable("provider not configured or disabled")160 p = dict(row)161 now = time.time()162 paced = provider == "jev" and p.get("rate_rpm") is not None163 if paced:164 # A lease not yet dispatched excludes other dispatch admissions.165 # This closes the acquire -> identity check -> send gap across166 # Gateway instances sharing this database, without a new queue.167 pending = c.execute(168 "SELECT 1 FROM provider_calls WHERE provider=? AND ended_at IS NULL AND expires_at>? AND sent_at IS NULL LIMIT 1",169 (provider, now),170 ).fetchone()171 last_sent = c.execute(172 "SELECT MAX(sent_at) FROM provider_calls WHERE provider=?", (provider,)173 ).fetchone()[0]174 if pending or (last_sent is not None and now < last_sent + 60.0 / p["rate_rpm"]):175 return None176 active = c.execute(177 "SELECT COUNT(*) FROM provider_calls WHERE provider=? AND ended_at IS NULL AND expires_at>?",178 (provider, time.time()),179 ).fetchone()[0]180 if active < p["limit_count"]:181 # A lost upstream connection retains its lease to the configured upstream deadline.182 c.execute(183 "INSERT INTO provider_calls(id,provider,run_id,started_at,expires_at,received_at,body_sha256) VALUES(?,?,?,?,?,?,?)",184 (185 call_id,186 provider,187 run_id,188 time.time(),189 time.time() + p["timeout"] + 60,190 received_at,191 body_sha256,192 ),193 )194 return p195 return None196197 def _queue_timeout(self, provider, run_id, call_id, received_at, body_sha256):198 now = time.time()199 with self.db.connect(transaction=True) as c:200 c.execute(201 "INSERT INTO events(run_id,at,type,payload) VALUES(?,?,?,?)",202 (run_id, now, "provider_queue_timeout", self.db.dumps({203 "provider": provider, "call_id": call_id, "received_at": received_at,204 "queue_seconds": max(0, now - received_at), "body_sha256": body_sha256,205 })),206 )207208 def _mark_sent(self, call_id):209 with self.db.connect(transaction=True) as c:210 c.execute("UPDATE provider_calls SET sent_at=? WHERE id=?", (time.time(), call_id))211212 def _finish(self, call_id, completed, status, error, response_limits, upstream_error):213 with self.db.connect(transaction=True) as c:214 c.execute(215 "UPDATE provider_calls SET ended_at=?,status=?,error=?,completed_at=?,response_limits=?,upstream_error=? WHERE id=?",216 (time.time() if completed else None, status, error, time.time() if completed else None,217 self.db.dumps(response_limits), self.db.dumps(upstream_error), call_id),218 )219220 async def request(self, provider, run_id, body, headers, call=None):221 received_at = time.time()222 await self.check_identity(run_id)223 deadline = time.monotonic() + 120224 call_id = uuid.uuid4().hex225 body_sha256 = hashlib.sha256(body).hexdigest()226 while True:227 if provider == "jev" and time.monotonic() >= deadline:228 await asyncio.to_thread(self._queue_timeout, provider, run_id, call_id, received_at, body_sha256)229 raise GatewayUnavailable("provider request queue timeout")230 p = await asyncio.to_thread(231 self._acquire, provider, run_id, call_id, received_at, body_sha256,232 deadline if provider == "jev" else None,233 )234 if p is not None:235 break236 if time.monotonic() > deadline:237 if provider == "jev":238 await asyncio.to_thread(self._queue_timeout, provider, run_id, call_id, received_at, body_sha256)239 raise GatewayUnavailable("provider request queue timeout")240 await asyncio.sleep(0.05)241242 async def send():243 request_headers = {244 "content-type": headers.get("content-type", "application/json")245 }246 credential = p.get("credential_path")247 if credential:248 request_headers["Authorization"] = (249 "Bearer " + Path(credential).read_text().strip()250 )251 elif headers.get("authorization"):252 request_headers["Authorization"] = headers["authorization"]253 request_headers["X-Experiment-Run-ID"] = run_id254 if headers.get("x-startlux-request-class"):255 request_headers["X-StartLux-Request-Class"] = headers[256 "x-startlux-request-class"257 ]258 status = None259 error = None260 completed = False261 response_limits = {}262 upstream_error = None263 try:264 await self.check_identity(run_id)265 paced = provider == "jev" and p.get("rate_rpm") is not None266 dispatched = asyncio.Event()267 initiated = False268269 async def invoke():270 nonlocal initiated271 try:272 if call is not None:273 initiated = True274 dispatched.set()275 return await call()276 async with httpx.AsyncClient(timeout=p["timeout"], trust_env=False) as client:277 initiated = True278 dispatched.set()279 return await client.post(p["endpoint"], content=body, headers=request_headers)280 finally:281 # Wake the sender even if client construction fails.282 dispatched.set()283284 if paced:285 # Keep the unsent lease until the transport coroutine has286 # entered call/post. Recording afterwards is conservative:287 # later dispatches wait a full interval from this observed288 # initiation, including arbitrary pre-dispatch delays.289 upstream = asyncio.create_task(invoke())290 try:291 await dispatched.wait()292 if initiated:293 await asyncio.to_thread(self._mark_sent, call_id)294 except BaseException:295 # A dispatched request still owns its concurrency lease.296 await asyncio.gather(upstream, return_exceptions=True)297 raise298 response = await upstream299 else:300 await asyncio.to_thread(self._mark_sent, call_id)301 response = await invoke()302 if call is not None:303 status = response[0] if isinstance(response, tuple) and response else 200304 completed = True305 return response306 status = response.status_code307 response_limits = {308 key: response.headers[key][:256]309 for key in (310 "retry-after", "x-request-id", "x-ratelimit-limit-requests",311 "x-ratelimit-remaining-requests", "x-ratelimit-reset-requests",312 "x-ratelimit-limit-tokens", "x-ratelimit-remaining-tokens",313 "x-ratelimit-reset-tokens",314 )315 if key in response.headers316 }317 if status >= 400:318 try:319 detail = response.json().get("error", {})320 if isinstance(detail, dict):321 upstream_error = {322 key: str(detail[key])[:128]323 for key in ("code", "type") if key in detail324 }325 except (ValueError, AttributeError):326 pass327 completed = True328 return (329 status,330 response.content,331 response.headers.get("content-type", "application/json"),332 )333 except Exception as exc:334 error = type(exc).__name__335 raise GatewayUnavailable(error) from exc336 finally:337 await asyncio.to_thread(self._finish, call_id, completed, status, error, response_limits, upstream_error)338339 task = asyncio.create_task(send())340 self.active.add(task)341 task.add_done_callback(self.active.discard)342 return await asyncio.shield(task)343344 def _preflight_cached(self, provider, key):345 with self.db.connect() as c:346 row = c.execute(347 "SELECT * FROM preflights WHERE key=? AND provider=?", (key, provider)348 ).fetchone()349 return dict(row) if row else None350351 def _preflight_store(self, provider, key, result):352 with self.db.connect(transaction=True) as c:353 c.execute(354 "INSERT OR REPLACE INTO preflights VALUES(?,?,?,?)",355 (key, provider, time.time(), json.dumps(result)),356 )357358 async def preflight(self, provider, key, call):359 lock = self.preflight_locks.setdefault(provider, asyncio.Lock())360 async with lock:361 row = await asyncio.to_thread(self._preflight_cached, provider, key)362 if row and time.time() - row["checked_at"] < 120:363 return {**json.loads(row["result"]), "cached": True}364 result = await call()365 if result.get("ok"):366 await asyncio.to_thread(self._preflight_store, provider, key, result)367 return {**result, "cached": False}368369 def overview(self):370 with self.db.connect() as c:371 rows = c.execute(372 "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",373 (time.time(),),374 ).fetchall()375 return [376 {k: v for k, v in dict(r).items() if k != "credential_path"} for r in rows377 ]