Claude: support custom models in gateway discovery

src/ucode/gateway_proxy.py

@@ -26,6 +26,7 @@
 import time
 import uuid
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
 
 import httpx
 
@@ -70,6 +71,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:
@@ -189,10 +191,81 @@ def _forwarded_request_headers(handler: BaseHTTPRequestHandler, token: str) -> d
     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
+    model_aliases: _ModelAliases
 
     def log_message(self, format: str, *args: object) -> None:
         return
@@ -210,7 +283,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,
@@ -229,7 +303,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)
@@ -380,11 +474,12 @@ 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},
+        {"cache": cache, "client": client, "model_aliases": model_aliases},
     )
     try:
         server = ThreadingHTTPServer(("127.0.0.1", port), handler)

tests/test_gateway_proxy.py

@@ -341,6 +341,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"
@@ -352,6 +353,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()