diff --git a/.env.example b/.env.example index 8d5e9f6..0ec92fe 100644 --- a/.env.example +++ b/.env.example @@ -106,6 +106,13 @@ ANTSEED_MAX_OUTPUT=1000 # NOTE this endpoint is UNTRUSTED — the funding keeper below treats its readings # as vetoes and alerts only, never as permission to spend. ANTSEED_WALLET_RPC_URL= + +# Optional external control plane ("bring your own control plane"): a separate +# service owning consumer keys / tenants (e.g. unhardcoded-cloud). Both must be +# set to enable; otherwise the router runs operator-only. The secret gates the +# ingress /internal/usage surface AND authenticates outbound resolve calls. +CONTROL_PLANE_URL= +CONTROL_PLANE_INTERNAL_SECRET= # --- funding autonomy (wallet_keeper.py) ------------------------------------- # The keeper reclaims stuck payment channels and tops the escrow up on its own. # It SHIPS DARK: 0 = off (default). Arm it deliberately, from the dashboard diff --git a/antseed/byo-gateway.js b/antseed/byo-gateway.js new file mode 100644 index 0000000..18f6fac --- /dev/null +++ b/antseed/byo-gateway.js @@ -0,0 +1,85 @@ +// Optional customer-owned BYO gateway. Separate from wallet control: this token +// can spend through inference, but cannot deposit, withdraw or export keys. +// Terminate HTTPS at the customer's reverse proxy; do not expose the raw proxy. +'use strict'; +const http = require('http'); +const { timingSafeEqual } = require('crypto'); + +function createGateway({ token, snapshot, proxyPort = 8377 }) { + if (!token || token.length < 32) throw new Error('ANTSEED_BYO_TOKEN must have at least 32 characters'); + const expected = Buffer.from('Bearer ' + token); + return http.createServer(async (req, res) => { + const auth = Buffer.from(req.headers.authorization || ''); + const send = (code, data) => { + res.writeHead(code, { 'content-type': 'application/json', 'cache-control': 'no-store' }); + res.end(JSON.stringify(data)); + }; + if (auth.length !== expected.length || !timingSafeEqual(auth, expected)) { + return send(401, { error: 'unauthorized' }); + } + if (req.method === 'GET' && req.url === '/snapshot') { + try { return send(200, await snapshot()); } + catch (_) { return send(503, { error: 'buyer_state_unavailable' }); } + } + if (req.method !== 'POST' || req.url !== '/v1/chat/completions') { + return send(404, { error: 'unsupported_endpoint' }); + } + // Fixed local upstream, no caller-chosen URL or forwarded credentials. + const chunks = []; + let bytes = 0; + try { + for await (const chunk of req) { + bytes += chunk.length; + if (bytes > 1024 * 1024) return send(413, { error: 'request_too_large' }); + chunks.push(chunk); + } + } catch (_) { return res.destroy(); } + const headers = { 'content-type': 'application/json' }; + const peer = req.headers['x-antseed-pin-peer']; + if (typeof peer !== 'string' || !peer || peer.length > 256) { + return send(400, { error: 'peer_pin_required' }); + } + headers['x-antseed-pin-peer'] = peer; + const upstream = http.request({ hostname: '127.0.0.1', port: proxyPort, + path: '/v1/chat/completions', method: 'POST', headers, timeout: 120000 }, response => { + res.writeHead(response.statusCode, { + 'content-type': response.headers['content-type'] || 'application/json', + 'cache-control': 'no-store', + }); + response.on('error', () => res.destroy()); + response.pipe(res); + }); + upstream.on('timeout', () => upstream.destroy()); + upstream.on('error', () => { + if (!res.headersSent) send(502, { error: 'buyer_unavailable' }); + else res.destroy(); + }); + res.on('close', () => upstream.destroy()); + upstream.end(Buffer.concat(chunks)); + }); +} + +if (require.main === module && process.env.ANTSEED_BYO_TOKEN) { + const { Pool } = require('pg'); + const { pgConfig } = require('./db.js'); + const pool = new Pool({ ...pgConfig(process.env.DATABASE_URL), statement_timeout: 5000 }); + const pid = process.env.ANTSEED_BUYER_PID || 'antseed'; + const snapshot = async () => { + const [offers, status] = await Promise.all([ + pool.query(`SELECT peer_id, service, price_in, price_out, price_cached_in, + max_concurrency, reputation, last_seen, last_reached_at, observed_at + FROM peer_offers WHERE observed_at >= $1 LIMIT 10000`, [Date.now() - 900000]), + pool.query(`SELECT pinned_peer_id, deposits_available, fetched_at + FROM buyer_status WHERE pid = $1`, [pid]), + ]); + return { peer_offers: offers.rows.map(row => ({ ...row, observed_at: Number(row.observed_at), + last_reached_at: row.last_reached_at == null ? null : Number(row.last_reached_at), + last_seen: row.last_seen == null ? null : Number(row.last_seen) })), + buyer_status: status.rows[0] ? { ...status.rows[0], fetched_at: Number(status.rows[0].fetched_at) } : null }; + }; + createGateway({ token: process.env.ANTSEED_BYO_TOKEN, snapshot, + proxyPort: Number(process.env.ANTSEED_PROXY_PORT || 8377) }) + .listen(Number(process.env.ANTSEED_BYO_PORT || 8380), '0.0.0.0'); +} + +module.exports = { createGateway }; diff --git a/antseed/byo-gateway.test.js b/antseed/byo-gateway.test.js new file mode 100644 index 0000000..8cbbf0f --- /dev/null +++ b/antseed/byo-gateway.test.js @@ -0,0 +1,35 @@ +const { test } = require('node:test'); +const assert = require('node:assert/strict'); +const http = require('node:http'); +const { createGateway } = require('./byo-gateway.js'); + +test('BYO token permits snapshot/inference, never wallet ops or unpinned traffic', async () => { + const received = []; + const proxy = http.createServer((req, res) => { + received.push(req.headers); + res.setHeader('content-type', 'text/event-stream'); + res.end('data: {"ok":true}\n\n'); + }); + await new Promise(resolve => proxy.listen(0, '127.0.0.1', resolve)); + const token = 't'.repeat(40); + const gateway = createGateway({token, proxyPort: proxy.address().port, + snapshot: async () => ({ peer_offers: [], buyer_status: { deposits_available: 3 } })}); + await new Promise(resolve => gateway.listen(0, '127.0.0.1', resolve)); + const base = 'http://127.0.0.1:' + gateway.address().port; + const headers = { authorization: 'Bearer ' + token }; + try { + assert.equal((await fetch(base + '/snapshot')).status, 401); + assert.equal((await (await fetch(base + '/snapshot', {headers})).json()).buyer_status.deposits_available, 3); + for (const path of ['/deposit','/withdraw','/reclaim','/status','/v1/models']) { + assert.equal((await fetch(base + path, {method:'POST', headers})).status, 404); + } + assert.equal((await fetch(base + '/v1/chat/completions', {method:'POST',headers,body:'{}'})).status, 400); + const result = await fetch(base + '/v1/chat/completions', {method:'POST', + headers:{...headers,'x-antseed-pin-peer':'peer-a'},body:'{}'}); + assert.equal(await result.text(), 'data: {"ok":true}\n\n'); + assert.equal(received[0]['x-antseed-pin-peer'], 'peer-a'); + assert.equal(received[0].authorization, undefined); + } finally { + await Promise.all([new Promise(resolve => gateway.close(resolve)), new Promise(resolve => proxy.close(resolve))]); + } +}); diff --git a/antseed/entrypoint.sh b/antseed/entrypoint.sh index be6078a..54f1805 100644 --- a/antseed/entrypoint.sh +++ b/antseed/entrypoint.sh @@ -62,6 +62,9 @@ socat "TCP-LISTEN:${PORT_PUBLIC},fork,reuseaddr" "TCP:127.0.0.1:${PORT_PROXY}" & # Self-disables when ANTSEED_CONTROL_TOKEN is unset. See antseed/control.js. node "$LIB/control.js" & +# Separate inference + read-only discovery token; no wallet control authority. +node "$LIB/byo-gateway.js" & + write_market() { raw="$MARKET_DIR/.market.raw.$$" timeout -k 5 "$CLI_TIMEOUT" antseed network browse --services --top "$TOP" --json > "$raw" 2>/dev/null || true diff --git a/auth_proxy.py b/auth_proxy.py index bc070eb..1a8ffc1 100644 --- a/auth_proxy.py +++ b/auth_proxy.py @@ -24,7 +24,9 @@ from fastapi.responses import HTMLResponse, JSONResponse, Response, StreamingResponse from starlette.background import BackgroundTask +import control_plane_client import host_store +import internal_api from env_secrets import load_env_secrets from shim import _CACHE_READ_FACTOR # the billing cache-read discount — one source @@ -85,6 +87,10 @@ SYNTHETIC_PROBE_INITIAL_DELAY_S = float(os.getenv("DASHBOARD_SYNTHETIC_PROBE_INITIAL_DELAY_S", "45")) SYNTHETIC_PROBE_TIMEOUT_S = float(os.getenv("DASHBOARD_SYNTHETIC_PROBE_TIMEOUT_S", "45")) SYNTHETIC_PROBE_CALLER = os.getenv("DASHBOARD_SYNTHETIC_PROBE_CALLER", "dashboard-probe") +# Optional namespace for control-plane-resolved callers (e.g. "cp:") so tenant +# slugs can't merge attribution with operator-minted consumer names. Default off: +# calls.caller == tenant slug, which is what the control plane's usage reads key on. +CONTROL_PLANE_CALLER_PREFIX = os.getenv("CONTROL_PLANE_CALLER_PREFIX", "") @@ -119,6 +125,10 @@ def _bootstrap_caller_key_hashes(raw: str) -> Dict[str, str]: logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"), format="%(message)s") log = logging.getLogger("llm-router-auth-proxy") app = FastAPI(title="llm-router auth proxy", docs_url=None, redoc_url=None) +# Control-plane metering surface (/internal/usage[/recent]). Registered before +# any route definition below so Starlette matches it ahead of the catch-all +# proxy; hidden (404) unless CONTROL_PLANE_INTERNAL_SECRET is configured. +app.include_router(internal_api.router) _client: httpx.AsyncClient | None = None _probe_task: asyncio.Task[None] | None = None _started_wall = time.time() @@ -734,6 +744,42 @@ def _caller_auth(token: str | None) -> dict[str, Any]: return {"ok": True, "caller": caller, "digest": digest, "storage": storage, "meta": meta} +async def _caller_auth_async(token: str | None) -> dict[str, Any]: + """_caller_auth plus the external-control-plane fallback. + + Local key stores stay authoritative: a locally known key — including a + revoked/inactive one — never falls through to the control plane; only a + pure local miss consults it (and only when the feature is configured, see + control_plane_client.enabled()). On a resolve hit the caller becomes the + control-plane consumer name (tenant slug) and the plan's rate limits ride + the returned meta; an explicit local record for that same name keeps + precedence (operator kill-switch: mark the slug inactive locally).""" + auth = await asyncio.to_thread(_caller_auth, token) + if auth.get("ok") or auth.get("caller") or not token or not control_plane_client.enabled(): + return auth + digest = hashlib.sha256(token.encode()).hexdigest() + resolved = await control_plane_client.resolve_key(digest) + if resolved is None or not resolved.active or not resolved.consumer: + return auth + caller = f"{CONTROL_PLANE_CALLER_PREFIX}{resolved.consumer}" + meta = await asyncio.to_thread(_consumer_meta, caller) + if meta.get("keys"): + # A local consumer already answers to this name — attribution merges. + control_plane_client.log_collision_once(caller) + if meta.get("status") != "active": + return {"ok": False, "caller": caller, "digest": digest, + "storage": "control_plane", "error_code": "caller_inactive"} + if meta.get("rate_per_min") is None: + meta["rate_per_min"] = resolved.rate_per_min + if meta.get("burst") is None: + meta["burst"] = resolved.burst + out = {"ok": True, "caller": caller, "digest": digest, + "storage": "control_plane", "meta": meta} + if resolved.tenant_id is not None: + out["tenant_id"] = resolved.tenant_id + return out + + def _rate_ok(caller: str, meta: dict[str, Any] | None = None) \ -> tuple[bool, bool, float]: meta = meta or _consumer_meta(caller) @@ -773,7 +819,8 @@ def _route_matches(pattern: str, route: str) -> bool: def _route_allowed(caller: str, route: str | None, meta: dict[str, Any] | None = None) -> bool: - allowed_routes = (meta or _consumer_meta(caller)).get("allowed_routes") or [] + meta = meta if meta is not None else _consumer_meta(caller) + allowed_routes = meta.get("allowed_routes") or [] if not allowed_routes: return True if not route: @@ -1251,6 +1298,7 @@ async def shutdown() -> None: _probe_task = None if _client: await _client.aclose() + await control_plane_client.close() @app.get("/healthz") @@ -2029,7 +2077,7 @@ async def consumer_skill(request: Request) -> Response: and the same key auth (revoked/expired/inactive keys are rejected).""" auth = request.headers.get("authorization") or "" token = auth[7:].strip() if auth.lower().startswith("bearer ") else auth.strip() - if not _caller_auth(token).get("ok"): + if not (await _caller_auth_async(token)).get("ok"): return JSONResponse(status_code=401, content={"error": { "message": "invalid or inactive API key", "type": "auth_error", "code": "consumer_auth"}}) @@ -2048,7 +2096,7 @@ async def consumer_skill(request: Request) -> Response: # or mutate state: they admit/normalize a term, preview the ranking, or list the # field vocabulary — all derived from the live catalog the skill already shows. async def _consumer_x_proxy(request: Request, method: str, path: str) -> Response: - if not _caller_auth(_extract_token(request)).get("ok"): + if not (await _caller_auth_async(_extract_token(request))).get("ok"): return JSONResponse(status_code=401, content={"error": { "message": "invalid or inactive API key", "type": "auth_error", "code": "consumer_auth"}}) @@ -2689,7 +2737,7 @@ async def dashboard_key_usage(request: Request) -> Response: @app.get("/api/usage") async def key_usage(request: Request) -> Response: token = _extract_token(request) - auth = _caller_auth(token) + auth = await _caller_auth_async(token) if not auth.get("ok"): code = auth.get("error_code") or "caller_auth" status_code = 403 if code in {"caller_inactive", "caller_key_revoked", "caller_key_expired"} else 401 @@ -2717,7 +2765,7 @@ async def session_view(sid: str, request: Request) -> Response: itself sends as X-Unhardcoded-Session. Lets a harness (e.g. the opencode plugin) show live router economics without operator access to /x/*.""" token = _extract_token(request) - auth = _caller_auth(token) + auth = await _caller_auth_async(token) if not auth.get("ok"): code = auth.get("error_code") or "caller_auth" status_code = 403 if code in {"caller_inactive", "caller_key_revoked", "caller_key_expired"} else 401 @@ -3857,7 +3905,7 @@ async def proxy(path: str, request: Request) -> Response: return JSONResponse(status_code=404, content={"error": {"message": "not found", "type": "invalid_request_error", "code": None}}) started = time.perf_counter() token = _extract_token(request) - auth = await asyncio.to_thread(_caller_auth, token) + auth = await _caller_auth_async(token) caller = auth.get("caller") if not auth.get("ok"): code = auth.get("error_code") or "caller_auth" @@ -3882,6 +3930,28 @@ async def proxy(path: str, request: Request) -> Response: consumer_meta = auth.get("meta") or {} body = await request.body() requested_route = _requested_route_from(path, body) + published_route = None + if auth.get("tenant_id") is not None: + # SaaS keys address published routes only. Raw policies, profiles and + # pins cannot replace the workspace's published restrictions. + import re + if (request.method != "POST" or path not in {"v1/chat/completions", "v1/responses"} + or not re.fullmatch(r"route:[a-z0-9][a-z0-9-]{0,79}", requested_route or "")): + return JSONResponse(status_code=400, content={"error": { + "message": "Use a published route, for example model='route:production'.", + "type": "invalid_request_error", "code": "route_required"}}) + from route_contract import apply_contract, PreferenceNotAllowed + try: + published_route = await control_plane_client.resolve_route( + auth["tenant_id"], requested_route[6:]) + payload, published_route = apply_contract(json.loads(body), published_route) + body = json.dumps(payload, separators=(",", ":")).encode() + except PreferenceNotAllowed as exc: + return JSONResponse(status_code=400, content={'error': { + 'message': str(exc), 'type': 'invalid_request_error', 'code': 'preference_not_allowed'}}) + except control_plane_client.RouteUnavailable as exc: + return JSONResponse(status_code=503, content={"error": { + "message": str(exc), "type": "server_error", "code": "route_unavailable"}}) if not _route_allowed(caller, requested_route, consumer_meta): _record_reject(reason="route_not_allowed", path="/" + path, caller=caller, status=403, route=requested_route) _log({"event": "reject", "reason": "route_not_allowed", "caller": caller, "path": "/" + path, "route": requested_route}) @@ -3953,9 +4023,21 @@ async def proxy(path: str, request: Request) -> Response: headers = { k: v for k, v in request.headers.items() - if k.lower() not in {"authorization", "host", "connection", "content-length"} + if k.lower() not in {"authorization", "host", "connection", "content-length", + "x-llm-router-tenant", "x-internal-secret", + "x-unhardcoded-route", "x-unhardcoded-revision", "x-unhardcoded-policy-id", "x-unhardcoded-preference"} } headers["x-llm-router-caller"] = caller + # Tenant identity for per-tenant provider credentials (BYO keys): set ONLY + # from the authenticated resolve — the client-sent header is stripped above, + # so it can never be smuggled past auth. + if auth.get("tenant_id") is not None: + headers["x-llm-router-tenant"] = str(auth["tenant_id"]) + headers["x-internal-secret"] = control_plane_client.CONTROL_PLANE_INTERNAL_SECRET + headers["x-unhardcoded-route"] = requested_route + headers["x-unhardcoded-revision"] = str(published_route["revision"]) + headers["x-unhardcoded-policy-id"] = published_route["policy_id"] + headers['x-unhardcoded-preference'] = published_route.get('routing_preference', 'default') status = 502 provider = None @@ -3979,11 +4061,16 @@ async def proxy(path: str, request: Request) -> Response: session_id = request.headers.get("x-unhardcoded-session") def _finish(): - nonlocal capacity_released + nonlocal capacity_released, decision_trace if capacity_released: return capacity_released = True latency_ms = round((time.perf_counter() - started) * 1000, 1) + if published_route: + decision_trace = {**(decision_trace or {}), 'route': requested_route, + 'route_revision': published_route['revision'], + 'policy_id': published_route['policy_id'], + 'routing_preference': published_route.get('routing_preference', 'default')} try: _record_request(caller=caller, method=request.method, path="/" + path, status=status, latency_ms=latency_ms, provider=provider, model_family=model_family, served_model_id=served_model_id, served_by=served_by, requested_model=requested_model, session=session_id, tokens_in=tokens_in, tokens_out=tokens_out, tokens_total=tokens_total, tokens_cached=tokens_cached, cost_usd=cost_usd, cost_basis=cost_basis, decision_trace=decision_trace, error_type=error_type, error_code=error_code, error_message=error_message, key_sha256=auth.get("digest")) _metric_request(status, latency_ms) diff --git a/byo_http.py b/byo_http.py new file mode 100644 index 0000000..f66b9c5 --- /dev/null +++ b/byo_http.py @@ -0,0 +1,48 @@ +"""HTTPS egress for tenant-supplied buyer endpoints. + +Resolve and pin a public IP on every new request. Preserve TLS hostname/SNI; +never follow redirects or use process HTTP proxies. DNS rebinding cannot turn +the validated hostname into a subsequent private-address connection. +""" +import asyncio +import ipaddress +import socket + +import httpx + + +class PublicHTTPS(httpx.AsyncBaseTransport): + def __init__(self): + # Pooling by pinned IP must not reuse a TLS session across two distinct + # hostnames that resolve to the same address. + self.inner = httpx.AsyncHTTPTransport(retries=0, limits=httpx.Limits(max_keepalive_connections=0)) + + async def handle_async_request(self, request): + url = request.url + if url.scheme != 'https' or url.username or url.password or url.fragment: + raise httpx.ConnectError('Buyer gateway requires a public HTTPS origin', request=request) + try: + rows = await asyncio.wait_for(asyncio.get_running_loop().getaddrinfo( + url.host, url.port or 443, type=socket.SOCK_STREAM), 5) + addresses = [ipaddress.ip_address(row[4][0]) for row in rows] + if not addresses or any(not ip.is_global for ip in addresses): + raise ValueError('non-public endpoint') + except (ValueError, OSError, TimeoutError) as exc: + raise httpx.ConnectError('Buyer gateway is not a public HTTPS endpoint', request=request) from exc + pinned = httpx.Request(request.method, url.copy_with(host=str(addresses[0])), + headers=request.headers, stream=request.stream, + extensions={**request.extensions, 'sni_hostname': url.host}) + return await self.inner.handle_async_request(pinned) + + async def aclose(self): + await self.inner.aclose() + + +def buyer_client(): + return httpx.AsyncClient(transport=PublicHTTPS(), follow_redirects=False, + trust_env=False, timeout=10) + + +def is_byo_buyer(request, env_get): + return bool(env_get('SAAS_TENANT_SCOPE') and env_get('ANTSEED_BYO_TOKEN') + and request.get('provider_id') == 'antseed') diff --git a/compose.yml b/compose.yml index 6df7e4e..8df3ed0 100644 --- a/compose.yml +++ b/compose.yml @@ -27,6 +27,12 @@ services: ANTSEED_CONTROL_TOKEN: ${ANTSEED_CONTROL_TOKEN:-} # Operational store (router reads operator config from it). Prod = RDS. DATABASE_URL: ${DATABASE_URL:-postgresql://hoststore:hoststore@postgres:5432/hoststore} + # Optional external control plane: per-tenant BYO provider credentials + # (fetched by tenant id from the trusted x-llm-router-tenant header the + # ingress stamps). Both unset -> feature off, platform keys only. + CONTROL_PLANE_URL: ${CONTROL_PLANE_URL:-} + CONTROL_PLANE_INTERNAL_SECRET: ${CONTROL_PLANE_INTERNAL_SECRET:-} + SAAS_SHARED_PROVIDERS: ${SAAS_SHARED_PROVIDERS:-} depends_on: postgres: condition: service_healthy @@ -84,6 +90,11 @@ services: DATABASE_URL: ${DATABASE_URL:-postgresql://hoststore:hoststore@postgres:5432/hoststore} RATE_PER_MIN: ${RATE_PER_MIN:-600} BURST: ${BURST:-200} + # Optional external control plane: consumer-key resolve fallback plus the + # /internal/usage metering surface. Both unset -> feature off (/internal/* + # answers 404, unknown keys just 401). + CONTROL_PLANE_URL: ${CONTROL_PLANE_URL:-} + CONTROL_PLANE_INTERNAL_SECRET: ${CONTROL_PLANE_INTERNAL_SECRET:-} LOG_LEVEL: ${LOG_LEVEL:-INFO} DASHBOARD_KEY_ENV_PATH: /run/llm-router/.env.secrets # env_file carries the HOST path; inside the container the file is the diff --git a/control_plane_client.py b/control_plane_client.py new file mode 100644 index 0000000..1e209e3 --- /dev/null +++ b/control_plane_client.py @@ -0,0 +1,327 @@ +"""Client for an external control plane ("bring your own control plane"). + +An operator can front this router with a separate control plane that owns +consumer keys and per-tenant provider credentials (any service speaking the +small HTTP contract below). The feature is OFF unless both CONTROL_PLANE_URL +and CONTROL_PLANE_INTERNAL_SECRET are set; nothing in this module assumes any +particular control-plane implementation. + +Contract (all requests carry the shared secret in `x-internal-secret`): + GET {CONTROL_PLANE_URL}/internal/keys/resolve?sha256=<64hex> + -> {"active": bool, "consumer": str, "tenant_id": int, + "rate_per_min": int|null, "burst": int|null} + GET {CONTROL_PLANE_URL}/internal/tenants//provider-env + -> {"env": {ENV_NAME: secret, ...}} + +The module is a leaf (no imports from auth_proxy/shim) shared by the ingress +(key resolution) and the router (per-tenant provider env). Secrets are never +logged — events carry env NAMES and tenant ids only. +""" +from __future__ import annotations + +import asyncio +import contextvars +import hashlib +import hmac +import json +import logging +import os +import time +from dataclasses import dataclass +from typing import Any, Mapping + +import httpx + +log = logging.getLogger("llm-router-control-plane") + +CONTROL_PLANE_URL = os.getenv("CONTROL_PLANE_URL", "").rstrip("/") +CONTROL_PLANE_INTERNAL_SECRET = os.getenv("CONTROL_PLANE_INTERNAL_SECRET", "") +RESOLVE_TTL_S = float(os.getenv("CP_RESOLVE_TTL_S", "60")) +NEGATIVE_TTL_S = float(os.getenv("CP_NEGATIVE_TTL_S", "15")) +RESOLVE_STALE_GRACE_S = float(os.getenv("CP_RESOLVE_STALE_GRACE_S", "300")) +TENANT_ENV_TTL_S = float(os.getenv("CP_TENANT_ENV_TTL_S", "120")) +TENANT_ENV_STALE_GRACE_S = float(os.getenv("CP_TENANT_ENV_STALE_GRACE_S", "600")) +ENV_ALLOWLIST = { + name.strip() + for name in os.getenv( + "CP_ENV_ALLOWLIST", + "OPENAI_API_KEY,OPENROUTER_API_KEY,ANTHROPIC_API_KEY,GEMINI_API_KEY", + ).split(",") + if name.strip() +} + +_RESOLVE_CACHE_MAX = 4096 +_TIMEOUT = httpx.Timeout(3.0, connect=1.5) + + +def enabled() -> bool: + return bool(CONTROL_PLANE_URL and CONTROL_PLANE_INTERNAL_SECRET) + + +def internal_secret_ok(headers: Mapping[str, str]) -> bool: + """Validate an inbound `x-internal-secret` header. False when the shared + secret is unconfigured — callers must treat that as 'feature hidden'.""" + if not CONTROL_PLANE_INTERNAL_SECRET: + return False + presented = headers.get("x-internal-secret") or "" + return hmac.compare_digest(presented, CONTROL_PLANE_INTERNAL_SECRET) + + +@dataclass +class ResolvedKey: + active: bool + consumer: str | None + tenant_id: int | None + rate_per_min: int | None + burst: int | None + fetched_at: float # time.monotonic() + + +_resolve_cache: dict[str, ResolvedKey] = {} # sha256 hex -> entry (positive AND negative) +_resolve_inflight: dict[str, asyncio.Future] = {} +_tenant_env_cache: dict[int, tuple[dict[str, str], float]] = {} +_TENANT_ENV: contextvars.ContextVar[dict[str, str] | None] = contextvars.ContextVar( + "cp_tenant_env", default=None +) +_client: httpx.AsyncClient | None = None +_collision_logged: set[str] = set() + + +def _get_client() -> httpx.AsyncClient: + global _client + if _client is None: + _client = httpx.AsyncClient(timeout=_TIMEOUT) + return _client + + +def sha256_hex(token: str) -> str: + return hashlib.sha256(token.encode()).hexdigest() + + +def _entry_ttl(entry: ResolvedKey) -> float: + return RESOLVE_TTL_S if entry.active else NEGATIVE_TTL_S + + +def _evict_if_full() -> None: + if len(_resolve_cache) < _RESOLVE_CACHE_MAX: + return + now = time.monotonic() + expired = [k for k, e in _resolve_cache.items() if now - e.fetched_at > _entry_ttl(e)] + for k in expired: + _resolve_cache.pop(k, None) + if len(_resolve_cache) >= _RESOLVE_CACHE_MAX: + _resolve_cache.clear() + + +def _parse_resolved(data: Any) -> ResolvedKey: + if not isinstance(data, dict): + data = {} + + def _opt_int(value: Any) -> int | None: + try: + out = int(value) + except (TypeError, ValueError): + return None + return out if out > 0 else None + + consumer = str(data.get("consumer") or "").strip() or None + active = bool(data.get("active")) and consumer is not None and _opt_int(data.get("tenant_id")) is not None + return ResolvedKey( + active=active, + consumer=consumer if active else None, + tenant_id=_opt_int(data.get("tenant_id")) if active else None, + rate_per_min=_opt_int(data.get("rate_per_min")) if active else None, + burst=_opt_int(data.get("burst")) if active else None, + fetched_at=time.monotonic(), + ) + + +async def _fetch_resolve(digest: str) -> ResolvedKey | None: + """One HTTP resolve. Returns the parsed entry (positive or negative) on a + definitive control-plane answer, None on transport error / 5xx.""" + try: + resp = await _get_client().get( + f"{CONTROL_PLANE_URL}/internal/keys/resolve", + params={"sha256": digest}, + headers={"x-internal-secret": CONTROL_PLANE_INTERNAL_SECRET}, + ) + except httpx.HTTPError as exc: + log.warning(json.dumps({"event": "cp_resolve_error", "error": type(exc).__name__})) + return None + if resp.status_code >= 500: + log.warning(json.dumps({"event": "cp_resolve_error", "status": resp.status_code})) + return None + if resp.status_code != 200: + # 403 (secret mismatch) etc. — a definitive "no": cache as negative so a + # misconfigured secret can't turn into a per-request CP hammer. + log.warning(json.dumps({"event": "cp_resolve_rejected", "status": resp.status_code})) + return _parse_resolved({}) + try: + return _parse_resolved(resp.json()) + except ValueError: + return _parse_resolved({}) + + +async def resolve_key(digest: str) -> ResolvedKey | None: + """Resolve a key digest against the control plane, with caching. + + Returns None when the feature is off or the CP is unreachable with no + usable cache (the caller should 401). A stale positive entry is served for + up to RESOLVE_STALE_GRACE_S past its TTL, but ONLY when the CP is + unreachable — a definitive answer always replaces the cache. Negative + entries never get grace. + """ + if not enabled(): + return None + now = time.monotonic() + cached = _resolve_cache.get(digest) + if cached is not None and now - cached.fetched_at <= _entry_ttl(cached): + return cached + + pending = _resolve_inflight.get(digest) + if pending is not None: + return await asyncio.shield(pending) + + future: asyncio.Future = asyncio.get_running_loop().create_future() + _resolve_inflight[digest] = future + try: + fresh = await _fetch_resolve(digest) + if fresh is not None: + _evict_if_full() + _resolve_cache[digest] = fresh + result: ResolvedKey | None = fresh + elif ( + cached is not None + and cached.active + and now - cached.fetched_at <= RESOLVE_TTL_S + RESOLVE_STALE_GRACE_S + ): + log.warning(json.dumps({"event": "cp_resolve_stale_grace", "consumer": cached.consumer})) + result = cached + else: + _resolve_cache.pop(digest, None) + result = None + future.set_result(result) + return result + except BaseException as exc: + future.set_exception(exc) + raise + finally: + _resolve_inflight.pop(digest, None) + + +async def tenant_env(tenant_id: int) -> dict[str, str]: + """Cached BYO provider env for a tenant, filtered through ENV_ALLOWLIST. + A missing map is an empty tenant credential set, never platform credentials.""" + if not enabled(): + return {} + now = time.monotonic() + cached = _tenant_env_cache.get(tenant_id) + if cached is not None and now - cached[1] <= TENANT_ENV_TTL_S: + return cached[0] + try: + resp = await _get_client().get( + f"{CONTROL_PLANE_URL}/internal/tenants/{int(tenant_id)}/provider-env", + headers={"x-internal-secret": CONTROL_PLANE_INTERNAL_SECRET}, + ) + resp.raise_for_status() + raw = resp.json().get("env") + if not isinstance(raw, dict): + raise ValueError("invalid credential map") + except (httpx.HTTPError, ValueError, AttributeError) as exc: + if cached is not None and now - cached[1] <= TENANT_ENV_TTL_S + TENANT_ENV_STALE_GRACE_S: + log.warning(json.dumps({"event": "tenant_env_stale_grace", "tenant_id": tenant_id})) + return cached[0] + log.warning(json.dumps({ + "event": "tenant_env_fallback", "tenant_id": tenant_id, "error": type(exc).__name__, + })) + return {} + env = { + str(k): str(v) + for k, v in (raw or {}).items() + if str(k) in ENV_ALLOWLIST and isinstance(v, str) and v + } + _tenant_env_cache[tenant_id] = (env, now) + return env + + +async def tenant_connections(tenant_id: int, allowed_env: set[str]) -> tuple[dict, dict]: + """Fresh closed credential scope and encrypted structured BYO connections. + The loaded catalog, not a second hardcoded list, declares credential names. + An unavailable control plane grants nothing; no stale authorization.""" + if not enabled(): + return {}, {} + try: + response = await _get_client().get( + f"{CONTROL_PLANE_URL}/internal/tenants/{int(tenant_id)}/provider-env", + headers={"x-internal-secret": CONTROL_PLANE_INTERNAL_SECRET}) + response.raise_for_status() + data = response.json() + raw, connections = data.get('env'), data.get('connections', {}) + if not isinstance(raw, dict) or not isinstance(connections, dict): + raise ValueError('invalid connection scope') + return ({k: v for k, v in raw.items() if k in allowed_env and isinstance(v, str) and v}, + {p: c for p, c in connections.items() if p in {'bedrock', 'antseed'} and isinstance(c, dict)}) + except (httpx.HTTPError, ValueError, AttributeError): + return {}, {} + + +def activate_tenant_env(env: dict[str, str] | None) -> contextvars.Token: + return _TENANT_ENV.set(env) + + +def reset_tenant_env(token: contextvars.Token) -> None: + _TENANT_ENV.reset(token) + + +def env_get(name: str) -> str | None: + """Tenant credentials are a closed set. Only operator calls use process env.""" + override = _TENANT_ENV.get() + if override is not None: + return override.get(name) + return os.environ.get(name) + + +class RouteUnavailable(RuntimeError): + pass + + +async def resolve_route(tenant_id: int, name: str) -> dict: + """Resolve the published revision on every call, so publish/pause is immediate.""" + try: + response = await _get_client().get( + f"{CONTROL_PLANE_URL}/internal/tenants/{tenant_id}/routes/{name}", + headers={"x-internal-secret": CONTROL_PLANE_INTERNAL_SECRET}, + ) + response.raise_for_status() + data = response.json() + if (not isinstance(data.get("policy_ir"), list) + or not isinstance(data.get("policy_id"), str) or len(data["policy_id"]) != 64 + or not isinstance(data.get("revision"), int) or data["revision"] <= 0 + or not isinstance(data.get("execution", {}), dict)): + raise ValueError("invalid route contract") + return data + except (httpx.HTTPError, ValueError, AttributeError) as exc: + raise RouteUnavailable("This route is unavailable or has not been published.") from exc + + +def log_collision_once(consumer: str) -> None: + if consumer in _collision_logged: + return + _collision_logged.add(consumer) + log.warning(json.dumps({"event": "cp_caller_collision", "caller": consumer})) + + +async def close() -> None: + global _client + if _client is not None: + await _client.aclose() + _client = None + + +def reset_for_tests() -> None: + global _client + _resolve_cache.clear() + _resolve_inflight.clear() + _tenant_env_cache.clear() + _collision_logged.clear() + _client = None diff --git a/host_store.py b/host_store.py index 3950b8d..fc1d99c 100644 --- a/host_store.py +++ b/host_store.py @@ -118,6 +118,7 @@ def _retention_days() -> int: # tautological; 'subscription' = $0). Lets the cost-accuracy panel flag only # rows with real signal instead of training the operator to ignore drift. "ALTER TABLE calls ADD COLUMN IF NOT EXISTS cost_basis TEXT", + "ALTER TABLE calls ADD COLUMN IF NOT EXISTS routing_summary JSONB", # Per-ATTEMPT route observations (one row per provider call the engine made, # including failed fallback tries — a grain `calls` does NOT have: `calls` is # per-REQUEST, final route only). The RAW from which reliability/latency are @@ -380,6 +381,22 @@ def _route_key(provider: "str | None", family: "str | None", # ---- calls ledger (best-effort telemetry) -------------------------------------- +def routing_summary(trace) -> dict | None: + """Bounded explanation only: never persist prompts, output, raw error bodies + or the full candidate catalog in the tenant activity ledger.""" + if not isinstance(trace, dict): + return None + summary = {key: str(trace[key])[:100] for key in ('route', 'route_revision', 'policy_id', 'routing_preference') + if trace.get(key) is not None} + summary['attempts'] = [ + {key: str(step[key])[:160] for key in ('provider_id', 'model_family', 'error_kind') + if step.get(key) is not None} + for step in (trace.get('decision_path') or [])[:32] + if isinstance(step, dict) and step.get('event') == 'attempted' + ] + summary['deadline_exceeded'] = trace.get('request_deadline_exceeded') is True + return summary + def insert_call(row: dict[str, Any]) -> None: """Record one call into the ledger from a usage-history-shaped row. Fail-soft: never raises into the request path. Best-effort telemetry.""" @@ -406,6 +423,7 @@ def insert_call(row: dict[str, Any]) -> None: # route stats from. Raw here; combining into a route key is a later step. row.get("served_by"), row.get("cost_basis"), # how cost_usd was determined (reported/computed/…) + json.dumps(routing_summary(row.get("decision_trace"))), ) with _get_pool().connection() as conn: # one transaction, auto commit/rollback conn.execute( @@ -413,8 +431,8 @@ def insert_call(row: dict[str, Any]) -> None: " caller, route_key, provider_id, model_family, served_model_id," " requested_model, status, error_type, latency_ms, tokens_in," " tokens_out, tokens_total, tokens_cached, cost_usd, served_by," - " cost_basis)" - " VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)", values) + " cost_basis, routing_summary)" + " VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s::jsonb)", values) caller = row.get("caller") cost = row.get("cost_usd") if caller and isinstance(cost, (int, float)) and float(cost) > 0: @@ -511,13 +529,15 @@ def observe_route_call_async(row: dict[str, Any]) -> None: _enqueue(lambda: _insert_route_observation(snap)) -def recent_calls(limit: int = 100) -> list[dict[str, Any]]: - """The most recent calls, newest first (operator view / verification).""" +def recent_calls(limit: int = 100, caller: "str | None" = None) -> list[dict[str, Any]]: + """The most recent calls, newest first (operator view / verification). + Optionally scoped to one caller (control-plane activity feed).""" try: + where, params = ("", []) if caller is None else (" WHERE caller = %s", [caller]) with _get_pool().connection() as conn: with conn.cursor(row_factory=dict_row) as cur: - cur.execute("SELECT * FROM calls ORDER BY id DESC LIMIT %s", - (int(limit),)) + cur.execute(f"SELECT * FROM calls{where} ORDER BY id DESC LIMIT %s", + params + [int(limit)]) return cur.fetchall() except Exception as exc: # noqa: BLE001 _log.warning("host_store recent_calls failed: %s", exc) @@ -1446,6 +1466,38 @@ def usage_aggregate(since_ts: "int | None" = None, caller: "str | None" = None, return _empty_usage_aggregate() +def usage_totals(since_ts: "int | None" = None, + caller: "str | None" = None) -> dict[str, Any]: + """One-row window totals over `calls`, including cached tokens (which the + dashboard aggregate doesn't sum) — the control-plane metering read. + Fail-soft -> zeros (same keys).""" + zeros = {"requests": 0, "errors": 0, "tokens_in": 0, "tokens_out": 0, + "tokens_cached": 0, "tokens_total": 0, "cost_usd": 0.0, "priced": 0} + try: + where, params = _usage_where(since_ts=since_ts, caller=caller) + sql = ( + "SELECT count(*)," + " count(*) FILTER (WHERE COALESCE(status,0) >= 400)," + " COALESCE(sum(COALESCE(tokens_in,0)),0)," + " COALESCE(sum(COALESCE(tokens_out,0)),0)," + " COALESCE(sum(COALESCE(tokens_cached,0)),0)," + " COALESCE(sum(CASE WHEN COALESCE(tokens_total,0) <> 0 THEN tokens_total" + " ELSE COALESCE(tokens_in,0)+COALESCE(tokens_out,0) END),0)," + " round(COALESCE(sum(GREATEST(cost_usd,0)),0)::numeric,6)::float8," + " count(cost_usd)" + f" FROM calls{where}") + with _get_pool().connection() as conn: + row = conn.execute(sql, params).fetchone() + requests, errors, tin, tout, tcached, ttotal, cost, priced = row + return {"requests": int(requests), "errors": int(errors), + "tokens_in": int(tin), "tokens_out": int(tout), + "tokens_cached": int(tcached), "tokens_total": int(ttotal), + "cost_usd": float(cost), "priced": int(priced)} + except Exception as exc: # noqa: BLE001 + _log.warning("host_store usage_totals failed: %s", exc) + return zeros + + def policy_backtest_groups(since_ts: "int | None" = None, caller: "str | None" = None, limit: int = 50) -> dict[str, Any]: diff --git a/internal_api.py b/internal_api.py new file mode 100644 index 0000000..9375a18 --- /dev/null +++ b/internal_api.py @@ -0,0 +1,98 @@ +"""Internal metering API for an external control plane. + +Served by the ingress, gated by the same shared secret the control-plane +client sends outbound (`x-internal-secret`); see control_plane_client for the +overall contract. Hidden entirely (404) while the secret is unconfigured, so +the surface does not exist on operator-only deployments. + + GET /internal/usage?caller=&since_ts=[&bucket=day] + GET /internal/usage/recent?caller=&limit= + +Rows come from the `calls` ledger, which only records LLM calls that reached +the router — ingress-level rejects (401/429) are not counted here. +""" +from __future__ import annotations + +import time +from typing import Any + +from fastapi import APIRouter, Request +from fastapi.responses import JSONResponse + +import control_plane_client +import host_store + +router = APIRouter() + +_RECENT_LIMIT_MAX = 500 + + +def _gate(request: Request) -> JSONResponse | None: + if not control_plane_client.CONTROL_PLANE_INTERNAL_SECRET: + return JSONResponse({"error": "not_found"}, status_code=404) + if not control_plane_client.internal_secret_ok(request.headers): + return JSONResponse({"error": "forbidden"}, status_code=403) + return None + + +@router.get("/internal/usage") +async def internal_usage(request: Request, caller: str = "", + since_ts: int | None = None, + bucket: str | None = None) -> JSONResponse: + denied = _gate(request) + if denied is not None: + return denied + caller = caller.strip() + if not caller: + return JSONResponse({"error": "caller_required"}, status_code=400) + totals = host_store.usage_totals(since_ts=since_ts, caller=caller) + out: dict[str, Any] = { + "caller": caller, + "window": {"since_ts": since_ts, "until_ts": int(time.time())}, + "runs": totals["requests"], + "errors": totals["errors"], + "tokens_in": totals["tokens_in"], + "tokens_out": totals["tokens_out"], + "tokens_cached": totals["tokens_cached"], + "tokens_total": totals["tokens_total"], + "cost_usd": totals["cost_usd"], + } + if bucket == "day": + by_day = host_store.usage_aggregate(since_ts=since_ts, caller=caller)["by_day"] + out["buckets"] = [ + {"date": day, "runs": counter["requests"], "cost_usd": counter["cost_usd"]} + for day, counter in sorted(by_day.items()) + ] + return JSONResponse(out) + + +@router.get("/internal/usage/recent") +async def internal_usage_recent(request: Request, caller: str = "", + limit: int = 50) -> JSONResponse: + denied = _gate(request) + if denied is not None: + return denied + caller = caller.strip() + if not caller: + return JSONResponse({"error": "caller_required"}, status_code=400) + limit = max(1, min(int(limit), _RECENT_LIMIT_MAX)) + calls = [] + for row in host_store.recent_calls(limit=limit, caller=caller): + calls.append({ + "ts": row.get("ts"), + "status": row.get("status"), + "requested_model": row.get("requested_model"), + "model_family": row.get("model_family"), + "provider": row.get("provider_id"), + "served_model_id": row.get("served_model_id"), + "latency_ms": row.get("latency_ms"), + "tokens_in": row.get("tokens_in"), + "tokens_out": row.get("tokens_out"), + "tokens_total": row.get("tokens_total"), + "tokens_cached": row.get("tokens_cached"), + "cost_usd": row.get("cost_usd"), + "error_type": row.get("error_type"), + "routing_summary": row.get("routing_summary"), + "key_sha256_prefix": (row.get("consumer_sha") or "")[:12] or None, + }) + return JSONResponse({"caller": caller, "calls": calls}) diff --git a/llm_router_host.py b/llm_router_host.py index eff0551..a2bc6f0 100644 --- a/llm_router_host.py +++ b/llm_router_host.py @@ -126,6 +126,9 @@ def __init__( logger: Logger | None = None, enforce_provider_auth: bool | None = None, ): + self._source_paths = (router_path, config_path, metrics_path) + self._tenant_id = None + self._tenant_allowed = None self.lua = LuaRuntime(unpack_returned_tuples=True) self._custom_call_hook = call_provider is not None @@ -232,19 +235,77 @@ def build_policy(self, spec: dict) -> dict: "version": self.router.ir.VERSION, } - def normalize_policy(self, policy_ir: list) -> dict: + def normalize_policy(self, policy_ir: list, *, admit: bool = False) -> dict: """Normalize a raw Σ_pol term and stamp its identity — the builder's - download/identify step when the frontend composes the IR directly - (rather than via the declarative elaborate surface). Like build_policy, - the term is canonicalized but NOT admitted; admission happens where it - is used (rank preview / execution).""" - nf = self.router.ir.term.normalize(_to_lua(self.lua, policy_ir)) + download/identify step when the frontend composes the IR directly. + Optional admission uses the engine's field schema; execution additionally + applies the host envelope. The default retains the existing normalize-only + API. SHA-256 hashes the engine's canonical encoding.""" + import hashlib + term = _to_lua(self.lua, policy_ir) + if admit: + self._flow_module() + nf = self.router.ir.compile(term, _to_lua(self.lua, {"schema": self._flow_schema})).term + else: + nf = self.router.ir.term.normalize(term) + encoded = self.router.ir.term.encode(nf) return { "policy_ir": _to_py(nf), "fingerprint": self.router.ir.term.fingerprint(nf), "version": self.router.ir.VERSION, + "policy_id": hashlib.sha256(encoded.encode()).hexdigest(), } + def for_tenant(self, tenant_id: int, env: dict, managed_providers=(), connections=None): + """A request-local engine: credential failures cannot disable other tenants. + + Reuses the same engine, catalog files, live discovery and HTTP adapters. + Only the Lua runtime and credential set are isolated. No tenant VM cache + retains secrets after the request finishes. + """ + from provider_connections import auth_env, shared_provider_ids + catalog = self.catalog() + managed = set(managed_providers).intersection(shared_provider_ids(catalog)) + scoped_env = dict(env) + scoped_env['SAAS_TENANT_SCOPE'] = str(tenant_id) + allowed = set(managed) + for pid, provider in (catalog.get('providers') or {}).items(): + key = auth_env(provider) + if key and env.get(key): + allowed.add(pid) + managed.discard(pid) # BYO credentials take precedence for this provider. + elif key and pid in managed and self._env.get(key): + scoped_env[key] = self._env[key] + from tenant_providers import configure + byo = connections or {} + byo_ids = configure(catalog, scoped_env, byo) + allowed.update(byo_ids) + managed.difference_update(byo_ids) + child = type(self)(*self._source_paths, env=scoped_env, now_ms=self._now_ms, + call_provider_async=self._async_call_hook, + call_provider=self._call_hook, discover=self._discover_hook, + enforce_provider_auth=True) + # Include runtime-added providers/models, not just the initial files. + child.config.providers = _to_lua(child.lua, catalog.get('providers') or {}) + child.config.models = _to_lua(child.lua, catalog.get('models') or {}) + child._tenant_id = tenant_id + child._tenant_allowed = allowed + child._tenant_managed = managed + child._tenant_connections = byo + child._tenant_offers = {} + child._mock_responses = self._mock_responses + child.init() + # Copy model observations only; operator credential health is unrelated. + state = child.dump_state() + live = self.dump_state() + state["ema_metrics"] = {key: value for key, value in (live.get("ema_metrics") or {}).items() + if not key.startswith('__credits|') or key.split('|', 1)[1] in managed} + for slot in ('circuit_breakers', 'disabled_providers'): + state.setdefault(slot, {}).update({pid: value for pid, value in (live.get(slot) or {}).items() + if pid in managed}) + child.restore_state(state) + return child + # ---- Σ_flow: composition over Σ_pol --------------------------------- def _flow_module(self): @@ -536,7 +597,9 @@ async def _resolve_call_async(self, request: dict, call_override=None, # observation too, the host-owned perf the algebra reads (derived) and the # market view surfaces (#15/#4a). Mocks record as well, so a mocked call # is measured exactly like a live one. - _fold_route_outcome(request, result, session=session) + if not (self._tenant_id is not None and result.get("error_kind") in + {"auth_error", "rate_limit", "payment_required"}): + _fold_route_outcome(request, result, session=session) return result def dump_state(self) -> dict: @@ -561,6 +624,8 @@ def _provider_auth_configured(self, provider_id: str, provider: dict) -> bool: OAuth providers that do not expose an env-backed token are left to their adapter because their readiness is backend-specific and refreshable. """ + if self._tenant_allowed is not None and provider_id not in self._tenant_allowed: + return False auth = provider.get("auth") if isinstance(provider.get("auth"), dict) else None kind = auth.get("kind") if auth else None env = provider.get("auth_env") or (auth.get("env") if auth else None) @@ -709,11 +774,21 @@ def _h_call_provider(self, request): return _to_lua(self.lua, resp) def _h_discover(self, discovery_id): + if discovery_id in getattr(self, '_tenant_offers', {}): + return _to_lua(self.lua, {'ok': True, 'offers': self._tenant_offers[discovery_id]}) + if self._tenant_allowed is not None: + allowed = any(pid in self._tenant_allowed and p.discovery_id == discovery_id + for pid, p in self.config.providers.items()) + if not allowed: + return _to_lua(self.lua, {"ok": False, "error": "provider_not_connected"}) if not self._discover_hook: return _to_lua(self.lua, {"ok": False, "error": "no_discover_hook"}) return _to_lua(self.lua, self._discover_hook(discovery_id)) def _h_price_multiplier(self, provider_id, source_name=None) -> float: + if self._tenant_id is not None: + # Operator subsidies/shadow prices do not describe BYO token bills. + return 1.0 import settings for name in (provider_id, source_name): diff --git a/provider_adapters/aws_credentials.py b/provider_adapters/aws_credentials.py new file mode 100644 index 0000000..53890a5 --- /dev/null +++ b/provider_adapters/aws_credentials.py @@ -0,0 +1,13 @@ +"""AWS credentials from the request scope, never the operator chain for BYO.""" + + +def client(service, region, env_get, **kwargs): + import boto3 + if env_get('SAAS_TENANT_SCOPE'): + key, secret = env_get('AWS_ACCESS_KEY_ID'), env_get('AWS_SECRET_ACCESS_KEY') + if not key or not secret: + raise ValueError('Workspace AWS credentials are missing') + # Explicit credentials disable profile, container-role and IMDS fallback. + kwargs.update(aws_access_key_id=key, aws_secret_access_key=secret, + aws_session_token=env_get('AWS_SESSION_TOKEN') or None) + return boto3.client(service, region_name=region, **kwargs) diff --git a/provider_adapters/bedrock.py b/provider_adapters/bedrock.py index c142d19..3da6187 100644 --- a/provider_adapters/bedrock.py +++ b/provider_adapters/bedrock.py @@ -44,6 +44,8 @@ def _aws_region(request: dict, env_get: Callable[[str], str | None]) -> str: + if env_get('SAAS_TENANT_SCOPE'): + return env_get('BEDROCK_REGION') or 'us-east-1' return ( request.get("region") or request.get("aws_region") @@ -235,12 +237,11 @@ def _next_stream_event(events): return None -def _bedrock_client(region: str, timeout_s: float): - import boto3 +def _bedrock_client(region: str, timeout_s: float, env_get=os.environ.get): + from provider_adapters.aws_credentials import client from botocore.config import Config - return boto3.client( - "bedrock-runtime", - region_name=region, + return client( + "bedrock-runtime", region, env_get, config=Config(read_timeout=timeout_s, connect_timeout=min(timeout_s, 10)), ) @@ -263,7 +264,7 @@ async def stream_bedrock( _env_get = env_get or os.environ.get region = _aws_region(request, _env_get) bedrock = client or (client_factory(region) if client_factory - else _bedrock_client(region, timeout_s)) + else _bedrock_client(region, timeout_s, _env_get)) body = _bedrock_request(request) t0 = time.monotonic() saw_output = False @@ -398,7 +399,7 @@ def _client(region: str): return client if client_factory is not None: return client_factory(region) - return _bedrock_client(region, timeout_s) + return _bedrock_client(region, timeout_s, _env_get) async def call(request: dict) -> dict: api_kind = request.get("api_kind") diff --git a/provider_adapters/openai_compatible.py b/provider_adapters/openai_compatible.py index 4be4a9f..70a1ec8 100644 --- a/provider_adapters/openai_compatible.py +++ b/provider_adapters/openai_compatible.py @@ -163,6 +163,11 @@ def _prepare_openai_call( auth_headers = {} headers = {"Content-Type": "application/json", **auth_headers, **extra} + from byo_http import is_byo_buyer + if is_byo_buyer(request, env_get): + # Never trust an offer's endpoint to choose where a tenant secret goes. + url = env_get('ANTSEED_BYO_URL').rstrip('/') + '/v1/chat/completions' + headers['Authorization'] = 'Bearer ' + env_get('ANTSEED_BYO_TOKEN') peer_id = offer.get("peer_id") if peer_id: headers["x-antseed-pin-peer"] = peer_id @@ -441,7 +446,11 @@ async def _ignore_delta(_delta: str) -> None: provider_rules=provider_rules, ) else: - if client is not None: + from byo_http import buyer_client, is_byo_buyer + if is_byo_buyer(request, _env_get): + async with buyer_client() as buyer: + resp = await buyer.post(url, json=body, headers=headers, timeout=timeout) + elif client is not None: resp = await client.post( url, json=body, headers=headers, timeout=timeout) else: @@ -512,10 +521,12 @@ async def stream_openai_compatible( return _peer_capacity_error( str(peer_id or ""), int(cap or 0), gate_error, t0) - _owns_client = client is None + from byo_http import buyer_client, is_byo_buyer + _byo = is_byo_buyer(request, env_get or os.environ.get) + _owns_client = client is None or _byo if _owns_client: import httpx - client = httpx.AsyncClient() + client = buyer_client() if _byo else httpx.AsyncClient() emitted = False text_parts: list[str] = [] diff --git a/provider_connections.py b/provider_connections.py new file mode 100644 index 0000000..8587f0c --- /dev/null +++ b/provider_connections.py @@ -0,0 +1,56 @@ +"""SaaS connection metadata derived from the actual dataplane catalog. + +This is not a second provider registry. Credentials and wire behavior remain +owned by catalog declarations and providers.py's existing adapters/sources. +""" +import os +import re + +LABELS = {'openai': 'OpenAI', 'openrouter': 'OpenRouter', 'anthropic': 'Anthropic', + 'gemini': 'Gemini', 'bedrock': 'Amazon Bedrock', 'antseed': 'AntSeed', + 'io_net': 'IO.net', 'openai_codex': 'ChatGPT / Codex', 'ollama': 'Ollama'} + + +def label(pid): + base = pid.removesuffix('_market') + return LABELS.get(base, base.replace('_', ' ').title()) + + +def auth_env(provider): + return provider.get('auth_env') or (provider.get('auth') or {}).get('env') + + +def shared_provider_ids(catalog): + """Explicit operator opt-in, independent of tenant-supplied credentials.""" + exposed = {s.strip() for s in os.getenv('SAAS_SHARED_PROVIDERS', '').split(',') if s.strip()} + return exposed.intersection((catalog.get('providers') or {}).keys()) + + +def credential_names(catalog): + return {name for p in (catalog.get('providers') or {}).values() + if isinstance(p, dict) and (name := auth_env(p)) + and re.fullmatch(r'[A-Z][A-Z0-9_]{0,79}', name)} + + +def connections(host): + catalog = host.catalog() + exposed = shared_provider_ids(catalog) + groups = {} + for pid, p in (catalog.get('providers') or {}).items(): + key = auth_env(p) + group = p.get('source') or pid.removesuffix('_market') + mode = ('aws' if p.get('api_kind') == 'bedrock' else + 'buyer' if str(p.get('discovery_id', '')).startswith('antseed') else + 'key' if key else 'managed') + if mode in ('aws', 'buyer'): + key = None + group_key = (group, key, mode) + item = groups.setdefault(group_key, {'id': group, 'label': label(group), + 'mode': mode, 'auth_env': key, 'providers': [], 'shared_providers': [], + 'connected': False}) + item['providers'].append(pid) + if pid in exposed: + item['shared_providers'].append(pid) + if pid in (host._tenant_allowed or set()): + item['connected'] = True + return list(groups.values()) diff --git a/providers.py b/providers.py index 9483c71..50192c7 100644 --- a/providers.py +++ b/providers.py @@ -80,34 +80,42 @@ def _bedrock_source(catalog, env_get): return BedrockSource(catalog, env_get=env_get) -def _anthropic_adapter(timeout_s): +def _anthropic_adapter(timeout_s, env_get=None): from provider_adapters.anthropic import make_anthropic_async_call_provider - return make_anthropic_async_call_provider(timeout_s=timeout_s) + if env_get is None: + return make_anthropic_async_call_provider(timeout_s=timeout_s) + return make_anthropic_async_call_provider(timeout_s=timeout_s, env_get=env_get) -def _anthropic_stream_adapter(timeout_s): +def _anthropic_stream_adapter(timeout_s, env_get=None): from provider_adapters.anthropic import stream_anthropic - return functools.partial(stream_anthropic, timeout_s=timeout_s) + if env_get is None: + return functools.partial(stream_anthropic, timeout_s=timeout_s) + return functools.partial(stream_anthropic, timeout_s=timeout_s, env_get=env_get) -def _bedrock_adapter(timeout_s): +def _bedrock_adapter(timeout_s, env_get=None): from provider_adapters.bedrock import make_bedrock_async_call_provider - return make_bedrock_async_call_provider(timeout_s=timeout_s) + return make_bedrock_async_call_provider(timeout_s=timeout_s, env_get=env_get) -def _bedrock_stream_adapter(timeout_s): +def _bedrock_stream_adapter(timeout_s, env_get=None): from provider_adapters.bedrock import stream_bedrock - return functools.partial(stream_bedrock, timeout_s=timeout_s) + return functools.partial(stream_bedrock, timeout_s=timeout_s, env_get=env_get) -def _google_adapter(timeout_s): +def _google_adapter(timeout_s, env_get=None): from provider_adapters.google import make_google_async_call_provider - return make_google_async_call_provider(timeout_s=timeout_s) + if env_get is None: + return make_google_async_call_provider(timeout_s=timeout_s) + return make_google_async_call_provider(timeout_s=timeout_s, env_get=env_get) -def _google_stream_adapter(timeout_s): +def _google_stream_adapter(timeout_s, env_get=None): from provider_adapters.google import stream_google - return functools.partial(stream_google, timeout_s=timeout_s) + if env_get is None: + return functools.partial(stream_google, timeout_s=timeout_s) + return functools.partial(stream_google, timeout_s=timeout_s, env_get=env_get) def _codex_source(catalog, env_get): @@ -349,18 +357,21 @@ def build_source_registry(catalog: dict, env_get=os.environ.get) -> list: return out -def native_adapter_handlers(timeout_s: float) -> "dict[str, Any]": +def native_adapter_handlers(timeout_s: float, env_get=None) -> "dict[str, Any]": """api_kind -> wire backend for the providers with a dedicated adapter - (codex is wired separately in serve.py because it needs `observe`).""" - return {p.api_kind: p.adapter(timeout_s) + (codex is wired separately in serve.py because it needs `observe`). + `env_get` threads a request-scoped credential lookup (per-tenant BYO keys) + into the adapters that resolve auth from env; platform-only adapters + (bedrock) ignore it.""" + return {p.api_kind: p.adapter(timeout_s, env_get) for p in PROVIDERS if not p.special and p.api_kind and p.adapter} -def native_streaming_adapter_handlers(timeout_s: float) -> "dict[str, Any]": +def native_streaming_adapter_handlers(timeout_s: float, env_get=None) -> "dict[str, Any]": """api_kind -> true streaming backend for native providers that support it (codex remains wired separately in serve.py because it needs `observe`).""" - return {p.api_kind: p.stream_adapter(timeout_s) + return {p.api_kind: p.stream_adapter(timeout_s, env_get) for p in PROVIDERS if not p.special and p.api_kind and p.stream_adapter} diff --git a/route_contract.py b/route_contract.py new file mode 100644 index 0000000..41b3907 --- /dev/null +++ b/route_contract.py @@ -0,0 +1,27 @@ +"""Select only an owner-published preference variant, never a caller policy.""" +import re + + +class PreferenceNotAllowed(ValueError): + pass + + +def apply_contract(payload, published): + payload, selected = dict(payload), dict(published) + preference = payload.pop('routing_preference', None) + selected['routing_preference'] = preference if isinstance(preference, str) else 'default' + if preference is not None: + variants = published.get('preferences') or {} + if not isinstance(preference, str) or preference not in variants: + raise PreferenceNotAllowed('This preference is not authorized by the published contract.') + variant = variants[preference] + if (not isinstance(variant, dict) or not isinstance(variant.get('policy_ir'), list) + or not re.fullmatch(r'[0-9a-f]{64}', str(variant.get('policy_id', '')))): + raise PreferenceNotAllowed('The published preference is unavailable. Republish the contract.') + selected.update(policy_ir=variant['policy_ir'], policy_id=variant['policy_id']) + for field in ('policy_ir', 'flow_ir', 'timeout_ms', 'first_token_timeout_ms', 'task_policies'): + payload.pop(field, None) + execution = published.get('execution') or {} + payload.update({key: execution[key] for key in ('timeout_ms', 'first_token_timeout_ms') if key in execution}) + payload['policy_ir'] = selected['policy_ir'] + return payload, selected diff --git a/saas_routes.py b/saas_routes.py new file mode 100644 index 0000000..4bf7bbe --- /dev/null +++ b/saas_routes.py @@ -0,0 +1,247 @@ +"""Intent-to-term adapter. All admission, ranking and execution belong to Σ_pol.""" +from __future__ import annotations + +import asyncio +import contextvars +import math +import re + +from fastapi import Request +from fastapi.responses import JSONResponse + +import control_plane_client as cp + +from provider_connections import connections, credential_names, label + +_active = contextvars.ContextVar("saas_host", default=None) +_revision = contextvars.ContextVar("saas_revision", default=None) + + +class ScopedHost: + def __init__(self, base): + self.base = base + + def __getattr__(self, name): + return getattr(_active.get() or self.base, name) + + async def execute_async(self, contract, **kwargs): + host = _active.get() or self.base + if host._tenant_id is not None and contract.get("session"): + contract = {**contract, "session": f"tenant:{host._tenant_id}:{contract['session']}"} + result = await host.execute_async(contract, **kwargs) + if revision := _revision.get(): + result.setdefault("trace", {}).update(revision) + return result + + +def choices(host): + """Read candidate identities from the engine, never a second catalog.""" + term = ["policy", ["and", ["meets_req"], ["not", ["is", "disabled"]]], + ["zero"], ["ordered"], ["id"], ["always", {"action": "abort"}]] + ranked, _ = host.rank({"policy_ir": term}) + out = {} + for row in ranked: + c = row["candidate"] + pid, family = c["provider_id"], c["model_family"] + if pid not in (host._tenant_allowed or set()): + continue + key = f"{pid}|{family}" + out[key] = {"id": key, "provider": pid, "family": family, + "label": f"{family} · {label(pid)}", + "tools": bool((c.get("capabilities") or {}).get("supports_tools")), + "price_in": finite(c.get("raw_price_in", c.get("price_in"))), + "price_out": finite(c.get("raw_price_out", c.get("price_out")))} + return sorted(out.values(), key=lambda c: (c["provider"], c["family"])) + + +def finite(value): + return value if isinstance(value, (float, int)) and math.isfinite(value) else None + + +def compile_intent(intent): + if not isinstance(intent, dict): + raise ValueError("Choose a route configuration.") + goal = intent.get("goal", "reliability") + preferences = intent.get('allowed_preferences', []) + if (not isinstance(preferences, list) or len(preferences) > 3 + or any(p not in ('cost', 'speed', 'reliability') for p in preferences)): + raise ValueError('Choose supported task preferences.') + workload = intent.get("workload", "chat") + if goal not in {"reliability", "cost", "speed"} or workload not in {"chat", "agent", "extraction"}: + raise ValueError("Choose a supported task and goal.") + targets = intent.get("targets") + if not isinstance(targets, list) or not 1 <= len(targets) <= 4: + raise ValueError("Choose a primary model and up to three alternatives.") + pairs = [] + for value in targets: + if not isinstance(value, str) or len(value) > 240 or value.count("|") != 1: + raise ValueError("Choose models from your connected providers.") + provider, family = value.split("|", 1) + if not re.fullmatch(r'[a-zA-Z0-9_-]{1,100}', provider) or not family: + raise ValueError("Choose a valid provider and model from the dataplane catalog.") + pairs.append({"provider": provider, "model": family}) + if len(set(targets)) != len(targets): + raise ValueError("Choose each model only once.") + timeout = intent.get("timeout_seconds", 8) + if isinstance(timeout, bool) or not isinstance(timeout, int) or not 1 <= timeout <= 20: + raise ValueError("Response timeout must be between 1 and 20 seconds.") + predicates = [["and", ["provider_eq", p["provider"]], ["family_eq", p["model"]]] for p in pairs] + gates = ["and", ["meets_req"], ["not", ["is", "disabled"]], ["or", *predicates]] + if workload == "agent": + gates.append(["has_cap", "supports_tools"]) + if workload == "extraction": + gates.append(["has_cap", "supports_json_mode"]) + for field in ("price_in", "price_out"): + value = intent.get(f"max_{field}") + if value is not None: + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or not 0 < value <= 1000: + raise ValueError("Price limits must be positive numbers up to $1,000 per million tokens.") + gates.append(["cmp", field, "le", value]) + scorer, selector = ["zero"], ["chain", pairs] + if goal == "cost": + # Token-price estimate for a stated 80% input / 20% output mix. + scorer = ["neg", ["normalize", ["add", + ["scale", .8, ["field", "price_in"]], + ["scale", .2, ["field", "price_out"]]]]] + gates.extend([["cmp", "price_in", "le", 1e9], ["cmp", "price_out", "le", 1e9]]) + selector = ["argmax"] + elif goal == "speed": + scorer = ["neg", ["normalize", ["field", "latency_ms"]]] + selector = ["argmax"] + # Never retry refusals, invalid requests or a stream already delivered. + failure = ["always", {"action": "abort"}] + for reason in ("timeout", "server_error", "network_error", "rate_limit", + "auth_error", "payment_required", "model_unavailable", "context_overflow"): + failure = ["override", failure, reason, {"action": "next_candidate"}] + return ["policy", gates, scorer, selector, + ["set_param", "timeout_ms", timeout * 1000], failure] + + +def preview(host, intent): + term = compile_intent(intent) + normalized = host.normalize_policy(term, admit=True) + ranked, rejected = host.rank({"policy_ir": normalized["policy_ir"]}) + rows = [] + for r in ranked: + c = r["candidate"] + rows.append({"id": f"{c['provider_id']}|{c['model_family']}", + "provider": c["provider_id"], "family": c["model_family"], + "label": f"{c['model_family']} · {label(c['provider_id'])}", + "price_in": finite(c.get("raw_price_in", c.get("price_in"))), + "price_out": finite(c.get("raw_price_out", c.get("price_out")))}) + selected = set(intent["targets"]) + reasons = {f"{r.get('provider')}|{r.get('model', r.get('model_family'))}": r.get("reason", "requirements") for r in rejected} + exclusions = [] + for target in selected - {r["id"] for r in rows}: + reason = reasons.get(target, "not_available") + if "price" in reason: + explanation = "Above your price limit, or its price is unknown." + elif "capability" in reason or "tools" in reason: + explanation = "Does not support the capabilities required by this task." + elif "disabled" in reason or reason in {"not_available", "not"}: + explanation = "Not available through your connected provider accounts." + else: + explanation = "Does not meet the mandatory requirements." + exclusions.append({"id": target, "label": target.replace("|", " · "), "reason": explanation}) + warnings = [] + if len(rows) == 1: + warnings.append("Only one model qualifies. There is no fallback if it fails.") + if intent.get("goal") == "speed": + warnings.append("Speed uses observed response latency. New models may have no measurements yet; this is not a latency guarantee.") + if intent.get("goal") == "cost": + warnings.append("Price ranking assumes 80% input and 20% output tokens. Your actual token mix and provider discounts can change the cost.") + task_policies, task_previews = {}, [] + for preference in intent.get('allowed_preferences', []): + variant = host.normalize_policy(compile_intent({**intent, 'goal': preference, + 'allowed_preferences': []}), admit=True) + task_policies[preference] = {'policy_ir': variant['policy_ir'], 'policy_id': variant['policy_id']} + candidates, _ = host.rank({'policy_ir': variant['policy_ir']}) + task_previews.append({'preference': preference, 'eligible': len(candidates), + 'first': (label(candidates[0]['candidate']['provider_id']) + ' · ' + candidates[0]['candidate']['model_family']) if candidates else None}) + return {**normalized, "ranked": rows, "excluded": exclusions, "warnings": warnings, 'task_previews': task_previews, + "execution": {"timeout_ms": intent.get("timeout_seconds", 8) * 1000, + "first_token_timeout_ms": intent.get("timeout_seconds", 8) * 1000, + "task_policies": task_policies}} + + +def install(app, scoped, handle_chat, chat_request): + @app.middleware("http") + async def tenant_context(request: Request, call_next): + raw = request.headers.get("x-llm-router-tenant") + if raw is None: + return await call_next(request) + if not cp.enabled() or not cp.internal_secret_ok(request.headers): + return JSONResponse({"error": {"message": "Untrusted tenant context"}}, status_code=403) + try: + tenant_id = int(raw) + if tenant_id <= 0: + raise ValueError() + except ValueError: + return JSONResponse({"error": {"message": "Invalid tenant"}}, status_code=400) + allowed = {"/v1/chat/completions", "/v1/responses", "/x/saas/catalog", "/x/saas/preview", "/x/saas/test", "/x/saas/connections"} + if request.url.path not in allowed: + return JSONResponse({"error": {"message": "Unsupported tenant endpoint"}}, status_code=404) + # Read current credentials for each request: revocation/rotation is immediate. + # The bounded request-local VM cannot retain stale tenant secrets/health. + env, connections_config = await cp.tenant_connections(tenant_id, credential_names(scoped.base.catalog())) + child = await asyncio.to_thread(scoped.base.for_tenant, tenant_id, env, (), connections_config) + if request.url.path != '/x/saas/connections': + from tenant_providers import prepare + await prepare(child) + env_token = cp.activate_tenant_env(child._env) + try: + host_token = _active.set(child) + rev_token = _revision.set({ + "route": request.headers.get("x-unhardcoded-route"), + "route_revision": request.headers.get("x-unhardcoded-revision"), + "policy_id": request.headers.get("x-unhardcoded-policy-id"), + 'routing_preference': request.headers.get('x-unhardcoded-preference'), + }) + try: + return await call_next(request) + finally: + _revision.reset(rev_token) + _active.reset(host_token) + finally: + cp.reset_tenant_env(env_token) + + def gate(request): + return cp.internal_secret_ok(request.headers) and _active.get() is not None + + @app.get("/x/saas/catalog") + def catalog(request: Request): + if not gate(request): + return JSONResponse({"error": "forbidden"}, status_code=403) + return {"models": choices(scoped), "connection_errors": getattr(scoped, '_connection_errors', {})} + + @app.get("/x/saas/connections") + def connection_catalog(request: Request): + if not gate(request): + return JSONResponse({"error": "forbidden"}, status_code=403) + return {"connections": connections(scoped)} + + @app.post("/x/saas/preview") + async def preview_route(request: Request): + if not gate(request): + return JSONResponse({"error": "forbidden"}, status_code=403) + try: + return await asyncio.to_thread(preview, scoped, await request.json()) + except (ValueError, TypeError, KeyError) as exc: + return JSONResponse({"error": {"message": str(exc)}}, status_code=400) + + @app.post("/x/saas/test") + async def test_route(request: Request): + if not gate(request): + return JSONResponse({"error": "forbidden"}, status_code=403) + body = await request.json() + prompt = body.get("prompt") + if not isinstance(prompt, str) or not prompt.strip() or len(prompt) > 8000: + return JSONResponse({"error": {"message": "Enter a test prompt of up to 8,000 characters."}}, status_code=400) + try: + built = await asyncio.to_thread(preview, scoped, body.get("intent")) + req = chat_request(model="", policy_ir=built["policy_ir"], max_tokens=512, + messages=[{"role": "user", "content": prompt}], + **{k: v for k, v in built['execution'].items() if k != 'task_policies'}) + return await handle_chat(req) + except (ValueError, TypeError, KeyError) as exc: + return JSONResponse({"error": {"message": str(exc)}}, status_code=400) diff --git a/serve.py b/serve.py index 409325f..7f92aee 100644 --- a/serve.py +++ b/serve.py @@ -169,10 +169,17 @@ async def codex_stream(request, emit): # The native api_kind adapters (anthropic/bedrock/google) come from the # modular provider registry; codex is wired here because its backend takes # the `observe` hook that feeds its scarcity-price source. - _native = providers.native_adapter_handlers(args.timeout_s) + # Credential lookup for env-auth adapters: control_plane_client.env_get + # consults the request-scoped tenant BYO map first (set by the shim from the + # trusted x-llm-router-tenant header), then the process env — identical to + # os.environ.get while the control-plane feature is off. + import control_plane_client + _native = providers.native_adapter_handlers(args.timeout_s, + env_get=control_plane_client.env_get) call_async = make_api_kind_dispatcher( default=make_async_call_provider(timeout_s=args.timeout_s, provider_rules=provider_rules, + env_get=control_plane_client.env_get, client=provider_http), handlers={**_native, "openai_codex": codex_call}, @@ -186,11 +193,13 @@ async def codex_stream(request, emit): make_streaming_dispatcher, stream_openai_compatible, ) - _native_streaming = providers.native_streaming_adapter_handlers(args.timeout_s) + _native_streaming = providers.native_streaming_adapter_handlers( + args.timeout_s, env_get=control_plane_client.env_get) streaming_call = make_streaming_dispatcher( default=functools.partial(stream_openai_compatible, timeout_s=args.timeout_s, provider_rules=provider_rules, + env_get=control_plane_client.env_get, client=provider_http), # Native providers and Codex have real streaming twins; Codex also feeds # `observe` for quota/scarcity pricing. diff --git a/shim.py b/shim.py index e6a34e5..a14289e 100644 --- a/shim.py +++ b/shim.py @@ -37,6 +37,7 @@ from fastapi.responses import JSONResponse, StreamingResponse from pydantic import BaseModel, ConfigDict +import control_plane_client from env_coerce import env_int import host_store from policy_templates import ( @@ -297,6 +298,8 @@ async def _execute_with_deadline(awaitable): }, } + from saas_routes import ScopedHost, install as install_saas + host = ScopedHost(host) app = FastAPI(title="llm-router shim", docs_url=None, redoc_url=None) # subscription backends (codex) are billed $0 per request — their ranking @@ -1109,9 +1112,14 @@ def _costed(body: dict) -> dict: sealed = {"role": "system", "content": _SEAL_PREFIX + summary} return _costed({"messages": frozen + [sealed] + recent, "compacted": True}) + async def _activate_tenant(request: Request) -> None: + # Tenant context is established and verified once by SaaS middleware. + pass + @app.post("/v1/chat/completions") async def chat_completions(req: ChatRequest, request: Request): _session_from_header(req, request) + await _activate_tenant(request) return await _handle_chat(req) @app.post("/{profile_name}/v1/chat/completions") @@ -1124,6 +1132,7 @@ async def chat_completions_profiled(profile_name: str, req: ChatRequest, request instead of a `profile:` model prefix. """ _session_from_header(req, request) + await _activate_tenant(request) return await _handle_chat(req, profile_name=profile_name) def _session_from_header(req: ChatRequest, request: Request) -> None: @@ -1152,12 +1161,14 @@ def _session_from_header(req: ChatRequest, request: Request) -> None: @app.post("/v1/responses") async def responses(req: ResponsesRequest, request: Request): _session_from_header(req, request) + await _activate_tenant(request) return await _handle_responses(req) @app.post("/{profile_name}/v1/responses") async def responses_profiled(profile_name: str, req: ResponsesRequest, request: Request): _session_from_header(req, request) + await _activate_tenant(request) return await _handle_responses(req, profile_name=profile_name) def _responses_object_with_router(result: dict, req: ResponsesRequest, @@ -1475,6 +1486,7 @@ async def _sse_gen(queue: "asyncio.Queue", task: "asyncio.Task", req: ChatReques yield _streaming.encode_error_event(err, err) yield _streaming.DONE_EVENT + install_saas(app, host, _handle_chat, ChatRequest) return app diff --git a/sources/antseed.py b/sources/antseed.py index f8cc2bd..dfe3a65 100644 --- a/sources/antseed.py +++ b/sources/antseed.py @@ -287,7 +287,8 @@ class AntSeedSource: name = "antseed" poll_interval_s = 300 - def __init__(self, catalog: dict): + def __init__(self, catalog: dict, store=None): + self._store = store if store is not None else host_store self._models = catalog.get("models") or {} # provider_id -> its marketplace config (cap, aliases, endpoint) self._providers: dict[str, dict] = { @@ -326,7 +327,7 @@ def _load_market(self) -> list[dict]: (max_concurrency), reputation admission and reachability ranking (last_seen/last_reached_at) are applied downstream in offers_sync / market_book.""" - rows = host_store.peer_offers(STALE_AFTER_S * 1000) + rows = self._store.peer_offers(STALE_AFTER_S * 1000) self._stats["stale"] = not rows return rows @@ -335,7 +336,7 @@ def _pinned_peer(self, provider_id: str) -> str | None: Browse mode leaves it null and the host pins per request instead (the offer carries peer_id -> x-antseed-pin-peer); when a session pin IS set, restrict offers to that peer's services to match what the proxy serves.""" - data = host_store.buyer_status(provider_id) + data = self._store.buyer_status(provider_id) return (data or {}).get("pinned_peer_id") or None def _family_vendor(self, fam: str, canon: str) -> str | None: @@ -666,7 +667,7 @@ def offers_sync(self, provider_id: str) -> list[dict]: cap_in = float(cap.get("input", float("inf"))) cap_out = float(cap.get("output", float("inf"))) # ONE buyer_status read serves both the session pin and the funds gate. - status = host_store.buyer_status(provider_id) or {} + status = self._store.buyer_status(provider_id) or {} pinned = status.get("pinned_peer_id") or None available = as_float(status.get("deposits_available")) self._stats["deposits_available"] = available @@ -712,7 +713,7 @@ def offers_sync(self, provider_id: str) -> list[dict]: half_open = 0 stale_reachability = 0 now_ms = int(time.time() * 1000) - durable_health = host_store.marketplace_route_health( + durable_health = self._store.marketplace_route_health( provider_id, window_ms=ROUTE_HEALTH_WINDOW_MS) route_health = durable_health.get("routes") or {} peer_health = durable_health.get("peers") or {} @@ -819,8 +820,8 @@ def offers_sync(self, provider_id: str) -> list[dict]: # #4a/#4c: reliability + latency + learned tool-incapability are derived on # the fly from route_observations (one query each per offers_sync, not per # candidate), keyed by route identity. - stats = host_store.route_stats() - incapable = host_store.tool_incapable_routes() + stats = self._store.route_stats() + incapable = self._store.tool_incapable_routes() offers = [] for row in kept_rows: family = row["family"] @@ -919,7 +920,7 @@ def _refresh_wallet_health(self) -> str: recent: list[bool] = [] health = None for pid in self.provider_ids: - rows = host_store.provider_recent_ok( + rows = self._store.provider_recent_ok( pid, limit=WEDGE_CONSECUTIVE_FAILURES) if len(rows) >= WEDGE_CONSECUTIVE_FAILURES and not any(rows): health = "wedged" # any wedged proxy wedges the source @@ -1059,7 +1060,7 @@ def credits_seed(self) -> dict[str, float]: and the first healthy refresh tick opens it a few minutes later.""" out: dict[str, float] = {} for pid in self.provider_ids: - status = host_store.buyer_status(pid) + status = self._store.buyer_status(pid) if not self._status_is_fresh(status): if status: _log.warning("antseed: not seeding credits for %s — its " @@ -1089,7 +1090,7 @@ async def balances(self) -> dict[str, Balance]: # offer counts it reads are current. self._refresh_wallet_health() for pid in self.provider_ids: - data = host_store.buyer_status(pid) + data = self._store.buyer_status(pid) fresh = self._status_is_fresh(data) available: "float | None" = None if fresh: diff --git a/sources/bedrock.py b/sources/bedrock.py index 08937d7..10b313d 100644 --- a/sources/bedrock.py +++ b/sources/bedrock.py @@ -150,8 +150,10 @@ def _aws_client(self, region: str): return self._bedrock_client if self._bedrock_client_factory is not None: return self._bedrock_client_factory(region) - import boto3 - return boto3.client("bedrock", region_name=region) + from provider_adapters.aws_credentials import client + from botocore.config import Config + return client("bedrock", region, self._env_get, + config=Config(connect_timeout=5, read_timeout=10, retries={'max_attempts': 1})) async def _bedrock_catalog(self, region: str) -> list[dict]: client = self._aws_client(region) diff --git a/tenant_providers.py b/tenant_providers.py new file mode 100644 index 0000000..e334518 --- /dev/null +++ b/tenant_providers.py @@ -0,0 +1,136 @@ +"""BYO configuration and request-local views over the existing provider sources. + +Never read operator buyer state or AWS entitlements for a tenant connection. +Bounded cache contains only derived discovery data, partitioned by tenant and +credential digest. Rotation/revocation changes the key immediately. +""" +import asyncio +from collections import OrderedDict +import hashlib +import json +import math +import time + + +_cache = OrderedDict() + + +def configure(catalog, env, connections): + allowed = set() + aws = connections.get('bedrock') or {} + if aws.get('access_key_id') and aws.get('secret_access_key') and aws.get('region'): + env.update(AWS_ACCESS_KEY_ID=aws['access_key_id'], AWS_SECRET_ACCESS_KEY=aws['secret_access_key'], + AWS_SESSION_TOKEN=aws.get('session_token', ''), BEDROCK_REGION=aws['region']) + for pid, p in catalog.get('providers', {}).items(): + if p.get('api_kind') == 'bedrock': + allowed.add(pid) + # Both native and market names discover THIS account's profiles. + p.update(discovery='marketplace', discovery_id=pid, aws_region=aws['region']) + buyer = connections.get('antseed') or {} + if buyer.get('gateway_url') and buyer.get('token') and 'antseed' in catalog.get('providers', {}): + allowed.add('antseed') + env.update(ANTSEED_BYO_URL=buyer['gateway_url'], ANTSEED_BYO_TOKEN=buyer['token']) + p = catalog['providers']['antseed'] + p.update(base_url=buyer['gateway_url'].rstrip('/') + '/v1', + auth={'kind': 'bearer', 'env': 'ANTSEED_BYO_TOKEN'}) + return allowed + + +class BuyerSnapshot: + """Read-only source store; no global wallet, pin or account-health fallback.""" + def __init__(self, data): + self.data = data + + def peer_offers(self, window_ms): + cutoff = time.time() * 1000 - window_ms + return [r for r in self.data.get('peer_offers', []) + if isinstance(r, dict) and isinstance(r.get('observed_at'), (int, float)) + and cutoff <= r['observed_at'] <= time.time() * 1000 + 5000] + + def buyer_status(self, pid): + return self.data.get('buyer_status') if pid == 'antseed' else None + + def marketplace_route_health(self, *args, **kwargs): + return {} + + def route_stats(self): + return {} + + def tool_incapable_routes(self): + return set() + + def provider_recent_ok(self, *args, **kwargs): + return [] + + +async def _buyer(host): + from byo_http import buyer_client + from sources.antseed import AntSeedSource, STALE_AFTER_S + async with buyer_client() as client: + async with client.stream('GET', host._env['ANTSEED_BYO_URL'].rstrip('/') + '/snapshot', + headers={'Authorization': 'Bearer ' + host._env['ANTSEED_BYO_TOKEN']}) as response: + response.raise_for_status() + body = bytearray() + async for chunk in response.aiter_bytes(): + body.extend(chunk) + if len(body) > 4 * 1024 * 1024: + raise ValueError('Buyer snapshot too large') + data = json.loads(body) + if not isinstance(data, dict) or not isinstance(data.get('peer_offers'), list): + raise ValueError('Invalid buyer snapshot') + store = BuyerSnapshot(data) + status = store.buyer_status('antseed') or {} + stamp = float(status.get('fetched_at') or 0) + credits = float(status.get('deposits_available') or 0) + if not math.isfinite(credits) or not 0 <= time.time() * 1000 - stamp <= STALE_AFTER_S * 1000: + credits = 0 + source = AntSeedSource(host.catalog(), store=store) + offers = await asyncio.to_thread(source.offers_sync, 'antseed') if credits > 0 else [] + return {'antseed': offers}, max(0, credits) + + +async def _aws(host): + from sources.bedrock import BedrockSource + source = BedrockSource(host.catalog(), env_get=host._env.get) + try: + await source.pricing() + return {pid: source.offers_sync(pid) for pid in source.provider_ids}, None + finally: + if source._client is not None: + await source._client.aclose() + + +async def prepare(host): + configs = getattr(host, '_tenant_connections', {}) + host._connection_errors = {} + + async def one(name, load, ids): + # Default to no offers before any network I/O: never fall back to the + # operator source after timeout, malformed data or revoked credentials. + host._tenant_offers.update({pid: [] for pid in ids}) + digest = hashlib.sha256(json.dumps(configs[name], sort_keys=True).encode()).hexdigest() + key = (host._tenant_id, name, digest) + cached = _cache.get(key) + try: + if cached and time.monotonic() - cached[0] < (15 if name == 'antseed' else 300): + offers, credits = cached[1] + else: + offers, credits = await asyncio.wait_for(load(host), timeout=35) + _cache[key] = (time.monotonic(), (offers, credits)) + _cache.move_to_end(key) + while len(_cache) > 128: + _cache.popitem(last=False) + host._tenant_offers.update(offers) + if credits is not None: + host.update_metrics('__credits', 'antseed', {'free_credits_remaining_usd': credits}) + except Exception: + # Never render exceptions containing a connection URL or credentials. + host._connection_errors[name] = 'Could not discover models. Check your connection, permissions and funding.' + + tasks = [] + if configs.get('bedrock') and host._env.get('AWS_ACCESS_KEY_ID'): + ids = [pid for pid, p in host.catalog()['providers'].items() if p.get('api_kind') == 'bedrock'] + tasks.append(one('bedrock', _aws, ids)) + if configs.get('antseed') and host._env.get('ANTSEED_BYO_URL'): + tasks.append(one('antseed', _buyer, ['antseed'])) + await asyncio.gather(*tasks) diff --git a/tests/fixtures/managed.lua b/tests/fixtures/managed.lua new file mode 100644 index 0000000..72392e5 --- /dev/null +++ b/tests/fixtures/managed.lua @@ -0,0 +1,13 @@ +return { + providers = { + antseed = {discovery='marketplace',discovery_id='antseed',base_url='http://buyer.test/v1',api_kind='openai_compatible',auth={kind='none'},tier='marketplace'}, + bedrock = {discovery='static',base_url='bedrock://us-east-1',api_kind='bedrock',aws_region='us-east-1',source='bedrock',tier='partner'}, + bedrock_market = {discovery='marketplace',discovery_id='bedrock_market',api_kind='bedrock',aws_region='us-east-1',source='bedrock',tier='partner'}, + custom_cloud = {discovery='static',base_url='http://custom.test/v1',api_kind='openai_compatible',auth_env='CUSTOM_CLOUD_SECRET',tier='partner'}, + }, + models = { + ['shared-model'] = {served_by={{provider='bedrock',provider_model_id='aws.profile.model'},{provider='custom_cloud'}},capabilities={context=128000,supports_tools=true,supports_json_mode=true}}, + }, + profiles={default={scorer={'zero'}}}, + policy_envelope={'and',{'meets_req'},{'not',{'is','disabled'}},{'or',{'not',{'provider_eq','antseed'}},{'cmp','credits','ge',1}}}, +} diff --git a/tests/fixtures/saas.lua b/tests/fixtures/saas.lua new file mode 100644 index 0000000..4e685fd --- /dev/null +++ b/tests/fixtures/saas.lua @@ -0,0 +1,15 @@ +return { + providers = { + openai = {discovery='static', base_url='http://provider.test', api_kind='openai_compatible', auth_env='OPENAI_API_KEY', tier='partner'}, + anthropic = {discovery='static', base_url='http://provider.test', api_kind='anthropic', auth_env='ANTHROPIC_API_KEY', tier='partner'}, + platform_only = {discovery='static', base_url='http://provider.test', api_kind='openai_compatible', tier='partner'}, + }, + models = { + primary = {served_by={{provider='openai'}}, capabilities={context=128000,supports_tools=true,supports_json_mode=true}}, + backup = {served_by={{provider='anthropic'}}, capabilities={context=128000,supports_tools=true,supports_json_mode=true}}, + basic = {served_by={{provider='openai'}}, capabilities={context=128000,supports_tools=false,supports_json_mode=false}}, + forbidden = {served_by={{provider='platform_only'}}, capabilities={context=128000}}, + }, + profiles = {default={scorer={'zero'}}}, + policy_envelope = {'and', {'meets_req'}, {'not', {'is','disabled'}}}, +} diff --git a/tests/test_antseed_node.py b/tests/test_antseed_node.py index 45aa4a1..aa5c6ee 100644 --- a/tests/test_antseed_node.py +++ b/tests/test_antseed_node.py @@ -37,6 +37,10 @@ def test_antseed_node_unit_tests(): _run_node_test("antseed/db.test.js") +def test_antseed_byo_gateway(): + _run_node_test('antseed/byo-gateway.test.js') + + def test_antseed_control_amount_cap(): """The control server's deposit-amount guard (antseed/amount.js). /deposit is now called autonomously by the router's wallet keeper, so the per-deposit diff --git a/tests/test_auth_proxy_control_plane.py b/tests/test_auth_proxy_control_plane.py new file mode 100644 index 0000000..83abe53 --- /dev/null +++ b/tests/test_auth_proxy_control_plane.py @@ -0,0 +1,401 @@ +"""Ingress <-> external control plane integration (auth fallback + /internal/*). + +The control plane's HTTP side is a fake client on control_plane_client._client; +the upstream router is a fake on auth_proxy._client. Store-backed tests use the +shared Postgres fixture (host_store_clean). +""" +from __future__ import annotations + +import hashlib +import json +import os +import sys +import time +from pathlib import Path + +import httpx +import pytest +from fastapi.testclient import TestClient + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +# auth_proxy reads the caller-key env at import; this module can be the first +# in the session to import it, so mirror the suite-wide fixture env here or the +# later-collected dashboard tests would see an empty CALLER_KEYS map. +os.environ.setdefault("CALLER_KEYS_JSON", '{"internal":"default"}') +os.environ.setdefault("CALLER_KEYS_SHA256_JSON", "{}") + +import auth_proxy # noqa: E402 +import control_plane_client as cpc # noqa: E402 +import host_store # noqa: E402 + +from conftest import require_host_store # noqa: E402 + + +class _FakeCPResp: + def __init__(self, status_code: int, payload): + self.status_code = status_code + self._payload = payload + + def json(self): + return self._payload + + +class _FakeCPClient: + def __init__(self, payloads): + self.payloads = list(payloads) + self.calls = 0 + + async def get(self, url, params=None, headers=None): + if '/routes/' in url: + return httpx.Response(200, request=httpx.Request('GET', url), json={ + 'route':'route:production', 'revision':2, 'policy_id':'a'*64, + 'policy_ir':['policy', ['top'], ['zero'], ['ordered'], ['id'], ['always', {'action':'abort'}]], + 'execution':{'timeout_ms':8000, 'first_token_timeout_ms':8000}}) + self.calls += 1 + item = self.payloads.pop(0) + if isinstance(item, Exception): + raise item + return _FakeCPResp(200, item) + + +class _FakeUpstreamResp: + status_code = 200 + headers = {"content-type": "application/json"} + + async def aread(self): + return b'{"ok": true}' + + async def aclose(self): + pass + + +class _FakeUpstream: + """Captures the headers the proxy forwards to the router.""" + + def __init__(self): + self.requests: list[dict] = [] + + def build_request(self, method, url, content=None, headers=None): + self.requests.append({"method": method, "url": url, "headers": headers or {}, "body": content}) + return object() + + async def send(self, req, stream=True): + return _FakeUpstreamResp() + + +@pytest.fixture(autouse=True) +def _clean(monkeypatch): + cpc.reset_for_tests() + monkeypatch.setattr(cpc, "CONTROL_PLANE_URL", "http://cp.test") + monkeypatch.setattr(cpc, "CONTROL_PLANE_INTERNAL_SECRET", "s3cret") + yield + cpc.reset_for_tests() + + +def _cp(monkeypatch, payloads) -> _FakeCPClient: + fake = _FakeCPClient(payloads) + monkeypatch.setattr(cpc, "_client", fake) + return fake + + +def _upstream(monkeypatch) -> _FakeUpstream: + fake = _FakeUpstream() + monkeypatch.setattr(auth_proxy, "_client", fake) + return fake + + +def _post_chat(client, token: str, extra_headers: dict | None = None): + headers = {"Authorization": f"Bearer {token}"} + headers.update(extra_headers or {}) + return client.post("/v1/chat/completions", headers=headers, + json={"model": "route:production", "messages": []}) + + +# ---- key resolution ---------------------------------------------------------- + +def test_feature_off_unknown_key_401_without_cp_call(monkeypatch): + require_host_store() + monkeypatch.setattr(cpc, "CONTROL_PLANE_URL", "") + fake_cp = _cp(monkeypatch, []) + _upstream(monkeypatch) + r = _post_chat(TestClient(auth_proxy.app), "tok-unknown") + assert r.status_code == 401 + assert fake_cp.calls == 0 + + +def test_cp_resolved_key_proxies_with_caller_and_tenant_headers(monkeypatch): + require_host_store() + _cp(monkeypatch, [{"active": True, "consumer": "acme", "tenant_id": 7, + "rate_per_min": 600, "burst": 200}]) + upstream = _upstream(monkeypatch) + r = _post_chat(TestClient(auth_proxy.app), "tok-tenant", + extra_headers={"x-llm-router-tenant": "999", # smuggle attempt + "x-internal-secret": "leak"}) + assert r.status_code == 200 + fwd = upstream.requests[0]["headers"] + assert fwd["x-llm-router-caller"] == "acme" + assert fwd["x-llm-router-tenant"] == "7" # authed value, not the smuggled 999 + assert fwd["x-internal-secret"] == "s3cret" + assert fwd["x-unhardcoded-revision"] == "2" + assert fwd["x-unhardcoded-policy-id"] == 'a'*64 + assert json.loads(upstream.requests[0]['body'])['policy_ir'][0] == 'policy' + + +def test_second_request_served_from_resolve_cache(monkeypatch): + require_host_store() + fake_cp = _cp(monkeypatch, [{"active": True, "consumer": "acme", "tenant_id": 7}]) + _upstream(monkeypatch) + client = TestClient(auth_proxy.app) + assert _post_chat(client, "tok-cache").status_code == 200 + assert _post_chat(client, "tok-cache").status_code == 200 + assert fake_cp.calls == 1 + + +def test_cp_plan_rate_limits_enforced(monkeypatch): + require_host_store() + _cp(monkeypatch, [{"active": True, "consumer": "tiny-plan", "tenant_id": 3, + "rate_per_min": 1, "burst": 1}]) + _upstream(monkeypatch) + client = TestClient(auth_proxy.app) + assert _post_chat(client, "tok-limited").status_code == 200 + r = _post_chat(client, "tok-limited") + assert r.status_code == 429 + assert r.json()["error"]["code"] == "caller_rate_limit" + + +def test_inactive_resolve_is_401(monkeypatch): + require_host_store() + _cp(monkeypatch, [{"active": False}]) + _upstream(monkeypatch) + assert _post_chat(TestClient(auth_proxy.app), "tok-revoked").status_code == 401 + + +def test_local_plaintext_key_never_consults_cp(monkeypatch): + require_host_store() + fake_cp = _cp(monkeypatch, []) + _upstream(monkeypatch) + monkeypatch.setattr(auth_proxy, "CALLER_KEYS", {"tok-local": "operator-app"}) + r = _post_chat(TestClient(auth_proxy.app), "tok-local") + assert r.status_code == 200 + assert fake_cp.calls == 0 + + +def test_locally_revoked_hash_key_never_falls_through_to_cp(monkeypatch): + require_host_store() + fake_cp = _cp(monkeypatch, []) + _upstream(monkeypatch) + digest = hashlib.sha256(b"tok-revoked-local").hexdigest() + monkeypatch.setattr(auth_proxy, "CALLER_KEY_HASHES", {digest: "operator-app"}) + host_store.set_consumer_keys({"operator-app": { + "status": "active", "keys": [{"sha256_prefix": digest[:12], "status": "revoked"}]}}) + r = _post_chat(TestClient(auth_proxy.app), "tok-revoked-local") + assert r.status_code == 403 + assert r.json()["error"]["code"] == "caller_key_revoked" + assert fake_cp.calls == 0 + + +def test_operator_kill_switch_blocks_cp_slug(monkeypatch): + require_host_store() + _cp(monkeypatch, [{"active": True, "consumer": "banned-tenant", "tenant_id": 4}]) + _upstream(monkeypatch) + host_store.set_consumer_keys({"banned-tenant": {"status": "inactive"}}) + r = _post_chat(TestClient(auth_proxy.app), "tok-banned") + assert r.status_code == 403 + assert r.json()["error"]["code"] == "caller_inactive" + + +def test_cp_caller_lands_in_the_ledger_under_the_tenant_slug(monkeypatch): + require_host_store() + _cp(monkeypatch, [{"active": True, "consumer": "acme", "tenant_id": 7}]) + _upstream(monkeypatch) + assert _post_chat(TestClient(auth_proxy.app), "tok-ledger").status_code == 200 + host_store._write_q.join() + rows = host_store.recent_calls(caller="acme") + assert len(rows) == 1 + assert rows[0]["caller"] == "acme" + + +def test_saas_policy_override_and_endpoint_escape_are_blocked(monkeypatch): + require_host_store() + _cp(monkeypatch, [{'active':True, 'consumer':'acme', 'tenant_id':7}]) + upstream = _upstream(monkeypatch) + client = TestClient(auth_proxy.app) + headers = {'Authorization':'Bearer tok-route', 'x-unhardcoded-policy-id':'forged'} + for model in ('pin:platform/forbidden', 'profile:default'): + assert client.post('/v1/chat/completions', headers=headers, json={'model':model}).status_code == 400 + assert client.get('/v1/models', headers=headers).status_code == 400 + r = client.post('/v1/chat/completions', headers=headers, json={ + 'model':'route:production', 'messages':[], 'policy_ir':['forged'], + 'flow_ir':['forged'], 'timeout_ms':999999, 'first_token_timeout_ms':999999}) + assert r.status_code == 200 + body = json.loads(upstream.requests[0]['body']) + assert body['policy_ir'][0] == 'policy' and 'flow_ir' not in body + assert body['timeout_ms'] == body['first_token_timeout_ms'] == 8000 + assert upstream.requests[0]['headers']['x-unhardcoded-policy-id'] == 'a'*64 + + +@pytest.mark.parametrize('path', ['/v1/chat/completions', '/v1/responses']) +def test_only_published_preferences_reach_upstream(monkeypatch, path): + require_host_store() + _cp(monkeypatch, [{'active':True, 'consumer':'acme', 'tenant_id':7}]) + upstream = _upstream(monkeypatch) + default = ['policy', ['top'], ['zero'], ['ordered'], ['id'], ['always', {'action':'abort'}]] + variant = ['policy', ['bottom'], ['zero'], ['ordered'], ['id'], ['always', {'action':'abort'}]] + async def resolve(*args): + return {'revision': 1, 'policy_ir': default, 'policy_id': 'a'*64, + 'execution': {'timeout_ms': 8000, 'first_token_timeout_ms': 8000}, + 'preferences': {'cost': {'policy_ir': variant, 'policy_id': 'b'*64}}} + monkeypatch.setattr(cpc, 'resolve_route', resolve) + client = TestClient(auth_proxy.app) + headers = {'Authorization':'Bearer preferences', 'x-unhardcoded-preference':'forged'} + body = {'model':'route:production', 'messages':[], 'routing_preference':'speed'} + response = client.post(path, headers=headers, json=body) + assert response.status_code == 400 and response.json()['error']['code'] == 'preference_not_allowed' + assert upstream.requests == [] + response = client.post(path, headers=headers, json={**body, 'routing_preference':'cost', + 'policy_ir':['forged'], 'timeout_ms': 999999, 'flow_ir':['forged']}) + assert response.status_code == 200 + routed = json.loads(upstream.requests[0]['body']) + assert routed['policy_ir'] == variant + assert routed['timeout_ms'] == 8000 and 'flow_ir' not in routed and 'routing_preference' not in routed + assert upstream.requests[0]['headers']['x-unhardcoded-policy-id'] == 'b'*64 + assert upstream.requests[0]['headers']['x-unhardcoded-preference'] == 'cost' + + +def test_route_resolution_failure_does_not_call_router(monkeypatch): + require_host_store() + _cp(monkeypatch, [{'active':True, 'consumer':'acme', 'tenant_id':7}]) + upstream = _upstream(monkeypatch) + async def unavailable(*args): + raise cpc.RouteUnavailable('Route unavailable') + monkeypatch.setattr(cpc, 'resolve_route', unavailable) + assert _post_chat(TestClient(auth_proxy.app), 'tok-route').status_code == 503 + assert upstream.requests == [] + + +@pytest.mark.parametrize('path', ['/v1/chat/completions', '/v1/responses']) +def test_full_ingress_to_engine_alias_and_failover(monkeypatch, path): + """Actual ingress + shim + Lua, with only CP HTTP and providers faked.""" + require_host_store() + import asyncio + from llm_router_host import LLMRouterHost + from saas_routes import compile_intent + from shim import create_app + term = compile_intent({'targets':['openai|primary', 'anthropic|backup']}) + host = LLMRouterHost(ROOT/'core/router.lua', ROOT/'tests/fixtures/saas.lua') + host.init() + seen = [] + async def call(request): + seen.append((request['provider_id'], cpc.env_get('OPENAI_API_KEY'))) + if request['provider_id'] == 'openai': + return {'ok':False, 'error_kind':'server_error'} + return {'ok':True, 'latency_ms':10, 'response':{'text':'Fallback works', 'tokens_in':2, 'tokens_out':3}} + host.set_async_call_hook(call) + def cp_http(request): + if request.url.path.endswith('/provider-env'): + data = {'env':{'OPENAI_API_KEY':'sk-tenant', 'ANTHROPIC_API_KEY':'sk-backup'}} + elif '/routes/' in request.url.path: + data = {'policy_ir':term, 'policy_id':host.normalize_policy(term)['policy_id'], + 'revision':3, 'execution':{'timeout_ms':8000}, 'route':'route:production'} + else: + data = {'active':True, 'consumer':'acme', 'tenant_id':7} + return httpx.Response(200, json=data) + cp_client = httpx.AsyncClient(transport=httpx.MockTransport(cp_http)) + shim_client = httpx.AsyncClient(transport=httpx.ASGITransport(app=create_app(host)), base_url='http://router.test') + monkeypatch.setattr(cpc, '_client', cp_client) + monkeypatch.setattr(auth_proxy, '_client', shim_client) + body = {'model':'route:production', 'messages':[{'role':'user','content':'hi'}]} if 'chat' in path else {'model':'route:production','input':'hi'} + try: + r = TestClient(auth_proxy.app).post(path, headers={'Authorization':'Bearer full-chain'}, json=body) + assert r.status_code == 200, r.text + trace = r.json()['x_router']['decision_trace'] + assert trace['route_revision'] == '3' + assert trace['route'] == 'route:production' + assert [s['provider_id'] for s in trace['decision_path'] if s['event'] == 'attempted'] == ['openai', 'anthropic'] + assert seen == [('openai','sk-tenant'), ('anthropic','sk-tenant')] + host_store._write_q.join() + row = host_store.recent_calls(caller='acme')[0] + assert row['routing_summary']['route_revision'] == '3' + assert row['routing_summary']['attempts'][0]['error_kind'] == 'server_error' + finally: + asyncio.run(cp_client.aclose()) + asyncio.run(shim_client.aclose()) + + +# ---- /internal/usage surface --------------------------------------------------- + +def _seed_calls(): + now = int(time.time()) + base = {"session": "s", "key_sha256": "c" * 64, "provider": "openrouter", + "model_family": "fam", "served_model_id": "m", "requested_model": "profile:default", + "latency_ms": 100.0} + host_store.insert_call({**base, "ts": now - 10, "caller": "acme", "status": 200, + "tokens_in": 70, "tokens_out": 30, "tokens_total": 100, + "tokens_cached": 25, "cost_usd": 0.01}) + host_store.insert_call({**base, "ts": now - 5, "caller": "acme", "status": 500, + "tokens_in": 10, "tokens_out": 0, "tokens_total": 10, + "cost_usd": 0.0}) + host_store.insert_call({**base, "ts": now - 5, "caller": "other", "status": 200, + "tokens_in": 1000, "tokens_out": 1000, "tokens_total": 2000, + "cost_usd": 9.99}) + return now + + +def test_internal_usage_hidden_without_secret(monkeypatch): + monkeypatch.setattr(cpc, "CONTROL_PLANE_INTERNAL_SECRET", "") + r = TestClient(auth_proxy.app).get("/internal/usage", params={"caller": "acme"}) + assert r.status_code == 404 + + +def test_internal_usage_wrong_secret_403(): + r = TestClient(auth_proxy.app).get("/internal/usage", params={"caller": "acme"}, + headers={"x-internal-secret": "wrong"}) + assert r.status_code == 403 + + +def test_internal_usage_totals_and_daily_buckets(): + require_host_store() + now = _seed_calls() + client = TestClient(auth_proxy.app) + r = client.get("/internal/usage", + params={"caller": "acme", "since_ts": now - 3600, "bucket": "day"}, + headers={"x-internal-secret": "s3cret"}) + assert r.status_code == 200 + data = r.json() + assert data["caller"] == "acme" + assert data["runs"] == 2 and data["errors"] == 1 + assert data["tokens_in"] == 80 and data["tokens_out"] == 30 + assert data["tokens_cached"] == 25 and data["tokens_total"] == 110 + assert data["cost_usd"] == pytest.approx(0.01) + assert data["window"]["since_ts"] == now - 3600 + assert sum(b["runs"] for b in data["buckets"]) == 2 + + # missing caller -> 400, and no cross-tenant bleed + assert client.get("/internal/usage", + headers={"x-internal-secret": "s3cret"}).status_code == 400 + + +def test_internal_usage_recent_scopes_to_caller(): + require_host_store() + _seed_calls() + r = TestClient(auth_proxy.app).get("/internal/usage/recent", + params={"caller": "acme", "limit": 10}, + headers={"x-internal-secret": "s3cret"}) + assert r.status_code == 200 + calls = r.json()["calls"] + assert len(calls) == 2 + assert {c["status"] for c in calls} == {200, 500} + assert all(c["key_sha256_prefix"] == "c" * 12 for c in calls) + assert all("consumer_sha" not in c for c in calls) # only the prefix leaves + assert calls[0]["latency_ms"] == 100.0 + + +def test_internal_usage_is_not_proxied_upstream(monkeypatch): + """Regression: the /internal router must match BEFORE the catch-all proxy.""" + upstream = _upstream(monkeypatch) + r = TestClient(auth_proxy.app).get("/internal/usage", + headers={"x-internal-secret": "s3cret"}) + assert r.status_code == 400 # caller_required, answered locally + assert upstream.requests == [] diff --git a/tests/test_control_plane_client.py b/tests/test_control_plane_client.py new file mode 100644 index 0000000..4cd1ab4 --- /dev/null +++ b/tests/test_control_plane_client.py @@ -0,0 +1,250 @@ +"""Unit tests for control_plane_client (no Postgres, no network). + +The HTTP boundary is a fake async client monkeypatched onto the module-level +`_client` (suite convention — no respx). Feature flags are monkeypatched module +attributes; every test starts from reset_for_tests(). +""" +from __future__ import annotations + +import asyncio +import os +import sys +from pathlib import Path + +import httpx +import pytest + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +import control_plane_client as cpc # noqa: E402 + + +class _FakeResp: + def __init__(self, status_code: int, payload): + self.status_code = status_code + self._payload = payload + + def json(self): + if isinstance(self._payload, Exception): + raise self._payload + return self._payload + + def raise_for_status(self): + if self.status_code >= 400: + raise httpx.HTTPStatusError("boom", request=None, response=None) + + +class _FakeClient: + """Scripted responses per URL prefix; records every call.""" + + def __init__(self, responses): + # responses: list of _FakeResp | Exception, consumed in order + self.responses = list(responses) + self.calls: list[dict] = [] + + async def get(self, url, params=None, headers=None): + self.calls.append({"url": url, "params": params, "headers": headers}) + item = self.responses.pop(0) + if isinstance(item, Exception): + raise item + return item + + +@pytest.fixture(autouse=True) +def _clean(monkeypatch): + cpc.reset_for_tests() + monkeypatch.setattr(cpc, "CONTROL_PLANE_URL", "http://cp.test") + monkeypatch.setattr(cpc, "CONTROL_PLANE_INTERNAL_SECRET", "s3cret") + yield + cpc.reset_for_tests() + + +def _install(monkeypatch, responses) -> _FakeClient: + fake = _FakeClient(responses) + monkeypatch.setattr(cpc, "_client", fake) + return fake + + +def _active(consumer="acme", tenant_id=7, rate=None, burst=None): + return {"active": True, "consumer": consumer, "tenant_id": tenant_id, + "rate_per_min": rate, "burst": burst} + + +DIGEST = "ab" * 32 + + +# ---- resolve_key ------------------------------------------------------------- + +def test_feature_off_returns_none_without_http(monkeypatch): + monkeypatch.setattr(cpc, "CONTROL_PLANE_URL", "") + fake = _install(monkeypatch, []) + out = asyncio.run(cpc.resolve_key(DIGEST)) + assert out is None + assert fake.calls == [] + + +def test_resolve_fetches_then_serves_from_cache(monkeypatch): + fake = _install(monkeypatch, [_FakeResp(200, _active(rate=60, burst=20))]) + first = asyncio.run(cpc.resolve_key(DIGEST)) + assert first.active and first.consumer == "acme" and first.tenant_id == 7 + assert first.rate_per_min == 60 and first.burst == 20 + second = asyncio.run(cpc.resolve_key(DIGEST)) + assert second is first + assert len(fake.calls) == 1 + call = fake.calls[0] + assert call["url"].endswith("/internal/keys/resolve") + assert call["params"] == {"sha256": DIGEST} + assert call["headers"] == {"x-internal-secret": "s3cret"} + + +def test_negative_answer_is_cached(monkeypatch): + fake = _install(monkeypatch, [_FakeResp(200, {"active": False})]) + first = asyncio.run(cpc.resolve_key(DIGEST)) + assert first is not None and first.active is False + second = asyncio.run(cpc.resolve_key(DIGEST)) + assert second.active is False + assert len(fake.calls) == 1 + + +def test_ttl_expiry_revalidates_and_definitive_no_replaces_positive(monkeypatch): + fake = _install(monkeypatch, [_FakeResp(200, _active()), + _FakeResp(200, {"active": False})]) + assert asyncio.run(cpc.resolve_key(DIGEST)).active is True + monkeypatch.setattr(cpc, "RESOLVE_TTL_S", 0.0) # everything positive expired + out = asyncio.run(cpc.resolve_key(DIGEST)) + assert out.active is False # revoked upstream wins + assert len(fake.calls) == 2 + + +def test_stale_grace_serves_positive_on_transport_error_only(monkeypatch): + fake = _install(monkeypatch, [_FakeResp(200, _active()), + httpx.ConnectError("down"), + httpx.ConnectError("down")]) + assert asyncio.run(cpc.resolve_key(DIGEST)).active is True + monkeypatch.setattr(cpc, "RESOLVE_TTL_S", 0.0) + # CP unreachable inside the grace window -> stale positive keeps working + assert asyncio.run(cpc.resolve_key(DIGEST)).consumer == "acme" + # grace exhausted -> None (caller 401s) + monkeypatch.setattr(cpc, "RESOLVE_STALE_GRACE_S", 0.0) + assert asyncio.run(cpc.resolve_key(DIGEST)) is None + assert len(fake.calls) == 3 + + +def test_negative_entries_get_no_grace(monkeypatch): + _install(monkeypatch, [_FakeResp(200, {"active": False}), + httpx.ConnectError("down")]) + assert asyncio.run(cpc.resolve_key(DIGEST)).active is False + monkeypatch.setattr(cpc, "NEGATIVE_TTL_S", 0.0) + assert asyncio.run(cpc.resolve_key(DIGEST)) is None + + +def test_5xx_counts_as_unreachable(monkeypatch): + _install(monkeypatch, [_FakeResp(200, _active()), _FakeResp(500, {})]) + assert asyncio.run(cpc.resolve_key(DIGEST)).active is True + monkeypatch.setattr(cpc, "RESOLVE_TTL_S", 0.0) + assert asyncio.run(cpc.resolve_key(DIGEST)).consumer == "acme" # stale grace + + +def test_single_flight_coalesces_concurrent_resolves(monkeypatch): + fake = _FakeClient([]) + started = asyncio.Event() + + async def slow_get(url, params=None, headers=None): + fake.calls.append({"url": url}) + started.set() + await asyncio.sleep(0.02) + return _FakeResp(200, _active()) + + fake.get = slow_get + monkeypatch.setattr(cpc, "_client", fake) + + async def run(): + return await asyncio.gather(cpc.resolve_key(DIGEST), cpc.resolve_key(DIGEST)) + + a, b = asyncio.run(run()) + assert a.consumer == b.consumer == "acme" + assert len(fake.calls) == 1 + + +def test_malformed_body_is_a_definitive_negative(monkeypatch): + _install(monkeypatch, [_FakeResp(200, ValueError("not json"))]) + out = asyncio.run(cpc.resolve_key(DIGEST)) + assert out is not None and out.active is False + + +# ---- tenant_env -------------------------------------------------------------- + +def test_tenant_env_filters_through_allowlist_and_caches(monkeypatch): + fake = _install(monkeypatch, [_FakeResp(200, {"env": { + "OPENAI_API_KEY": "sk-tenant", "EVIL_PATH_OVERRIDE": "x", + "ANTHROPIC_API_KEY": ""}})]) + env = asyncio.run(cpc.tenant_env(7)) + assert env == {"OPENAI_API_KEY": "sk-tenant"} # allowlist + empty dropped + assert asyncio.run(cpc.tenant_env(7)) == env + assert len(fake.calls) == 1 + assert fake.calls[0]["url"].endswith("/internal/tenants/7/provider-env") + + +def test_tenant_env_outage_returns_empty_scope(monkeypatch): + _install(monkeypatch, [httpx.ConnectError("down")]) + assert asyncio.run(cpc.tenant_env(9)) == {} + + +def test_tenant_env_stale_grace_then_empty(monkeypatch): + _install(monkeypatch, [_FakeResp(200, {"env": {"OPENAI_API_KEY": "sk-t"}}), + httpx.ConnectError("down"), + httpx.ConnectError("down")]) + assert asyncio.run(cpc.tenant_env(7)) == {"OPENAI_API_KEY": "sk-t"} + monkeypatch.setattr(cpc, "TENANT_ENV_TTL_S", 0.0) + assert asyncio.run(cpc.tenant_env(7)) == {"OPENAI_API_KEY": "sk-t"} # grace + monkeypatch.setattr(cpc, "TENANT_ENV_STALE_GRACE_S", 0.0) + assert asyncio.run(cpc.tenant_env(7)) == {} + + +# ---- env_get / context isolation --------------------------------------------- + +def test_env_get_never_falls_back_to_platform_in_tenant_scope(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-platform") + monkeypatch.setenv("OPENROUTER_API_KEY", "sk-platform-router") + assert cpc.env_get("OPENAI_API_KEY") == "sk-platform" + token = cpc.activate_tenant_env({"OPENAI_API_KEY": "sk-tenant"}) + try: + assert cpc.env_get("OPENAI_API_KEY") == "sk-tenant" + assert cpc.env_get("OPENROUTER_API_KEY") is None + finally: + cpc.reset_tenant_env(token) + assert cpc.env_get("OPENAI_API_KEY") == "sk-platform" + + +def test_env_get_isolation_across_tasks(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-platform") + seen: dict[str, str | None] = {} + + async def tenant_task(name: str, key: str | None): + if key is not None: + cpc.activate_tenant_env({"OPENAI_API_KEY": key}) + await asyncio.sleep(0.01) + + async def child(): + seen[name] = cpc.env_get("OPENAI_API_KEY") + + # tasks created AFTER activation copy the context (streaming/flow nodes) + await asyncio.create_task(child()) + + async def run(): + await asyncio.gather(tenant_task("a", "sk-a"), tenant_task("b", "sk-b"), + tenant_task("none", None)) + + asyncio.run(run()) + assert seen == {"a": "sk-a", "b": "sk-b", "none": "sk-platform"} + + +# ---- internal_secret_ok -------------------------------------------------------- + +def test_internal_secret_ok(monkeypatch): + assert cpc.internal_secret_ok({"x-internal-secret": "s3cret"}) is True + assert cpc.internal_secret_ok({"x-internal-secret": "wrong"}) is False + assert cpc.internal_secret_ok({}) is False + monkeypatch.setattr(cpc, "CONTROL_PLANE_INTERNAL_SECRET", "") + assert cpc.internal_secret_ok({"x-internal-secret": ""}) is False diff --git a/tests/test_host_store.py b/tests/test_host_store.py index c91dd62..6100e77 100644 --- a/tests/test_host_store.py +++ b/tests/test_host_store.py @@ -420,6 +420,45 @@ def test_route_stats_window_excludes_old_observations(store): assert set(store.route_stats(window_ms=30 * 60 * 1000)) == {"p|m|fresh", "p|m|stale"} +# ---- control-plane metering helpers --------------------------------------- + +def test_usage_totals_window_and_caller(store): + store.insert_call(_row(usage_event_id="a", caller="acme", ts=2_000_000, + tokens_in=70, tokens_out=30, tokens_total=100, + cost_usd=0.01)) + store.insert_call(_row(usage_event_id="b", caller="acme", ts=2_000_100, + status=500, tokens_in=10, tokens_out=0, + tokens_total=0, cost_usd=None)) + store.insert_call(_row(usage_event_id="c", caller="other", ts=2_000_100, + tokens_in=999, tokens_out=999, tokens_total=1998, + cost_usd=5.0)) + store.insert_call(_row(usage_event_id="old", caller="acme", ts=1_000)) + t = store.usage_totals(since_ts=1_500_000, caller="acme") + assert t["requests"] == 2 and t["errors"] == 1 + assert t["tokens_in"] == 80 and t["tokens_out"] == 30 + # tokens_total falls back to in+out when the stamped total is 0 + assert t["tokens_total"] == 110 + assert t["cost_usd"] == pytest.approx(0.01) + assert t["priced"] == 1 # the NULL-cost row is not counted as priced + + +def test_usage_totals_includes_cached_tokens_and_fails_soft(store, monkeypatch): + store.insert_call(_row(tokens_cached=40, caller="acme")) + assert store.usage_totals(caller="acme")["tokens_cached"] == 40 + monkeypatch.setattr(hs, "_get_pool", lambda: (_ for _ in ()).throw(RuntimeError("down"))) + zeros = store.usage_totals(caller="acme") + assert zeros["requests"] == 0 and zeros["cost_usd"] == 0.0 + assert set(zeros) == {"requests", "errors", "tokens_in", "tokens_out", + "tokens_cached", "tokens_total", "cost_usd", "priced"} + + +def test_recent_calls_caller_filter(store): + store.insert_call(_row(usage_event_id="a", caller="acme", ts=1)) + store.insert_call(_row(usage_event_id="b", caller="other", ts=2)) + store.insert_call(_row(usage_event_id="c", caller="acme", ts=3)) + rows = store.recent_calls(caller="acme") + assert [r["usage_event_id"] for r in rows] == ["c", "a"] # newest first + assert [r["usage_event_id"] for r in store.recent_calls()][0] == "c" def test_marketplace_health_tracks_route_and_peer_fault_scope(store): from conftest import seed_route_obs now = int(time.time() * 1000) diff --git a/tests/test_saas_byo.py b/tests/test_saas_byo.py new file mode 100644 index 0000000..1b9b641 --- /dev/null +++ b/tests/test_saas_byo.py @@ -0,0 +1,191 @@ +import asyncio +from pathlib import Path +import socket +import time +from unittest.mock import Mock + +import httpx +import pytest + +from llm_router_host import LLMRouterHost +from saas_routes import choices +import tenant_providers as tp + +ROOT = Path(__file__).resolve().parents[1] + + +@pytest.fixture +def base(host_store_clean): + tp._cache.clear() + host = LLMRouterHost(ROOT/'core/router.lua', ROOT/'tests/fixtures/managed.lua', + discover=lambda _: {'ok': True, 'offers': [{'model_family': 'OPERATOR-ONLY'}]}) + host.init() + host.update_metrics('__credits', 'antseed', {'free_credits_remaining_usd': 999}) + return host + + +def aws(key, region='eu-west-1'): + return {'bedrock': {'access_key_id': key, 'secret_access_key': 'secret-' + key, + 'session_token': 'session-' + key, 'region': region}} + + +def buyer(token, host='buyer.example.com'): + return {'antseed': {'gateway_url': 'https://' + host, 'token': token}} + + +def test_concurrent_bedrock_invocations_use_own_credentials_region_and_session(base, monkeypatch): + import boto3 + import control_plane_client as cp + from providers import _bedrock_adapter, _bedrock_stream_adapter + seen = [] + def make(service, **kwargs): + seen.append((service, kwargs)) + response = {'output': {'message': {'content': [{'text': kwargs['aws_access_key_id']}]}}, + 'usage': {'inputTokens': 1, 'outputTokens': 1}} + return Mock(converse=lambda **_: response, converse_stream=lambda **_: {'stream': [ + {'contentBlockDelta': {'delta': {'text': kwargs['aws_access_key_id']}}}]}) + monkeypatch.setattr(boto3, 'client', make) + monkeypatch.setenv('AWS_ACCESS_KEY_ID', 'OPERATOR') + call = _bedrock_adapter(10, cp.env_get) + stream = _bedrock_stream_adapter(10, cp.env_get) + async def run(tenant, key, region): + child = base.for_tenant(tenant, {}, connections=aws(key, region)) + token = cp.activate_tenant_env(child._env) + try: + request = {'api_kind': 'bedrock', 'served_model_id': 'my-model', + 'aws_region': 'us-east-1', 'messages': [{'role':'user','content':'hello'}]} + result = await call(request) + deltas = [] + async def emit(s): deltas.append(s) + assert (await stream(request, emit))['ok'] + assert deltas == [key] + return result['response']['text'] + finally: + cp.reset_tenant_env(token) + async def both(): + return await asyncio.gather(run(1, 'A', 'eu-west-1'), run(2, 'B', 'us-west-2')) + assert asyncio.run(both()) == ['A', 'B'] + assert len(seen) == 4 + for _, config in seen: + key = config['aws_access_key_id'] + assert config['aws_secret_access_key'] == 'secret-' + key + assert config['aws_session_token'] == 'session-' + key + assert config['region_name'] == {'A':'eu-west-1','B':'us-west-2'}[key] + + +def test_missing_aws_credentials_never_uses_operator_chain(monkeypatch): + import boto3 + from provider_adapters.aws_credentials import client + native = Mock() + monkeypatch.setattr(boto3, 'client', native) + with pytest.raises(ValueError): + client('bedrock', 'us-east-1', {'SAAS_TENANT_SCOPE': '1'}.get) + native.assert_not_called() + + +def test_bedrock_discovery_is_account_scoped_and_failure_never_uses_operator(base, monkeypatch): + from sources.bedrock import BedrockSource + seen = [] + async def pricing(self): + key = self._env_get('AWS_ACCESS_KEY_ID') + seen.append(key) + if key == 'revoked': + raise RuntimeError('secret must not leak') + self._offers_by_provider = {pid: [{'model_family':'shared-model','wire_model_id': key + '.profile', + 'price_in_usd_per_mtok':1, 'price_out_usd_per_mtok':2}] for pid in self.provider_ids} + return [] + monkeypatch.setattr(BedrockSource, 'pricing', pricing) + async def run(): + for tid, key in [(1,'A'), (2,'B'), (1,'revoked')]: + child = base.for_tenant(tid, {}, connections=aws(key)) + await tp.prepare(child) + if key == 'revoked': + assert choices(child) == [] + assert 'secret' not in str(child._connection_errors) + else: + assert child._tenant_offers['bedrock'][0]['wire_model_id'] == key + '.profile' + assert {r['provider'] for r in choices(child)} == {'bedrock','bedrock_market'} + asyncio.run(run()) + assert seen == ['A','B','revoked'] + assert choices(base.for_tenant(1, {})) == [] + + +def test_antseed_buyer_offers_funding_and_wire_are_tenant_scoped(base, monkeypatch): + import byo_http + from provider_adapters.openai_compatible import make_async_call_provider, stream_openai_compatible + calls = [] + def network(request): + token = request.headers['authorization'].removeprefix('Bearer ') + calls.append((request.url.host, token, request.url.path, request.headers.get('x-antseed-pin-peer'))) + if request.url.path == '/snapshot': + return httpx.Response(200, json={'buyer_status': {'fetched_at': 0 if token == 'stale' else int(time.time()*1000), + 'deposits_available': 0 if token == 'empty' else 5}, 'peer_offers': [{ + 'observed_at': int(time.time()*1000), 'peer_id': 'peer-' + token, + 'service': 'shared-model', 'price_in': 1, 'price_out': 2}]}) + import json + if json.loads(request.content).get('stream'): + chunk = json.dumps({'choices': [{'delta': {'content': token}}]}) + return httpx.Response(200, headers={'content-type':'text/event-stream'}, + content=f'data: {chunk}\n\ndata: [DONE]\n\n') + return httpx.Response(200, json={'choices': [{'message': {'content': token}}]}) + monkeypatch.setattr(byo_http, 'buyer_client', lambda: httpx.AsyncClient(transport=httpx.MockTransport(network))) + async def run(): + for tenant, token in [(1, 'A'), (2, 'B'), (3, 'empty'), (4, 'stale')]: + child = base.for_tenant(tenant, {}, connections=buyer(token, token.lower() + '.example.com')) + await tp.prepare(child) + if token in ('empty', 'stale'): + assert choices(child) == [] # Operator's 999 credits must not help. + continue + assert {r['provider'] for r in choices(child)} == {'antseed'} + offer = child._tenant_offers['antseed'][0] + assert offer['peer_id'] == 'peer-' + token + # Even a forged offer cannot exfiltrate a buyer's connection token. + offer['seller_endpoint'] = 'https://evil.example.com/v1' + result = await make_async_call_provider(env_get=child._env.get)({ + 'provider_id':'antseed','served_model_id':'shared-model','offer':offer, + 'base_url': 'https://evil.example.com/v1','messages':[{'role':'user','content':'hello'}]}) + assert result['ok'] and result['response']['text'] == token + deltas = [] + async def emit(s): deltas.append(s) + result = await stream_openai_compatible({'provider_id':'antseed', + 'served_model_id':'shared-model','offer':offer,'base_url':'https://evil.example.com/v1', + 'messages':[{'role':'user','content':'hello'}]}, emit, env_get=child._env.get) + assert result['ok'] and deltas == [token] + asyncio.run(run()) + assert [(host, token, peer) for host, token, path, peer in calls if path.endswith('completions')] == [ + ('a.example.com','A','peer-A'), ('a.example.com','A','peer-A'), + ('b.example.com','B','peer-B'), ('b.example.com','B','peer-B')] + + +@pytest.mark.parametrize('ip', ['127.0.0.1', '169.254.169.254', '10.0.0.1', '::1', '::ffff:127.0.0.1']) +def test_buyer_egress_rejects_private_dns_even_after_save(ip, monkeypatch): + from byo_http import PublicHTTPS + async def run(): + async def resolve(*a, **kw): + return [(socket.AF_INET, socket.SOCK_STREAM, 6, '', (ip, 443))] + monkeypatch.setattr(asyncio.get_running_loop(), 'getaddrinfo', resolve) + async with httpx.AsyncClient(transport=PublicHTTPS()) as client: + with pytest.raises(httpx.ConnectError): + await client.get('https://buyer.example.com/snapshot') + asyncio.run(run()) + + +def test_public_dns_is_pinned_and_tls_hostname_preserved(monkeypatch): + from byo_http import PublicHTTPS + seen = [] + async def run(): + async def resolve(*a, **kw): + return [(socket.AF_INET, socket.SOCK_STREAM, 6, '', ('8.8.8.8', 443))] + monkeypatch.setattr(asyncio.get_running_loop(), 'getaddrinfo', resolve) + transport = PublicHTTPS() + async def wire(request): + seen.append(request) + return httpx.Response(200, json={}) + await transport.inner.aclose() + transport.inner = httpx.MockTransport(wire) + async with httpx.AsyncClient(transport=transport) as client: + await client.get('https://buyer.example.com/snapshot') + asyncio.run(run()) + assert seen[0].url.host == '8.8.8.8' + assert seen[0].headers['host'] == 'buyer.example.com' + assert seen[0].extensions['sni_hostname'] == 'buyer.example.com' diff --git a/tests/test_saas_managed.py b/tests/test_saas_managed.py new file mode 100644 index 0000000..69b5e29 --- /dev/null +++ b/tests/test_saas_managed.py @@ -0,0 +1,89 @@ +"""Managed native and marketplace providers use the existing engine/adapters.""" +import asyncio +from pathlib import Path + +import httpx +import pytest + +from llm_router_host import LLMRouterHost +from provider_connections import connections, credential_names +from saas_routes import choices, compile_intent, preview + +ROOT = Path(__file__).resolve().parents[1] + + +@pytest.fixture +def managed_host(monkeypatch): + monkeypatch.setenv('SAAS_SHARED_PROVIDERS', 'antseed,bedrock,bedrock_market') + def discover(pid): + if pid == 'antseed': + return {'ok':True,'offers':[{'model_family':'shared-model','wire_model_id':'peer.model', + 'peer_id':'peer-one','seller_endpoint':'http://buyer.test/v1', + 'price_in_usd_per_mtok':1,'price_out_usd_per_mtok':2}]} + return {'ok':True,'offers':[]} + base = LLMRouterHost(ROOT/'core/router.lua', ROOT/'tests/fixtures/managed.lua', + discover=discover, now_ms=lambda:1000) + base.init() + base.update_metrics('__credits', 'antseed', {'free_credits_remaining_usd':5}) + return base + + +def test_managed_catalog_and_actual_engine_admission(managed_host): + child = managed_host.for_tenant(7, {}, ['antseed','bedrock','bedrock_market']) + assert {r['provider'] for r in choices(child)} == {'antseed','bedrock'} + info = {c['id']:c for c in connections(child)} + assert info['antseed']['mode'] == 'buyer' and info['antseed']['connected'] + assert info['bedrock']['mode'] == 'aws' + assert info['bedrock']['shared_providers'] == ['bedrock', 'bedrock_market'] or set(info['bedrock']['shared_providers']) == {'bedrock','bedrock_market'} + out = preview(child, {'targets':['antseed|shared-model','bedrock|shared-model']}) + assert [r['provider'] for r in out['ranked']] == ['antseed','bedrock'] + + +def test_no_implicit_operator_access_and_grants_must_be_exposed(managed_host, monkeypatch): + assert choices(managed_host.for_tenant(8, {})) == [] + monkeypatch.setenv('SAAS_SHARED_PROVIDERS', '') + assert choices(managed_host.for_tenant(7, {}, ['antseed','bedrock'])) == [] + + +def test_antseed_funding_gate_is_preserved(managed_host): + managed_host.update_metrics('__credits', 'antseed', {'free_credits_remaining_usd':0}) + # Core EMA updates can smooth observations; force the snapshot to zero. + state = managed_host.dump_state() + state['ema_metrics']['__credits|antseed']['free_credits_remaining_usd'] = 0 + managed_host.restore_state(state) + child = managed_host.for_tenant(7, {}, ['antseed','bedrock']) + assert {r['provider'] for r in choices(child)} == {'bedrock'} + + +def test_catalog_declares_new_byok_provider_without_saas_whitelist(managed_host): + assert 'CUSTOM_CLOUD_SECRET' in credential_names(managed_host.catalog()) + child = managed_host.for_tenant(7, {'CUSTOM_CLOUD_SECRET':'tenant-secret'}) + assert {r['provider'] for r in choices(child)} == {'custom_cloud'} + assert preview(child, {'targets':['custom_cloud|shared-model']})['ranked'] + + +def test_antseed_to_native_bedrock_fallback_uses_real_adapters(managed_host): + from provider_adapters.openai_compatible import make_async_call_provider + from provider_adapters.bedrock import make_bedrock_async_call_provider + from provider_adapters.dispatcher import make_api_kind_dispatcher + seen = [] + def buyer(request): + seen.append(('buyer', request.headers.get('x-antseed-pin-peer'))) + return httpx.Response(503, json={'error':{'message':'unavailable'}}) + class AWS: + def converse(self, **request): + seen.append(('aws', request['modelId'])) + return {'output':{'message':{'content':[{'text':'Bedrock fallback'}]}}, + 'usage':{'inputTokens':2,'outputTokens':2},'stopReason':'end_turn'} + async def run(): + async with httpx.AsyncClient(transport=httpx.MockTransport(buyer)) as http: + managed_host.set_async_call_hook(make_api_kind_dispatcher( + make_async_call_provider(client=http), + {'bedrock':make_bedrock_async_call_provider(client=AWS())})) + child = managed_host.for_tenant(7, {}, ['antseed','bedrock']) + return await child.execute_async({'policy_ir':compile_intent({'targets':['antseed|shared-model','bedrock|shared-model']}), + 'messages':[{'role':'user','content':'hello'}]}) + result = asyncio.run(run()) + assert result['ok'], result + assert result['response']['text'] == 'Bedrock fallback' + assert seen == [('buyer','peer-one'),('aws','aws.profile.model')] diff --git a/tests/test_saas_routes.py b/tests/test_saas_routes.py new file mode 100644 index 0000000..65dab7f --- /dev/null +++ b/tests/test_saas_routes.py @@ -0,0 +1,129 @@ +"""Guided intent goes through the real Lua engine; no upstream network calls.""" +import asyncio +from pathlib import Path + +import pytest + +from llm_router_host import LLMRouterHost +from saas_routes import choices, compile_intent, preview + +ROOT = Path(__file__).resolve().parents[1] + + +@pytest.fixture +def tenant_host(): + base = LLMRouterHost(ROOT / 'core/router.lua', ROOT / 'tests/fixtures/saas.lua', + now_ms=lambda: 1000, enforce_provider_auth=False) + base.init() + host = base.for_tenant(7, {'OPENAI_API_KEY': 'test', 'ANTHROPIC_API_KEY': 'test'}) + for provider, model, price in [('openai', 'primary', 8), ('anthropic', 'backup', 2), ('openai', 'basic', 1)]: + host.update_metrics(provider, model, {'ok': True, 'latency_ms': price * 100, + 'price_in': price, 'price_out': price * 2}) + return host + + +def intent(**kwargs): + return {'targets': ['openai|primary', 'anthropic|backup'], 'goal': 'reliability', + 'workload': 'chat', 'timeout_seconds': 8, **kwargs} + + +def test_real_engine_preview_and_identity(tenant_host): + out = preview(tenant_host, intent()) + assert [r['family'] for r in out['ranked']] == ['primary', 'backup'] + assert len(out['policy_id']) == 64 + assert preview(tenant_host, intent())['policy_id'] == out['policy_id'] + assert all(r['provider'] != 'platform_only' for r in choices(tenant_host)) + + +def test_cost_and_hard_price_limits(tenant_host): + assert preview(tenant_host, intent(goal='cost'))['ranked'][0]['family'] == 'backup' + out = preview(tenant_host, intent(max_price_in=3)) + assert [r['family'] for r in out['ranked']] == ['backup'] + assert out['excluded'] and out['warnings'] + assert preview(tenant_host, intent(max_price_out=.001))['ranked'] == [] + + +def test_capability_requirements(tenant_host): + out = preview(tenant_host, intent(targets=['openai|basic'], workload='agent')) + assert out['ranked'] == [] + assert out['excluded'] + + +def test_every_authorized_preference_preserves_hard_limits(tenant_host): + from route_contract import apply_contract + built = preview(tenant_host, intent(max_price_in=3, workload='agent', + allowed_preferences=['cost', 'speed', 'reliability'])) + published = {**built, 'preferences': built['execution']['task_policies']} + for preference in ('cost', 'speed', 'reliability'): + payload, selected = apply_contract({'routing_preference': preference, 'policy_ir': ['forged'], + 'flow_ir': ['forged'], 'timeout_ms': 999999}, published) + ranked, _ = tenant_host.rank(payload) + assert [r['candidate']['model_family'] for r in ranked] == ['backup'] + assert payload['timeout_ms'] == 8000 and 'flow_ir' not in payload + assert payload['policy_ir'] == built['execution']['task_policies'][preference]['policy_ir'] + assert selected['policy_id'] == built['execution']['task_policies'][preference]['policy_id'] + + +def test_preferences_cannot_authorize_a_call_when_nothing_qualifies(tenant_host): + from route_contract import apply_contract + built = preview(tenant_host, intent(max_price_out=.001, allowed_preferences=['cost', 'speed'])) + called = [] + async def provider(request): + called.append(request) + raise AssertionError('No provider should be called') + tenant_host.set_async_call_hook(provider) + for preference in ('cost', 'speed'): + payload, _ = apply_contract({'routing_preference': preference, 'messages': []}, + {**built, 'preferences': built['execution']['task_policies']}) + result = asyncio.run(tenant_host.execute_async(payload)) + assert not result['ok'] + assert called == [] + + +@pytest.mark.parametrize('preference', ['ignore_limits', 'speed', [], {}, 1]) +def test_unpublished_preferences_are_rejected(preference): + from route_contract import apply_contract, PreferenceNotAllowed + with pytest.raises(PreferenceNotAllowed): + apply_contract({'routing_preference': preference}, {'policy_ir': ['policy'], 'preferences': {}}) + + +def test_real_failover_stays_within_approved_set(tenant_host): + calls = [] + async def provider(request): + calls.append(request['provider_id']) + if request['provider_id'] == 'openai': + return {'ok': False, 'error': 'unavailable', 'error_kind': 'server_error'} + return {'ok': True, 'latency_ms': 10, 'response': {'text': 'backup response', 'tokens_in': 2, 'tokens_out': 2}} + tenant_host.set_async_call_hook(provider) + result = asyncio.run(tenant_host.execute_async({'policy_ir': compile_intent(intent()), + 'messages': [{'role':'user','content':'hello'}]})) + assert result['ok'], result + assert calls == ['openai', 'anthropic'] + + +@pytest.mark.parametrize('change', [{'targets': []}, {'targets': ['invalid/provider|forbidden']}, + {'timeout_seconds': 0}, {'max_price_in': float('nan')}, + {'targets': ['openai|primary'] * 2}]) +def test_invalid_intents_refused(change): + with pytest.raises(ValueError): + compile_intent(intent(**change)) + + +def test_activity_summary_never_includes_prompt_output_or_upstream_error(): + from host_store import routing_summary + result = routing_summary({'route_revision':3, 'messages':['secret prompt'], + 'response':'secret output', 'decision_path':[{'event':'attempted', + 'provider_id':'openai', 'error_kind':'auth_error', 'error_message':'secret key'}]}) + assert result['attempts'] == [{'provider_id':'openai','error_kind':'auth_error'}] + assert 'secret' not in str(result) + + +def test_live_catalog_accepts_guided_policy_without_platform_credentials(): + base = LLMRouterHost(ROOT/'core/router.lua', ROOT/'config.live.lua', ROOT/'metrics.live.lua') + base.init() + child = base.for_tenant(7, {'OPENAI_API_KEY':'sk-test', 'ANTHROPIC_API_KEY':'sk-test'}) + models = choices(child) + assert models + assert {m['provider'] for m in models} <= {'openai','anthropic'} + result = preview(child, intent(targets=[models[0]['id']])) + assert result['ranked'][0]['id'] == models[0]['id'] diff --git a/tests/test_tenant_env_shim.py b/tests/test_tenant_env_shim.py new file mode 100644 index 0000000..e7fc5c3 --- /dev/null +++ b/tests/test_tenant_env_shim.py @@ -0,0 +1,255 @@ +"""Per-tenant BYO provider credentials through the router (shim + adapters). + +Covers the full chain: the trusted x-llm-router-tenant header -> _activate_tenant +-> ContextVar -> control_plane_client.env_get inside the provider call, including +task-context propagation (streaming/timeout paths create tasks) and isolation +between concurrent tenants. The control plane is faked at the HTTP boundary. +""" +from __future__ import annotations + +import asyncio +import sys +from pathlib import Path + +import pytest +from fastapi.testclient import TestClient + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +import control_plane_client as cpc # noqa: E402 +from llm_router_host import LLMRouterHost # noqa: E402 +from shim import create_app # noqa: E402 + + +class _FakeCPResp: + status_code = 200 + + def __init__(self, payload): + self._payload = payload + + def json(self): + return self._payload + + def raise_for_status(self): + pass + + +class _FakeCPClient: + def __init__(self, env_by_tenant): + self.env_by_tenant = env_by_tenant + self.calls: list[str] = [] + + async def get(self, url, params=None, headers=None): + self.calls.append(url) + tenant_id = int(url.rstrip("/").split("/")[-2]) + return _FakeCPResp({"env": self.env_by_tenant.get(tenant_id, {})}) + + +@pytest.fixture(autouse=True) +def _clean(monkeypatch): + cpc.reset_for_tests() + monkeypatch.setattr(cpc, "CONTROL_PLANE_URL", "http://cp.test") + monkeypatch.setattr(cpc, "CONTROL_PLANE_INTERNAL_SECRET", "s3cret") + yield + cpc.reset_for_tests() + + +@pytest.fixture +def host(): + h = LLMRouterHost( + router_path=ROOT / "core" / "router.lua", + config_path=ROOT / "tests" / "fixtures" / "saas.lua", + now_ms=lambda: 1_000_000, + # The custom call hook below would otherwise turn on auth enforcement + # and pre-disable every example provider (their env keys are unset here). + enforce_provider_auth=False, + ) + h.init() + return h + + +def _ok_result(): + return {"ok": True, "latency_ms": 10, + "response": {"text": "hi", "tool_calls": None, "finish_reason": "stop", + "tokens_in": 7, "tokens_out": 3, "tokens_total": 10, + "raw_model": "mock-model-id"}} + + +def _capture_hook(seen: list): + """An async provider hook that records what env_get resolves AT CALL TIME — + i.e. inside host.execute_async, past any create_task boundaries.""" + + async def hook(request: dict) -> dict: + await asyncio.sleep(0.005) # let concurrent requests interleave + seen.append(cpc.env_get("OPENAI_API_KEY")) + return _ok_result() + + return hook + + +def _chat(client, tenant: int | None): + headers = {"x-llm-router-tenant": str(tenant), "x-internal-secret": "s3cret"} if tenant is not None else {} + return client.post("/v1/chat/completions", headers=headers, + json={"model": "profile:default", + "messages": [{"role": "user", "content": "hi"}]}) + + +def test_tenant_header_activates_byo_key(host, monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-platform") + monkeypatch.setattr(cpc, "_client", _FakeCPClient({7: {"OPENAI_API_KEY": "sk-tenant"}})) + seen: list = [] + host.set_async_call_hook(_capture_hook(seen)) + client = TestClient(create_app(host, default_profile="default")) + assert _chat(client, tenant=7).status_code == 200 + assert seen == ["sk-tenant"] + + +def test_no_header_uses_platform_key(host, monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-platform") + fake = _FakeCPClient({}) + monkeypatch.setattr(cpc, "_client", fake) + seen: list = [] + host.set_async_call_hook(_capture_hook(seen)) + client = TestClient(create_app(host, default_profile="default")) + assert _chat(client, tenant=None).status_code == 200 + assert seen == ["sk-platform"] + assert fake.calls == [] + + +def test_missing_tenant_credentials_never_use_platform_or_other_providers(host, monkeypatch): + monkeypatch.setenv('OPENAI_API_KEY', 'sk-platform') + monkeypatch.setattr(cpc, '_client', _FakeCPClient({})) + seen = [] + host.set_async_call_hook(_capture_hook(seen)) + client = TestClient(create_app(host, default_profile='default')) + assert _chat(client, tenant=7).status_code != 200 + assert seen == [] + assert cpc.env_get('OPENAI_API_KEY') == 'sk-platform' + + +def test_untrusted_context_and_internal_catalog_are_rejected(host): + client = TestClient(create_app(host)) + assert client.get('/x/saas/catalog').status_code == 403 + assert client.post('/v1/chat/completions', headers={'x-llm-router-tenant':'7'}, json={}).status_code == 403 + assert client.get('/x/runtime', headers={'x-llm-router-tenant':'7', 'x-internal-secret':'s3cret'}).status_code == 404 + + +def test_tenant_preview_uses_engine_without_provider_call(host, monkeypatch): + monkeypatch.setattr(cpc, '_client', _FakeCPClient({7:{'OPENAI_API_KEY':'sk-tenant'}})) + seen = [] + host.set_async_call_hook(_capture_hook(seen)) + client = TestClient(create_app(host)) + headers = {'x-llm-router-tenant':'7', 'x-internal-secret':'s3cret'} + catalog = client.get('/x/saas/catalog', headers=headers).json()['models'] + assert {m['provider'] for m in catalog} == {'openai'} + preview = client.post('/x/saas/preview', headers=headers, json={'targets':['openai|primary']}) + assert preview.status_code == 200, preview.text + assert preview.json()['ranked'][0]['family'] == 'primary' + assert len(preview.json()['policy_id']) == 64 + assert seen == [] + + +def test_feature_off_rejects_tenant_header(host, monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-platform") + monkeypatch.setattr(cpc, "CONTROL_PLANE_URL", "") + fake = _FakeCPClient({7: {"OPENAI_API_KEY": "sk-tenant"}}) + monkeypatch.setattr(cpc, "_client", fake) + seen: list = [] + host.set_async_call_hook(_capture_hook(seen)) + client = TestClient(create_app(host, default_profile="default")) + assert _chat(client, tenant=7).status_code == 403 + assert seen == [] + assert fake.calls == [] + + +def test_concurrent_tenants_are_isolated(host, monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "sk-platform") + monkeypatch.setattr(cpc, "_client", _FakeCPClient({ + 1: {"OPENAI_API_KEY": "sk-one"}, 2: {"OPENAI_API_KEY": "sk-two"}})) + seen: list = [] + host.set_async_call_hook(_capture_hook(seen)) + app = create_app(host, default_profile="default") + + # Drive the ASGI app directly so both requests share one event loop and + # genuinely interleave inside the provider hook. + import httpx + + async def run(): + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://t") as c: + body = {"model": "profile:default", + "messages": [{"role": "user", "content": "hi"}]} + r1, r2, r3 = await asyncio.gather( + c.post("/v1/chat/completions", json=body, + headers={"x-llm-router-tenant": "1", "x-internal-secret": "s3cret"}), + c.post("/v1/chat/completions", json=body, + headers={"x-llm-router-tenant": "2", "x-internal-secret": "s3cret"}), + c.post("/v1/chat/completions", json=body), + ) + assert r1.status_code == r2.status_code == r3.status_code == 200 + + asyncio.run(run()) + assert sorted(seen, key=str) == ["sk-one", "sk-platform", "sk-two"] + + +def test_streaming_request_carries_tenant_env(host, monkeypatch): + """stream:true goes through asyncio.create_task in the shim — the context + must propagate into the task.""" + monkeypatch.setenv("OPENAI_API_KEY", "sk-platform") + monkeypatch.setattr(cpc, "_client", _FakeCPClient({7: {"OPENAI_API_KEY": "sk-tenant"}})) + seen: list = [] + + async def streaming_call(request, emit): + seen.append(cpc.env_get("OPENAI_API_KEY")) + await emit({"delta": "hi"}) + return _ok_result() + + host.set_async_call_hook(_capture_hook(seen)) + client = TestClient(create_app(host, default_profile="default", + streaming_call=streaming_call)) + r = client.post("/v1/chat/completions", + headers={"x-llm-router-tenant": "7", "x-internal-secret": "s3cret"}, + json={"model": "profile:default", "stream": True, + "messages": [{"role": "user", "content": "hi"}]}) + assert r.status_code == 200 + assert "sk-tenant" in seen + + +def test_adapter_authorization_header_uses_tenant_key(monkeypatch): + """End of the chain: the OpenAI-compatible adapter builds Authorization from + control_plane_client.env_get, so an active tenant map changes the wire key.""" + from provider_adapters.openai_compatible import make_async_call_provider + + captured: dict = {} + + class _FakeHTTPResp: + status_code = 200 + + def json(self): + return {"choices": [{"message": {"content": "ok"}, + "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, + "total_tokens": 2}} + + class _FakeHTTPClient: + async def post(self, url, json=None, headers=None, timeout=None): + captured["headers"] = headers + return _FakeHTTPResp() + + call = make_async_call_provider(env_get=cpc.env_get, client=_FakeHTTPClient()) + request = {"provider_id": "openai", "served_model_id": "gpt-x", + "base_url": "https://api.test/v1", + "auth": {"kind": "bearer", "env": "OPENAI_API_KEY"}, + "messages": [{"role": "user", "content": "hi"}]} + monkeypatch.setenv("OPENAI_API_KEY", "sk-platform") + token = cpc.activate_tenant_env({"OPENAI_API_KEY": "sk-tenant"}) + try: + result = asyncio.run(call(dict(request))) + finally: + cpc.reset_tenant_env(token) + assert result["ok"] is True + assert captured["headers"]["Authorization"] == "Bearer sk-tenant" + + result = asyncio.run(call(dict(request))) + assert captured["headers"]["Authorization"] == "Bearer sk-platform"