diff --git a/src/benchflow/acp/runtime.py b/src/benchflow/acp/runtime.py index a4cd802a8..41d9f584b 100644 --- a/src/benchflow/acp/runtime.py +++ b/src/benchflow/acp/runtime.py @@ -34,7 +34,11 @@ find_provider_for_bare_model, strip_provider_prefix, ) -from benchflow.agents.registry import AGENTS, OPENCODE_PROXY_PROVIDER_ID +from benchflow.agents.registry import ( + ACPX_KEY_PREFIX, + AGENTS, + OPENCODE_PROXY_PROVIDER_ID, +) from benchflow.diagnostics import ( AgentPromptTimeoutDiagnostic, AgentPromptTimeoutError, @@ -268,6 +272,13 @@ def _format_acp_model(model: str, agent: str) -> str: models.dev provider prefix when the agent requires it. """ bare = strip_provider_prefix(model) + # Gemini CLI accepts bare Google model IDs. ``google/`` is a models.dev + # provider prefix rather than a registered BenchFlow provider, so the + # generic normalizer intentionally leaves it alone. + if agent.removeprefix(ACPX_KEY_PREFIX) == "gemini" and model.startswith( + ("google/gemini-", "google/gemma-") + ): + bare = model.removeprefix("google/") agent_cfg = AGENTS.get(agent) if agent_cfg and agent_cfg.acp_model_format == "registered-provider/model": # Proxy mode: BenchFlow's LiteLLM proxy serves the model under the alias diff --git a/src/benchflow/agents/registry.py b/src/benchflow/agents/registry.py index 76c7ed378..127ef423f 100644 --- a/src/benchflow/agents/registry.py +++ b/src/benchflow/agents/registry.py @@ -656,7 +656,11 @@ class AgentConfig: description="Google Gemini CLI via ACP", skill_paths=["$HOME/.gemini/skills"], install_cmd=_js_agent_install("gemini", "@google/gemini-cli@0.42.0"), - launch_cmd=_js_agent_launch("gemini", "--acp --yolo"), + # BenchFlow already isolates the agent inside the task sandbox and + # deliberately grants tool approval with --yolo. Gemini CLI 0.42 also + # requires the workspace to be trusted before it will honor that mode; + # headless ACP runs cannot answer the interactive trust prompt. + launch_cmd=_js_agent_launch("gemini", "--acp --yolo --skip-trust"), protocol="acp", # The Gemini CLI reads GEMINI_API_KEY natively. GOOGLE_API_KEY is # accepted as an alias: auto_inherit_env mirrors it both ways so users diff --git a/src/benchflow/providers/litellm_config.py b/src/benchflow/providers/litellm_config.py index acee4756c..f703b32d2 100644 --- a/src/benchflow/providers/litellm_config.py +++ b/src/benchflow/providers/litellm_config.py @@ -375,6 +375,9 @@ def resolve_litellm_route(model: str, env: dict[str, str]) -> LiteLLMRoute: elif lower.startswith("gemini/"): upstream = model required = ("GEMINI_API_KEY",) + elif lower.startswith(("google/gemini-", "google/gemma-")): + upstream = f"gemini/{model.split('/', 1)[1]}" + required = ("GEMINI_API_KEY",) elif "gemini" in lower: upstream = f"gemini/{bare}" required = ("GEMINI_API_KEY",) diff --git a/src/benchflow/providers/litellm_gemini_passthrough_patch.py b/src/benchflow/providers/litellm_gemini_passthrough_patch.py new file mode 100644 index 000000000..75a0f5e37 --- /dev/null +++ b/src/benchflow/providers/litellm_gemini_passthrough_patch.py @@ -0,0 +1,116 @@ +"""Backport Gemini pass-through streaming logging for LiteLLM 1.89/1.91. + +Those releases classify every native ``generateContent`` URL as Vertex AI and +send its streamed response through the Vertex logger. That happens to work for +Gemini model names present in both catalogs, but it rejects Google AI Studio +Gemma names and drops the entire callback — including token usage and provider +trajectory evidence. LiteLLM added a dedicated Gemini endpoint type later. + +The proxy imports this module from ``sitecustomize``. It patches only the old +vendor shape (no ``EndpointType.GEMINI``) and only Google AI Studio-style model +paths, leaving Vertex and every other provider untouched. +""" + +from __future__ import annotations + +import re +from typing import Any +from urllib.parse import urlparse + +_GOOGLE_AI_STUDIO_GENERATE_PATH_RE = re.compile( + r"/v1(?:beta)?/models/[^/:]+:(?:streamGenerateContent|generateContent)$" +) + + +def _is_google_ai_studio_generate_content(url_route: str) -> bool: + """Distinguish native AI Studio model paths from Vertex project paths.""" + try: + path = urlparse(url_route).path + except Exception: + return False + return bool( + _GOOGLE_AI_STUDIO_GENERATE_PATH_RE.search(path) + and "/projects/" not in path + and "/locations/" not in path + ) + + +def _apply_patch() -> None: + try: + from litellm.proxy.pass_through_endpoints.llm_provider_handlers.gemini_passthrough_logging_handler import ( + GeminiPassthroughLoggingHandler, + ) + from litellm.proxy.pass_through_endpoints.streaming_handler import ( + PassThroughStreamingHandler, + ) + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + EndpointType, + ) + except Exception: + return + + # Newer LiteLLM releases have their own Gemini dispatch branch. Do not + # replace supported vendor behavior after the pinned dependency advances. + if hasattr(EndpointType, "GEMINI"): + return + + original = PassThroughStreamingHandler._build_passthrough_logging_result + if getattr(original, "__benchflow_gemini_passthrough_patch__", False): + return + + def build_passthrough_logging_result( + *, + litellm_logging_obj: Any, + passthrough_success_handler_obj: Any, + url_route: str, + request_body: dict[str, Any], + endpoint_type: Any, + start_time: Any, + raw_bytes: list[bytes], + end_time: Any, + model: str | None, + ) -> tuple[Any, dict[str, Any]]: + if not _is_google_ai_studio_generate_content(url_route): + return original( + litellm_logging_obj=litellm_logging_obj, + passthrough_success_handler_obj=passthrough_success_handler_obj, + url_route=url_route, + request_body=request_body, + endpoint_type=endpoint_type, + start_time=start_time, + raw_bytes=raw_bytes, + end_time=end_time, + model=model, + ) + + all_chunks = PassThroughStreamingHandler._convert_raw_bytes_to_str_lines( + raw_bytes + ) + result = ( + GeminiPassthroughLoggingHandler._handle_logging_gemini_collected_chunks( + litellm_logging_obj=litellm_logging_obj, + passthrough_success_handler_obj=passthrough_success_handler_obj, + url_route=url_route, + request_body=request_body, + endpoint_type=endpoint_type, + start_time=start_time, + all_chunks=all_chunks, + end_time=end_time, + model=model, + ) + ) + return result["result"], result["kwargs"] + + setattr( # noqa: B010 - marker on the monkey-patched vendor function + build_passthrough_logging_result, + "__benchflow_gemini_passthrough_patch__", + True, + ) + setattr( # noqa: B010 - avoids static narrowing on the vendor API + PassThroughStreamingHandler, + "_build_passthrough_logging_result", + staticmethod(build_passthrough_logging_result), + ) + + +_apply_patch() diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index fc53f1186..f93d74e94 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -64,6 +64,7 @@ LITELLM_SANDBOX_ROOT = "/tmp/benchflow-litellm" _CALLBACK_MODULE = "benchflow_litellm_callback" _PATCH_MODULE = "benchflow_litellm_bedrock_patch" +_GEMINI_PATCH_MODULE = "benchflow_litellm_gemini_passthrough_patch" # The proxy is an internal single-route gateway — it must never register the # FastAPI Swagger docs route. litellm's `_get_docs_url()` honours an inherited @@ -768,12 +769,18 @@ def _write_runtime_files( runtime_dir.mkdir(parents=True, exist_ok=True) callback_path = runtime_dir / f"{_CALLBACK_MODULE}.py" patch_path = runtime_dir / f"{_PATCH_MODULE}.py" + gemini_patch_path = runtime_dir / f"{_GEMINI_PATCH_MODULE}.py" sitecustomize_path = runtime_dir / "sitecustomize.py" config_path = runtime_dir / "config.yaml" callback_path.write_text(callback_module_source()) patch_source = Path(__file__).with_name("litellm_bedrock_patch.py").read_text() patch_path.write_text(patch_source) - sitecustomize_path.write_text(f"import {_PATCH_MODULE}\n") + gemini_patch_path.write_text( + Path(__file__).with_name("litellm_gemini_passthrough_patch.py").read_text() + ) + sitecustomize_path.write_text( + f"import {_PATCH_MODULE}\nimport {_GEMINI_PATCH_MODULE}\n" + ) config_path.write_text(yaml.safe_dump(config, sort_keys=False)) return config_path, callback_path, patch_path @@ -987,6 +994,7 @@ async def _upload_runtime_files_to_sandbox( "config": f"{runtime_dir}/config.yaml", "callback": f"{runtime_dir}/{_CALLBACK_MODULE}.py", "patch": f"{runtime_dir}/{_PATCH_MODULE}.py", + "gemini_patch": f"{runtime_dir}/{_GEMINI_PATCH_MODULE}.py", "sitecustomize": f"{runtime_dir}/sitecustomize.py", "launcher": f"{runtime_dir}/launcher.py", "stdout": f"{runtime_dir}/stdout.log", @@ -1012,7 +1020,16 @@ async def _upload_runtime_files_to_sandbox( ".py", ) await _upload_text( - sandbox, f"import {_PATCH_MODULE}\n", paths["sitecustomize"], ".py" + sandbox, + Path(__file__).with_name("litellm_gemini_passthrough_patch.py").read_text(), + paths["gemini_patch"], + ".py", + ) + await _upload_text( + sandbox, + f"import {_PATCH_MODULE}\nimport {_GEMINI_PATCH_MODULE}\n", + paths["sitecustomize"], + ".py", ) await _upload_text(sandbox, _sandbox_launcher_source(), paths["launcher"], ".py") await _upload_text( @@ -1503,7 +1520,16 @@ def _wire_litellm_agent_env( # in the upstream Gemini key server-side. updated.pop(LITELLM_MODEL_ALIAS_ENV, None) updated["GOOGLE_GEMINI_BASE_URL"] = f"{base_url.rstrip('/')}/gemini" - updated["GEMINI_API_KEY"] = master_key + # 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 + # key from the sandbox and silently bypass usage/evidence capture. + for key in ( + "GEMINI_API_KEY", + "GOOGLE_API_KEY", + "GOOGLE_GENERATIVE_AI_API_KEY", + ): + updated[key] = master_key return updated if agent == "claude-agent-acp": updated["ANTHROPIC_BASE_URL"] = base_url.rstrip("/") diff --git a/tests/test_acp.py b/tests/test_acp.py index 34bdf9116..766b79f02 100644 --- a/tests/test_acp.py +++ b/tests/test_acp.py @@ -1402,22 +1402,41 @@ def _make_mocks(): @pytest.mark.asyncio @pytest.mark.parametrize( - "model_in, expected_model", + "agent, model_in, expected_model", [ # Registered vllm/ prefix stripped; HF org/model intact — this is # what pi-acp and other ACP agents need for downstream routing. - ("vllm/Qwen/Qwen3.5-35B-A3B", "Qwen/Qwen3.5-35B-A3B"), - ("zai/glm-5", "glm-5"), + ("test-agent", "vllm/Qwen/Qwen3.5-35B-A3B", "Qwen/Qwen3.5-35B-A3B"), + ("test-agent", "zai/glm-5", "glm-5"), # Bare HF ID (no registered prefix) passes through unchanged. - ("Qwen/Qwen3-Coder", "Qwen/Qwen3-Coder"), + ("test-agent", "Qwen/Qwen3-Coder", "Qwen/Qwen3-Coder"), # Vertex ADC provider — prefix stripped like any other registered one. - ("anthropic-vertex/claude-sonnet-4-6", "claude-sonnet-4-6"), + ("test-agent", "anthropic-vertex/claude-sonnet-4-6", "claude-sonnet-4-6"), # No prefix at all — unchanged. - ("claude-sonnet-4-6", "claude-sonnet-4-6"), + ("test-agent", "claude-sonnet-4-6", "claude-sonnet-4-6"), + # Gemini CLI expects bare Google Gemini and Gemma model IDs. + ( + "gemini", + "google/gemini-3.1-flash-lite-preview", + "gemini-3.1-flash-lite-preview", + ), + ("gemini", "google/gemma-3-27b-it", "gemma-3-27b-it"), + ("acpx:gemini", "google/gemma-3-27b-it", "gemma-3-27b-it"), + ("gemini", "google/text-bison", "google/text-bison"), + ], + ids=[ + "vllm-hf", + "zai", + "bare-hf", + "vertex", + "no-prefix", + "gemini-google", + "gemma-google", + "acpx-gemma-google", + "unrelated-google", ], - ids=["vllm-hf", "zai", "bare-hf", "vertex", "no-prefix"], ) - async def test_model_id_selection(self, model_in, expected_model, tmp_path): + async def test_model_id_selection(self, agent, model_in, expected_model, tmp_path): from benchflow.acp.runtime import connect_acp mock_acp = self._make_mocks() @@ -1429,7 +1448,7 @@ async def test_model_id_selection(self, model_in, expected_model, tmp_path): ): await connect_acp( env=mock_env, - agent="test-agent", + agent=agent, agent_launch="test-agent", agent_env={}, sandbox_user=None, diff --git a/tests/test_agent_gemini_defaults.py b/tests/test_agent_gemini_defaults.py index 33e06b38d..fb1258c3d 100644 --- a/tests/test_agent_gemini_defaults.py +++ b/tests/test_agent_gemini_defaults.py @@ -77,6 +77,15 @@ def test_gemini_requires_env_matches_cli_native(self): "GEMINI_API_KEY (what the CLI reads); GOOGLE_API_KEY is an alias" ) + def test_gemini_headless_launch_trusts_the_sandbox_workspace(self): + """Guards PR #1030's E2E fix for Gemini CLI 0.42 folder trust. + + BenchFlow runs Gemini non-interactively with privileged tool approval. + Without ``--skip-trust``, the pinned CLI rejects every tool call before + the normalized model route can complete a task. + """ + assert "--skip-trust" in AGENTS["gemini"].launch_cmd.split() + if __name__ == "__main__": pytest.main([__file__, "-xvs"]) diff --git a/tests/test_litellm_config.py b/tests/test_litellm_config.py index 2611eecf0..b829c4021 100644 --- a/tests/test_litellm_config.py +++ b/tests/test_litellm_config.py @@ -274,3 +274,24 @@ def test_proxy_config_no_responses_bridge_for_non_openai_upstream(): config = litellm_proxy_config(route, master_key="sk-local") names = [entry["model_name"] for entry in config["model_list"]] assert not any(n.endswith("-responses-bridge") for n in names) + + +@pytest.mark.parametrize( + "model,expected", + [ + ( + "google/gemini-3.1-flash-lite-preview", + "gemini/gemini-3.1-flash-lite-preview", + ), + ("google/gemma-3-27b-it", "gemini/gemma-3-27b-it"), + ("google/geminix-1", "gemini/google/geminix-1"), + ], +) +def test_google_model_normalizes_only_gemini_families(model, expected): + """Google Gemini/Gemma IDs get one LiteLLM provider prefix.""" + route = resolve_litellm_route( + model, + {"GEMINI_API_KEY": "key"}, + ) + assert route.upstream_model == expected + assert route.litellm_params["model"] == expected diff --git a/tests/test_litellm_gemini_passthrough_patch.py b/tests/test_litellm_gemini_passthrough_patch.py new file mode 100644 index 000000000..07f47d88f --- /dev/null +++ b/tests/test_litellm_gemini_passthrough_patch.py @@ -0,0 +1,63 @@ +"""Regression coverage for the PR #1030 Gemma pass-through logging fix.""" + +from datetime import datetime +from types import SimpleNamespace + +from benchflow.providers.litellm_gemini_passthrough_patch import ( + _is_google_ai_studio_generate_content, +) + + +def test_google_ai_studio_path_detection_excludes_vertex() -> None: + assert _is_google_ai_studio_generate_content( + "https://generativelanguage.googleapis.com/v1beta/models/" + "gemma-4-31b-it:streamGenerateContent?alt=sse" + ) + assert not _is_google_ai_studio_generate_content( + "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/l/" + "publishers/google/models/gemma-4-31b-it:streamGenerateContent" + ) + + +def test_old_litellm_routes_gemma_streams_through_gemini_logger(monkeypatch) -> None: + """Guards PR #1030: Gemma usage must not be dropped by the Vertex logger.""" + from litellm.proxy.pass_through_endpoints.llm_provider_handlers.gemini_passthrough_logging_handler import ( + GeminiPassthroughLoggingHandler, + ) + from litellm.proxy.pass_through_endpoints.streaming_handler import ( + PassThroughStreamingHandler, + ) + from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType + + if hasattr(EndpointType, "GEMINI"): + return + + seen: dict[str, object] = {} + + def fake_gemini_handler(**kwargs): + seen.update(kwargs) + return {"result": "gemma-response", "kwargs": {"model": "gemma-4-31b-it"}} + + monkeypatch.setattr( + GeminiPassthroughLoggingHandler, + "_handle_logging_gemini_collected_chunks", + staticmethod(fake_gemini_handler), + ) + response, kwargs = PassThroughStreamingHandler._build_passthrough_logging_result( + litellm_logging_obj=SimpleNamespace(), + passthrough_success_handler_obj=SimpleNamespace(), + url_route=( + "https://generativelanguage.googleapis.com/v1beta/models/" + "gemma-4-31b-it:streamGenerateContent?alt=sse" + ), + request_body={"contents": []}, + endpoint_type=EndpointType.VERTEX_AI, + start_time=datetime.now(), + raw_bytes=[b'data: {"usageMetadata":{"promptTokenCount":1}}\n\n'], + end_time=datetime.now(), + model=None, + ) + + assert response == "gemma-response" + assert kwargs["model"] == "gemma-4-31b-it" + assert seen["all_chunks"] == ['data: {"usageMetadata":{"promptTokenCount":1}}'] diff --git a/tests/test_litellm_hardening.py b/tests/test_litellm_hardening.py index c39ee1e82..f8663296c 100644 --- a/tests/test_litellm_hardening.py +++ b/tests/test_litellm_hardening.py @@ -583,6 +583,15 @@ def test_bedrock_patch_preflight_passes_when_runtime_files_on_pythonpath(tmp_pat assert result.returncode == 0, result.stdout + result.stderr +def test_runtime_files_load_gemini_passthrough_patch(tmp_path): + """Guards PR #1030: host proxies load the Gemma usage-capture backport.""" + runtime_mod._write_runtime_files(tmp_path, config={"model_list": []}) + + sitecustomize = (tmp_path / "sitecustomize.py").read_text() + assert "import benchflow_litellm_gemini_passthrough_patch" in sitecustomize + assert (tmp_path / "benchflow_litellm_gemini_passthrough_patch.py").is_file() + + def test_bedrock_patch_preflight_fails_closed_when_patch_not_loaded(tmp_path): """THE regression test for issue #602's fail-open (fixed in PR #668): when the patch never loads (sitecustomize missing from PYTHONPATH — the exact silent-failure diff --git a/tests/test_litellm_runtime.py b/tests/test_litellm_runtime.py index 21f934b77..23efb3d61 100644 --- a/tests/test_litellm_runtime.py +++ b/tests/test_litellm_runtime.py @@ -693,7 +693,7 @@ async def fail_start(**_kwargs): @pytest.mark.asyncio async def test_gemini_uses_native_generate_content_through_sandbox_proxy(monkeypatch): - """Guards PR #942 remediation: Gemini review keeps no-web isolation.""" + """Guards PR #942 and PR #1030: every Google key alias stays proxied.""" starts = [] @@ -710,7 +710,11 @@ async def unexpected_host_start(**_kwargs): updated, provider_runtime = await ensure_litellm_runtime( agent="gemini", - agent_env={"GEMINI_API_KEY": "upstream-gemini-key"}, + 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", runtime=None, environment="docker", @@ -722,8 +726,12 @@ 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" - assert updated["GEMINI_API_KEY"] == provider_runtime.master_key - assert updated["GEMINI_API_KEY"] != "upstream-gemini-key" + for key in ( + "GEMINI_API_KEY", + "GOOGLE_API_KEY", + "GOOGLE_GENERATIVE_AI_API_KEY", + ): + assert updated[key] == provider_runtime.master_key assert LITELLM_MODEL_ALIAS_ENV not in updated