diff --git a/src/ucode/anthropic_model_discovery_proxy.py b/src/ucode/anthropic_model_discovery_proxy.py index 989616f6..ada123bf 100644 --- a/src/ucode/anthropic_model_discovery_proxy.py +++ b/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 diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index fc7ab3bd..3350e99c 100644 --- a/src/ucode/databricks.py +++ b/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,6 +2784,15 @@ 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. @@ -2738,8 +2800,7 @@ def list_anthropic_models(workspace: str, token: str) -> tuple[list[str], str | 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 diff --git a/tests/test_anthropic_model_discovery_proxy.py b/tests/test_anthropic_model_discovery_proxy.py index d44c782a..0e6b93c1 100644 --- a/tests/test_anthropic_model_discovery_proxy.py +++ b/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) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index c88cfa8e..392e59da 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -173,8 +173,8 @@ def test_lists_all_anthropic_model_ids_without_legacy_validation(self, monkeypat ] } - def fake_get(url, token): - captured["request"] = (url, token) + def fake_get(url, token, **kwargs): + captured["request"] = (url, token, kwargs) return payload, None monkeypatch.setattr(db_mod, "_http_get_json", fake_get) @@ -190,6 +190,7 @@ def fake_get(url, token): assert captured["request"] == ( f"{WS}/ai-gateway/anthropic/v1/models", "token", + {"max_retries": 2}, ) def test_selects_opus_4_8_when_advertised(self, monkeypatch): @@ -200,7 +201,7 @@ def test_selects_opus_4_8_when_advertised(self, monkeypatch): {"id": "databricks-claude-sonnet-4-6"}, ] } - monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token, **kwargs: (payload, None)) models, reason = db_mod.discover_claude_models(WS, "token") @@ -214,7 +215,7 @@ def test_buckets_fable_family(self, monkeypatch): {"id": "databricks-claude-opus-4-8"}, ] } - monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token, **kwargs: (payload, None)) models, reason = db_mod.discover_claude_models(WS, "token") @@ -2020,6 +2021,103 @@ def test_reason_without_body_is_status_only(self): assert reason == "HTTP 404 Not Found" +class TestHttpGetJsonRetries: + @staticmethod + def _http_error(code: int, message: str, headers: dict[str, str] | None = None): + from urllib.error import HTTPError + + return HTTPError(url="", code=code, msg=message, hdrs=headers or {}, fp=None) + + def test_retries_429_after_retry_after_delay(self, monkeypatch): + outcomes = iter( + [ + self._http_error(429, "Too Many Requests", {"Retry-After": "1"}), + _FakeResponse({"data": []}), + ] + ) + calls = [] + sleeps = [] + + def fake_urlopen(request, timeout=None): + calls.append(request) + outcome = next(outcomes) + if isinstance(outcome, Exception): + raise outcome + return outcome + + monkeypatch.setattr(db_mod.urllib_request, "urlopen", fake_urlopen) + monkeypatch.setattr(db_mod.random, "uniform", lambda low, high: high) + monkeypatch.setattr(db_mod.time, "sleep", sleeps.append) + + payload, reason = db_mod._http_get_json("https://x/y", "tok", max_retries=2) + + assert payload == {"data": []} + assert reason is None + assert len(calls) == 2 + assert sleeps == [1.25] + + def test_retries_network_error_with_exponential_backoff(self, monkeypatch): + from urllib.error import URLError + + outcomes = iter([URLError("connection reset"), _FakeResponse({"ok": True})]) + sleeps = [] + + def fake_urlopen(request, timeout=None): + outcome = next(outcomes) + if isinstance(outcome, Exception): + raise outcome + return outcome + + monkeypatch.setattr(db_mod.urllib_request, "urlopen", fake_urlopen) + monkeypatch.setattr(db_mod.random, "uniform", lambda low, high: high) + monkeypatch.setattr(db_mod.time, "sleep", sleeps.append) + + payload, reason = db_mod._http_get_json("https://x/y", "tok", max_retries=2) + + assert payload == {"ok": True} + assert reason is None + assert sleeps == [1.25] + + def test_stops_after_configured_retries(self, monkeypatch): + calls = [] + sleeps = [] + + def fake_urlopen(request, timeout=None): + calls.append(request) + raise self._http_error(429, "Too Many Requests") + + monkeypatch.setattr(db_mod.urllib_request, "urlopen", fake_urlopen) + monkeypatch.setattr(db_mod.random, "uniform", lambda low, high: high) + monkeypatch.setattr(db_mod.time, "sleep", sleeps.append) + + payload, reason = db_mod._http_get_json("https://x/y", "tok", max_retries=2) + + assert payload is None + assert reason == "HTTP 429 Too Many Requests" + assert len(calls) == 3 + assert sleeps == [1.25, 2.5] + + def test_does_not_retry_other_http_error(self, monkeypatch): + calls = [] + + def fake_urlopen(request, timeout=None): + calls.append(request) + raise self._http_error(503, "Service Unavailable") + + monkeypatch.setattr(db_mod.urllib_request, "urlopen", fake_urlopen) + monkeypatch.setattr( + db_mod.time, + "sleep", + lambda delay: pytest.fail(f"unexpected retry delay: {delay}"), + ) + + payload, reason = db_mod._http_get_json("https://x/y", "tok", max_retries=2) + + assert payload is None + assert reason == "HTTP 503 Service Unavailable" + assert len(calls) == 1 + + class TestParseDatabricksCliVersion: def test_parses_standard_format(self): assert _parse_databricks_cli_version("Databricks CLI v0.299.2") == (0, 299, 2)