From 03269bbbf828eaa8528f57d50a67d3a732348420 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:46:14 +0000 Subject: [PATCH 01/15] claude: authenticate gateway model discovery --- src/ucode/agents/claude.py | 4 +++- tests/test_agent_claude.py | 2 ++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 9f4df255..3cc29304 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -1024,7 +1024,9 @@ def launch(state: dict, tool_args: list[str]) -> None: _launch_relayed(state, binary, tool_args) return if workspace: - os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) + token = get_databricks_token(workspace, state.get("profile")) + os.environ["OAUTH_TOKEN"] = token + os.environ["ANTHROPIC_AUTH_TOKEN"] = token exec_or_spawn(_build_claude_argv(binary, tool_args)) diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 37b74012..faeef0fb 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -662,6 +662,7 @@ def fake_execvp(binary: str, args: list[str]) -> None: raise RuntimeError("stop") monkeypatch.delenv("OAUTH_TOKEN", raising=False) + monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False) monkeypatch.setattr( claude, "get_databricks_token", lambda workspace, profile=None: "fresh-token" ) @@ -673,6 +674,7 @@ def fake_execvp(binary: str, args: list[str]) -> None: assert str(exc) == "stop" assert os.environ["OAUTH_TOKEN"] == "fresh-token" + assert os.environ["ANTHROPIC_AUTH_TOKEN"] == "fresh-token" assert exec_calls == [ ( "claude", From a427c1d5072c302a9a9fba7dc58f500f13c93589 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Thu, 20 Aug 2026 19:46:23 +0000 Subject: [PATCH 02/15] claude: initialize gateway before bootstrap --- src/ucode/agents/claude.py | 2 ++ tests/test_agent_claude.py | 4 ++++ 2 files changed, 6 insertions(+) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 3cc29304..1dc9caa9 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -1027,6 +1027,8 @@ def launch(state: dict, tool_args: list[str]) -> None: token = get_databricks_token(workspace, state.get("profile")) os.environ["OAUTH_TOKEN"] = token os.environ["ANTHROPIC_AUTH_TOKEN"] = token + os.environ["ANTHROPIC_BASE_URL"] = build_tool_base_url("claude", workspace) + os.environ["CLAUDE_CODE_USE_GATEWAY"] = "1" exec_or_spawn(_build_claude_argv(binary, tool_args)) diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index faeef0fb..f899e7c2 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -663,6 +663,8 @@ def fake_execvp(binary: str, args: list[str]) -> None: monkeypatch.delenv("OAUTH_TOKEN", raising=False) 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, "get_databricks_token", lambda workspace, profile=None: "fresh-token" ) @@ -675,6 +677,8 @@ def fake_execvp(binary: str, args: list[str]) -> None: assert os.environ["OAUTH_TOKEN"] == "fresh-token" assert os.environ["ANTHROPIC_AUTH_TOKEN"] == "fresh-token" + assert os.environ["ANTHROPIC_BASE_URL"] == f"{WS}/ai-gateway/anthropic" + assert os.environ["CLAUDE_CODE_USE_GATEWAY"] == "1" assert exec_calls == [ ( "claude", From 4a2fdede9d17f3f18d9fd3498a6cd7f0900309a9 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Thu, 20 Aug 2026 21:41:03 +0000 Subject: [PATCH 03/15] Refresh Claude gateway credentials through proxy --- src/ucode/agents/claude.py | 52 ++++++++++++++++++++++----- src/ucode/gateway_proxy.py | 25 +++++++------ tests/test_agent_claude.py | 71 +++++++++++++++++++++++++++---------- tests/test_gateway_proxy.py | 7 ++++ 4 files changed, 118 insertions(+), 37 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 1dc9caa9..e6996ca8 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -27,7 +27,6 @@ from ucode.databricks import ( build_auth_shell_command, build_tool_base_url, - get_databricks_token, ) from ucode.launcher import exec_or_spawn from ucode.managed_files import OS, current_os, write_managed_file @@ -875,7 +874,12 @@ def _merge_claude_settings(base: dict, overlay: dict) -> dict: return merged -def _build_claude_argv(binary: str, tool_args: list[str], relayed: bool = False) -> list[str]: +def _build_claude_argv( + binary: str, + tool_args: list[str], + relayed: bool = False, + settings_override: dict | None = None, +) -> list[str]: """Build the ``claude`` argv, composing any caller ``--settings`` with ucode's managed settings. @@ -898,7 +902,7 @@ def _build_claude_argv(binary: str, tool_args: list[str], relayed: bool = False) """ source_args = ["--setting-sources", _RELAYED_SETTING_SOURCES] if relayed else [] caller_values, remaining = _extract_caller_settings(tool_args) - if not caller_values: + if not caller_values and settings_override is None: # No caller --settings: hand Claude ucode's settings file directly (the # common path; behavior unchanged). return [binary, *source_args, "--settings", str(CLAUDE_SETTINGS_PATH), *tool_args] @@ -908,6 +912,8 @@ def _build_claude_argv(binary: str, tool_args: list[str], relayed: bool = False) # ucode wins over the caller for conflicting keys (protects gateway auth); # hooks from both sides survive. merged = _merge_claude_settings(caller_settings, read_json_safe(CLAUDE_SETTINGS_PATH)) + if settings_override is not None: + merged = _merge_claude_settings(merged, settings_override) return [ binary, *source_args, @@ -1017,6 +1023,39 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: raise SystemExit(returncode) +def _launch_gateway(state: dict, binary: str, tool_args: list[str]) -> None: + from ucode.gateway_proxy import AUTHORIZATION_HEADER, start_proxy + + workspace = state["workspace"] + server, cache, client = start_proxy( + workspace, state.get("profile"), 0, token_header=AUTHORIZATION_HEADER + ) + 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["CLAUDE_CODE_USE_GATEWAY"] = "1" + + server_thread = threading.Thread(target=server.serve_forever, daemon=True) + server_thread.start() + settings_override = { + "env": {"ANTHROPIC_BASE_URL": os.environ["ANTHROPIC_BASE_URL"]}, + } + proc = subprocess.Popen( + _build_claude_argv(binary, tool_args, settings_override=settings_override) + ) + try: + returncode = proc.wait() + except KeyboardInterrupt: + proc.send_signal(signal.SIGINT) + returncode = proc.wait() + finally: + cache.stop() + server.shutdown() + client.close() + raise SystemExit(returncode) + + def launch(state: dict, tool_args: list[str]) -> None: binary = SPEC["binary"] workspace = state.get("workspace") @@ -1024,11 +1063,8 @@ def launch(state: dict, tool_args: list[str]) -> None: _launch_relayed(state, binary, tool_args) return if workspace: - token = get_databricks_token(workspace, state.get("profile")) - os.environ["OAUTH_TOKEN"] = token - os.environ["ANTHROPIC_AUTH_TOKEN"] = token - os.environ["ANTHROPIC_BASE_URL"] = build_tool_base_url("claude", workspace) - os.environ["CLAUDE_CODE_USE_GATEWAY"] = "1" + _launch_gateway(state, binary, tool_args) + return exec_or_spawn(_build_claude_argv(binary, tool_args)) diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index de1589df..aae71073 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -34,6 +34,7 @@ # Header we overwrite with the freshly-minted Databricks credential. Any # client-supplied value is replaced, so a stale settings.json value can't leak. _SWAP_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() @@ -50,9 +51,6 @@ "content-length", ) ) -# Request headers the proxy manages itself and must never forward on: hop-by-hop -# plus the swap header (replaced with a freshly-minted value per request). -_STRIP_ON_FORWARD = _HOP_BY_HOP | {_SWAP_HEADER.lower()} # 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. @@ -181,11 +179,14 @@ def stop(self) -> None: self._stop.set() -def _forwarded_request_headers(handler: BaseHTTPRequestHandler, token: str) -> dict[str, str]: +def _forwarded_request_headers( + handler: BaseHTTPRequestHandler, token: str, token_header: str = _SWAP_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 + key: value for key, value in handler.headers.items() if key.lower() not in strip_on_forward } - headers[_SWAP_HEADER] = f"Bearer {token}" + headers[token_header] = f"Bearer {token}" return headers @@ -193,6 +194,7 @@ class _ProxyHandler(BaseHTTPRequestHandler): # Set by the server factory. cache: _TokenCache client: httpx.Client + token_header = _SWAP_HEADER def log_message(self, format: str, *args: object) -> None: return @@ -219,7 +221,7 @@ def _handle(self) -> None: ) try: # First attempt with the current token. - headers = _forwarded_request_headers(self, self.cache.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", @@ -249,7 +251,7 @@ def _handle(self) -> None: # 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) + 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", @@ -362,7 +364,10 @@ def __getattr__(self, name: str): def start_proxy( - workspace: str, profile: str | None, port: int + workspace: str, + profile: str | None, + port: int, + token_header: str = _SWAP_HEADER, ) -> tuple[ThreadingHTTPServer, _TokenCache, httpx.Client]: """Start the loopback refresh proxy + its background token refresher. @@ -384,7 +389,7 @@ def start_proxy( handler = type( "BoundProxyHandler", (_ProxyHandler,), - {"cache": cache, "client": client}, + {"cache": cache, "client": client, "token_header": token_header}, ) try: server = ThreadingHTTPServer(("127.0.0.1", port), handler) diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index f899e7c2..7c7d2096 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -654,36 +654,69 @@ def boom(name, entry, scope=mcp_mod.MCP_USER_SCOPE): class TestClaudeLaunch: - def test_sets_oauth_token_before_exec(self, monkeypatch): - exec_calls: list[tuple[str, list[str]]] = [] + def test_runs_through_refresh_proxy(self, monkeypatch): + import ucode.gateway_proxy as gateway_proxy - def fake_execvp(binary: str, args: list[str]) -> None: - exec_calls.append((binary, args)) - raise RuntimeError("stop") + calls: list[tuple] = [] + + class Server: + server_address = ("127.0.0.1", 12345) + + def serve_forever(self): + calls.append(("serve",)) + + def shutdown(self): + calls.append(("shutdown",)) + + class Cache: + token = "fresh-token" + + def stop(self): + calls.append(("stop",)) + + class Client: + def close(self): + calls.append(("close",)) + + class Process: + def __init__(self, argv): + calls.append(("popen", argv)) + + def wait(self): + return 0 + + def start_proxy(workspace, profile, port, token_header): + calls.append(("proxy", workspace, profile, port, token_header)) + return Server(), Cache(), Client() monkeypatch.delenv("OAUTH_TOKEN", raising=False) 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, "get_databricks_token", lambda workspace, profile=None: "fresh-token" - ) - monkeypatch.setattr(os, "execvp", fake_execvp) + monkeypatch.setattr(gateway_proxy, "start_proxy", start_proxy) + monkeypatch.setattr(claude.subprocess, "Popen", Process) - try: - claude.launch({"workspace": WS}, ["--debug"]) - except RuntimeError as exc: - assert str(exc) == "stop" + with pytest.raises(SystemExit) as exc: + claude.launch({"workspace": WS, "profile": "test"}, ["--debug"]) + assert exc.value.code == 0 assert os.environ["OAUTH_TOKEN"] == "fresh-token" assert os.environ["ANTHROPIC_AUTH_TOKEN"] == "fresh-token" - assert os.environ["ANTHROPIC_BASE_URL"] == f"{WS}/ai-gateway/anthropic" + assert os.environ["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:12345" assert os.environ["CLAUDE_CODE_USE_GATEWAY"] == "1" - assert exec_calls == [ - ( - "claude", - ["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"], - ) + assert calls[:2] == [ + ("proxy", WS, "test", 0, gateway_proxy.AUTHORIZATION_HEADER), + ("serve",), + ] + assert calls[2][0] == "popen" + argv = calls[2][1] + assert argv[:2] == ["claude", "--settings"] + assert json.loads(argv[2])["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:12345" + assert argv[3:] == ["--debug"] + assert calls[3:] == [ + ("stop",), + ("shutdown",), + ("close",), ] diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py index f103aca2..c438bc7e 100644 --- a/tests/test_gateway_proxy.py +++ b/tests/test_gateway_proxy.py @@ -47,6 +47,13 @@ def test_overwrites_client_supplied_swap_header(self): out = gateway_proxy._forwarded_request_headers(handler, "fresh") assert out["X-Databricks-AI-Gateway-Token"] == "Bearer fresh" + def test_overwrites_authorization_header(self): + handler = _FakeHandler({"Authorization": "Bearer stale"}) + out = gateway_proxy._forwarded_request_headers( + handler, "fresh", gateway_proxy.AUTHORIZATION_HEADER + ) + assert out["Authorization"] == "Bearer fresh" + def test_strips_hop_by_hop_headers(self): handler = _FakeHandler( {"Host": "localhost:9", "Content-Length": "5", "Connection": "keep-alive"} From 45665109494c79b5416b58528e17e214d2797f55 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Thu, 20 Aug 2026 23:53:04 +0000 Subject: [PATCH 04/15] Force refresh expiring gateway tokens --- src/ucode/gateway_proxy.py | 14 ++++++-------- tests/test_gateway_proxy.py | 13 ++++++------- 2 files changed, 12 insertions(+), 15 deletions(-) diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index aae71073..806eae0f 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -126,13 +126,11 @@ def __init__(self, workspace: str, profile: str | None) -> None: # Force on start so we begin on a full-TTL token rather than inheriting a # near-expiry one cached from an earlier CLI call. Raises if auth is dead # (surfaced by the caller at launch, before Claude Code starts). - self._refresh(force=True) + self._refresh() - def _refresh(self, *, force: bool) -> None: - """Mint a token and record its expiry. Caller holds `_refresh_lock` (or is - __init__). Non-force lets a token another process just refreshed satisfy - this call from the shared cache with no write — shrinking lock contention.""" - token = get_databricks_token(self._workspace, self._profile, force_refresh=force) + def _refresh(self) -> None: + """Force-mint a token and record its expiry.""" + token = get_databricks_token(self._workspace, self._profile, force_refresh=True) expiry = _jwt_exp(token) or (time.time() + _DEFAULT_TTL_S) with self._state_lock: self._token = token @@ -149,7 +147,7 @@ def _ensure_fresh(self) -> None: if self._fresh_enough(): # another thread refreshed while we waited return try: - self._refresh(force=False) + self._refresh() except RuntimeError as exc: # Keep serving the current token; a request that then 401s triggers # a forced refresh + retry (see _ProxyHandler._handle). @@ -164,7 +162,7 @@ def token(self) -> str: def refresh(self) -> None: """Force a fresh mint now (used by the retry-on-401 path).""" with self._refresh_lock: - self._refresh(force=True) + self._refresh() def run_refresher(self) -> None: while not self._stop.wait(_REFRESHER_POLL_S): diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py index c438bc7e..eaff0ad2 100644 --- a/tests/test_gateway_proxy.py +++ b/tests/test_gateway_proxy.py @@ -247,15 +247,14 @@ def test_fresh_token_is_not_refreshed(self, monkeypatch): _ = cache.token assert state["forces"] == [True] # no extra mint while fresh - def test_near_expiry_triggers_nonforce_refresh(self, monkeypatch): - # First mint expires within the buffer -> reading .token refreshes once, - # non-force (so a token another process just wrote can satisfy it). + def test_near_expiry_triggers_forced_refresh(self, monkeypatch): + # First mint expires within the buffer, so reading .token force-refreshes it. state = _install_fake_token(monkeypatch, [100, 5000]) cache = gateway_proxy._TokenCache("ws", None) _ = cache.token - assert state["forces"] == [True, False] + assert state["forces"] == [True, True] _ = cache.token # now fresh again - assert state["forces"] == [True, False] + assert state["forces"] == [True, True] def test_refresh_is_single_flighted(self, monkeypatch): # A burst of concurrent requests at the expiry boundary must trigger ONE @@ -267,8 +266,8 @@ def test_refresh_is_single_flighted(self, monkeypatch): t.start() for t in threads: t.join() - # 1 forced init + exactly 1 non-force refresh shared by all 10 readers. - assert state["forces"] == [True, False] + # 1 forced init + exactly 1 forced refresh shared by all 10 readers. + assert state["forces"] == [True, True] def test_ensure_fresh_keeps_token_when_refresh_fails(self, monkeypatch): _install_fake_token(monkeypatch, [5000]) From 5b1e481d27d9e8eb0a4680f55e4a7a8bc026af7e Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Fri, 21 Aug 2026 16:45:25 +0000 Subject: [PATCH 05/15] Gate Claude gateway refresh proxy --- src/ucode/agents/claude.py | 15 +++++++++---- src/ucode/gateway_proxy.py | 45 +++++++++++++++++++++++-------------- tests/test_agent_claude.py | 34 +++++++++++++++++++++++----- tests/test_gateway_proxy.py | 19 +++++++++++----- 4 files changed, 81 insertions(+), 32 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index e6996ca8..cf0614eb 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -27,7 +27,9 @@ from ucode.databricks import ( build_auth_shell_command, build_tool_base_url, + get_databricks_token, ) +from ucode.gateway_proxy import 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 ( @@ -43,6 +45,7 @@ CLAUDE_CONFIG_DIR = Path.home() / ".claude" CLAUDE_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "ucode-settings.json" CLAUDE_BACKUP_PATH = APP_DIR / "claude-ucode-settings.backup.json" +GATEWAY_MODEL_DISCOVERY_ENV_VAR = "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY" SPEC: ToolSpec = { "binary": "claude", @@ -1024,11 +1027,13 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: def _launch_gateway(state: dict, binary: str, tool_args: list[str]) -> None: - from ucode.gateway_proxy import AUTHORIZATION_HEADER, start_proxy - workspace = state["workspace"] server, cache, client = start_proxy( - workspace, state.get("profile"), 0, token_header=AUTHORIZATION_HEADER + workspace, + state.get("profile"), + 0, + token_header=AUTHORIZATION_HEADER, + force_refresh_near_expiry=True, ) token = cache.token os.environ["OAUTH_TOKEN"] = token @@ -1062,9 +1067,11 @@ def launch(state: dict, tool_args: list[str]) -> None: if state.get("claude_relayed"): _launch_relayed(state, binary, tool_args) return - if workspace: + if workspace and os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1": _launch_gateway(state, binary, tool_args) return + if workspace: + os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) exec_or_spawn(_build_claude_argv(binary, tool_args)) diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index 806eae0f..f5a58bae 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -1,18 +1,17 @@ -"""Loopback refresh proxy for relayed Anthropic (Claude Max/Team/Enterprise). +"""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. That Databricks token is short-lived and a static settings.json header -can't be refreshed, so `ucode claude` points `ANTHROPIC_BASE_URL` at this proxy -instead: it forwards every request to the workspace gateway unchanged except for -adding a freshly-minted swap header, and streams the response back verbatim. +header. Native gateway discovery instead carries the Databricks credential in +`Authorization`. The proxy refreshes the applicable header and streams responses +back verbatim. 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 and never read, stored, or logged. + passed through untouched in relayed mode and never logged. """ from __future__ import annotations @@ -115,9 +114,16 @@ class _TokenCache: boundary triggers exactly one CLI call, not a thundering herd on the shared token cache.""" - def __init__(self, workspace: str, profile: str | None) -> None: + 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() @@ -126,11 +132,11 @@ def __init__(self, workspace: str, profile: str | None) -> None: # Force on start so we begin on a full-TTL token rather than inheriting a # near-expiry one cached from an earlier CLI call. Raises if auth is dead # (surfaced by the caller at launch, before Claude Code starts). - self._refresh() + self._refresh(force=True) - def _refresh(self) -> None: - """Force-mint a token and record its expiry.""" - token = get_databricks_token(self._workspace, self._profile, force_refresh=True) + 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 @@ -147,7 +153,7 @@ def _ensure_fresh(self) -> None: if self._fresh_enough(): # another thread refreshed while we waited return try: - self._refresh() + 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). @@ -162,7 +168,7 @@ def token(self) -> str: def refresh(self) -> None: """Force a fresh mint now (used by the retry-on-401 path).""" with self._refresh_lock: - self._refresh() + self._refresh(force=True) def run_refresher(self) -> None: while not self._stop.wait(_REFRESHER_POLL_S): @@ -234,9 +240,9 @@ def _handle(self) -> None: # Auth rejected. Drain the (small) error body so the pooled # connection can be reused, then fall through to one retry. resp.read() - # A 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 swap token and retry once. If it was the + # 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. @@ -366,6 +372,7 @@ def start_proxy( profile: str | None, port: int, token_header: str = _SWAP_HEADER, + force_refresh_near_expiry: bool = False, ) -> tuple[ThreadingHTTPServer, _TokenCache, httpx.Client]: """Start the loopback refresh proxy + its background token refresher. @@ -378,7 +385,11 @@ def start_proxy( thread) and calls shutdown()/cache.stop()/client.close() on exit. """ upstream_base = f"{workspace.rstrip('/')}/ai-gateway/anthropic/" - cache = _TokenCache(workspace, profile) + 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. diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 7c7d2096..597c62ab 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -654,9 +654,21 @@ def boom(name, entry, scope=mcp_mod.MCP_USER_SCOPE): class TestClaudeLaunch: - def test_runs_through_refresh_proxy(self, monkeypatch): - import ucode.gateway_proxy as gateway_proxy + def test_default_launch_keeps_existing_auth_path(self, monkeypatch): + calls: list[list[str]] = [] + monkeypatch.delenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, raising=False) + monkeypatch.delenv("OAUTH_TOKEN", raising=False) + monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") + monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) + claude.launch({"workspace": WS, "profile": "test"}, ["--debug"]) + + assert os.environ["OAUTH_TOKEN"] == "token" + assert calls == [ + ["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"] + ] + + def test_runs_through_refresh_proxy(self, monkeypatch): calls: list[tuple] = [] class Server: @@ -685,15 +697,25 @@ def __init__(self, argv): def wait(self): return 0 - def start_proxy(workspace, profile, port, token_header): - calls.append(("proxy", workspace, profile, port, token_header)) + def start_proxy(workspace, profile, port, token_header, force_refresh_near_expiry): + calls.append( + ( + "proxy", + workspace, + profile, + port, + token_header, + force_refresh_near_expiry, + ) + ) return Server(), Cache(), Client() + monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "1") monkeypatch.delenv("OAUTH_TOKEN", raising=False) monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False) monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False) monkeypatch.delenv("CLAUDE_CODE_USE_GATEWAY", raising=False) - monkeypatch.setattr(gateway_proxy, "start_proxy", start_proxy) + monkeypatch.setattr(claude, "start_proxy", start_proxy) monkeypatch.setattr(claude.subprocess, "Popen", Process) with pytest.raises(SystemExit) as exc: @@ -705,7 +727,7 @@ def start_proxy(workspace, profile, port, token_header): assert os.environ["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:12345" assert os.environ["CLAUDE_CODE_USE_GATEWAY"] == "1" assert calls[:2] == [ - ("proxy", WS, "test", 0, gateway_proxy.AUTHORIZATION_HEADER), + ("proxy", WS, "test", 0, claude.AUTHORIZATION_HEADER, True), ("serve",), ] assert calls[2][0] == "popen" diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py index eaff0ad2..77ecdda8 100644 --- a/tests/test_gateway_proxy.py +++ b/tests/test_gateway_proxy.py @@ -247,20 +247,25 @@ def test_fresh_token_is_not_refreshed(self, monkeypatch): _ = cache.token assert state["forces"] == [True] # no extra mint while fresh - def test_near_expiry_triggers_forced_refresh(self, monkeypatch): - # First mint expires within the buffer, so reading .token force-refreshes it. + def test_near_expiry_preserves_default_nonforce_refresh(self, monkeypatch): state = _install_fake_token(monkeypatch, [100, 5000]) cache = gateway_proxy._TokenCache("ws", None) _ = cache.token - assert state["forces"] == [True, True] + assert state["forces"] == [True, False] _ = cache.token # now fresh again + assert state["forces"] == [True, False] + + def test_near_expiry_can_force_refresh(self, monkeypatch): + state = _install_fake_token(monkeypatch, [100, 5000]) + cache = gateway_proxy._TokenCache("ws", None, force_refresh_near_expiry=True) + _ = cache.token assert state["forces"] == [True, True] def test_refresh_is_single_flighted(self, monkeypatch): # A burst of concurrent requests at the expiry boundary must trigger ONE # refresh, not a thundering herd on the shared token cache. state = _install_fake_token(monkeypatch, [100, 5000], delay=0.05) - cache = gateway_proxy._TokenCache("ws", None) + cache = gateway_proxy._TokenCache("ws", None, force_refresh_near_expiry=True) threads = [threading.Thread(target=lambda: cache.token) for _ in range(10)] for t in threads: t.start() @@ -411,7 +416,11 @@ class _StubCache: def run_refresher(self): return None - monkeypatch.setattr(gateway_proxy, "_TokenCache", lambda workspace, profile: _StubCache()) + monkeypatch.setattr( + gateway_proxy, + "_TokenCache", + lambda workspace, profile, **_kwargs: _StubCache(), + ) # Occupy a port to simulate the leftover proxy holding it. occupied = socket.socket(socket.AF_INET, socket.SOCK_STREAM) occupied.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) From e907b40eb50f13d665f1bb769f31326403cae4b2 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Fri, 21 Aug 2026 17:14:10 +0000 Subject: [PATCH 06/15] Format Claude launch test --- tests/test_agent_claude.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 597c62ab..6b497823 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -664,9 +664,7 @@ def test_default_launch_keeps_existing_auth_path(self, monkeypatch): claude.launch({"workspace": WS, "profile": "test"}, ["--debug"]) assert os.environ["OAUTH_TOKEN"] == "token" - assert calls == [ - ["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"] - ] + assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"]] def test_runs_through_refresh_proxy(self, monkeypatch): calls: list[tuple] = [] From 3c90a00c42248c686883878854cd294a961c1252 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Fri, 21 Aug 2026 19:45:37 +0000 Subject: [PATCH 07/15] Preserve relayed token refresh behavior --- src/ucode/gateway_proxy.py | 7 +++---- tests/test_gateway_proxy.py | 10 +++++----- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index f5a58bae..272a80a0 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -129,10 +129,9 @@ def __init__( self._stop = threading.Event() self._token = "" self._expiry = 0.0 - # Force on start so we begin on a full-TTL token rather than inheriting a - # near-expiry one cached from an earlier CLI call. Raises if auth is dead - # (surfaced by the caller at launch, before Claude Code starts). - self._refresh(force=True) + # 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.""" diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py index 77ecdda8..a9851747 100644 --- a/tests/test_gateway_proxy.py +++ b/tests/test_gateway_proxy.py @@ -235,25 +235,25 @@ def fake(_ws, _profile, force_refresh=False): class TestTokenCache: - def test_initial_mint_is_forced(self, monkeypatch): + def test_initial_mint_preserves_default_nonforce_refresh(self, monkeypatch): state = _install_fake_token(monkeypatch, [5000]) gateway_proxy._TokenCache("ws", None) - assert state["forces"] == [True] # full-TTL start + assert state["forces"] == [False] def test_fresh_token_is_not_refreshed(self, monkeypatch): state = _install_fake_token(monkeypatch, [5000]) cache = gateway_proxy._TokenCache("ws", None) _ = cache.token _ = cache.token - assert state["forces"] == [True] # no extra mint while fresh + assert state["forces"] == [False] # no extra mint while fresh def test_near_expiry_preserves_default_nonforce_refresh(self, monkeypatch): state = _install_fake_token(monkeypatch, [100, 5000]) cache = gateway_proxy._TokenCache("ws", None) _ = cache.token - assert state["forces"] == [True, False] + assert state["forces"] == [False, False] _ = cache.token # now fresh again - assert state["forces"] == [True, False] + assert state["forces"] == [False, False] def test_near_expiry_can_force_refresh(self, monkeypatch): state = _install_fake_token(monkeypatch, [100, 5000]) From 4f60a222148093e588c63b518ec70925b2370627 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:34:54 +0000 Subject: [PATCH 08/15] Use Claude model discovery flag name --- src/ucode/agents/claude.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index cf0614eb..ba2bbcde 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -45,7 +45,7 @@ CLAUDE_CONFIG_DIR = Path.home() / ".claude" CLAUDE_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "ucode-settings.json" CLAUDE_BACKUP_PATH = APP_DIR / "claude-ucode-settings.backup.json" -GATEWAY_MODEL_DISCOVERY_ENV_VAR = "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY" +GATEWAY_MODEL_DISCOVERY_ENV_VAR = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" SPEC: ToolSpec = { "binary": "claude", From f9e8b36f045758394fba23ccd4966a2b745cab69 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Mon, 24 Aug 2026 20:01:24 +0000 Subject: [PATCH 09/15] Restore ucode model discovery opt-in flag --- src/ucode/agents/claude.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index ba2bbcde..cf0614eb 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -45,7 +45,7 @@ CLAUDE_CONFIG_DIR = Path.home() / ".claude" CLAUDE_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "ucode-settings.json" CLAUDE_BACKUP_PATH = APP_DIR / "claude-ucode-settings.backup.json" -GATEWAY_MODEL_DISCOVERY_ENV_VAR = "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" +GATEWAY_MODEL_DISCOVERY_ENV_VAR = "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY" SPEC: ToolSpec = { "binary": "claude", From dabbd5ac5940160b1423f5bdface815f24d3695f 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:28:20 +0000 Subject: [PATCH 10/15] fix(claude): require proxy auth options --- src/ucode/agents/claude.py | 13 ++++++++----- src/ucode/gateway_proxy.py | 12 +++++++----- tests/test_gateway_proxy.py | 8 ++++++-- 3 files changed, 21 insertions(+), 12 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index cf0614eb..a4b14359 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -29,7 +29,7 @@ build_tool_base_url, get_databricks_token, ) -from ucode.gateway_proxy import AUTHORIZATION_HEADER, start_proxy +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 ( @@ -41,7 +41,6 @@ from ucode.tracing import tracing_env from ucode.ui import print_err, print_note, print_success, print_warning -GATEWAY_MODEL_DISCOVERY_ENV_VAR = "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY" CLAUDE_CONFIG_DIR = Path.home() / ".claude" CLAUDE_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "ucode-settings.json" CLAUDE_BACKUP_PATH = APP_DIR / "claude-ucode-settings.backup.json" @@ -975,8 +974,6 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: """Relayed launch: sign into the Claude subscription, start the loopback refresh proxy, then run Claude Code alongside it (the proxy must outlive the exec, so we spawn-and-wait rather than replacing the process).""" - from ucode.gateway_proxy import start_proxy - conflict = _managed_relayed_conflicts() if conflict is not None: managed_path, keys = conflict @@ -1002,7 +999,13 @@ 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(workspace, state.get("profile"), port) + server, cache, client = start_proxy( + workspace, + state.get("profile"), + port, + token_header=AI_GATEWAY_TOKEN_HEADER, + force_refresh_near_expiry=False, + ) # start_proxy falls back to an OS-assigned port when the cached one is taken # (stale proxy from a killed session). Reconcile settings + state to whatever # it actually bound, so Claude Code connects to the live port. diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index 272a80a0..9f78b6ff 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -32,7 +32,7 @@ # Header we overwrite with the freshly-minted Databricks credential. Any # client-supplied value is replaced, so a stale settings.json value can't leak. -_SWAP_HEADER = "X-Databricks-AI-Gateway-Token" +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( @@ -183,7 +183,9 @@ def stop(self) -> None: def _forwarded_request_headers( - handler: BaseHTTPRequestHandler, token: str, token_header: str = _SWAP_HEADER + handler: BaseHTTPRequestHandler, + token: str, + token_header: str = AI_GATEWAY_TOKEN_HEADER, ) -> dict[str, str]: strip_on_forward = _HOP_BY_HOP | {token_header.lower()} headers = { @@ -197,7 +199,7 @@ class _ProxyHandler(BaseHTTPRequestHandler): # Set by the server factory. cache: _TokenCache client: httpx.Client - token_header = _SWAP_HEADER + token_header = AI_GATEWAY_TOKEN_HEADER def log_message(self, format: str, *args: object) -> None: return @@ -370,8 +372,8 @@ def start_proxy( workspace: str, profile: str | None, port: int, - token_header: str = _SWAP_HEADER, - force_refresh_near_expiry: bool = False, + token_header: str, + force_refresh_near_expiry: bool, ) -> tuple[ThreadingHTTPServer, _TokenCache, httpx.Client]: """Start the loopback refresh proxy + its background token refresher. diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py index a9851747..3b514666 100644 --- a/tests/test_gateway_proxy.py +++ b/tests/test_gateway_proxy.py @@ -330,7 +330,7 @@ def __init__(self, responses): self.sent_tokens: list[str | None] = [] def stream(self, _method, _url, headers, content): - self.sent_tokens.append(headers.get(gateway_proxy._SWAP_HEADER)) + self.sent_tokens.append(headers.get(gateway_proxy.AI_GATEWAY_TOKEN_HEADER)) return self._responses.pop(0) @@ -429,7 +429,11 @@ def run_refresher(self): busy_port = occupied.getsockname()[1] try: server, _cache, client = gateway_proxy.start_proxy( - "https://x.staging.cloud.databricks.com", None, busy_port + "https://x.staging.cloud.databricks.com", + None, + busy_port, + token_header=gateway_proxy.AI_GATEWAY_TOKEN_HEADER, + force_refresh_near_expiry=False, ) try: bound = server.server_address[1] From 0ecb32a682a628ff7dc63805c46896558f07ff9f 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 11/15] 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 9f78b6ff..b8bcab0e 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 3b514666..3856894d 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 86edc0185e029364ca6f8aaf111cbd051cc98895 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 12/15] 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 b8bcab0e..2fc0310f 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 3856894d..800b92fa 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 a7969f9a30e9a5fc022ddf8fa7df220b42959e5f 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 13/15] 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 2fc0310f..79a7362b 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 800b92fa..0aa431dd 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 8ce6d9412621810610ae6463405014b03ef92c4a 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 14/15] 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 a4b14359..81ce0467 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 79a7362b..8a4eb749 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 00000000..3e9ff673 --- /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 0aa431dd..96f3320f 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 a65df484b4cfec86adbebdc109dad08b5dd3a8eb 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 15/15] 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 8a4eb749..00152f68 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 96f3320f..9c6ab9b7 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):