Claude: separate model discovery proxy

src/ucode/agents/claude.py

@@ -16,6 +16,9 @@
 from typing import cast
 
 from ucode.agent_updates import available_npm_package_update
+from ucode.anthropic_model_discovery_proxy import (
+    start_proxy as start_anthropic_model_discovery_proxy,
+)
 from ucode.config_io import (
     APP_DIR,
     ToolSpec,
@@ -24,12 +27,16 @@
     read_json_safe,
     write_json_file,
 )
+from ucode.constants import LOOPBACK_HOST
 from ucode.databricks import (
     build_auth_shell_command,
     build_tool_base_url,
     get_databricks_token,
 )
-from ucode.gateway_proxy import AI_GATEWAY_TOKEN_HEADER, AUTHORIZATION_HEADER, start_proxy
+from ucode.gateway_proxy import (
+    AI_GATEWAY_TOKEN_HEADER,
+    AUTHORIZATION_HEADER,
+)
 from ucode.launcher import exec_or_spawn
 from ucode.managed_files import OS, current_os, write_managed_file
 from ucode.smart_routing.claude_hooks import (
@@ -214,10 +221,10 @@ def relayed_proxy_base_url(state: dict) -> str:
     port = state.get("relayed_proxy_port")
     if not isinstance(port, int):
         with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-            sock.bind(("127.0.0.1", 0))
+            sock.bind((LOOPBACK_HOST, 0))
             port = sock.getsockname()[1]
         state["relayed_proxy_port"] = port
-    return f"http://127.0.0.1:{port}"
+    return f"http://{LOOPBACK_HOST}:{port}"
 
 
 def _web_search_mcp_entry(workspace: str, search_model: str, profile: str | None = None) -> dict:
@@ -966,7 +973,7 @@ def _rewrite_relayed_port(state: dict, port: int) -> None:
     settings = read_json_safe(CLAUDE_SETTINGS_PATH)
     env = settings.get("env")
     if isinstance(env, dict):
-        env["ANTHROPIC_BASE_URL"] = f"http://127.0.0.1:{port}"
+        env["ANTHROPIC_BASE_URL"] = f"http://{LOOPBACK_HOST}:{port}"
         write_json_file(CLAUDE_SETTINGS_PATH, settings)
 
 
@@ -999,7 +1006,7 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None:
     if not isinstance(port, int):
         raise RuntimeError("Relayed proxy port was not configured; re-run `ucode claude`.")
 
-    server, cache, client = start_proxy(
+    server, cache, client = start_anthropic_model_discovery_proxy(
         workspace,
         state.get("profile"),
         port,
@@ -1031,7 +1038,7 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None:
 
 def _launch_gateway(state: dict, binary: str, tool_args: list[str]) -> None:
     workspace = state["workspace"]
-    server, cache, client = start_proxy(
+    server, cache, client = start_anthropic_model_discovery_proxy(
         workspace,
         state.get("profile"),
         0,
@@ -1041,7 +1048,7 @@ def _launch_gateway(state: dict, binary: str, tool_args: list[str]) -> None:
     token = cache.token
     os.environ["OAUTH_TOKEN"] = token
     os.environ["ANTHROPIC_AUTH_TOKEN"] = token
-    os.environ["ANTHROPIC_BASE_URL"] = f"http://127.0.0.1:{server.server_address[1]}"
+    os.environ["ANTHROPIC_BASE_URL"] = f"http://{LOOPBACK_HOST}:{server.server_address[1]}"
     os.environ["CLAUDE_CODE_USE_GATEWAY"] = "1"
 
     server_thread = threading.Thread(target=server.serve_forever, daemon=True)

src/ucode/anthropic_model_discovery_proxy.py

@@ -0,0 +1,124 @@
+"""Anthropic model discovery transformations for the gateway proxy."""
+
+from __future__ import annotations
+
+import json
+import threading
+from http import HTTPStatus
+from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
+
+import httpx
+
+from ucode import gateway_proxy
+
+_MODEL_ALIAS_PREFIX = "anthropic-aigw-"
+_ANTHROPIC_MODELS_PATH = "/v1/models"
+_ANTHROPIC_MESSAGES_PATH = "/v1/messages"
+
+
+class _AnthropicModelAliases:
+    """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 prefix_model_ids(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"]
+            lowered = model_id.lower()
+            if "claude" in lowered or "anthropic" in lowered:
+                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 != _ANTHROPIC_MODELS_PATH:
+            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 != _ANTHROPIC_MESSAGES_PATH 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 _AnthropicModelDiscoveryHandler(gateway_proxy._ProxyHandler):
+    anthropic_model_aliases: _AnthropicModelAliases
+
+    def _transform_request(self, body: bytes | None) -> tuple[str, bytes | None]:
+        body = self.anthropic_model_aliases.rewrite_body(self.path, body)
+        url = self.anthropic_model_aliases.rewrite_path(self.path).lstrip("/")
+        return url, body
+
+    def _transform_response(self, resp: httpx.Response) -> bytes | None:
+        should_prefix_model_ids = (
+            self.command == "GET"
+            and urlsplit(self.path).path == _ANTHROPIC_MODELS_PATH
+            and HTTPStatus.OK <= resp.status_code < HTTPStatus.MULTIPLE_CHOICES
+        )
+        if not should_prefix_model_ids:
+            return None
+        return self.anthropic_model_aliases.prefix_model_ids(resp.read())
+
+
+def start_proxy(
+    workspace: str,
+    profile: str | None,
+    port: int,
+    token_header: str,
+    force_refresh_near_expiry: bool,
+):
+    return gateway_proxy._start_proxy(
+        workspace,
+        profile,
+        port,
+        token_header,
+        force_refresh_near_expiry,
+        handler_class=_AnthropicModelDiscoveryHandler,
+        handler_attributes={"anthropic_model_aliases": _AnthropicModelAliases()},
+    )

src/ucode/constants.py

@@ -0,0 +1,3 @@
+"""Shared UCode constants."""
+
+LOOPBACK_HOST = "127.0.0.1"

src/ucode/gateway_proxy.py

@@ -25,9 +25,11 @@
 import time
 import uuid
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from typing import cast
 
 import httpx
 
+from ucode.constants import LOOPBACK_HOST
 from ucode.databricks import get_databricks_token
 
 # Header we overwrite with the freshly-minted Databricks credential. Any
@@ -212,12 +214,18 @@ def _safe_send_error(self, code: int, message: str) -> None:
         except OSError:
             pass
 
+    def _transform_request(self, body: bytes | None) -> tuple[str, bytes | None]:
+        return self.path.lstrip("/"), body
+
+    def _transform_response(self, resp: httpx.Response) -> bytes | None:
+        return None
+
     def _handle(self) -> None:
         diagnostic_id = uuid.uuid4().hex[:12]
         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("/")
+        url, body = self._transform_request(body)
         _diagnostic_log(
             "request_start",
             request_id=diagnostic_id,
@@ -304,9 +312,15 @@ def _relay_response(
         bytes_relayed = 0
         first_byte_ms: int | None = None
         try:
+            transformed_body = self._transform_response(resp)
             self.send_response(resp.status_code)
             for key, value in resp.headers.items():
-                if key.lower() not in _HOP_BY_HOP:
+                header_name = key.lower()
+                drop_stale_content_encoding = (
+                    transformed_body is not None and header_name == "content-encoding"
+                )
+                # resp.read() decodes compression; rewritten JSON is uncompressed.
+                if header_name not in _HOP_BY_HOP and not drop_stale_content_encoding:
                     self.send_header(key, value)
             self.end_headers()
             # Do not pass a fixed chunk size here. httpx accumulates bytes until
@@ -315,7 +329,10 @@ 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 = (
+                [transformed_body] if transformed_body is not None 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)
@@ -368,12 +385,15 @@ def __getattr__(self, name: str):
         raise AttributeError(name)
 
 
-def start_proxy(
+def _start_proxy(
     workspace: str,
     profile: str | None,
     port: int,
     token_header: str,
     force_refresh_near_expiry: bool,
+    *,
+    handler_class: type[_ProxyHandler],
+    handler_attributes: dict[str, object],
 ) -> tuple[ThreadingHTTPServer, _TokenCache, httpx.Client]:
     """Start the loopback refresh proxy + its background token refresher.
 
@@ -395,19 +415,44 @@ 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)
-
-    handler = type(
-        "BoundProxyHandler",
-        (_ProxyHandler,),
-        {"cache": cache, "client": client, "token_header": token_header},
+    handler = cast(
+        type[BaseHTTPRequestHandler],
+        type(
+            "BoundProxyHandler",
+            (handler_class,),
+            {
+                "cache": cache,
+                "client": client,
+                "token_header": token_header,
+                **handler_attributes,
+            },
+        ),
     )
     try:
-        server = ThreadingHTTPServer(("127.0.0.1", port), handler)
+        server = ThreadingHTTPServer((LOOPBACK_HOST, port), handler)
     except OSError:
         # Cached port is occupied (stale proxy from a killed session). Port 0 lets
         # the OS pick any free port; the caller reconciles the base URL to it.
-        server = ThreadingHTTPServer(("127.0.0.1", 0), handler)
+        server = ThreadingHTTPServer((LOOPBACK_HOST, 0), handler)
 
     refresher = threading.Thread(target=cache.run_refresher, daemon=True)
     refresher.start()
     return server, cache, client
+
+
+def start_proxy(
+    workspace: str,
+    profile: str | None,
+    port: int,
+    token_header: str,
+    force_refresh_near_expiry: bool,
+) -> tuple[ThreadingHTTPServer, _TokenCache, httpx.Client]:
+    return _start_proxy(
+        workspace,
+        profile,
+        port,
+        token_header,
+        force_refresh_near_expiry,
+        handler_class=_ProxyHandler,
+        handler_attributes={},
+    )

tests/test_agent_claude.py

@@ -666,7 +666,7 @@ def test_default_launch_keeps_existing_auth_path(self, monkeypatch):
         assert os.environ["OAUTH_TOKEN"] == "token"
         assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"]]
 
-    def test_runs_through_refresh_proxy(self, monkeypatch):
+    def test_gateway_discovery_uses_anthropic_proxy(self, monkeypatch):
         calls: list[tuple] = []
 
         class Server:
@@ -713,7 +713,7 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir
         monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
         monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False)
         monkeypatch.delenv("CLAUDE_CODE_USE_GATEWAY", raising=False)
-        monkeypatch.setattr(claude, "start_proxy", start_proxy)
+        monkeypatch.setattr(claude, "start_anthropic_model_discovery_proxy", start_proxy)
         monkeypatch.setattr(claude.subprocess, "Popen", Process)
 
         with pytest.raises(SystemExit) as exc: