From 3273aba9c0507e89997d98997b06be3b3242375f Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Fri, 11 Sep 2026 14:04:57 +0000 Subject: [PATCH 1/5] Proxy: route Databricks-hosted models to gateway auth in relayed sessions Lets a relayed (subscription-relay) Claude Code session also reach Databricks-hosted OSS / system.ai models. The loopback proxy now inspects each request's `model`: a namespace-qualified Databricks id is re-routed to gateway auth (Databricks token in `Authorization`, MPS + swap headers dropped), while a bare Anthropic subscription id keeps today's OAuth-passthrough relay. So one session can switch between the relayed Enterprise subscription and Databricks models via `/model`. Gated by start_proxy(hybrid_oss_routing=...), on only for the relayed launch; a pure-relay session sends only bare Anthropic ids, so behavior is unchanged. Users surface Databricks models through Claude Code's existing modelPicker / --model. Toward AIGTWY-4490. Co-authored-by: Isaac --- src/ucode/agents/claude.py | 5 ++ src/ucode/gateway_proxy.py | 76 ++++++++++++++++++-- tests/test_agent_claude.py | 6 +- tests/test_gateway_proxy.py | 139 ++++++++++++++++++++++++++++++++++++ 4 files changed, 221 insertions(+), 5 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index a77013d1..c695ca1e 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -1283,6 +1283,11 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: port, token_header=gateway_proxy.AI_GATEWAY_TOKEN_HEADER, force_refresh_near_expiry=False, + # Let a relayed session also reach Databricks-hosted (OSS / system.ai) models: + # the proxy re-routes those requests to gateway auth while relayed subscription + # models keep the OAuth passthrough. Bare Anthropic ids are unaffected, so a + # pure-relay session behaves exactly as before. + hybrid_oss_routing=True, ) # 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 diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index e5cd4788..d5de7ea8 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -7,6 +7,11 @@ `Authorization`. The proxy refreshes the applicable header and streams responses back verbatim. +In a hybrid relayed session the proxy picks per request by the requested model: +Databricks-hosted ids (system.ai / OSS) take the gateway-auth path while relayed +subscription models keep the OAuth passthrough, so one Claude Code session can use +both. + 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, @@ -34,6 +39,9 @@ # 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" +# Header that routes a request to a specific Model Provider Service. Dropped when a +# request is re-routed to a Databricks-hosted model so the gateway serves it directly. +MODEL_PROVIDER_SERVICE_HEADER = "Databricks-Model-Provider-Service" # Hop-by-hop headers must not be forwarded across a proxy. HOP_BY_HOP_HEADERS = frozenset( h.lower() @@ -186,8 +194,9 @@ def forwarded_request_headers( handler: BaseHTTPRequestHandler, token: str, token_header: str = AI_GATEWAY_TOKEN_HEADER, + extra_strip: frozenset[str] = frozenset(), ) -> dict[str, str]: - strip_on_forward = HOP_BY_HOP_HEADERS | {token_header.lower()} + strip_on_forward = HOP_BY_HOP_HEADERS | {token_header.lower()} | extra_strip headers = { key: value for key, value in handler.headers.items() if key.lower() not in strip_on_forward } @@ -195,11 +204,47 @@ def forwarded_request_headers( return headers +# On the Databricks-hosted path the gateway credential goes in `Authorization` (so the +# caller's Anthropic OAuth is replaced), and the swap + MPS headers are dropped so the +# gateway serves the model directly instead of relaying to the subscription MPS. +_DATABRICKS_ROUTE_STRIP = frozenset( + {AI_GATEWAY_TOKEN_HEADER.lower(), MODEL_PROVIDER_SERVICE_HEADER.lower()} +) + + +def is_databricks_routed_model(model: str | None) -> bool: + """True when ``model`` is a Databricks-hosted (gateway-served) id rather than a model + the relayed Anthropic subscription serves. + + Databricks ids are namespace-qualified (``system.ai.*``, ``catalog.schema.model``, + ``databricks-*``); the relayed subscription uses Anthropic's bare canonical names + (``claude-opus-4-1``, ``claude-sonnet-4-5``, ...), which never carry a dot.""" + if not model: + return False + return "." in model or model.startswith("databricks-") + + +def _request_model(body: bytes | None) -> str | None: + """The ``model`` field of a JSON request body, or None when absent/unparseable.""" + if not body: + return None + try: + payload = json.loads(body) + except (ValueError, TypeError): + return None + model = payload.get("model") if isinstance(payload, dict) else None + return model if isinstance(model, str) else None + + class _ProxyHandler(BaseHTTPRequestHandler): # Set by the server factory. cache: TokenCache client: httpx.Client token_header = AI_GATEWAY_TOKEN_HEADER + # When True, requests for a Databricks-hosted model are re-routed to gateway auth + # (Databricks token in `Authorization`) so a relayed session can also reach OSS / + # system.ai models; relayed subscription models keep the OAuth-passthrough path. + hybrid_oss_routing = False def log_message(self, format: str, *args: object) -> None: return @@ -218,15 +263,32 @@ def _handle(self) -> None: length = int(self.headers.get("Content-Length", 0) or 0) body = self.rfile.read(length) if length else None url = self.path.lstrip("/") + # Databricks-hosted models (in a hybrid relayed session) authenticate with the + # gateway token in `Authorization`; everything else keeps the relay path. + route_databricks = self.hybrid_oss_routing and is_databricks_routed_model( + _request_model(body) + ) log_proxy_diagnostic( "request_start", request_id=diagnostic_id, method=self.command, path=self.path.split("?", 1)[0], + route="databricks" if route_databricks else "relay", ) + + def request_headers() -> dict[str, str]: + if route_databricks: + return forwarded_request_headers( + self, + self.cache.token, + AUTHORIZATION_HEADER, + extra_strip=_DATABRICKS_ROUTE_STRIP, + ) + return forwarded_request_headers(self, self.cache.token, self.token_header) + try: # First attempt with the current token. - headers = forwarded_request_headers(self, self.cache.token, self.token_header) + headers = request_headers() with self.client.stream(self.command, url, headers=headers, content=body) as resp: log_proxy_diagnostic( "upstream_headers", @@ -256,7 +318,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_token_refresh_failure(exc) - headers = forwarded_request_headers(self, self.cache.token, self.token_header) + headers = request_headers() with self.client.stream(self.command, url, headers=headers, content=body) as resp: log_proxy_diagnostic( "upstream_headers", @@ -374,6 +436,7 @@ def start_proxy( port: int, token_header: str, force_refresh_near_expiry: bool, + hybrid_oss_routing: bool = False, ) -> tuple[ThreadingHTTPServer, TokenCache, httpx.Client]: """Start the loopback refresh proxy + its background token refresher. @@ -399,7 +462,12 @@ def start_proxy( handler = type( "BoundProxyHandler", (_ProxyHandler,), - {"cache": cache, "client": client, "token_header": token_header}, + { + "cache": cache, + "client": client, + "token_header": token_header, + "hybrid_oss_routing": hybrid_oss_routing, + }, ) try: server = ThreadingHTTPServer(("127.0.0.1", port), handler) diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 92f99fe2..f42ae5d8 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -1094,7 +1094,9 @@ def __init__(self, argv): def wait(self): return 0 - def start_proxy(workspace, profile, port, token_header, force_refresh_near_expiry): + def start_proxy( + workspace, profile, port, token_header, force_refresh_near_expiry, hybrid_oss_routing + ): calls.append( ( "proxy", @@ -1103,6 +1105,7 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir port, token_header, force_refresh_near_expiry, + hybrid_oss_routing, ) ) return Server(), Cache(), Client() @@ -1132,6 +1135,7 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir 12345, claude.gateway_proxy.AI_GATEWAY_TOKEN_HEADER, False, + True, # hybrid_oss_routing — relayed sessions also reach Databricks-hosted models ) assert calls[-3:] == [("stop",), ("shutdown",), ("close",)] diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py index 28600045..d06e1c38 100644 --- a/tests/test_gateway_proxy.py +++ b/tests/test_gateway_proxy.py @@ -63,6 +63,68 @@ def test_strips_hop_by_hop_headers(self): assert "Content-Length" not in out assert "Connection" not in out + def test_extra_strip_removes_named_headers(self): + handler = _FakeHandler({"Databricks-Model-Provider-Service": "cat.s.mps", "Keep": "me"}) + out = gateway_proxy.forwarded_request_headers( + handler, "t", extra_strip=frozenset({"databricks-model-provider-service"}) + ) + assert "Databricks-Model-Provider-Service" not in out + assert out["Keep"] == "me" + + def test_databricks_route_swaps_authorization_and_drops_relay_headers(self): + # OSS path: the Databricks token replaces the caller's OAuth in Authorization, and + # the swap + MPS headers are dropped so the gateway serves the model directly. + handler = _FakeHandler( + { + "Authorization": "Bearer anthropic-oauth", + "X-Databricks-AI-Gateway-Token": "Bearer stale-swap", + "Databricks-Model-Provider-Service": "cat.s.relayed_mps", + } + ) + out = gateway_proxy.forwarded_request_headers( + handler, + "dbx-token", + gateway_proxy.AUTHORIZATION_HEADER, + extra_strip=gateway_proxy._DATABRICKS_ROUTE_STRIP, + ) + assert out["Authorization"] == "Bearer dbx-token" + assert "X-Databricks-AI-Gateway-Token" not in out + assert "Databricks-Model-Provider-Service" not in out + + +class TestIsDatabricksRoutedModel: + def test_namespace_qualified_ids_route_to_databricks(self): + for model in ( + "system.ai.claude-opus-4-8", + "system.ai.gpt-oss-120b", + "my_catalog.my_schema.my_model", + "databricks-meta-llama-3-3-70b-instruct", + ): + assert gateway_proxy.is_databricks_routed_model(model), model + + def test_bare_anthropic_ids_relay(self): + for model in ("claude-opus-4-1", "claude-sonnet-4-5", "claude-3-7-sonnet-20250219"): + assert not gateway_proxy.is_databricks_routed_model(model), model + + def test_missing_model_relays(self): + assert not gateway_proxy.is_databricks_routed_model(None) + assert not gateway_proxy.is_databricks_routed_model("") + + +class TestRequestModel: + def test_extracts_model_from_json_body(self): + assert gateway_proxy._request_model(b'{"model": "system.ai.x", "n": 1}') == "system.ai.x" + + def test_none_for_missing_body(self): + assert gateway_proxy._request_model(None) is None + + def test_none_for_invalid_json(self): + assert gateway_proxy._request_model(b"not json") is None + + def test_none_for_non_object_or_non_string_model(self): + assert gateway_proxy._request_model(b"[1, 2]") is None + assert gateway_proxy._request_model(b'{"model": 5}') is None + class _FakeResponse: """Stand-in for httpx.Response exposing only what `_relay_response` reads.""" @@ -328,9 +390,11 @@ class _FakeClient: def __init__(self, responses): self._responses = list(responses) self.sent_tokens: list[str | None] = [] + self.sent_headers: list[dict] = [] def stream(self, _method, _url, headers, content): self.sent_tokens.append(headers.get(gateway_proxy.AI_GATEWAY_TOKEN_HEADER)) + self.sent_headers.append(headers) return self._responses.pop(0) @@ -473,3 +537,78 @@ def run_refresher(self): client.close() finally: occupied.close() + + +def _hybrid_handler(client, cache, wfile, *, headers, body, hybrid) -> gateway_proxy._ProxyHandler: + h = object.__new__(gateway_proxy._ProxyHandler) + h.client = client + h.cache = cache + h.hybrid_oss_routing = hybrid + hdrs = dict(headers) + hdrs["Content-Length"] = str(len(body)) + h.headers = hdrs + h.rfile = io.BytesIO(body) + h.path = "/v1/messages" + h.command = "POST" + h.wfile = wfile + h.request_version = "HTTP/1.1" + h.requestline = "POST /v1/messages HTTP/1.1" + h._headers_buffer = [] + return h + + +class TestHybridOssRouting: + _CLIENT_HEADERS = { + "Authorization": "Bearer anthropic-oauth", + "Databricks-Model-Provider-Service": "cat.s.relayed_mps", + } + + def test_databricks_model_routes_to_gateway_auth(self): + # A Databricks-hosted model in a hybrid relayed session: the gateway token + # replaces the OAuth in Authorization, and the swap + MPS headers are dropped. + client = _FakeClient([_FakeResp(200, b"ok")]) + _hybrid_handler( + client, + _FakeCache(), + _Collect(), + headers=self._CLIENT_HEADERS, + body=b'{"model": "system.ai.gpt-oss-120b"}', + hybrid=True, + )._handle() + sent = client.sent_headers[0] + assert sent["Authorization"] == "Bearer tok1" + assert gateway_proxy.AI_GATEWAY_TOKEN_HEADER not in sent + assert "Databricks-Model-Provider-Service" not in sent + + def test_relayed_model_keeps_oauth_passthrough(self): + # A subscription model still relays: the OAuth is untouched, the swap header + # carries the Databricks token, and the MPS header is preserved. + client = _FakeClient([_FakeResp(200, b"ok")]) + _hybrid_handler( + client, + _FakeCache(), + _Collect(), + headers=self._CLIENT_HEADERS, + body=b'{"model": "claude-opus-4-1"}', + hybrid=True, + )._handle() + sent = client.sent_headers[0] + assert sent["Authorization"] == "Bearer anthropic-oauth" + assert sent[gateway_proxy.AI_GATEWAY_TOKEN_HEADER] == "Bearer tok1" + assert sent["Databricks-Model-Provider-Service"] == "cat.s.relayed_mps" + + def test_routing_off_relays_even_a_databricks_model(self): + # With hybrid routing disabled (a pure-relay session) nothing is re-routed, + # so behavior is identical to before this feature. + client = _FakeClient([_FakeResp(200, b"ok")]) + _hybrid_handler( + client, + _FakeCache(), + _Collect(), + headers=self._CLIENT_HEADERS, + body=b'{"model": "system.ai.gpt-oss-120b"}', + hybrid=False, + )._handle() + sent = client.sent_headers[0] + assert sent["Authorization"] == "Bearer anthropic-oauth" + assert sent[gateway_proxy.AI_GATEWAY_TOKEN_HEADER] == "Bearer tok1" From 33f220953f05a7b88b94d4de2ab27de42b9cc567 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Fri, 11 Sep 2026 14:14:30 +0000 Subject: [PATCH 2/5] Ignore the mlflow claude autolog runtime log Co-authored-by: Isaac --- .gitignore | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.gitignore b/.gitignore index 007899c4..79005ff6 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,5 @@ dist/ .venv/ .DS_Store .isaac/ +# mlflow `autolog claude` Stop-hook runtime log (personal session ids/paths, not source). +.claude/mlflow/ From 0117740fa4dedfa6c58728add5f3bd449c5060bf Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Fri, 11 Sep 2026 17:13:03 +0000 Subject: [PATCH 3/5] Rename hybrid_oss_routing to relayed_oss_routing Co-authored-by: Isaac --- src/ucode/agents/claude.py | 2 +- src/ucode/gateway_proxy.py | 12 ++++++------ tests/test_agent_claude.py | 6 +++--- tests/test_gateway_proxy.py | 24 +++++++++++++----------- 4 files changed, 23 insertions(+), 21 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index c695ca1e..6ac184b6 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -1287,7 +1287,7 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: # the proxy re-routes those requests to gateway auth while relayed subscription # models keep the OAuth passthrough. Bare Anthropic ids are unaffected, so a # pure-relay session behaves exactly as before. - hybrid_oss_routing=True, + relayed_oss_routing=True, ) # 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 diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index d5de7ea8..1f69aabd 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -7,7 +7,7 @@ `Authorization`. The proxy refreshes the applicable header and streams responses back verbatim. -In a hybrid relayed session the proxy picks per request by the requested model: +With relayed OSS-routing on, the proxy picks per request by the requested model: Databricks-hosted ids (system.ai / OSS) take the gateway-auth path while relayed subscription models keep the OAuth passthrough, so one Claude Code session can use both. @@ -244,7 +244,7 @@ class _ProxyHandler(BaseHTTPRequestHandler): # When True, requests for a Databricks-hosted model are re-routed to gateway auth # (Databricks token in `Authorization`) so a relayed session can also reach OSS / # system.ai models; relayed subscription models keep the OAuth-passthrough path. - hybrid_oss_routing = False + relayed_oss_routing = False def log_message(self, format: str, *args: object) -> None: return @@ -263,9 +263,9 @@ def _handle(self) -> None: length = int(self.headers.get("Content-Length", 0) or 0) body = self.rfile.read(length) if length else None url = self.path.lstrip("/") - # Databricks-hosted models (in a hybrid relayed session) authenticate with the + # Databricks-hosted models (when relayed OSS-routing is on) authenticate with the # gateway token in `Authorization`; everything else keeps the relay path. - route_databricks = self.hybrid_oss_routing and is_databricks_routed_model( + route_databricks = self.relayed_oss_routing and is_databricks_routed_model( _request_model(body) ) log_proxy_diagnostic( @@ -436,7 +436,7 @@ def start_proxy( port: int, token_header: str, force_refresh_near_expiry: bool, - hybrid_oss_routing: bool = False, + relayed_oss_routing: bool = False, ) -> tuple[ThreadingHTTPServer, TokenCache, httpx.Client]: """Start the loopback refresh proxy + its background token refresher. @@ -466,7 +466,7 @@ def start_proxy( "cache": cache, "client": client, "token_header": token_header, - "hybrid_oss_routing": hybrid_oss_routing, + "relayed_oss_routing": relayed_oss_routing, }, ) try: diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index f42ae5d8..046a992a 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -1095,7 +1095,7 @@ def wait(self): return 0 def start_proxy( - workspace, profile, port, token_header, force_refresh_near_expiry, hybrid_oss_routing + workspace, profile, port, token_header, force_refresh_near_expiry, relayed_oss_routing ): calls.append( ( @@ -1105,7 +1105,7 @@ def start_proxy( port, token_header, force_refresh_near_expiry, - hybrid_oss_routing, + relayed_oss_routing, ) ) return Server(), Cache(), Client() @@ -1135,7 +1135,7 @@ def start_proxy( 12345, claude.gateway_proxy.AI_GATEWAY_TOKEN_HEADER, False, - True, # hybrid_oss_routing — relayed sessions also reach Databricks-hosted models + True, # relayed_oss_routing — relayed sessions also reach Databricks-hosted models ) assert calls[-3:] == [("stop",), ("shutdown",), ("close",)] diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py index d06e1c38..184d6e4b 100644 --- a/tests/test_gateway_proxy.py +++ b/tests/test_gateway_proxy.py @@ -539,11 +539,13 @@ def run_refresher(self): occupied.close() -def _hybrid_handler(client, cache, wfile, *, headers, body, hybrid) -> gateway_proxy._ProxyHandler: +def _relayed_oss_handler( + client, cache, wfile, *, headers, body, enabled +) -> gateway_proxy._ProxyHandler: h = object.__new__(gateway_proxy._ProxyHandler) h.client = client h.cache = cache - h.hybrid_oss_routing = hybrid + h.relayed_oss_routing = enabled hdrs = dict(headers) hdrs["Content-Length"] = str(len(body)) h.headers = hdrs @@ -557,23 +559,23 @@ def _hybrid_handler(client, cache, wfile, *, headers, body, hybrid) -> gateway_p return h -class TestHybridOssRouting: +class TestRelayedOssRouting: _CLIENT_HEADERS = { "Authorization": "Bearer anthropic-oauth", "Databricks-Model-Provider-Service": "cat.s.relayed_mps", } def test_databricks_model_routes_to_gateway_auth(self): - # A Databricks-hosted model in a hybrid relayed session: the gateway token + # A Databricks-hosted model with relayed OSS-routing on: the gateway token # replaces the OAuth in Authorization, and the swap + MPS headers are dropped. client = _FakeClient([_FakeResp(200, b"ok")]) - _hybrid_handler( + _relayed_oss_handler( client, _FakeCache(), _Collect(), headers=self._CLIENT_HEADERS, body=b'{"model": "system.ai.gpt-oss-120b"}', - hybrid=True, + enabled=True, )._handle() sent = client.sent_headers[0] assert sent["Authorization"] == "Bearer tok1" @@ -584,13 +586,13 @@ def test_relayed_model_keeps_oauth_passthrough(self): # A subscription model still relays: the OAuth is untouched, the swap header # carries the Databricks token, and the MPS header is preserved. client = _FakeClient([_FakeResp(200, b"ok")]) - _hybrid_handler( + _relayed_oss_handler( client, _FakeCache(), _Collect(), headers=self._CLIENT_HEADERS, body=b'{"model": "claude-opus-4-1"}', - hybrid=True, + enabled=True, )._handle() sent = client.sent_headers[0] assert sent["Authorization"] == "Bearer anthropic-oauth" @@ -598,16 +600,16 @@ def test_relayed_model_keeps_oauth_passthrough(self): assert sent["Databricks-Model-Provider-Service"] == "cat.s.relayed_mps" def test_routing_off_relays_even_a_databricks_model(self): - # With hybrid routing disabled (a pure-relay session) nothing is re-routed, + # With relayed OSS-routing off (a pure-relay session) nothing is re-routed, # so behavior is identical to before this feature. client = _FakeClient([_FakeResp(200, b"ok")]) - _hybrid_handler( + _relayed_oss_handler( client, _FakeCache(), _Collect(), headers=self._CLIENT_HEADERS, body=b'{"model": "system.ai.gpt-oss-120b"}', - hybrid=False, + enabled=False, )._handle() sent = client.sent_headers[0] assert sent["Authorization"] == "Bearer anthropic-oauth" From 26b62945652878006abc62d4259d799044dc7a37 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Fri, 11 Sep 2026 19:02:47 +0000 Subject: [PATCH 4/5] Extend relayed e2e test to also serve a Databricks-hosted model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Runs the relayed proxy with relayed_oss_routing on and, after the subscription relay check, POSTs a namespace-qualified Databricks-hosted model (a discovered OSS id when offered, else a system.ai one) through the same proxy — asserting it serves (200). Proves one relayed session reaches both the subscription and Databricks models. Skips gracefully when the workspace advertises no such model. Co-authored-by: Isaac --- tests/test_e2e.py | 57 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 831e6603..a3ef9f64 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -587,6 +587,27 @@ def _first_relayed_service(tool: str, workspace: str, token: str) -> str: pytest.skip(f"no relayed {tool} model provider services available on this workspace") return names[0] + @staticmethod + def _first_databricks_hosted_model(workspace: str, token: str) -> str | None: + """A namespace-qualified (Databricks-hosted) model id from the anthropic gateway + catalog — one the relayed subscription doesn't serve, so it exercises the proxy's + per-model Databricks re-route. Prefers a non-Claude (OSS) id when one is offered.""" + from ucode.gateway_proxy import is_databricks_routed_model + + try: + resp = httpx.get( + f"{build_tool_base_url('claude', workspace)}/v1/models", + headers={"Authorization": f"Bearer {token}"}, + timeout=30, + ) + resp.raise_for_status() + ids = [m.get("id") for m in resp.json().get("data", [])] + except (httpx.HTTPError, ValueError, KeyError): + return None + qualified = [i for i in ids if i and is_databricks_routed_model(i)] + oss = [i for i in qualified if "claude" not in i] + return (oss or qualified or [None])[0] + @staticmethod def _skip_if_provider_unusable(combined: str, provider: str) -> None: # Environmental provider-account conditions, not ucode bugs: the test only proves routing @@ -700,6 +721,11 @@ def test_launch_claude_through_relayed_provider( via CLAUDE_CODE_OAUTH_TOKEN (`claude setup-token` output); without it the launch would open an interactive browser login, so the test skips. Also needs a relayed MPS on the workspace, so it stays inert until both exist. + + Also asserts the hybrid path: through the same proxy (relayed_oss_routing on), + a Databricks-hosted model the subscription doesn't serve is re-routed to gateway + auth and served — so one relayed session reaches both the subscription and + Databricks models. """ import ucode.config_io as config_io_mod from ucode import gateway_proxy @@ -725,15 +751,21 @@ def test_launch_claude_through_relayed_provider( # Start the real loopback refresh proxy exactly as `_launch_relayed` does, # so the request is credential-swapped and relayed like a live session. + # relayed_oss_routing lets the same proxy also serve Databricks-hosted models. server, cache, client = gateway_proxy.start_proxy( e2e_workspace, None, 0, token_header=gateway_proxy.AI_GATEWAY_TOKEN_HEADER, force_refresh_near_expiry=False, + relayed_oss_routing=True, ) port = server.server_address[1] threading.Thread(target=server.serve_forever, daemon=True).start() + # A Databricks-hosted model the subscription doesn't serve, to exercise the + # proxy's per-model Databricks re-route from within the relayed session. + oss_model = self._first_databricks_hosted_model(e2e_workspace, e2e_token) + oss_response = None try: state = {**e2e_state, "workspace": e2e_workspace, "relayed_proxy_port": port} with pytest.MonkeyPatch().context() as mp: @@ -745,6 +777,23 @@ def test_launch_claude_through_relayed_provider( "ANTHROPIC_BASE_URL": f"http://127.0.0.1:{port}", } result = _run_agent(claude.validate_cmd("claude"), env=env, timeout=90) + if oss_model is not None: + # A deliberately fake Authorization proves the Databricks route replaced it: + # were this wrongly relayed to the subscription, the bad OAuth would 401. + oss_response = httpx.post( + f"http://127.0.0.1:{port}/v1/messages", + headers={ + "Authorization": "Bearer not-a-real-oauth", + "content-type": "application/json", + "anthropic-version": "2023-06-01", + }, + json={ + "model": oss_model, + "max_tokens": 16, + "messages": [{"role": "user", "content": "say hi in 3 words"}], + }, + timeout=60, + ) finally: cache.stop() server.shutdown() @@ -755,6 +804,14 @@ def test_launch_claude_through_relayed_provider( f"relayed provider={provider} rc={result.returncode} " f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}" ) + # The Databricks-hosted model must serve (200) through the relayed proxy — proof the + # session reaches Databricks models via per-model routing, not just the subscription. + if oss_response is not None: + self._skip_if_provider_unusable(oss_response.text, oss_model) + assert oss_response.status_code == 200, ( + f"Databricks-hosted model {oss_model} via the relayed proxy: " + f"HTTP {oss_response.status_code}: {oss_response.text[:300]}" + ) def test_launch_codex_through_provider( self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token From 1d6edc178b03deb6330f600dc3c4cbf331a340a2 Mon Sep 17 00:00:00 2001 From: Mason Cao Date: Fri, 11 Sep 2026 20:54:22 +0000 Subject: [PATCH 5/5] Fix relayed e2e OSS-route check: use a natively-servable model, tolerate access MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The check picked an `anthropic-aigw-*` catalog alias (deepseek) that is listed but 404s on a direct call — those need their provider-service header, which the Databricks route drops. Exclude those aliases and pick a natively-servable `system.ai.*` id (prefer a non-Claude one). Assertion now fails only on 401 (the fake OAuth was relayed = routing regression); any other non-200 means the route reached the gateway but the CI principal can't serve the model (environmental) and skips; 200 proves it end to end. Co-authored-by: Isaac --- tests/test_e2e.py | 39 +++++++++++++++++++++++++++------------ 1 file changed, 27 insertions(+), 12 deletions(-) diff --git a/tests/test_e2e.py b/tests/test_e2e.py index a3ef9f64..ab35ba08 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -589,9 +589,11 @@ def _first_relayed_service(tool: str, workspace: str, token: str) -> str: @staticmethod def _first_databricks_hosted_model(workspace: str, token: str) -> str | None: - """A namespace-qualified (Databricks-hosted) model id from the anthropic gateway - catalog — one the relayed subscription doesn't serve, so it exercises the proxy's - per-model Databricks re-route. Prefers a non-Claude (OSS) id when one is offered.""" + """A natively-servable Databricks-hosted model id from the anthropic gateway catalog: + namespace-qualified (so the proxy re-routes it to gateway auth) and served directly. + Excludes `anthropic-aigw-*` aliases — they're listed but need their provider-service + header to route, which the Databricks route drops (they 404 on a direct call). Prefers + a non-Claude (OSS) native id when the workspace serves one.""" from ucode.gateway_proxy import is_databricks_routed_model try: @@ -604,9 +606,13 @@ def _first_databricks_hosted_model(workspace: str, token: str) -> str | None: ids = [m.get("id") for m in resp.json().get("data", [])] except (httpx.HTTPError, ValueError, KeyError): return None - qualified = [i for i in ids if i and is_databricks_routed_model(i)] - oss = [i for i in qualified if "claude" not in i] - return (oss or qualified or [None])[0] + native = [ + i + for i in ids + if i and is_databricks_routed_model(i) and not i.startswith("anthropic-aigw-") + ] + oss = [i for i in native if "claude" not in i] + return (oss or native or [None])[0] @staticmethod def _skip_if_provider_unusable(combined: str, provider: str) -> None: @@ -804,14 +810,23 @@ def test_launch_claude_through_relayed_provider( f"relayed provider={provider} rc={result.returncode} " f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}" ) - # The Databricks-hosted model must serve (200) through the relayed proxy — proof the - # session reaches Databricks models via per-model routing, not just the subscription. + # Databricks re-route check. The fake OAuth is the tell: 401 means the proxy relayed + # this to the subscription (routing regression) instead of swapping in the gateway + # token — so 401 fails. The relay check above already passed, so the gateway token is + # valid and a 401 here can only be the fake OAuth. 200 proves the model served; any + # other status means the route reached the gateway but the CI principal can't serve + # this model (environmental, not a routing bug) — skip. if oss_response is not None: - self._skip_if_provider_unusable(oss_response.text, oss_model) - assert oss_response.status_code == 200, ( - f"Databricks-hosted model {oss_model} via the relayed proxy: " - f"HTTP {oss_response.status_code}: {oss_response.text[:300]}" + assert oss_response.status_code != 401, ( + f"relayed_oss_routing regressed: {oss_model} was relayed to the subscription " + f"(401) instead of routed to the gateway. Body: {oss_response.text[:200]}" ) + if oss_response.status_code != 200: + pytest.skip( + f"gateway did not serve {oss_model} for the CI principal " + f"(HTTP {oss_response.status_code}); routing reached the gateway but " + f"model access is environmental: {oss_response.text[:200]}" + ) def test_launch_codex_through_provider( self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token