diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 222e780d..9cfefac2 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -15,6 +15,7 @@ from pathlib import Path from typing import cast +from ucode import gateway_proxy from ucode.agent_updates import available_npm_package_update from ucode.anthropic_model_discovery_proxy import ( start_proxy as start_anthropic_model_discovery_proxy, @@ -33,10 +34,6 @@ build_tool_base_url, get_databricks_token, ) -from ucode.gateway_proxy import ( - AI_GATEWAY_TOKEN_HEADER, - AUTHORIZATION_HEADER, -) from ucode.launcher import exec_or_spawn from ucode.managed_files import OS, current_os, write_managed_file from ucode.smart_routing import v2 as smart_routing_v2 @@ -58,6 +55,8 @@ CLAUDE_USER_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "settings.json" CLAUDE_BACKUP_PATH = APP_DIR / "claude-ucode-settings.backup.json" WEB_SEARCH_MCP_STATE_KEY = "claude_web_search_mcp" +MINIMUM_CLAUDE_VERSION = (2, 1, 248) +MINIMUM_CLAUDE_VERSION_TEXT = "2.1.248" SPEC: ToolSpec = { "binary": "claude", @@ -98,6 +97,49 @@ def is_update_available() -> tuple[str, str] | None: return available_npm_package_update(SPEC["package"]) +def _parse_version(value: str) -> tuple[int, int, int] | None: + match = re.search(r"(\d+)\.(\d+)\.(\d+)", value) + if not match: + return None + major, minor, patch = match.groups() + return int(major), int(minor), int(patch) + + +def _installed_version_status() -> tuple[str, bool] | None: + version = agent_version(SPEC["binary"]) + parsed = _parse_version(version) + if parsed is None: + return None + return version, parsed < MINIMUM_CLAUDE_VERSION + + +def minimum_version_error() -> str | None: + status = _installed_version_status() + if status is None: + return None + version, is_too_old = status + if not is_too_old: + return None + return ( + f"Claude Code {version} is too old for gateway model discovery. " + f"Claude Code must be updated to {MINIMUM_CLAUDE_VERSION_TEXT} or newer; " + f"run `npm install -g {SPEC['package']}` or `ucode configure`." + ) + + +def required_update_message() -> str | None: + status = _installed_version_status() + if status is None: + return None + version, is_too_old = status + if not is_too_old: + return None + return ( + f"Claude Code {version} is older than required {MINIMUM_CLAUDE_VERSION_TEXT}; " + "updating Claude Code is required for gateway model discovery." + ) + + def _resolve_web_search_model(state: dict) -> str | None: """Pick the model the web_search MCP server should call. Prefers an explicit override in state, otherwise the first endpoint discovered as @@ -1122,11 +1164,11 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: if not isinstance(port, int): raise RuntimeError("Relayed proxy port was not configured; re-run `ucode claude`.") - server, cache, client = start_anthropic_model_discovery_proxy( + server, cache, client = gateway_proxy.start_proxy( workspace, state.get("profile"), port, - token_header=AI_GATEWAY_TOKEN_HEADER, + token_header=gateway_proxy.AI_GATEWAY_TOKEN_HEADER, force_refresh_near_expiry=False, ) # start_proxy falls back to an OS-assigned port when the cached one is taken @@ -1155,18 +1197,9 @@ 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 ) -> None: - """Launch Claude through a refreshing gateway proxy.""" + """Launch Claude through the gateway model-alias proxy.""" workspace = state["workspace"] - server, cache, client = start_anthropic_model_discovery_proxy( - workspace, - state.get("profile"), - 0, - token_header=AUTHORIZATION_HEADER, - force_refresh_near_expiry=True, - ) - token = cache.token - os.environ["OAUTH_TOKEN"] = token - os.environ["ANTHROPIC_AUTH_TOKEN"] = token + server, client = start_anthropic_model_discovery_proxy(workspace, 0) os.environ["ANTHROPIC_BASE_URL"] = f"http://{LOOPBACK_HOST}:{server.server_address[1]}" os.environ["CLAUDE_CODE_USE_GATEWAY"] = "1" @@ -1201,7 +1234,6 @@ def compose_gateway_settings(args: list[str]) -> tuple[dict, list[str]]: proc.send_signal(signal.SIGINT) returncode = proc.wait() finally: - cache.stop() server.shutdown() client.close() raise SystemExit(returncode) diff --git a/src/ucode/anthropic_model_discovery_proxy.py b/src/ucode/anthropic_model_discovery_proxy.py index ada123bf..61a66065 100644 --- a/src/ucode/anthropic_model_discovery_proxy.py +++ b/src/ucode/anthropic_model_discovery_proxy.py @@ -1,12 +1,11 @@ """Loopback proxy for Claude gateway model discovery. -The proxy refreshes the Databricks credential, streams inference responses -verbatim, and rewrites model discovery responses when needed. +The proxy forwards Claude Code's apiKeyHelper credential, streams inference +responses verbatim, and rewrites model discovery responses when needed. 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 - and is refreshed off the request path. + - Never logs header values or bodies. """ from __future__ import annotations @@ -26,13 +25,9 @@ from ucode.constants import LOOPBACK_HOST from ucode.databricks import _http_get_retry_delay from ucode.gateway_proxy import ( - AI_GATEWAY_TOKEN_HEADER, HOP_BY_HOP_HEADERS, UPSTREAM_TIMEOUT, - TokenCache, - forwarded_request_headers, log_proxy_diagnostic, - log_token_refresh_failure, ) # Claude Code abandons model discovery after roughly three seconds. One retry @@ -42,9 +37,7 @@ class _ProxyHandler(BaseHTTPRequestHandler): # Set by the server factory. - cache: TokenCache client: httpx.Client - token_header = AI_GATEWAY_TOKEN_HEADER def log_message(self, format: str, *args: object) -> None: return @@ -87,7 +80,11 @@ def _retry_model_discovery( delay_ms=round(delay * 1000), ) time.sleep(delay) - headers = forwarded_request_headers(self, self.cache.token, self.token_header) + headers = { + key: value + for key, value in self.headers.items() + if key.lower() not in HOP_BY_HOP_HEADERS + } with self.client.stream(self.command, url, headers=headers, content=body) as resp: log_proxy_diagnostic( "model_discovery_upstream_headers", @@ -118,8 +115,11 @@ def _handle(self) -> None: path=self.path.split("?", 1)[0], ) try: - # First attempt with the current token. - headers = forwarded_request_headers(self, self.cache.token, self.token_header) + headers = { + key: value + for key, value in self.headers.items() + if key.lower() not in HOP_BY_HOP_HEADERS + } with self.client.stream(self.command, url, headers=headers, content=body) as resp: log_proxy_diagnostic( "model_discovery_upstream_headers", @@ -139,27 +139,6 @@ def _handle(self) -> None: retry_after, ) return - if resp.status_code not in (401, 403): - self._relay_response(resp, diagnostic_id=diagnostic_id, started=started) - return - # Auth rejected. Drain the (small) error body so the pooled - # connection can be reused, then fall through to one retry. - resp.read() - # Force-refresh the Databricks token and retry once. - try: - self.cache.refresh() - 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) - with self.client.stream(self.command, url, headers=headers, content=body) as resp: - log_proxy_diagnostic( - "model_discovery_upstream_headers", - request_id=diagnostic_id, - attempt=2, - status=resp.status_code, - elapsed_ms=round((time.monotonic() - started) * 1000), - ) self._relay_response(resp, diagnostic_id=diagnostic_id, started=started) except (BrokenPipeError, ConnectionResetError): # Client closed before/while we relayed headers — routine on cancel. @@ -366,18 +345,10 @@ def _response_chunks(self, resp: httpx.Response) -> tuple[Iterable[bytes], froze def start_proxy( workspace: str, - profile: str | None, port: int, - token_header: str, - force_refresh_near_expiry: bool, -) -> tuple[ThreadingHTTPServer, TokenCache, httpx.Client]: - """Start the Anthropic model discovery proxy and token refresher.""" +) -> tuple[ThreadingHTTPServer, httpx.Client]: + """Start the Anthropic model discovery proxy.""" upstream_base = f"{workspace.rstrip('/')}/ai-gateway/anthropic/" - cache = TokenCache( - workspace, - profile, - force_refresh_near_expiry=force_refresh_near_expiry, - ) client = httpx.Client(base_url=upstream_base, timeout=UPSTREAM_TIMEOUT, follow_redirects=False) handler = cast( type[BaseHTTPRequestHandler], @@ -385,9 +356,7 @@ def start_proxy( "BoundProxyHandler", (_AnthropicModelDiscoveryHandler,), { - "cache": cache, "client": client, - "token_header": token_header, "anthropic_model_aliases": _AnthropicModelAliases(), }, ), @@ -397,6 +366,4 @@ def start_proxy( except OSError: server = ThreadingHTTPServer((LOOPBACK_HOST, 0), handler) - refresher = threading.Thread(target=cache.run_refresher, daemon=True) - refresher.start() - return server, cache, client + return server, client diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index dc99ab21..579fc0bd 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -26,6 +26,34 @@ def test_display(self): assert claude.SPEC["display"] == "Claude Code" +class TestMinimumVersion: + @pytest.mark.parametrize("version", ["2.1.248", "2.1.250", "3.0.0"]) + def test_supported_version(self, monkeypatch, version): + monkeypatch.setattr(claude, "agent_version", lambda _binary: version) + + assert claude.minimum_version_error() is None + assert claude.required_update_message() is None + + def test_older_version_requires_update(self, monkeypatch): + monkeypatch.setattr(claude, "agent_version", lambda _binary: "2.1.247 (Claude Code)") + + assert claude.minimum_version_error() == ( + "Claude Code 2.1.247 (Claude Code) is too old for gateway model discovery. " + "Claude Code must be updated to 2.1.248 or newer; run " + "`npm install -g @anthropic-ai/claude-code` or `ucode configure`." + ) + assert claude.required_update_message() == ( + "Claude Code 2.1.247 (Claude Code) is older than required 2.1.248; " + "updating Claude Code is required for gateway model discovery." + ) + + def test_unknown_version_does_not_block(self, monkeypatch): + monkeypatch.setattr(claude, "agent_version", lambda _binary: "unknown") + + assert claude.minimum_version_error() is None + assert claude.required_update_message() is None + + class TestRenderOverlay: def test_long_context_suffix_supports_major_only_claude_versions(self): assert claude._maybe_add_1m_suffix("system.ai.claude-sonnet-5") == ( @@ -682,6 +710,79 @@ def boom(name, entry, scope=mcp_mod.MCP_USER_SCOPE): class TestClaudeLaunch: + def test_relayed_launch_uses_refresh_proxy_not_discovery_proxy(self, monkeypatch): + 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: + 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): + calls.append( + ( + "proxy", + workspace, + profile, + port, + token_header, + force_refresh_near_expiry, + ) + ) + return Server(), Cache(), Client() + + monkeypatch.setattr(claude, "_managed_relayed_conflicts", lambda: None) + monkeypatch.setattr(claude, "_managed_pinned_model", lambda: None) + monkeypatch.setattr(claude, "_ensure_subscription_login", lambda: None) + monkeypatch.setattr(claude.gateway_proxy, "start_proxy", start_proxy) + monkeypatch.setattr( + claude, + "start_anthropic_model_discovery_proxy", + lambda *_args: pytest.fail("relayed auth must not use the discovery proxy"), + ) + monkeypatch.setattr(claude.subprocess, "Popen", Process) + + with pytest.raises(SystemExit) as exc: + claude.launch( + { + "workspace": WS, + "profile": "test", + "claude_relayed": True, + "relayed_proxy_port": 12345, + }, + ["--debug"], + ) + + assert exc.value.code == 0 + assert calls[0] == ( + "proxy", + WS, + "test", + 12345, + claude.gateway_proxy.AI_GATEWAY_TOKEN_HEADER, + False, + ) + assert calls[-3:] == [("stop",), ("shutdown",), ("close",)] + def test_smart_routing_on_windows_is_not_supported(self, monkeypatch): monkeypatch.setenv(v2.ENV_VAR, "1") monkeypatch.setattr(claude.os, "name", "nt") @@ -798,12 +899,6 @@ def serve_forever(self): 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",)) @@ -815,24 +910,20 @@ def __init__(self, argv): def wait(self): return 0 - def start_proxy(workspace, profile, port, token_header, force_refresh_near_expiry): - calls.append( - ( - "proxy", - workspace, - profile, - port, - token_header, - force_refresh_near_expiry, - ) - ) - return Server(), Cache(), Client() + def start_proxy(workspace, port): + calls.append(("proxy", workspace, port)) + return Server(), Client() monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "1") monkeypatch.delenv("OAUTH_TOKEN", raising=False) monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False) monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False) monkeypatch.delenv("CLAUDE_CODE_USE_GATEWAY", raising=False) + monkeypatch.setattr( + claude, + "get_databricks_token", + lambda *_args: pytest.fail("model discovery must rely on apiKeyHelper"), + ) monkeypatch.setattr( claude, "start_anthropic_model_discovery_proxy", @@ -844,12 +935,12 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir claude.launch({"workspace": WS, "profile": "test"}, ["--debug"]) assert exc.value.code == 0 - assert os.environ["OAUTH_TOKEN"] == "fresh-token" - assert os.environ["ANTHROPIC_AUTH_TOKEN"] == "fresh-token" + assert "OAUTH_TOKEN" not in os.environ + assert "ANTHROPIC_AUTH_TOKEN" not in os.environ 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, 0), ("serve",), ] assert calls[2][0] == "popen" @@ -857,11 +948,7 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir 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",), - ] + assert calls[3:] == [("shutdown",), ("close",)] def test_smart_routing_uses_anthropic_proxy(self, monkeypatch): calls: list[tuple] = [] @@ -876,27 +963,21 @@ def serve_forever(self): 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",)) - def start_proxy(workspace, profile, port, token_header, force_refresh_near_expiry): - calls.append( - ("proxy", workspace, profile, port, token_header, force_refresh_near_expiry) - ) - return Server(), Cache(), Client() + def start_proxy(workspace, port): + calls.append(("proxy", workspace, port)) + return Server(), Client() def launch_v2(state, tool_args, **kwargs): captured["settings"] = kwargs["compose_settings"](["--debug"]) raise SystemExit(0) monkeypatch.setenv(v2.ENV_VAR, "1") + monkeypatch.delenv("OAUTH_TOKEN", raising=False) + monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False) monkeypatch.setattr(claude, "start_anthropic_model_discovery_proxy", start_proxy) monkeypatch.setattr( claude, @@ -909,14 +990,16 @@ def launch_v2(state, tool_args, **kwargs): claude.launch({"workspace": WS, "profile": "test"}, ["--debug"]) assert exc.value.code == 0 + assert "OAUTH_TOKEN" not in os.environ + assert "ANTHROPIC_AUTH_TOKEN" not in os.environ assert calls[:2] == [ - ("proxy", WS, "test", 0, claude.AUTHORIZATION_HEADER, True), + ("proxy", WS, 0), ("serve",), ] settings, remaining = captured["settings"] assert settings["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:12345" assert remaining == ["--debug"] - assert calls[2:] == [("stop",), ("shutdown",), ("close",)] + assert calls[2:] == [("shutdown",), ("close",)] class TestWriteToolConfigPrunesStaleModelEnv: diff --git a/tests/test_anthropic_model_discovery_proxy.py b/tests/test_anthropic_model_discovery_proxy.py index 0e6b93c1..1d8996e1 100644 --- a/tests/test_anthropic_model_discovery_proxy.py +++ b/tests/test_anthropic_model_discovery_proxy.py @@ -43,13 +43,6 @@ def stream(self, method, url, headers, content): return self.responses.pop(0) -class _FakeCache: - token = "databricks-token" - - def refresh(self): - return None - - class _Collect(io.RawIOBase): def __init__(self): self.data = bytearray() @@ -154,28 +147,43 @@ def test_leaves_malformed_discovery_response_unchanged(self): class TestAnthropicModelDiscoveryHandler: - def test_inherits_relayed_auth_and_prefixes_models(self): + def test_forwards_api_key_helper_credential_without_refresh_auth(self): + out = _Collect() + handler = _handler(out) + handler.headers = {"X-Api-Key": "api-key-helper-token"} + handler.rfile = io.BytesIO() + handler.client = _FakeClient(_FakeResponse(200, {}, b'{"data":[]}')) + + handler._handle() + + _method, _url, headers, _body = handler.client.request + assert headers["X-Api-Key"] == "api-key-helper-token" + assert "Authorization" not in headers + assert "X-Databricks-AI-Gateway-Token" not in headers + + def test_prefixes_models_and_strips_hop_by_hop_headers(self): out = _Collect() handler = _handler(out) - handler.headers = {"Authorization": "Bearer subscription-token"} + handler.headers = { + "X-Api-Key": "api-key-helper-token", + "Connection": "keep-alive", + "Host": "127.0.0.1", + } handler.rfile = io.BytesIO() - handler.cache = _FakeCache() handler.client = _FakeClient(_FakeResponse(200, {}, b'{"data":[{"id":"custom-model"}]}')) handler._handle() method, url, headers, body = handler.client.request assert (method, url, body) == ("GET", "v1/models", None) - assert headers["Authorization"] == "Bearer subscription-token" - assert headers["X-Databricks-AI-Gateway-Token"] == "Bearer databricks-token" + assert headers == {"X-Api-Key": "api-key-helper-token"} assert b"anthropic-aigw-custom-model" in bytes(out.data) def test_retries_rate_limited_model_discovery(self, monkeypatch): out = _Collect() handler = _handler(out) - handler.headers = {"Authorization": "Bearer subscription-token"} + handler.headers = {"X-Api-Key": "api-key-helper-token"} handler.rfile = io.BytesIO() - handler.cache = _FakeCache() rate_limited = _FakeResponse(429, {"Retry-After": "0"}, b"rate limited") success = _FakeResponse(200, {}, b'{"data":[{"id":"custom-model"}]}') handler.client = _FakeClient([rate_limited, success]) @@ -184,6 +192,10 @@ def test_retries_rate_limited_model_discovery(self, monkeypatch): handler._handle() assert len(handler.client.requests) == 2 + assert all( + request[2] == {"X-Api-Key": "api-key-helper-token"} + for request in handler.client.requests + ) assert rate_limited.read_calls == 1 assert b"429 Too Many Requests" not in bytes(out.data) assert b"anthropic-aigw-custom-model" in bytes(out.data) @@ -191,9 +203,8 @@ def test_retries_rate_limited_model_discovery(self, monkeypatch): def test_relays_rate_limit_after_model_discovery_retries_are_exhausted(self, monkeypatch): out = _Collect() handler = _handler(out) - handler.headers = {"Authorization": "Bearer subscription-token"} + handler.headers = {"X-Api-Key": "api-key-helper-token"} handler.rfile = io.BytesIO() - handler.cache = _FakeCache() responses = [_FakeResponse(429, {"Retry-After": "0"}, b"rate limited") for _ in range(2)] handler.client = _FakeClient(responses) monkeypatch.setattr(anthropic_model_discovery_proxy.time, "sleep", lambda _delay: None) @@ -232,20 +243,18 @@ def test_keeps_content_encoding_for_unchanged_error(self): assert response.read_calls == 0 assert response.iter_raw_calls == 1 - def test_streams_relayed_inference_response_without_buffering(self): + def test_streams_inference_response_without_buffering(self): out = _Collect() handler = _handler(out, path="/v1/messages", command="POST") - handler.headers = {"Authorization": "Bearer subscription-token", "Content-Length": "2"} + handler.headers = {"X-Api-Key": "api-key-helper-token", "Content-Length": "2"} handler.rfile = io.BytesIO(b"{}") - handler.cache = _FakeCache() response = _FakeResponse(200, {"Content-Type": "text/event-stream"}, b"data: event\n\n") handler.client = _FakeClient(response) handler._handle() _method, _url, headers, _body = handler.client.request - assert headers["Authorization"] == "Bearer subscription-token" - assert headers["X-Databricks-AI-Gateway-Token"] == "Bearer databricks-token" + assert headers == {"X-Api-Key": "api-key-helper-token"} assert response.read_calls == 0 assert response.iter_raw_calls == 1 assert b"data: event\n\n" in bytes(out.data) @@ -262,33 +271,18 @@ def test_strips_known_alias_from_message_request(self): assert json.loads(body) == {"model": "catalog.schema.custom"} -def test_start_proxy_uses_discovery_handler(monkeypatch): - class _StubCache: - def run_refresher(self): - return None - - cache = _StubCache() - monkeypatch.setattr( - anthropic_model_discovery_proxy, - "TokenCache", - lambda *_args, **_kwargs: cache, - ) - - server, actual_cache, client = anthropic_model_discovery_proxy.start_proxy( - "https://workspace.example.com", "profile", 0, "header", False - ) +def test_start_proxy_uses_discovery_handler(): + server, client = anthropic_model_discovery_proxy.start_proxy("https://workspace.example.com", 0) try: handler = server.RequestHandlerClass assert issubclass( handler, anthropic_model_discovery_proxy._AnthropicModelDiscoveryHandler, ) - assert handler.cache is cache assert isinstance( handler.anthropic_model_aliases, anthropic_model_discovery_proxy._AnthropicModelAliases, ) - assert actual_cache is cache finally: server.server_close() client.close()