feat(proxy): extend X-Headroom-Upstream override to Gemini + /v1beta normalization

Reviewing headroomlabs-ai#1089 (the launcher-side opencode wrap) surfaced two gaps in
the initial override implementation:

1. Gemini's version segment is /v1beta (not /v1). The proxy's gemini
   handlers append /v1beta/models/... themselves, so a caller passing the
   versioned URL (e.g. matching headroom's _KNOWN_UPSTREAMS
   'generativelanguage.googleapis.com/v1beta') would have produced a
   doubled /v1beta/v1beta path. The resolver now strips a trailing /v1beta
   as well as /v1.

2. The three /v1beta gemini routes did not thread the override through.
   handle_gemini_generate_content and handle_gemini_count_tokens already
   accepted upstream_base_url; handle_gemini_stream_generate_content did
   not. All three routes now pass request_upstream_override(request), and
   the stream handler gained the parameter.

Coverage is now OpenAI (/v1/chat/completions, /v1/responses), Anthropic
(/v1/messages), Gemini (/v1beta generateContent/streamGenerateContent/
countTokens), and every passthrough / catch-all route.

Author: JavaGT

CHANGELOG.md

@@ -10,7 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Features
 
-* **proxy:** per-request upstream override via the `X-Headroom-Upstream` header. A single proxy instance can now fan out to many upstreams (one per provider) instead of one proxy per upstream — the caller tags each request with its real upstream base and the proxy forwards there, overriding the startup default for that provider. The value is normalized like `*_TARGET_API_URL` (trailing slash and trailing `/v1` stripped) and the proxy appends the incoming request path. Honored on `/v1/chat/completions`, `/v1/responses`, `/v1/messages`, and every passthrough / catch-all route; stripped before the upstream call so it never leaks. Enables single-proxy multi-provider setups such as OpenCode's 75+ providers (each provider configured with `baseURL` = the proxy + an `X-Headroom-Upstream` header = its real upstream).
+* **proxy:** per-request upstream override via the `X-Headroom-Upstream` header. A single proxy instance can now fan out to many upstreams (one per provider) instead of one proxy per upstream — the caller tags each request with its real upstream base and the proxy forwards there, overriding the startup default for that provider. The value is normalized like `*_TARGET_API_URL` (trailing slash and a trailing `/v1` or `/v1beta` version segment stripped) and the proxy appends the incoming request path. Honored on `/v1/chat/completions`, `/v1/responses`, `/v1/messages`, `/v1beta/models/{model}:generateContent`/`:streamGenerateContent`/`:countTokens`, and every passthrough / catch-all route; stripped before the upstream call so it never leaks. Enables single-proxy multi-provider setups such as OpenCode's 75+ providers (each provider configured with `baseURL` = the proxy + an `X-Headroom-Upstream` header = its real upstream).
 
 * **proxy:** measure and surface rolling and current token throughput metrics (active/wall-clock input, compression, effective forward, and streamed generation) in `headroom perf` CLI and the dashboard ([#959](https://github.com/chopratejas/headroom/issues/959)).
 * **vibe:** add Mistral Vibe CLI support with `headroom wrap vibe`.

docs/content/docs/proxy.mdx

@@ -266,16 +266,17 @@ Rewriting the request body invalidates the caller's **SigV4** signature (it cove
 
 A single Headroom proxy normally forwards to one configured upstream per provider (`OPENAI_TARGET_API_URL`, `ANTHROPIC_TARGET_API_URL`, …). The `X-Headroom-Upstream` request header overrides that upstream **per request**, so one proxy instance can fan out to many upstreams — no need to run one proxy per provider.
 
-Set the header to the upstream base URL. It is normalized the same way as the `*_TARGET_API_URL` env vars (trailing slash and a trailing `/v1` segment are stripped), then the proxy appends the incoming request path:
+Set the header to the upstream base URL. It is normalized the same way as the `*_TARGET_API_URL` env vars (trailing slash and a trailing API-version segment — `/v1` or `/v1beta` — are stripped), then the proxy appends the incoming request path:
 
 | Header value | Request path | Forwarded to |
 | --- | --- | --- |
 | `https://api.deepseek.com` | `/v1/chat/completions` | `https://api.deepseek.com/v1/chat/completions` |
 | `https://api.deepseek.com/v1` | `/v1/chat/completions` | `https://api.deepseek.com/v1/chat/completions` |
 | `https://api.groq.com/openai/v1` | `/v1/chat/completions` | `https://api.groq.com/openai/v1/chat/completions` |
 | `https://api.anthropic.com` | `/v1/messages` | `https://api.anthropic.com/v1/messages` |
+| `https://generativelanguage.googleapis.com/v1beta` | `/v1beta/models/gemini-2.0-flash:generateContent` | `https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent` |
 
-The header is an internal `x-headroom-*` control flag: it is consumed by the proxy and stripped before the upstream call, so it never leaks to the provider. Honored on the OpenAI (`/v1/chat/completions`, `/v1/responses`), Anthropic (`/v1/messages`), and every passthrough / catch-all route.
+The header is an internal `x-headroom-*` control flag: it is consumed by the proxy and stripped before the upstream call, so it never leaks to the provider. Honored on the OpenAI (`/v1/chat/completions`, `/v1/responses`), Anthropic (`/v1/messages`), Gemini (`/v1beta/models/{model}:generateContent`, `:streamGenerateContent`, `:countTokens`), and every passthrough / catch-all route.
 
 ### Use case: many providers through one proxy
 

headroom/providers/proxy_routes.py

@@ -576,15 +576,21 @@ async def cancel_batch(request: Request, batch_id: str):
 
     @app.post("/v1beta/models/{model}:generateContent")
     async def gemini_generate_content(request: Request, model: str):
-        return await proxy.handle_gemini_generate_content(request, model)
+        return await proxy.handle_gemini_generate_content(
+            request, model, upstream_base_url=request_upstream_override(request)
+        )
 
     @app.post("/v1beta/models/{model}:streamGenerateContent")
     async def gemini_stream_generate_content(request: Request, model: str):
-        return await proxy.handle_gemini_stream_generate_content(request, model)
+        return await proxy.handle_gemini_stream_generate_content(
+            request, model, upstream_base_url=request_upstream_override(request)
+        )
 
     @app.post("/v1beta/models/{model}:countTokens")
     async def gemini_count_tokens(request: Request, model: str):
-        return await proxy.handle_gemini_count_tokens(request, model)
+        return await proxy.handle_gemini_count_tokens(
+            request, model, upstream_base_url=request_upstream_override(request)
+        )
 
     @app.post("/v1internal:streamGenerateContent")
     async def google_cloudcode_stream_generate_content(request: Request):

headroom/proxy/handlers/gemini.py

@@ -901,11 +901,12 @@ async def handle_gemini_stream_generate_content(
         self,
         request: Request,
         model: str,
+        upstream_base_url: str | None = None,
     ) -> StreamingResponse | JSONResponse:
         """Handle Gemini streaming endpoint /v1beta/models/{model}:streamGenerateContent."""
         from fastapi.responses import JSONResponse
 
-        from headroom.proxy.helpers import _read_request_json
+        from headroom.proxy.helpers import _read_request_json, request_upstream_override
         from headroom.tokenizers import get_tokenizer
 
         start_time = time.time()
@@ -957,9 +958,10 @@ async def handle_gemini_stream_generate_content(
 
         # Build URL with SSE param
         query_params = dict(request.query_params)
-        url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:streamGenerateContent?alt=sse"
+        _gemini_base = upstream_base_url or request_upstream_override(request) or self.GEMINI_API_URL
+        url = f"{_gemini_base}/v1beta/models/{model}:streamGenerateContent?alt=sse"
         if "key" in query_params:
-            url = f"{self.GEMINI_API_URL}/v1beta/models/{model}:streamGenerateContent?key={query_params['key']}&alt=sse"
+            url = f"{_gemini_base}/v1beta/models/{model}:streamGenerateContent?key={query_params['key']}&alt=sse"
 
         return await self._stream_response(
             url,

headroom/proxy/helpers.py

@@ -1574,10 +1574,12 @@ def request_upstream_override(request: Request) -> str | None:
     providers) and the proxy forwards there instead of its startup default.
 
     The value is normalized to match the proxy's internal ``*_API_URL``
-    format (trailing slash and a trailing ``/v1`` segment are stripped), so
-    both ``https://api.deepseek.com`` and ``https://api.deepseek.com/v1``
-    resolve to ``https://api.deepseek.com``. The proxy then appends the
-    incoming request path (e.g. ``/v1/chat/completions``).
+    format (trailing slash and a trailing API-version segment are stripped),
+    so all of ``https://api.deepseek.com``, ``https://api.deepseek.com/v1``,
+    and ``https://generativelanguage.googleapis.com/v1beta`` resolve to their
+    bare host. The proxy then appends the incoming request path —
+    ``/v1/chat/completions`` for OpenAI/Anthropic, ``/v1beta/models/...`` for
+    Gemini — so passing the version in the header does not double it up.
 
     Returns ``None`` when the header is unset or empty. The header is an
     ``x-headroom-*`` control flag and is stripped from upstream-bound
@@ -1589,8 +1591,13 @@ def request_upstream_override(request: Request) -> str | None:
     normalized = raw.strip().rstrip("/")
     if not normalized:
         return None
-    if normalized.endswith("/v1"):
-        normalized = normalized[:-3]
+    # Strip a trailing API-version segment the proxy's handlers add
+    # themselves (/v1 for OpenAI/Anthropic/Vertex, /v1beta for Gemini), so
+    # callers may pass either the bare host or the versioned URL.
+    for suffix in ("/v1", "/v1beta"):
+        if normalized.endswith(suffix):
+            normalized = normalized[: -len(suffix)]
+            break
     return normalized
 
 

tests/test_proxy_upstream_override.py

@@ -94,6 +94,28 @@ def test_override_strips_trailing_v1_with_slash() -> None:
     )
 
 
+def test_override_strips_trailing_v1beta() -> None:
+    # Gemini's version segment is /v1beta (not /v1). The proxy's gemini
+    # handlers append /v1beta/models/... themselves, so a caller passing the
+    # versioned URL (e.g. matching headroom's _KNOWN_UPSTREAMS) must have it
+    # stripped to avoid a doubled /v1beta/v1beta path.
+    assert (
+        request_upstream_override(
+            _stub_request({"x-headroom-upstream": "https://generativelanguage.googleapis.com/v1beta"})
+        )
+        == "https://generativelanguage.googleapis.com"
+    )
+
+
+def test_override_strips_trailing_v1beta_with_slash() -> None:
+    assert (
+        request_upstream_override(
+            _stub_request({"x-headroom-upstream": "https://generativelanguage.googleapis.com/v1beta/"})
+        )
+        == "https://generativelanguage.googleapis.com"
+    )
+
+
 def test_override_preserves_path_prefix_before_v1() -> None:
     # OpenRouter / Groq style: the /v1 is the API version, the prefix is real.
     assert (
@@ -256,6 +278,73 @@ def test_catchall_forwards_to_override_upstream() -> None:
     assert url == "https://api.deepseek.com/some/custom/path"
 
 
+# ── /v1beta gemini routes thread the override through upstream_base_url ─
+
+
+def test_v1beta_generate_content_threads_override(monkeypatch) -> None:
+    captured: list[Any] = []
+
+    async def fake_gemini_generate(
+        self, request, model, upstream_base_url=None, provider_name="gemini"
+    ):  # type: ignore[no-untyped-def]
+        captured.append(upstream_base_url)
+        return JSONResponse({"upstream_base_url": upstream_base_url, "model": model})
+
+    monkeypatch.setattr(HeadroomProxy, "handle_gemini_generate_content", fake_gemini_generate)
+
+    with TestClient(_app()) as client:
+        assert client.post(
+            "/v1beta/models/gemini-mock:generateContent",
+            json={"contents": [{"parts": [{"text": "hi"}]}]},
+            headers={"X-Headroom-Upstream": "https://generativelanguage.googleapis.com/v1beta"},
+        ).json()["upstream_base_url"] == "https://generativelanguage.googleapis.com"
+
+    assert captured == ["https://generativelanguage.googleapis.com"]
+
+
+def test_v1beta_count_tokens_threads_override(monkeypatch) -> None:
+    captured: list[Any] = []
+
+    async def fake_gemini_count(
+        self, request, model, upstream_base_url=None, provider_name="gemini"
+    ):  # type: ignore[no-untyped-def]
+        captured.append(upstream_base_url)
+        return JSONResponse({"upstream_base_url": upstream_base_url})
+
+    monkeypatch.setattr(HeadroomProxy, "handle_gemini_count_tokens", fake_gemini_count)
+
+    with TestClient(_app()) as client:
+        assert client.post(
+            "/v1beta/models/gemini-mock:countTokens",
+            json={"contents": [{"parts": [{"text": "hi"}]}]},
+            headers={"X-Headroom-Upstream": "https://generativelanguage.googleapis.com"},
+        ).json()["upstream_base_url"] == "https://generativelanguage.googleapis.com"
+
+    assert captured == ["https://generativelanguage.googleapis.com"]
+
+
+def test_v1beta_stream_generate_content_threads_override(monkeypatch) -> None:
+    """handle_gemini_stream_generate_content now accepts upstream_base_url."""
+    captured: list[Any] = []
+
+    async def fake_gemini_stream(
+        self, request, model, upstream_base_url=None
+    ):  # type: ignore[no-untyped-def]
+        captured.append(upstream_base_url)
+        return JSONResponse({"upstream_base_url": upstream_base_url})
+
+    monkeypatch.setattr(HeadroomProxy, "handle_gemini_stream_generate_content", fake_gemini_stream)
+
+    with TestClient(_app()) as client:
+        assert client.post(
+            "/v1beta/models/gemini-mock:streamGenerateContent",
+            json={"contents": [{"parts": [{"text": "hi"}]}]},
+            headers={"X-Headroom-Upstream": "https://generativelanguage.googleapis.com/v1beta"},
+        ).json()["upstream_base_url"] == "https://generativelanguage.googleapis.com"
+
+    assert captured == ["https://generativelanguage.googleapis.com"]
+
+
 def test_override_header_stripped_before_upstream_call() -> None:
     """The x-headroom-upstream control flag must not leak to the upstream."""
     response = httpx.Response(