diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index aef1a675..b0fa2034 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -25,6 +25,7 @@ import time import uuid from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit import httpx @@ -67,6 +68,7 @@ # and exception class names — never headers, bodies, or credentials. _DIAGNOSTICS_ENV = "UCODE_RELAYED_PROXY_DIAGNOSTICS" _DIAGNOSTICS_TRUE = frozenset({"1", "true", "yes", "on"}) +_MODEL_ALIAS_PREFIX = "anthropic-aigw-" def _diagnostics_enabled() -> bool: @@ -193,11 +195,82 @@ def _forwarded_request_headers( return headers +class _ModelAliases: + """Maps Claude-compatible discovery IDs back to their gateway model IDs.""" + + def __init__(self) -> None: + self._original_by_alias: dict[str, str] = {} + self._lock = threading.Lock() + + def advertise_models(self, body: bytes) -> bytes: + try: + payload = json.loads(body) + models = payload["data"] + if not isinstance(models, list): + return body + except (UnicodeDecodeError, json.JSONDecodeError, KeyError, TypeError): + return body + + aliases: dict[str, str] = {} + for model in models: + if not isinstance(model, dict) or not isinstance(model.get("id"), str): + continue + model_id = model["id"] + if "claude" in model_id.lower() or "anthropic" in model_id.lower(): + continue + alias = f"{_MODEL_ALIAS_PREFIX}{model_id}" + model["id"] = alias + aliases[alias] = model_id + + with self._lock: + self._original_by_alias.update(aliases) + + for cursor in ("first_id", "last_id"): + model_id = payload.get(cursor) + alias = f"{_MODEL_ALIAS_PREFIX}{model_id}" + if alias in aliases: + payload[cursor] = alias + return json.dumps(payload, separators=(",", ":")).encode() + + def original_id(self, model_id: str) -> str: + with self._lock: + return self._original_by_alias.get(model_id, model_id) + + def rewrite_path(self, path: str) -> str: + parsed = urlsplit(path) + if parsed.path != "/v1/models": + return path + query = [ + (key, self.original_id(value) if key == "after_id" else value) + for key, value in parse_qsl(parsed.query, keep_blank_values=True) + ] + return urlunsplit( + (parsed.scheme, parsed.netloc, parsed.path, urlencode(query), parsed.fragment) + ) + + def rewrite_body(self, path: str, body: bytes | None) -> bytes | None: + if urlsplit(path).path != "/v1/messages" or body is None: + return body + try: + payload = json.loads(body) + model_id = payload.get("model") + if not isinstance(model_id, str): + return body + except (UnicodeDecodeError, json.JSONDecodeError, AttributeError): + return body + original_id = self.original_id(model_id) + if original_id == model_id: + return body + payload["model"] = original_id + return json.dumps(payload, separators=(",", ":")).encode() + + class _ProxyHandler(BaseHTTPRequestHandler): # Set by the server factory. cache: _TokenCache client: httpx.Client token_header = _SWAP_HEADER + model_aliases: _ModelAliases def log_message(self, format: str, *args: object) -> None: return @@ -215,7 +288,8 @@ def _handle(self) -> None: started = time.monotonic() length = int(self.headers.get("Content-Length", 0) or 0) body = self.rfile.read(length) if length else None - url = self.path.lstrip("/") + body = self.model_aliases.rewrite_body(self.path, body) + url = self.model_aliases.rewrite_path(self.path).lstrip("/") _diagnostic_log( "request_start", request_id=diagnostic_id, @@ -234,7 +308,13 @@ def _handle(self) -> None: elapsed_ms=round((time.monotonic() - started) * 1000), ) if resp.status_code not in (401, 403): - self._relay_response(resp, diagnostic_id=diagnostic_id, started=started) + self._relay_response( + resp, + transform_models=self.command == "GET" + and urlsplit(self.path).path == "/v1/models", + diagnostic_id=diagnostic_id, + started=started, + ) return # Auth rejected. Drain the (small) error body so the pooled # connection can be reused, then fall through to one retry. @@ -258,7 +338,13 @@ def _handle(self) -> None: status=resp.status_code, elapsed_ms=round((time.monotonic() - started) * 1000), ) - self._relay_response(resp, diagnostic_id=diagnostic_id, started=started) + self._relay_response( + resp, + transform_models=self.command == "GET" + and urlsplit(self.path).path == "/v1/models", + diagnostic_id=diagnostic_id, + started=started, + ) except (BrokenPipeError, ConnectionResetError): # Client closed before/while we relayed headers — routine on cancel. _diagnostic_log( @@ -289,6 +375,7 @@ def _relay_response( self, resp: httpx.Response, *, + transform_models: bool = False, diagnostic_id: str | None = None, started: float | None = None, ) -> None: @@ -299,7 +386,9 @@ def _relay_response( try: self.send_response(resp.status_code) for key, value in resp.headers.items(): - if key.lower() not in _HOP_BY_HOP: + if key.lower() not in _HOP_BY_HOP and not ( + transform_models and key.lower() == "content-encoding" + ): self.send_header(key, value) self.end_headers() # Do not pass a fixed chunk size here. httpx accumulates bytes until @@ -308,7 +397,12 @@ def _relay_response( # With ``chunk_size=None`` (the default), raw upstream chunks are # yielded as they arrive and pings keep the downstream connection # alive even before the model produces a large content block. - for chunk in resp.iter_raw(): + response_chunks = ( + [self.model_aliases.advertise_models(resp.read())] + if transform_models and 200 <= resp.status_code < 300 + else resp.iter_raw() + ) + for chunk in response_chunks: if chunk: if first_byte_ms is None: first_byte_ms = round((time.monotonic() - started) * 1000) @@ -388,11 +482,17 @@ def start_proxy( # to the gateway instead of a fresh handshake per request. Don't follow # redirects — a proxy relays 3xx verbatim. client = httpx.Client(base_url=upstream_base, timeout=_UPSTREAM_TIMEOUT, follow_redirects=False) + model_aliases = _ModelAliases() handler = type( "BoundProxyHandler", (_ProxyHandler,), - {"cache": cache, "client": client, "token_header": token_header}, + { + "cache": cache, + "client": client, + "token_header": token_header, + "model_aliases": model_aliases, + }, ) try: server = ThreadingHTTPServer(("127.0.0.1", port), handler) diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py index a9851747..3af743d0 100644 --- a/tests/test_gateway_proxy.py +++ b/tests/test_gateway_proxy.py @@ -352,6 +352,7 @@ def _handle_handler(client, cache, wfile) -> gateway_proxy._ProxyHandler: h = object.__new__(gateway_proxy._ProxyHandler) h.client = client h.cache = cache + h.model_aliases = gateway_proxy._ModelAliases() h.headers = {} h.rfile = io.BytesIO(b"") h.path = "/v1/messages" @@ -363,6 +364,74 @@ def _handle_handler(client, cache, wfile) -> gateway_proxy._ProxyHandler: return h +class TestModelAliases: + def test_advertises_custom_models_without_changing_display_name(self): + aliases = gateway_proxy._ModelAliases() + body = json.dumps( + { + "data": [ + {"id": "catalog.schema.custom", "display_name": "Custom model"}, + {"id": "system.ai.claude-sonnet"}, + {"id": "catalog.schema.anthropic-provider"}, + ], + "first_id": "catalog.schema.custom", + "last_id": "catalog.schema.anthropic-provider", + } + ).encode() + + payload = json.loads(aliases.advertise_models(body)) + + assert payload == { + "data": [ + { + "id": "anthropic-aigw-catalog.schema.custom", + "display_name": "Custom model", + }, + {"id": "system.ai.claude-sonnet"}, + {"id": "catalog.schema.anthropic-provider"}, + ], + "first_id": "anthropic-aigw-catalog.schema.custom", + "last_id": "catalog.schema.anthropic-provider", + } + + def test_rewrites_known_alias_in_messages_body(self): + aliases = gateway_proxy._ModelAliases() + aliases.advertise_models(b'{"data":[{"id":"catalog.schema.custom"}]}') + + body = aliases.rewrite_body( + "/v1/messages", b'{"model":"anthropic-aigw-catalog.schema.custom","messages":[]}' + ) + + assert json.loads(body) == {"model": "catalog.schema.custom", "messages": []} + + def test_rewrites_known_alias_in_pagination_cursor(self): + aliases = gateway_proxy._ModelAliases() + aliases.advertise_models(b'{"data":[{"id":"catalog.schema.custom"}]}') + + assert ( + aliases.rewrite_path( + "/v1/models?limit=1000&after_id=anthropic-aigw-catalog.schema.custom" + ) + == "/v1/models?limit=1000&after_id=catalog.schema.custom" + ) + + def test_does_not_strip_unknown_prefixed_id(self): + aliases = gateway_proxy._ModelAliases() + unknown = "anthropic-aigw-legitimate-upstream-id" + + assert aliases.rewrite_path(f"/v1/models?after_id={unknown}") == ( + f"/v1/models?after_id={unknown}" + ) + assert ( + aliases.rewrite_body("/v1/messages", json.dumps({"model": unknown}).encode()) + == json.dumps({"model": unknown}).encode() + ) + + def test_leaves_malformed_discovery_response_unchanged(self): + aliases = gateway_proxy._ModelAliases() + assert aliases.advertise_models(b"not-json") == b"not-json" + + class _Collect(io.RawIOBase): def __init__(self): self.data = bytearray()