Fix Gemini CLI Vertex routing and authentication

src/benchflow/providers/litellm_config.py

@@ -434,6 +434,10 @@ def litellm_proxy_config(
 ) -> dict[str, object]:
     """Build the LiteLLM ``config.yaml`` payload for one route."""
     params = dict(route.litellm_params)
+    if route.provider_name == "google-vertex":
+        # Gemini gateway keys omit project/location from request paths. Native
+        # pass-through registration resolves both and uses the proxy's ADC.
+        params["use_in_pass_through"] = True
     cost = custom_cost_per_token(route.upstream_model)
     if cost is not None:
         params.setdefault("input_cost_per_token", cost[0])

src/benchflow/providers/litellm_runtime.py

@@ -1399,6 +1399,7 @@ def _provider_models_for_proxy_alias(
         "ANTHROPIC_BEDROCK_BASE_URL",
         "GEMINI_BASE_URL",
         "GOOGLE_GEMINI_BASE_URL",
+        "GOOGLE_VERTEX_BASE_URL",
     }
 )
 
@@ -1538,10 +1539,14 @@ def _wire_litellm_agent_env(
         # Gemini CLI speaks Google's native GenerateContent protocol. Use
         # LiteLLM's byte-preserving Gemini pass-through route: its translated
         # GenerateContent route can corrupt streamed, multi-tool responses.
-        # The gateway authenticates the reviewer with ``master_key`` and swaps
-        # in the upstream Gemini key server-side.
+        # Vertex keeps its CLI auth mode and uses the ADC-backed pass-through;
+        # its gateway route accepts bearer auth rather than x-goog-api-key.
         updated.pop(LITELLM_MODEL_ALIAS_ENV, None)
-        updated["GOOGLE_GEMINI_BASE_URL"] = f"{base_url.rstrip('/')}/gemini"
+        if route.provider_name == "google-vertex":
+            updated["GOOGLE_VERTEX_BASE_URL"] = f"{base_url.rstrip('/')}/vertex_ai"
+            updated["GEMINI_API_KEY_AUTH_MECHANISM"] = "bearer"
+        else:
+            updated["GOOGLE_GEMINI_BASE_URL"] = f"{base_url.rstrip('/')}/gemini"
         # Gemini CLI recognizes several equivalent credential names, with the
         # selected alias varying by model family and CLI release. Point every
         # accepted alias at the gateway so Gemma cannot inherit a real Google

tests/test_litellm_config.py

@@ -9,6 +9,80 @@
 )
 
 
