diff --git a/.env.sample b/.env.sample index e7c8749f5..07b61963e 100644 --- a/.env.sample +++ b/.env.sample @@ -20,6 +20,7 @@ # OPENAI_API_KEY= # codex-acp # GEMINI_API_KEY= # gemini # ZAI_API_KEY= # zai/ provider (openclaw, claude-agent-acp) +# AI_GATEWAY_API_KEY= # vercel/ provider (Vercel AI Gateway) # ── Provider options ──────────────────────────────────────── # Optional OpenAI prompt cache retention policy. Allowed: in_memory, 24h. diff --git a/src/benchflow/acp/client.py b/src/benchflow/acp/client.py index 98d033766..0020e0132 100644 --- a/src/benchflow/acp/client.py +++ b/src/benchflow/acp/client.py @@ -54,6 +54,39 @@ def _auto_approve_option_id(options: list[dict[str, Any]]) -> str: return option_id +# Non-spec ACP stop reasons and usage keys seen in the wild (fx's dialect), +# mapped to the spec values the SDK models accept. +_STOP_REASON_ALIASES = { + "refused": "refusal", + "max_output_tokens": "max_tokens", + "max_model_turns": "max_turn_requests", +} +_USAGE_KEY_ALIASES = { + "cacheReadTokens": "cachedReadTokens", + "cacheWriteTokens": "cachedWriteTokens", + "reasoningTokens": "thoughtTokens", +} + + +def _normalize_prompt_result(result: dict) -> None: + """Map known non-spec ACP values in-place onto their spec equivalents.""" + stop = result.get("stopReason") + if stop in _STOP_REASON_ALIASES: + result["stopReason"] = _STOP_REASON_ALIASES[stop] + usage = result.get("usage") + if isinstance(usage, dict): + for theirs, ours in _USAGE_KEY_ALIASES.items(): + if theirs in usage: + usage[ours] = usage.pop(theirs) + if "inputTokens" in usage and "outputTokens" in usage: + usage.setdefault( + "totalTokens", usage["inputTokens"] + usage["outputTokens"] + ) + else: + # The SDK requires input/output/total; drop unusable usage. + result.pop("usage") + + class ACPClient: """Client that speaks ACP to an agent process. @@ -399,6 +432,8 @@ async def prompt(self, text: str) -> PromptResult: result = await self._send_request( "session/prompt", params.model_dump(by_alias=True, exclude_none=True) ) + if isinstance(result, dict): + _normalize_prompt_result(result) prompt_result = PromptResult.model_validate(result) # The SDK exposes ``stop_reason`` as a plain string; coerce it to the # vendored ``StopReason`` enum so consumers keep ``.value`` / member diff --git a/src/benchflow/agents/env.py b/src/benchflow/agents/env.py index ffa5e5f02..d8f1080e9 100644 --- a/src/benchflow/agents/env.py +++ b/src/benchflow/agents/env.py @@ -472,6 +472,7 @@ def resolve_provider_env( ZAI_CODING_REGISTRY_BASE_ENV, find_provider, find_provider_for_bare_model, + is_native_provider_model, resolve_base_url, strip_provider_prefix, ) @@ -486,6 +487,15 @@ def resolve_provider_env( # is never emitted and harnesses that rely on it misroute (openhands' litellm # saw no provider; openclaw defaulted to anthropic/). Mirrors acp/runtime.py. _prov = find_provider(model) or find_provider_for_bare_model(model) + if ( + agent_cfg + and agent_cfg.native_provider + and not is_native_provider_model(agent_cfg.native_provider, model) + ): + raise ValueError( + f"Agent {agent!r} only runs {agent_cfg.native_provider}/ models; " + f"got {model!r}." + ) if _prov: _prov_name, _prov_cfg = _prov agent_env.setdefault("BENCHFLOW_PROVIDER_NAME", _prov_name) diff --git a/src/benchflow/agents/manifest.py b/src/benchflow/agents/manifest.py index e4f676ae2..abc49bf90 100644 --- a/src/benchflow/agents/manifest.py +++ b/src/benchflow/agents/manifest.py @@ -93,6 +93,8 @@ "disallow_hosted_search_launch_suffix", "task_mcp_transport", "task_mcp_config_path", + # Proxy-bypass is a key-isolation decision; not manifest-declarable. + "native_provider", } ) diff --git a/src/benchflow/agents/providers.py b/src/benchflow/agents/providers.py index d7feeedc4..db3ac4249 100644 --- a/src/benchflow/agents/providers.py +++ b/src/benchflow/agents/providers.py @@ -218,6 +218,21 @@ def all_endpoints(self) -> dict[str, str]: auth_type="api_key", auth_env="OPENROUTER_API_KEY", ), + # Vercel AI Gateway: OpenAI-compatible endpoints under /v1; the Anthropic + # Messages surface sits at the root (clients append /v1/messages). Model ids + # keep the gateway's creator/model form: vercel/anthropic/claude-sonnet-4.5. + "vercel": ProviderConfig( + name="vercel", + base_url="https://ai-gateway.vercel.sh/v1", + api_protocol="openai-completions", + auth_type="api_key", + auth_env="AI_GATEWAY_API_KEY", + endpoints={ + "openai-completions": "https://ai-gateway.vercel.sh/v1", + "openai-responses": "https://ai-gateway.vercel.sh/v1", + "anthropic-messages": "https://ai-gateway.vercel.sh", + }, + ), # TODO: add eu-openai (https://eu.api.openai.com/v1) when needed. # ── OpenAI-compatible inference servers (user-supplied base_url) ── "vllm": ProviderConfig( @@ -420,6 +435,12 @@ def find_provider(model: str) -> tuple[str, ProviderConfig] | None: return name, cfg +def is_native_provider_model(native_provider: str, model: str) -> bool: + """True when ``model`` carries the given registered provider's prefix.""" + provider = find_provider(model) + return provider is not None and provider[0] == native_provider + + def _bare_model_matches_token(model: str, token: str) -> bool: """True if a bare model id belongs to the family named by *token*. diff --git a/src/benchflow/agents/registry.py b/src/benchflow/agents/registry.py index badc3fff3..2745d54da 100644 --- a/src/benchflow/agents/registry.py +++ b/src/benchflow/agents/registry.py @@ -50,7 +50,7 @@ import base64 import shlex -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path from benchflow._utils.text import describe_exception @@ -123,6 +123,8 @@ def _apt_install(*packages: str) -> str: _OPENHANDS_CLI_GIT_REV = "2df8a2835d3f1bd2f2eadf5a7a2e1ad0dfb0d271" _OPENHANDS_SDK_VERSION = "1.28.1" _OPENHANDS_TOOLS_VERSION = "1.28.1" +# fx release pin, passed to the fx.sh installer (`bash -s -- `). +_FX_VERSION = "v0.0.8" _JS_AGENT_PATH = ( f"{_BENCHFLOW_BIN_PREFIX}:{_BENCHFLOW_JS_AGENT_PREFIX}/bin:" f"{_BENCHFLOW_NODE_PREFIX}/bin:$PATH" @@ -487,6 +489,9 @@ class AgentConfig: supports_acp_set_model: bool = True # Some ACP agents configure the model through env/config at launch time and # do not implement session/set_model (e.g. OpenHands CLI ACP). + native_provider: str = "" + # Provider whose models this agent serves natively; its traffic cannot be + # proxied, so runs bypass the proxy for those models and reject others. # ACP session config option id used for model selection when an agent # exposes model as a session option instead of implementing set_model. acp_model_config_id: str = "" @@ -1019,6 +1024,35 @@ class AgentConfig: ), disallow_web_tools_owned_paths=["$HOME/.openhands"], ), + "fx": AgentConfig( + name="fx", + description="Vercel fx agent via ACP (native binary; models served " + "through the Vercel AI Gateway)", + install_cmd=( + "export DEBIAN_FRONTEND=noninteractive && " + "( command -v curl >/dev/null 2>&1 || " + f" {_apt_install('curl', 'ca-certificates')} ) && " + # Shared prefix so the sandbox user inherits the binary. + "export FX_INSTALL_DIR=/usr/local/bin && " + f"curl -fsSL https://fx.sh/setup.sh | bash -s -- {_FX_VERSION} && " + "command -v fx >/dev/null 2>&1" + ), + launch_cmd="fx acp", + protocol="acp", + requires_env=["AI_GATEWAY_API_KEY"], + env_mapping={ + "BENCHFLOW_PROVIDER_MODEL": "FX_MODEL", + }, + supports_acp_set_model=False, + # fx speaks the AI SDK gateway wire protocol. + native_provider="vercel", + disallow_web_tools_setup_cmd=_json_settings_merge( + "$BENCHFLOW_AGENT_HOME/.fx/settings.json", + 'perm=d.setdefault("permission",{});' + 'perm["web_fetch"]={"*":"deny"};perm["web_search"]={"*":"deny"}', + ), + disallow_web_tools_owned_paths=["$HOME/.fx"], + ), } @@ -1178,13 +1212,11 @@ def _acpx_wrap(config: AgentConfig) -> AgentConfig: acpx_agent_name = alias break - # The acpx wrapper only overrides name/install_cmd/launch_cmd. Every other - # AgentConfig field must pass through from the underlying agent so that - # routing-relevant attributes (api_protocol, default_model, env_mapping, - # requires_env, credentials, …) survive when the wrapped config is cached - # into AGENTS and later read by resolve_provider_env. ``protocol`` stays - # "acp" because acpx itself speaks ACP regardless of the inner agent. - return AgentConfig( + # Only name/install/launch/protocol/description change; every other field + # passes through so routing-relevant attributes survive when the wrapped + # config is cached into AGENTS and later read by resolve_provider_env. + return replace( + config, # ``acpx:`` runtime key — see acpx_runtime_key / module-level contract. name=acpx_runtime_key(config.name), install_cmd=f"{config.install_cmd} && {_ACPX_INSTALL}", @@ -1192,28 +1224,7 @@ def _acpx_wrap(config: AgentConfig) -> AgentConfig: f'export PATH="{_JS_AGENT_PATH}" && acpx {acpx_agent_name} --approve-all' ), protocol="acp", - session_factory=config.session_factory, - requires_env=config.requires_env, description=f"{config.description} (via acpx)", - skill_paths=config.skill_paths, - install_timeout=config.install_timeout, - default_model=config.default_model, - api_protocol=config.api_protocol, - env_mapping=config.env_mapping, - credential_files=config.credential_files, - home_dirs=config.home_dirs, - acp_model_format=config.acp_model_format, - subscription_auth=config.subscription_auth, - supports_acp_set_model=config.supports_acp_set_model, - acp_model_config_id=config.acp_model_config_id, - acp_effort_config_id=config.acp_effort_config_id, - disallow_web_tools_setup_cmd=config.disallow_web_tools_setup_cmd, - disallow_web_tools_owned_paths=config.disallow_web_tools_owned_paths, - disallow_web_tools_launch_suffix=config.disallow_web_tools_launch_suffix, - disallow_hosted_search_setup_cmd=config.disallow_hosted_search_setup_cmd, - disallow_hosted_search_launch_suffix=config.disallow_hosted_search_launch_suffix, - task_mcp_transport=config.task_mcp_transport, - task_mcp_config_path=config.task_mcp_config_path, ) diff --git a/src/benchflow/continue_run/orchestrator.py b/src/benchflow/continue_run/orchestrator.py index 3b9a92310..9ff8db404 100644 --- a/src/benchflow/continue_run/orchestrator.py +++ b/src/benchflow/continue_run/orchestrator.py @@ -4,10 +4,10 @@ new run produces a standard HF-compatible folder) but injects a record-replay proxy in front of OpenHands: -- Host proxy mode uses ``usage_tracking="off"`` so ``ensure_litellm_runtime`` is - a no-op and the - ``agent_env`` we pass through (``LLM_BASE_URL`` → the replay proxy) is left - untouched (``providers/litellm_runtime.py``). +- Host proxy mode passes ``model=None`` so ``ensure_litellm_runtime`` is a + no-op and the ``agent_env`` we pass through (``LLM_BASE_URL`` → the replay + proxy) is left untouched (``providers/litellm_runtime.py``); + ``usage_tracking="off"`` marks telemetry off in result metadata. - Sandbox proxy mode starts the provider LiteLLM proxy inside the sandbox and keeps the replay proxy on sandbox loopback, so remote environments such as Daytona do not need host-loopback connectivity. @@ -184,8 +184,8 @@ def build_rollout_config( environment=run.environment, sandbox_user=run.sandbox_user, agent_env=agent_env, - # The seam that stops benchflow starting its own gateway / rewriting - # LLM_BASE_URL — our replay proxy stays in front of the agent. + # model=None above keeps benchflow's gateway out (the replay proxy + # stays in front of the agent); "off" marks telemetry off in results. usage_tracking="off", timeout=timeout, agent_idle_timeout=run.agent_idle_timeout_sec, @@ -656,7 +656,6 @@ async def _write_artifacts_before_cleanup() -> None: runtime=None, environment=run.environment, session_id=rollout_name, - usage_tracking="required", sandbox=rollout.env, ) replay_proxy = await SandboxReplayProxy.start( diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index 60c9953c3..c50dc4884 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -28,6 +28,7 @@ from benchflow._utils.text import describe_exception from benchflow.agents.codex_config import apply_codex_provider_config from benchflow.agents.env import uses_native_subscription_auth +from benchflow.agents.providers import PROVIDERS, is_native_provider_model from benchflow.agents.registry import AGENTS from benchflow.providers.litellm_bedrock_preflight import ( BEDROCK_PATCH_PREFLIGHT_SOURCE, @@ -53,7 +54,7 @@ from benchflow.sandbox.providers import SANDBOX_MODEL_PROXY_PROVIDERS from benchflow.trajectories._llm_capture import LiveLLMTrajectoryWriter from benchflow.trajectories.types import Trajectory -from benchflow.usage_tracking import UsageTrackingConfig, usage_unavailable +from benchflow.usage_tracking import usage_unavailable if TYPE_CHECKING: from benchflow.contracts.planes import LiveUsageGateway @@ -105,10 +106,6 @@ # frozen-but-plausible token count. _LIVE_CAPTURE_STALL_WARN_TICKS = 30 -# Agents that cannot make model calls through LiteLLM. ``oracle`` has no model -# at all. Gemini is routable through LiteLLM's native Google GenerateContent -# endpoints; keeping it behind the proxy is required for no-web reviewer runs. -_NATIVE_PROTOCOL_AGENTS = frozenset({"oracle"}) # Providers whose mandatory LiteLLM proxy runs inside the sandbox. Keeping this # placement policy in the canonical provider registry prevents a new backend from # accidentally handing an in-sandbox agent a host-loopback endpoint. @@ -646,8 +643,19 @@ async def log_tail(self) -> str: def needs_litellm_runtime(agent: str, model: str | None) -> bool: - """True when an agent/model pair should be routed through LiteLLM.""" - return bool(model) and agent not in _NATIVE_PROTOCOL_AGENTS + """True when an agent/model pair should be routed through LiteLLM. + + ``oracle`` has no model; ``native_provider`` agents bypass the proxy only + for their own provider's models (a wire protocol the proxy does not + expose). Gemini stays behind the proxy (LiteLLM's native GenerateContent + endpoints) — required for no-web reviewer runs. + """ + if not model or agent == "oracle": + return False + cfg = AGENTS.get(agent) + if cfg is None or not cfg.native_provider: + return True + return not is_native_provider_model(cfg.native_provider, model) def _find_free_port() -> int: @@ -1315,10 +1323,16 @@ def _missing_required_env(route: LiteLLMRoute, env: dict[str, str]) -> list[str] return missing +def _scrub_foreign_provider_secrets( + env: dict[str, str], keep: str | None +) -> dict[str, str]: + """Copy of ``env`` without raw provider secrets, except ``keep``.""" + secrets = _provider_secret_env_names() - {keep} + return {k: v for k, v in env.items() if k not in secrets} + + def _provider_secret_env_names() -> set[str]: """Upstream provider credentials the proxy owns and the agent must not see.""" - from benchflow.agents.providers import PROVIDERS - names = { "ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", @@ -1626,7 +1640,6 @@ async def ensure_litellm_runtime( runtime: Any | None, environment: str, session_id: str = "", - usage_tracking: UsageTrackingConfig | dict[str, Any] | str | None = None, sandbox: Any | None = None, sandbox_setup_timeout: int = 120, required_skill_names: tuple[str, ...] = (), @@ -1637,12 +1650,15 @@ async def ensure_litellm_runtime( Every LiteLLM-routable agent is *always* routed through the proxy so provider traffic is metered and captured (``llm_trajectory.jsonl``) and the - raw provider key never reaches the agent. ``usage_tracking`` no longer gates - whether the proxy runs — it only governs whether trusted telemetry is - *required* (``required`` fails closed when usage cannot be captured at all). + raw provider key never reaches the agent. Required-mode usage tracking is + enforced end-of-run by ``Rollout._enforce_required_usage_tracking``. The only agents that skip the proxy are those that physically cannot be - routed through it: ``oracle`` (no model) and native-subscription auth (no - API key to proxy). Gemini uses LiteLLM's native GenerateContent endpoints. + routed through it: ``oracle`` (no model), native-subscription auth (no + API key to proxy), and ``native_provider`` agents on their own provider's + models (wire protocol the proxy does not expose; that provider's key + reaches the agent, other provider secrets are scrubbed, and usage is + ACP-reported — ``required`` tracking is enforced at end of run). Gemini + uses LiteLLM's native GenerateContent endpoints. """ # Re-entrant connects pass back proxy-owned env, which cannot reconstruct # upstream routing or credentials. Restore controller-held source config. @@ -1657,8 +1673,6 @@ async def ensure_litellm_runtime( ): agent_env = dict(runtime.source_env) - usage_cfg = UsageTrackingConfig.coerce(usage_tracking).with_env_defaults() - if uses_native_subscription_auth(agent, model, agent_env): return await _skip_litellm_runtime( agent_env, @@ -1667,10 +1681,11 @@ async def ensure_litellm_runtime( ) if not needs_litellm_runtime(agent, model): - if usage_cfg.mode == "required" and agent != "oracle": - raise RuntimeError( - "Token usage tracking is required, but agent " - f"{agent!r} cannot be routed through LiteLLM." + cfg = AGENTS.get(agent) + if cfg is not None and cfg.native_provider: + # The agent only needs its own provider's key; scrub the rest. + agent_env = _scrub_foreign_provider_secrets( + agent_env, PROVIDERS[cfg.native_provider].auth_env ) return await _skip_litellm_runtime(agent_env, runtime) assert model is not None diff --git a/src/benchflow/providers/runtime.py b/src/benchflow/providers/runtime.py index 55ab8591a..7023d9b41 100644 --- a/src/benchflow/providers/runtime.py +++ b/src/benchflow/providers/runtime.py @@ -11,8 +11,6 @@ from pathlib import Path from typing import TYPE_CHECKING, Any -from benchflow.usage_tracking import UsageTrackingConfig - if TYPE_CHECKING: from benchflow.providers.litellm_runtime import LiteLLMProcess @@ -49,7 +47,6 @@ async def ensure_litellm_runtime( runtime: ProviderRuntime | None, environment: str, session_id: str = "", - usage_tracking: UsageTrackingConfig | dict[str, Any] | str | None = None, sandbox: Any | None = None, sandbox_setup_timeout: int = 120, required_skill_names: tuple[str, ...] = (), @@ -67,7 +64,6 @@ async def ensure_litellm_runtime( runtime=runtime, environment=environment, session_id=session_id, - usage_tracking=usage_tracking, sandbox=sandbox, sandbox_setup_timeout=sandbox_setup_timeout, required_skill_names=required_skill_names, diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index c06a584b2..6a4b2d71c 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -1315,7 +1315,6 @@ async def connect(self) -> None: runtime=getattr(self, "_usage_runtime", None), environment=cfg.environment, session_id=getattr(self, "_rollout_name", "") or "", - usage_tracking=cfg.usage_tracking, sandbox=self._env, sandbox_setup_timeout=cfg.sandbox_setup_timeout, required_skill_names=getattr(self, "_required_skill_names", ()), @@ -2341,7 +2340,6 @@ async def connect_as(self, role: Role) -> None: runtime=getattr(self, "_usage_runtime", None), environment=cfg.environment, session_id=getattr(self, "_rollout_name", "") or "", - usage_tracking=cfg.usage_tracking, sandbox=self._env, sandbox_setup_timeout=cfg.sandbox_setup_timeout, required_skill_names=getattr(self, "_required_skill_names", ()), diff --git a/tests/test_agent_env_resolution.py b/tests/test_agent_env_resolution.py index ea4e5a8f9..2dc648603 100644 --- a/tests/test_agent_env_resolution.py +++ b/tests/test_agent_env_resolution.py @@ -128,5 +128,17 @@ def test_explicit_agent_env_overrides_host(self, monkeypatch, tmp_path): assert result["GEMINI_API_KEY"] == "explicit" +class TestResolveAgentEnvFx: + def test_non_vercel_model_rejected(self, monkeypatch, tmp_path): + """Guards PR #1052: fx serves only vercel/ models; others fail fast.""" + _clear_keys(monkeypatch) + _patch_no_subscription(monkeypatch, tmp_path) + monkeypatch.setenv("BENCHFLOW_DOTENV_PATH", str(tmp_path / "no.env")) + monkeypatch.setenv("DEEPSEEK_API_KEY", "dk-test") + + with pytest.raises(ValueError, match="vercel/"): + resolve_agent_env(agent="fx", model="deepseek/deepseek-v4", agent_env=None) + + if __name__ == "__main__": pytest.main([__file__, "-xvs"]) diff --git a/tests/test_agent_spec.py b/tests/test_agent_spec.py index c2bd7bda5..69215a47e 100644 --- a/tests/test_agent_spec.py +++ b/tests/test_agent_spec.py @@ -71,6 +71,16 @@ def test_resolve_acpx_preserves_web_policy_owned_paths(self): config = resolve_agent("acpx/gemini") assert config.disallow_web_tools_owned_paths == ["$HOME/.gemini"] + def test_resolve_acpx_preserves_native_provider(self): + """Guards PR #1052: acpx-wrapped fx keeps its gateway-native routing.""" + from benchflow.agents.registry import resolve_agent_key + from benchflow.providers.litellm_runtime import needs_litellm_runtime + + config = resolve_agent("acpx/fx") + assert config.native_provider == "vercel" + key = resolve_agent_key("acpx/fx") + assert not needs_litellm_runtime(key, "vercel/anthropic/claude-sonnet-4.5") + def test_acpx_wrap_carries_routing_fields(self): """Regression for PR #322: _acpx_wrap must inherit api_protocol and default_model (and all other non-overridden fields) from the underlying diff --git a/tests/test_internet_policy.py b/tests/test_internet_policy.py index 246fae660..47ad81e2e 100644 --- a/tests/test_internet_policy.py +++ b/tests/test_internet_policy.py @@ -197,6 +197,11 @@ def test_agent_registry_has_supported_hard_web_disable_snippets(): assert "webfetch" in mimo_cmd assert "mimocode" in mimo_cmd + fx_cmd = AGENTS["fx"].disallow_web_tools_setup_cmd + assert "web_fetch" in fx_cmd + assert "web_search" in fx_cmd + assert "deny" in fx_cmd + @pytest.mark.asyncio async def test_connect_as_applies_hard_web_policy_to_role_agent(tmp_path): diff --git a/tests/test_litellm_config.py b/tests/test_litellm_config.py index 188c1a2e6..cdea2353d 100644 --- a/tests/test_litellm_config.py +++ b/tests/test_litellm_config.py @@ -240,6 +240,18 @@ def test_openrouter_route_uses_openai_compatible_endpoint(): assert route.required_env == ("OPENROUTER_API_KEY",) +def test_vercel_route_uses_openai_compatible_endpoint(): + route = resolve_litellm_route( + "vercel/anthropic/claude-sonnet-4.5", + {"AI_GATEWAY_API_KEY": "vck-test"}, + ) + + assert route.upstream_model == "openai/anthropic/claude-sonnet-4.5" + assert route.litellm_params["api_base"] == "https://ai-gateway.vercel.sh/v1" + assert route.litellm_params["api_key"] == "os.environ/AI_GATEWAY_API_KEY" + assert route.required_env == ("AI_GATEWAY_API_KEY",) + + def test_proxy_config_registers_plain_and_openai_aliases(): route = resolve_litellm_route( "aws-bedrock/us.anthropic.claude-opus-4-8", diff --git a/tests/test_litellm_hardening.py b/tests/test_litellm_hardening.py index d87984ee4..ae49d5129 100644 --- a/tests/test_litellm_hardening.py +++ b/tests/test_litellm_hardening.py @@ -53,6 +53,8 @@ async def stop(self) -> None: [ ("gemini", "gemini-3.5-flash", True), ("oracle", "openai/gpt-4.1-mini", False), + ("fx", "vercel/anthropic/claude-sonnet-4.5", False), + ("fx", "deepseek/deepseek-v4-flash", True), ("openhands", "gemini-3.5-flash", True), ("codex-acp", "openai/gpt-4.1-mini", True), ("openhands", None, False), @@ -169,6 +171,29 @@ async def fake_start(**kwargs): assert updated["LLM_BASE_URL"] == "http://127.0.0.1:4000/v1" +@pytest.mark.asyncio +async def test_native_provider_bypass_scrubs_foreign_secrets(): + """Guards PR #1052: foreign secrets must not reach a native_provider agent.""" + real = { + "AI_GATEWAY_API_KEY": "real-gateway-secret", + "OPENAI_API_KEY": "real-openai-secret", + "DEEPSEEK_API_KEY": "real-deepseek-secret", + } + updated, provider_runtime = await ensure_litellm_runtime( + agent="fx", + agent_env=real, + model="vercel/anthropic/claude-sonnet-4.5", + runtime=None, + environment="local", + session_id="s", + ) + + assert provider_runtime is None + assert updated["AI_GATEWAY_API_KEY"] == "real-gateway-secret" + assert "OPENAI_API_KEY" not in updated + assert "DEEPSEEK_API_KEY" not in updated + + # # # Networking: host proxy binds loopback locally, bridge IP for docker # # # diff --git a/tests/test_litellm_runtime.py b/tests/test_litellm_runtime.py index 80d7d1b94..cd0452c5f 100644 --- a/tests/test_litellm_runtime.py +++ b/tests/test_litellm_runtime.py @@ -63,7 +63,6 @@ async def fake_start(**kwargs): runtime=None, environment="docker", session_id="run-1", - usage_tracking="required", ) assert provider_runtime is not None @@ -227,7 +226,6 @@ async def fake_start(**kwargs): runtime=None, environment="docker", session_id="run-1", - usage_tracking="required", ) assert provider_runtime is not None @@ -272,7 +270,6 @@ async def fake_start(**kwargs): runtime=None, environment="docker", session_id="run-1", - usage_tracking="required", ) assert provider_runtime is not None @@ -387,7 +384,6 @@ async def test_required_usage_fails_when_litellm_lacks_provider_key(monkeypatch) model="openai/gpt-4.1-mini", runtime=None, environment="docker", - usage_tracking="required", ) @@ -407,7 +403,6 @@ async def fail_start(**_kwargs): model="openai/gpt-4.1-mini", runtime=None, environment="docker", - usage_tracking="required", ) assert updated == env @@ -430,7 +425,6 @@ async def fail_start(**_kwargs): model="claude-sonnet-4-6", runtime=None, environment="docker", - usage_tracking="required", ) assert updated == env @@ -454,7 +448,6 @@ async def fake_start(**kwargs): model="openai/gpt-4.1-mini", runtime=None, environment="docker", - usage_tracking="off", ) assert provider_runtime is not None @@ -488,7 +481,6 @@ async def fake_start(**kwargs): model="openai/gpt-4.1-mini", runtime=existing, environment="docker", - usage_tracking="off", ) assert old_server.stopped is True @@ -520,7 +512,6 @@ async def fake_start(**kwargs): runtime=None, environment="daytona", session_id="run-1", - usage_tracking="off", sandbox=SimpleNamespace(), ) @@ -582,7 +573,6 @@ async def fail_start(**_kwargs): model="openai/gpt-4.1-mini", runtime=None, environment="docker", - usage_tracking="auto", ) @@ -606,7 +596,6 @@ async def fail_start(**_kwargs): model="azure-foundry-openai/gpt-4.1-mini", runtime=None, environment="docker", - usage_tracking="auto", ) @@ -626,7 +615,6 @@ async def fail_start(**_kwargs): model="openai/gpt-4.1-mini", runtime=None, environment="docker", - usage_tracking="auto", ) @@ -662,7 +650,6 @@ async def fail_start(**_kwargs): model="openai/gpt-4.1-mini", runtime=None, environment="docker", - usage_tracking="auto", ) message = str(excinfo.value) @@ -691,7 +678,6 @@ async def fail_start(**_kwargs): model="aws-bedrock/us.anthropic.claude-opus-4-8-20251101-v1:0", runtime=None, environment="docker", - usage_tracking="auto", ) @@ -706,7 +692,6 @@ async def test_auto_usage_requires_sandbox_handle_for_sandbox_local_litellm(): model="openai/gpt-4.1-mini", runtime=None, environment="daytona", - usage_tracking="auto", sandbox=None, ) @@ -730,7 +715,6 @@ def fail_route(*_args, **_kwargs): model="azure-foundry-openai/gpt-4.1-mini", runtime=None, environment="daytona", - usage_tracking="auto", sandbox=None, ) @@ -753,7 +737,6 @@ async def fail_start(**_kwargs): model="openai/gpt-4.1-mini", runtime=None, environment="docker", - usage_tracking="required", ) @@ -784,7 +767,6 @@ async def unexpected_host_start(**_kwargs): model="gemini-2.5-flash", runtime=None, environment="docker", - usage_tracking="required", sandbox=sandbox, force_sandbox_local=True, ) diff --git a/tests/test_native_acp_usage.py b/tests/test_native_acp_usage.py index 7bf0d04e6..a37889b33 100644 --- a/tests/test_native_acp_usage.py +++ b/tests/test_native_acp_usage.py @@ -80,6 +80,69 @@ async def test_disconnect_preserves_native_usage_in_final_metrics(): assert rollout._usage_metrics["total_tokens"] == 14 +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fx_reason,spec_reason", + [ + ("refused", "refusal"), + ("max_output_tokens", "max_tokens"), + ("max_model_turns", "max_turn_requests"), + ], +) +async def test_acp_client_normalizes_fx_stop_reasons(fx_reason, spec_reason): + """Guards PR #1052: fx's non-spec stop reasons parse as spec values.""" + from benchflow.acp.client import ACPClient + from benchflow.acp.session import ACPSession + + client = ACPClient.__new__(ACPClient) + client._session = ACPSession("session-1") + + async def fake_send_request(method, params): + return {"stopReason": fx_reason} + + client._send_request = fake_send_request + + result = await client.prompt("solve") + + assert result.stop_reason == spec_reason + + +@pytest.mark.asyncio +async def test_acp_client_normalizes_fx_usage_keys(): + """Guards PR #1052: fx's usage keys (cacheReadTokens, reasoningTokens, no + totalTokens) are captured instead of failing PromptResponse validation.""" + from benchflow.acp.client import ACPClient + from benchflow.acp.session import ACPSession + + client = ACPClient.__new__(ACPClient) + client._session = ACPSession("session-1") + + async def fake_send_request(method, params): + return { + "stopReason": "end_turn", + "usage": { + "inputTokens": 10, + "outputTokens": 4, + "cacheReadTokens": 2, + "cacheWriteTokens": 1, + "reasoningTokens": 3, + }, + } + + client._send_request = fake_send_request + + await client.prompt("solve") + + assert client._session.latest_usage_totals() == { + "input_tokens": 10, + "output_tokens": 4, + "total_tokens": 14, + "cached_read_tokens": 2, + "cached_write_tokens": 1, + "thought_tokens": 3, + } + + def test_rollout_native_acp_usage_uses_cumulative_deltas(): """Guards PR #613 follow-up: ACP cumulative usage is not double-counted.""" from benchflow.acp.session import ACPSession diff --git a/tests/test_providers.py b/tests/test_providers.py index 387da1264..f10f97f7c 100644 --- a/tests/test_providers.py +++ b/tests/test_providers.py @@ -89,6 +89,23 @@ def test_openrouter_prefix(self): == "qwen/qwen3.5-397b-a17b" ) + def test_vercel_prefix(self): + name, cfg = find_provider("vercel/anthropic/claude-sonnet-4.5") + + assert name == "vercel" + assert cfg.api_protocol == "openai-completions" + assert cfg.auth_env == "AI_GATEWAY_API_KEY" + assert cfg.base_url == "https://ai-gateway.vercel.sh/v1" + assert cfg.endpoints["anthropic-messages"] == "https://ai-gateway.vercel.sh" + assert ( + resolve_auth_env("vercel/anthropic/claude-sonnet-4.5") + == "AI_GATEWAY_API_KEY" + ) + assert ( + strip_provider_prefix("vercel/anthropic/claude-sonnet-4.5") + == "anthropic/claude-sonnet-4.5" + ) + @pytest.mark.parametrize( ("model", "expected_protocol"), [ diff --git a/tests/test_registry_invariants.py b/tests/test_registry_invariants.py index 962d764ee..983c3e1b9 100644 --- a/tests/test_registry_invariants.py +++ b/tests/test_registry_invariants.py @@ -506,6 +506,7 @@ def test_provider_model_prefixes_unique_and_resolvable(): ("azure-foundry-anthropic/claude-opus-4-5", "azure-foundry-anthropic"), ("aws-bedrock/openai.gpt-oss-20b-1:0", "aws-bedrock"), ("github-models/openai/gpt-4.1-mini", "github-models"), + ("vercel/anthropic/claude-sonnet-4.5", "vercel"), ("zai/glm-5", "zai"), ("zai-coding/glm-5.4-flash", "zai-coding"), ("vllm/local-model", "vllm"), diff --git a/tests/test_usage_tracking.py b/tests/test_usage_tracking.py index aded746f4..083700308 100644 --- a/tests/test_usage_tracking.py +++ b/tests/test_usage_tracking.py @@ -9,7 +9,6 @@ async def test_daytona_required_usage_tracking_requires_sandbox_handle(): """Guards the LiteLLM sandbox-local path: required still fails closed.""" from benchflow.providers.runtime import ensure_litellm_runtime - from benchflow.usage_tracking import UsageTrackingConfig with pytest.raises(RuntimeError, match="sandbox-local LiteLLM"): await ensure_litellm_runtime( @@ -19,7 +18,6 @@ async def test_daytona_required_usage_tracking_requires_sandbox_handle(): runtime=None, environment="daytona", session_id="rollout-1", - usage_tracking=UsageTrackingConfig(mode="required"), ) @@ -28,7 +26,6 @@ async def test_daytona_usage_tracking_starts_sandbox_local_litellm(monkeypatch): """Daytona auto telemetry should use LiteLLM inside the agent sandbox.""" from benchflow.providers import litellm_runtime as runtime_mod from benchflow.providers.runtime import ensure_litellm_runtime - from benchflow.usage_tracking import UsageTrackingConfig class FakeSandboxLiteLLM: def __init__(self, sandbox, route): @@ -56,7 +53,6 @@ async def fake_start(**kwargs): runtime=None, environment="daytona", session_id="rollout-1", - usage_tracking=UsageTrackingConfig(mode="required"), sandbox=sandbox, )