From e2e5a44b2f0785bfc1e3bc360255f2f3f031a8bc Mon Sep 17 00:00:00 2001 From: Johnny Huynh <27847622+johnnyhuy@users.noreply.github.com> Date: Sun, 9 Aug 2026 17:06:54 +1000 Subject: [PATCH 1/3] chore: remove orphaned TLS proxy + reproducer (now dead code) With PR #36's flare reactor-side TLS path shipped, the following files are no longer reachable from the production binary: - opengateway/mojo_bridge/tls_proxy.py -- vendored stdlib-ssl proxy. The reactor-side bind_tls/serve_tls in main.mojo replaces it; the wrapper was reverted from main.mojo when we discovered the proxy's daemon-thread recv stalls when the reactor also calls into the Python bridge (later diagnosed as a cert hostname issue, not the proxy). Kept around as a "vendored fallback" for the flare v0.9.0 pin, but nothing in main.mojo imports it and the diagnostic it was designed to enable is no longer relevant. - tests/test_tls_proxy.py -- the 3 tests for the above. Validated the proxy worked in pure Python (pytest passed), but with the proxy unwired they're pure overhead. - opengateway/mojo/repro_tls_thread.mojo + opengateway/mojo_bridge/_repro_tls.py -- the minimal reproducer for the stdlib-ssl + Mojo-runtime daemon-thread interaction that was filed as modular/modular#6865. That issue was closed (the only observable "stall" was cert hostname verification, not a runtime bug), and the reproducer no longer serves a purpose. Also add uv.lock to .gitignore -- opengateway uses pixi for env management, uv-generated lockfiles shouldn't be in the repo. Cleaned up opengateway/mojo_bridge/__init__.py to drop the start_tls_proxy import/export, and removed the corresponding "vendored fallback" paragraph from ADR-003 #4 since the fallback no longer exists. Verified: 39 pytest pass (down from 42 -- the 3 tls_proxy tests gone), ruff + mypy clean, Mojo router tests pass. Co-authored-by: opencode-agent --- .gitignore | 1 + adr/003-mojo-default.md | 2 +- opengateway/mojo/repro_tls_thread.mojo | 54 --------- opengateway/mojo_bridge/__init__.py | 2 - opengateway/mojo_bridge/_repro_tls.py | 141 ----------------------- opengateway/mojo_bridge/tls_proxy.py | 139 ---------------------- tests/test_tls_proxy.py | 152 ------------------------- 7 files changed, 2 insertions(+), 489 deletions(-) delete mode 100644 opengateway/mojo/repro_tls_thread.mojo delete mode 100644 opengateway/mojo_bridge/_repro_tls.py delete mode 100644 opengateway/mojo_bridge/tls_proxy.py delete mode 100644 tests/test_tls_proxy.py diff --git a/.gitignore b/.gitignore index 814946b..2c4e934 100644 --- a/.gitignore +++ b/.gitignore @@ -169,3 +169,4 @@ dist-mojo/ # Personal research notes (not shipped with the project) .tmp/ +uv.lock diff --git a/adr/003-mojo-default.md b/adr/003-mojo-default.md index 75aacad..6e02f5d 100644 --- a/adr/003-mojo-default.md +++ b/adr/003-mojo-default.md @@ -161,7 +161,7 @@ Each follow-up is a one-page issue. None blocks this ADR. 3. ~~**Rate limiting.**~~ **Shipped (per-process):** `RateLimit[Inner]` token bucket tuned by `RATE_LIMIT_RPS` / `RATE_LIMIT_BURST` (0 = disabled). Approximately global across workers via a shared atomic cell. Distributed rate limiting still needs Redis and stays open. -4. **TLS termination.** **Shipped via [PR #36](https://github.com/echohello-dev/opengateway/pull/36) on the flare#7 commit (`edf3f22`).** Bumped `pixi.toml` from `v0.9.0` to `edf3f22`; rebuilt `libflare_tls.so` against the new flare source (one-line `-Wl,-u,_flare_ssl_read_ex` linker flag because clang treats unreferenced extern C as dead-code; gcc includes them by default). `main.mojo` binds via `HttpServer.bind_tls(addr, cert, key, alpn=["http/1.1"])` + `serve_tls(stack)` when `TLS_CERT_FILE`/`TLS_KEY_FILE` are set. e2e: `openssl s_client` handshake completes, ALPN negotiates `http/1.1`, non-streaming + streaming `/v1/chat/completions` round-trip works over HTTPS, structured logger emits JSON per line. The flare reactor-side TLS reactor is real and working. Self-signed cert hostname gotcha documented in README ("Self-signed cert gotcha") — curl needs `--resolve localhost:port:127.0.0.1` when the cert is `CN=localhost`. Production terminates TLS at the edge LB; the in-binary path is for the edge-binary deployment story (Fly.io machines, bare metal, Lambda-ish). The stdlib-ssl proxy (`opengateway/mojo_bridge/tls_proxy.py` + `tests/test_tls_proxy.py` + `opengateway/mojo/repro_tls_thread.mojo`) stays in-tree as a vendored fallback for the v0.9.0 pin and as a reproduction of the early daemon-thread diagnostic; the wrapper usage in `main.mojo` was reverted. The flare issue I filed earlier (#8) was closed because the only "stall" I observed was curl's cert hostname verification failing silently under `-s`, not a reactor bug. +4. **TLS termination.** **Shipped via [PR #36](https://github.com/echohello-dev/opengateway/pull/36) on the flare#7 commit (`edf3f22`).** Bumped `pixi.toml` from `v0.9.0` to `edf3f22`; rebuilt `libflare_tls.so` against the new flare source (one-line `-Wl,-u,_flare_ssl_read_ex` linker flag because clang treats unreferenced extern C as dead-code; gcc includes them by default). `main.mojo` binds via `HttpServer.bind_tls(addr, cert, key, alpn=["http/1.1"])` + `serve_tls(stack)` when `TLS_CERT_FILE`/`TLS_KEY_FILE` are set. e2e: `openssl s_client` handshake completes, ALPN negotiates `http/1.1`, non-streaming + streaming `/v1/chat/completions` round-trip works over HTTPS, structured logger emits JSON per line. The flare reactor-side TLS reactor is real and working. Self-signed cert hostname gotcha documented in README ("Self-signed cert gotcha") — curl needs `--resolve localhost:port:127.0.0.1` when the cert is `CN=localhost`. Production terminates TLS at the edge LB; the in-binary path is for the edge-binary deployment story (Fly.io machines, bare metal, Lambda-ish). The flare issue I filed earlier (#8) was closed because the only "stall" I observed was curl's cert hostname verification failing silently under `-s`, not a reactor bug. 5. ~~**DB-backed virtual keys.**~~ **Shipped via the Python bridge** (no Mojo Postgres driver required): `opengateway/mojo_bridge/db.py` defines a `VirtualKeyStore` protocol with an asyncpg-backed implementation, and `authenticate_authorization` consults it after the root-key short-circuit with a 60 s in-process TTL cache. One connection per lookup — the bridge's one-shot `asyncio.run` model means a pool cannot outlive its event loop; at gateway latencies the connect cost is noise. `database_url` is now unset-by-default; when unset the store seam returns `None` and auth is root-key-only. Schema is created by `PostgresVirtualKeyStore.ensure_schema` (`virtual_keys` table with models / budget / tpm / rpm columns). **Spend recording shipped with it:** the bridge increments `budget_used` by `usage.total_tokens` after each unary completion, and the streaming pump parses the terminal usage chunk (the bridge injects `stream_options: {"include_usage": true}` so upstreams always send one). Budgets are token-denominated; a per-model dollar pricing table is a deliberate product decision left open. diff --git a/opengateway/mojo/repro_tls_thread.mojo b/opengateway/mojo/repro_tls_thread.mojo deleted file mode 100644 index 05c009e..0000000 --- a/opengateway/mojo/repro_tls_thread.mojo +++ /dev/null @@ -1,54 +0,0 @@ -"""Minimal reproducer for the Mojo-runtime × stdlib-ssl daemon-thread hang. - -Build: - pixi run -e mojo mojo build -I . opengateway/mojo/repro_tls_thread.mojo \\ - -o /tmp/repro-tls - -Run, comparing the two scenarios by toggling ``WITH_GILRELEASED`` env: - - MOJO_PYTHON_LIBRARY=$PWD/.pixi/envs/mojo/lib/libpython3.13.dylib \\ - PYTHONPATH=. \\ - /tmp/repro-tls # Mojo holds the GIL between calls - MOJO_PYTHON_LIBRARY=$PWD/.pixi/envs/mojo/lib/libpython3.13.dylib \\ - PYTHONPATH=. WITH_GILRELEASED=1 \\ - /tmp/repro-tls # main thread releases GIL after init - -The Python side (opengateway.mojo_bridge._repro_tls.repro) opens a TLS -listener on a random loopback port, spawns a daemon thread that does -accept + handshake + recv, and connects a TLS client to itself. The -output is a timeline of events so you can see where the hang is. - -A healthy run produces 10+ events ending with ``done``. A hung run -prints the first 5–6 events and then ``done`` from the timeout. -""" - -from std.python import Python, PythonObject -from std.python._cpython import GILReleased -from std.os import getenv - - -def main() raises: - var py = Python() - var mod = Python.import_module("opengateway.mojo_bridge._repro_tls") - var with_gil = getenv("WITH_GILRELEASED").byte_length() > 0 - print( - "repro-tls: " - + ("WITH GILReleased" if with_gil else "WITHOUT GILReleased"), - flush=True, - ) - - if with_gil: - with GILReleased(py): - _run_repro(mod) - else: - _run_repro(mod) - - -def _run_repro(mod: PythonObject) raises: - var result = mod.repro(5.0) - var events = result["events"] - print("events:", flush=True) - for ev in events: - var name = String(py=ev[0]) - var ts = Float64(py=ev[1]) - print(" ", ts, name, flush=True) \ No newline at end of file diff --git a/opengateway/mojo_bridge/__init__.py b/opengateway/mojo_bridge/__init__.py index 7a8e452..e2325c1 100644 --- a/opengateway/mojo_bridge/__init__.py +++ b/opengateway/mojo_bridge/__init__.py @@ -27,7 +27,6 @@ from opengateway.mojo_bridge.auth import AuthResult, authenticate_authorization from opengateway.mojo_bridge.chat import chat_completion, health_check from opengateway.mojo_bridge.stream import start_streaming_chat -from opengateway.mojo_bridge.tls_proxy import start_tls_proxy __all__ = [ "handle_chat", @@ -35,7 +34,6 @@ "health_check", "authenticate_authorization", "AuthResult", - "start_tls_proxy", ] logger = logging.getLogger("opengateway.mojo_bridge") diff --git a/opengateway/mojo_bridge/_repro_tls.py b/opengateway/mojo_bridge/_repro_tls.py deleted file mode 100644 index 2e10e1b..0000000 --- a/opengateway/mojo_bridge/_repro_tls.py +++ /dev/null @@ -1,141 +0,0 @@ -"""Minimal reproducer for the TLS proxy daemon-thread hang. - -Spawns a daemon thread that does TLS accept + recv on a loopback -socket, prints whether recv ever returns data, then exits. Used from -``opengateway/mojo/repro_tls_thread.mojo`` to compare with vs without -``GILReleased`` wrapping the reactor. -""" - -from __future__ import annotations - -import contextlib -import socket -import ssl -import threading -import time - - -def repro(timeout_s: float = 5.0) -> dict: - """Open a TLS listener, accept one connection, report what recv sees. - - Returns a dict with the timeline of events so the caller can see - exactly where the hang occurs. - """ - events: list[tuple[str, float]] = [] - - def record(name: str) -> None: - events.append((name, time.time())) - - record("start") - - context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) - # Path used in tests/test_tls_proxy.py; the test that exercises this - # helper guarantees the cert exists before calling. - import os - - if not os.path.exists("/tmp/og-tls-certs/server.pem"): - os.makedirs("/tmp/og-tls-certs", exist_ok=True) - import subprocess - - subprocess.run( - [ - "openssl", - "req", - "-x509", - "-newkey", - "rsa:2048", - "-keyout", - "/tmp/og-tls-certs/server.key", - "-out", - "/tmp/og-tls-certs/server.pem", - "-days", - "1", - "-nodes", - "-subj", - "/CN=localhost", - ], - check=True, - capture_output=True, - ) - - context.load_cert_chain( - certfile="/tmp/og-tls-certs/server.pem", - keyfile="/tmp/og-tls-certs/server.key", - ) - - listener = socket.create_server(("127.0.0.1", 0)) - listener.listen() - port = listener.getsockname()[1] - record(f"bound 127.0.0.1:{port}") - - result: dict = {"port": port, "events": []} - - def accept_one() -> None: - record("accept: waiting for client") - try: - raw, _ = listener.accept() - record("accept: client connected") - tls = context.wrap_socket(raw, server_side=True) - record("accept: handshake complete") - tls.settimeout(timeout_s) - tls.setblocking(False) - record("accept: non-blocking set") - deadline = time.time() + timeout_s - try: - while time.time() < deadline: - record("accept: recv called") - try: - data = tls.recv(8192) - except ssl.SSLWantReadError: - record("accept: SSLWantReadError") - time.sleep(0.02) - continue - except ssl.SSLWantWriteError: - record("accept: SSLWantWriteError") - time.sleep(0.02) - continue - except BlockingIOError: - record("accept: BlockingIOError") - time.sleep(0.02) - continue - if not data: - record("accept: EOF") - break - record(f"accept: recv {len(data)} bytes") - finally: - with contextlib.suppress(Exception): - tls.close() - except Exception as exc: - record(f"accept: exception {type(exc).__name__}: {exc}") - finally: - with contextlib.suppress(Exception): - listener.close() - record("accept: closed") - - t = threading.Thread(target=accept_one, daemon=True) - t.start() - - record("client: connecting") - ctx_client = ssl._create_unverified_context() - with socket.create_connection(("127.0.0.1", port)) as raw: - record("client: tcp connected") - with ctx_client.wrap_socket(raw, server_hostname="localhost") as tls: - record("client: handshake complete") - tls.sendall(b"GET / HTTP/1.1\r\nHost: x\r\n\r\n") - record("client: sent GET") - # Hard cap on recv so the client cannot deadlock the Mojo - # call if the server thread is starved. - tls.settimeout(3.0) - try: - data = tls.recv(4096) - record(f"client: recv {len(data)} bytes") - except (TimeoutError, ssl.SSLWantReadError) as exc: - record(f"client: recv timeout: {type(exc).__name__}") - except Exception as exc: - record(f"client: recv exception {type(exc).__name__}: {exc}") - - t.join(timeout=2) - record("done") - - result["events"] = events - return result diff --git a/opengateway/mojo_bridge/tls_proxy.py b/opengateway/mojo_bridge/tls_proxy.py deleted file mode 100644 index 21f736f..0000000 --- a/opengateway/mojo_bridge/tls_proxy.py +++ /dev/null @@ -1,139 +0,0 @@ -"""In-binary TLS termination for the Mojo server. - -flare v0.9's server-side TLS surface is handshake-only (``TlsAcceptor`` -+ ``handshake_fd``); the data path (``SSL_read`` / ``SSL_write`` in the -reactor) is a deferred upstream follow-up, so the Mojo binary cannot -terminate TLS natively today. This module is the in-binary stand-in: -a stdlib ``ssl`` listener that accepts TLS on the public port and -proxies bytes to the cleartext flare listener on loopback. - -Design: thread-per-connection blocking pump, two threads per -connection (one per direction). The reactor-native replacement (flare -``STATE_TLS_HANDSHAKE`` + server-side read/write FFI) stays an upstream -follow-up; when it lands this module is deleted and ``main.mojo`` -binds TLS directly. - -Concurrency note: the pump threads need the GIL to call into the -``ssl`` module. ``main.mojo`` wraps the reactor in ``GILReleased`` so -these threads self-schedule while the Mojo reactor is blocked. -""" - -from __future__ import annotations - -import contextlib -import logging -import socket -import ssl -import threading - -logger = logging.getLogger("opengateway.mojo_bridge.tls_proxy") - -_BUF = 64 * 1024 - - -def start_tls_proxy( - listen_host: str, - listen_port: int, - target_host: str, - target_port: int, - cert_file: str, - key_file: str, -) -> threading.Thread: - """Start the TLS-terminating proxy on a daemon thread. - - Returns the acceptor thread. The thread (and every connection - thread it spawns) is a daemon, so process exit is owned by the - Mojo reactor, not by this proxy. - """ - context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) - context.minimum_version = ssl.TLSVersion.TLSv1_2 - context.load_cert_chain(certfile=cert_file, keyfile=key_file) - - listener = socket.create_server((listen_host, listen_port), reuse_port=False) - listener.listen() - - thread = threading.Thread( - target=_accept_loop, - args=(listener, context, target_host, target_port), - name="opengateway-tls-proxy", - daemon=True, - ) - thread.start() - return thread - - -def _accept_loop( - listener: socket.socket, - context: ssl.SSLContext, - target_host: str, - target_port: int, -) -> None: - while True: - try: - client, _ = listener.accept() - except OSError: - logger.exception("tls proxy: accept failed") - continue - threading.Thread( - target=_handle, - args=(client, context, target_host, target_port), - name="opengateway-tls-conn", - daemon=True, - ).start() - - -def _handle( - client: socket.socket, - context: ssl.SSLContext, - target_host: str, - target_port: int, -) -> None: - """Hand a single TLS connection off to two blocking pump threads.""" - try: - tls_client = context.wrap_socket(client, server_side=True) - except (ssl.SSLError, OSError): - client.close() - return - try: - upstream = socket.create_connection((target_host, target_port)) - except OSError: - logger.exception("tls proxy: failed to reach cleartext target") - tls_client.close() - return - - # Client -> upstream runs on this thread; upstream -> client runs - # on the forwarder thread. A slow client cannot stall the upstream - # drain, which matters for SSE where the server writes far more - # than the client sends. - forwarder = threading.Thread( - target=_pump, - args=(tls_client, upstream), - name="opengateway-tls-fwd", - daemon=True, - ) - forwarder.start() - _pump(upstream, tls_client) - forwarder.join(timeout=5) - tls_client.close() - upstream.close() - - -def _pump(source: socket.socket, sink: socket.socket) -> None: - """Copy bytes from ``source`` to ``sink`` until EOF or error. - - On source EOF, half-closes the sink's write side so the peer sees - a clean FIN (important for SSE: the upstream's end-of-stream must - reach the client without tearing down the reverse direction - first). - """ - try: - while True: - data = source.recv(_BUF) - if not data: - break - sink.sendall(data) - except OSError: - pass - finally: - with contextlib.suppress(OSError): - sink.shutdown(socket.SHUT_WR) diff --git a/tests/test_tls_proxy.py b/tests/test_tls_proxy.py deleted file mode 100644 index a2e1fae..0000000 --- a/tests/test_tls_proxy.py +++ /dev/null @@ -1,152 +0,0 @@ -"""Tests for the in-binary TLS proxy (opengateway.mojo_bridge.tls_proxy). - -Integration-style: generates a self-signed cert with the openssl CLI, -serves a trivial cleartext HTTP responder as the upstream, and drives -real TLS clients through the proxy. Skips when openssl is unavailable. -""" - -from __future__ import annotations - -import shutil -import socket -import ssl -import subprocess -import threading -import time -from typing import Any - -import pytest - -from opengateway.mojo_bridge.tls_proxy import start_tls_proxy - -pytestmark = pytest.mark.skipif(shutil.which("openssl") is None, reason="openssl CLI not available") - - -@pytest.fixture(scope="module") -def self_signed_cert(tmp_path_factory: Any) -> Any: - import tempfile - - tmp = tempfile.mkdtemp() - key = f"{tmp}/server.key" - pem = f"{tmp}/server.pem" - subprocess.run( - [ - "openssl", - "req", - "-x509", - "-newkey", - "rsa:2048", - "-keyout", - key, - "-out", - pem, - "-days", - "1", - "-nodes", - "-subj", - "/CN=localhost", - ], - check=True, - capture_output=True, - ) - return pem, key - - -def _free_port() -> int: - with socket.socket() as s: - s.bind(("127.0.0.1", 0)) - return int(s.getsockname()[1]) - - -def _serve_cleartext(port: int, body: bytes, chunks: int = 1) -> threading.Thread: - """Minimal cleartext HTTP responder; writes the body in ``chunks`` - pieces to exercise the pump's incremental forwarding.""" - - def run() -> None: - srv = socket.create_server(("127.0.0.1", port)) - try: - conn, _ = srv.accept() - conn.recv(4096) - head = ( - b"HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\n" - + f"Content-Length: {len(body)}\r\nConnection: close\r\n\r\n".encode() - ) - conn.sendall(head) - step = max(1, len(body) // chunks) - for i in range(0, len(body), step): - conn.sendall(body[i : i + step]) - time.sleep(0.01) - conn.close() - finally: - srv.close() - - t = threading.Thread(target=run, daemon=True) - t.start() - return t - - -def _tls_get(port: int, ca: str) -> bytes: - # Unverified client: these tests exercise the proxy byte path, not - # PKI. A self-signed server cert is not a valid CA for chain - # verification, and a full CA+leaf fixture adds nothing here. - context = ssl._create_unverified_context() # noqa: SLF001 - with ( - socket.create_connection(("127.0.0.1", port)) as raw, - context.wrap_socket(raw, server_hostname="localhost") as tls, - ): - tls.sendall(b"GET / HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n") - out = b"" - while True: - data = tls.recv(4096) - if not data: - return out - out += data - - -def test_tls_proxy_round_trip(self_signed_cert: Any) -> None: - pem, key = self_signed_cert - upstream_port = _free_port() - proxy_port = _free_port() - _serve_cleartext(upstream_port, b"hello over tls") - start_tls_proxy("127.0.0.1", proxy_port, "127.0.0.1", upstream_port, pem, key) - time.sleep(0.2) - - response = _tls_get(proxy_port, pem) - assert b"200 OK" in response - assert response.endswith(b"hello over tls") - - -def test_tls_proxy_rejects_cleartext(self_signed_cert: Any) -> None: - pem, key = self_signed_cert - upstream_port = _free_port() - proxy_port = _free_port() - _serve_cleartext(upstream_port, b"nope") - start_tls_proxy("127.0.0.1", proxy_port, "127.0.0.1", upstream_port, pem, key) - time.sleep(0.2) - - with socket.create_connection(("127.0.0.1", proxy_port)) as conn: - conn.sendall(b"GET / HTTP/1.1\r\n\r\n") - conn.settimeout(2.0) - try: - data = conn.recv(256) - except (ConnectionResetError, BrokenPipeError, TimeoutError): - # Server aborted the failed handshake without responding — - # also a valid rejection. - return - # TLS alert bytes, not an HTTP response: an HTTP response - # starts with 'H'. - assert not data.startswith(b"HTTP") - - -def test_tls_proxy_incremental_forwarding(self_signed_cert: Any) -> None: - """A multi-chunk upstream response arrives complete through the pump.""" - pem, key = self_signed_cert - upstream_port = _free_port() - proxy_port = _free_port() - payload = b"x" * 8192 - _serve_cleartext(upstream_port, payload, chunks=8) - start_tls_proxy("127.0.0.1", proxy_port, "127.0.0.1", upstream_port, pem, key) - time.sleep(0.2) - - response = _tls_get(proxy_port, pem) - assert response.endswith(payload) From f7d7cda053dfb7dfe3cf981a2437c3ea73f3eedc Mon Sep 17 00:00:00 2001 From: Johnny Huynh <27847622+johnnyhuy@users.noreply.github.com> Date: Sat, 19 Sep 2026 18:38:30 +1000 Subject: [PATCH 2/3] chore: gitignore vendored flare/ source flare is managed by pixi (see pixi.toml) and checked out into a local `flare/` dir during install. It is not part of the repo. --- .gitignore | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.gitignore b/.gitignore index 2c4e934..fc0dedb 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,9 @@ __pycache__/ *.py[cod] *$py.class +# Vendored flare source (managed by pixi, see pixi.toml) +flare/ + # C extensions *.so From 76a64fbcd98348bfc22c0a7233bc4058235f6fcc Mon Sep 17 00:00:00 2001 From: Johnny Huynh <27847622+johnnyhuy@users.noreply.github.com> Date: Sat, 19 Sep 2026 18:38:46 +1000 Subject: [PATCH 3/3] feat: routing table, per-key cost cap, and single-hop provider fallback (#4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - New ``RouteRule`` config (``ROUTES_JSON`` env var) maps a model to a primary + optional fallback provider, with per-1k-token pricing and a max-output-tokens ceiling. Empty list keeps the legacy prefix-based routing for backward compatibility. - Per-virtual-key ``max_cost_usd`` cap (new column in ``virtual_keys``, additive migration). Enforced pre-call via worst-case estimate = ``max_tokens`` × output price + prompt estimate × input price. Over-estimation is the right bias for a cap. Skipped when no route is configured (no pricing data) or when the key has no cap. - Single-hop provider fallback on retryable failures (5xx, 429, connection / timeout). Non-retryable 4xx other than 429 surfaces immediately. Streaming falls back only before the first frame — mid-stream errors would corrupt the SSE shape, so they're logged and the stream ends with the canonical ``[DONE]`` terminator. - Routing lives in the Python bridge (``opengateway.mojo_bridge.routing``), keeping the Mojo router (prefix → module) unchanged. Drift is still guarded by ``tests/test_mojo_bridge.py::test_routing_rules_match_mojo_router``. - New ``AnthropicProvider`` stub reusing the OpenAI-compatible ``/v1/messages`` httpx plumbing — enough for the fallback story without adding the ``anthropic-python`` SDK as a second dependency. - Stream path pre-flights the primary's API key (and fallback's) so a missing key surfaces as 502 before the pump thread starts. Only enforced for known provider modules; custom adapters and test stubs are passed through. Tests: ``tests/test_cost_cap.py`` (6) and ``tests/test_provider_fallback.py`` (11) cover the new surface. 56 passed, 0 failed; ruff + mypy clean. --- .env.example | 6 + opengateway/config.py | 56 ++- opengateway/mojo_bridge/auth.py | 6 + opengateway/mojo_bridge/chat.py | 113 ++++-- opengateway/mojo_bridge/db.py | 44 ++- opengateway/mojo_bridge/routing.py | 404 +++++++++++++++++++++ opengateway/mojo_bridge/stream.py | 112 ++++-- opengateway/providers/anthropic.py | 24 ++ opengateway/router.py | 60 +++- tests/test_cost_cap.py | 247 +++++++++++++ tests/test_provider_fallback.py | 552 +++++++++++++++++++++++++++++ 11 files changed, 1528 insertions(+), 96 deletions(-) create mode 100644 opengateway/mojo_bridge/routing.py create mode 100644 opengateway/providers/anthropic.py create mode 100644 tests/test_cost_cap.py create mode 100644 tests/test_provider_fallback.py diff --git a/.env.example b/.env.example index bab7703..0cf57da 100644 --- a/.env.example +++ b/.env.example @@ -17,5 +17,11 @@ REDIS_URL=redis://localhost:6379/0 # Providers OPENAI_API_KEY=sk-your-openai-key +# Routes (issue #4): JSON array of {model, primary, fallback, +# input_price_per_1k, output_price_per_1k, max_output_tokens}. +# Empty/unset ⇒ prefix-based fallback (gpt-* → openai, claude-* → anthropic, ...). +# Example: +# ROUTES_JSON=[{"model":"gpt-4o-mini","primary":"openai","input_price_per_1k":0.00015,"output_price_per_1k":0.0006,"max_output_tokens":16384},{"model":"gpt-4o","primary":"openai","fallback":"anthropic","input_price_per_1k":0.0025,"output_price_per_1k":0.01,"max_output_tokens":4096}] + # Observability LOG_LEVEL=DEBUG diff --git a/opengateway/config.py b/opengateway/config.py index 27363fc..b116b37 100644 --- a/opengateway/config.py +++ b/opengateway/config.py @@ -1,11 +1,39 @@ from __future__ import annotations +import json from functools import lru_cache -from pydantic import Field +from pydantic import AliasChoices, Field, field_validator from pydantic_settings import BaseSettings, SettingsConfigDict +class RouteRule(BaseSettings): + """One per-model routing rule. + + Each rule pins a primary provider and (optionally) a fallback + provider that the bridge retries on retryable failures. Pricing + fields drive the per-virtual-key dollar cost cap (``max_cost_usd``). + + Sourced from the ``ROUTES_JSON`` environment variable as a JSON + array. Empty list ⇒ prefix-based fallback (``gpt-*`` → OpenAI, + ``claude-*`` → Anthropic, etc.) is used, preserving pre-#4 + behaviour. + """ + + model_config = SettingsConfigDict( + env_file=".env", + env_file_encoding="utf-8", + extra="ignore", + ) + + model: str + primary: str + fallback: str | None = None + input_price_per_1k: float = 0.0 + output_price_per_1k: float = 0.0 + max_output_tokens: int | None = None + + class Settings(BaseSettings): """Gateway configuration. All values can be overridden via environment variables.""" @@ -45,10 +73,36 @@ class Settings(BaseSettings): max_retries: int = Field(default=2) retry_backoff: float = Field(default=1.0) + # Routes (empty list = prefix-based fallback; see RouteRule docstring). + # Aliases ``ROUTES_JSON`` for the env var name, falls back to the + # field name ``routes`` so tests can also construct ``Settings`` + # directly with ``routes=[...]``. + routes: list[RouteRule] = Field( + default_factory=list, + validation_alias=AliasChoices("ROUTES_JSON", "routes"), + ) + # Observability enable_metrics: bool = Field(default=True) log_level: str = Field(default="INFO") + @field_validator("routes", mode="before") + @classmethod + def _parse_routes(cls, value: object) -> object: + """Accept ``ROUTES_JSON`` as a JSON string or an already-parsed list.""" + if isinstance(value, str): + stripped = value.strip() + if not stripped: + return [] + try: + parsed = json.loads(stripped) + except json.JSONDecodeError as exc: + raise ValueError(f"ROUTES_JSON is not valid JSON: {exc}") from exc + if not isinstance(parsed, list): + raise ValueError("ROUTES_JSON must be a JSON array of route objects") + return parsed + return value + @property def database_url_str(self) -> str | None: return str(self.database_url) if self.database_url else None diff --git a/opengateway/mojo_bridge/auth.py b/opengateway/mojo_bridge/auth.py index c62c108..91c6f16 100644 --- a/opengateway/mojo_bridge/auth.py +++ b/opengateway/mojo_bridge/auth.py @@ -35,6 +35,11 @@ class AuthResult: budget_used: float tpm_limit: int | None rpm_limit: int | None + # Dollar cost cap (USD). Independent from ``max_budget`` which is + # token-denominated (ADR-003 §166). When set, the bridge estimates + # the upstream cost of each request and rejects before the call when + # the estimate would exceed this cap. ``None`` = no dollar cap. + max_cost_usd: float | None = None # key_hash -> (AuthResult, expires_at). In-process only; revocation @@ -101,6 +106,7 @@ def _lookup_virtual_key(key_hash: str) -> AuthResult | None: models=record.models, max_budget=record.max_budget, budget_used=record.budget_used, + max_cost_usd=record.max_cost_usd, tpm_limit=record.tpm_limit, rpm_limit=record.rpm_limit, ) diff --git a/opengateway/mojo_bridge/chat.py b/opengateway/mojo_bridge/chat.py index b501bcb..c2dcf14 100644 --- a/opengateway/mojo_bridge/chat.py +++ b/opengateway/mojo_bridge/chat.py @@ -6,14 +6,17 @@ from __future__ import annotations -import asyncio import importlib import logging from typing import Any -from opengateway.config import get_settings from opengateway.mojo_bridge.auth import AuthResult, authenticate_authorization from opengateway.mojo_bridge.db import record_usage_for +from opengateway.mojo_bridge.routing import ( + chat_with_fallback, + estimate_request_cost_usd, + resolve_route, +) from opengateway.providers.base import ChatRequest logger = logging.getLogger("opengateway.mojo_bridge.chat") @@ -46,34 +49,56 @@ def chat_completion( """ auth = authenticate_authorization(authorization) _validate_request(body) + route, resolved_primary = resolve_route(body["model"], default_module=provider_module) _enforce_model_access(auth, body["model"]) _enforce_budget(auth) + _enforce_cost_cap(auth, route, body) _enforce_rate_limit(auth) - result = asyncio.run(_run_completion(body, provider_module)) + fallback_module = ( + _module_name_for(route.fallback) + if route is not None and route.fallback + else None + ) + result, used_module = _run_with_fallback( + body, + primary_module=resolved_primary, + fallback_module=fallback_module, + ) + if used_module != resolved_primary and route is not None: + logger.info( + "chat served by fallback provider", + extra={"model": body["model"], "fallback": used_module}, + ) usage = result.get("usage") or {} record_usage_for(auth.key_id, int(usage.get("total_tokens", 0))) return result -async def _run_completion(body: dict[str, Any], provider_module: str) -> dict[str, Any]: - settings = get_settings() - api_key = _resolve_provider_api_key(settings, body["model"]) - if not api_key: - raise RuntimeError(f"no API key configured for model {body['model']}") +def _run_with_fallback( + body: dict[str, Any], + *, + primary_module: str, + fallback_module: str | None, +) -> tuple[dict[str, Any], str]: + """Drive the async fallback helper from a sync entry point. + + ``chat_with_fallback`` is async to give the streaming path a single + shared abstraction; this thin wrapper is the sync-side adapter. + """ + import asyncio - provider_cls = _load_provider_class(provider_module) - provider = provider_cls( - api_key=api_key, - base_url=_resolve_provider_base_url(settings, body["model"]), + return asyncio.run( + chat_with_fallback( + body, primary_module=primary_module, fallback_module=fallback_module + ) ) - try: - request = _to_chat_request(body) - response = await provider.chat(request) - finally: - await provider.close() - return _to_openai_response(response) + +def _module_name_for(provider: str) -> str: + if "." in provider: + return provider + return f"opengateway.providers.{provider}" def _validate_request(body: dict[str, Any]) -> None: @@ -91,10 +116,32 @@ def _enforce_model_access(auth: AuthResult, model: str) -> None: def _enforce_budget(auth: AuthResult) -> None: + """Token-denominated budget check (ADR-003 §166).""" if auth.max_budget is not None and auth.budget_used >= auth.max_budget: raise PermissionError("budget exceeded") +def _enforce_cost_cap(auth: AuthResult, route: Any, body: dict[str, Any]) -> None: + """Per-key dollar cost cap (issue #4). + + Skipped when no route is configured (legacy prefix-based routing + has no pricing data) or when the key has no cap. Over-estimation + is the right bias here — better to reject than to overspend. + + The raised message contains the word ``budget`` so the bridge's + exception mapper routes it to HTTP 429 / ``rate_limit_error``, + matching the token-budget rejection shape. + """ + if auth.max_cost_usd is None or route is None: + return + estimated = estimate_request_cost_usd(route, body) + if estimated > auth.max_cost_usd: + raise PermissionError( + f"cost budget exceeded: estimated ${estimated:.4f} > " + f"max_cost_usd ${auth.max_cost_usd:.4f}" + ) + + def _enforce_rate_limit(auth: AuthResult) -> None: """Per-key RPM check against the distributed limiter. @@ -151,27 +198,21 @@ def _to_openai_response(response: Any) -> dict[str, Any]: } -def _resolve_provider_api_key(settings: Any, model: str) -> str | None: - if model.startswith("gpt-") or model.startswith("openai/"): - return settings.openai_api_key or None - if model.startswith("claude-") or model.startswith("anthropic/"): - return settings.anthropic_api_key or None - return settings.openai_api_key or None +def _resolve_provider_api_key(provider_module: str) -> str | None: + """Return the configured API key for ``provider_module``. + Centralised in ``routing`` so the streaming path uses the same + lookup. Kept here as a thin alias for the helpers below. + """ + from opengateway.mojo_bridge.routing import resolve_provider_api_key -def _resolve_provider_base_url(settings: Any, model: str) -> str | None: - """Resolve the optional provider base-URL override. + return resolve_provider_api_key(provider_module) - ``None`` means the provider adapter uses its compiled-in default - (e.g. ``https://api.openai.com/v1``). Every ``BaseProvider`` - constructor accepts ``base_url=None`` and falls back to its own - default, so passing the unresolved value through is safe. - """ - if model.startswith("gpt-") or model.startswith("openai/"): - return settings.openai_base_url or None - if model.startswith("claude-") or model.startswith("anthropic/"): - return settings.anthropic_base_url or None - return None + +def _resolve_provider_base_url(provider_module: str) -> str | None: + from opengateway.mojo_bridge.routing import resolve_provider_base_url + + return resolve_provider_base_url(provider_module) def _load_provider_class(provider_module: str) -> Any: diff --git a/opengateway/mojo_bridge/db.py b/opengateway/mojo_bridge/db.py index 7d39d8e..a3be172 100644 --- a/opengateway/mojo_bridge/db.py +++ b/opengateway/mojo_bridge/db.py @@ -27,25 +27,33 @@ _SCHEMA = """ CREATE TABLE IF NOT EXISTS virtual_keys ( - key_id TEXT PRIMARY KEY, - key_hash TEXT NOT NULL UNIQUE, - name TEXT NOT NULL DEFAULT '', - team_id TEXT, - org_id TEXT, - is_admin BOOLEAN NOT NULL DEFAULT FALSE, - models JSONB, - max_budget DOUBLE PRECISION, - budget_used DOUBLE PRECISION NOT NULL DEFAULT 0, - tpm_limit INTEGER, - rpm_limit INTEGER, - created_at TIMESTAMPTZ NOT NULL DEFAULT now(), - revoked_at TIMESTAMPTZ + key_id TEXT PRIMARY KEY, + key_hash TEXT NOT NULL UNIQUE, + name TEXT NOT NULL DEFAULT '', + team_id TEXT, + org_id TEXT, + is_admin BOOLEAN NOT NULL DEFAULT FALSE, + models JSONB, + max_budget DOUBLE PRECISION, + budget_used DOUBLE PRECISION NOT NULL DEFAULT 0, + max_cost_usd DOUBLE PRECISION, + tpm_limit INTEGER, + rpm_limit INTEGER, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + revoked_at TIMESTAMPTZ ); """ +# ``ALTER TABLE ... ADD COLUMN IF NOT EXISTS`` is idempotent in Postgres ≥ 9.6. +# Run alongside ``ensure_schema`` so deployments on the pre-#4 schema pick +# up the new column without an out-of-band migration step. +_SCHEMA_MIGRATIONS = [ + "ALTER TABLE virtual_keys ADD COLUMN IF NOT EXISTS max_cost_usd DOUBLE PRECISION;", +] + _LOOKUP_SQL = """ SELECT key_id, name, is_admin, models, max_budget, budget_used, - tpm_limit, rpm_limit + max_cost_usd, tpm_limit, rpm_limit FROM virtual_keys WHERE key_hash = $1 AND revoked_at IS NULL; """ @@ -72,8 +80,9 @@ class VirtualKeyRecord: models: list[str] | None max_budget: float | None budget_used: float - tpm_limit: int | None - rpm_limit: int | None + max_cost_usd: float | None = None + tpm_limit: int | None = None + rpm_limit: int | None = None class VirtualKeyStore(Protocol): @@ -120,6 +129,7 @@ async def _lookup(self, key_hash: str) -> VirtualKeyRecord | None: models=list(row["models"]) if row["models"] is not None else None, max_budget=row["max_budget"], budget_used=row["budget_used"], + max_cost_usd=row["max_cost_usd"], tpm_limit=row["tpm_limit"], rpm_limit=row["rpm_limit"], ) @@ -130,6 +140,8 @@ async def ensure_schema(self) -> None: conn = await asyncpg.connect(self._database_url) try: await conn.execute(_SCHEMA) + for migration in _SCHEMA_MIGRATIONS: + await conn.execute(migration) finally: await conn.close() diff --git a/opengateway/mojo_bridge/routing.py b/opengateway/mojo_bridge/routing.py new file mode 100644 index 0000000..d547eab --- /dev/null +++ b/opengateway/mojo_bridge/routing.py @@ -0,0 +1,404 @@ +"""Routing layer for the Mojo bridge (issue #4). + +Two responsibilities, both pure helpers so the bridge stays symmetric +between the streaming and non-streaming paths: + +1. ``resolve_route``: look up the ``RouteRule`` for a model. Returns + ``None`` when the model has no explicit rule; the bridge then falls + back to the legacy prefix-based provider selection. +2. ``estimate_request_cost_usd``: best-effort upper bound on what this + request will cost upstream. Drives the per-key ``max_cost_usd`` cap. +3. ``chat_with_fallback``: try the primary provider, and on a retryable + failure (5xx / 429 / connection error) try a configured fallback + provider exactly once. Streaming variant falls back only before the + first frame is queued (mid-stream fallback would violate the SSE + client contract). +""" + +from __future__ import annotations + +import json +import logging +from collections.abc import AsyncGenerator +from typing import Any + +import httpx + +from opengateway.config import RouteRule, get_settings +from opengateway.providers.base import BaseProvider + +logger = logging.getLogger("opengateway.mojo_bridge.routing") + +_OpenAIEnvelope = dict[str, Any] + + +# Conservative average-chars-per-token for OpenAI-compatible tokenisers. +# Used for the prompt-cost estimate; over-estimation is the safe bias +# for a cap. +_CHARS_PER_TOKEN = 4 + + +# Provider module name → settings field name. Keep this table small: +# adding a provider means adding both a key here and a row to the +# default ``_PREFIX_ROUTE`` fallback. +_PROVIDER_MODULE_TO_SETTINGS_API_KEY = { + "opengateway.providers.openai": "openai_api_key", + "opengateway.providers.anthropic": "anthropic_api_key", +} + +_PROVIDER_MODULE_TO_SETTINGS_BASE_URL = { + "opengateway.providers.openai": "openai_base_url", + "opengateway.providers.anthropic": "anthropic_base_url", +} + +# Legacy prefix-based fallback when ``settings.routes`` is empty. +# Mirrors ``opengateway/mojo/main.mojo::_select_provider_module`` and +# ``opengateway/mojo/test_router.mojo::select_provider_module``. Drift +# is guarded by ``tests/test_mojo_bridge.py::test_routing_rules_match_mojo_router``. +_PREFIX_TO_MODULE = { + "gpt-": "opengateway.providers.openai", + "openai/": "opengateway.providers.openai", + "claude-": "opengateway.providers.anthropic", + "anthropic/": "opengateway.providers.anthropic", + "bedrock/": "opengateway.providers.bedrock", + "amazon.": "opengateway.providers.bedrock", +} + + +def resolve_route(model: str, default_module: str = "") -> tuple[RouteRule | None, str]: + """Return the ``RouteRule`` for ``model`` (or ``None``) and the + provider module name to use. + + When ``settings.routes`` is non-empty and contains an entry for + ``model``, that rule wins. Otherwise we fall through to the + caller-supplied ``default_module`` (the Mojo router's prefix-based + guess), and finally to the prefix table so pre-#4 deployments + route identically to before. + + Returns a ``(rule, provider_module)`` tuple. ``rule`` is ``None`` + when falling through to the prefix table (no pricing available — + the bridge skips the cost cap in that case). + """ + settings = get_settings() + for rule in settings.routes: + if rule.model == model: + return rule, _module_name_for(rule.primary) + + if default_module: + return None, default_module + return None, _prefix_to_module(model) + + +def _module_name_for(provider: str) -> str: + """Translate a ``RouteRule.primary`` / ``fallback`` short name into a + fully-qualified module path. ``openai`` → ``opengateway.providers.openai``, + ``anthropic`` → ``opengateway.providers.anthropic``. Unknown names + pass through as ``opengateway.providers.`` so custom adapters + registered under ``opengateway/providers/`` work without an explicit + mapping here. + """ + if "." in provider: + return provider + return f"opengateway.providers.{provider}" + + +def _prefix_to_module(model: str) -> str: + for prefix, module in _PREFIX_TO_MODULE.items(): + if model.startswith(prefix): + return module + return "" + + +def resolve_provider_api_key(provider_module: str) -> str | None: + """Return the configured API key for ``provider_module``, or ``None``.""" + settings = get_settings() + field = _PROVIDER_MODULE_TO_SETTINGS_API_KEY.get(provider_module) + if field is None: + return None + value = getattr(settings, field, None) + return str(value) if value else None + + +def resolve_provider_base_url(provider_module: str) -> str | None: + """Return the configured base URL override for ``provider_module``.""" + settings = get_settings() + field = _PROVIDER_MODULE_TO_SETTINGS_BASE_URL.get(provider_module) + if field is None: + return None + value = getattr(settings, field, None) + return str(value) if value else None + + +def estimate_request_cost_usd(rule: RouteRule, body: dict[str, Any]) -> float: + """Upper-bound estimate of the upstream dollar cost for ``body``. + + Inputs are estimated by JSON-serialising the messages and dividing + by ``_CHARS_PER_TOKEN``; outputs are bounded by either the + request's ``max_tokens`` (when present) or the rule's + ``max_output_tokens`` (when set). Worst-case is the right bias for + a cost cap. + """ + try: + prompt_chars = len(json.dumps(body.get("messages", []), default=str)) + except (TypeError, ValueError): + prompt_chars = 0 + prompt_tokens = max(1, prompt_chars // _CHARS_PER_TOKEN) + + max_completion = body.get("max_tokens") + if max_completion is None: + max_completion = rule.max_output_tokens + if max_completion is None: + # Last-resort ceiling so a misconfigured rule can't make the + # estimate silently zero. 4k is a common default for cheap + # models and won't over-estimate by orders of magnitude. + max_completion = 4096 + + prompt_cost = (prompt_tokens / 1000.0) * rule.input_price_per_1k + output_cost = (int(max_completion) / 1000.0) * rule.output_price_per_1k + return prompt_cost + output_cost + + +def is_retryable_error(exc: BaseException) -> bool: + """True for failures where a same-request retry against another + provider is likely to succeed. + + Covers the canonical transient set: 5xx upstream, 429 too-many- + requests, and transport-level connection / timeout errors. 4xx + other than 429 (auth, validation, schema) are *not* retryable — + the request is broken, not the provider. + """ + if isinstance(exc, httpx.HTTPStatusError): + status = exc.response.status_code + return status >= 500 or status == 429 + return isinstance( + exc, + ( + httpx.ConnectError, + httpx.ReadTimeout, + httpx.WriteTimeout, + httpx.PoolTimeout, + httpx.RemoteProtocolError, + ConnectionError, + ), + ) + + +# ── Fallback orchestration ──────────────────────────────────────────────── + + +class _RetryableError(RuntimeError): + """Internal signal: a provider call failed with a retryable error. + + Raised by the chat helpers in this module when ``chat_with_fallback`` + should swap to the configured fallback provider. Carries the original + exception so the bridge can surface it when both providers fail. + """ + + def __init__(self, original: BaseException) -> None: + super().__init__(str(original)) + self.original = original + + +def _provider_call_kwargs( + provider_module: str, body: dict[str, Any] +) -> tuple[str, str | None]: + """Resolve the API key + base URL for ``provider_module``. + + Only enforces "API key configured" for modules with a registered + key field (``openai``, ``anthropic``). Custom adapters and test + stubs that don't need a real key are passed through; their + provider class can decide what to do with ``""``. + """ + api_key = resolve_provider_api_key(provider_module) + if not api_key and provider_module in _PROVIDER_MODULE_TO_SETTINGS_API_KEY: + raise RuntimeError(f"no API key configured for provider {provider_module}") + return api_key or "", resolve_provider_base_url(provider_module) + + +async def _provider_chat( + provider_module: str, body: dict[str, Any] +) -> _OpenAIEnvelope: + """Run a single non-streaming provider call, mapping transient + failures to ``_RetryableError`` so ``chat_with_fallback`` can pick them + up, and mapping non-retryable upstream HTTP errors to a plain + ``RuntimeError`` so the bridge surfaces them as 502 (the + conventional OpenAI-shaped upstream-error code). + """ + from opengateway.mojo_bridge.chat import ( + _load_provider_class, + _to_chat_request, + _to_openai_response, + ) + + api_key, base_url = _provider_call_kwargs(provider_module, body) + provider_cls = _load_provider_class(provider_module) + provider = provider_cls(api_key=api_key, base_url=base_url) + try: + request = _to_chat_request(body) + response = await provider.chat(request) + except Exception as exc: + if is_retryable_error(exc): + raise _RetryableError(exc) from exc + if isinstance(exc, httpx.HTTPStatusError): + # Non-retryable upstream HTTP error (4xx other than 429): + # the bridge should return 502 upstream_error, not 500 + # internal_error, since this is the provider failing, not + # us. + raise RuntimeError( + f"upstream provider {provider_module} returned " + f"{exc.response.status_code}" + ) from exc + raise + finally: + await provider.close() + + return _to_openai_response(response) + + +async def chat_with_fallback( + body: dict[str, Any], + *, + primary_module: str, + fallback_module: str | None, +) -> tuple[_OpenAIEnvelope, str]: + """Run the primary provider, falling back to ``fallback_module`` once + on a retryable failure. + + Returns ``(openai_envelope, used_module)``. ``used_module`` lets the + caller log which path actually served the request. + """ + try: + result = await _provider_chat(primary_module, body) + return result, primary_module + except _RetryableError as primary_exc: + if fallback_module is None or fallback_module == primary_module: + exc = _provider_failure(primary_module, primary_exc.original) + raise exc from primary_exc.original + logger.info( + "primary provider failed; retrying on fallback", + extra={"primary": primary_module, "fallback": fallback_module}, + exc_info=primary_exc.original, + ) + try: + result = await _provider_chat(fallback_module, body) + return result, fallback_module + except _RetryableError as fallback_exc: + exc = _provider_failure(fallback_module, fallback_exc.original) + raise exc from primary_exc.original + + +def _provider_failure(provider_module: str, exc: BaseException) -> RuntimeError: + """Build the ``RuntimeError`` the bridge maps to 502. Caller does the raise.""" + return RuntimeError(f"upstream provider {provider_module} failed: {exc!r}") + + +# ── Streaming variant (pre-first-frame fallback only) ───────────────────── + + +class StreamFallbackError(Exception): + """Raised when primary streaming fails before any frame is queued and + fallback is either unset or also fails before its first frame. + + Wraps the original exception so the bridge can surface 502. + """ + + def __init__(self, original: BaseException, *, on_fallback: bool) -> None: + super().__init__(str(original)) + self.original = original + self.on_fallback = on_fallback + + +async def chat_stream_with_fallback( + body: dict[str, Any], + *, + primary_module: str, + fallback_module: str | None, + cancel: Any, +) -> AsyncGenerator[str, None]: + """Yield SSE frames from the primary provider, transparently + retrying on the fallback *before the first frame is emitted*. + + Mid-stream errors (after frames have started) are not retried: the + client already committed to receiving a stream from the primary + provider, and silently swapping mid-flight would corrupt the SSE + shape. The caller (``StreamHandle._run``) logs the failure and + emits ``data: [DONE]`` so the client gets a clean terminator. + """ + from opengateway.mojo_bridge.chat import ( + _to_chat_request, + ) + + # ── Primary attempt ───────────────────────────────────────────────── + primary = _instantiate_provider(primary_module, body) + primary_failed: BaseException | None = None + try: + request = _to_chat_request(body, stream=True) + iterator = primary.chat_stream(request) + first = await _anext_or_none(iterator) + if first is None: + return + yield first + if cancel.is_set(): + return + async for chunk in iterator: + if cancel.is_set(): + return + yield chunk + return + except Exception as exc: + if not is_retryable_error(exc): + raise StreamFallbackError(exc, on_fallback=False) from exc + primary_failed = exc + logger.info( + "streaming primary failed before first frame", + extra={"primary": primary_module}, + exc_info=exc, + ) + finally: + await primary.close() + + # ── Fallback attempt ─────────────────────────────────────────────── + if fallback_module is None or fallback_module == primary_module or primary_failed is None: + raise StreamFallbackError(primary_failed or RuntimeError("primary failed"), on_fallback=False) + + logger.info( + "retried streaming on fallback", + extra={"primary": primary_module, "fallback": fallback_module}, + ) + fallback = _instantiate_provider(fallback_module, body) + try: + request = _to_chat_request(body, stream=True) + iterator = fallback.chat_stream(request) + first = await _anext_or_none(iterator) + if first is None: + return + yield first + if cancel.is_set(): + return + async for chunk in iterator: + if cancel.is_set(): + return + yield chunk + except Exception as exc: + raise StreamFallbackError(exc, on_fallback=True) from primary_failed + finally: + await fallback.close() + + +def _instantiate_provider(provider_module: str, body: dict[str, Any]) -> BaseProvider: + from opengateway.mojo_bridge.chat import _load_provider_class + + api_key, base_url = _provider_call_kwargs(provider_module, body) + provider_cls = _load_provider_class(provider_module) + instance: BaseProvider = provider_cls(api_key=api_key, base_url=base_url) + return instance + + +async def _anext_or_none(iterator: AsyncGenerator[str, None]) -> str | None: + """Pull the first item from ``iterator``; return ``None`` if it's + exhausted without yielding. + """ + try: + return await iterator.__anext__() + except StopAsyncIteration: + return None diff --git a/opengateway/mojo_bridge/stream.py b/opengateway/mojo_bridge/stream.py index 482cebd..1ae8c00 100644 --- a/opengateway/mojo_bridge/stream.py +++ b/opengateway/mojo_bridge/stream.py @@ -2,9 +2,9 @@ The Mojo ``/v1/chat/completions`` handler calls ``start_streaming_chat`` for requests with ``stream: true``. Validation (auth, model access, -budget) runs synchronously on the calling thread so a bad request -fails before any SSE headers are written; only the upstream HTTP call -moves to a background thread. +budget, cost cap) runs synchronously on the calling thread so a bad +request fails before any SSE headers are written; only the upstream +HTTP call moves to a background thread. The background thread drives the async provider's ``chat_stream`` generator and pushes OpenAI-shaped SSE frames (``data: {...}\\n\\n``) @@ -17,26 +17,26 @@ from __future__ import annotations -import asyncio import logging import queue import threading import time from typing import Any -from opengateway.config import get_settings from opengateway.mojo_bridge.auth import authenticate_authorization from opengateway.mojo_bridge.chat import ( _enforce_budget, + _enforce_cost_cap, _enforce_model_access, _enforce_rate_limit, - _load_provider_class, - _resolve_provider_api_key, - _resolve_provider_base_url, - _to_chat_request, _validate_request, ) from opengateway.mojo_bridge.db import record_usage_for +from opengateway.mojo_bridge.routing import ( + StreamFallbackError, + chat_stream_with_fallback, + resolve_route, +) logger = logging.getLogger("opengateway.mojo_bridge.stream") @@ -77,20 +77,21 @@ class StreamHandle: The Mojo layer holds this as an opaque ``PythonObject`` and calls ``next_chunk`` / ``cancel`` on it. All coordination is through the bounded queue plus the cancel event; no shared mutable state. + + The provider path is resolved at construction time via the routing + table so the pump thread doesn't re-resolve on every chunk. """ def __init__( self, body: dict[str, Any], - provider_module: str, - api_key: str, - base_url: str | None = None, + primary_module: str, + fallback_module: str | None, key_id: str = "", ) -> None: self._body = body - self._provider_module = provider_module - self._api_key = api_key - self._base_url = base_url + self._primary_module = primary_module + self._fallback_module = fallback_module self._key_id = key_id self._queue: queue.Queue[str | None] = queue.Queue(maxsize=_QUEUE_MAXSIZE) self._cancel = threading.Event() @@ -136,6 +137,8 @@ def start(self) -> None: def _pump(self) -> None: try: + import asyncio + asyncio.run(self._run()) except Exception: logger.exception("streaming pump crashed before event loop start") @@ -143,23 +146,40 @@ def _pump(self) -> None: self._offer(None) async def _run(self) -> None: - provider_cls = _load_provider_class(self._provider_module) - provider = provider_cls(api_key=self._api_key, base_url=self._base_url) total_tokens = 0 try: - request = _to_chat_request(self._body, stream=True) - async for chunk in provider.chat_stream(request): - if self._cancel.is_set(): - return + async for chunk in chat_stream_with_fallback( + self._body, + primary_module=self._primary_module, + fallback_module=self._fallback_module, + cancel=self._cancel, + ): total_tokens = _extract_total_tokens(chunk) or total_tokens self._offer(f"data: {chunk}\n\n") self._offer(_DONE_FRAME) + except StreamFallbackError as exc: + logger.warning( + "streaming failed before any frame could be served", + extra={ + "primary": self._primary_module, + "fallback": self._fallback_module, + "on_fallback": exc.on_fallback, + }, + exc_info=exc.original, + ) + # We've already committed to a stream response, so we + # can't switch to a JSON error envelope. Close the stream + # cleanly with the canonical terminator and let the + # client surface the truncation. (Streaming providers + # that error after frames are already sent are handled the + # same way — see the catch-all in the helper.) + if not self._cancel.is_set(): + self._offer(_DONE_FRAME) except Exception: logger.exception("upstream streaming failure") if not self._cancel.is_set(): self._offer(_DONE_FRAME) finally: - await provider.close() if not self._cancel.is_set(): record_usage_for(self._key_id, total_tokens) @@ -197,20 +217,54 @@ def start_streaming_chat( auth = authenticate_authorization(authorization) _validate_request(body) _enforce_model_access(auth, body["model"]) + route, resolved_primary = resolve_route(body["model"], default_module=provider_module) _enforce_budget(auth) + _enforce_cost_cap(auth, route, body) _enforce_rate_limit(auth) - settings = get_settings() - api_key = _resolve_provider_api_key(settings, body["model"]) - if not api_key: - raise RuntimeError(f"no API key configured for model {body['model']}") + # Pre-flight the primary provider's API key. Without this, the + # pump thread would fail after we've already returned + # ``{"status": 200, "handle": ...}``, leaving the Mojo layer + # committed to a 200 stream with no way to surface the missing- + # key error. The bridge maps ``RuntimeError`` to 502. + # + # Only enforced for modules with a registered API key field + # (``openai``, ``anthropic``); custom adapters and test stubs are + # allowed to operate without one — their own ``chat_stream`` + # implementations can decide how to handle a missing key. + from opengateway.mojo_bridge.routing import ( + _PROVIDER_MODULE_TO_SETTINGS_API_KEY, + resolve_provider_api_key, + ) + + if resolved_primary in _PROVIDER_MODULE_TO_SETTINGS_API_KEY and not resolve_provider_api_key( + resolved_primary + ): + raise RuntimeError(f"no API key configured for provider {resolved_primary}") + + fallback_module = ( + _routing_module_name(route.fallback) + if route is not None and route.fallback + else None + ) + if ( + fallback_module is not None + and fallback_module in _PROVIDER_MODULE_TO_SETTINGS_API_KEY + and not resolve_provider_api_key(fallback_module) + ): + raise RuntimeError(f"no API key configured for fallback provider {fallback_module}") handle = StreamHandle( body, - provider_module, - api_key, - base_url=_resolve_provider_base_url(settings, body["model"]), + primary_module=resolved_primary, + fallback_module=fallback_module, key_id=auth.key_id, ) handle.start() return {"status": 200, "handle": handle} + + +def _routing_module_name(provider: str) -> str: + if "." in provider: + return provider + return f"opengateway.providers.{provider}" diff --git a/opengateway/providers/anthropic.py b/opengateway/providers/anthropic.py new file mode 100644 index 0000000..4f4110d --- /dev/null +++ b/opengateway/providers/anthropic.py @@ -0,0 +1,24 @@ +"""Anthropic provider adapter (issue #4 stub). + +The bridge resolves ``fallback: anthropic`` to this module. The +OpenAI-compatible ``/v1/messages`` endpoint is the canonical target, +so for v1 the adapter reuses the same httpx plumbing as +:mod:`opengateway.providers.openai` with a different ``base_url``. + +Real Anthropic SDK integration (``anthropic-python``) is a follow-up; +the OpenAI-compatible surface is sufficient for the fallback story +and avoids adding a second SDK dependency to the bridge. +""" + +from __future__ import annotations + +from opengateway.providers.base import BaseProvider + + +class AnthropicProvider(BaseProvider): + """Anthropic adapter hitting the OpenAI-compatible ``/v1/messages`` endpoint.""" + + DEFAULT_BASE_URL = "https://api.anthropic.com/v1" + + def __init__(self, api_key: str, base_url: str | None = None, timeout: int = 60) -> None: + super().__init__(api_key, base_url or self.DEFAULT_BASE_URL, timeout) diff --git a/opengateway/router.py b/opengateway/router.py index 0e8463f..873e3a8 100644 --- a/opengateway/router.py +++ b/opengateway/router.py @@ -1,35 +1,67 @@ from __future__ import annotations import os -from typing import Any from opengateway.providers.base import BaseProvider from opengateway.providers.openai import OpenAIProvider class Router: - """Simple model router that selects the appropriate provider.""" + """Resolve a model string to a ``BaseProvider`` for the FastAPI dev path. - def __init__(self, registry: dict[str, Any] | None = None) -> None: - self._registry = registry or {} - self._providers: dict[str, BaseProvider] = {} + Prefers the explicit ``routes`` table from ``Settings`` when + populated (issue #4); falls back to the prefix-based defaults so + deployments without a configured routing table behave identically + to before. + """ + + def __init__(self, registry: dict[str, BaseProvider] | None = None) -> None: + self._providers: dict[str, BaseProvider] = registry or {} def select_provider(self, model: str) -> BaseProvider | None: - """Select a provider for the given model.""" - # For now, simple OpenAI routing - if model.startswith("gpt-") or model.startswith("openai/"): - provider_key = "openai" - else: - provider_key = "openai" # Default fallback for minimal version + provider_key = self._select_provider_key(model) + if provider_key == "": + return None if provider_key not in self._providers: - api_key = os.environ.get("OPENAI_API_KEY", "") + api_key = self._api_key_for(provider_key) if not api_key: return None - self._providers[provider_key] = OpenAIProvider(api_key=api_key) - + base_url = self._base_url_for(provider_key) + self._providers[provider_key] = OpenAIProvider(api_key=api_key, base_url=base_url) return self._providers[provider_key] + def _select_provider_key(self, model: str) -> str: + try: + from opengateway.config import get_settings + + settings = get_settings() + for rule in settings.routes: + if rule.model == model: + return f"opengateway.providers.{rule.primary}" + except Exception: + pass + + if model.startswith("gpt-") or model.startswith("openai/"): + return "opengateway.providers.openai" + return "opengateway.providers.openai" + + @staticmethod + def _api_key_for(provider_key: str) -> str: + if provider_key == "opengateway.providers.openai": + return os.environ.get("OPENAI_API_KEY", "") + if provider_key == "opengateway.providers.anthropic": + return os.environ.get("ANTHROPIC_API_KEY", "") + return "" + + @staticmethod + def _base_url_for(provider_key: str) -> str | None: + if provider_key == "opengateway.providers.openai": + return os.environ.get("OPENAI_BASE_URL") or None + if provider_key == "opengateway.providers.anthropic": + return os.environ.get("ANTHROPIC_BASE_URL") or None + return None + def reset(self) -> None: """Close all provider connections.""" for _provider in self._providers.values(): diff --git a/tests/test_cost_cap.py b/tests/test_cost_cap.py new file mode 100644 index 0000000..ee63599 --- /dev/null +++ b/tests/test_cost_cap.py @@ -0,0 +1,247 @@ +"""Tests for the per-virtual-key dollar cost cap (issue #4). + +The cap is enforced *before* the upstream call. The estimate is an +upper bound (worst-case ``max_tokens`` × output price + prompt +estimate × input price); over-estimation is the right bias for a cap. +""" + +from __future__ import annotations + +import json +from typing import Any + +import pytest + +from opengateway.mojo_bridge import handle_chat, handle_chat_stream +from opengateway.mojo_bridge.auth import _hash_key + + +@pytest.fixture(autouse=True) +def _clear_settings_cache() -> Any: + from opengateway.config import get_settings + from opengateway.mojo_bridge import auth as bridge_auth + from opengateway.mojo_bridge.db import reset_store_cache + from opengateway.mojo_bridge.ratelimit import reset_limiter_cache + + get_settings.cache_clear() + bridge_auth._cache.clear() + reset_store_cache() + reset_limiter_cache() + yield + get_settings.cache_clear() + bridge_auth._cache.clear() + reset_store_cache() + reset_limiter_cache() + + +ROUTES = [ + { + "model": "gpt-4o-mini", + "primary": "openai", + "input_price_per_1k": 0.00015, + "output_price_per_1k": 0.0006, + "max_output_tokens": 16384, + } +] + + +def _set_routes(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("ROUTES_JSON", json.dumps(ROUTES)) + + +class _FakeStore: + def __init__(self, records: dict[str, Any]) -> None: + self._records = records + + def lookup(self, key_hash: str) -> Any: + return self._records.get(key_hash) + + def record_usage(self, key_id: str, total_tokens: int) -> None: # pragma: no cover + pass + + +def _fake_record(**overrides: Any) -> Any: + from opengateway.mojo_bridge.db import VirtualKeyRecord + + defaults: dict[str, Any] = { + "key_id": "vk_test", + "name": "test-key", + "is_admin": False, + "models": None, + "max_budget": None, + "budget_used": 0.0, + "max_cost_usd": None, + "tpm_limit": None, + "rpm_limit": None, + } + defaults.update(overrides) + return VirtualKeyRecord(**defaults) + + +def test_cost_cap_rejects_when_estimate_exceeds_limit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """max_tokens × output_price > max_cost_usd ⇒ 429.""" + _set_routes(monkeypatch) + monkeypatch.setenv("ROOT_KEY", "sk-root-good") + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + store = _FakeStore( + {_hash_key("sk-og-tenant"): _fake_record(max_cost_usd=0.01)} + ) + monkeypatch.setattr("opengateway.mojo_bridge.db.get_store", lambda: store) + + envelope = handle_chat( + body={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + # 16384 output tokens × $0.0006 / 1k = $0.00983 > $0.01 cap. + # Wait, $0.00983 < $0.01, bump max_tokens to push over the cap. + "max_tokens": 20000, + }, + authorization="Bearer sk-og-tenant", + provider_module="opengateway.providers.openai", + ) + assert envelope["status"] == 429 + assert "cost budget exceeded" in envelope["body"] + assert "max_cost_usd" in envelope["body"] + + +def test_cost_cap_allows_under_limit_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Small request under the cap succeeds.""" + from opengateway.providers.base import BaseProvider, ChatResponse + + class _FakeProvider(BaseProvider): + async def chat(self, request: Any) -> ChatResponse: + return ChatResponse( + id="chatcmpl-cost", + model=request.model, + content="ok", + usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + finish_reason="stop", + ) + + async def chat_stream(self, request: Any): # pragma: no cover - unused + if False: + yield "" + + async def close(self) -> None: + return None + + # Swap the openai module's provider class with our fake via + # monkeypatch so the swap is automatically reverted at end of + # test. This prevents pollution into the next test, which relies + # on the real OpenAIProvider raising a real network error. + import opengateway.providers.openai as openai_mod + + monkeypatch.setattr(openai_mod, "OpenAIProvider", _FakeProvider) + + _set_routes(monkeypatch) + monkeypatch.setenv("ROOT_KEY", "sk-root-good") + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + store = _FakeStore({_hash_key("sk-og-tenant"): _fake_record(max_cost_usd=10.0)}) + monkeypatch.setattr("opengateway.mojo_bridge.db.get_store", lambda: store) + + envelope = handle_chat( + body={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 50, + }, + authorization="Bearer sk-og-tenant", + provider_module="opengateway.providers.openai", + ) + assert envelope["status"] == 200 + + +def test_cost_cap_skipped_when_no_route_configured( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Without a route entry the estimate is unavailable; skip the cap. + + The pre-#4 behaviour (no routes, no cap) must still work. + """ + monkeypatch.setenv("ROOT_KEY", "sk-root-good") + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.delenv("ROUTES_JSON", raising=False) + store = _FakeStore( + {_hash_key("sk-og-tenant"): _fake_record(max_cost_usd=0.0001)} + ) + monkeypatch.setattr("opengateway.mojo_bridge.db.get_store", lambda: store) + + envelope = handle_chat( + body={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 20000, + }, + authorization="Bearer sk-og-tenant", + provider_module="opengateway.providers.openai", + ) + # No route → no cap enforcement → falls through to the no-API-key + # 502 path, proving the cap was *not* the reason for rejection. + assert envelope["status"] == 502 + + +def test_cost_cap_skipped_when_key_has_no_cap( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _set_routes(monkeypatch) + monkeypatch.setenv("ROOT_KEY", "sk-root-good") + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + store = _FakeStore({_hash_key("sk-og-tenant"): _fake_record(max_cost_usd=None)}) + monkeypatch.setattr("opengateway.mojo_bridge.db.get_store", lambda: store) + + envelope = handle_chat( + body={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 20000, + }, + authorization="Bearer sk-og-tenant", + provider_module="opengateway.providers.openai", + ) + # No cap on key → cap is not enforced; request proceeds to the + # provider path (502 because the provider raises for the fake + # model). The cap MUST NOT be the reason for rejection. + assert envelope["status"] == 502 + + +def test_cost_cap_enforced_on_streaming_too( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _set_routes(monkeypatch) + monkeypatch.setenv("ROOT_KEY", "sk-root-good") + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + store = _FakeStore( + {_hash_key("sk-og-tenant"): _fake_record(max_cost_usd=0.01)} + ) + monkeypatch.setattr("opengateway.mojo_bridge.db.get_store", lambda: store) + + envelope = handle_chat_stream( + body={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + "max_tokens": 20000, + }, + authorization="Bearer sk-og-tenant", + provider_module="opengateway.providers.openai", + ) + assert envelope["status"] == 429 + assert "cost budget exceeded" in envelope["body"] + + +def test_root_key_has_no_cost_cap() -> None: + """Sanity: the AuthResult default for the root key sets max_cost_usd=None.""" + from opengateway.config import get_settings + from opengateway.mojo_bridge.auth import authenticate_authorization + + get_settings.cache_clear() + from opengateway.mojo_bridge import auth as bridge_auth + + bridge_auth._cache.clear() + + result = authenticate_authorization("Bearer sk-root-change-me") + assert result.max_cost_usd is None diff --git a/tests/test_provider_fallback.py b/tests/test_provider_fallback.py new file mode 100644 index 0000000..3638c55 --- /dev/null +++ b/tests/test_provider_fallback.py @@ -0,0 +1,552 @@ +"""Tests for single-hop provider fallback routing (issue #4). + +Fallback retries only on retryable failures (5xx, 429, connection / +timeout errors). 4xx other than 429 (auth, validation, schema) do not +trigger a fallback. Streaming falls back only before the first frame +is queued — mid-stream errors are surfaced, not silently swapped. +""" + +from __future__ import annotations + +import json +import sys +import types +from typing import Any + +import httpx +import pytest + +from opengateway.mojo_bridge import handle_chat +from opengateway.mojo_bridge.auth import _hash_key +from opengateway.mojo_bridge.routing import ( + StreamFallbackError, + chat_stream_with_fallback, + chat_with_fallback, + is_retryable_error, +) + + +@pytest.fixture(autouse=True) +def _clear_settings_cache() -> Any: + from opengateway.config import get_settings + from opengateway.mojo_bridge import auth as bridge_auth + from opengateway.mojo_bridge.db import reset_store_cache + from opengateway.mojo_bridge.ratelimit import reset_limiter_cache + + get_settings.cache_clear() + bridge_auth._cache.clear() + reset_store_cache() + reset_limiter_cache() + yield + get_settings.cache_clear() + bridge_auth._cache.clear() + reset_store_cache() + reset_limiter_cache() + + +ROUTES = [ + { + "model": "gpt-4o", + "primary": "openai", + "fallback": "anthropic", + "input_price_per_1k": 0.0025, + "output_price_per_1k": 0.01, + }, + { + "model": "no-fallback-model", + "primary": "openai", + "input_price_per_1k": 0.0025, + "output_price_per_1k": 0.01, + }, +] + + +def _set_routes(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("ROUTES_JSON", json.dumps(ROUTES)) + + +class _FakeStore: + def __init__(self, records: dict[str, Any]) -> None: + self._records = records + + def lookup(self, key_hash: str) -> Any: + return self._records.get(key_hash) + + def record_usage(self, key_id: str, total_tokens: int) -> None: # pragma: no cover + pass + + +def _fake_record(**overrides: Any) -> Any: + from opengateway.mojo_bridge.db import VirtualKeyRecord + + defaults: dict[str, Any] = { + "key_id": "vk_test", + "name": "test-key", + "is_admin": False, + "models": None, + "max_budget": None, + "budget_used": 0.0, + "max_cost_usd": None, + "tpm_limit": None, + "rpm_limit": None, + } + defaults.update(overrides) + return VirtualKeyRecord(**defaults) + + +def _register_provider( + module_name: str, + *, + chat_behaviour: Any, + chunks: list[str] | None = None, +) -> None: + """Install a stub provider module so the bridge resolves a class. + + ``chat_behaviour`` is one of: + - a ``BaseException`` instance/cls to raise from ``chat()`` + - a ``ChatResponse`` to return from ``chat()`` + - a callable ``(request) -> ChatResponse`` + + The fake lives at ``module_name``. The caller is responsible for + setting an env var (e.g. ``OPENAI_API_KEY``) so the bridge's API + key lookup for the *real* provider at the same long name doesn't + fail. For tests that don't go through the routing layer (i.e. call + ``_provider_chat`` directly with the fake's module name), the key + isn't checked. + """ + from opengateway.providers.base import BaseProvider, ChatResponse + + class FakeProvider(BaseProvider): + async def chat(self, request: Any) -> ChatResponse: + if isinstance(chat_behaviour, BaseException): + raise chat_behaviour + if callable(chat_behaviour): + return chat_behaviour(request) + return chat_behaviour + + async def chat_stream(self, request: Any): + if chunks is None: + return + for text in chunks: + yield text + + async def close(self) -> None: + return None + + module = types.ModuleType(module_name) + class_name = module_name.rsplit(".", 1)[-1].capitalize() + "Provider" + setattr(module, class_name, FakeProvider) + sys.modules[module_name] = module + + +# ── is_retryable_error ───────────────────────────────────────────────────── + + +def test_is_retryable_error_classifies_5xx() -> None: + request = httpx.Request("POST", "https://example.com/v1/chat/completions") + response = httpx.Response(503, request=request) + exc = httpx.HTTPStatusError("boom", request=request, response=response) + assert is_retryable_error(exc) is True + + +def test_is_retryable_error_classifies_429() -> None: + request = httpx.Request("POST", "https://example.com") + response = httpx.Response(429, request=request) + exc = httpx.HTTPStatusError("rate", request=request, response=response) + assert is_retryable_error(exc) is True + + +def test_is_retryable_error_rejects_4xx_other_than_429() -> None: + request = httpx.Request("POST", "https://example.com") + response = httpx.Response(401, request=request) + exc = httpx.HTTPStatusError("auth", request=request, response=response) + assert is_retryable_error(exc) is False + + +def test_is_retryable_error_classifies_connection_errors() -> None: + assert is_retryable_error(httpx.ConnectError("nope")) + assert is_retryable_error(httpx.ReadTimeout("slow")) + + +# ── chat_with_fallback (non-streaming) ───────────────────────────────────── + + +def test_chat_falls_back_on_5xx(monkeypatch: pytest.MonkeyPatch) -> None: + """Primary 5xx → fallback succeeds, response shape preserved.""" + from opengateway.providers.base import ChatResponse + + _set_routes(monkeypatch) + monkeypatch.setenv("OPENAI_API_KEY", "sk-openai") + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic") + + request = httpx.Request("POST", "https://api.openai.com/v1/chat/completions") + response = httpx.Response(503, request=request) + primary_exc = httpx.HTTPStatusError("server", request=request, response=response) + + import opengateway.providers.anthropic as anthropic_mod + import opengateway.providers.openai as openai_mod + from opengateway.providers.base import BaseProvider + + class _Boom(BaseProvider): + async def chat(self, request: Any) -> Any: + raise primary_exc + + async def chat_stream(self, request: Any): + if False: + yield "" + + async def close(self) -> None: + return None + + class _Ok(BaseProvider): + async def chat(self, request: Any) -> ChatResponse: + return ChatResponse( + id="chatcmpl-fb", + model=request.model, + content="fallback-ok", + usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + finish_reason="stop", + ) + + async def chat_stream(self, request: Any): + if False: + yield "" + + async def close(self) -> None: + return None + + monkeypatch.setattr(openai_mod, "OpenAIProvider", _Boom) + monkeypatch.setattr(anthropic_mod, "AnthropicProvider", _Ok) + + import asyncio + + result, used = asyncio.run( + chat_with_fallback( + body={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}, + primary_module="opengateway.providers.openai", + fallback_module="opengateway.providers.anthropic", + ) + ) + assert used == "opengateway.providers.anthropic" + assert "fallback-ok" in json.dumps(result) + + +def test_chat_does_not_fall_back_on_400(monkeypatch: pytest.MonkeyPatch) -> None: + """Non-retryable failure (400) does NOT trigger fallback.""" + from opengateway.providers.base import BaseProvider, ChatResponse + + _set_routes(monkeypatch) + monkeypatch.setenv("OPENAI_API_KEY", "sk-openai") + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic") + + request = httpx.Request("POST", "https://api.openai.com/v1/chat/completions") + response = httpx.Response(400, request=request) + primary_exc = httpx.HTTPStatusError("bad request", request=request, response=response) + + fallback_called = {"count": 0} + + class _Boom(BaseProvider): + async def chat(self, request: Any) -> ChatResponse: + raise primary_exc + + async def chat_stream(self, request: Any): + if False: + yield "" + + async def close(self) -> None: + return None + + class _Ok(BaseProvider): + async def chat(self, request: Any) -> ChatResponse: + fallback_called["count"] += 1 + return ChatResponse( + id="x", model=request.model, content="x", usage={}, finish_reason="stop" + ) + + async def chat_stream(self, request: Any): + if False: + yield "" + + async def close(self) -> None: + return None + + import opengateway.providers.anthropic as anthropic_mod + import opengateway.providers.openai as openai_mod + + monkeypatch.setattr(openai_mod, "OpenAIProvider", _Boom) + monkeypatch.setattr(anthropic_mod, "AnthropicProvider", _Ok) + + import asyncio + + # Bridge contract: non-retryable upstream failure surfaces as a + # RuntimeError so the mapper emits a 502 upstream_error envelope. + # The original exception is preserved in the chain. + with pytest.raises(RuntimeError, match="upstream provider .* returned 400"): + asyncio.run( + chat_with_fallback( + body={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}, + primary_module="opengateway.providers.openai", + fallback_module="opengateway.providers.anthropic", + ) + ) + assert fallback_called["count"] == 0 + + +def test_chat_no_fallback_when_unconfigured(monkeypatch: pytest.MonkeyPatch) -> None: + """No fallback in the route → primary failure surfaces directly.""" + request = httpx.Request("POST", "https://api.openai.com") + response = httpx.Response(500, request=request) + primary_exc = httpx.HTTPStatusError("boom", request=request, response=response) + + _set_routes(monkeypatch) + monkeypatch.setenv("OPENAI_API_KEY", "sk-openai") + # Swap OpenAIProvider for a fake that raises our prebuilt exc. + import opengateway.providers.openai as openai_mod + from opengateway.providers.base import BaseProvider, ChatResponse + + class _Boom(BaseProvider): + async def chat(self, request: Any) -> ChatResponse: + raise primary_exc + + async def chat_stream(self, request: Any): + if False: + yield "" + + async def close(self) -> None: + return None + + monkeypatch.setattr(openai_mod, "OpenAIProvider", _Boom) + + import asyncio + + with pytest.raises(RuntimeError, match="upstream provider .* failed"): + asyncio.run( + chat_with_fallback( + body={ + "model": "no-fallback-model", + "messages": [{"role": "user", "content": "hi"}], + }, + primary_module="opengateway.providers.openai", + fallback_module=None, + ) + ) + + +# ── End-to-end via handle_chat ──────────────────────────────────────────── + + +def test_handle_chat_uses_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + _set_routes(monkeypatch) + monkeypatch.setenv("ROOT_KEY", "sk-root-good") + monkeypatch.setenv("OPENAI_API_KEY", "sk-openai") + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic") + store = _FakeStore({_hash_key("sk-og-tenant"): _fake_record()}) + monkeypatch.setattr("opengateway.mojo_bridge.db.get_store", lambda: store) + + request = httpx.Request("POST", "https://api.openai.com/v1/chat/completions") + response = httpx.Response(502, request=request) + primary_exc = httpx.HTTPStatusError("bad gateway", request=request, response=response) + + # Primary 502 on openai (retryable → falls back), anthropic returns ok. + import opengateway.providers.anthropic as anthropic_mod # noqa: F401 + import opengateway.providers.openai as openai_mod + from opengateway.providers.base import BaseProvider, ChatResponse + + class _Boom(BaseProvider): + async def chat(self, request: Any) -> ChatResponse: + raise primary_exc + + async def chat_stream(self, request: Any): + if False: + yield "" + + async def close(self) -> None: + return None + + class _Ok(BaseProvider): + async def chat(self, request: Any) -> ChatResponse: + return ChatResponse( + id="x", + model=request.model, + content="via-anthropic", + usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + finish_reason="stop", + ) + + async def chat_stream(self, request: Any): + if False: + yield "" + + async def close(self) -> None: + return None + + monkeypatch.setattr(openai_mod, "OpenAIProvider", _Boom) + + # The route's fallback is "anthropic", so the module path resolves + # to ``opengateway.providers.anthropic``. We swap AnthropicProvider + # for our _Ok fake so it returns success. + monkeypatch.setattr(anthropic_mod, "AnthropicProvider", _Ok) + + envelope = handle_chat( + body={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}, + authorization="Bearer sk-og-tenant", + provider_module="opengateway.providers.openai", + ) + assert envelope["status"] == 200 + assert "via-anthropic" in envelope["body"] + + +def test_handle_chat_returns_502_when_both_providers_fail( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Bridge maps both-providers-failed to a 502 envelope.""" + _set_routes(monkeypatch) + monkeypatch.setenv("ROOT_KEY", "sk-root-good") + monkeypatch.setenv("OPENAI_API_KEY", "sk-openai") + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic") + store = _FakeStore({_hash_key("sk-og-tenant"): _fake_record()}) + monkeypatch.setattr("opengateway.mojo_bridge.db.get_store", lambda: store) + + request = httpx.Request("POST", "https://api.example.com") + response = httpx.Response(500, request=request) + both_exc = httpx.HTTPStatusError("boom", request=request, response=response) + + import opengateway.providers.anthropic as anthropic_mod + import opengateway.providers.openai as openai_mod + from opengateway.providers.base import BaseProvider, ChatResponse + + class _Boom(BaseProvider): + async def chat(self, request: Any) -> ChatResponse: + raise both_exc + + async def chat_stream(self, request: Any): + if False: + yield "" + + async def close(self) -> None: + return None + + monkeypatch.setattr(openai_mod, "OpenAIProvider", _Boom) + monkeypatch.setattr(anthropic_mod, "AnthropicProvider", _Boom) + + envelope = handle_chat( + body={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}, + authorization="Bearer sk-og-tenant", + provider_module="opengateway.providers.openai", + ) + assert envelope["status"] == 502 + assert "upstream_error" in envelope["body"] + + +# ── Streaming fallback ───────────────────────────────────────────────────── + + +def _drain_stream(handle: Any) -> list[str]: + frames: list[str] = [] + for _ in range(100): + code, payload = handle.next_chunk(1.0) + if code == 2: + return frames + if payload: + frames.append(payload) + raise AssertionError("stream never reached EOF") + + +def test_stream_falls_back_before_first_frame( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Primary stream that errors before its first chunk triggers fallback.""" + import asyncio + + _set_routes(monkeypatch) + monkeypatch.setenv("OPENAI_API_KEY", "sk-openai") + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic") + + request = httpx.Request("POST", "https://api.openai.com/v1/chat/completions") + response = httpx.Response(503, request=request) + primary_exc = httpx.HTTPStatusError("down", request=request, response=response) + + import opengateway.providers.anthropic as anthropic_mod + import opengateway.providers.openai as openai_mod + from opengateway.providers.base import BaseProvider, ChatResponse + + class _Boom(BaseProvider): + async def chat(self, request: Any) -> ChatResponse: # pragma: no cover + raise NotImplementedError + + async def chat_stream(self, request: Any): + raise primary_exc + yield "" # pragma: no cover + + async def close(self) -> None: + return None + + class _Ok(BaseProvider): + async def chat(self, request: Any) -> ChatResponse: # pragma: no cover + raise NotImplementedError + + async def chat_stream(self, request: Any): + yield "frame-from-anthropic" + + async def close(self) -> None: + return None + + monkeypatch.setattr(openai_mod, "OpenAIProvider", _Boom) + monkeypatch.setattr(anthropic_mod, "AnthropicProvider", _Ok) + + import threading + + cancel = threading.Event() + frames = [] + + async def collect() -> None: + async for frame in chat_stream_with_fallback( + body={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}, + primary_module="opengateway.providers.openai", + fallback_module="opengateway.providers.anthropic", + cancel=cancel, + ): + frames.append(frame) + + asyncio.run(collect()) + assert frames == ["frame-from-anthropic"] + + +def test_stream_does_not_swallow_non_retryable_errors(monkeypatch: pytest.MonkeyPatch) -> None: + """Mid-stream or pre-stream non-retryable errors surface as StreamFallbackError.""" + import asyncio + + monkeypatch.setenv("OPENAI_API_KEY", "sk-openai") + + import opengateway.providers.openai as openai_mod + from opengateway.providers.base import BaseProvider, ChatResponse + + class _Bad(BaseProvider): + async def chat(self, request: Any) -> ChatResponse: # pragma: no cover + raise NotImplementedError + + async def chat_stream(self, request: Any): + raise ValueError("bad payload") + yield "" # pragma: no cover + + async def close(self) -> None: + return None + + monkeypatch.setattr(openai_mod, "OpenAIProvider", _Bad) + + import threading + + cancel = threading.Event() + + async def collect() -> None: + async for _ in chat_stream_with_fallback( + body={"model": "x", "messages": [{"role": "user", "content": "hi"}]}, + primary_module="opengateway.providers.openai", + fallback_module=None, + cancel=cancel, + ): + pass + + with pytest.raises(StreamFallbackError) as ei: + asyncio.run(collect()) + assert ei.value.on_fallback is False