From 75427bbd4ad14fcb1b205f3253ec4e24ae1d2118 Mon Sep 17 00:00:00 2001 From: Vinay Gupta <58447456+aivinay@users.noreply.github.com> Date: Tue, 30 Jun 2026 14:15:57 -0500 Subject: [PATCH] fix(wrap): preserve custom Vertex base URL (#1477) ## Description Fixes `headroom wrap claude` in Vertex mode when the user has configured a custom Vertex-compatible gateway through `ANTHROPIC_VERTEX_BASE_URL`. Before this change, wrap mode redirected Claude Code's `ANTHROPIC_VERTEX_BASE_URL` to the local Headroom proxy, but the original custom upstream was not forwarded to the proxy as `VERTEX_TARGET_API_URL`. The proxy therefore fell back to the default Google Vertex endpoints and custom gateways could return 404 or auth/model errors. Closes #1476 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Capture the original Vertex upstream before `wrap claude` redirects Claude Code to the local proxy. - Pass custom Vertex upstreams to the proxy as `--vertex-api-url` / `VERTEX_TARGET_API_URL`. - Let explicit `VERTEX_TARGET_API_URL` take precedence over `ANTHROPIC_VERTEX_BASE_URL`. - Guard against accidentally using the local Headroom proxy URL as the proxy's own Vertex upstream. - Restart idle running proxies when their configured Vertex upstream does not match the requested Vertex mode state. - Persist and restore `ANTHROPIC_VERTEX_BASE_URL` for Vertex-mode Claude daemon workers, and clean it up during `unwrap claude`. - Expose `vertex_api_url` in loopback health config so wrapper reuse checks can detect mismatches. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text $ rtk gh pr checks 1477 --repo headroomlabs-ai/headroom CI Checks Summary: [ok] Passed: 22 [FAIL] Failed: 0 $ rtk pytest tests/test_cli/test_wrap_claude_vertex_proxy_env.py tests/test_cli/test_unwrap_claude.py tests/test_azure_foundry_claude_compression.py tests/test_cli/test_wrap_persistent.py tests/test_provider_registry.py -q Pytest: 64 passed $ rtk uvx ruff==0.15.17 check headroom/cli/wrap.py headroom/proxy/server.py tests/test_cli/test_wrap_claude_vertex_proxy_env.py tests/test_cli/test_unwrap_claude.py All checks passed! $ rtk uvx ruff==0.15.17 format --check headroom/cli/wrap.py headroom/proxy/server.py tests/test_cli/test_wrap_claude_vertex_proxy_env.py tests/test_cli/test_unwrap_claude.py 4 files already formatted $ rtk python3 -m py_compile headroom/cli/wrap.py headroom/proxy/server.py tests/test_cli/test_wrap_claude_vertex_proxy_env.py tests/test_cli/test_unwrap_claude.py # passed, no output $ rtk uv run --python 3.13 pytest tests/test_vertex_claude_compression.py -q Failed before test collection while building the local editable package: esaxx-rs build failed with fatal error: 'cstdint' file not found. ``` ## Real Behavior Proof - Environment: GitHub Actions CI on PR #1477 plus local macOS worktree `fix/1476-vertex-base-url`. - Exact command / steps: CI ran lint, type checking, build, unit-test shards, native wrapper checks, wrap-native e2e, and Docker e2e jobs; locally ran focused wrapper, unwrap, Foundry, persistent-proxy, and provider-registry tests. - Observed result: CI passed 22 checks with 0 failures; local focused tests passed; Ruff check/format passed; Python compile passed. - Not tested: broader proxy-route tests that import `headroom.proxy.server` through a local editable build could not run locally because the native `esaxx-rs` build fails before test collection with missing `cstdint`. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) N/A ## Additional Notes - Documentation, CHANGELOG, and extra code-comment checklist items are N/A for this narrow wrapper bug fix. - Full local unit test execution is limited by the existing native extension build issue described above; focused Python-only coverage passes and GitHub CI is green. --- headroom/cli/wrap.py | 112 +++++- headroom/proxy/server.py | 1 + tests/test_cli/test_unwrap_claude.py | 20 + .../test_wrap_claude_vertex_proxy_env.py | 373 ++++++++++++++++++ 4 files changed, 494 insertions(+), 12 deletions(-) create mode 100644 tests/test_cli/test_wrap_claude_vertex_proxy_env.py diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index 1c40b0154..07b42e6ed 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -388,6 +388,8 @@ def _start_proxy( region: str | None = None, openai_api_url: str | None = None, anthropic_api_url: str | None = None, + vertex_api_url: str | None = None, + clear_vertex_api_url: bool = False, copilot_api_token: str | None = None, ) -> subprocess.Popen: """Start Headroom proxy as a background subprocess. @@ -434,6 +436,9 @@ def _start_proxy( if anthropic_api_url: cmd.extend(["--anthropic-api-url", anthropic_api_url]) + if vertex_api_url: + cmd.extend(["--vertex-api-url", vertex_api_url]) + timeout_seconds = _resolve_wrap_proxy_timeout_seconds() log_path = _get_log_path() stdio_log_path = _get_proxy_stdio_log_path() @@ -453,6 +458,10 @@ def _start_proxy( proxy_env["OPENAI_TARGET_API_URL"] = openai_api_url if anthropic_api_url: proxy_env["ANTHROPIC_TARGET_API_URL"] = anthropic_api_url + if clear_vertex_api_url: + proxy_env.pop("VERTEX_TARGET_API_URL", None) + if vertex_api_url: + proxy_env["VERTEX_TARGET_API_URL"] = vertex_api_url # Pin the wrapper-validated Copilot token for this proxy instance only. # Injected into the subprocess env here (not the parent's os.environ) so it # never leaks into shared state. The proxy's CopilotTokenProvider honours @@ -798,22 +807,55 @@ def _foundry_proxy_url(proxy_url: str) -> str: return proxy_url.rstrip("/") + "/anthropic" +def _vertex_target_api_url_from_claude_env(proxy_url: str) -> str | None: + """Return the Vertex upstream that the proxy should use for Claude Code.""" + explicit_target = os.environ.get("VERTEX_TARGET_API_URL", "").strip() + if explicit_target: + return ( + None + if _normalize_proxy_api_url(explicit_target) == _normalize_proxy_api_url(proxy_url) + else explicit_target + ) + + vertex_url = os.environ.get("ANTHROPIC_VERTEX_BASE_URL", "").strip() + if not vertex_url: + return None + + from headroom.providers.registry import DEFAULT_VERTEX_API_URL + + normalized_vertex_url = _normalize_proxy_api_url(vertex_url) + if normalized_vertex_url == _normalize_proxy_api_url(DEFAULT_VERTEX_API_URL): + return None + if normalized_vertex_url == _normalize_proxy_api_url(proxy_url): + return None + return vertex_url + + +def _claude_wrap_base_url_env_key(*, foundry_mode: bool = False, vertex_mode: bool = False) -> str: + if vertex_mode: + return "ANTHROPIC_VERTEX_BASE_URL" + if foundry_mode: + return "ANTHROPIC_FOUNDRY_BASE_URL" + return "ANTHROPIC_BASE_URL" + + def _write_claude_wrap_base_url( proxy_url: str, *, foundry_mode: bool = False, + vertex_mode: bool = False, settings_path: Path | None = None, ) -> str | None: """Persist proxy URL into project-local settings env key for daemon child inheritance. Claude Code's cc-daemon pre-forks conversation workers using spawn (not fork), so those workers read settings.json fresh rather than inheriting - the daemon's environment. Writing env.ANTHROPIC_BASE_URL into the - project-local settings file (.claude/settings.local.json in cwd) ensures - every new conversation — including those started after the initial launch — - routes through the Headroom proxy without touching the global user settings - file or affecting sessions in other projects. Returns the previous value - so the caller can restore it on exit (issue #951). + the daemon's environment. Writing the mode-specific Claude base URL env + key into the project-local settings file (.claude/settings.local.json in + cwd) ensures every new conversation — including those started after the + initial launch — routes through the Headroom proxy without touching the + global user settings file or affecting sessions in other projects. Returns + the previous value so the caller can restore it on exit (issue #951). """ path = settings_path or (Path.cwd() / ".claude" / "settings.local.json") payload: dict[str, Any] = {} @@ -825,7 +867,7 @@ def _write_claude_wrap_base_url( if not isinstance(payload, dict): payload = {} env_map = dict(payload.get("env") or {}) if isinstance(payload.get("env"), dict) else {} - key = "ANTHROPIC_FOUNDRY_BASE_URL" if foundry_mode else "ANTHROPIC_BASE_URL" + key = _claude_wrap_base_url_env_key(foundry_mode=foundry_mode, vertex_mode=vertex_mode) previous = env_map.get(key) env_map[key] = proxy_url payload["env"] = env_map @@ -838,6 +880,7 @@ def _restore_claude_wrap_base_url( previous: str | None, *, foundry_mode: bool = False, + vertex_mode: bool = False, settings_path: Path | None = None, ) -> None: """Restore (or remove) the env key written by _write_claude_wrap_base_url. @@ -859,7 +902,7 @@ def _restore_claude_wrap_base_url( env_map = payload.get("env") if not isinstance(env_map, dict): return - key = "ANTHROPIC_FOUNDRY_BASE_URL" if foundry_mode else "ANTHROPIC_BASE_URL" + key = _claude_wrap_base_url_env_key(foundry_mode=foundry_mode, vertex_mode=vertex_mode) if previous is None: if key not in env_map: return @@ -2624,6 +2667,8 @@ def _ensure_proxy( region: str | None = None, openai_api_url: str | None = None, anthropic_api_url: str | None = None, + vertex_api_url: str | None = None, + clear_vertex_api_url: bool = False, copilot_api_token: str | None = None, ) -> subprocess.Popen | None: """Start or verify proxy. Returns process handle if we started it.""" @@ -2836,6 +2881,13 @@ def _ensure_proxy( requested_openai_url = _normalize_proxy_api_url(openai_api_url) if running_openai_url != requested_openai_url: missing.append("openai-api-url") + if vertex_api_url or clear_vertex_api_url: + running_vertex_url = _normalize_proxy_api_url( + running_config.get("vertex_api_url") + ) + requested_vertex_url = _normalize_proxy_api_url(vertex_api_url) + if running_vertex_url != requested_vertex_url: + missing.append("vertex-api-url") if missing: flags_str = ", ".join( @@ -2909,6 +2961,8 @@ def _ensure_proxy( region=region, openai_api_url=openai_api_url, anthropic_api_url=anthropic_api_url, + vertex_api_url=vertex_api_url, + clear_vertex_api_url=clear_vertex_api_url, copilot_api_token=copilot_api_token, ), ) @@ -2921,6 +2975,23 @@ def _ensure_proxy( else: if not helpers._check_proxy(port): click.echo(f" Warning: No proxy detected on port {port}") + elif vertex_api_url or clear_vertex_api_url: + health_payload = helpers._query_proxy_health(port) + running_config = helpers._proxy_health_config(health_payload) + if running_config is None: + running_config = helpers._query_proxy_config(port) + running_vertex_url = ( + _normalize_proxy_api_url(running_config.get("vertex_api_url")) + if running_config is not None + else None + ) + requested_vertex_url = _normalize_proxy_api_url(vertex_api_url) + if running_vertex_url != requested_vertex_url: + click.echo( + " Warning: --no-proxy is set, but the running proxy does not " + "advertise the requested Vertex target. Requests may still go " + "to the proxy's existing Vertex upstream." + ) return None @@ -3532,6 +3603,7 @@ def claude( proxy_holder: list[subprocess.Popen | None] = [None] _saved_base_url: list[str | None] = [None] # previous settings.json value for restore _settings_foundry: list[bool] = [False] + _settings_vertex: list[bool] = [False] cleanup = _make_cleanup(proxy_holder, port) _register_proxy_client(port) signal.signal(signal.SIGINT, _ignore_child_sigint) @@ -3606,6 +3678,8 @@ def claude( # location) using Claude Code's own ADC token — no API key, no creds held # by Headroom. This is the turnkey Vertex compression path. use_vertex = bool(os.environ.get("CLAUDE_CODE_USE_VERTEX")) + proxy_url = _claude_proxy_base_url(port) + vertex_upstream = _vertex_target_api_url_from_claude_env(proxy_url) if use_vertex else None proxy_holder[0] = _ensure_proxy( port, @@ -3617,6 +3691,8 @@ def claude( backend=backend, region=region, anthropic_api_url=foundry_upstream, + vertex_api_url=vertex_upstream, + clear_vertex_api_url=use_vertex and vertex_upstream is None, ) _push_runtime_env(port, no_proxy) @@ -3652,7 +3728,6 @@ def claude( if code_graph: _setup_code_graph(verbose=verbose) - proxy_url = _claude_proxy_base_url(port) click.echo() click.echo(" Launching Claude Code (API routed through Headroom)...") if use_vertex: @@ -3688,10 +3763,18 @@ def claude( # Issue #951: write to settings.json so daemon-spawned conversation # workers (which read settings.json fresh rather than inheriting the # daemon's environment) also route through Headroom. - _settings_foundry[0] = bool(foundry_upstream) + _settings_vertex[0] = bool(use_vertex) + _settings_foundry[0] = bool(foundry_upstream) and not _settings_vertex[0] _saved_base_url[0] = _write_claude_wrap_base_url( - _foundry_proxy_url(proxy_url) if _settings_foundry[0] else proxy_url, + ( + _foundry_proxy_url(proxy_url) + if _settings_foundry[0] + else env["ANTHROPIC_VERTEX_BASE_URL"] + if _settings_vertex[0] + else proxy_url + ), foundry_mode=_settings_foundry[0], + vertex_mode=_settings_vertex[0], ) # Per-project savings attribution: tag every request with the launch @@ -3731,7 +3814,11 @@ def claude( click.echo(f" Error: {e}") raise SystemExit(1) from e finally: - _restore_claude_wrap_base_url(_saved_base_url[0], foundry_mode=_settings_foundry[0]) + _restore_claude_wrap_base_url( + _saved_base_url[0], + foundry_mode=_settings_foundry[0], + vertex_mode=_settings_vertex[0], + ) cleanup() @@ -3800,6 +3887,7 @@ def unwrap_claude( _restore_claude_wrap_base_url(None) _restore_claude_wrap_base_url(None, foundry_mode=True) + _restore_claude_wrap_base_url(None, vertex_mode=True) click.echo() click.echo("✓ Claude is no longer durably wrapped by Headroom.") diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 03749011b..5b35afd57 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -2277,6 +2277,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: "openai_api_url": config.openai_api_url, "gemini_api_url": config.gemini_api_url, "cloudcode_api_url": config.cloudcode_api_url, + "vertex_api_url": config.vertex_api_url, "savings_profile": config.savings_profile, "target_ratio": effective_target_ratio, "target_savings_percent": ( diff --git a/tests/test_cli/test_unwrap_claude.py b/tests/test_cli/test_unwrap_claude.py index 0af53809a..00003567a 100644 --- a/tests/test_cli/test_unwrap_claude.py +++ b/tests/test_cli/test_unwrap_claude.py @@ -210,6 +210,26 @@ def test_unwrap_claude_keep_flags_skip_cleanup( stop_proxy.assert_not_called() +def test_unwrap_claude_restores_all_base_url_modes(runner: CliRunner) -> None: + restore_calls: list[dict[str, object]] = [] + + def restore_base_url(previous: str | None, **kwargs: object) -> None: + restore_calls.append({"previous": previous, **kwargs}) + + with patch("headroom.cli.wrap._restore_claude_wrap_base_url", side_effect=restore_base_url): + result = runner.invoke( + main, + ["unwrap", "claude", "--keep-mcp", "--keep-rtk", "--no-stop-proxy"], + ) + + assert result.exit_code == 0, result.output + assert restore_calls == [ + {"previous": None}, + {"previous": None, "foundry_mode": True}, + {"previous": None, "vertex_mode": True}, + ] + + def test_remove_claude_rtk_hooks_removes_init_hooks_and_env(tmp_path: Path) -> None: settings = tmp_path / "settings.json" settings.write_text( diff --git a/tests/test_cli/test_wrap_claude_vertex_proxy_env.py b/tests/test_cli/test_wrap_claude_vertex_proxy_env.py new file mode 100644 index 000000000..33f2eb239 --- /dev/null +++ b/tests/test_cli/test_wrap_claude_vertex_proxy_env.py @@ -0,0 +1,373 @@ +"""Claude wrap Vertex upstream handoff tests.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pytest +from click.testing import CliRunner + +from headroom.cli import wrap as wrap_mod +from headroom.cli.main import main +from headroom.providers.registry import DEFAULT_VERTEX_API_URL + + +class _Completed: + returncode = 0 + + +class _FakeProxyProcess: + returncode = None + + def poll(self) -> None: + return None + + def kill(self) -> None: + return None + + +@pytest.fixture +def runner() -> CliRunner: + return CliRunner() + + +def _clear_claude_mode_env(monkeypatch: pytest.MonkeyPatch) -> None: + for key in ( + "ANTHROPIC_VERTEX_BASE_URL", + "ANTHROPIC_FOUNDRY_BASE_URL", + "ANTHROPIC_FOUNDRY_RESOURCE", + "CLAUDE_CODE_USE_VERTEX", + "CLAUDE_CODE_USE_FOUNDRY", + "VERTEX_TARGET_API_URL", + ): + monkeypatch.delenv(key, raising=False) + + +def _invoke_wrap_claude( + runner: CliRunner, + monkeypatch: pytest.MonkeyPatch, + *, + env: dict[str, str], +) -> tuple[dict[str, Any], str]: + captured: dict[str, Any] = {} + + _clear_claude_mode_env(monkeypatch) + monkeypatch.setattr(wrap_mod.shutil, "which", lambda _name: "/usr/bin/claude") + monkeypatch.setattr(wrap_mod, "_register_proxy_client", lambda _port: None) + monkeypatch.setattr(wrap_mod, "_make_cleanup", lambda _holder, _port: lambda: None) + monkeypatch.setattr(wrap_mod.signal, "signal", lambda *_args, **_kwargs: None) + monkeypatch.setattr(wrap_mod, "_push_runtime_env", lambda *_args, **_kwargs: None) + monkeypatch.setattr(wrap_mod, "_setup_coding_compressor", lambda *_args, **_kwargs: None) + + def fake_write_base_url(*args: object, **kwargs: object) -> None: + captured["write_base_url_args"] = args + captured["write_base_url_kwargs"] = kwargs + + monkeypatch.setattr(wrap_mod, "_write_claude_wrap_base_url", fake_write_base_url) + monkeypatch.setattr(wrap_mod, "_restore_claude_wrap_base_url", lambda *_args, **_kwargs: None) + monkeypatch.setattr(wrap_mod, "_print_telemetry_notice", lambda: None) + + def fake_ensure_proxy(*args: object, **kwargs: object) -> None: + captured["ensure_args"] = args + captured["ensure_kwargs"] = kwargs + + def fake_run(cmd: list[str], *, env: dict[str, str]) -> _Completed: + captured["child_cmd"] = cmd + captured["child_env"] = env + return _Completed() + + monkeypatch.setattr(wrap_mod, "_ensure_proxy", fake_ensure_proxy) + monkeypatch.setattr(wrap_mod.subprocess, "run", fake_run) + + result = runner.invoke( + main, + [ + "wrap", + "claude", + "--no-context-tool", + "--no-mcp", + "--no-tokensave", + "--no-serena", + ], + env=env, + ) + + assert result.exit_code == 0, result.output + return captured, result.output + + +def test_wrap_claude_vertex_passes_custom_base_url_to_proxy_before_child_redirect( + runner: CliRunner, monkeypatch: pytest.MonkeyPatch +) -> None: + custom_vertex_url = "https://vertex-gateway.internal/custom/v1" + + captured, _output = _invoke_wrap_claude( + runner, + monkeypatch, + env={ + "CLAUDE_CODE_USE_VERTEX": "1", + "ANTHROPIC_VERTEX_BASE_URL": custom_vertex_url, + }, + ) + + ensure_kwargs = captured["ensure_kwargs"] + child_env = captured["child_env"] + write_kwargs = captured["write_base_url_kwargs"] + assert ensure_kwargs["vertex_api_url"] == custom_vertex_url + assert ensure_kwargs["clear_vertex_api_url"] is False + assert ensure_kwargs["anthropic_api_url"] is None + assert child_env["ANTHROPIC_VERTEX_BASE_URL"] == "http://127.0.0.1:8787" + assert write_kwargs["vertex_mode"] is True + assert write_kwargs["foundry_mode"] is False + + +def test_wrap_claude_vertex_target_env_beats_anthropic_vertex_base_url( + runner: CliRunner, monkeypatch: pytest.MonkeyPatch +) -> None: + captured, _output = _invoke_wrap_claude( + runner, + monkeypatch, + env={ + "CLAUDE_CODE_USE_VERTEX": "1", + "ANTHROPIC_VERTEX_BASE_URL": "https://client-gateway.example.com/vertex/v1", + "VERTEX_TARGET_API_URL": "https://proxy-gateway.example.com/vertex/v1", + }, + ) + + ensure_kwargs = captured["ensure_kwargs"] + child_env = captured["child_env"] + assert ensure_kwargs["vertex_api_url"] == "https://proxy-gateway.example.com/vertex/v1" + assert ensure_kwargs["clear_vertex_api_url"] is False + assert child_env["ANTHROPIC_VERTEX_BASE_URL"] == "http://127.0.0.1:8787" + + +@pytest.mark.parametrize( + "env", + [ + {"CLAUDE_CODE_USE_VERTEX": "1"}, + { + "CLAUDE_CODE_USE_VERTEX": "1", + "ANTHROPIC_VERTEX_BASE_URL": DEFAULT_VERTEX_API_URL, + }, + { + "CLAUDE_CODE_USE_VERTEX": "1", + "ANTHROPIC_VERTEX_BASE_URL": "http://127.0.0.1:8787", + }, + { + "CLAUDE_CODE_USE_VERTEX": "1", + "VERTEX_TARGET_API_URL": "http://127.0.0.1:8787", + }, + ], +) +def test_wrap_claude_vertex_default_or_absent_base_url_does_not_force_vertex_target( + runner: CliRunner, monkeypatch: pytest.MonkeyPatch, env: dict[str, str] +) -> None: + captured, _output = _invoke_wrap_claude(runner, monkeypatch, env=env) + + ensure_kwargs = captured["ensure_kwargs"] + child_env = captured["child_env"] + assert ensure_kwargs["vertex_api_url"] is None + assert ensure_kwargs["clear_vertex_api_url"] is True + assert child_env["ANTHROPIC_VERTEX_BASE_URL"] == "http://127.0.0.1:8787" + + +def test_wrap_claude_foundry_proxy_env_behavior_is_unchanged( + runner: CliRunner, monkeypatch: pytest.MonkeyPatch +) -> None: + foundry_url = "https://my-resource.services.ai.azure.com/anthropic" + + captured, _output = _invoke_wrap_claude( + runner, + monkeypatch, + env={ + "CLAUDE_CODE_USE_FOUNDRY": "1", + "ANTHROPIC_FOUNDRY_BASE_URL": foundry_url, + }, + ) + + ensure_kwargs = captured["ensure_kwargs"] + child_env = captured["child_env"] + assert ensure_kwargs["anthropic_api_url"] == foundry_url + assert ensure_kwargs["vertex_api_url"] is None + assert child_env["ANTHROPIC_FOUNDRY_BASE_URL"] == "http://127.0.0.1:8787/anthropic" + assert captured["write_base_url_kwargs"]["foundry_mode"] is True + assert captured["write_base_url_kwargs"]["vertex_mode"] is False + + +def test_write_vertex_mode_sets_vertex_key(tmp_path: Path) -> None: + path = tmp_path / ".claude" / "settings.local.json" + + previous = wrap_mod._write_claude_wrap_base_url( + "http://127.0.0.1:8787", + vertex_mode=True, + settings_path=path, + ) + + assert previous is None + payload = path.read_text(encoding="utf-8") + assert '"ANTHROPIC_VERTEX_BASE_URL": "http://127.0.0.1:8787"' in payload + assert "ANTHROPIC_BASE_URL" not in payload + + +def test_restore_vertex_mode_restores_previous_vertex_key(tmp_path: Path) -> None: + path = tmp_path / ".claude" / "settings.local.json" + wrap_mod._write_claude_wrap_base_url( + "http://127.0.0.1:8787", + vertex_mode=True, + settings_path=path, + ) + + wrap_mod._restore_claude_wrap_base_url( + "https://existing-gateway.example.com/vertex/v1", + vertex_mode=True, + settings_path=path, + ) + + payload = path.read_text(encoding="utf-8") + assert ( + '"ANTHROPIC_VERTEX_BASE_URL": "https://existing-gateway.example.com/vertex/v1"' in payload + ) + + +def test_start_proxy_sets_vertex_target_env_for_proxy_subprocess( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + fake_proc = _FakeProxyProcess() + captured: dict[str, Any] = {} + + monkeypatch.setattr(wrap_mod, "_get_log_path", lambda: tmp_path / "proxy.log") + monkeypatch.setattr(wrap_mod, "_check_proxy", lambda _port: True) + monkeypatch.setattr(wrap_mod.time, "sleep", lambda _seconds: None) + + def fake_popen(cmd: list[str], **kwargs: object) -> _FakeProxyProcess: + captured["cmd"] = cmd + captured["kwargs"] = kwargs + return fake_proc + + monkeypatch.setattr(wrap_mod.subprocess, "Popen", fake_popen) + + proc = wrap_mod._start_proxy( + 8787, + agent_type="claude", + vertex_api_url="https://vertex-gateway.internal/custom", + ) + + assert proc is fake_proc + assert captured["cmd"][-2:] == [ + "--vertex-api-url", + "https://vertex-gateway.internal/custom", + ] + proxy_env = captured["kwargs"]["env"] + assert proxy_env["VERTEX_TARGET_API_URL"] == "https://vertex-gateway.internal/custom" + + +def test_start_proxy_clears_inherited_vertex_target_env( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + fake_proc = _FakeProxyProcess() + captured: dict[str, Any] = {} + + monkeypatch.setenv("VERTEX_TARGET_API_URL", "http://127.0.0.1:8787") + monkeypatch.setattr(wrap_mod, "_get_log_path", lambda: tmp_path / "proxy.log") + monkeypatch.setattr(wrap_mod, "_check_proxy", lambda _port: True) + monkeypatch.setattr(wrap_mod.time, "sleep", lambda _seconds: None) + + def fake_popen(cmd: list[str], **kwargs: object) -> _FakeProxyProcess: + captured["cmd"] = cmd + captured["kwargs"] = kwargs + return fake_proc + + monkeypatch.setattr(wrap_mod.subprocess, "Popen", fake_popen) + + proc = wrap_mod._start_proxy(8787, agent_type="claude", clear_vertex_api_url=True) + + assert proc is fake_proc + assert "--vertex-api-url" not in captured["cmd"] + proxy_env = captured["kwargs"]["env"] + assert "VERTEX_TARGET_API_URL" not in proxy_env + + +def test_ensure_proxy_restarts_idle_proxy_for_vertex_api_url_mismatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[object] = [] + health = { + "version": wrap_mod._HEADROOM_VERSION, + "runtime": {"websocket_sessions": {"active_sessions": 0, "active_relay_tasks": 0}}, + "config": { + "pid": "12345", + "memory": False, + "learn": False, + "code_graph": False, + "vertex_api_url": "https://old-gateway.example.com/vertex/v1", + }, + } + + monkeypatch.setattr(wrap_mod, "_find_persistent_manifest", lambda _port: None) + monkeypatch.setattr(wrap_mod, "_check_proxy", lambda _port: len(calls) == 0) + monkeypatch.setattr(wrap_mod, "_query_proxy_health", lambda _port: health) + monkeypatch.setattr(wrap_mod, "_port_bind_error", lambda _port: None) + monkeypatch.setattr(wrap_mod, "_live_proxy_clients", lambda *args, **kwargs: []) + monkeypatch.setattr( + wrap_mod, + "_kill_proxy_by_pid", + lambda pid, port: calls.append(("kill", pid, port)) or True, + ) + monkeypatch.setattr( + wrap_mod, + "_start_proxy", + lambda *args, **kwargs: calls.append(("start", args, kwargs)), + ) + + result = wrap_mod._ensure_proxy( + 8787, + False, + vertex_api_url="https://new-gateway.example.com/vertex/v1", + ) + + assert result is None + assert calls[0] == ("kill", 12345, 8787) + assert calls[1][0] == "start" + assert calls[1][2]["vertex_api_url"] == "https://new-gateway.example.com/vertex/v1" + + +def test_ensure_proxy_restarts_idle_proxy_to_clear_vertex_api_url( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[object] = [] + health = { + "version": wrap_mod._HEADROOM_VERSION, + "runtime": {"websocket_sessions": {"active_sessions": 0, "active_relay_tasks": 0}}, + "config": { + "pid": "12345", + "memory": False, + "learn": False, + "code_graph": False, + "vertex_api_url": "https://old-gateway.example.com/vertex/v1", + }, + } + + monkeypatch.setattr(wrap_mod, "_find_persistent_manifest", lambda _port: None) + monkeypatch.setattr(wrap_mod, "_check_proxy", lambda _port: len(calls) == 0) + monkeypatch.setattr(wrap_mod, "_query_proxy_health", lambda _port: health) + monkeypatch.setattr(wrap_mod, "_port_bind_error", lambda _port: None) + monkeypatch.setattr(wrap_mod, "_live_proxy_clients", lambda *args, **kwargs: []) + monkeypatch.setattr( + wrap_mod, + "_kill_proxy_by_pid", + lambda pid, port: calls.append(("kill", pid, port)) or True, + ) + monkeypatch.setattr( + wrap_mod, + "_start_proxy", + lambda *args, **kwargs: calls.append(("start", args, kwargs)), + ) + + result = wrap_mod._ensure_proxy(8787, False, clear_vertex_api_url=True) + + assert result is None + assert calls[0] == ("kill", 12345, 8787) + assert calls[1][0] == "start" + assert calls[1][2]["vertex_api_url"] is None + assert calls[1][2]["clear_vertex_api_url"] is True