From 48b8bf6239c4b7f874fa825a9fb695a7608a1d01 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Mon, 24 Aug 2026 19:08:51 +0000 Subject: [PATCH 1/9] Add Claude gateway model aliases --- src/ucode/gateway_proxy.py | 112 ++++++++++++++++++++++++++++++++++-- tests/test_gateway_proxy.py | 69 ++++++++++++++++++++++ 2 files changed, 175 insertions(+), 6 deletions(-) diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index 9f78b6f..b8bcab0 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -25,6 +25,7 @@ import time import uuid from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit import httpx @@ -67,6 +68,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: @@ -195,11 +197,82 @@ def _forwarded_request_headers( 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 token_header = AI_GATEWAY_TOKEN_HEADER + model_aliases: _ModelAliases def log_message(self, format: str, *args: object) -> None: return @@ -217,7 +290,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, @@ -236,7 +310,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. @@ -265,7 +345,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( @@ -296,6 +382,7 @@ def _relay_response( self, resp: httpx.Response, *, + transform_models: bool = False, diagnostic_id: str | None = None, started: float | None = None, ) -> None: @@ -306,7 +393,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 @@ -315,7 +404,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) @@ -395,11 +489,17 @@ 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, "token_header": token_header}, + { + "cache": cache, + "client": client, + "token_header": token_header, + "model_aliases": model_aliases, + }, ) try: server = ThreadingHTTPServer(("127.0.0.1", port), handler) diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py index 3b51466..3856894 100644 --- a/tests/test_gateway_proxy.py +++ b/tests/test_gateway_proxy.py @@ -352,6 +352,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" @@ -363,6 +364,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() From b17a0a8373c9b1dc3f29a69ad8079304aed6b20a Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Tue, 25 Aug 2026 21:30:48 +0000 Subject: [PATCH 2/9] fix(claude): scope gateway model aliases --- src/ucode/gateway_proxy.py | 48 ++++++++++++++++++------------- tests/test_gateway_proxy.py | 57 +++++++++++++++++++++++-------------- 2 files changed, 64 insertions(+), 41 deletions(-) diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index b8bcab0..2fc0310 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -24,6 +24,7 @@ import threading import time import uuid +from http import HTTPStatus from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit @@ -69,6 +70,8 @@ _DIAGNOSTICS_ENV = "UCODE_RELAYED_PROXY_DIAGNOSTICS" _DIAGNOSTICS_TRUE = frozenset({"1", "true", "yes", "on"}) _MODEL_ALIAS_PREFIX = "anthropic-aigw-" +_ANTHROPIC_MODELS_PATH = "/anthropic/v1/models" +_ANTHROPIC_MESSAGES_PATH = "/anthropic/v1/messages" def _diagnostics_enabled() -> bool: @@ -197,7 +200,7 @@ def _forwarded_request_headers( return headers -class _ModelAliases: +class _AnthropicModelAliases: """Maps Claude-compatible discovery IDs back to their gateway model IDs.""" def __init__(self) -> None: @@ -218,7 +221,7 @@ def advertise_models(self, body: bytes) -> bytes: 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(): + if model_id.lower().startswith(("claude", "anthropic")): continue alias = f"{_MODEL_ALIAS_PREFIX}{model_id}" model["id"] = alias @@ -240,7 +243,7 @@ def original_id(self, model_id: str) -> str: def rewrite_path(self, path: str) -> str: parsed = urlsplit(path) - if parsed.path != "/v1/models": + if parsed.path != _ANTHROPIC_MODELS_PATH: return path query = [ (key, self.original_id(value) if key == "after_id" else value) @@ -251,7 +254,7 @@ def rewrite_path(self, path: str) -> str: ) def rewrite_body(self, path: str, body: bytes | None) -> bytes | None: - if urlsplit(path).path != "/v1/messages" or body is None: + if urlsplit(path).path != _ANTHROPIC_MESSAGES_PATH or body is None: return body try: payload = json.loads(body) @@ -272,7 +275,7 @@ class _ProxyHandler(BaseHTTPRequestHandler): cache: _TokenCache client: httpx.Client token_header = AI_GATEWAY_TOKEN_HEADER - model_aliases: _ModelAliases + anthropic_model_aliases: _AnthropicModelAliases def log_message(self, format: str, *args: object) -> None: return @@ -290,8 +293,12 @@ 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 - body = self.model_aliases.rewrite_body(self.path, body) - url = self.model_aliases.rewrite_path(self.path).lstrip("/") + upstream_path = f"/anthropic{self.path}" + body = self.anthropic_model_aliases.rewrite_body(upstream_path, body) + url = self.anthropic_model_aliases.rewrite_path(upstream_path).lstrip("/") + should_transform_model_names = ( + self.command == "GET" and urlsplit(upstream_path).path == _ANTHROPIC_MODELS_PATH + ) _diagnostic_log( "request_start", request_id=diagnostic_id, @@ -312,8 +319,7 @@ def _handle(self) -> None: if resp.status_code not in (401, 403): self._relay_response( resp, - transform_models=self.command == "GET" - and urlsplit(self.path).path == "/v1/models", + transform_model_names=should_transform_model_names, diagnostic_id=diagnostic_id, started=started, ) @@ -347,8 +353,7 @@ def _handle(self) -> None: ) self._relay_response( resp, - transform_models=self.command == "GET" - and urlsplit(self.path).path == "/v1/models", + transform_model_names=should_transform_model_names, diagnostic_id=diagnostic_id, started=started, ) @@ -382,7 +387,7 @@ def _relay_response( self, resp: httpx.Response, *, - transform_models: bool = False, + transform_model_names: bool, diagnostic_id: str | None = None, started: float | None = None, ) -> None: @@ -393,9 +398,11 @@ 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 and not ( - transform_models and key.lower() == "content-encoding" - ): + header_name = key.lower() + drop_stale_content_encoding = ( + transform_model_names and header_name == "content-encoding" + ) + 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 @@ -405,8 +412,9 @@ def _relay_response( # yielded as they arrive and pings keep the downstream connection # alive even before the model produces a large content block. response_chunks = ( - [self.model_aliases.advertise_models(resp.read())] - if transform_models and 200 <= resp.status_code < 300 + [self.anthropic_model_aliases.advertise_models(resp.read())] + if transform_model_names + and HTTPStatus.OK <= resp.status_code < HTTPStatus.MULTIPLE_CHOICES else resp.iter_raw() ) for chunk in response_chunks: @@ -479,7 +487,7 @@ def start_proxy( Returns (server, cache, client); the caller runs the server (e.g. in a thread) and calls shutdown()/cache.stop()/client.close() on exit. """ - upstream_base = f"{workspace.rstrip('/')}/ai-gateway/anthropic/" + upstream_base = f"{workspace.rstrip('/')}/ai-gateway/" cache = _TokenCache( workspace, profile, @@ -489,7 +497,7 @@ 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() + anthropic_model_aliases = _AnthropicModelAliases() handler = type( "BoundProxyHandler", @@ -498,7 +506,7 @@ def start_proxy( "cache": cache, "client": client, "token_header": token_header, - "model_aliases": model_aliases, + "anthropic_model_aliases": anthropic_model_aliases, }, ) try: diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py index 3856894..800b92f 100644 --- a/tests/test_gateway_proxy.py +++ b/tests/test_gateway_proxy.py @@ -105,7 +105,7 @@ def test_relay_swallows_broken_pipe_on_headers(self): handler = _relay_handler(_BrokenPipeWriter()) resp = _FakeResponse(200, {}, [b'{"ok":true}']) # Must not raise — a dead client is a routine teardown, not an error. - handler._relay_response(resp) + handler._relay_response(resp, transform_model_names=False) def test_relay_swallows_connection_reset_mid_stream(self): # Headers flush ok, then the client resets while streaming body chunks. @@ -123,7 +123,7 @@ def flush(self): handler = _relay_handler(_ResetAfterHeaders()) resp = _FakeResponse(200, {}, [b"chunk-of-sse-data"]) - handler._relay_response(resp) + handler._relay_response(resp, transform_model_names=False) def test_relay_swallows_upstream_error_mid_stream(self): # Upstream drops mid-body after headers are already sent — we can't signal @@ -141,7 +141,7 @@ def _chunks(): handler = _relay_handler(_Ok()) resp = _FakeResponse(200, {}, _chunks()) - handler._relay_response(resp) # must not raise + handler._relay_response(resp, transform_model_names=False) # must not raise def test_relay_forwards_status_and_skips_hop_by_hop_headers(self): # A non-200 status (e.g. 429 rate limit) and content headers are relayed; @@ -162,7 +162,7 @@ def flush(self): {"Content-Type": "application/json", "Transfer-Encoding": "chunked"}, [b'{"type":"error"}'], ) - handler._relay_response(resp) + handler._relay_response(resp, transform_model_names=False) blob = b"".join(chunks_written) assert b"429" in blob assert b"Content-Type: application/json" in blob @@ -186,6 +186,7 @@ def _chunks(): handler = _relay_handler(_Ok()) handler._relay_response( _FakeResponse(200, {}, _chunks()), + transform_model_names=False, diagnostic_id="local-id", started=time.monotonic(), ) @@ -202,7 +203,7 @@ def _chunks(): def test_diagnostics_are_silent_by_default(self, monkeypatch, capsys): monkeypatch.delenv(gateway_proxy._DIAGNOSTICS_ENV, raising=False) handler = _relay_handler(_Collect()) - handler._relay_response(_FakeResponse(200, {}, [b"ok"])) + handler._relay_response(_FakeResponse(200, {}, [b"ok"]), transform_model_names=False) assert capsys.readouterr().err == "" @@ -328,9 +329,11 @@ class _FakeClient: def __init__(self, responses): self._responses = list(responses) self.sent_tokens: list[str | None] = [] + self.sent_urls: list[str] = [] def stream(self, _method, _url, headers, content): self.sent_tokens.append(headers.get(gateway_proxy.AI_GATEWAY_TOKEN_HEADER)) + self.sent_urls.append(_url) return self._responses.pop(0) @@ -352,7 +355,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.anthropic_model_aliases = gateway_proxy._AnthropicModelAliases() h.headers = {} h.rfile = io.BytesIO(b"") h.path = "/v1/messages" @@ -364,15 +367,17 @@ def _handle_handler(client, cache, wfile) -> gateway_proxy._ProxyHandler: return h -class TestModelAliases: +class TestAnthropicModelAliases: def test_advertises_custom_models_without_changing_display_name(self): - aliases = gateway_proxy._ModelAliases() + aliases = gateway_proxy._AnthropicModelAliases() body = json.dumps( { "data": [ {"id": "catalog.schema.custom", "display_name": "Custom model"}, {"id": "system.ai.claude-sonnet"}, + {"id": "claude-sonnet"}, {"id": "catalog.schema.anthropic-provider"}, + {"id": "anthropic-provider"}, ], "first_id": "catalog.schema.custom", "last_id": "catalog.schema.anthropic-provider", @@ -387,48 +392,57 @@ def test_advertises_custom_models_without_changing_display_name(self): "id": "anthropic-aigw-catalog.schema.custom", "display_name": "Custom model", }, - {"id": "system.ai.claude-sonnet"}, - {"id": "catalog.schema.anthropic-provider"}, + {"id": "anthropic-aigw-system.ai.claude-sonnet"}, + {"id": "claude-sonnet"}, + {"id": "anthropic-aigw-catalog.schema.anthropic-provider"}, + {"id": "anthropic-provider"}, ], "first_id": "anthropic-aigw-catalog.schema.custom", - "last_id": "catalog.schema.anthropic-provider", + "last_id": "anthropic-aigw-catalog.schema.anthropic-provider", } def test_rewrites_known_alias_in_messages_body(self): - aliases = gateway_proxy._ModelAliases() + aliases = gateway_proxy._AnthropicModelAliases() aliases.advertise_models(b'{"data":[{"id":"catalog.schema.custom"}]}') body = aliases.rewrite_body( - "/v1/messages", b'{"model":"anthropic-aigw-catalog.schema.custom","messages":[]}' + "/anthropic/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 = gateway_proxy._AnthropicModelAliases() aliases.advertise_models(b'{"data":[{"id":"catalog.schema.custom"}]}') assert ( aliases.rewrite_path( - "/v1/models?limit=1000&after_id=anthropic-aigw-catalog.schema.custom" + "/anthropic/v1/models?limit=1000&after_id=anthropic-aigw-catalog.schema.custom" ) - == "/v1/models?limit=1000&after_id=catalog.schema.custom" + == "/anthropic/v1/models?limit=1000&after_id=catalog.schema.custom" ) + def test_ignores_non_anthropic_models_path(self): + aliases = gateway_proxy._AnthropicModelAliases() + path = "/codex/v1/models?after_id=anthropic-aigw-catalog.schema.custom" + + assert aliases.rewrite_path(path) == path + def test_does_not_strip_unknown_prefixed_id(self): - aliases = gateway_proxy._ModelAliases() + aliases = gateway_proxy._AnthropicModelAliases() 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_path(f"/anthropic/v1/models?after_id={unknown}") == ( + f"/anthropic/v1/models?after_id={unknown}" ) assert ( - aliases.rewrite_body("/v1/messages", json.dumps({"model": unknown}).encode()) + aliases.rewrite_body("/anthropic/v1/messages", json.dumps({"model": unknown}).encode()) == json.dumps({"model": unknown}).encode() ) def test_leaves_malformed_discovery_response_unchanged(self): - aliases = gateway_proxy._ModelAliases() + aliases = gateway_proxy._AnthropicModelAliases() assert aliases.advertise_models(b"not-json") == b"not-json" @@ -474,6 +488,7 @@ def test_success_first_try_does_not_refresh(self): _handle_handler(client, cache, out)._handle() assert cache.refreshed == 0 assert client.sent_tokens == ["Bearer tok1"] + assert client.sent_urls == ["anthropic/v1/messages"] assert b"hi" in bytes(out.data) From 68fe0c659b7618da2ffb466743f207fcdbc2becc Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Tue, 25 Aug 2026 21:46:12 +0000 Subject: [PATCH 3/9] fix(claude): match gateway discovery filter --- src/ucode/gateway_proxy.py | 3 ++- tests/test_gateway_proxy.py | 6 +++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index 2fc0310..79a7362 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -221,7 +221,8 @@ def advertise_models(self, body: bytes) -> bytes: if not isinstance(model, dict) or not isinstance(model.get("id"), str): continue model_id = model["id"] - if model_id.lower().startswith(("claude", "anthropic")): + lowered = model_id.lower() + if "claude" in lowered or "anthropic" in lowered: continue alias = f"{_MODEL_ALIAS_PREFIX}{model_id}" model["id"] = alias diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py index 800b92f..0aa431d 100644 --- a/tests/test_gateway_proxy.py +++ b/tests/test_gateway_proxy.py @@ -392,13 +392,13 @@ def test_advertises_custom_models_without_changing_display_name(self): "id": "anthropic-aigw-catalog.schema.custom", "display_name": "Custom model", }, - {"id": "anthropic-aigw-system.ai.claude-sonnet"}, + {"id": "system.ai.claude-sonnet"}, {"id": "claude-sonnet"}, - {"id": "anthropic-aigw-catalog.schema.anthropic-provider"}, + {"id": "catalog.schema.anthropic-provider"}, {"id": "anthropic-provider"}, ], "first_id": "anthropic-aigw-catalog.schema.custom", - "last_id": "anthropic-aigw-catalog.schema.anthropic-provider", + "last_id": "catalog.schema.anthropic-provider", } def test_rewrites_known_alias_in_messages_body(self): From d642760de3b89958b000a65e9e974eb558121bd1 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Wed, 26 Aug 2026 01:24:22 +0000 Subject: [PATCH 4/9] fix(claude): clarify Anthropic proxy scope --- src/ucode/agents/claude.py | 15 +++-- ...ay_proxy.py => anthropic_gateway_proxy.py} | 32 ++++++----- src/ucode/constants.py | 3 + ...oxy.py => test_anthropic_gateway_proxy.py} | 57 ++++++++++++++----- 4 files changed, 75 insertions(+), 32 deletions(-) rename src/ucode/{gateway_proxy.py => anthropic_gateway_proxy.py} (95%) create mode 100644 src/ucode/constants.py rename tests/{test_gateway_proxy.py => test_anthropic_gateway_proxy.py} (90%) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index a4b1435..81ce046 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -16,6 +16,11 @@ from typing import cast from ucode.agent_updates import available_npm_package_update +from ucode.anthropic_gateway_proxy import ( + AI_GATEWAY_TOKEN_HEADER, + AUTHORIZATION_HEADER, + start_proxy, +) from ucode.config_io import ( APP_DIR, ToolSpec, @@ -24,12 +29,12 @@ 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.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 +219,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 +971,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) @@ -1041,7 +1046,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) diff --git a/src/ucode/gateway_proxy.py b/src/ucode/anthropic_gateway_proxy.py similarity index 95% rename from src/ucode/gateway_proxy.py rename to src/ucode/anthropic_gateway_proxy.py index 79a7362..8a4eb74 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/anthropic_gateway_proxy.py @@ -30,6 +30,7 @@ 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 @@ -70,8 +71,8 @@ _DIAGNOSTICS_ENV = "UCODE_RELAYED_PROXY_DIAGNOSTICS" _DIAGNOSTICS_TRUE = frozenset({"1", "true", "yes", "on"}) _MODEL_ALIAS_PREFIX = "anthropic-aigw-" -_ANTHROPIC_MODELS_PATH = "/anthropic/v1/models" -_ANTHROPIC_MESSAGES_PATH = "/anthropic/v1/messages" +_ANTHROPIC_MODELS_PATH = "/v1/models" +_ANTHROPIC_MESSAGES_PATH = "/v1/messages" def _diagnostics_enabled() -> bool: @@ -207,7 +208,7 @@ def __init__(self) -> None: self._original_by_alias: dict[str, str] = {} self._lock = threading.Lock() - def advertise_models(self, body: bytes) -> bytes: + def rewrite_discovery_response(self, body: bytes) -> bytes: try: payload = json.loads(body) models = payload["data"] @@ -294,11 +295,10 @@ 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 - upstream_path = f"/anthropic{self.path}" - body = self.anthropic_model_aliases.rewrite_body(upstream_path, body) - url = self.anthropic_model_aliases.rewrite_path(upstream_path).lstrip("/") + body = self.anthropic_model_aliases.rewrite_body(self.path, body) + url = self.anthropic_model_aliases.rewrite_path(self.path).lstrip("/") should_transform_model_names = ( - self.command == "GET" and urlsplit(upstream_path).path == _ANTHROPIC_MODELS_PATH + self.command == "GET" and urlsplit(self.path).path == _ANTHROPIC_MODELS_PATH ) _diagnostic_log( "request_start", @@ -397,12 +397,17 @@ def _relay_response( bytes_relayed = 0 first_byte_ms: int | None = None try: + rewrite_model_response = ( + transform_model_names + and HTTPStatus.OK <= resp.status_code < HTTPStatus.MULTIPLE_CHOICES + ) self.send_response(resp.status_code) for key, value in resp.headers.items(): header_name = key.lower() drop_stale_content_encoding = ( - transform_model_names and header_name == "content-encoding" + rewrite_model_response 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() @@ -413,9 +418,8 @@ def _relay_response( # yielded as they arrive and pings keep the downstream connection # alive even before the model produces a large content block. response_chunks = ( - [self.anthropic_model_aliases.advertise_models(resp.read())] - if transform_model_names - and HTTPStatus.OK <= resp.status_code < HTTPStatus.MULTIPLE_CHOICES + [self.anthropic_model_aliases.rewrite_discovery_response(resp.read())] + if rewrite_model_response else resp.iter_raw() ) for chunk in response_chunks: @@ -488,7 +492,7 @@ def start_proxy( Returns (server, cache, client); the caller runs the server (e.g. in a thread) and calls shutdown()/cache.stop()/client.close() on exit. """ - upstream_base = f"{workspace.rstrip('/')}/ai-gateway/" + upstream_base = f"{workspace.rstrip('/')}/ai-gateway/anthropic" cache = _TokenCache( workspace, profile, @@ -511,11 +515,11 @@ def start_proxy( }, ) 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() diff --git a/src/ucode/constants.py b/src/ucode/constants.py new file mode 100644 index 0000000..3e9ff67 --- /dev/null +++ b/src/ucode/constants.py @@ -0,0 +1,3 @@ +"""Shared UCode constants.""" + +LOOPBACK_HOST = "127.0.0.1" diff --git a/tests/test_gateway_proxy.py b/tests/test_anthropic_gateway_proxy.py similarity index 90% rename from tests/test_gateway_proxy.py rename to tests/test_anthropic_gateway_proxy.py index 0aa431d..96f3320 100644 --- a/tests/test_gateway_proxy.py +++ b/tests/test_anthropic_gateway_proxy.py @@ -11,7 +11,7 @@ import httpx -from ucode import gateway_proxy +from ucode import anthropic_gateway_proxy as gateway_proxy def _make_jwt(exp: float | None) -> str: @@ -77,6 +77,9 @@ def iter_raw(self, chunk_size=None): self.chunk_sizes.append(chunk_size) yield from self._chunks + def read(self): + return b"".join(self._chunks) + class _BrokenPipeWriter(io.RawIOBase): """A wfile stand-in that raises BrokenPipeError on write, mimicking a client @@ -96,6 +99,7 @@ def _relay_handler(wfile) -> gateway_proxy._ProxyHandler: handler.requestline = "POST /v1/messages HTTP/1.1" handler.command = "POST" handler._headers_buffer = [] + handler.anthropic_model_aliases = gateway_proxy._AnthropicModelAliases() return handler @@ -169,6 +173,30 @@ def flush(self): assert b"Transfer-Encoding" not in blob # hop-by-hop, stripped assert resp.chunk_sizes == [None] # relay each upstream SSE/network chunk immediately + def test_rewritten_response_drops_content_encoding(self): + out = _Collect() + handler = _relay_handler(out) + resp = _FakeResponse( + 200, + {"Content-Encoding": "gzip"}, + [b'{"data":[{"id":"custom-model"}]}'], + ) + + handler._relay_response(resp, transform_model_names=True) + + assert b"Content-Encoding" not in bytes(out.data) + assert b"anthropic-aigw-custom-model" in bytes(out.data) + + def test_unchanged_error_response_keeps_content_encoding(self): + out = _Collect() + handler = _relay_handler(out) + resp = _FakeResponse(400, {"Content-Encoding": "gzip"}, [b"compressed-error"]) + + handler._relay_response(resp, transform_model_names=True) + + assert b"Content-Encoding: gzip" in bytes(out.data) + assert b"compressed-error" in bytes(out.data) + def test_diagnostics_identify_upstream_mid_stream_drop(self, monkeypatch, capsys): monkeypatch.setenv(gateway_proxy._DIAGNOSTICS_ENV, "1") @@ -368,7 +396,7 @@ def _handle_handler(client, cache, wfile) -> gateway_proxy._ProxyHandler: class TestAnthropicModelAliases: - def test_advertises_custom_models_without_changing_display_name(self): + def test_rewrites_custom_model_ids_without_changing_display_name(self): aliases = gateway_proxy._AnthropicModelAliases() body = json.dumps( { @@ -384,7 +412,7 @@ def test_advertises_custom_models_without_changing_display_name(self): } ).encode() - payload = json.loads(aliases.advertise_models(body)) + payload = json.loads(aliases.rewrite_discovery_response(body)) assert payload == { "data": [ @@ -403,10 +431,10 @@ def test_advertises_custom_models_without_changing_display_name(self): def test_rewrites_known_alias_in_messages_body(self): aliases = gateway_proxy._AnthropicModelAliases() - aliases.advertise_models(b'{"data":[{"id":"catalog.schema.custom"}]}') + aliases.rewrite_discovery_response(b'{"data":[{"id":"catalog.schema.custom"}]}') body = aliases.rewrite_body( - "/anthropic/v1/messages", + "/v1/messages", b'{"model":"anthropic-aigw-catalog.schema.custom","messages":[]}', ) @@ -414,13 +442,13 @@ def test_rewrites_known_alias_in_messages_body(self): def test_rewrites_known_alias_in_pagination_cursor(self): aliases = gateway_proxy._AnthropicModelAliases() - aliases.advertise_models(b'{"data":[{"id":"catalog.schema.custom"}]}') + aliases.rewrite_discovery_response(b'{"data":[{"id":"catalog.schema.custom"}]}') assert ( aliases.rewrite_path( - "/anthropic/v1/models?limit=1000&after_id=anthropic-aigw-catalog.schema.custom" + "/v1/models?limit=1000&after_id=anthropic-aigw-catalog.schema.custom" ) - == "/anthropic/v1/models?limit=1000&after_id=catalog.schema.custom" + == "/v1/models?limit=1000&after_id=catalog.schema.custom" ) def test_ignores_non_anthropic_models_path(self): @@ -433,17 +461,17 @@ def test_does_not_strip_unknown_prefixed_id(self): aliases = gateway_proxy._AnthropicModelAliases() unknown = "anthropic-aigw-legitimate-upstream-id" - assert aliases.rewrite_path(f"/anthropic/v1/models?after_id={unknown}") == ( - f"/anthropic/v1/models?after_id={unknown}" + assert aliases.rewrite_path(f"/v1/models?after_id={unknown}") == ( + f"/v1/models?after_id={unknown}" ) assert ( - aliases.rewrite_body("/anthropic/v1/messages", json.dumps({"model": unknown}).encode()) + 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._AnthropicModelAliases() - assert aliases.advertise_models(b"not-json") == b"not-json" + assert aliases.rewrite_discovery_response(b"not-json") == b"not-json" class _Collect(io.RawIOBase): @@ -488,7 +516,7 @@ def test_success_first_try_does_not_refresh(self): _handle_handler(client, cache, out)._handle() assert cache.refreshed == 0 assert client.sent_tokens == ["Bearer tok1"] - assert client.sent_urls == ["anthropic/v1/messages"] + assert client.sent_urls == ["v1/messages"] assert b"hi" in bytes(out.data) @@ -523,6 +551,9 @@ def run_refresher(self): bound = server.server_address[1] assert bound != busy_port # fell back to a different, free port assert bound != 0 + assert str(client.base_url) == ( + "https://x.staging.cloud.databricks.com/ai-gateway/anthropic/" + ) finally: server.server_close() client.close() From 5fa0bfbc67c267416712fd79f3b5ab34387fd37d Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Wed, 26 Aug 2026 01:34:23 +0000 Subject: [PATCH 5/9] refactor(claude): clarify model ID prefixing --- src/ucode/anthropic_gateway_proxy.py | 4 ++-- tests/test_anthropic_gateway_proxy.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/ucode/anthropic_gateway_proxy.py b/src/ucode/anthropic_gateway_proxy.py index 8a4eb74..00152f6 100644 --- a/src/ucode/anthropic_gateway_proxy.py +++ b/src/ucode/anthropic_gateway_proxy.py @@ -208,7 +208,7 @@ def __init__(self) -> None: self._original_by_alias: dict[str, str] = {} self._lock = threading.Lock() - def rewrite_discovery_response(self, body: bytes) -> bytes: + def prefix_model_ids(self, body: bytes) -> bytes: try: payload = json.loads(body) models = payload["data"] @@ -418,7 +418,7 @@ def _relay_response( # yielded as they arrive and pings keep the downstream connection # alive even before the model produces a large content block. response_chunks = ( - [self.anthropic_model_aliases.rewrite_discovery_response(resp.read())] + [self.anthropic_model_aliases.prefix_model_ids(resp.read())] if rewrite_model_response else resp.iter_raw() ) diff --git a/tests/test_anthropic_gateway_proxy.py b/tests/test_anthropic_gateway_proxy.py index 96f3320..9c6ab9b 100644 --- a/tests/test_anthropic_gateway_proxy.py +++ b/tests/test_anthropic_gateway_proxy.py @@ -412,7 +412,7 @@ def test_rewrites_custom_model_ids_without_changing_display_name(self): } ).encode() - payload = json.loads(aliases.rewrite_discovery_response(body)) + payload = json.loads(aliases.prefix_model_ids(body)) assert payload == { "data": [ @@ -431,7 +431,7 @@ def test_rewrites_custom_model_ids_without_changing_display_name(self): def test_rewrites_known_alias_in_messages_body(self): aliases = gateway_proxy._AnthropicModelAliases() - aliases.rewrite_discovery_response(b'{"data":[{"id":"catalog.schema.custom"}]}') + aliases.prefix_model_ids(b'{"data":[{"id":"catalog.schema.custom"}]}') body = aliases.rewrite_body( "/v1/messages", @@ -442,7 +442,7 @@ def test_rewrites_known_alias_in_messages_body(self): def test_rewrites_known_alias_in_pagination_cursor(self): aliases = gateway_proxy._AnthropicModelAliases() - aliases.rewrite_discovery_response(b'{"data":[{"id":"catalog.schema.custom"}]}') + aliases.prefix_model_ids(b'{"data":[{"id":"catalog.schema.custom"}]}') assert ( aliases.rewrite_path( @@ -471,7 +471,7 @@ def test_does_not_strip_unknown_prefixed_id(self): def test_leaves_malformed_discovery_response_unchanged(self): aliases = gateway_proxy._AnthropicModelAliases() - assert aliases.rewrite_discovery_response(b"not-json") == b"not-json" + assert aliases.prefix_model_ids(b"not-json") == b"not-json" class _Collect(io.RawIOBase): From 2ad807dece47d2bbbdcd01c217beddda9f375204 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Wed, 26 Aug 2026 01:55:03 +0000 Subject: [PATCH 6/9] refactor(claude): separate model discovery proxy --- src/ucode/agents/claude.py | 14 +- src/ucode/anthropic_model_discovery_proxy.py | 124 ++++++++++ ...opic_gateway_proxy.py => gateway_proxy.py} | 164 ++++--------- tests/test_agent_claude.py | 4 +- tests/test_anthropic_model_discovery_proxy.py | 224 ++++++++++++++++++ ...gateway_proxy.py => test_gateway_proxy.py} | 114 +-------- 6 files changed, 417 insertions(+), 227 deletions(-) create mode 100644 src/ucode/anthropic_model_discovery_proxy.py rename src/ucode/{anthropic_gateway_proxy.py => gateway_proxy.py} (79%) create mode 100644 tests/test_anthropic_model_discovery_proxy.py rename tests/{test_anthropic_gateway_proxy.py => test_gateway_proxy.py} (79%) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 81ce046..af6faf1 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -16,10 +16,8 @@ from typing import cast from ucode.agent_updates import available_npm_package_update -from ucode.anthropic_gateway_proxy import ( - AI_GATEWAY_TOKEN_HEADER, - AUTHORIZATION_HEADER, - start_proxy, +from ucode.anthropic_model_discovery_proxy import ( + start_proxy as start_anthropic_model_discovery_proxy, ) from ucode.config_io import ( APP_DIR, @@ -35,6 +33,10 @@ build_tool_base_url, get_databricks_token, ) +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 ( @@ -1004,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, @@ -1036,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, diff --git a/src/ucode/anthropic_model_discovery_proxy.py b/src/ucode/anthropic_model_discovery_proxy.py new file mode 100644 index 0000000..826cbb5 --- /dev/null +++ b/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()}, + ) diff --git a/src/ucode/anthropic_gateway_proxy.py b/src/ucode/gateway_proxy.py similarity index 79% rename from src/ucode/anthropic_gateway_proxy.py rename to src/ucode/gateway_proxy.py index 00152f6..a2d287c 100644 --- a/src/ucode/anthropic_gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -24,9 +24,8 @@ import threading import time import uuid -from http import HTTPStatus from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit +from typing import cast import httpx @@ -70,9 +69,6 @@ # 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-" -_ANTHROPIC_MODELS_PATH = "/v1/models" -_ANTHROPIC_MESSAGES_PATH = "/v1/messages" def _diagnostics_enabled() -> bool: @@ -201,83 +197,11 @@ def _forwarded_request_headers( return headers -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 _ProxyHandler(BaseHTTPRequestHandler): # Set by the server factory. cache: _TokenCache client: httpx.Client token_header = AI_GATEWAY_TOKEN_HEADER - anthropic_model_aliases: _AnthropicModelAliases def log_message(self, format: str, *args: object) -> None: return @@ -290,16 +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 - body = self.anthropic_model_aliases.rewrite_body(self.path, body) - url = self.anthropic_model_aliases.rewrite_path(self.path).lstrip("/") - should_transform_model_names = ( - self.command == "GET" and urlsplit(self.path).path == _ANTHROPIC_MODELS_PATH - ) + url, body = self._transform_request(body) _diagnostic_log( "request_start", request_id=diagnostic_id, @@ -318,12 +244,7 @@ def _handle(self) -> None: elapsed_ms=round((time.monotonic() - started) * 1000), ) if resp.status_code not in (401, 403): - self._relay_response( - resp, - transform_model_names=should_transform_model_names, - diagnostic_id=diagnostic_id, - started=started, - ) + self._relay_response(resp, 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. @@ -352,12 +273,7 @@ def _handle(self) -> None: status=resp.status_code, elapsed_ms=round((time.monotonic() - started) * 1000), ) - self._relay_response( - resp, - transform_model_names=should_transform_model_names, - diagnostic_id=diagnostic_id, - started=started, - ) + self._relay_response(resp, diagnostic_id=diagnostic_id, started=started) except (BrokenPipeError, ConnectionResetError): # Client closed before/while we relayed headers — routine on cancel. _diagnostic_log( @@ -388,7 +304,6 @@ def _relay_response( self, resp: httpx.Response, *, - transform_model_names: bool, diagnostic_id: str | None = None, started: float | None = None, ) -> None: @@ -397,15 +312,12 @@ def _relay_response( bytes_relayed = 0 first_byte_ms: int | None = None try: - rewrite_model_response = ( - transform_model_names - and HTTPStatus.OK <= resp.status_code < HTTPStatus.MULTIPLE_CHOICES - ) + transformed_body = self._transform_response(resp) self.send_response(resp.status_code) for key, value in resp.headers.items(): header_name = key.lower() drop_stale_content_encoding = ( - rewrite_model_response and header_name == "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: @@ -418,9 +330,7 @@ def _relay_response( # yielded as they arrive and pings keep the downstream connection # alive even before the model produces a large content block. response_chunks = ( - [self.anthropic_model_aliases.prefix_model_ids(resp.read())] - if rewrite_model_response - else resp.iter_raw() + [transformed_body] if transformed_body is not None else resp.iter_raw() ) for chunk in response_chunks: if chunk: @@ -475,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. @@ -492,7 +405,7 @@ def start_proxy( Returns (server, cache, client); the caller runs the server (e.g. in a thread) and calls shutdown()/cache.stop()/client.close() on exit. """ - upstream_base = f"{workspace.rstrip('/')}/ai-gateway/anthropic" + upstream_base = f"{workspace.rstrip('/')}/ai-gateway/anthropic/" cache = _TokenCache( workspace, profile, @@ -502,17 +415,18 @@ 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) - anthropic_model_aliases = _AnthropicModelAliases() - - handler = type( - "BoundProxyHandler", - (_ProxyHandler,), - { - "cache": cache, - "client": client, - "token_header": token_header, - "anthropic_model_aliases": anthropic_model_aliases, - }, + handler = cast( + type[BaseHTTPRequestHandler], + type( + "BoundProxyHandler", + (handler_class,), + { + "cache": cache, + "client": client, + "token_header": token_header, + **handler_attributes, + }, + ), ) try: server = ThreadingHTTPServer((LOOPBACK_HOST, port), handler) @@ -524,3 +438,21 @@ def start_proxy( 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={}, + ) diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 6b49782..1c1bb18 100644 --- a/tests/test_agent_claude.py +++ b/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: diff --git a/tests/test_anthropic_model_discovery_proxy.py b/tests/test_anthropic_model_discovery_proxy.py new file mode 100644 index 0000000..66441fd --- /dev/null +++ b/tests/test_anthropic_model_discovery_proxy.py @@ -0,0 +1,224 @@ +"""Tests for Anthropic model discovery transformations.""" + +from __future__ import annotations + +import io +import json + +from ucode import anthropic_model_discovery_proxy + + +class _FakeResponse: + def __init__(self, status_code: int, headers: dict[str, str], body: bytes): + self.status_code = status_code + self.headers = headers + self._body = body + + def read(self): + return self._body + + def iter_raw(self): + yield self._body + + def __enter__(self): + return self + + def __exit__(self, *_args): + return False + + +class _FakeClient: + def __init__(self, response): + self.response = response + self.request = None + + def stream(self, method, url, headers, content): + self.request = (method, url, headers, content) + return self.response + + +class _FakeCache: + token = "databricks-token" + + def refresh(self): + return None + + +class _Collect(io.RawIOBase): + def __init__(self): + self.data = bytearray() + + def write(self, body): # type: ignore[override] + self.data += bytes(body) + return len(body) + + def flush(self): + return None + + +def _handler(wfile, path="/v1/models", command="GET"): + handler = object.__new__(anthropic_model_discovery_proxy._AnthropicModelDiscoveryHandler) + handler.wfile = wfile + handler.request_version = "HTTP/1.1" + handler.requestline = f"{command} {path} HTTP/1.1" + handler.command = command + handler.path = path + handler._headers_buffer = [] + handler.anthropic_model_aliases = anthropic_model_discovery_proxy._AnthropicModelAliases() + return handler + + +class TestAnthropicModelAliases: + def test_prefixes_custom_model_ids_without_changing_display_name(self): + aliases = anthropic_model_discovery_proxy._AnthropicModelAliases() + body = json.dumps( + { + "data": [ + {"id": "catalog.schema.custom", "display_name": "Custom model"}, + {"id": "system.ai.claude-sonnet"}, + {"id": "claude-sonnet"}, + {"id": "catalog.schema.anthropic-provider"}, + {"id": "anthropic-provider"}, + ], + "first_id": "catalog.schema.custom", + "last_id": "catalog.schema.anthropic-provider", + } + ).encode() + + payload = json.loads(aliases.prefix_model_ids(body)) + + assert payload == { + "data": [ + { + "id": "anthropic-aigw-catalog.schema.custom", + "display_name": "Custom model", + }, + {"id": "system.ai.claude-sonnet"}, + {"id": "claude-sonnet"}, + {"id": "catalog.schema.anthropic-provider"}, + {"id": "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 = anthropic_model_discovery_proxy._AnthropicModelAliases() + aliases.prefix_model_ids(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 = anthropic_model_discovery_proxy._AnthropicModelAliases() + aliases.prefix_model_ids(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_ignores_non_anthropic_models_path(self): + aliases = anthropic_model_discovery_proxy._AnthropicModelAliases() + path = "/codex/v1/models?after_id=anthropic-aigw-catalog.schema.custom" + + assert aliases.rewrite_path(path) == path + + def test_does_not_strip_unknown_prefixed_id(self): + aliases = anthropic_model_discovery_proxy._AnthropicModelAliases() + 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 = anthropic_model_discovery_proxy._AnthropicModelAliases() + assert aliases.prefix_model_ids(b"not-json") == b"not-json" + + +class TestAnthropicModelDiscoveryHandler: + def test_inherits_relayed_auth_and_prefixes_models(self): + out = _Collect() + handler = _handler(out) + handler.headers = {"Authorization": "Bearer subscription-token"} + handler.rfile = io.BytesIO() + handler.cache = _FakeCache() + handler.client = _FakeClient(_FakeResponse(200, {}, b'{"data":[{"id":"custom-model"}]}')) + + handler._handle() + + method, url, headers, body = handler.client.request + assert (method, url, body) == ("GET", "v1/models", None) + assert headers["Authorization"] == "Bearer subscription-token" + assert headers["X-Databricks-AI-Gateway-Token"] == "Bearer databricks-token" + assert b"anthropic-aigw-custom-model" in bytes(out.data) + + def test_prefixes_successful_model_response_and_drops_content_encoding(self): + out = _Collect() + handler = _handler(out) + response = _FakeResponse( + 200, + {"Content-Encoding": "gzip"}, + b'{"data":[{"id":"custom-model"}]}', + ) + + handler._relay_response(response) + + assert b"Content-Encoding" not in bytes(out.data) + assert b"anthropic-aigw-custom-model" in bytes(out.data) + + def test_keeps_content_encoding_for_unchanged_error(self): + out = _Collect() + handler = _handler(out) + response = _FakeResponse(400, {"Content-Encoding": "gzip"}, b"compressed-error") + + handler._relay_response(response) + + assert b"Content-Encoding: gzip" in bytes(out.data) + assert b"compressed-error" in bytes(out.data) + + def test_strips_known_alias_from_message_request(self): + handler = _handler(_Collect(), path="/v1/messages", command="POST") + handler.anthropic_model_aliases.prefix_model_ids( + b'{"data":[{"id":"catalog.schema.custom"}]}' + ) + + url, body = handler._transform_request(b'{"model":"anthropic-aigw-catalog.schema.custom"}') + + assert url == "v1/messages" + assert json.loads(body) == {"model": "catalog.schema.custom"} + + +def test_start_proxy_uses_discovery_handler(monkeypatch): + call = {} + + def start(*args, **kwargs): + call["args"] = args + call["kwargs"] = kwargs + return "server", "cache", "client" + + monkeypatch.setattr(anthropic_model_discovery_proxy.gateway_proxy, "_start_proxy", start) + + result = anthropic_model_discovery_proxy.start_proxy("workspace", "profile", 1, "header", False) + + assert result == ("server", "cache", "client") + assert call["args"][:5] == ("workspace", "profile", 1, "header", False) + assert ( + call["kwargs"]["handler_class"] + is anthropic_model_discovery_proxy._AnthropicModelDiscoveryHandler + ) + assert isinstance( + call["kwargs"]["handler_attributes"]["anthropic_model_aliases"], + anthropic_model_discovery_proxy._AnthropicModelAliases, + ) diff --git a/tests/test_anthropic_gateway_proxy.py b/tests/test_gateway_proxy.py similarity index 79% rename from tests/test_anthropic_gateway_proxy.py rename to tests/test_gateway_proxy.py index 9c6ab9b..498aab4 100644 --- a/tests/test_anthropic_gateway_proxy.py +++ b/tests/test_gateway_proxy.py @@ -11,7 +11,7 @@ import httpx -from ucode import anthropic_gateway_proxy as gateway_proxy +from ucode import gateway_proxy def _make_jwt(exp: float | None) -> str: @@ -99,7 +99,6 @@ def _relay_handler(wfile) -> gateway_proxy._ProxyHandler: handler.requestline = "POST /v1/messages HTTP/1.1" handler.command = "POST" handler._headers_buffer = [] - handler.anthropic_model_aliases = gateway_proxy._AnthropicModelAliases() return handler @@ -109,7 +108,7 @@ def test_relay_swallows_broken_pipe_on_headers(self): handler = _relay_handler(_BrokenPipeWriter()) resp = _FakeResponse(200, {}, [b'{"ok":true}']) # Must not raise — a dead client is a routine teardown, not an error. - handler._relay_response(resp, transform_model_names=False) + handler._relay_response(resp) def test_relay_swallows_connection_reset_mid_stream(self): # Headers flush ok, then the client resets while streaming body chunks. @@ -127,7 +126,7 @@ def flush(self): handler = _relay_handler(_ResetAfterHeaders()) resp = _FakeResponse(200, {}, [b"chunk-of-sse-data"]) - handler._relay_response(resp, transform_model_names=False) + handler._relay_response(resp) def test_relay_swallows_upstream_error_mid_stream(self): # Upstream drops mid-body after headers are already sent — we can't signal @@ -145,7 +144,7 @@ def _chunks(): handler = _relay_handler(_Ok()) resp = _FakeResponse(200, {}, _chunks()) - handler._relay_response(resp, transform_model_names=False) # must not raise + handler._relay_response(resp) # must not raise def test_relay_forwards_status_and_skips_hop_by_hop_headers(self): # A non-200 status (e.g. 429 rate limit) and content headers are relayed; @@ -166,36 +165,26 @@ def flush(self): {"Content-Type": "application/json", "Transfer-Encoding": "chunked"}, [b'{"type":"error"}'], ) - handler._relay_response(resp, transform_model_names=False) + handler._relay_response(resp) blob = b"".join(chunks_written) assert b"429" in blob assert b"Content-Type: application/json" in blob assert b"Transfer-Encoding" not in blob # hop-by-hop, stripped assert resp.chunk_sizes == [None] # relay each upstream SSE/network chunk immediately - def test_rewritten_response_drops_content_encoding(self): + def test_base_handler_does_not_transform_model_response(self): out = _Collect() handler = _relay_handler(out) - resp = _FakeResponse( + response = _FakeResponse( 200, {"Content-Encoding": "gzip"}, - [b'{"data":[{"id":"custom-model"}]}'], + [b"compressed-model-response"], ) - handler._relay_response(resp, transform_model_names=True) - - assert b"Content-Encoding" not in bytes(out.data) - assert b"anthropic-aigw-custom-model" in bytes(out.data) - - def test_unchanged_error_response_keeps_content_encoding(self): - out = _Collect() - handler = _relay_handler(out) - resp = _FakeResponse(400, {"Content-Encoding": "gzip"}, [b"compressed-error"]) - - handler._relay_response(resp, transform_model_names=True) + handler._relay_response(response) assert b"Content-Encoding: gzip" in bytes(out.data) - assert b"compressed-error" in bytes(out.data) + assert b"compressed-model-response" in bytes(out.data) def test_diagnostics_identify_upstream_mid_stream_drop(self, monkeypatch, capsys): monkeypatch.setenv(gateway_proxy._DIAGNOSTICS_ENV, "1") @@ -214,7 +203,6 @@ def _chunks(): handler = _relay_handler(_Ok()) handler._relay_response( _FakeResponse(200, {}, _chunks()), - transform_model_names=False, diagnostic_id="local-id", started=time.monotonic(), ) @@ -231,7 +219,7 @@ def _chunks(): def test_diagnostics_are_silent_by_default(self, monkeypatch, capsys): monkeypatch.delenv(gateway_proxy._DIAGNOSTICS_ENV, raising=False) handler = _relay_handler(_Collect()) - handler._relay_response(_FakeResponse(200, {}, [b"ok"]), transform_model_names=False) + handler._relay_response(_FakeResponse(200, {}, [b"ok"])) assert capsys.readouterr().err == "" @@ -383,7 +371,6 @@ def _handle_handler(client, cache, wfile) -> gateway_proxy._ProxyHandler: h = object.__new__(gateway_proxy._ProxyHandler) h.client = client h.cache = cache - h.anthropic_model_aliases = gateway_proxy._AnthropicModelAliases() h.headers = {} h.rfile = io.BytesIO(b"") h.path = "/v1/messages" @@ -395,85 +382,6 @@ def _handle_handler(client, cache, wfile) -> gateway_proxy._ProxyHandler: return h -class TestAnthropicModelAliases: - def test_rewrites_custom_model_ids_without_changing_display_name(self): - aliases = gateway_proxy._AnthropicModelAliases() - body = json.dumps( - { - "data": [ - {"id": "catalog.schema.custom", "display_name": "Custom model"}, - {"id": "system.ai.claude-sonnet"}, - {"id": "claude-sonnet"}, - {"id": "catalog.schema.anthropic-provider"}, - {"id": "anthropic-provider"}, - ], - "first_id": "catalog.schema.custom", - "last_id": "catalog.schema.anthropic-provider", - } - ).encode() - - payload = json.loads(aliases.prefix_model_ids(body)) - - assert payload == { - "data": [ - { - "id": "anthropic-aigw-catalog.schema.custom", - "display_name": "Custom model", - }, - {"id": "system.ai.claude-sonnet"}, - {"id": "claude-sonnet"}, - {"id": "catalog.schema.anthropic-provider"}, - {"id": "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._AnthropicModelAliases() - aliases.prefix_model_ids(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._AnthropicModelAliases() - aliases.prefix_model_ids(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_ignores_non_anthropic_models_path(self): - aliases = gateway_proxy._AnthropicModelAliases() - path = "/codex/v1/models?after_id=anthropic-aigw-catalog.schema.custom" - - assert aliases.rewrite_path(path) == path - - def test_does_not_strip_unknown_prefixed_id(self): - aliases = gateway_proxy._AnthropicModelAliases() - 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._AnthropicModelAliases() - assert aliases.prefix_model_ids(b"not-json") == b"not-json" - - class _Collect(io.RawIOBase): def __init__(self): self.data = bytearray() From 290b84e16483f03b9a94339f69a2e72057af60f8 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Wed, 26 Aug 2026 17:03:20 +0000 Subject: [PATCH 7/9] Address model discovery proxy review --- src/ucode/anthropic_model_discovery_proxy.py | 11 +++-- src/ucode/gateway_proxy.py | 44 +++++-------------- tests/test_anthropic_model_discovery_proxy.py | 26 ++++++++++- 3 files changed, 43 insertions(+), 38 deletions(-) diff --git a/src/ucode/anthropic_model_discovery_proxy.py b/src/ucode/anthropic_model_discovery_proxy.py index 826cbb5..c22e17b 100644 --- a/src/ucode/anthropic_model_discovery_proxy.py +++ b/src/ucode/anthropic_model_discovery_proxy.py @@ -4,6 +4,7 @@ import json import threading +from collections.abc import Iterable from http import HTTPStatus from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit @@ -95,15 +96,17 @@ def _transform_request(self, body: bytes | None) -> tuple[str, bytes | None]: url = self.anthropic_model_aliases.rewrite_path(self.path).lstrip("/") return url, body - def _transform_response(self, resp: httpx.Response) -> bytes | None: + def _response_chunks(self, resp: httpx.Response) -> tuple[Iterable[bytes], frozenset[str]]: 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()) + return super()._response_chunks(resp) + body = self.anthropic_model_aliases.prefix_model_ids(resp.read()) + # resp.read() decodes compression; rewritten JSON is uncompressed. + return (body,), frozenset({"content-encoding"}) def start_proxy( @@ -113,7 +116,7 @@ def start_proxy( token_header: str, force_refresh_near_expiry: bool, ): - return gateway_proxy._start_proxy( + return gateway_proxy.start_proxy( workspace, profile, port, diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index a2d287c..9d232fc 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -24,6 +24,7 @@ import threading import time import uuid +from collections.abc import Iterable from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import cast @@ -217,8 +218,8 @@ def _safe_send_error(self, code: int, message: str) -> None: 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 _response_chunks(self, resp: httpx.Response) -> tuple[Iterable[bytes], frozenset[str]]: + return resp.iter_raw(), frozenset() def _handle(self) -> None: diagnostic_id = uuid.uuid4().hex[:12] @@ -312,15 +313,13 @@ def _relay_response( bytes_relayed = 0 first_byte_ms: int | None = None try: - transformed_body = self._transform_response(resp) + # The upstream request has completed through response headers before + # this hook selects raw streaming or a buffered response body. + response_chunks, dropped_headers = self._response_chunks(resp) self.send_response(resp.status_code) for key, value in resp.headers.items(): 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: + if header_name not in _HOP_BY_HOP and header_name not in dropped_headers: self.send_header(key, value) self.end_headers() # Do not pass a fixed chunk size here. httpx accumulates bytes until @@ -329,9 +328,6 @@ 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. - 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: @@ -385,15 +381,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], + handler_class: type[_ProxyHandler] = _ProxyHandler, + handler_attributes: dict[str, object] | None = None, ) -> tuple[ThreadingHTTPServer, _TokenCache, httpx.Client]: """Start the loopback refresh proxy + its background token refresher. @@ -424,7 +420,7 @@ def _start_proxy( "cache": cache, "client": client, "token_header": token_header, - **handler_attributes, + **(handler_attributes or {}), }, ), ) @@ -438,21 +434,3 @@ def _start_proxy( 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={}, - ) diff --git a/tests/test_anthropic_model_discovery_proxy.py b/tests/test_anthropic_model_discovery_proxy.py index 66441fd..7fff900 100644 --- a/tests/test_anthropic_model_discovery_proxy.py +++ b/tests/test_anthropic_model_discovery_proxy.py @@ -13,11 +13,15 @@ def __init__(self, status_code: int, headers: dict[str, str], body: bytes): self.status_code = status_code self.headers = headers self._body = body + self.read_calls = 0 + self.iter_raw_calls = 0 def read(self): + self.read_calls += 1 return self._body def iter_raw(self): + self.iter_raw_calls += 1 yield self._body def __enter__(self): @@ -187,6 +191,26 @@ def test_keeps_content_encoding_for_unchanged_error(self): assert b"Content-Encoding: gzip" in bytes(out.data) assert b"compressed-error" in bytes(out.data) + assert response.read_calls == 0 + assert response.iter_raw_calls == 1 + + def test_streams_relayed_inference_response_without_buffering(self): + out = _Collect() + handler = _handler(out, path="/v1/messages", command="POST") + handler.headers = {"Authorization": "Bearer subscription-token", "Content-Length": "2"} + handler.rfile = io.BytesIO(b"{}") + handler.cache = _FakeCache() + response = _FakeResponse(200, {"Content-Type": "text/event-stream"}, b"data: event\n\n") + handler.client = _FakeClient(response) + + handler._handle() + + _method, _url, headers, _body = handler.client.request + assert headers["Authorization"] == "Bearer subscription-token" + assert headers["X-Databricks-AI-Gateway-Token"] == "Bearer databricks-token" + assert response.read_calls == 0 + assert response.iter_raw_calls == 1 + assert b"data: event\n\n" in bytes(out.data) def test_strips_known_alias_from_message_request(self): handler = _handler(_Collect(), path="/v1/messages", command="POST") @@ -208,7 +232,7 @@ def start(*args, **kwargs): call["kwargs"] = kwargs return "server", "cache", "client" - monkeypatch.setattr(anthropic_model_discovery_proxy.gateway_proxy, "_start_proxy", start) + monkeypatch.setattr(anthropic_model_discovery_proxy.gateway_proxy, "start_proxy", start) result = anthropic_model_discovery_proxy.start_proxy("workspace", "profile", 1, "header", False) From a63ddd6e5e4b6de41fb6600517122292dde4a56e Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Wed, 26 Aug 2026 17:19:04 +0000 Subject: [PATCH 8/9] Refactor Claude gateway proxy isolation --- src/ucode/agents/claude.py | 8 +- src/ucode/anthropic_gateway_proxy.py | 553 ++++++++++++++++++ src/ucode/anthropic_model_discovery_proxy.py | 127 ---- src/ucode/gateway_proxy.py | 43 +- tests/test_agent_claude.py | 2 +- ...oxy.py => test_anthropic_gateway_proxy.py} | 31 +- tests/test_gateway_proxy.py | 23 - 7 files changed, 582 insertions(+), 205 deletions(-) create mode 100644 src/ucode/anthropic_gateway_proxy.py delete mode 100644 src/ucode/anthropic_model_discovery_proxy.py rename tests/{test_anthropic_model_discovery_proxy.py => test_anthropic_gateway_proxy.py} (87%) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index af6faf1..f64b536 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -16,8 +16,8 @@ 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.anthropic_gateway_proxy import ( + start_proxy as start_anthropic_gateway_proxy, ) from ucode.config_io import ( APP_DIR, @@ -1006,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_anthropic_model_discovery_proxy( + server, cache, client = start_anthropic_gateway_proxy( workspace, state.get("profile"), port, @@ -1038,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_anthropic_model_discovery_proxy( + server, cache, client = start_anthropic_gateway_proxy( workspace, state.get("profile"), 0, diff --git a/src/ucode/anthropic_gateway_proxy.py b/src/ucode/anthropic_gateway_proxy.py new file mode 100644 index 0000000..c7cebb1 --- /dev/null +++ b/src/ucode/anthropic_gateway_proxy.py @@ -0,0 +1,553 @@ +"""Loopback refresh proxy for Claude gateway requests. + +A relayed Model Provider Service authenticates the caller's own Anthropic +subscription OAuth (which Claude Code owns in the `Authorization` header) and +carries a Databricks credential in the `X-Databricks-AI-Gateway-Token` swap +header. Native gateway discovery instead carries the Databricks credential in +`Authorization`. The proxy refreshes the applicable header, streams inference +responses verbatim, and rewrites model discovery responses when needed. + +Security invariants (mirroring `databricks.py` token handling): + - Binds 127.0.0.1 only; never exposed off-host. + - Never logs header values or bodies. The Databricks token lives in memory, + refreshed off the request path; the Anthropic OAuth in `Authorization` is + passed through untouched in relayed mode and never logged. +""" + +from __future__ import annotations + +import base64 +import binascii +import json +import os +import sys +import threading +import time +import uuid +from collections.abc import Iterable +from http import HTTPStatus +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import cast +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit + +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 +# client-supplied value is replaced, so a stale settings.json value can't leak. +AI_GATEWAY_TOKEN_HEADER = "X-Databricks-AI-Gateway-Token" +AUTHORIZATION_HEADER = "Authorization" +# Hop-by-hop headers must not be forwarded across the proxy. +_HOP_BY_HOP = frozenset( + h.lower() + for h in ( + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailers", + "transfer-encoding", + "upgrade", + "host", + "content-length", + ) +) +# Per-operation upstream timeouts. `read` is generous because model turns stream +# over a single response and Anthropic emits SSE pings, so inter-chunk gaps stay +# small; `connect`/`pool` fail fast when the gateway is unreachable. +_UPSTREAM_TIMEOUT = httpx.Timeout(connect=10.0, read=600.0, write=600.0, pool=10.0) +# Refresh once the token has less than this many seconds of life left. Databricks +# access tokens live ~1h; a 10-min buffer leaves ample headroom for a retry. +_REFRESH_BUFFER_S = 600 +# How often the background thread re-checks freshness. Cheap: it only shells out +# to the CLI when actually within the buffer, otherwise it's a bare clock compare. +_REFRESHER_POLL_S = 120 +# Assumed lifetime when a token carries no decodable `exp` (defensive fallback). +_DEFAULT_TTL_S = 3600 +# Opt-in transport diagnostics for intermittent streaming failures. Events only +# contain locally-generated request ids, timings, status codes, byte counts, +# and exception class names — never headers, bodies, or credentials. +_DIAGNOSTICS_ENV = "UCODE_RELAYED_PROXY_DIAGNOSTICS" +_DIAGNOSTICS_TRUE = frozenset({"1", "true", "yes", "on"}) + + +def _diagnostics_enabled() -> bool: + return os.environ.get(_DIAGNOSTICS_ENV, "").strip().lower() in _DIAGNOSTICS_TRUE + + +def _diagnostic_log(event: str, **fields: object) -> None: + if not _diagnostics_enabled(): + return + payload = {"event": event, **fields} + sys.stderr.write( + f"[ucode-relay] {json.dumps(payload, sort_keys=True, separators=(',', ':'))}\n" + ) + sys.stderr.flush() + + +def _jwt_exp(token: str) -> float | None: + """Best-effort `exp` (epoch seconds) from a JWT access token, else None.""" + try: + payload = token.split(".")[1] + payload += "=" * (-len(payload) % 4) # restore base64 padding + return float(json.loads(base64.urlsafe_b64decode(payload))["exp"]) + except (IndexError, ValueError, KeyError, binascii.Error, json.JSONDecodeError): + return None + + +def _log_refresh_failure(exc: BaseException) -> None: + """Surface (never silently swallow) a refresh failure, without leaking any + token or header value.""" + sys.stderr.write( + f"[ucode] Databricks token refresh failed: {exc}. If the session stalls, " + "run `databricks auth login` for your workspace profile.\n" + ) + + +class _TokenCache: + """Holds the current Databricks token and its expiry, refreshing lazily as it + nears expiry. + + A background thread refreshes proactively so the request path rarely blocks, + but the request path also refreshes on demand — which is what carries the + token across events the timer can't (laptop sleep suspends the monotonic + clock, so a fixed interval silently stops advancing). All refreshes are + single-flighted through ``_refresh_lock`` so a burst of requests at the expiry + boundary triggers exactly one CLI call, not a thundering herd on the shared + token cache.""" + + def __init__( + self, + workspace: str, + profile: str | None, + *, + force_refresh_near_expiry: bool = False, + ) -> None: + self._workspace = workspace + self._profile = profile + self._force_refresh_near_expiry = force_refresh_near_expiry + self._state_lock = threading.Lock() # guards _token / _expiry (brief) + self._refresh_lock = threading.Lock() # single-flights the CLI refresh + self._stop = threading.Event() + self._token = "" + self._expiry = 0.0 + # Preserve the existing non-forced relayed-auth fetch. Gateway discovery + # opts into a forced fetch so its static client token starts with a full TTL. + self._refresh(force=force_refresh_near_expiry) + + def _refresh(self, *, force: bool) -> None: + """Mint a token and record its expiry.""" + token = get_databricks_token(self._workspace, self._profile, force_refresh=force) + expiry = _jwt_exp(token) or (time.time() + _DEFAULT_TTL_S) + with self._state_lock: + self._token = token + self._expiry = expiry + + def _fresh_enough(self) -> bool: + with self._state_lock: + return bool(self._token) and time.time() < self._expiry - _REFRESH_BUFFER_S + + def _ensure_fresh(self) -> None: + if self._fresh_enough(): + return + with self._refresh_lock: + if self._fresh_enough(): # another thread refreshed while we waited + return + try: + self._refresh(force=self._force_refresh_near_expiry) + except RuntimeError as exc: + # Keep serving the current token; a request that then 401s triggers + # a forced refresh + retry (see _ProxyHandler._handle). + _log_refresh_failure(exc) + + @property + def token(self) -> str: + self._ensure_fresh() + with self._state_lock: + return self._token + + def refresh(self) -> None: + """Force a fresh mint now (used by the retry-on-401 path).""" + with self._refresh_lock: + self._refresh(force=True) + + def run_refresher(self) -> None: + while not self._stop.wait(_REFRESHER_POLL_S): + try: + self._ensure_fresh() + except Exception as exc: # noqa: BLE001 - a stray error must NOT kill the thread + # If this thread dies, nothing refreshes and the session lapses at + # the ~1h mark until restart. Log and keep looping instead. + _log_refresh_failure(exc) + + def stop(self) -> None: + self._stop.set() + + +def _forwarded_request_headers( + handler: BaseHTTPRequestHandler, + token: str, + token_header: str = AI_GATEWAY_TOKEN_HEADER, +) -> dict[str, str]: + strip_on_forward = _HOP_BY_HOP | {token_header.lower()} + headers = { + key: value for key, value in handler.headers.items() if key.lower() not in strip_on_forward + } + headers[token_header] = f"Bearer {token}" + return headers + + +class _ProxyHandler(BaseHTTPRequestHandler): + # Set by the server factory. + cache: _TokenCache + client: httpx.Client + token_header = AI_GATEWAY_TOKEN_HEADER + + def log_message(self, format: str, *args: object) -> None: + return + + def _safe_send_error(self, code: int, message: str) -> None: + # The client (Claude Code) may already have disconnected, in which case + # reporting the error writes to a dead socket and raises again; swallow it. + try: + self.send_error(code, message) + except OSError: + pass + + def _transform_request(self, body: bytes | None) -> tuple[str, bytes | None]: + return self.path.lstrip("/"), body + + def _response_chunks(self, resp: httpx.Response) -> tuple[Iterable[bytes], frozenset[str]]: + return resp.iter_raw(), frozenset() + + 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, body = self._transform_request(body) + _diagnostic_log( + "request_start", + request_id=diagnostic_id, + method=self.command, + path=self.path.split("?", 1)[0], + ) + try: + # First attempt with the current token. + headers = _forwarded_request_headers(self, self.cache.token, self.token_header) + with self.client.stream(self.command, url, headers=headers, content=body) as resp: + _diagnostic_log( + "upstream_headers", + request_id=diagnostic_id, + attempt=1, + status=resp.status_code, + 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) + return + # Auth rejected. Drain the (small) error body so the pooled + # connection can be reused, then fall through to one retry. + resp.read() + # A relayed 401/403 may be a stale Databricks swap token rather than a + # bad Anthropic OAuth — the two are indistinguishable from the status + # alone. Force-refresh the Databricks token and retry once. If it was the + # Anthropic layer, the retry still 401s and we relay it verbatim, so a + # genuine re-auth is triggered; a stale-Databricks 401 self-heals here + # instead of surfacing to Claude Code as a spurious Anthropic prompt. + try: + self.cache.refresh() + except RuntimeError as exc: + # Refresh failed: the Databricks OAuth session is dead (not just the + # access token) and can't be re-minted non-interactively. Surface the + # `databricks auth login` hint rather than silently relaying a bare 401, + # which otherwise reads as an Anthropic `/login` prompt and sends the + # user to the wrong re-auth. Still retry + relay with the existing token. + _log_refresh_failure(exc) + headers = _forwarded_request_headers(self, self.cache.token, self.token_header) + with self.client.stream(self.command, url, headers=headers, content=body) as resp: + _diagnostic_log( + "upstream_headers", + request_id=diagnostic_id, + attempt=2, + status=resp.status_code, + elapsed_ms=round((time.monotonic() - started) * 1000), + ) + self._relay_response(resp, diagnostic_id=diagnostic_id, started=started) + except (BrokenPipeError, ConnectionResetError): + # Client closed before/while we relayed headers — routine on cancel. + _diagnostic_log( + "client_disconnect", + request_id=diagnostic_id, + phase="request", + elapsed_ms=round((time.monotonic() - started) * 1000), + ) + return + except httpx.HTTPError as exc: + # Upstream failed before any bytes reached the client; a 502 is still + # sendable. (An HTTP *status* like 429 is not an error here — httpx + # only raises for transport failures — so real gateway errors are + # relayed verbatim by `_relay_response`.) + _diagnostic_log( + "upstream_request_error", + request_id=diagnostic_id, + error_type=type(exc).__name__, + elapsed_ms=round((time.monotonic() - started) * 1000), + ) + self._safe_send_error(502, "gateway proxy upstream error") + + # Streaming passthrough: forward chunks as they arrive so SSE token streaming + # is not buffered (buffering would add full-response latency to first token). + # `iter_raw` preserves any Content-Encoding verbatim (we relay that header), + # so the proxy stays byte-transparent. + def _relay_response( + self, + resp: httpx.Response, + *, + diagnostic_id: str | None = None, + started: float | None = None, + ) -> None: + started = started if started is not None else time.monotonic() + chunks = 0 + bytes_relayed = 0 + first_byte_ms: int | None = None + try: + # The upstream request has completed through response headers before + # this hook selects raw streaming or a buffered response body. + response_chunks, dropped_headers = self._response_chunks(resp) + self.send_response(resp.status_code) + for key, value in resp.headers.items(): + header_name = key.lower() + if header_name not in _HOP_BY_HOP and header_name not in dropped_headers: + self.send_header(key, value) + self.end_headers() + # Do not pass a fixed chunk size here. httpx accumulates bytes until + # that size is reached, which can hide small SSE heartbeat frames + # from Claude Code for minutes during a slow artifact/tool call. + # 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 response_chunks: + if chunk: + if first_byte_ms is None: + first_byte_ms = round((time.monotonic() - started) * 1000) + self.wfile.write(chunk) + self.wfile.flush() + chunks += 1 + bytes_relayed += len(chunk) + _diagnostic_log( + "response_complete", + request_id=diagnostic_id, + status=resp.status_code, + chunks=chunks, + bytes=bytes_relayed, + first_byte_ms=first_byte_ms, + elapsed_ms=round((time.monotonic() - started) * 1000), + ) + except (BrokenPipeError, ConnectionResetError): + # Client (Claude Code) closed the connection mid-response — routine on + # cancelled turns / SSE teardown. Nothing left to relay to, so stop + # quietly rather than crashing the handler thread. + _diagnostic_log( + "client_disconnect", + request_id=diagnostic_id, + phase="response", + chunks=chunks, + bytes=bytes_relayed, + elapsed_ms=round((time.monotonic() - started) * 1000), + ) + return + except httpx.HTTPError as exc: + # Upstream dropped mid-stream. Headers (and status) may already be + # sent, so we can't reliably signal a fresh error — stop and let the + # client see a truncated stream rather than corrupt the framing. + _diagnostic_log( + "upstream_stream_error", + request_id=diagnostic_id, + error_type=type(exc).__name__, + status=resp.status_code, + chunks=chunks, + bytes=bytes_relayed, + elapsed_ms=round((time.monotonic() - started) * 1000), + ) + return + + # Forward every method: this is a transparent pass-through, so routing any + # `do_` lookup to `_handle` lets the gateway reject unsupported methods. + def __getattr__(self, name: str): + if name.startswith("do_"): + return self._handle + raise AttributeError(name) + + +def _start_proxy( + workspace: str, + profile: str | None, + port: int, + token_header: str, + force_refresh_near_expiry: bool, + *, + handler_class: type[_ProxyHandler] = _ProxyHandler, + handler_attributes: dict[str, object] | None = None, +) -> tuple[ThreadingHTTPServer, _TokenCache, httpx.Client]: + """Start the loopback refresh proxy + its background token refresher. + + Binds ``port``, falling back to a fresh OS-assigned port when it is already + in use (e.g. a prior session's proxy that was killed before its teardown ran + still holds the socket). The caller reads ``server.server_address[1]`` for the + actual port and points Claude Code at it. + + Returns (server, cache, client); the caller runs the server (e.g. in a + thread) and calls shutdown()/cache.stop()/client.close() on exit. + """ + upstream_base = f"{workspace.rstrip('/')}/ai-gateway/anthropic/" + cache = _TokenCache( + workspace, + profile, + force_refresh_near_expiry=force_refresh_near_expiry, + ) + # One pooled, keep-alive client shared across handler threads: reuses TCP+TLS + # 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 = cast( + type[BaseHTTPRequestHandler], + type( + "BoundProxyHandler", + (handler_class,), + { + "cache": cache, + "client": client, + "token_header": token_header, + **(handler_attributes or {}), + }, + ), + ) + try: + 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((LOOPBACK_HOST, 0), handler) + + refresher = threading.Thread(target=cache.run_refresher, daemon=True) + refresher.start() + return server, cache, client + + +_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 _AnthropicGatewayHandler(_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 _response_chunks(self, resp: httpx.Response) -> tuple[Iterable[bytes], frozenset[str]]: + 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 super()._response_chunks(resp) + body = self.anthropic_model_aliases.prefix_model_ids(resp.read()) + # resp.read() decodes compression; rewritten JSON is uncompressed. + return (body,), frozenset({"content-encoding"}) + + +def start_proxy( + workspace: str, + profile: str | None, + port: int, + token_header: str, + force_refresh_near_expiry: bool, +): + return _start_proxy( + workspace, + profile, + port, + token_header, + force_refresh_near_expiry, + handler_class=_AnthropicGatewayHandler, + handler_attributes={"anthropic_model_aliases": _AnthropicModelAliases()}, + ) diff --git a/src/ucode/anthropic_model_discovery_proxy.py b/src/ucode/anthropic_model_discovery_proxy.py deleted file mode 100644 index c22e17b..0000000 --- a/src/ucode/anthropic_model_discovery_proxy.py +++ /dev/null @@ -1,127 +0,0 @@ -"""Anthropic model discovery transformations for the gateway proxy.""" - -from __future__ import annotations - -import json -import threading -from collections.abc import Iterable -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 _response_chunks(self, resp: httpx.Response) -> tuple[Iterable[bytes], frozenset[str]]: - 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 super()._response_chunks(resp) - body = self.anthropic_model_aliases.prefix_model_ids(resp.read()) - # resp.read() decodes compression; rewritten JSON is uncompressed. - return (body,), frozenset({"content-encoding"}) - - -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()}, - ) diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index 9d232fc..9f78b6f 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -24,13 +24,10 @@ import threading import time import uuid -from collections.abc import Iterable 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 @@ -215,18 +212,12 @@ 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 _response_chunks(self, resp: httpx.Response) -> tuple[Iterable[bytes], frozenset[str]]: - return resp.iter_raw(), frozenset() - 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, body = self._transform_request(body) + url = self.path.lstrip("/") _diagnostic_log( "request_start", request_id=diagnostic_id, @@ -313,13 +304,9 @@ def _relay_response( bytes_relayed = 0 first_byte_ms: int | None = None try: - # The upstream request has completed through response headers before - # this hook selects raw streaming or a buffered response body. - response_chunks, dropped_headers = self._response_chunks(resp) self.send_response(resp.status_code) for key, value in resp.headers.items(): - header_name = key.lower() - if header_name not in _HOP_BY_HOP and header_name not in dropped_headers: + if key.lower() not in _HOP_BY_HOP: self.send_header(key, value) self.end_headers() # Do not pass a fixed chunk size here. httpx accumulates bytes until @@ -328,7 +315,7 @@ 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 response_chunks: + for chunk in resp.iter_raw(): if chunk: if first_byte_ms is None: first_byte_ms = round((time.monotonic() - started) * 1000) @@ -387,9 +374,6 @@ def start_proxy( port: int, token_header: str, force_refresh_near_expiry: bool, - *, - handler_class: type[_ProxyHandler] = _ProxyHandler, - handler_attributes: dict[str, object] | None = None, ) -> tuple[ThreadingHTTPServer, _TokenCache, httpx.Client]: """Start the loopback refresh proxy + its background token refresher. @@ -411,25 +395,18 @@ 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 = cast( - type[BaseHTTPRequestHandler], - type( - "BoundProxyHandler", - (handler_class,), - { - "cache": cache, - "client": client, - "token_header": token_header, - **(handler_attributes or {}), - }, - ), + + handler = type( + "BoundProxyHandler", + (_ProxyHandler,), + {"cache": cache, "client": client, "token_header": token_header}, ) try: - server = ThreadingHTTPServer((LOOPBACK_HOST, port), handler) + server = ThreadingHTTPServer(("127.0.0.1", 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((LOOPBACK_HOST, 0), handler) + server = ThreadingHTTPServer(("127.0.0.1", 0), handler) refresher = threading.Thread(target=cache.run_refresher, daemon=True) refresher.start() diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 1c1bb18..e66ed27 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -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_anthropic_model_discovery_proxy", start_proxy) + monkeypatch.setattr(claude, "start_anthropic_gateway_proxy", start_proxy) monkeypatch.setattr(claude.subprocess, "Popen", Process) with pytest.raises(SystemExit) as exc: diff --git a/tests/test_anthropic_model_discovery_proxy.py b/tests/test_anthropic_gateway_proxy.py similarity index 87% rename from tests/test_anthropic_model_discovery_proxy.py rename to tests/test_anthropic_gateway_proxy.py index 7fff900..bff28cc 100644 --- a/tests/test_anthropic_model_discovery_proxy.py +++ b/tests/test_anthropic_gateway_proxy.py @@ -5,7 +5,7 @@ import io import json -from ucode import anthropic_model_discovery_proxy +from ucode import anthropic_gateway_proxy class _FakeResponse: @@ -61,20 +61,20 @@ def flush(self): def _handler(wfile, path="/v1/models", command="GET"): - handler = object.__new__(anthropic_model_discovery_proxy._AnthropicModelDiscoveryHandler) + handler = object.__new__(anthropic_gateway_proxy._AnthropicGatewayHandler) handler.wfile = wfile handler.request_version = "HTTP/1.1" handler.requestline = f"{command} {path} HTTP/1.1" handler.command = command handler.path = path handler._headers_buffer = [] - handler.anthropic_model_aliases = anthropic_model_discovery_proxy._AnthropicModelAliases() + handler.anthropic_model_aliases = anthropic_gateway_proxy._AnthropicModelAliases() return handler class TestAnthropicModelAliases: def test_prefixes_custom_model_ids_without_changing_display_name(self): - aliases = anthropic_model_discovery_proxy._AnthropicModelAliases() + aliases = anthropic_gateway_proxy._AnthropicModelAliases() body = json.dumps( { "data": [ @@ -107,7 +107,7 @@ def test_prefixes_custom_model_ids_without_changing_display_name(self): } def test_rewrites_known_alias_in_messages_body(self): - aliases = anthropic_model_discovery_proxy._AnthropicModelAliases() + aliases = anthropic_gateway_proxy._AnthropicModelAliases() aliases.prefix_model_ids(b'{"data":[{"id":"catalog.schema.custom"}]}') body = aliases.rewrite_body( @@ -118,7 +118,7 @@ def test_rewrites_known_alias_in_messages_body(self): assert json.loads(body) == {"model": "catalog.schema.custom", "messages": []} def test_rewrites_known_alias_in_pagination_cursor(self): - aliases = anthropic_model_discovery_proxy._AnthropicModelAliases() + aliases = anthropic_gateway_proxy._AnthropicModelAliases() aliases.prefix_model_ids(b'{"data":[{"id":"catalog.schema.custom"}]}') assert ( @@ -129,13 +129,13 @@ def test_rewrites_known_alias_in_pagination_cursor(self): ) def test_ignores_non_anthropic_models_path(self): - aliases = anthropic_model_discovery_proxy._AnthropicModelAliases() + aliases = anthropic_gateway_proxy._AnthropicModelAliases() path = "/codex/v1/models?after_id=anthropic-aigw-catalog.schema.custom" assert aliases.rewrite_path(path) == path def test_does_not_strip_unknown_prefixed_id(self): - aliases = anthropic_model_discovery_proxy._AnthropicModelAliases() + aliases = anthropic_gateway_proxy._AnthropicModelAliases() unknown = "anthropic-aigw-legitimate-upstream-id" assert aliases.rewrite_path(f"/v1/models?after_id={unknown}") == ( @@ -147,11 +147,11 @@ def test_does_not_strip_unknown_prefixed_id(self): ) def test_leaves_malformed_discovery_response_unchanged(self): - aliases = anthropic_model_discovery_proxy._AnthropicModelAliases() + aliases = anthropic_gateway_proxy._AnthropicModelAliases() assert aliases.prefix_model_ids(b"not-json") == b"not-json" -class TestAnthropicModelDiscoveryHandler: +class TestAnthropicGatewayHandler: def test_inherits_relayed_auth_and_prefixes_models(self): out = _Collect() handler = _handler(out) @@ -232,17 +232,14 @@ def start(*args, **kwargs): call["kwargs"] = kwargs return "server", "cache", "client" - monkeypatch.setattr(anthropic_model_discovery_proxy.gateway_proxy, "start_proxy", start) + monkeypatch.setattr(anthropic_gateway_proxy, "_start_proxy", start) - result = anthropic_model_discovery_proxy.start_proxy("workspace", "profile", 1, "header", False) + result = anthropic_gateway_proxy.start_proxy("workspace", "profile", 1, "header", False) assert result == ("server", "cache", "client") assert call["args"][:5] == ("workspace", "profile", 1, "header", False) - assert ( - call["kwargs"]["handler_class"] - is anthropic_model_discovery_proxy._AnthropicModelDiscoveryHandler - ) + assert call["kwargs"]["handler_class"] is anthropic_gateway_proxy._AnthropicGatewayHandler assert isinstance( call["kwargs"]["handler_attributes"]["anthropic_model_aliases"], - anthropic_model_discovery_proxy._AnthropicModelAliases, + anthropic_gateway_proxy._AnthropicModelAliases, ) diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py index 498aab4..3b51466 100644 --- a/tests/test_gateway_proxy.py +++ b/tests/test_gateway_proxy.py @@ -77,9 +77,6 @@ def iter_raw(self, chunk_size=None): self.chunk_sizes.append(chunk_size) yield from self._chunks - def read(self): - return b"".join(self._chunks) - class _BrokenPipeWriter(io.RawIOBase): """A wfile stand-in that raises BrokenPipeError on write, mimicking a client @@ -172,20 +169,6 @@ def flush(self): assert b"Transfer-Encoding" not in blob # hop-by-hop, stripped assert resp.chunk_sizes == [None] # relay each upstream SSE/network chunk immediately - def test_base_handler_does_not_transform_model_response(self): - out = _Collect() - handler = _relay_handler(out) - response = _FakeResponse( - 200, - {"Content-Encoding": "gzip"}, - [b"compressed-model-response"], - ) - - handler._relay_response(response) - - assert b"Content-Encoding: gzip" in bytes(out.data) - assert b"compressed-model-response" in bytes(out.data) - def test_diagnostics_identify_upstream_mid_stream_drop(self, monkeypatch, capsys): monkeypatch.setenv(gateway_proxy._DIAGNOSTICS_ENV, "1") @@ -345,11 +328,9 @@ class _FakeClient: def __init__(self, responses): self._responses = list(responses) self.sent_tokens: list[str | None] = [] - self.sent_urls: list[str] = [] def stream(self, _method, _url, headers, content): self.sent_tokens.append(headers.get(gateway_proxy.AI_GATEWAY_TOKEN_HEADER)) - self.sent_urls.append(_url) return self._responses.pop(0) @@ -424,7 +405,6 @@ def test_success_first_try_does_not_refresh(self): _handle_handler(client, cache, out)._handle() assert cache.refreshed == 0 assert client.sent_tokens == ["Bearer tok1"] - assert client.sent_urls == ["v1/messages"] assert b"hi" in bytes(out.data) @@ -459,9 +439,6 @@ def run_refresher(self): bound = server.server_address[1] assert bound != busy_port # fell back to a different, free port assert bound != 0 - assert str(client.base_url) == ( - "https://x.staging.cloud.databricks.com/ai-gateway/anthropic/" - ) finally: server.server_close() client.close() From fb38030b1d5c213dbcceaa92042db255a9eb822e Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Wed, 26 Aug 2026 17:35:41 +0000 Subject: [PATCH 9/9] Inline Anthropic proxy startup --- src/ucode/anthropic_gateway_proxy.py | 89 ++++++++------------------- tests/test_anthropic_gateway_proxy.py | 34 +++++----- 2 files changed, 46 insertions(+), 77 deletions(-) diff --git a/src/ucode/anthropic_gateway_proxy.py b/src/ucode/anthropic_gateway_proxy.py index c7cebb1..39798f5 100644 --- a/src/ucode/anthropic_gateway_proxy.py +++ b/src/ucode/anthropic_gateway_proxy.py @@ -383,61 +383,6 @@ def __getattr__(self, name: str): raise AttributeError(name) -def _start_proxy( - workspace: str, - profile: str | None, - port: int, - token_header: str, - force_refresh_near_expiry: bool, - *, - handler_class: type[_ProxyHandler] = _ProxyHandler, - handler_attributes: dict[str, object] | None = None, -) -> tuple[ThreadingHTTPServer, _TokenCache, httpx.Client]: - """Start the loopback refresh proxy + its background token refresher. - - Binds ``port``, falling back to a fresh OS-assigned port when it is already - in use (e.g. a prior session's proxy that was killed before its teardown ran - still holds the socket). The caller reads ``server.server_address[1]`` for the - actual port and points Claude Code at it. - - Returns (server, cache, client); the caller runs the server (e.g. in a - thread) and calls shutdown()/cache.stop()/client.close() on exit. - """ - upstream_base = f"{workspace.rstrip('/')}/ai-gateway/anthropic/" - cache = _TokenCache( - workspace, - profile, - force_refresh_near_expiry=force_refresh_near_expiry, - ) - # One pooled, keep-alive client shared across handler threads: reuses TCP+TLS - # 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 = cast( - type[BaseHTTPRequestHandler], - type( - "BoundProxyHandler", - (handler_class,), - { - "cache": cache, - "client": client, - "token_header": token_header, - **(handler_attributes or {}), - }, - ), - ) - try: - 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((LOOPBACK_HOST, 0), handler) - - refresher = threading.Thread(target=cache.run_refresher, daemon=True) - refresher.start() - return server, cache, client - - _MODEL_ALIAS_PREFIX = "anthropic-aigw-" _ANTHROPIC_MODELS_PATH = "/v1/models" _ANTHROPIC_MESSAGES_PATH = "/v1/messages" @@ -541,13 +486,33 @@ def start_proxy( port: int, token_header: str, force_refresh_near_expiry: bool, -): - return _start_proxy( +) -> tuple[ThreadingHTTPServer, _TokenCache, httpx.Client]: + """Start the Anthropic loopback proxy and token refresher.""" + upstream_base = f"{workspace.rstrip('/')}/ai-gateway/anthropic/" + cache = _TokenCache( workspace, profile, - port, - token_header, - force_refresh_near_expiry, - handler_class=_AnthropicGatewayHandler, - handler_attributes={"anthropic_model_aliases": _AnthropicModelAliases()}, + force_refresh_near_expiry=force_refresh_near_expiry, + ) + client = httpx.Client(base_url=upstream_base, timeout=_UPSTREAM_TIMEOUT, follow_redirects=False) + handler = cast( + type[BaseHTTPRequestHandler], + type( + "BoundProxyHandler", + (_AnthropicGatewayHandler,), + { + "cache": cache, + "client": client, + "token_header": token_header, + "anthropic_model_aliases": _AnthropicModelAliases(), + }, + ), ) + try: + server = ThreadingHTTPServer((LOOPBACK_HOST, port), handler) + except OSError: + server = ThreadingHTTPServer((LOOPBACK_HOST, 0), handler) + + refresher = threading.Thread(target=cache.run_refresher, daemon=True) + refresher.start() + return server, cache, client diff --git a/tests/test_anthropic_gateway_proxy.py b/tests/test_anthropic_gateway_proxy.py index bff28cc..6a2f7e9 100644 --- a/tests/test_anthropic_gateway_proxy.py +++ b/tests/test_anthropic_gateway_proxy.py @@ -225,21 +225,25 @@ def test_strips_known_alias_from_message_request(self): def test_start_proxy_uses_discovery_handler(monkeypatch): - call = {} + class _StubCache: + def run_refresher(self): + return None - def start(*args, **kwargs): - call["args"] = args - call["kwargs"] = kwargs - return "server", "cache", "client" + cache = _StubCache() + monkeypatch.setattr(anthropic_gateway_proxy, "_TokenCache", lambda *_args, **_kwargs: cache) - monkeypatch.setattr(anthropic_gateway_proxy, "_start_proxy", start) - - result = anthropic_gateway_proxy.start_proxy("workspace", "profile", 1, "header", False) - - assert result == ("server", "cache", "client") - assert call["args"][:5] == ("workspace", "profile", 1, "header", False) - assert call["kwargs"]["handler_class"] is anthropic_gateway_proxy._AnthropicGatewayHandler - assert isinstance( - call["kwargs"]["handler_attributes"]["anthropic_model_aliases"], - anthropic_gateway_proxy._AnthropicModelAliases, + server, actual_cache, client = anthropic_gateway_proxy.start_proxy( + "https://workspace.example.com", "profile", 0, "header", False ) + try: + handler = server.RequestHandlerClass + assert issubclass(handler, anthropic_gateway_proxy._AnthropicGatewayHandler) + assert handler.cache is cache + assert isinstance( + handler.anthropic_model_aliases, + anthropic_gateway_proxy._AnthropicModelAliases, + ) + assert actual_cache is cache + finally: + server.server_close() + client.close()