From ae23ab6b43bbdec5b42aea4f136cab0eecb63c72 Mon Sep 17 00:00:00 2001 From: kywch Date: Sun, 30 Aug 2026 17:16:20 -0700 Subject: [PATCH 1/4] feat(providers): add Z.ai Coding Plan routing --- src/benchflow/agents/providers.py | 24 +++++++++++ src/benchflow/providers/litellm_config.py | 22 ++++++++-- tests/test_litellm_config.py | 52 +++++++++++++++++++++++ tests/test_providers.py | 11 +++++ tests/test_registry_invariants.py | 1 + 5 files changed, 107 insertions(+), 3 deletions(-) diff --git a/src/benchflow/agents/providers.py b/src/benchflow/agents/providers.py index 9a576cbbc..38f54c561 100644 --- a/src/benchflow/agents/providers.py +++ b/src/benchflow/agents/providers.py @@ -275,6 +275,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", "glm-5.1", "glm-5.2", "glm-5-turbo") + ], + ), "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..78ac3a5c2 100644 --- a/src/benchflow/providers/litellm_config.py +++ b/src/benchflow/providers/litellm_config.py @@ -291,7 +291,24 @@ 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: + registry_api_bases = set() + for endpoint_protocol in provider_cfg.all_endpoints: + try: + endpoint = resolve_base_url( + provider_cfg, env, protocol=endpoint_protocol + ).rstrip("/") + except KeyError: + continue + registry_api_bases.add(endpoint) + native_api_key = ( + (env.get(provider_cfg.auth_env) or "").strip() if provider_cfg.auth_env else "" + ) + registry_api_base = ( + explicit_api_base.rstrip("/") in registry_api_bases + and bool(native_api_key) + and explicit_api_key == native_api_key + ) + if explicit_api_base and explicit_api_key and not registry_api_base: api_base = explicit_api_base else: try: @@ -335,7 +352,7 @@ def _route_registered_provider( params["api_base"] = api_base api_key_ref = ( _env_ref("BENCHFLOW_PROVIDER_API_KEY") - if explicit_api_base and explicit_api_key + if explicit_api_base and explicit_api_key and not registry_api_base else _registered_api_key_ref(provider_cfg) ) if api_key_ref: @@ -344,7 +361,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/tests/test_litellm_config.py b/tests/test_litellm_config.py index b829c4021..333911024 100644 --- a/tests/test_litellm_config.py +++ b/tests/test_litellm_config.py @@ -120,6 +120,58 @@ def test_registered_provider_route_honors_explicit_generic_proxy_env(): assert route.required_env == ("BENCHFLOW_PROVIDER_API_KEY",) +@pytest.mark.parametrize( + ("model", "key", "derived_base", "expected_base", "expected_upstream"), + [ + ( + "zai-coding/glm-5.9", + "ZAI_API_KEY", + "https://api.z.ai/api/anthropic/", + "https://api.z.ai/api/coding/paas/v4", + "openai/glm-5.9", + ), + ( + "openrouter/qwen/qwen3.5-397b-a17b", + "OPENROUTER_API_KEY", + "https://openrouter.ai/api/v1", + "https://openrouter.ai/api/v1", + "openai/qwen/qwen3.5-397b-a17b", + ), + ], +) +def test_registered_provider_route_ignores_registry_derived_generic_proxy_env( + model, key, derived_base, expected_base, expected_upstream +): + route = resolve_litellm_route( + model, + { + key: "native-key", + "BENCHFLOW_PROVIDER_BASE_URL": derived_base, + "BENCHFLOW_PROVIDER_API_KEY": "native-key", + }, + ) + + assert route.litellm_params["api_base"] == expected_base + assert route.litellm_params["api_key"] == f"os.environ/{key}" + assert route.required_env == (key,) + assert route.upstream_model == expected_upstream + + +def test_registered_endpoint_with_generic_key_remains_explicit_override(): + route = resolve_litellm_route( + "openrouter/qwen/qwen3.5-397b-a17b", + { + "OPENROUTER_API_KEY": "native-key", + "BENCHFLOW_PROVIDER_BASE_URL": "https://openrouter.ai/api/v1", + "BENCHFLOW_PROVIDER_API_KEY": "generic-key", + }, + ) + + assert route.litellm_params["api_base"] == "https://openrouter.ai/api/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_providers.py b/tests/test_providers.py index 3d6c486e7..4c7ad1dcd 100644 --- a/tests/test_providers.py +++ b/tests/test_providers.py @@ -226,6 +226,17 @@ 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): + assert resolve_base_url(PROVIDERS["zai-coding"], {}, protocol) == expected + 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..5ff505b73 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.2", "zai-coding"), ("vllm/local-model", "vllm"), ("kimi/kimi-k2.6", "kimi"), ("qwen-dashscope/qwen3.6-max-preview", "qwen-dashscope"), From ee4a2fdc3ace0ada0d6af85d526f050708b18275 Mon Sep 17 00:00:00 2001 From: kywch Date: Sun, 30 Aug 2026 17:17:52 -0700 Subject: [PATCH 2/4] fix(litellm): forward validated generation parameters --- src/benchflow/providers/litellm_config.py | 38 +++++++++++++++- tests/test_litellm_config.py | 55 +++++++++++++++++++++++ 2 files changed, 92 insertions(+), 1 deletion(-) diff --git a/src/benchflow/providers/litellm_config.py b/src/benchflow/providers/litellm_config.py index 78ac3a5c2..848f6a5ab 100644 --- a/src/benchflow/providers/litellm_config.py +++ b/src/benchflow/providers/litellm_config.py @@ -4,6 +4,7 @@ import hashlib import json +import math import re from dataclasses import dataclass from urllib.parse import urlparse @@ -177,6 +178,38 @@ def _registered_api_key_ref(cfg: ProviderConfig) -> str | None: return None +def _apply_model_generation_params( + params: dict[str, str | int | float | bool | list[str]], env: dict[str, str] +) -> None: + for env_name, param, upper_bound in ( + ("BENCHFLOW_MODEL_TEMPERATURE", "temperature", None), + ("BENCHFLOW_MODEL_TOP_P", "top_p", 1.0), + ): + value = (env.get(env_name) or "").strip() + try: + parsed = float(value) + except ValueError: + continue + if ( + math.isfinite(parsed) + and parsed >= 0 + and (upper_bound is None or parsed <= upper_bound) + ): + params[param] = parsed + + try: + max_tokens = int((env.get("BENCHFLOW_MODEL_MAX_TOKENS") or "").strip()) + except ValueError: + return + if max_tokens > 0: + token_param = ( + "max_output_tokens" + if env.get("BENCHFLOW_PROVIDER_PROTOCOL") == "openai-responses" + else "max_tokens" + ) + params[token_param] = max_tokens + + def _provider_reasoning_effort(env: dict[str, str]) -> str | None: """Return an explicitly requested gateway-side reasoning effort.""" raw = ( @@ -376,12 +409,14 @@ def resolve_litellm_route(model: str, env: dict[str, str]) -> LiteLLMRoute: provider = find_provider(model) if provider is not None: provider_name, provider_cfg = provider - return _route_registered_provider( + route = _route_registered_provider( model=model, provider_name=provider_name, provider_cfg=provider_cfg, env=env, ) + _apply_model_generation_params(route.litellm_params, env) + return route lower = model.lower() bare = strip_provider_prefix(model) @@ -420,6 +455,7 @@ def resolve_litellm_route(model: str, env: dict[str, str]) -> LiteLLMRoute: key = required[0] if required else None if key and "api_key" not in params: params["api_key"] = _env_ref(key) + _apply_model_generation_params(params, env) return LiteLLMRoute( requested_model=model, model_alias=safe_model_alias(model), diff --git a/tests/test_litellm_config.py b/tests/test_litellm_config.py index 333911024..b42db3dcd 100644 --- a/tests/test_litellm_config.py +++ b/tests/test_litellm_config.py @@ -172,6 +172,61 @@ def test_registered_endpoint_with_generic_key_remains_explicit_override(): assert route.required_env == ("BENCHFLOW_PROVIDER_API_KEY",) +@pytest.mark.parametrize( + ("model", "key", "protocol", "token_param"), + [ + ("zai-coding/glm-5.2", "ZAI_API_KEY", "openai-responses", "max_output_tokens"), + ("gemini-3.5-flash", "GEMINI_API_KEY", "openai-completions", "max_tokens"), + ], +) +def test_litellm_route_generation_overrides(model, key, protocol, token_param): + params = resolve_litellm_route( + model, + { + key: "key", + "BENCHFLOW_PROVIDER_PROTOCOL": protocol, + "BENCHFLOW_MODEL_TEMPERATURE": "1.0", + "BENCHFLOW_MODEL_TOP_P": "0.95", + "BENCHFLOW_MODEL_MAX_TOKENS": "131072", + }, + ).litellm_params + expected = {"temperature": 1.0, "top_p": 0.95, token_param: 131072} + assert {name: params[name] for name in expected} == expected + assert ({"max_tokens", "max_output_tokens"} - {token_param}).isdisjoint(params) + + +@pytest.mark.parametrize( + ("env_name", "param", "value"), + [ + ("BENCHFLOW_MODEL_TEMPERATURE", "temperature", "nan"), + ("BENCHFLOW_MODEL_TEMPERATURE", "temperature", "inf"), + ("BENCHFLOW_MODEL_TEMPERATURE", "temperature", "-0.1"), + ("BENCHFLOW_MODEL_TOP_P", "top_p", "1.1"), + ("BENCHFLOW_MODEL_TOP_P", "top_p", "-0.1"), + ("BENCHFLOW_MODEL_MAX_TOKENS", "max_tokens", "0"), + ("BENCHFLOW_MODEL_MAX_TOKENS", "max_tokens", "-1"), + ("BENCHFLOW_MODEL_MAX_TOKENS", "max_tokens", "1.5"), + ], +) +def test_litellm_route_rejects_invalid_generation_overrides(env_name, param, value): + params = resolve_litellm_route( + "zai-coding/glm-5.2", + { + "ZAI_API_KEY": "key", + env_name: value, + }, + ).litellm_params + assert param not in params + + +def test_special_registered_provider_generation_overrides(): + route = resolve_litellm_route( + "aws-bedrock/us.anthropic.claude-opus-4-8", + {"BENCHFLOW_MODEL_MAX_TOKENS": "4096"}, + ) + assert route.litellm_params["max_tokens"] == 4096 + + @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.""" From e8a9b6ce819c8948809d3aa66172b9fe07b99262 Mon Sep 17 00:00:00 2001 From: kywch Date: Sun, 30 Aug 2026 19:17:13 -0700 Subject: [PATCH 3/4] fix(litellm): preserve explicit provider endpoints --- src/benchflow/agents/env.py | 10 ++++++ src/benchflow/agents/providers.py | 2 ++ src/benchflow/providers/litellm_config.py | 23 +++----------- tests/test_litellm_config.py | 37 +++++++++++++++-------- 4 files changed, 41 insertions(+), 31 deletions(-) diff --git a/src/benchflow/agents/env.py b/src/benchflow/agents/env.py index ebf83e15c..0e1c5a6b3 100644 --- a/src/benchflow/agents/env.py +++ b/src/benchflow/agents/env.py @@ -469,12 +469,17 @@ def resolve_provider_env( ) -> None: """Detect provider for model, inject BENCHFLOW_PROVIDER_* and env_mapping.""" from benchflow.agents.providers import ( + PROVIDER_ENV_SOURCE_ENV, find_provider, find_provider_for_bare_model, resolve_base_url, strip_provider_prefix, ) + provider_route_missing = not any( + key in agent_env + for key in ("BENCHFLOW_PROVIDER_BASE_URL", "BENCHFLOW_PROVIDER_API_KEY") + ) agent_env.setdefault("BENCHFLOW_PROVIDER_MODEL", strip_provider_prefix(model)) agent_cfg = AGENTS.get(agent) # Agent-declared protocol takes precedence over provider's primary so @@ -532,6 +537,11 @@ def resolve_provider_env( "BENCHFLOW_PROVIDER_API_KEY", _BEDROCK_PROVIDER_PLACEHOLDER_API_KEY, ) + if provider_route_missing and all( + agent_env.get(key) + for key in ("BENCHFLOW_PROVIDER_BASE_URL", "BENCHFLOW_PROVIDER_API_KEY") + ): + agent_env[PROVIDER_ENV_SOURCE_ENV] = "registry" else: # No registered provider prefix — bridge the model's well-known API key # to BENCHFLOW_PROVIDER_API_KEY so env_mapping can translate it to diff --git a/src/benchflow/agents/providers.py b/src/benchflow/agents/providers.py index 38f54c561..4277b2c9e 100644 --- a/src/benchflow/agents/providers.py +++ b/src/benchflow/agents/providers.py @@ -61,6 +61,8 @@ from dataclasses import dataclass, field +PROVIDER_ENV_SOURCE_ENV = "_BENCHFLOW_PROVIDER_ENV_SOURCE" + @dataclass class ProviderConfig: diff --git a/src/benchflow/providers/litellm_config.py b/src/benchflow/providers/litellm_config.py index 848f6a5ab..db6c1696b 100644 --- a/src/benchflow/providers/litellm_config.py +++ b/src/benchflow/providers/litellm_config.py @@ -10,6 +10,7 @@ from urllib.parse import urlparse from benchflow.agents.providers import ( + PROVIDER_ENV_SOURCE_ENV, ProviderConfig, find_provider, resolve_base_url, @@ -324,24 +325,8 @@ 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() - registry_api_bases = set() - for endpoint_protocol in provider_cfg.all_endpoints: - try: - endpoint = resolve_base_url( - provider_cfg, env, protocol=endpoint_protocol - ).rstrip("/") - except KeyError: - continue - registry_api_bases.add(endpoint) - native_api_key = ( - (env.get(provider_cfg.auth_env) or "").strip() if provider_cfg.auth_env else "" - ) - registry_api_base = ( - explicit_api_base.rstrip("/") in registry_api_bases - and bool(native_api_key) - and explicit_api_key == native_api_key - ) - if explicit_api_base and explicit_api_key and not registry_api_base: + provider_env_is_registry = env.get(PROVIDER_ENV_SOURCE_ENV) == "registry" + if explicit_api_base and explicit_api_key and not provider_env_is_registry: api_base = explicit_api_base else: try: @@ -385,7 +370,7 @@ def _route_registered_provider( params["api_base"] = api_base api_key_ref = ( _env_ref("BENCHFLOW_PROVIDER_API_KEY") - if explicit_api_base and explicit_api_key and not registry_api_base + if explicit_api_base and explicit_api_key and not provider_env_is_registry else _registered_api_key_ref(provider_cfg) ) if api_key_ref: diff --git a/tests/test_litellm_config.py b/tests/test_litellm_config.py index b42db3dcd..57e5d3714 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_provider_env from benchflow.providers.litellm_config import ( litellm_proxy_config, resolve_litellm_route, @@ -121,35 +122,30 @@ def test_registered_provider_route_honors_explicit_generic_proxy_env(): @pytest.mark.parametrize( - ("model", "key", "derived_base", "expected_base", "expected_upstream"), + ("model", "key", "agent", "expected_base", "expected_upstream"), [ ( "zai-coding/glm-5.9", "ZAI_API_KEY", - "https://api.z.ai/api/anthropic/", + "claude-agent-acp", "https://api.z.ai/api/coding/paas/v4", "openai/glm-5.9", ), ( "openrouter/qwen/qwen3.5-397b-a17b", "OPENROUTER_API_KEY", - "https://openrouter.ai/api/v1", + "pi-acp", "https://openrouter.ai/api/v1", "openai/qwen/qwen3.5-397b-a17b", ), ], ) def test_registered_provider_route_ignores_registry_derived_generic_proxy_env( - model, key, derived_base, expected_base, expected_upstream + model, key, agent, expected_base, expected_upstream ): - route = resolve_litellm_route( - model, - { - key: "native-key", - "BENCHFLOW_PROVIDER_BASE_URL": derived_base, - "BENCHFLOW_PROVIDER_API_KEY": "native-key", - }, - ) + env = {key: "native-key"} + resolve_provider_env(env, model, agent) + route = resolve_litellm_route(model, env) assert route.litellm_params["api_base"] == expected_base assert route.litellm_params["api_key"] == f"os.environ/{key}" @@ -157,6 +153,23 @@ def test_registered_provider_route_ignores_registry_derived_generic_proxy_env( assert route.upstream_model == expected_upstream +@pytest.mark.parametrize("resolve_first", [False, True], ids=["direct", "resolved"]) +def test_registered_provider_route_preserves_explicit_endpoint_with_native_key( + resolve_first, +): + env = { + "ZAI_API_KEY": "same-key", + "BENCHFLOW_PROVIDER_BASE_URL": "https://api.z.ai/api/anthropic", + "BENCHFLOW_PROVIDER_API_KEY": "same-key", + } + if resolve_first: + resolve_provider_env(env, "zai-coding/glm-5.2", "claude-agent-acp") + route = resolve_litellm_route("zai-coding/glm-5.2", env) + + assert route.litellm_params["api_base"] == "https://api.z.ai/api/anthropic" + assert route.required_env == ("BENCHFLOW_PROVIDER_API_KEY",) + + def test_registered_endpoint_with_generic_key_remains_explicit_override(): route = resolve_litellm_route( "openrouter/qwen/qwen3.5-397b-a17b", From d4708464d1c814c7b99f87349c1083cf2dc5162b Mon Sep 17 00:00:00 2001 From: kywch Date: Tue, 1 Sep 2026 20:33:47 -0700 Subject: [PATCH 4/4] fix(providers): harden Z.ai Coding Plan routing --- src/benchflow/agents/env.py | 13 +- src/benchflow/agents/providers.py | 4 +- src/benchflow/providers/litellm_config.py | 53 ++------ src/benchflow/providers/litellm_runtime.py | 16 +++ src/benchflow/providers/runtime.py | 6 +- tests/test_litellm_config.py | 140 ++++++--------------- tests/test_litellm_runtime.py | 58 +++++++++ tests/test_providers.py | 8 ++ tests/test_registry_invariants.py | 2 +- 9 files changed, 140 insertions(+), 160 deletions(-) diff --git a/src/benchflow/agents/env.py b/src/benchflow/agents/env.py index 0e1c5a6b3..ffa5e5f02 100644 --- a/src/benchflow/agents/env.py +++ b/src/benchflow/agents/env.py @@ -469,17 +469,13 @@ def resolve_provider_env( ) -> None: """Detect provider for model, inject BENCHFLOW_PROVIDER_* and env_mapping.""" from benchflow.agents.providers import ( - PROVIDER_ENV_SOURCE_ENV, + ZAI_CODING_REGISTRY_BASE_ENV, find_provider, find_provider_for_bare_model, resolve_base_url, strip_provider_prefix, ) - provider_route_missing = not any( - key in agent_env - for key in ("BENCHFLOW_PROVIDER_BASE_URL", "BENCHFLOW_PROVIDER_API_KEY") - ) agent_env.setdefault("BENCHFLOW_PROVIDER_MODEL", strip_provider_prefix(model)) agent_cfg = AGENTS.get(agent) # Agent-declared protocol takes precedence over provider's primary so @@ -520,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, @@ -537,11 +535,6 @@ def resolve_provider_env( "BENCHFLOW_PROVIDER_API_KEY", _BEDROCK_PROVIDER_PLACEHOLDER_API_KEY, ) - if provider_route_missing and all( - agent_env.get(key) - for key in ("BENCHFLOW_PROVIDER_BASE_URL", "BENCHFLOW_PROVIDER_API_KEY") - ): - agent_env[PROVIDER_ENV_SOURCE_ENV] = "registry" else: # No registered provider prefix — bridge the model's well-known API key # to BENCHFLOW_PROVIDER_API_KEY so env_mapping can translate it to diff --git a/src/benchflow/agents/providers.py b/src/benchflow/agents/providers.py index 4277b2c9e..d7feeedc4 100644 --- a/src/benchflow/agents/providers.py +++ b/src/benchflow/agents/providers.py @@ -61,7 +61,7 @@ from dataclasses import dataclass, field -PROVIDER_ENV_SOURCE_ENV = "_BENCHFLOW_PROVIDER_ENV_SOURCE" +ZAI_CODING_REGISTRY_BASE_ENV = "_BENCHFLOW_ZAI_CODING_REGISTRY_BASE" @dataclass @@ -298,7 +298,7 @@ def all_endpoints(self) -> dict[str, str]: "contextWindow": 200000, "maxTokens": 131072, } - for model in ("glm-5", "glm-5.1", "glm-5.2", "glm-5-turbo") + for model in ("glm-5.3", "glm-5.3-flash") ], ), "kimi": ProviderConfig( diff --git a/src/benchflow/providers/litellm_config.py b/src/benchflow/providers/litellm_config.py index db6c1696b..127df0caa 100644 --- a/src/benchflow/providers/litellm_config.py +++ b/src/benchflow/providers/litellm_config.py @@ -4,13 +4,12 @@ import hashlib import json -import math import re from dataclasses import dataclass from urllib.parse import urlparse from benchflow.agents.providers import ( - PROVIDER_ENV_SOURCE_ENV, + ZAI_CODING_REGISTRY_BASE_ENV, ProviderConfig, find_provider, resolve_base_url, @@ -179,38 +178,6 @@ def _registered_api_key_ref(cfg: ProviderConfig) -> str | None: return None -def _apply_model_generation_params( - params: dict[str, str | int | float | bool | list[str]], env: dict[str, str] -) -> None: - for env_name, param, upper_bound in ( - ("BENCHFLOW_MODEL_TEMPERATURE", "temperature", None), - ("BENCHFLOW_MODEL_TOP_P", "top_p", 1.0), - ): - value = (env.get(env_name) or "").strip() - try: - parsed = float(value) - except ValueError: - continue - if ( - math.isfinite(parsed) - and parsed >= 0 - and (upper_bound is None or parsed <= upper_bound) - ): - params[param] = parsed - - try: - max_tokens = int((env.get("BENCHFLOW_MODEL_MAX_TOKENS") or "").strip()) - except ValueError: - return - if max_tokens > 0: - token_param = ( - "max_output_tokens" - if env.get("BENCHFLOW_PROVIDER_PROTOCOL") == "openai-responses" - else "max_tokens" - ) - params[token_param] = max_tokens - - def _provider_reasoning_effort(env: dict[str, str]) -> str | None: """Return an explicitly requested gateway-side reasoning effort.""" raw = ( @@ -325,8 +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() - provider_env_is_registry = env.get(PROVIDER_ENV_SOURCE_ENV) == "registry" - if explicit_api_base and explicit_api_key and not provider_env_is_registry: + 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: @@ -368,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 and not provider_env_is_registry + if explicit_route or explicit_zai_key else _registered_api_key_ref(provider_cfg) ) if api_key_ref: @@ -394,14 +368,12 @@ def resolve_litellm_route(model: str, env: dict[str, str]) -> LiteLLMRoute: provider = find_provider(model) if provider is not None: provider_name, provider_cfg = provider - route = _route_registered_provider( + return _route_registered_provider( model=model, provider_name=provider_name, provider_cfg=provider_cfg, env=env, ) - _apply_model_generation_params(route.litellm_params, env) - return route lower = model.lower() bare = strip_provider_prefix(model) @@ -440,7 +412,6 @@ def resolve_litellm_route(model: str, env: dict[str, str]) -> LiteLLMRoute: key = required[0] if required else None if key and "api_key" not in params: params["api_key"] = _env_ref(key) - _apply_model_generation_params(params, 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 fc53f1186..af7c57924 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -1596,6 +1596,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): @@ -1719,6 +1732,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 57e5d3714..b56b1b0a7 100644 --- a/tests/test_litellm_config.py +++ b/tests/test_litellm_config.py @@ -2,7 +2,7 @@ import pytest -from benchflow.agents.env import resolve_provider_env +from benchflow.agents.env import resolve_agent_env, resolve_provider_env from benchflow.providers.litellm_config import ( litellm_proxy_config, resolve_litellm_route, @@ -122,124 +122,54 @@ def test_registered_provider_route_honors_explicit_generic_proxy_env(): @pytest.mark.parametrize( - ("model", "key", "agent", "expected_base", "expected_upstream"), + ("agent", "agent_base"), [ - ( - "zai-coding/glm-5.9", - "ZAI_API_KEY", - "claude-agent-acp", - "https://api.z.ai/api/coding/paas/v4", - "openai/glm-5.9", - ), - ( - "openrouter/qwen/qwen3.5-397b-a17b", - "OPENROUTER_API_KEY", - "pi-acp", - "https://openrouter.ai/api/v1", - "openai/qwen/qwen3.5-397b-a17b", - ), + ("claude-agent-acp", "https://api.z.ai/api/anthropic"), + ("openclaw", "https://api.z.ai/api/coding/paas/v4"), ], ) -def test_registered_provider_route_ignores_registry_derived_generic_proxy_env( - model, key, agent, expected_base, expected_upstream -): - env = {key: "native-key"} - resolve_provider_env(env, model, agent) - route = resolve_litellm_route(model, env) - - assert route.litellm_params["api_base"] == expected_base - assert route.litellm_params["api_key"] == f"os.environ/{key}" - assert route.required_env == (key,) - assert route.upstream_model == expected_upstream - - -@pytest.mark.parametrize("resolve_first", [False, True], ids=["direct", "resolved"]) -def test_registered_provider_route_preserves_explicit_endpoint_with_native_key( - resolve_first, -): - env = { - "ZAI_API_KEY": "same-key", - "BENCHFLOW_PROVIDER_BASE_URL": "https://api.z.ai/api/anthropic", - "BENCHFLOW_PROVIDER_API_KEY": "same-key", - } - if resolve_first: - resolve_provider_env(env, "zai-coding/glm-5.2", "claude-agent-acp") - route = resolve_litellm_route("zai-coding/glm-5.2", env) - - assert route.litellm_params["api_base"] == "https://api.z.ai/api/anthropic" +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_registered_endpoint_with_generic_key_remains_explicit_override(): +def test_zai_coding_preserves_explicit_proxy_route(): + """Guards PR #1074: explicit Z.AI-compatible proxies remain authoritative.""" route = resolve_litellm_route( - "openrouter/qwen/qwen3.5-397b-a17b", + "zai-coding/glm-5.3-flash", { - "OPENROUTER_API_KEY": "native-key", - "BENCHFLOW_PROVIDER_BASE_URL": "https://openrouter.ai/api/v1", - "BENCHFLOW_PROVIDER_API_KEY": "generic-key", + "BENCHFLOW_PROVIDER_BASE_URL": "https://proxy.example.test/v1", + "BENCHFLOW_PROVIDER_API_KEY": "proxy-key", }, ) - assert route.litellm_params["api_base"] == "https://openrouter.ai/api/v1" - assert route.litellm_params["api_key"] == "os.environ/BENCHFLOW_PROVIDER_API_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", "key", "protocol", "token_param"), - [ - ("zai-coding/glm-5.2", "ZAI_API_KEY", "openai-responses", "max_output_tokens"), - ("gemini-3.5-flash", "GEMINI_API_KEY", "openai-completions", "max_tokens"), - ], -) -def test_litellm_route_generation_overrides(model, key, protocol, token_param): - params = resolve_litellm_route( - model, - { - key: "key", - "BENCHFLOW_PROVIDER_PROTOCOL": protocol, - "BENCHFLOW_MODEL_TEMPERATURE": "1.0", - "BENCHFLOW_MODEL_TOP_P": "0.95", - "BENCHFLOW_MODEL_MAX_TOKENS": "131072", - }, - ).litellm_params - expected = {"temperature": 1.0, "top_p": 0.95, token_param: 131072} - assert {name: params[name] for name in expected} == expected - assert ({"max_tokens", "max_output_tokens"} - {token_param}).isdisjoint(params) - - -@pytest.mark.parametrize( - ("env_name", "param", "value"), - [ - ("BENCHFLOW_MODEL_TEMPERATURE", "temperature", "nan"), - ("BENCHFLOW_MODEL_TEMPERATURE", "temperature", "inf"), - ("BENCHFLOW_MODEL_TEMPERATURE", "temperature", "-0.1"), - ("BENCHFLOW_MODEL_TOP_P", "top_p", "1.1"), - ("BENCHFLOW_MODEL_TOP_P", "top_p", "-0.1"), - ("BENCHFLOW_MODEL_MAX_TOKENS", "max_tokens", "0"), - ("BENCHFLOW_MODEL_MAX_TOKENS", "max_tokens", "-1"), - ("BENCHFLOW_MODEL_MAX_TOKENS", "max_tokens", "1.5"), - ], -) -def test_litellm_route_rejects_invalid_generation_overrides(env_name, param, value): - params = resolve_litellm_route( - "zai-coding/glm-5.2", - { - "ZAI_API_KEY": "key", - env_name: value, - }, - ).litellm_params - assert param not in params - - -def test_special_registered_provider_generation_overrides(): - route = resolve_litellm_route( - "aws-bedrock/us.anthropic.claude-opus-4-8", - {"BENCHFLOW_MODEL_MAX_TOKENS": "4096"}, - ) - assert route.litellm_params["max_tokens"] == 4096 - - @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 21f934b77..f6d71f8c1 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 4c7ad1dcd..387da1264 100644 --- a/tests/test_providers.py +++ b/tests/test_providers.py @@ -235,8 +235,16 @@ def test_protocol_selects_endpoint(self): ], ) 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 5ff505b73..efeb025be 100644 --- a/tests/test_registry_invariants.py +++ b/tests/test_registry_invariants.py @@ -486,7 +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.2", "zai-coding"), + ("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"),