diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 01b3a1f9..28dcd06e 100644 --- a/src/ucode/agents/claude.py +++ b/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)) diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index e5cd4788..9326f6f3 100644 --- a/src/ucode/gateway_proxy.py +++ b/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) diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 61f07438..a5881e7a 100644 --- a/tests/test_agent_claude.py +++ b/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 diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py index 28600045..aeb5185b 100644 --- a/tests/test_gateway_proxy.py +++ b/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()