Retry transient Claude model discovery failures

src/ucode/anthropic_model_discovery_proxy.py

@@ -24,6 +24,7 @@
 import httpx
 
 from ucode.constants import LOOPBACK_HOST
+from ucode.databricks import _http_get_retry_delay
 from ucode.gateway_proxy import (
     AI_GATEWAY_TOKEN_HEADER,
     HOP_BY_HOP_HEADERS,
@@ -34,6 +35,10 @@
     log_token_refresh_failure,
 )
 
+# Claude Code abandons model discovery after roughly three seconds. One retry
+# leaves enough time for the normal one-second backoff and the upstream request.
+_ANTHROPIC_MODEL_DISCOVERY_MAX_RETRIES = 1
+
 
 class _ProxyHandler(BaseHTTPRequestHandler):
     # Set by the server factory.
@@ -58,6 +63,48 @@ def _transform_request(self, body: bytes | None) -> tuple[str, bytes | None]:
     def _response_chunks(self, resp: httpx.Response) -> tuple[Iterable[bytes], frozenset[str]]:
         return resp.iter_raw(), frozenset()
 
+    def _should_retry_model_discovery(self, resp: httpx.Response) -> bool:
+        return (
+            self.command == "GET"
+            and urlsplit(self.path).path == _ANTHROPIC_MODELS_PATH
+            and resp.status_code == HTTPStatus.TOO_MANY_REQUESTS
+        )
+
+    def _retry_model_discovery(
+        self,
+        url: str,
+        body: bytes | None,
+        diagnostic_id: str,
+        started: float,
+        retry_after: str | None,
+    ) -> None:
+        for retry_index in range(_ANTHROPIC_MODEL_DISCOVERY_MAX_RETRIES):
+            delay = _http_get_retry_delay(retry_after, retry_index)
+            log_proxy_diagnostic(
+                "model_discovery_retry_scheduled",
+                request_id=diagnostic_id,
+                attempt=retry_index + 2,
+                delay_ms=round(delay * 1000),
+            )
+            time.sleep(delay)
+            headers = forwarded_request_headers(self, self.cache.token, self.token_header)
+            with self.client.stream(self.command, url, headers=headers, content=body) as resp:
+                log_proxy_diagnostic(
+                    "model_discovery_upstream_headers",
+                    request_id=diagnostic_id,
+                    attempt=retry_index + 2,
+                    status=resp.status_code,
+                    elapsed_ms=round((time.monotonic() - started) * 1000),
+                )
+                if (
+                    not self._should_retry_model_discovery(resp)
+                    or retry_index == _ANTHROPIC_MODEL_DISCOVERY_MAX_RETRIES - 1
+                ):
+                    self._relay_response(resp, diagnostic_id=diagnostic_id, started=started)
+                    return
+                retry_after = resp.headers.get("Retry-After")
+                resp.read()
+
     def _handle(self) -> None:
         diagnostic_id = uuid.uuid4().hex[:12]
         started = time.monotonic()
@@ -81,6 +128,17 @@ def _handle(self) -> None:
                     status=resp.status_code,
                     elapsed_ms=round((time.monotonic() - started) * 1000),
                 )
+                if self._should_retry_model_discovery(resp):
+                    retry_after = resp.headers.get("Retry-After")
+                    resp.read()
+                    self._retry_model_discovery(
+                        url,
+                        body,
+                        diagnostic_id,
+                        started,
+                        retry_after,
+                    )
+                    return
                 if resp.status_code not in (401, 403):
                     self._relay_response(resp, diagnostic_id=diagnostic_id, started=started)
                     return

src/ucode/databricks.py

@@ -63,6 +63,11 @@
 # raced — so we retry rather than treat them as an expired session.
 _TOKEN_CACHE_LOCK_MARKERS = ("cache update", "exit status 45")
 _TOKEN_FETCH_MAX_ATTEMPTS = 4
+_HTTP_GET_RETRYABLE_STATUS_CODES = frozenset({429})
+_HTTP_GET_RETRY_BASE_SECONDS = 1.0
+_HTTP_GET_RETRY_MAX_SECONDS = 5.0
+_HTTP_GET_RETRY_AFTER_JITTER_SECONDS = 0.25
+_ANTHROPIC_MODEL_DISCOVERY_SETUP_MAX_RETRIES = 2
 
 
 def _debug_enabled() -> bool:
@@ -214,52 +219,100 @@ def _log_auth_diagnostics() -> None:
         _debug(f"databrickscfg ({cfg_path})", f"read error: {exc}")
 
 
