diff --git a/src/benchflow/providers/litellm_config.py b/src/benchflow/providers/litellm_config.py index cfa94a73d..b28140af9 100644 --- a/src/benchflow/providers/litellm_config.py +++ b/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]) diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index 60c9953c3..f0977a253 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/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,15 @@ 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.pop("GEMINI_API_KEY_AUTH_MECHANISM", None) + 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 diff --git a/tests/test_litellm_config.py b/tests/test_litellm_config.py index 188c1a2e6..7dbbb6072 100644 --- a/tests/test_litellm_config.py +++ b/tests/test_litellm_config.py @@ -9,6 +9,85 @@ ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model", ["gemini-3.1-pro-preview", "gemini-3.8-flash", "gemini-3.5-flash-lite"] +) +async def test_gemini_vertex_passthrough_exchanges_gateway_auth_for_adc( + monkeypatch, model +): + """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( + f"google-vertex/{model}", + {"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=f"v1beta1/publishers/google/models/{model}: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/" + f"publishers/google/models/{model}: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", diff --git a/tests/test_litellm_runtime.py b/tests/test_litellm_runtime.py index 80d7d1b94..ff82e329b 100644 --- a/tests/test_litellm_runtime.py +++ b/tests/test_litellm_runtime.py @@ -758,8 +758,16 @@ 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)] +) +@pytest.mark.parametrize( + "model", ["gemini-3.1-pro-preview", "gemini-3.8-flash", "gemini-3.5-flash-lite"] +) +async def test_gemini_uses_native_generate_content_through_sandbox_proxy( + monkeypatch, vertex, upstream_keys, model +): + """Guards PR #942/#1030 and Vertex auth regression in commit 28b82e33.""" starts = [] @@ -774,14 +782,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 vertex else "bearer", + } + 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=f"google-vertex/{model}" if vertex else model, runtime=None, environment="docker", usage_tracking="required", @@ -791,7 +810,14 @@ 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 + assert "GEMINI_API_KEY_AUTH_MECHANISM" not in updated for key in ( "GEMINI_API_KEY", "GOOGLE_API_KEY",