Fix Claude provider model discovery proxy

src/ucode/agents/claude.py

@@ -1337,7 +1337,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)."""
@@ -1353,6 +1355,7 @@ 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,
+        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
@@ -1377,11 +1380,59 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None:
     raise SystemExit(returncode)
 
 
+def _launch_claude_with_gateway_proxy(
+    state: dict,
+    binary: str,
+    tool_args: list[str],
+    *,
+    provider: str,
+) -> None:
+    """Launch provider-scoped model discovery through the refresh proxy."""
+    workspace = state["workspace"]
+    server, cache, client = gateway_proxy.start_proxy(
+        workspace,
+        state.get("profile"),
+        0,
+        token_header=gateway_proxy.AUTHORIZATION_HEADER,
+        force_refresh_near_expiry=True,
+        model_provider_service=provider,
+    )
+    token = cache.token
+    os.environ["OAUTH_TOKEN"] = token
+    os.environ["ANTHROPIC_AUTH_TOKEN"] = token
+    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)
+    server_thread.start()
+    settings_override = {"env": {"ANTHROPIC_BASE_URL": os.environ["ANTHROPIC_BASE_URL"]}}
+    try:
+        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")
+    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()
@@ -1409,14 +1460,13 @@ def launch(state: dict, tool_args: list[str]) -> None:
             model_name=_maybe_add_1m_suffix,
         )
         return
-    if (
-        workspace
-        and os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1"
-        and not _has_provider_launch(state)
-    ):
+    if workspace and os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1":
         # Discovery is launch-scoped. Pass it in the process environment rather
         # than persisting it in Claude's private or OS-managed settings.
         os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1"
+        if provider:
+            _launch_claude_with_gateway_proxy(state, binary, tool_args, provider=provider)
+            return
     if workspace:
         os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile"))
     exec_or_spawn(_build_claude_argv(binary, tool_args))

src/ucode/gateway_proxy.py

@@ -34,6 +34,7 @@
 # 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"
+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()
@@ -200,6 +201,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
@@ -212,6 +214,18 @@ def _safe_send_error(self, code: int, message: str) -> None:
         except OSError:
             pass
 
+    def _forwarded_request_headers(self) -> dict[str, str]:
+        headers = forwarded_request_headers(self, self.cache.token, self.token_header)
+        if (
+            self.command == "GET"
+            and self.path.partition("?")[0] == "/v1/models"
+            and self.model_provider_service
+        ):
+            # Claude Code omits ANTHROPIC_CUSTOM_HEADERS from native model
+            # discovery. Inject only the routing header that scopes the picker.
+            headers[MODEL_PROVIDER_SERVICE_HEADER] = self.model_provider_service
+        return headers
+
     def _handle(self) -> None:
         diagnostic_id = uuid.uuid4().hex[:12]
         started = time.monotonic()
@@ -226,7 +240,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(
                     "upstream_headers",
@@ -256,7 +270,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 = self._forwarded_request_headers()
             with self.client.stream(self.command, url, headers=headers, content=body) as resp:
                 log_proxy_diagnostic(
                     "upstream_headers",
@@ -374,6 +388,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 loopback refresh proxy + its background token refresher.
 
@@ -399,7 +414,12 @@ def start_proxy(
     handler = type(
         "BoundProxyHandler",
         (_ProxyHandler,),
-        {"cache": cache, "client": client, "token_header": token_header},
+        {
+            "cache": cache,
+            "client": client,
+            "token_header": token_header,
+            "model_provider_service": model_provider_service,
+        },
     )
     try:
         server = ThreadingHTTPServer(("127.0.0.1", port), handler)

tests/test_agent_claude.py

@@ -214,9 +214,7 @@ def test_smart_routing_does_not_persist_gateway_model_discovery(self, monkeypatc
         assert "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY" not in overlay["env"]
 
     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.
+        # Discovery is launch-scoped and should not be persisted in settings.
         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"]
@@ -1046,7 +1044,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",
@@ -1055,6 +1060,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()
@@ -1071,6 +1077,7 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir
                     "profile": "test",
                     "claude_relayed": True,
                     "relayed_proxy_port": 12345,
+                    "_claude_launch_provider": "main.default.anthropic",
                 },
                 ["--debug"],
             )
@@ -1083,6 +1090,7 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir
             12345,
             claude.gateway_proxy.AI_GATEWAY_TOKEN_HEADER,
             False,
+            "main.default.anthropic",
         )
         assert calls[-3:] == [("stop",), ("shutdown",), ("close",)]
 
@@ -1221,6 +1229,93 @@ def test_gateway_discovery_uses_direct_gateway(self, monkeypatch):
         assert os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1"
         assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"]]
 
+    @pytest.mark.parametrize(
+        "provider_state",
+        [
+            {"provider_services": {"claude": "main.default.anthropic"}},
+            {"_claude_launch_provider": "main.default.anthropic"},
+        ],
+    )
+    def test_gateway_discovery_enabled_for_provider_launch(self, monkeypatch, provider_state):
+        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,
+            force_refresh_near_expiry,
+            model_provider_service=None,
+        ):
+            calls.append(
+                (
+                    "proxy",
+                    workspace,
+                    profile,
+                    port,
+                    token_header,
+                    force_refresh_near_expiry,
+                    model_provider_service,
+                )
+            )
+            return Server(), Cache(), Client()
+
+        monkeypatch.delenv(v2.ENV_VAR, raising=False)
+        monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "1")
+        monkeypatch.delenv("CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY", raising=False)
+        monkeypatch.setattr(claude.gateway_proxy, "start_proxy", start_proxy)
+        monkeypatch.setattr(claude.subprocess, "Popen", Process)
+
+        with pytest.raises(SystemExit) as exc:
+            claude.launch({"workspace": WS, **provider_state}, ["--debug"])
+
+        assert exc.value.code == 0
+        assert os.environ["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] == "1"
+        assert calls[:2] == [
+            (
+                "proxy",
+                WS,
+                None,
+                0,
+                claude.gateway_proxy.AUTHORIZATION_HEADER,
+                True,
+                "main.default.anthropic",
+            ),
+            ("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",)]
+
 
 class TestWriteToolConfigPrunesStaleModelEnv:
     """Stale ucode-managed model env keys (ANTHROPIC_MODEL, etc.) from earlier