+def _http_get_retry_delay(retry_after: str | None, retry_index: int) -> float:
+    if retry_after is not None:
+        try:
+            retry_after_seconds = float(retry_after)
+        except ValueError:
+            pass
+        else:
+            if retry_after_seconds >= 0:
+                return min(retry_after_seconds, _HTTP_GET_RETRY_MAX_SECONDS) + random.uniform(
+                    0, _HTTP_GET_RETRY_AFTER_JITTER_SECONDS
+                )
+
+    backoff = min(
+        _HTTP_GET_RETRY_BASE_SECONDS * (2 ** min(retry_index, 10)),
+        _HTTP_GET_RETRY_MAX_SECONDS,
+    )
+    return backoff + random.uniform(0, min(backoff * 0.25, 0.5))
+
+
 def _http_get_json(
-    url: str, token: str, *, timeout: int = 10
+    url: str,
+    token: str,
+    *,
+    timeout: int = 10,
+    max_retries: int = 0,
 ) -> tuple[dict | list | None, str | None]:
     """GET a JSON endpoint. Returns (payload, None) on success, (None, reason) on failure.
 
+    ``max_retries`` opts individual callers into bounded retries for rate limits
+    and network failures. Other callers retain the original single-attempt
+    behavior.
+
     Honors UCODE_DEBUG=1 to append status + truncated body to ~/.ucode/debug.log.
     """
+    if max_retries < 0:
+        raise ValueError("max_retries must be non-negative")
+
     request = urllib_request.Request(
         url,
         headers={"Authorization": f"Bearer {token}", "Accept": "application/json"},
     )
-    try:
-        with urllib_request.urlopen(request, timeout=timeout) as response:
-            body = response.read().decode("utf-8")
-        _debug(f"GET {url}", f"HTTP 200, {len(body)} bytes")
-        if _debug_enabled():
-            _debug("body", body[:4000])
+    for attempt in range(max_retries + 1):
         try:
-            return json.loads(body), None
-        except json.JSONDecodeError as exc:
-            return None, f"response was not valid JSON ({exc.msg})"
-    except urllib_error.HTTPError as exc:
-        body = ""
-        try:
-            body = exc.read().decode("utf-8", errors="replace") if exc.fp else ""
-        except Exception:
+            with urllib_request.urlopen(request, timeout=timeout) as response:
+                body = response.read().decode("utf-8")
+            _debug(f"GET {url}", f"HTTP 200, {len(body)} bytes")
+            if _debug_enabled():
+                _debug("body", body[:4000])
+            try:
+                return json.loads(body), None
+            except json.JSONDecodeError as exc:
+                return None, f"response was not valid JSON ({exc.msg})"
+        except urllib_error.HTTPError as exc:
             body = ""
