feat(vertex): turnkey Claude Code + Vertex compression (+ fixes from the Vertex review) (#1113)

## Description

Makes **Claude Code on Google Vertex AI** actually receive Headroom's
prompt compression, and fixes the issues found in a deep review of the
Vertex path. The headline is a turnkey path: `headroom wrap claude`
(with the user's existing Vertex env) compresses each request and
forwards to Vertex using the client's own GCP ADC token — Headroom holds
no credentials.

_No linked issue — this addresses the internal Vertex code review
(`docs/proposals/vertex-claude-compression-review.md`)._

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [x] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [x] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- `cli/wrap.py`: `wrap claude` detects `CLAUDE_CODE_USE_VERTEX=1` and
points Claude Code's Vertex endpoint at the proxy via
`ANTHROPIC_VERTEX_BASE_URL` (Claude Code ignores `ANTHROPIC_BASE_URL` in
Vertex mode). Client keeps its own GCP ADC auth. Adds
`--backend`/`--region` flags (parity with `wrap aider`).
- `providers/registry.py`: alias `litellm-vertex` → provider
`vertex_ai`. Previously it resolved to `"vertex"` (not in the registry)
→ generic pass-through with the wrong model prefix, dropped region, and
mishandled auth, even though all help text advertises `litellm-vertex`.
- `providers/proxy_routes.py`: derive the Vertex upstream host
per-request from the path's `locations/{location}` (handles `global`)
instead of pinning the configured fixed-region host; explicit
`--vertex-api-url` overrides still win.
- `docs/content/docs/claude-code-vertex.mdx` (+ nav): simple user guide
for running Claude Code on Vertex through Headroom.
- `docs/proposals/vertex-claude-compression-review.md`: the deep-review
findings these fixes address.
- `tests/test_vertex_claude_compression.py`: new tests.

## 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
$ python -m pytest tests/test_vertex_claude_compression.py -q
8 passed

$ python -m pytest tests/test_provider_proxy_routes.py tests/test_provider_registry.py tests/test_cli_proxy_env.py tests/test_backend_bugs.py -q
108 passed

$ ruff check headroom/providers/registry.py headroom/providers/proxy_routes.py headroom/cli/wrap.py tests/test_vertex_claude_compression.py
All checks passed!

$ mypy headroom/providers/registry.py headroom/providers/proxy_routes.py headroom/cli/wrap.py
Success: no issues found in 3 source files
```

## Real Behavior Proof

- Environment: local macOS, Python 3.12 `.venv`, branch
`feat/vertex-claude-compression`.
- Exact command / steps: ran the test suite above; verified in code that
the native `:rawPredict` route (publisher=anthropic) delegates to
`handle_anthropic_messages` with the region-derived host, that
`create_proxy_backend("litellm-vertex")` resolves to provider
`vertex_ai`, and that `wrap claude` sets `ANTHROPIC_VERTEX_BASE_URL`
when `CLAUDE_CODE_USE_VERTEX` is set.
- Observed result: 8 new tests + 108 existing tests pass; ruff + mypy
clean; the alias, region derivation (incl. `global` and explicit
override), and rawPredict→compression-handler delegation all behave as
asserted.
- Not tested: a live end-to-end run of Claude Code against a real Google
Vertex project (no GCP credentials available in this environment).
Recommend one smoke test against a live Vertex project before announcing
GA. The Rust `headroom-proxy` Vertex path is intentionally out of scope
(separate, unwired binary).

## 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
- [ ] I have updated the CHANGELOG.md if applicable

## Additional Notes

- CHANGELOG.md left unchanged — releases are managed by release-please
from conventional commits.
- Follow-ups (not in this PR): wire or formally retire the Rust
`headroom-proxy` Vertex implementation; add a live-Vertex smoke test
once CI has GCP credentials.
This commit is contained in:
Tejas Chopra 2026-06-18 00:56:23 -07:00 committed by GitHub
parent e45cf4e061
commit 0e0591506c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 582 additions and 6 deletions

View file

@ -0,0 +1,117 @@
---
title: Claude Code on Vertex AI
description: Run Claude Code against Claude models on Google Vertex AI, with Headroom compressing your prompts — fewer input tokens, same answers, your own GCP login.
---
If your Claude models live on **Google Vertex AI**, you can still get Headroom's
prompt compression. Headroom sits between Claude Code and Vertex: it shrinks the
big stuff in each request (file reads, logs, tool output) and forwards the rest to
Vertex using **your own Google credentials**. You keep your GCP setup; Headroom
just makes each call cheaper.
## What you get
- **Fewer input tokens** on every Claude Code request to Vertex (often 3060% on
agent workloads), so you pay Vertex for less.
- **Same answers** — compression is reversible and content-aware.
- **No new secrets** — Headroom never holds your Google credentials. Claude Code
keeps authenticating to Vertex with its own ADC token; Headroom passes it through.
## Before you start
You should already have Claude Code working against Vertex **without** Headroom.
That means these are set in your shell:
```bash
export CLAUDE_CODE_USE_VERTEX=1
export ANTHROPIC_VERTEX_PROJECT_ID=<your-gcp-project>
export CLOUD_ML_REGION=us-east5 # your Vertex region (or "global")
gcloud auth application-default login # or set GOOGLE_APPLICATION_CREDENTIALS
```
No `ANTHROPIC_API_KEY` is needed — Vertex mode uses your Google login.
## Run it (one command)
```bash
pip install headroom
headroom wrap claude
```
That's it. Because `CLAUDE_CODE_USE_VERTEX=1` is set, `headroom wrap claude`
automatically:
1. starts the Headroom proxy,
2. points Claude Code's Vertex endpoint at it (`ANTHROPIC_VERTEX_BASE_URL`),
3. leaves your project, region, and Google login untouched.
You'll see a line like:
```
Vertex mode: ANTHROPIC_VERTEX_BASE_URL=http://127.0.0.1:8787
→ compress, then forward to Vertex with your GCP ADC token
```
Use Claude Code exactly as you normally would.
## How it works
```
Claude Code ──(Vertex request)──▶ Headroom ──(compressed)──▶ Vertex AI (Claude)
in Vertex mode compresses your project + region
(your ADC token) ───────────── passed through ───────────▶ authenticates you
```
Claude Code sends its normal Vertex `…:rawPredict` / `:streamRawPredict` request to
Headroom. Headroom compresses the messages (keeping the Vertex request shape
intact), then forwards to the correct regional Vertex host — derived from the
request itself, so multi-region and `global` both work — using the Google token
Claude Code already attached.
## Check that compression is working
1. Open the dashboard: [http://localhost:8787/dashboard](http://localhost:8787/dashboard).
"Tokens saved" should climb as you use Claude Code.
2. Or look at the response headers on a request: `x-headroom-tokens-before`,
`x-headroom-tokens-after`, `x-headroom-tokens-saved`.
If "tokens saved" stays at 0 on large prompts, see Troubleshooting below.
## Troubleshooting
- **It still goes straight to Google (no savings).** Make sure `CLAUDE_CODE_USE_VERTEX=1`
is exported *in the same shell* before `headroom wrap claude`. The wrapper only
switches to Vertex mode when it sees that variable.
- **Wrong region / 404 from Vertex.** Confirm `CLOUD_ML_REGION` matches a region
where your Claude model is enabled. `global` is supported and maps to the
non-regional host.
- **Auth errors.** Headroom forwards your token as-is — if `gcloud auth
application-default login` (or `GOOGLE_APPLICATION_CREDENTIALS`) works for Claude
Code without Headroom, it works with it.
## Alternative: let Headroom talk to Vertex for you
If you'd rather **not** run Claude Code in Vertex mode, you can have Headroom be the
translator instead: Claude Code speaks plain Anthropic to Headroom, and Headroom
calls Vertex on your behalf.
```bash
export HEADROOM_BACKEND=litellm-vertex_ai # note the _ai suffix
export HEADROOM_REGION=us-east5
export VERTEXAI_PROJECT=<your-gcp-project>
export GOOGLE_APPLICATION_CREDENTIALS=/path/sa.json # or gcloud ADC
export ANTHROPIC_API_KEY=placeholder # Claude Code needs *a* key to start
headroom wrap claude --backend litellm-vertex_ai --region us-east5
```
The native Vertex-mode flow above is recommended — it keeps your existing GCP auth
and has the smallest moving parts. Use this alternative only if you can't set
`CLAUDE_CODE_USE_VERTEX`.
## Notes
- Pick a Claude model that is enabled in your Vertex project/region
(e.g. `claude-sonnet-4-6`, `claude-haiku-4-5`).
- Streaming, tool use, and prompt caching all work through Headroom.
- Want to point at a private Vertex gateway instead of Google's host? Start the
proxy with `--vertex-api-url https://your-gateway` and Headroom will forward there.

View file

@ -32,6 +32,7 @@
"agno", "agno",
"strands", "strands",
"litellm", "litellm",
"claude-code-vertex",
"mcp", "mcp",
"---Configuration---", "---Configuration---",
"configuration", "configuration",

View file

@ -0,0 +1,255 @@
# Claude Code + Vertex AI + Headroom compression — does it work?
*A plain-English deep code review. Last updated 2026-06-18.*
## TL;DR (read this first)
**Yes, Headroom can compress Claude traffic that goes to Vertex AI — but only if you
set it up by hand, and only if you work around one real bug.** There is **no
turnkey path today**: `headroom wrap claude` does not know anything about Vertex,
and the one backend flag everybody is told to use (`--backend litellm-vertex`) is
broken.
Two setups actually work (both need manual environment variables):
- **Setup A — keep Claude Code in Vertex mode**, and point its Vertex URL at
Headroom. Headroom compresses, then forwards to real Vertex.
- **Setup B — run Claude Code in normal Anthropic mode**, and let Headroom be the
translator that talks to Vertex (`--backend litellm-vertex_ai`).
Everything in the *middle* (compression, request/response translation, streaming,
tool calls) is implemented correctly. The gaps are all at the **edges**: how the
client is pointed at Headroom, one mis-named backend, and a few env vars Headroom
never sets for you.
> Correction to an earlier claim: it is **not** true that "the Python proxy just
> passes Vertex through without compressing." For the Anthropic publisher it runs
> the full compression pipeline. That earlier statement was based on an incomplete
> read of the routing code; the verified behavior is in this doc.
---
## The thing we're trying to do
An enterprise runs **Claude Code**, but their Claude models live on **Google
Vertex AI** (not the direct Anthropic API). They want **Headroom** in the middle so
their prompts get compressed (fewer input tokens = lower cost), without changing
the answers.
For that to happen, three things must all be true:
1. **The client's traffic must actually reach Headroom** (the proxy must be in the path).
2. **Headroom must compress it.**
3. **Headroom must forward it to Vertex correctly** (right URL, right auth, right body shape) and translate the answer back so Claude Code understands it.
This review checks all three.
---
## The map: where Claude-on-Vertex can run, and what compresses
There are **two proxies** in this repo and **three** possible routes. Only some compress.
| Route | What it is | Compresses? | Notes |
|---|---|---|---|
| **Python proxy, native Vertex `:rawPredict`** (publisher = `anthropic`) | Client sends a real Vertex request to Headroom | ✅ **Yes** | Runs the full Anthropic compression pipeline, keeps the Vertex body shape, forwards the client's own Google token. `proxy_routes.py:648` |
| **Python proxy, `--backend litellm-vertex_ai`** | Client speaks plain Anthropic; Headroom translates to Vertex | ✅ **Yes** (correct string only) | Full Anthropic↔Vertex translation incl. streaming + tools. **`litellm-vertex` is broken — must use `litellm-vertex_ai`.** |
| **Rust proxy, native Vertex `:rawPredict`** | A separate `headroom-proxy` binary | ✅ **Yes** | Correct and well-built — **but never run by `headroom proxy`/`wrap`.** Dead code for normal users. |
| **Python proxy, passthrough** (any *other* publisher) | Generic verbatim forward | ❌ No | Only used for non-Anthropic, non-Google publishers. `openai.py:6014` |
**Key takeaway:** the *compression engine* for Vertex+Claude exists and works in the
Python proxy. The problems are getting traffic into it and one naming bug.
---
## Does it work end-to-end? The honest answer
**Through `headroom wrap claude` with zero extra setup: no.** `wrap claude` only
sets `ANTHROPIC_BASE_URL`. If Claude Code is in Vertex mode it ignores that and
talks straight to Google — Headroom is never in the path. And `wrap claude` has no
`--backend`/`--region` flags and sets no Vertex environment variables.
**With manual setup: yes, one of two ways.** Both are below. Both work *around*
issues, not because the product wires them for you.
---
## Setup A — Claude Code stays in Vertex mode (recommended for Vertex shops)
Idea: Claude Code keeps using its native Vertex mode and its own Google login.
You just tell it "send Vertex requests to Headroom instead of straight to Google,"
and you tell Headroom where the real Vertex endpoint is.
```bash
# 1) Run Headroom, telling it the real Vertex endpoint (match your region!)
headroom proxy --port 8787 \
--vertex-api-url https://us-east5-aiplatform.googleapis.com # use YOUR region
# 2) Run Claude Code in Vertex mode, but point its Vertex base URL at Headroom
export CLAUDE_CODE_USE_VERTEX=1
export ANTHROPIC_VERTEX_PROJECT_ID=<your-gcp-project>
export CLOUD_ML_REGION=us-east5
export ANTHROPIC_VERTEX_BASE_URL=http://127.0.0.1:8787 # <-- the redirect that makes it work
claude
```
What happens: Claude Code → `ANTHROPIC_VERTEX_BASE_URL` (Headroom) → Headroom
matches the `:rawPredict` route, sees `publisher=anthropic`, **compresses**, then
forwards to the real Vertex endpoint using Claude Code's own Google token.
Caveats: you must set `--vertex-api-url` to your region (see Issue #6), and
`wrap claude` won't set `ANTHROPIC_VERTEX_BASE_URL` for you (Issue #3).
---
## Setup B — Claude Code in normal Anthropic mode; Headroom talks to Vertex
Idea: Claude Code thinks it's talking to plain Anthropic. Headroom holds the Google
credentials and is the one that actually talks to Vertex.
```bash
# Headroom does the Vertex talking — note the backend name carefully
export HEADROOM_BACKEND=litellm-vertex_ai # NOT "litellm-vertex" (that's broken — Issue #1)
export HEADROOM_REGION=us-east5 # becomes the Vertex location
export VERTEXAI_PROJECT=<your-gcp-project> # Headroom does NOT set this for you (Issue #4)
export GOOGLE_APPLICATION_CREDENTIALS=/path/sa.json # or use `gcloud auth application-default login`
export ANTHROPIC_API_KEY=placeholder-not-used # Claude Code needs *a* key to start (Issue #5)
# Do NOT set CLAUDE_CODE_USE_VERTEX here — Claude Code must stay in normal mode
headroom wrap claude
```
What happens: Claude Code → Headroom (plain Anthropic `/v1/messages`) →
**compresses** → LiteLLM converts to Vertex and calls Claude on Vertex → converts
the answer back to Anthropic shape → Claude Code reads it.
Caveats: the backend-name bug (Issue #1), the missing project env (Issue #4), and
this path has **no automated tests** (Issue #8) — smoke-test it before relying on it.
---
## How to verify compression is really happening
1. Open the dashboard: `http://localhost:8787/dashboard` — "tokens saved" should
climb as you use Claude Code.
2. Or check response headers on a request: `x-headroom-tokens-before`,
`x-headroom-tokens-after`, `x-headroom-tokens-saved`.
3. Confirm it actually hit Vertex (proxy logs show a `vertex_ai/claude-…` model or
a Vertex host, not `api.anthropic.com`).
If `tokens-saved` is 0 on large prompts, compression isn't running — re-check the
setup against the issues below.
---
## Every issue we found (the full list)
Severity: **BROKEN** = doesn't work; **GAP** = works only with manual workaround;
**BUG** = wrong behavior in an edge case; **HOUSEKEEPING** = confusing but harmless.
### 1. BROKEN — `--backend litellm-vertex` never reaches Vertex
The backend name is turned into a provider by chopping off `litellm-`, so
`litellm-vertex` becomes the provider `vertex`. But the Vertex integration is keyed
on `vertex_ai`, not `vertex`. So Headroom falls back to a generic "unknown
provider" mode: it builds the wrong model name (`vertex/claude-…` instead of
`vertex_ai/claude-…`), **ignores the region**, and mishandles auth.
**You must use `--backend litellm-vertex_ai`.** Worse: every help message and the
`wrap` example tell users the broken `litellm-vertex`.
*Where:* `providers/registry.py:174-178`, `backends/litellm.py:291,326-336,681-682`;
help text at `cli/proxy.py:524`, `cli/wrap.py:3645`, `proxy/server.py:3913`.
*Fix (small):* alias `vertex``vertex_ai` in `create_proxy_backend`, or add a
`"vertex"` entry to the provider registry. Then fix the help text.
### 2. GAP — `headroom wrap claude` has no Vertex support
The `claude` wrap command has no `--backend` and no `--region` (the `aider` wrap
command has both). It only ever sets `ANTHROPIC_BASE_URL`. So there's no flag to
turn on a Vertex backend for Claude Code — you must pre-export env vars.
*Where:* `cli/wrap.py:2780-2819` (vs `cli/wrap.py:3612,3615` for aider).
*Fix:* add `--backend`/`--region` to `wrap claude`, mirroring `aider`.
### 3. GAP — Vertex-mode Claude Code bypasses the proxy, and Headroom never sets the fix
With `CLAUDE_CODE_USE_VERTEX=1`, Claude Code ignores `ANTHROPIC_BASE_URL` and goes
straight to Google. There **is** a documented override — `ANTHROPIC_VERTEX_BASE_URL`
— that points Claude Code's Vertex traffic at a gateway. But Headroom never sets it
(0 references in the codebase). So the proxy has the right routes, but nothing
connects the client to them automatically.
*Where:* repo-wide grep for `ANTHROPIC_VERTEX_BASE_URL` = 0 hits.
*Fix:* in a Vertex-aware `wrap claude`, set `ANTHROPIC_VERTEX_BASE_URL` to the proxy.
### 4. GAP — the GCP project is never passed to LiteLLM
For Setup B, Headroom passes the region to LiteLLM but never the project. You must
export `VERTEXAI_PROJECT` (or `GOOGLE_CLOUD_PROJECT`) yourself or it fails.
*Where:* `backends/litellm.py:682` sets only `vertex_location`.
*Fix:* thread a project config/env through to the LiteLLM call.
### 5. GAP — no placeholder API key for Claude Code
With a custom `ANTHROPIC_BASE_URL`, Claude Code needs *an* `ANTHROPIC_API_KEY` to
start, even though the proxy uses Google creds upstream. `wrap claude` never sets a
placeholder, so the user must.
*Where:* `cli/wrap.py:2984-2988`.
*Fix:* set a placeholder key (or `ANTHROPIC_AUTH_TOKEN`) when launching.
### 6. BUG — region/host mismatch
Headroom pins the Vertex host to one region (default `us-central1`) but throws away
the region in the client's request path. If your client targets, say,
`europe-west1` while Headroom is on the default host, the request goes to the wrong
region unless you set `--vertex-api-url` to match.
*Where:* `copilot_auth.py:936` (host = base + path, no region reconciliation),
`providers/proxy_routes.py:647` (path `location` is discarded), `registry.py:16`.
*Fix:* derive the upstream host from the request path's `location`, or validate they match.
### 7. HOUSEKEEPING — the Rust Vertex proxy is correct but unwired
There's a second, Rust proxy (`crates/headroom-proxy/src/vertex/`) that compresses
Vertex traffic correctly. But `headroom proxy` and `headroom wrap` run the **Python**
server and never call it — it's a separate binary you'd run by hand. This is a
frequent source of "but I thought Vertex compression was added" confusion: it was,
in Rust, on a path nobody runs by default.
*Where:* `crates/headroom-proxy/Cargo.toml` (`[[bin]]`), no Python→Rust bridge to it.
*Fix:* either document that the Rust proxy is separate, or wire/retire it.
### 8. GAP — no tests for the Vertex compression paths
No test instantiates the LiteLLM Vertex backend with a mocked Vertex call, and the
native-Vertex compression route isn't covered against a real Vertex shape. The
translation code is correct by inspection, but unproven by CI.
*Fix:* add a mocked round-trip test (Anthropic in → compressed → Vertex call asserted → Anthropic out).
### 9. HOUSEKEEPING — stale Rust doc comment
`crates/headroom-proxy/src/vertex/mod.rs:42-47` describes a "synthetic model
injection" strategy the code no longer implements. Doc only; behavior is correct.
---
## What's actually solid (so we don't over-correct)
These were verified and are **correct**:
- **Native Vertex `:rawPredict` for `publisher=anthropic` compresses** and preserves
the Vertex body shape (keeps `anthropic_version`, never injects `model`).
`proxy_routes.py:648`, `anthropic.py:604-606,1941-1949`.
- **LiteLLM translation is real and complete** (with the `vertex_ai` provider):
response is rebuilt into Anthropic shape (`backends/litellm.py:575-628`), streaming
emits proper Anthropic SSE events (`streaming.py:1344-1472`, `litellm.py:736-947`),
and tool calls round-trip both directions (`litellm.py:525-565,593-602`).
- **Compression runs before the backend dispatch** (`anthropic.py:1671` then
`:1781`), so the backend always gets the compressed body.
- **Auth is forwarded correctly** on the native path — the client's Google token is
passed through untouched (`copilot_auth.py:1156-1157`).
---
## Recommended fixes, smallest-first
1. **Fix the backend name (Issue #1)** — one-line alias `vertex``vertex_ai`, then
correct the help text. This is the highest-impact, lowest-effort fix; it turns the
*documented* command from broken to working.
2. **Add `--backend`/`--region` to `wrap claude` (Issue #2)** — copy from `aider`.
3. **Add a Vertex mode to `wrap claude` (Issues #3, #5)** — detect/set
`ANTHROPIC_VERTEX_BASE_URL` → proxy, set a placeholder API key, and configure the
proxy's Vertex upstream — so Setup A becomes one command.
4. **Pass the GCP project (Issue #4)** and **reconcile the region/host (Issue #6).**
5. **Add a mocked round-trip test (Issue #8).**
6. **Decide the Rust proxy's fate (Issue #7)** — document-as-separate or wire it in.
After #1#3, the honest customer message becomes: *"`headroom wrap claude` works
with Vertex out of the box."* Until then, it's *"works with a documented manual
setup."*

View file

@ -2814,6 +2814,18 @@ def unwrap() -> None:
"ENABLE_TOOL_SEARCH env var is respected." "ENABLE_TOOL_SEARCH env var is respected."
), ),
) )
@click.option(
"--backend",
default=None,
help="API backend for the proxy: 'anthropic' (default), 'litellm-vertex_ai', etc. "
"(env: HEADROOM_BACKEND). For Vertex, prefer CLAUDE_CODE_USE_VERTEX=1 (native, "
"keeps your GCP auth) over a litellm backend.",
)
@click.option(
"--region",
default=None,
help="Cloud region for Vertex/Bedrock backends (env: HEADROOM_REGION).",
)
@click.option("--verbose", "-v", is_flag=True, help="Verbose output") @click.option("--verbose", "-v", is_flag=True, help="Verbose output")
@click.option("--prepare-only", is_flag=True, hidden=True) @click.option("--prepare-only", is_flag=True, hidden=True)
@click.argument("claude_args", nargs=-1, type=click.UNPROCESSED) @click.argument("claude_args", nargs=-1, type=click.UNPROCESSED)
@ -2827,6 +2839,8 @@ def claude(
learn: bool, learn: bool,
memory: bool, memory: bool,
tool_search: str | None, tool_search: str | None,
backend: str | None,
region: str | None,
verbose: bool, verbose: bool,
prepare_only: bool, prepare_only: bool,
claude_args: tuple, claude_args: tuple,
@ -2928,6 +2942,16 @@ def claude(
if os.environ.get("CLAUDE_CODE_USE_FOUNDRY"): if os.environ.get("CLAUDE_CODE_USE_FOUNDRY"):
foundry_upstream = os.environ.get("ANTHROPIC_FOUNDRY_BASE_URL") foundry_upstream = os.environ.get("ANTHROPIC_FOUNDRY_BASE_URL")
# Detect Vertex mode: with CLAUDE_CODE_USE_VERTEX=1, Claude Code IGNORES
# ANTHROPIC_BASE_URL and authenticates to Google Vertex with GCP ADC. The
# documented way to route its Vertex :rawPredict / :streamRawPredict
# traffic through a gateway is ANTHROPIC_VERTEX_BASE_URL. Point it at
# Headroom and the proxy compresses the request, then forwards to the
# real regional Vertex host (derived per-request from the path's
# 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_holder[0] = _ensure_proxy( proxy_holder[0] = _ensure_proxy(
port, port,
no_proxy, no_proxy,
@ -2935,6 +2959,8 @@ def claude(
memory=memory, memory=memory,
agent_type="claude", agent_type="claude",
code_graph=code_graph, code_graph=code_graph,
backend=backend,
region=region,
anthropic_api_url=foundry_upstream, anthropic_api_url=foundry_upstream,
) )
@ -2970,7 +2996,12 @@ def claude(
proxy_url = _claude_proxy_base_url(port) proxy_url = _claude_proxy_base_url(port)
click.echo() click.echo()
click.echo(" Launching Claude Code (API routed through Headroom)...") click.echo(" Launching Claude Code (API routed through Headroom)...")
if foundry_upstream: if use_vertex:
click.echo(
f" Vertex mode: ANTHROPIC_VERTEX_BASE_URL={proxy_url} "
"→ compress, then forward to Vertex with your GCP ADC token"
)
elif foundry_upstream:
click.echo( click.echo(
f" Foundry mode: ANTHROPIC_FOUNDRY_BASE_URL={proxy_url} → upstream {foundry_upstream}" f" Foundry mode: ANTHROPIC_FOUNDRY_BASE_URL={proxy_url} → upstream {foundry_upstream}"
) )
@ -2982,7 +3013,12 @@ def claude(
click.echo() click.echo()
env = os.environ.copy() env = os.environ.copy()
if foundry_upstream: if use_vertex:
# Claude Code stays in Vertex mode (keeps CLAUDE_CODE_USE_VERTEX,
# ANTHROPIC_VERTEX_PROJECT_ID, CLOUD_ML_REGION, ADC — all inherited);
# we only redirect its Vertex endpoint to Headroom.
env["ANTHROPIC_VERTEX_BASE_URL"] = proxy_url
elif foundry_upstream:
env["ANTHROPIC_FOUNDRY_BASE_URL"] = proxy_url env["ANTHROPIC_FOUNDRY_BASE_URL"] = proxy_url
else: else:
env["ANTHROPIC_BASE_URL"] = proxy_url env["ANTHROPIC_BASE_URL"] = proxy_url

View file

@ -28,6 +28,27 @@ def _api_target(proxy: Any, provider_name: str) -> str:
return cast(str, getattr(proxy, legacy_attr, proxy.provider_runtime.api_target(provider_name))) return cast(str, getattr(proxy, legacy_attr, proxy.provider_runtime.api_target(provider_name)))
def _vertex_target_for_location(proxy: Any, location: str) -> str:
"""Resolve the Vertex upstream host for a request, region-aware.
The Vertex regional host must match the ``locations/{location}`` in the
request path (e.g. a ``europe-west1`` request cannot go to a
``us-central1`` host). The configured target is a single fixed-region host
(default ``us-central1``), so unless the operator pinned an explicit
non-default upstream (e.g. a private gateway), derive the host from the
request's own location. ``global`` maps to the unprefixed host.
"""
from headroom.providers.registry import DEFAULT_VERTEX_API_URL
configured = _api_target(proxy, "vertex")
if configured and configured != DEFAULT_VERTEX_API_URL:
# Operator pinned an explicit upstream (gateway / specific host) — honor it.
return configured
if not location or location == "global":
return "https://aiplatform.googleapis.com"
return f"https://{location}-aiplatform.googleapis.com"
def _select_passthrough_base_url(proxy: Any, headers: dict[str, str]) -> str: def _select_passthrough_base_url(proxy: Any, headers: dict[str, str]) -> str:
# Codex CLI subscription mode hits a wide surface under # Codex CLI subscription mode hits a wide surface under
# `/backend-api/*` (rate-limit polling, agent identity, JWT # `/backend-api/*` (rate-limit polling, agent identity, JWT
@ -644,11 +665,11 @@ def register_provider_routes(app: FastAPI, proxy: Any) -> None:
publisher: str, publisher: str,
model: str, model: str,
): ):
del api_version, project, location del api_version, project
if publisher == "anthropic": if publisher == "anthropic":
return await proxy.handle_anthropic_messages( return await proxy.handle_anthropic_messages(
request, request,
_api_target(proxy, "vertex"), _vertex_target_for_location(proxy, location),
"vertex:anthropic", "vertex:anthropic",
model, model,
) )
@ -665,11 +686,11 @@ def register_provider_routes(app: FastAPI, proxy: Any) -> None:
publisher: str, publisher: str,
model: str, model: str,
): ):
del api_version, project, location del api_version, project
if publisher == "anthropic": if publisher == "anthropic":
return await proxy.handle_anthropic_messages( return await proxy.handle_anthropic_messages(
request, request,
_api_target(proxy, "vertex"), _vertex_target_for_location(proxy, location),
"vertex:anthropic", "vertex:anthropic",
model, model,
True, True,

View file

@ -173,6 +173,12 @@ def create_proxy_backend(
normalized_backend = backend if backend.startswith("litellm-") else f"litellm-{backend}" normalized_backend = backend if backend.startswith("litellm-") else f"litellm-{backend}"
provider = normalized_backend.replace("litellm-", "") provider = normalized_backend.replace("litellm-", "")
# `litellm-vertex` is the name in our docs/help, but LiteLLM (and our
# provider registry) keys Google Vertex on `vertex_ai`. Without this alias
# the provider falls through to a generic pass-through: wrong model prefix
# (`vertex/…` instead of `vertex_ai/…`), region dropped, auth mishandled.
if provider in ("vertex", "google-vertex", "googlevertex"):
provider = "vertex_ai"
try: try:
backend_cls = litellm_backend_cls or _load_litellm_backend() backend_cls = litellm_backend_cls or _load_litellm_backend()
instance = cast("Backend", backend_cls(provider=provider, region=bedrock_region)) instance = cast("Backend", backend_cls(provider=provider, region=bedrock_region))

View file

@ -0,0 +1,140 @@
"""Turnkey Claude Code + Vertex compression wiring.
Covers the fixes that let `headroom wrap claude` + Vertex actually deliver
compression:
- `litellm-vertex` -> `vertex_ai` provider alias (registry),
- the Vertex upstream host derived per-request from the path's `location`
(so a europe-west1 request is not sent to a us-central1 host),
- the native `:rawPredict` route running the compression handler (not the
verbatim passthrough) for the `anthropic` publisher.
No real GCP/Vertex is contacted handlers and the backend class are stubbed.
"""
from __future__ import annotations
import logging
import types
from typing import Any
from fastapi.responses import JSONResponse
from fastapi.testclient import TestClient
from headroom.providers import proxy_routes, registry
from headroom.providers.registry import DEFAULT_VERTEX_API_URL
from headroom.proxy.server import HeadroomProxy, ProxyConfig, create_app
# --------------------------------------------------------------------------
# Region-aware Vertex upstream host
# --------------------------------------------------------------------------
def _stub_proxy(vertex_url: str) -> Any:
# `_api_target` reads `proxy.VERTEX_API_URL`, but its getattr default eagerly
# evaluates `proxy.provider_runtime.api_target(...)`, so the stub needs both.
return types.SimpleNamespace(
VERTEX_API_URL=vertex_url,
provider_runtime=types.SimpleNamespace(api_target=lambda name: vertex_url),
)
def test_vertex_target_derives_region_from_location_on_default() -> None:
proxy = _stub_proxy(DEFAULT_VERTEX_API_URL)
assert (
proxy_routes._vertex_target_for_location(proxy, "europe-west1")
== "https://europe-west1-aiplatform.googleapis.com"
)
def test_vertex_target_global_uses_unprefixed_host() -> None:
proxy = _stub_proxy(DEFAULT_VERTEX_API_URL)
assert (
proxy_routes._vertex_target_for_location(proxy, "global")
== "https://aiplatform.googleapis.com"
)
def test_vertex_target_empty_location_uses_unprefixed_host() -> None:
proxy = _stub_proxy(DEFAULT_VERTEX_API_URL)
assert (
proxy_routes._vertex_target_for_location(proxy, "") == "https://aiplatform.googleapis.com"
)
def test_vertex_target_honors_explicit_override() -> None:
# An operator who pinned a non-default upstream (private gateway) wins,
# regardless of the path's location.
proxy = _stub_proxy("https://vertex-gateway.internal")
assert (
proxy_routes._vertex_target_for_location(proxy, "europe-west1")
== "https://vertex-gateway.internal"
)
# --------------------------------------------------------------------------
# litellm-vertex -> vertex_ai provider alias
# --------------------------------------------------------------------------
def _capture_provider(backend: str) -> dict[str, Any]:
captured: dict[str, Any] = {}
class FakeLiteLLM:
def __init__(self, provider: str, region: str | None = None) -> None:
captured["provider"] = provider
captured["region"] = region
registry.create_proxy_backend(
backend=backend,
anyllm_provider="openai",
bedrock_region="us-east5",
logger=logging.getLogger("test-vertex"),
litellm_backend_cls=FakeLiteLLM,
)
return captured
def test_litellm_vertex_aliases_to_vertex_ai() -> None:
# The documented `litellm-vertex` must reach the `vertex_ai` provider, not
# a generic pass-through (which would drop the region and mangle the model).
assert _capture_provider("litellm-vertex")["provider"] == "vertex_ai"
def test_litellm_vertex_ai_unchanged() -> None:
assert _capture_provider("litellm-vertex_ai")["provider"] == "vertex_ai"
def test_litellm_bedrock_not_aliased() -> None:
assert _capture_provider("litellm-bedrock")["provider"] == "bedrock"
# --------------------------------------------------------------------------
# Native :rawPredict route → compression handler with region-derived host
# --------------------------------------------------------------------------
def _default_vertex_app() -> Any:
# No vertex_api_url override -> default us-central1 host -> route derives
# the regional host from the request path.
return create_app(ProxyConfig(optimize=True, cache_enabled=False, rate_limit_enabled=False))
def test_vertex_rawpredict_anthropic_runs_compression_handler(monkeypatch) -> None:
captured: dict[str, str] = {}
# The route calls handle_anthropic_messages(request, base_url, provider, model).
async def fake(self, request, base_url, provider, model, *rest): # type: ignore[no-untyped-def]
captured.update(base_url=str(base_url), provider=str(provider), model=str(model))
return JSONResponse({"ok": True})
monkeypatch.setattr(HeadroomProxy, "handle_anthropic_messages", fake)
with TestClient(_default_vertex_app()) as client:
resp = client.post(
"/v1/projects/p/locations/europe-west1/publishers/anthropic/models/"
"claude-sonnet-4-6:rawPredict",
json={"anthropic_version": "vertex-2023-10-16", "messages": []},
)
assert resp.status_code == 200
# The anthropic publisher is routed to the compression handler (not the
# verbatim passthrough), with the region-derived upstream host (exact match,
# not a substring check).
assert captured["provider"] == "vertex:anthropic"
assert captured["base_url"] == "https://europe-west1-aiplatform.googleapis.com"
assert captured["model"] == "claude-sonnet-4-6"