From 0845235188aaaad10c95d5fa2df12943d66f79c4 Mon Sep 17 00:00:00 2001 From: Kyoung Choe Date: Sun, 16 Aug 2026 20:49:52 +0000 Subject: [PATCH 1/7] Fix Gemini ACP model IDs --- src/benchflow/acp/runtime.py | 5 +++++ tests/test_litellm_hardening.py | 10 ++++++++++ 2 files changed, 15 insertions(+) diff --git a/src/benchflow/acp/runtime.py b/src/benchflow/acp/runtime.py index 0bf808155..7550090f2 100644 --- a/src/benchflow/acp/runtime.py +++ b/src/benchflow/acp/runtime.py @@ -255,6 +255,11 @@ 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 == "gemini" and model.startswith("google/"): + 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/tests/test_litellm_hardening.py b/tests/test_litellm_hardening.py index c39ee1e82..76fdd4f8a 100644 --- a/tests/test_litellm_hardening.py +++ b/tests/test_litellm_hardening.py @@ -83,6 +83,16 @@ def test_format_acp_model_passes_through_existing_provider_prefix(): ) +def test_gemini_acp_model_strips_google_prefix(): + """Guards Gemini model-resource 404s for current Google model IDs.""" + from benchflow.acp.runtime import _format_acp_model + + assert ( + _format_acp_model("google/gemini-3.5-flash-lite", "gemini") + == "gemini-3.5-flash-lite" + ) + + def test_mimo_registered_xiaomi_model_keeps_provider_prefix(): from benchflow.acp.runtime import _format_acp_model From b76f7334f5774928fc59653e442fd9f519f268a8 Mon Sep 17 00:00:00 2001 From: kywch Date: Sun, 30 Aug 2026 13:19:10 -0700 Subject: [PATCH 2/7] fix(litellm): normalize Google Gemini gateway routes --- src/benchflow/providers/litellm_config.py | 3 +++ tests/test_litellm_config.py | 10 ++++++++++ 2 files changed, 13 insertions(+) diff --git a/src/benchflow/providers/litellm_config.py b/src/benchflow/providers/litellm_config.py index acee4756c..c151b7579 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"): + 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/tests/test_litellm_config.py b/tests/test_litellm_config.py index 2611eecf0..d8f63700f 100644 --- a/tests/test_litellm_config.py +++ b/tests/test_litellm_config.py @@ -274,3 +274,13 @@ 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) + + +def test_google_gemini_model_uses_single_litellm_provider_prefix(): + """Guards models.dev Google IDs from becoming gemini/google/... routes.""" + route = resolve_litellm_route( + "google/gemini-3.1-flash-lite-preview", + {"GEMINI_API_KEY": "key"}, + ) + assert route.upstream_model == "gemini/gemini-3.1-flash-lite-preview" + assert route.litellm_params["model"] == "gemini/gemini-3.1-flash-lite-preview" From b4addc9f435824c5e3f5776a4d3f12c674996794 Mon Sep 17 00:00:00 2001 From: kywch Date: Sun, 30 Aug 2026 14:32:36 -0700 Subject: [PATCH 3/7] fix(gemini): support Gemma model routing --- src/benchflow/acp/runtime.py | 2 +- src/benchflow/providers/litellm_config.py | 2 +- tests/test_acp.py | 35 +++++++++++++++++------ tests/test_litellm_config.py | 21 ++++++++++---- 4 files changed, 44 insertions(+), 16 deletions(-) diff --git a/src/benchflow/acp/runtime.py b/src/benchflow/acp/runtime.py index 7550090f2..305f05e53 100644 --- a/src/benchflow/acp/runtime.py +++ b/src/benchflow/acp/runtime.py @@ -258,7 +258,7 @@ def _format_acp_model(model: str, agent: str) -> str: # 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 == "gemini" and model.startswith("google/"): + if agent == "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": diff --git a/src/benchflow/providers/litellm_config.py b/src/benchflow/providers/litellm_config.py index c151b7579..f703b32d2 100644 --- a/src/benchflow/providers/litellm_config.py +++ b/src/benchflow/providers/litellm_config.py @@ -375,7 +375,7 @@ 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"): + elif lower.startswith(("google/gemini-", "google/gemma-")): upstream = f"gemini/{model.split('/', 1)[1]}" required = ("GEMINI_API_KEY",) elif "gemini" in lower: diff --git a/tests/test_acp.py b/tests/test_acp.py index c20819c61..442b50ebe 100644 --- a/tests/test_acp.py +++ b/tests/test_acp.py @@ -1176,22 +1176,39 @@ 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"), + ("gemini", "google/text-bison", "google/text-bison"), + ], + ids=[ + "vllm-hf", + "zai", + "bare-hf", + "vertex", + "no-prefix", + "gemini-google", + "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() @@ -1203,7 +1220,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_litellm_config.py b/tests/test_litellm_config.py index d8f63700f..b829c4021 100644 --- a/tests/test_litellm_config.py +++ b/tests/test_litellm_config.py @@ -276,11 +276,22 @@ def test_proxy_config_no_responses_bridge_for_non_openai_upstream(): assert not any(n.endswith("-responses-bridge") for n in names) -def test_google_gemini_model_uses_single_litellm_provider_prefix(): - """Guards models.dev Google IDs from becoming gemini/google/... routes.""" +@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( - "google/gemini-3.1-flash-lite-preview", + model, {"GEMINI_API_KEY": "key"}, ) - assert route.upstream_model == "gemini/gemini-3.1-flash-lite-preview" - assert route.litellm_params["model"] == "gemini/gemini-3.1-flash-lite-preview" + assert route.upstream_model == expected + assert route.litellm_params["model"] == expected From f85120642fdafd52ad996e80c0cd5d4739005e6f Mon Sep 17 00:00:00 2001 From: kywch Date: Sun, 30 Aug 2026 19:17:12 -0700 Subject: [PATCH 4/7] fix(acp): normalize wrapped Gemini model IDs --- src/benchflow/acp/runtime.py | 10 ++++++++-- tests/test_acp.py | 2 ++ tests/test_litellm_hardening.py | 10 ---------- 3 files changed, 10 insertions(+), 12 deletions(-) diff --git a/src/benchflow/acp/runtime.py b/src/benchflow/acp/runtime.py index 05e022646..e3699072a 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, @@ -272,7 +276,9 @@ def _format_acp_model(model: str, agent: str) -> str: # 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 == "gemini" and model.startswith(("google/gemini-", "google/gemma-")): + 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": diff --git a/tests/test_acp.py b/tests/test_acp.py index 442b50ebe..293631d76 100644 --- a/tests/test_acp.py +++ b/tests/test_acp.py @@ -1195,6 +1195,7 @@ def _make_mocks(): "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=[ @@ -1205,6 +1206,7 @@ def _make_mocks(): "no-prefix", "gemini-google", "gemma-google", + "acpx-gemma-google", "unrelated-google", ], ) diff --git a/tests/test_litellm_hardening.py b/tests/test_litellm_hardening.py index 76fdd4f8a..c39ee1e82 100644 --- a/tests/test_litellm_hardening.py +++ b/tests/test_litellm_hardening.py @@ -83,16 +83,6 @@ def test_format_acp_model_passes_through_existing_provider_prefix(): ) -def test_gemini_acp_model_strips_google_prefix(): - """Guards Gemini model-resource 404s for current Google model IDs.""" - from benchflow.acp.runtime import _format_acp_model - - assert ( - _format_acp_model("google/gemini-3.5-flash-lite", "gemini") - == "gemini-3.5-flash-lite" - ) - - def test_mimo_registered_xiaomi_model_keeps_provider_prefix(): from benchflow.acp.runtime import _format_acp_model From 5c118a9d54198b41a5e5898f1f844decea6d5f27 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Tue, 1 Sep 2026 23:57:09 -0700 Subject: [PATCH 5/7] fix(gemini): trust sandbox workspace for headless runs --- src/benchflow/agents/registry.py | 6 +++++- tests/test_agent_gemini_defaults.py | 9 +++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) 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/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"]) From 28b82e336673bb58e55ef12c1597a914a42a3a7b Mon Sep 17 00:00:00 2001 From: Bingran You Date: Wed, 2 Sep 2026 00:07:54 -0700 Subject: [PATCH 6/7] fix(gemini): route every Google key alias through proxy --- src/benchflow/providers/litellm_runtime.py | 11 ++++++++++- tests/test_litellm_runtime.py | 16 ++++++++++++---- 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index fc53f1186..c1238fbcb 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -1503,7 +1503,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_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 From cc90842eabd77b3a9ed156d11ede0966933928f8 Mon Sep 17 00:00:00 2001 From: Bingran You Date: Wed, 2 Sep 2026 00:34:05 -0700 Subject: [PATCH 7/7] fix(gemini): capture Gemma pass-through usage --- .../litellm_gemini_passthrough_patch.py | 116 ++++++++++++++++++ src/benchflow/providers/litellm_runtime.py | 21 +++- .../test_litellm_gemini_passthrough_patch.py | 63 ++++++++++ tests/test_litellm_hardening.py | 9 ++ 4 files changed, 207 insertions(+), 2 deletions(-) create mode 100644 src/benchflow/providers/litellm_gemini_passthrough_patch.py create mode 100644 tests/test_litellm_gemini_passthrough_patch.py 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 c1238fbcb..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( 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