diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 222e780d..d7b62618 100644 --- a/src/ucode/agents/claude.py +++ b/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")) diff --git a/src/ucode/anthropic_model_discovery_proxy.py b/src/ucode/anthropic_model_discovery_proxy.py index ada123bf..f8c6089e 100644 --- a/src/ucode/anthropic_model_discovery_proxy.py +++ b/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, }, ), ) diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index dc99ab21..b80de66d 100644 --- a/tests/test_agent_claude.py +++ b/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,7 +850,14 @@ 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" @@ -849,7 +865,15 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir 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"] diff --git a/tests/test_anthropic_model_discovery_proxy.py b/tests/test_anthropic_model_discovery_proxy.py index 0e6b93c1..4edeca89 100644 --- a/tests/test_anthropic_model_discovery_proxy.py +++ b/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()