Fix Claude provider model discovery proxy

src/ucode/agents/claude.py

@@ -334,13 +334,11 @@ def render_overlay(
         "ENABLE_TOOL_SEARCH": "1",
         "CLAUDE_CODE_USE_GATEWAY": "1",
     }
-    # Native /model discovery: picker lists every gateway Messages-API endpoint,
-    # not just the family aliases. Skipped under a provider (its routing header
-    # would send a discovered gateway id to a provider that can't resolve it).
+    # Native /model discovery uses the MPS header to scope provider launches.
     discovery_enabled = (
         os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1" or smart_routing_v2.enabled()
     )
-    if discovery_enabled and not provider:
+    if discovery_enabled:
         env["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1"
     # Intentionally NOT setting ANTHROPIC_MODEL by default. Setting it produces a
     # duplicate catalog row in Claude Code's /model picker (e.g. "Opus 4.8 (1M
@@ -1093,7 +1091,9 @@ def _rewrite_relayed_port(state: dict, port: int) -> None:
         write_json_file(CLAUDE_SETTINGS_PATH, settings)
 
 
-def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None:
+def _launch_relayed(
+    state: dict, binary: str, tool_args: list[str], *, provider: str | None = None
+) -> 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)."""
@@ -1128,6 +1128,7 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None:
         port,
         token_header=AI_GATEWAY_TOKEN_HEADER,
         force_refresh_near_expiry=False,
+        model_provider_service=provider,
     )
     # 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
@@ -1153,7 +1154,12 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None:
 
 
 def _launch_claude_with_gateway_proxy(
-    state: dict, binary: str, tool_args: list[str], *, smart_routing: bool
+    state: dict,
+    binary: str,
+    tool_args: list[str],
+    *,
+    smart_routing: bool,
+    provider: str | None = None,
 ) -> None:
     """Launch Claude through a refreshing gateway proxy."""
     workspace = state["workspace"]
@@ -1163,6 +1169,7 @@ def _launch_claude_with_gateway_proxy(
         0,
         token_header=AUTHORIZATION_HEADER,
         force_refresh_near_expiry=True,
+        model_provider_service=provider,
     )
     token = cache.token
     os.environ["OAUTH_TOKEN"] = token
@@ -1210,8 +1217,14 @@ def compose_gateway_settings(args: list[str]) -> tuple[dict, list[str]]:
 def launch(state: dict, tool_args: list[str]) -> None:
     binary = SPEC["binary"]
     workspace = state.get("workspace")
+    transient_provider = state.get("_claude_launch_provider")
+    provider = (
+        transient_provider
+        if isinstance(transient_provider, str) and transient_provider
+        else get_provider_service(state, "claude")
+    )
     if state.get("claude_relayed"):
-        _launch_relayed(state, binary, tool_args)
+        _launch_relayed(state, binary, tool_args, provider=provider)
         return
     first_prompt_routing = (
         smart_routing_v2.enabled()
@@ -1228,10 +1241,14 @@ def launch(state: dict, tool_args: list[str]) -> None:
             "Please use Codex or disable smart routing."
         )
     if first_prompt_routing:
-        _launch_claude_with_gateway_proxy(state, binary, tool_args, smart_routing=True)
+        _launch_claude_with_gateway_proxy(
+            state, binary, tool_args, smart_routing=True, provider=provider
+        )
         return
     if workspace and os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1":
-        _launch_claude_with_gateway_proxy(state, binary, tool_args, smart_routing=False)
+        _launch_claude_with_gateway_proxy(
+            state, binary, tool_args, smart_routing=False, provider=provider
+        )
         return
     if workspace:
         os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile"))

src/ucode/anthropic_model_discovery_proxy.py

@@ -45,6 +45,7 @@ class _ProxyHandler(BaseHTTPRequestHandler):
     cache: TokenCache
     client: httpx.Client
     token_header = AI_GATEWAY_TOKEN_HEADER
+    model_provider_service: str | None = None
 
     def log_message(self, format: str, *args: object) -> None:
         return
@@ -63,6 +64,19 @@ def _transform_request(self, body: bytes | None) -> tuple[str, bytes | None]:
     def _response_chunks(self, resp: httpx.Response) -> tuple[Iterable[bytes], frozenset[str]]:
         return resp.iter_raw(), frozenset()
 
+    def _forwarded_request_headers(self) -> dict[str, str]:
+        headers = forwarded_request_headers(self, self.cache.token, self.token_header)
+        if (
+            self.command == "GET"
+            and urlsplit(self.path).path == _ANTHROPIC_MODELS_PATH
+            and self.model_provider_service
+        ):
+            # Claude Code does not apply ANTHROPIC_CUSTOM_HEADERS to its native
+            # model-discovery request. Add the configured MPS header at the
+            # loopback boundary so discovery is scoped like inference.
+            headers[_MODEL_PROVIDER_SERVICE_HEADER] = self.model_provider_service
+        return headers
+
     def _should_retry_model_discovery(self, resp: httpx.Response) -> bool:
         return (
             self.command == "GET"
@@ -87,7 +101,7 @@ def _retry_model_discovery(
                 delay_ms=round(delay * 1000),
             )
             time.sleep(delay)
-            headers = forwarded_request_headers(self, self.cache.token, self.token_header)
+            headers = self._forwarded_request_headers()
             with self.client.stream(self.command, url, headers=headers, content=body) as resp:
                 log_proxy_diagnostic(
                     "model_discovery_upstream_headers",
@@ -119,7 +133,7 @@ def _handle(self) -> None:
         )
         try:
             # First attempt with the current token.
-            headers = forwarded_request_headers(self, self.cache.token, self.token_header)
+            headers = self._forwarded_request_headers()
             with self.client.stream(self.command, url, headers=headers, content=body) as resp:
                 log_proxy_diagnostic(
                     "model_discovery_upstream_headers",
@@ -151,7 +165,7 @@ def _handle(self) -> None:
             except RuntimeError as exc:
                 # Still retry with the existing token after reporting the failure.
                 log_token_refresh_failure(exc)
-            headers = forwarded_request_headers(self, self.cache.token, self.token_header)
+            headers = self._forwarded_request_headers()
             with self.client.stream(self.command, url, headers=headers, content=body) as resp:
                 log_proxy_diagnostic(
                     "model_discovery_upstream_headers",
@@ -268,6 +282,7 @@ def __getattr__(self, name: str):
 
 
 _MODEL_ALIAS_PREFIX = "anthropic-aigw-"
+_MODEL_PROVIDER_SERVICE_HEADER = "Databricks-Model-Provider-Service"
 _ANTHROPIC_MODELS_PATH = "/v1/models"
 _ANTHROPIC_MESSAGES_PATH = "/v1/messages"
 
@@ -370,6 +385,7 @@ def start_proxy(
     port: int,
     token_header: str,
     force_refresh_near_expiry: bool,
+    model_provider_service: str | None = None,
 ) -> tuple[ThreadingHTTPServer, TokenCache, httpx.Client]:
     """Start the Anthropic model discovery proxy and token refresher."""
     upstream_base = f"{workspace.rstrip('/')}/ai-gateway/anthropic/"
@@ -389,6 +405,7 @@ def start_proxy(
                 "client": client,
                 "token_header": token_header,
                 "anthropic_model_aliases": _AnthropicModelAliases(),
+                "model_provider_service": model_provider_service,
             },
         ),
     )

tests/test_agent_claude.py

@@ -151,13 +151,14 @@ def test_enables_gateway_model_discovery_for_smart_routing_v2(self, monkeypatch)
         overlay, _ = claude.render_overlay(WS, "s4")
         assert overlay["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1"
 
-    def test_gateway_model_discovery_skipped_under_provider(self, monkeypatch):
-        # A Model Provider Service routes every request to the external provider,
-        # so a discovered gateway endpoint id would reach a provider that can't
-        # resolve it — discovery must be off in that mode.
+    def test_enables_gateway_model_discovery_under_provider(self, monkeypatch):
         monkeypatch.setenv("ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY", "1")
         overlay, _ = claude.render_overlay(WS, "s4", provider="main.x.claude-svc")
-        assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in overlay["env"]
+        assert overlay["env"]["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1"
+        assert (
+            "Databricks-Model-Provider-Service: main.x.claude-svc"
+            in overlay["env"]["ANTHROPIC_CUSTOM_HEADERS"]
+        )
 
     def test_sets_api_key_helper(self):
         overlay, _ = claude.render_overlay(WS, "s4")
@@ -815,7 +816,14 @@ 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,
+            model_provider_service=None,
+        ):
             calls.append(
                 (
                     "proxy",
@@ -824,6 +832,7 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir
                     port,
                     token_header,
                     force_refresh_near_expiry,
+                    model_provider_service,
                 )
             )
             return Server(), Cache(), Client()
@@ -841,15 +850,30 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir
         monkeypatch.setattr(claude.subprocess, "Popen", Process)
 
         with pytest.raises(SystemExit) as exc:
-            claude.launch({"workspace": WS, "profile": "test"}, ["--debug"])
+            claude.launch(
+                {
+                    "workspace": WS,
+                    "profile": "test",
+                    "_claude_launch_provider": "main.default.anthropic",
+                },
+                ["--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"] == "http://127.0.0.1:12345"
         assert os.environ["CLAUDE_CODE_USE_GATEWAY"] == "1"
         assert calls[:2] == [
-            ("proxy", WS, "test", 0, claude.AUTHORIZATION_HEADER, True),
+            (
+                "proxy",
+                WS,
+                "test",
+                0,
+                claude.AUTHORIZATION_HEADER,
+                True,
+                "main.default.anthropic",
+            ),
             ("serve",),
         ]
         assert calls[2][0] == "popen"
@@ -886,9 +910,24 @@ class Client:
             def close(self):
                 calls.append(("close",))
 
-        def start_proxy(workspace, profile, port, token_header, force_refresh_near_expiry):
+        def start_proxy(
+            workspace,
+            profile,
+            port,
+            token_header,
+            force_refresh_near_expiry,
+            model_provider_service=None,
+        ):
             calls.append(
-                ("proxy", workspace, profile, port, token_header, force_refresh_near_expiry)
+                (
+                    "proxy",
+                    workspace,
+                    profile,
+                    port,
+                    token_header,
+                    force_refresh_near_expiry,
+                    model_provider_service,
+                )
             )
             return Server(), Cache(), Client()
 
@@ -910,7 +949,7 @@ def launch_v2(state, tool_args, **kwargs):
 
         assert exc.value.code == 0
         assert calls[:2] == [
-            ("proxy", WS, "test", 0, claude.AUTHORIZATION_HEADER, True),
+            ("proxy", WS, "test", 0, claude.AUTHORIZATION_HEADER, True, None),
             ("serve",),
         ]
         settings, remaining = captured["settings"]

tests/test_anthropic_model_discovery_proxy.py

@@ -154,6 +154,20 @@ def test_leaves_malformed_discovery_response_unchanged(self):
 
 
 class TestAnthropicModelDiscoveryHandler:
+    def test_injects_provider_header_into_model_discovery(self):
+        out = _Collect()
+        handler = _handler(out)
+        handler.headers = {}
+        handler.rfile = io.BytesIO()
+        handler.cache = _FakeCache()
+        handler.model_provider_service = "main.default.anthropic"
+        handler.client = _FakeClient(_FakeResponse(200, {}, b'{"data":[]}'))
+
+        handler._handle()
+
+        _method, _url, headers, _body = handler.client.request
+        assert headers["Databricks-Model-Provider-Service"] == "main.default.anthropic"
+
     def test_inherits_relayed_auth_and_prefixes_models(self):
         out = _Collect()
         handler = _handler(out)
@@ -238,6 +252,7 @@ def test_streams_relayed_inference_response_without_buffering(self):
         handler.headers = {"Authorization": "Bearer subscription-token", "Content-Length": "2"}
         handler.rfile = io.BytesIO(b"{}")
         handler.cache = _FakeCache()
+        handler.model_provider_service = "main.default.anthropic"
         response = _FakeResponse(200, {"Content-Type": "text/event-stream"}, b"data: event\n\n")
         handler.client = _FakeClient(response)
 
@@ -246,6 +261,7 @@ def test_streams_relayed_inference_response_without_buffering(self):
         _method, _url, headers, _body = handler.client.request
         assert headers["Authorization"] == "Bearer subscription-token"
         assert headers["X-Databricks-AI-Gateway-Token"] == "Bearer databricks-token"
+        assert "Databricks-Model-Provider-Service" not in headers
         assert response.read_calls == 0
         assert response.iter_raw_calls == 1
         assert b"data: event\n\n" in bytes(out.data)
@@ -275,7 +291,12 @@ def run_refresher(self):
     )
 
     server, actual_cache, client = anthropic_model_discovery_proxy.start_proxy(
-        "https://workspace.example.com", "profile", 0, "header", False
+        "https://workspace.example.com",
+        "profile",
+        0,
+        "header",
+        False,
+        model_provider_service="main.default.anthropic",
     )
     try:
         handler = server.RequestHandlerClass
@@ -288,6 +309,7 @@ def run_refresher(self):
             handler.anthropic_model_aliases,
             anthropic_model_discovery_proxy._AnthropicModelAliases,
         )
+        assert handler.model_provider_service == "main.default.anthropic"
         assert actual_cache is cache
     finally:
         server.server_close()