Cache Claude gateway model discovery

Author: andy-xu-db

src/ucode/agents/claude.py

@@ -1044,6 +1044,7 @@ def _launch_gateway(state: dict, binary: str, tool_args: list[str]) -> None:
         0,
         token_header=AUTHORIZATION_HEADER,
         force_refresh_near_expiry=True,
+        prefetch_models=True,
     )
     token = cache.token
     os.environ["OAUTH_TOKEN"] = token

src/ucode/anthropic_model_discovery_proxy.py

@@ -3,7 +3,9 @@
 from __future__ import annotations
 
 import json
+import sys
 import threading
+import time
 from http import HTTPStatus
 from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
 
@@ -14,6 +16,8 @@
 _MODEL_ALIAS_PREFIX = "anthropic-aigw-"
 _ANTHROPIC_MODELS_PATH = "/v1/models"
 _ANTHROPIC_MESSAGES_PATH = "/v1/messages"
+_MODEL_DISCOVERY_LIMIT = 1000
+_MODEL_CACHE_REFRESH_S = 600
 
 
 class _AnthropicModelAliases:
@@ -87,8 +91,82 @@ def rewrite_body(self, path: str, body: bytes | None) -> bytes | None:
         return json.dumps(payload, separators=(",", ":")).encode()
 
 
+class _ModelCache:
+    """Caches the complete model list so Claude's discovery request is local."""
+
+    def __init__(self, aliases: _AnthropicModelAliases) -> None:
+        self._aliases = aliases
+        self._body: bytes | None = None
+        self._lock = threading.Lock()
+
+    def refresh(self, client: httpx.Client, token: str, token_header: str) -> None:
+        headers = {
+            token_header: f"Bearer {token}",
+            "anthropic-version": "2023-06-01",
+        }
+        models: list[object] = []
+        first_page: dict[str, object] | None = None
+        last_page: dict[str, object] | None = None
+        after_id: str | None = None
+        seen_cursors: set[str] = set()
+
+        while True:
+            params: dict[str, str | int] = {"limit": _MODEL_DISCOVERY_LIMIT}
+            if after_id is not None:
+                params["after_id"] = after_id
+            response = client.get("v1/models", headers=headers, params=params)
+            response.raise_for_status()
+            payload = response.json()
+            if not isinstance(payload, dict) or not isinstance(payload.get("data"), list):
+                raise ValueError("invalid model discovery response")
+            if first_page is None:
+                first_page = payload
+            last_page = payload
+            models.extend(payload["data"])
+            if not payload.get("has_more"):
+                break
+            after_id = payload.get("last_id")
+            if not isinstance(after_id, str) or after_id in seen_cursors:
+                raise ValueError("invalid model discovery cursor")
+            seen_cursors.add(after_id)
+
+        combined = dict(first_page or {})
+        combined["data"] = models
+        combined["has_more"] = False
+        if last_page is not None:
+            combined["last_id"] = last_page.get("last_id")
+        body = self._aliases.prefix_model_ids(json.dumps(combined, separators=(",", ":")).encode())
+        with self._lock:
+            self._body = body
+
+    def get(self, method: str, path: str) -> bytes | None:
+        parsed = urlsplit(path)
+        if method != "GET" or parsed.path != _ANTHROPIC_MODELS_PATH:
+            return None
+        if any(
+            key in {"after_id", "before_id"}
+            for key, _value in parse_qsl(parsed.query, keep_blank_values=True)
+        ):
+            return None
+        with self._lock:
+            return self._body
+
+    def run_refresher(
+        self,
+        client: httpx.Client,
+        token_cache: gateway_proxy._TokenCache,
+        token_header: str,
+    ) -> None:
+        while not token_cache.wait_until_stopped(_MODEL_CACHE_REFRESH_S):
+            try:
+                self.refresh(client, token_cache.token, token_header)
+            except Exception:  # noqa: BLE001 - refresh failure must not kill the thread
+                continue
+
+
 class _AnthropicModelDiscoveryHandler(gateway_proxy._ProxyHandler):
     anthropic_model_aliases: _AnthropicModelAliases
+    model_cache: _ModelCache
 
     def _transform_request(self, body: bytes | None) -> tuple[str, bytes | None]:
         body = self.anthropic_model_aliases.rewrite_body(self.path, body)
@@ -105,20 +183,62 @@ def _transform_response(self, resp: httpx.Response) -> bytes | None:
             return None
         return self.anthropic_model_aliases.prefix_model_ids(resp.read())
 
