diff --git a/src/benchflow/agents/env.py b/src/benchflow/agents/env.py index ebf83e15c..ffa5e5f02 100644 --- a/src/benchflow/agents/env.py +++ b/src/benchflow/agents/env.py @@ -469,6 +469,7 @@ def resolve_provider_env( ) -> None: """Detect provider for model, inject BENCHFLOW_PROVIDER_* and env_mapping.""" from benchflow.agents.providers import ( + ZAI_CODING_REGISTRY_BASE_ENV, find_provider, find_provider_for_bare_model, resolve_base_url, @@ -515,6 +516,8 @@ def resolve_provider_env( "BENCHFLOW_PROVIDER_BASE_URL", base_url, ) + if _prov_name == "zai-coding" and base_url: + agent_env[ZAI_CODING_REGISTRY_BASE_ENV] = "1" agent_env.setdefault( "BENCHFLOW_PROVIDER_PROTOCOL", agent_protocol or _prov_cfg.api_protocol, diff --git a/src/benchflow/agents/providers.py b/src/benchflow/agents/providers.py index 9a576cbbc..d7feeedc4 100644 --- a/src/benchflow/agents/providers.py +++ b/src/benchflow/agents/providers.py @@ -61,6 +61,8 @@ from dataclasses import dataclass, field +ZAI_CODING_REGISTRY_BASE_ENV = "_BENCHFLOW_ZAI_CODING_REGISTRY_BASE" + @dataclass class ProviderConfig: @@ -275,6 +277,30 @@ def all_endpoints(self) -> dict[str, str]: }, ], ), + "zai-coding": ProviderConfig( + name="zai-coding", + base_url="https://api.z.ai/api/coding/paas/v4", + api_protocol="openai-completions", + auth_type="api_key", + auth_env="ZAI_API_KEY", + endpoints={ + "openai-completions": "https://api.z.ai/api/coding/paas/v4", + "openai-responses": "https://api.z.ai/api/coding/paas/v4", + "anthropic-messages": "https://api.z.ai/api/anthropic", + }, + models=[ + { + "id": model, + "name": model.upper(), + "reasoning": True, + "input": ["text"], + "cost": {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0}, + "contextWindow": 200000, + "maxTokens": 131072, + } + for model in ("glm-5.3", "glm-5.3-flash") + ], + ), "kimi": ProviderConfig( name="kimi", base_url="{base_url}", diff --git a/src/benchflow/providers/litellm_config.py b/src/benchflow/providers/litellm_config.py index f703b32d2..127df0caa 100644 --- a/src/benchflow/providers/litellm_config.py +++ b/src/benchflow/providers/litellm_config.py @@ -9,6 +9,7 @@ from urllib.parse import urlparse from benchflow.agents.providers import ( + ZAI_CODING_REGISTRY_BASE_ENV, ProviderConfig, find_provider, resolve_base_url, @@ -291,7 +292,11 @@ def _route_registered_provider( ) explicit_api_base = (env.get("BENCHFLOW_PROVIDER_BASE_URL") or "").strip() explicit_api_key = (env.get("BENCHFLOW_PROVIDER_API_KEY") or "").strip() - if explicit_api_base and explicit_api_key: + zai_registry_base = ( + provider_name == "zai-coding" and env.get(ZAI_CODING_REGISTRY_BASE_ENV) == "1" + ) + explicit_route = explicit_api_base and explicit_api_key and not zai_registry_base + if explicit_api_base and not zai_registry_base: api_base = explicit_api_base else: try: @@ -333,9 +338,13 @@ def _route_registered_provider( params: dict[str, str | int | float | bool | list[str]] = {"model": upstream} if api_base: params["api_base"] = api_base + native_key = (env.get(provider_cfg.auth_env or "") or "").strip() + explicit_zai_key = bool( + zai_registry_base and explicit_api_key and explicit_api_key != native_key + ) api_key_ref = ( _env_ref("BENCHFLOW_PROVIDER_API_KEY") - if explicit_api_base and explicit_api_key + if explicit_route or explicit_zai_key else _registered_api_key_ref(provider_cfg) ) if api_key_ref: @@ -344,7 +353,6 @@ def _route_registered_provider( required_env.append("BENCHFLOW_PROVIDER_API_KEY") elif provider_cfg.auth_env: required_env.append(provider_cfg.auth_env) - return LiteLLMRoute( requested_model=model, model_alias=safe_model_alias(model), diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index f93d74e94..9c0619a20 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -1622,6 +1622,19 @@ async def ensure_litellm_runtime( routed through it: ``oracle`` (no model) and native-subscription auth (no API key to proxy). 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. + if ( + runtime is not None + and getattr(runtime, "kind", None) == "litellm" + and getattr(runtime, "source_agent", None) == agent + and getattr(runtime, "source_model", None) == model + and agent_env.get(LITELLM_MASTER_KEY_ENV) + == getattr(runtime, "master_key", None) + and getattr(runtime, "source_env", None) is not None + ): + agent_env = dict(runtime.source_env) + usage_cfg = UsageTrackingConfig.coerce(usage_tracking).with_env_defaults() if uses_native_subscription_auth(agent, model, agent_env): @@ -1745,6 +1758,9 @@ async def ensure_litellm_runtime( server=server, config_key=config_key, master_key=master_key, + source_agent=agent, + source_model=model, + source_env=dict(agent_env), ) if live_trajectory_path is not None: server.start_live_capture(live_trajectory_path) diff --git a/src/benchflow/providers/runtime.py b/src/benchflow/providers/runtime.py index e35d24b41..55ab8591a 100644 --- a/src/benchflow/providers/runtime.py +++ b/src/benchflow/providers/runtime.py @@ -7,7 +7,7 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any @@ -31,6 +31,10 @@ class ProviderRuntime: server: LiteLLMProcess | None = None config_key: str | None = None master_key: str | None = None + source_agent: str | None = None + source_model: str | None = None + # Controller-only upstream config for re-entrant connects and proxy restarts. + source_env: dict[str, str] | None = field(default=None, repr=False) @property def base_url(self) -> str: diff --git a/tests/test_litellm_config.py b/tests/test_litellm_config.py index b829c4021..b56b1b0a7 100644 --- a/tests/test_litellm_config.py +++ b/tests/test_litellm_config.py @@ -2,6 +2,7 @@ import pytest +from benchflow.agents.env import resolve_agent_env, resolve_provider_env from benchflow.providers.litellm_config import ( litellm_proxy_config, resolve_litellm_route, @@ -120,6 +121,55 @@ def test_registered_provider_route_honors_explicit_generic_proxy_env(): assert route.required_env == ("BENCHFLOW_PROVIDER_API_KEY",) +@pytest.mark.parametrize( + ("agent", "agent_base"), + [ + ("claude-agent-acp", "https://api.z.ai/api/anthropic"), + ("openclaw", "https://api.z.ai/api/coding/paas/v4"), + ], +) +def test_zai_coding_clawsbench_routes(agent, agent_base): + """Guards PR #1074: ClawsBench agents use each supported Z.AI surface.""" + env = resolve_agent_env(agent, "zai-coding/glm-5.3", {"ZAI_API_KEY": "native-key"}) + route = resolve_litellm_route("zai-coding/glm-5.3", env) + + assert env["BENCHFLOW_PROVIDER_BASE_URL"] == agent_base + assert route.litellm_params["api_base"] == ("https://api.z.ai/api/coding/paas/v4") + assert route.litellm_params["api_key"] == "os.environ/ZAI_API_KEY" + assert route.required_env == ("ZAI_API_KEY",) + assert route.upstream_model == "openai/glm-5.3" + + +@pytest.mark.parametrize("model", ["glm-5.4", "glm-5.4-flash"]) +def test_zai_coding_registry_base_preserves_explicit_generic_key(model): + """Guards PR #1074: mixed provenance must not retain Anthropic upstream URL.""" + env = {"BENCHFLOW_PROVIDER_API_KEY": "generic-key"} + model_id = f"zai-coding/{model}" + resolve_provider_env(env, model_id, "claude-agent-acp") + route = resolve_litellm_route(model_id, env) + + assert env["BENCHFLOW_PROVIDER_BASE_URL"] == "https://api.z.ai/api/anthropic" + assert route.litellm_params["api_base"] == ("https://api.z.ai/api/coding/paas/v4") + assert route.litellm_params["api_key"] == ("os.environ/BENCHFLOW_PROVIDER_API_KEY") + assert route.required_env == ("BENCHFLOW_PROVIDER_API_KEY",) + assert route.upstream_model == f"openai/{model}" + + +def test_zai_coding_preserves_explicit_proxy_route(): + """Guards PR #1074: explicit Z.AI-compatible proxies remain authoritative.""" + route = resolve_litellm_route( + "zai-coding/glm-5.3-flash", + { + "BENCHFLOW_PROVIDER_BASE_URL": "https://proxy.example.test/v1", + "BENCHFLOW_PROVIDER_API_KEY": "proxy-key", + }, + ) + + assert route.litellm_params["api_base"] == "https://proxy.example.test/v1" + assert route.litellm_params["api_key"] == ("os.environ/BENCHFLOW_PROVIDER_API_KEY") + assert route.required_env == ("BENCHFLOW_PROVIDER_API_KEY",) + + @pytest.mark.parametrize("model", ["gemini/gemini-2.5-flash", "gemini-2.5-flash"]) def test_gemini_native_route_honors_explicit_base_url(model): """Guards the fix from PR #881 for issue #672.""" diff --git a/tests/test_litellm_runtime.py b/tests/test_litellm_runtime.py index 23efb3d61..981dae632 100644 --- a/tests/test_litellm_runtime.py +++ b/tests/test_litellm_runtime.py @@ -6,6 +6,7 @@ import pytest from benchflow.agents.codex_config import CODEX_DEFAULT_AUTH_REQUEST_ENV +from benchflow.agents.env import resolve_agent_env from benchflow.providers import litellm_runtime as runtime_mod from benchflow.providers.litellm_bedrock_preflight import BedrockPatchPreflightError from benchflow.providers.litellm_config import LITELLM_MODEL_ALIAS_ENV @@ -310,6 +311,63 @@ async def fake_start(**kwargs): assert created[0].stopped is True +@pytest.mark.parametrize("agent", ["claude-agent-acp", "openclaw"]) +@pytest.mark.asyncio +async def test_zai_runtime_reconnect_preserves_upstream_route(monkeypatch, agent): + """Guards PR #1074: reconnects retain Z.AI upstream routing and auth.""" + starts = [] + + async def fake_start(**kwargs): + starts.append(kwargs) + return FakeLiteLLMServer("http://127.0.0.1:4000", kwargs["route"]) + + monkeypatch.setattr(runtime_mod, "_start_host_litellm", fake_start) + env = resolve_agent_env(agent, "zai-coding/glm-5.3", {"ZAI_API_KEY": "native-key"}) + updated, first = await ensure_litellm_runtime( + agent=agent, + agent_env=env, + model="zai-coding/glm-5.3", + runtime=None, + environment="local", + session_id="run-1", + ) + _updated, second = await ensure_litellm_runtime( + agent=agent, + agent_env=updated, + model="zai-coding/glm-5.3", + runtime=first, + environment="local", + session_id="run-1", + ) + + assert first is not None + assert second is first + assert len(starts) == 1 + expected_params = { + "model": "openai/glm-5.3", + "api_base": "https://api.z.ai/api/coding/paas/v4", + "api_key": "os.environ/ZAI_API_KEY", + } + assert starts[0]["agent_env"]["ZAI_API_KEY"] == "native-key" + assert starts[0]["route"].litellm_params == expected_params + + assert first.server is not None + await first.server.stop() + _updated, third = await ensure_litellm_runtime( + agent=agent, + agent_env=updated, + model="zai-coding/glm-5.3", + runtime=first, + environment="local", + session_id="run-1", + ) + + assert third is not first + assert len(starts) == 2 + assert starts[1]["agent_env"]["ZAI_API_KEY"] == "native-key" + assert starts[1]["route"].litellm_params == expected_params + + @pytest.mark.asyncio async def test_required_usage_fails_when_litellm_lacks_provider_key(monkeypatch): monkeypatch.setattr(runtime_mod, "uses_native_subscription_auth", lambda *_: False) diff --git a/tests/test_providers.py b/tests/test_providers.py index 3d6c486e7..387da1264 100644 --- a/tests/test_providers.py +++ b/tests/test_providers.py @@ -226,6 +226,25 @@ def test_protocol_selects_endpoint(self): == "https://api.z.ai/api/paas/v4" ) + @pytest.mark.parametrize( + ("protocol", "expected"), + [ + ("openai-completions", "https://api.z.ai/api/coding/paas/v4"), + ("openai-responses", "https://api.z.ai/api/coding/paas/v4"), + ("anthropic-messages", "https://api.z.ai/api/anthropic"), + ], + ) + def test_zai_coding_protocol_selects_endpoint(self, protocol, expected): + """Guards PR #1074: Coding Plan supports agent-specific API surfaces.""" + assert resolve_base_url(PROVIDERS["zai-coding"], {}, protocol) == expected + + def test_zai_coding_advertises_current_glm5_models(self): + """Guards PR #1074: advertise current Coding Plan models, not GLM-4.x.""" + assert [model["id"] for model in PROVIDERS["zai-coding"].models] == [ + "glm-5.3", + "glm-5.3-flash", + ] + def test_protocol_fallback_to_base_url(self): """Unknown protocol falls back to primary base_url.""" p = PROVIDERS["zai"] diff --git a/tests/test_registry_invariants.py b/tests/test_registry_invariants.py index a6dc7d03f..efeb025be 100644 --- a/tests/test_registry_invariants.py +++ b/tests/test_registry_invariants.py @@ -486,6 +486,7 @@ def test_provider_model_prefixes_unique_and_resolvable(): ("aws-bedrock/openai.gpt-oss-20b-1:0", "aws-bedrock"), ("github-models/openai/gpt-4.1-mini", "github-models"), ("zai/glm-5", "zai"), + ("zai-coding/glm-5.4-flash", "zai-coding"), ("vllm/local-model", "vllm"), ("kimi/kimi-k2.6", "kimi"), ("qwen-dashscope/qwen3.6-max-preview", "qwen-dashscope"),