Skip to content
Merged
13 changes: 12 additions & 1 deletion src/benchflow/acp/runtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
6 changes: 5 additions & 1 deletion src/benchflow/agents/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 3 additions & 0 deletions src/benchflow/providers/litellm_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",)
Expand Down
116 changes: 116 additions & 0 deletions src/benchflow/providers/litellm_gemini_passthrough_patch.py
Original file line number Diff line number Diff line change
@@ -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()
32 changes: 29 additions & 3 deletions src/benchflow/providers/litellm_runtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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",
Expand All @@ -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(
Expand Down Expand Up @@ -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("/")
Expand Down
37 changes: 28 additions & 9 deletions tests/test_acp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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,
Expand Down
9 changes: 9 additions & 0 deletions tests/test_agent_gemini_defaults.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"])
21 changes: 21 additions & 0 deletions tests/test_litellm_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading
Loading