-        _debug(f"GET {url}", f"HTTP {exc.code} {exc.reason}")
-        if _debug_enabled() and body:
-            _debug("body", body[:4000])
-        reason = f"HTTP {exc.code} {exc.reason}"
-        # Surface the response body too — gateway auth failures return 400
-        # with body `Invalid Token`, which is invisible without this.
-        body_excerpt = body.strip()[:200]
-        if body_excerpt:
-            reason = f"{reason}: {body_excerpt}"
-        return None, reason
-    except urllib_error.URLError as exc:
-        _debug(f"GET {url}", f"URLError: {exc.reason}")
-        return None, f"network error: {exc.reason}"
-    except OSError as exc:
-        # A socket read timeout raises a bare TimeoutError (an OSError), not a
-        # URLError, so it must be caught explicitly or it escapes the whole
-        # discovery flow. Surface it as a reason like every other failure.
-        _debug(f"GET {url}", f"OSError: {exc}")
-        return None, f"network error: {exc}"
+            try:
+                body = exc.read().decode("utf-8", errors="replace") if exc.fp else ""
+            except Exception:
+                body = ""
+            _debug(f"GET {url}", f"HTTP {exc.code} {exc.reason}")
+            if _debug_enabled() and body:
+                _debug("body", body[:4000])
+            reason = f"HTTP {exc.code} {exc.reason}"
+            # Surface the response body too — gateway auth failures return 400
+            # with body `Invalid Token`, which is invisible without this.
+            body_excerpt = body.strip()[:200]
+            if body_excerpt:
+                reason = f"{reason}: {body_excerpt}"
+            if exc.code not in _HTTP_GET_RETRYABLE_STATUS_CODES or attempt == max_retries:
+                return None, reason
+            retry_after = exc.headers.get("Retry-After") if exc.headers is not None else None
+        except urllib_error.URLError as exc:
+            _debug(f"GET {url}", f"URLError: {exc.reason}")
+            reason = f"network error: {exc.reason}"
+            if attempt == max_retries:
+                return None, reason
+            retry_after = None
+        except OSError as exc:
+            # A socket read timeout raises a bare TimeoutError (an OSError), not a
+            # URLError, so it must be caught explicitly or it escapes the whole
+            # discovery flow. Surface it as a reason like every other failure.
+            _debug(f"GET {url}", f"OSError: {exc}")
+            reason = f"network error: {exc}"
+            if attempt == max_retries:
+                return None, reason
+            retry_after = None
+
+        delay = _http_get_retry_delay(retry_after, attempt)
+        _debug(
+            f"GET {url}",
+            f"attempt {attempt + 1} failed: {reason}; retrying in {delay:.2f}s",
+        )
+        time.sleep(delay)
+
+    raise AssertionError("unreachable")
 
 
 def _http_send_json(
@@ -2731,15 +2784,23 @@ def collect_services(result, _ref):
     return sorted(names), None
 
 
+def _get_anthropic_models_json(workspace: str, token: str) -> tuple[dict | list | None, str | None]:
+    hostname = workspace_hostname(workspace)
+    return _http_get_json(
+        f"https://{hostname}{ANTHROPIC_MODELS_PATH}",
+        token,
+        max_retries=_ANTHROPIC_MODEL_DISCOVERY_SETUP_MAX_RETRIES,
+    )
+
+
 def list_anthropic_models(workspace: str, token: str) -> tuple[list[str], str | None]:
     """List every model id advertised by AI Gateway's Anthropic endpoint.
 
     Claude Code's native gateway discovery consumes this same catalog, so callers
     using that mode must not apply ucode's legacy ``databricks-claude-*`` family
     validation.
     """
-    hostname = workspace_hostname(workspace)
-    payload, reason = _http_get_json(f"https://{hostname}{ANTHROPIC_MODELS_PATH}", token)
+    payload, reason = _get_anthropic_models_json(workspace, token)
     if payload is None:
         return [], reason
 
@@ -2765,8 +2826,7 @@ def discover_claude_models(workspace: str, token: str) -> tuple[dict[str, str],
     describes why the dict is empty (HTTP error, network error, or no models
     matching the expected naming convention).
     """
-    hostname = workspace_hostname(workspace)
-    payload, reason = _http_get_json(f"https://{hostname}{ANTHROPIC_MODELS_PATH}", token)
+    payload, reason = _get_anthropic_models_json(workspace, token)
     if payload is None:
         return {}, reason
 

tests/test_anthropic_model_discovery_proxy.py

@@ -33,12 +33,14 @@ def __exit__(self, *_args):
 
 class _FakeClient:
     def __init__(self, response):
-        self.response = response
+        self.responses = list(response) if isinstance(response, list) else [response]
         self.request = None
+        self.requests = []
 
     def stream(self, method, url, headers, content):
         self.request = (method, url, headers, content)
-        return self.response
+        self.requests.append(self.request)
+        return self.responses.pop(0)
 
 
 class _FakeCache:
@@ -168,6 +170,42 @@ 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_retries_rate_limited_model_discovery(self, monkeypatch):
+        out = _Collect()
+        handler = _handler(out)
+        handler.headers = {"Authorization": "Bearer subscription-token"}
+        handler.rfile = io.BytesIO()
+        handler.cache = _FakeCache()
+        rate_limited = _FakeResponse(429, {"Retry-After": "0"}, b"rate limited")
+        success = _FakeResponse(200, {}, b'{"data":[{"id":"custom-model"}]}')
+        handler.client = _FakeClient([rate_limited, success])
+        monkeypatch.setattr(anthropic_model_discovery_proxy.time, "sleep", lambda _delay: None)
+
+        handler._handle()
+
+        assert len(handler.client.requests) == 2
+        assert rate_limited.read_calls == 1
+        assert b"429 Too Many Requests" not in bytes(out.data)
+        assert b"anthropic-aigw-custom-model" in bytes(out.data)
+
+    def test_relays_rate_limit_after_model_discovery_retries_are_exhausted(self, monkeypatch):
+        out = _Collect()
+        handler = _handler(out)
+        handler.headers = {"Authorization": "Bearer subscription-token"}
+        handler.rfile = io.BytesIO()
+        handler.cache = _FakeCache()
+        responses = [_FakeResponse(429, {"Retry-After": "0"}, b"rate limited") for _ in range(2)]
+        handler.client = _FakeClient(responses)
+        monkeypatch.setattr(anthropic_model_discovery_proxy.time, "sleep", lambda _delay: None)
+
+        handler._handle()
+
+        assert len(handler.client.requests) == 2
+        assert responses[0].read_calls == 1
+        assert responses[1].read_calls == 0
+        assert b"429 Too Many Requests" in bytes(out.data)
+        assert b"rate limited" in bytes(out.data)
+
     def test_prefixes_successful_model_response_and_drops_content_encoding(self):
         out = _Collect()
         handler = _handler(out)