+@pytest.mark.asyncio
+async def test_gemini_vertex_passthrough_exchanges_gateway_auth_for_adc(monkeypatch):
+    """Guards Vertex proxy auth against the regression in commit 28b82e33."""
+    from unittest.mock import AsyncMock, Mock
+
+    import litellm
+    from litellm.proxy import proxy_server
+    from litellm.proxy.pass_through_endpoints import llm_passthrough_endpoints as vendor
+    from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import (
+        PassthroughEndpointRouter,
+    )
+    from starlette.requests import Request
+    from starlette.responses import Response
+
+    route = resolve_litellm_route(
+        "google-vertex/gemini-3.1-flash-lite",
+        {"GOOGLE_CLOUD_PROJECT": "dummy-project", "GOOGLE_CLOUD_LOCATION": "global"},
+    )
+    config = litellm_proxy_config(route, master_key="gateway-key")
+    monkeypatch.setattr(
+        vendor, "passthrough_endpoint_router", PassthroughEndpointRouter()
+    )
+    model_router = litellm.Router(model_list=config["model_list"])
+    monkeypatch.setattr(proxy_server, "llm_router", model_router)
+    # Native registration must leave other harnesses' translated routes usable.
+    deployment = model_router.get_available_deployment(model=route.model_alias)
+    assert deployment["litellm_params"]["model"] == route.upstream_model
+
+    auth = AsyncMock()
+    adc = AsyncMock(return_value=("mock-oauth-token", "dummy-project"))
+    forward = Mock(return_value=AsyncMock())
+    monkeypatch.setattr(vendor, "user_api_key_auth", auth)
+    monkeypatch.setattr(vendor.VertexBase, "_ensure_access_token_async", adc)
+    monkeypatch.setattr(vendor, "create_pass_through_route", forward)
+    request = Request(
+        {
+            "type": "http",
+            "headers": [
+                (b"authorization", b"Bearer gateway-key"),
+                (b"x-goog-api-key", b"gateway-key"),
+            ],
+        }
+    )
+    # Actual CLI 0.42.0 Express-style path, captured using a dummy gateway key.
+    await vendor._base_vertex_proxy_route(
+        endpoint="v1beta1/publishers/google/models/gemini-3.1-flash-lite:streamGenerateContent",
+        request=request,
+        fastapi_response=Response(),
+        get_vertex_pass_through_handler=vendor.get_vertex_pass_through_handler(
+            call_type="aiplatform"
+        ),
+    )
+    auth.assert_awaited_once_with(request=request, api_key="Bearer gateway-key")
+    adc.assert_awaited_once_with(
+        credentials=None,
+        project_id="dummy-project",
+        custom_llm_provider="vertex_ai_beta",
+    )
+    sent = forward.call_args.kwargs
+    assert sent["target"] == (
+        "https://aiplatform.googleapis.com/v1beta1/projects/dummy-project/locations/global/"
+        "publishers/google/models/gemini-3.1-flash-lite:streamGenerateContent?alt=sse"
+    )
+    assert sent["is_streaming_request"] is True
+    # Use the forwarding layer's real merge rule, including incoming key headers.
+    headers = vendor.HttpPassThroughEndpointHelpers.forward_headers_from_request(
+        dict(request.headers),
+        sent["custom_headers"],
+        sent.get("_forward_headers", False),
+    )
+    assert headers["Authorization"] == "Bearer mock-oauth-token"
+    assert "gateway-key" not in str(headers)
+
+
 def test_bedrock_model_maps_to_litellm_bedrock_route():
     route = resolve_litellm_route(
         "aws-bedrock/us.anthropic.claude-opus-4-8",

tests/test_litellm_runtime.py

@@ -758,8 +758,13 @@ async def fail_start(**_kwargs):
 
 
 @pytest.mark.asyncio
-async def test_gemini_uses_native_generate_content_through_sandbox_proxy(monkeypatch):
-    """Guards PR #942 and PR #1030: every Google key alias stays proxied."""
+@pytest.mark.parametrize(
+    "vertex,upstream_keys", [(False, True), (True, False), (True, True)]
+)
+async def test_gemini_uses_native_generate_content_through_sandbox_proxy(
+    monkeypatch, vertex, upstream_keys
+):
+    """Guards PR #942/#1030 and Vertex auth regression in commit 28b82e33."""
 
     starts = []
 
@@ -774,14 +779,25 @@ async def unexpected_host_start(**_kwargs):
     monkeypatch.setattr(runtime_mod, "_start_host_litellm", unexpected_host_start)
     sandbox = SimpleNamespace()
 
+    agent_env = {
+        "GOOGLE_CLOUD_PROJECT": "dummy-project",
+        "GOOGLE_CLOUD_LOCATION": "global",
+        "GOOGLE_GEMINI_BASE_URL": "https://stale-gemini.example.test",
+        "GOOGLE_VERTEX_BASE_URL": "https://stale-vertex.example.test",
+        "GEMINI_API_KEY_AUTH_MECHANISM": "x-goog-api-key",
+    }
+    if upstream_keys or not vertex:
+        agent_env.update(
+            {
+                "GEMINI_API_KEY": "upstream-gemini-key",
+                "GOOGLE_API_KEY": "upstream-google-key",
+                "GOOGLE_GENERATIVE_AI_API_KEY": "upstream-generative-ai-key",
+            }
+        )
     updated, provider_runtime = await ensure_litellm_runtime(
         agent="gemini",
-        agent_env={
-            "GEMINI_API_KEY": "upstream-gemini-key",
-            "GOOGLE_API_KEY": "upstream-google-key",
-            "GOOGLE_GENERATIVE_AI_API_KEY": "upstream-generative-ai-key",
-        },
-        model="gemini-2.5-flash",
+        agent_env=agent_env,
+        model="google-vertex/gemini-3.1-flash-lite" if vertex else "gemini-2.5-flash",
         runtime=None,
         environment="docker",
         usage_tracking="required",
@@ -791,7 +807,13 @@ async def unexpected_host_start(**_kwargs):
 
     assert starts[0]["sandbox"] is sandbox
     assert provider_runtime is not None
-    assert updated["GOOGLE_GEMINI_BASE_URL"] == "http://127.0.0.1:45678/gemini"
+    if vertex:
+        assert updated["GOOGLE_VERTEX_BASE_URL"] == "http://127.0.0.1:45678/vertex_ai"
+        assert updated["GEMINI_API_KEY_AUTH_MECHANISM"] == "bearer"
+        assert "GOOGLE_GEMINI_BASE_URL" not in updated
+    else:
+        assert updated["GOOGLE_GEMINI_BASE_URL"] == "http://127.0.0.1:45678/gemini"
+        assert "GOOGLE_VERTEX_BASE_URL" not in updated
     for key in (
         "GEMINI_API_KEY",
         "GOOGLE_API_KEY",