tests/test_gateway_proxy.py

@@ -328,9 +328,11 @@ class _FakeClient:
     def __init__(self, responses):
         self._responses = list(responses)
         self.sent_tokens: list[str | None] = []
+        self.sent_headers: list[dict[str, str]] = []
 
     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)
 
 
@@ -392,6 +394,28 @@ def flush(self):
 
 
 class TestRetryOn401:
+    def test_injects_provider_header_into_model_discovery(self):
+        client = _FakeClient([_FakeResp(200, b'{"data":[]}')])
+        handler = _handle_handler(client, _FakeCache(), _Collect())
+        handler.command = "GET"
+        handler.path = "/v1/models?limit=20"
+        handler.model_provider_service = "main.default.anthropic"
+
+        handler._handle()
+
+        assert client.sent_headers[0][gateway_proxy.MODEL_PROVIDER_SERVICE_HEADER] == (
+            "main.default.anthropic"
+        )
+
+    def test_does_not_inject_provider_header_into_inference(self):
+        client = _FakeClient([_FakeResp(200, b"ok")])
+        handler = _handle_handler(client, _FakeCache(), _Collect())
+        handler.model_provider_service = "main.default.anthropic"
+
+        handler._handle()
+
+        assert gateway_proxy.MODEL_PROVIDER_SERVICE_HEADER not in client.sent_headers[0]
+
     def test_401_forces_refresh_and_retries(self):
         # A stale swap token yields 401; the proxy force-refreshes and retries,
         # this time succeeding, so Claude Code never sees the 401.
@@ -463,11 +487,13 @@ def run_refresher(self):
                 busy_port,
                 token_header=gateway_proxy.AI_GATEWAY_TOKEN_HEADER,
                 force_refresh_near_expiry=False,
+                model_provider_service="main.default.anthropic",
             )
             try:
                 bound = server.server_address[1]
                 assert bound != busy_port  # fell back to a different, free port
                 assert bound != 0
+                assert server.RequestHandlerClass.model_provider_service == "main.default.anthropic"
             finally:
                 server.server_close()
                 client.close()