Route Databricks-hosted models to gateway auth in relayed sessions (hybrid OSS + relayed)

.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/

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.
+        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

src/ucode/gateway_proxy.py

@@ -7,6 +7,11 @@
 `Authorization`. The proxy refreshes the applicable header and streams responses
 back verbatim.
 
+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.
+
 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,20 +194,57 @@ 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
     }
     headers[token_header] = f"Bearer {token}"
     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.
+    relayed_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 (when relayed OSS-routing is on) authenticate with the
+        # gateway token in `Authorization`; everything else keeps the relay path.
+        route_databricks = self.relayed_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,
+    relayed_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,
+            "relayed_oss_routing": relayed_oss_routing,
+        },
     )
     try:
         server = ThreadingHTTPServer(("127.0.0.1", port), handler)

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, relayed_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,
+                    relayed_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,  # relayed_oss_routing — relayed sessions also reach Databricks-hosted models
         )
         assert calls[-3:] == [("stop",), ("shutdown",), ("close",)]
 

tests/test_e2e.py

@@ -587,6 +587,33 @@ 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 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:
+            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
+        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:
         # Environmental provider-account conditions, not ucode bugs: the test only proves routing
@@ -700,6 +727,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 +757,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 +783,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 +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}"
         )
+        # 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:
+            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