+    def _handle_cached_response(self, diagnostic_id: str, started: float) -> bool:
+        cached_models = self.model_cache.get(self.command, self.path)
+        if cached_models is None:
+            return False
+        try:
+            self.send_response(HTTPStatus.OK)
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(cached_models)))
+            self.end_headers()
+            self.wfile.write(cached_models)
+            self.wfile.flush()
+            gateway_proxy._diagnostic_log(
+                "model_cache_hit",
+                request_id=diagnostic_id,
+                bytes=len(cached_models),
+                elapsed_ms=round((time.monotonic() - started) * 1000),
+            )
+        except (BrokenPipeError, ConnectionResetError):
+            pass
+        return True
+
 
 def start_proxy(
     workspace: str,
     profile: str | None,
     port: int,
     token_header: str,
     force_refresh_near_expiry: bool,
+    prefetch_models: bool = False,
 ):
-    return gateway_proxy._start_proxy(
+    aliases = _AnthropicModelAliases()
+    model_cache = _ModelCache(aliases)
+    server, token_cache, client = gateway_proxy._start_proxy(
         workspace,
         profile,
         port,
         token_header,
         force_refresh_near_expiry,
         handler_class=_AnthropicModelDiscoveryHandler,
-        handler_attributes={"anthropic_model_aliases": _AnthropicModelAliases()},
+        handler_attributes={
+            "anthropic_model_aliases": aliases,
+            "model_cache": model_cache,
+        },
     )
+    if prefetch_models:
+        try:
+            model_cache.refresh(client, token_cache.token, token_header)
+        except (httpx.HTTPError, TypeError, ValueError) as exc:
+            sys.stderr.write(
+                "[ucode] Claude model prefetch failed "
+                f"({type(exc).__name__}); falling back to live discovery.\n"
+            )
+        model_refresher = threading.Thread(
+            target=model_cache.run_refresher,
+            args=(client, token_cache, token_header),
+            daemon=True,
+        )
+        model_refresher.start()
+    return server, token_cache, client

src/ucode/gateway_proxy.py

@@ -183,6 +183,9 @@ def run_refresher(self) -> None:
     def stop(self) -> None:
         self._stop.set()
 
+    def wait_until_stopped(self, timeout: float) -> bool:
+        return self._stop.wait(timeout)
+
 
 def _forwarded_request_headers(
     handler: BaseHTTPRequestHandler,
@@ -220,6 +223,9 @@ def _transform_request(self, body: bytes | None) -> tuple[str, bytes | None]:
     def _transform_response(self, resp: httpx.Response) -> bytes | None:
         return None
 
+    def _handle_cached_response(self, diagnostic_id: str, started: float) -> bool:
+        return False
+
     def _handle(self) -> None:
         diagnostic_id = uuid.uuid4().hex[:12]
         started = time.monotonic()
@@ -232,6 +238,8 @@ def _handle(self) -> None:
             method=self.command,
             path=self.path.split("?", 1)[0],
         )
+        if self._handle_cached_response(diagnostic_id, started):
+            return
         try:
             # First attempt with the current token.
             headers = _forwarded_request_headers(self, self.cache.token, self.token_header)

tests/test_agent_claude.py

@@ -695,7 +695,14 @@ def __init__(self, argv):
             def wait(self):
                 return 0
 
-        def start_proxy(workspace, profile, port, token_header, force_refresh_near_expiry):
+        def start_proxy(
+            workspace,
+            profile,
+            port,
+            token_header,
+            force_refresh_near_expiry,
+            prefetch_models,
+        ):
             calls.append(
                 (
                     "proxy",
@@ -704,6 +711,7 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir
                     port,
                     token_header,
                     force_refresh_near_expiry,
+                    prefetch_models,
                 )
             )
             return Server(), Cache(), Client()
@@ -725,7 +733,7 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir
         assert os.environ["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:12345"
         assert os.environ["CLAUDE_CODE_USE_GATEWAY"] == "1"
         assert calls[:2] == [
-            ("proxy", WS, "test", 0, claude.AUTHORIZATION_HEADER, True),
+            ("proxy", WS, "test", 0, claude.AUTHORIZATION_HEADER, True, True),
             ("serve",),
         ]
         assert calls[2][0] == "popen"

tests/test_anthropic_model_discovery_proxy.py

@@ -65,6 +65,9 @@ def _handler(wfile, path="/v1/models", command="GET"):
     handler.path = path
     handler._headers_buffer = []
     handler.anthropic_model_aliases = anthropic_model_discovery_proxy._AnthropicModelAliases()
+    handler.model_cache = anthropic_model_discovery_proxy._ModelCache(
+        handler.anthropic_model_aliases
+    )
     return handler
 
 
@@ -147,6 +150,83 @@ def test_leaves_malformed_discovery_response_unchanged(self):
         assert aliases.prefix_model_ids(b"not-json") == b"not-json"
 
 
+class _DiscoveryResponse:
+    def __init__(self, payload: dict):
+        self._payload = payload
+
+    def raise_for_status(self):
+        return None
+
+    def json(self):
+        return self._payload
+
+
+class _DiscoveryClient:
+    def __init__(self, pages: list[dict]):
+        self._pages = iter(pages)
+        self.requests: list[tuple[str, dict[str, str], dict[str, str | int]]] = []
+
+    def get(self, path, *, headers, params):
+        self.requests.append((path, headers, params))
+        return _DiscoveryResponse(next(self._pages))
+
+
+class TestModelCache:
+    def test_prefetches_all_pages_and_serves_one_aliased_response(self):
+        aliases = anthropic_model_discovery_proxy._AnthropicModelAliases()
+        cache = anthropic_model_discovery_proxy._ModelCache(aliases)
+        client = _DiscoveryClient(
+            [
+                {
+                    "data": [{"id": "catalog.schema.one"}],
+                    "has_more": True,
+                    "first_id": "catalog.schema.one",
+                    "last_id": "catalog.schema.one",
+                },
+                {
+                    "data": [{"id": "system.ai.claude-sonnet"}],
+                    "has_more": False,
+                    "first_id": "system.ai.claude-sonnet",
+                    "last_id": "system.ai.claude-sonnet",
+                },
+            ]
+        )
+
+        cache.refresh(
+            client,
+            "token",
+            anthropic_model_discovery_proxy.gateway_proxy.AUTHORIZATION_HEADER,
+        )
+
+        payload = json.loads(cache.get("GET", "/v1/models?limit=1000"))
+        assert [model["id"] for model in payload["data"]] == [
+            "anthropic-aigw-catalog.schema.one",
+            "system.ai.claude-sonnet",
+        ]
+        assert payload["has_more"] is False
+        assert payload["last_id"] == "system.ai.claude-sonnet"
+        assert [request[2] for request in client.requests] == [
+            {"limit": 1000},
+            {"limit": 1000, "after_id": "catalog.schema.one"},
+        ]
+
+    def test_only_serves_first_page_get_requests(self):
+        cache = anthropic_model_discovery_proxy._ModelCache(
+            anthropic_model_discovery_proxy._AnthropicModelAliases()
+        )
+        cache.refresh(
+            _DiscoveryClient([{"data": [], "has_more": False}]),
+            "token",
+            anthropic_model_discovery_proxy.gateway_proxy.AUTHORIZATION_HEADER,
+        )
+
+        assert cache.get("GET", "/v1/models") is not None
+        assert cache.get("POST", "/v1/models") is None
+        assert cache.get("GET", "/v1/models?after_id=cursor") is None
+        assert cache.get("GET", "/v1/models?before_id=cursor") is None
+        assert cache.get("GET", "/v1/messages") is None
+
+
 class TestAnthropicModelDiscoveryHandler:
     def test_inherits_relayed_auth_and_prefixes_models(self):
         out = _Collect()
@@ -164,6 +244,24 @@ def test_inherits_relayed_auth_and_prefixes_models(self):
         assert headers["X-Databricks-AI-Gateway-Token"] == "Bearer databricks-token"
         assert b"anthropic-aigw-custom-model" in bytes(out.data)
 
+    def test_serves_cached_models_without_upstream_request(self):
+        out = _Collect()
+        handler = _handler(out)
+        handler.headers = {}
+        handler.rfile = io.BytesIO()
+        handler.cache = _FakeCache()
+        handler.client = _FakeClient(None)
+        handler.model_cache.refresh(
+            _DiscoveryClient([{"data": [{"id": "system.ai.claude"}], "has_more": False}]),
+            "token",
+            anthropic_model_discovery_proxy.gateway_proxy.AUTHORIZATION_HEADER,
+        )
+
+        handler._handle()
+
+        assert handler.client.request is None
+        assert b"system.ai.claude" in bytes(out.data)
+
     def test_prefixes_successful_model_response_and_drops_content_encoding(self):
         out = _Collect()
         handler = _handler(out)
@@ -222,3 +320,7 @@ def start(*args, **kwargs):
         call["kwargs"]["handler_attributes"]["anthropic_model_aliases"],
         anthropic_model_discovery_proxy._AnthropicModelAliases,
     )
+    assert isinstance(
+        call["kwargs"]["handler_attributes"]["model_cache"],
+        anthropic_model_discovery_proxy._ModelCache,
+    )