fix(proxy): forward model_info in /health/test_connection

test_model_connection hardcoded model_info={} when calling
_update_litellm_params_for_health_check, silently dropping flags like
health_check_supports_max_tokens: false and always injecting max_tokens: 5
for chat mode. The background /health endpoint already honored this flag;
only the UI Test Connection button was broken.

When the router resolves a deployment, use its server-side model_info so
routing-affecting fields (health_check_model, mode, health_check_voice)
cannot be overridden by the request body after authorization. When no
deployment is resolved (ad-hoc credentials), fall back to the caller-supplied
model_info.

Fixes #29266
This commit is contained in:
Sebin P Johnson 2026-05-30 09:35:04 +05:30
parent af17400c38
commit 4e286a4b8b
2 changed files with 73 additions and 1 deletions

View file

@ -1771,6 +1771,7 @@ async def test_model_connection(
# Look up model configuration from router if model name is provided
# This gets the litellm_params from proxy config (with resolved env vars)
config_litellm_params: dict = {}
resolved_model_info: Optional[dict] = None
if llm_router is not None:
# Prefer disambiguation by deployment id (`model_info.id`) when
# the caller supplies it. This is required when multiple
@ -1793,6 +1794,9 @@ async def test_model_connection(
config_litellm_params = deployment_by_id.litellm_params.model_dump(
exclude_none=True
)
resolved_model_info = deployment_by_id.model_info.model_dump(
exclude_none=True
)
elif model_name:
# Fall back to model_name lookup for callers (e.g. the
# "Add Model" wizard, or curl) that don't supply an id.
@ -1819,6 +1823,9 @@ async def test_model_connection(
config_litellm_params = dict(
deployments[0].get("litellm_params", {})
)
resolved_model_info = dict(
deployments[0].get("model_info", {}) or {}
)
except Exception as e:
verbose_proxy_logger.debug(
f"Could not find model {model_name} in router: {e}. "
@ -1842,7 +1849,7 @@ async def test_model_connection(
)
# Include health_check_params if provided
litellm_params = _update_litellm_params_for_health_check(
model_info={},
model_info=resolved_model_info if resolved_model_info is not None else (model_info or {}),
litellm_params=litellm_params,
)
mode = mode or litellm_params.pop("mode", None)

View file

@ -696,6 +696,71 @@ async def test_test_model_connection_falls_back_to_deployments_zero_without_id()
assert model_params.get("api_key") == "fake-key-A"
@pytest.mark.asyncio
async def test_test_model_connection_respects_health_check_supports_max_tokens():
"""
Regression test for https://github.com/BerriAI/litellm/issues/29266.
When model_info.health_check_supports_max_tokens is False, the outgoing
request must not include max_tokens. Previously model_info={} was hardcoded
in the call to _update_litellm_params_for_health_check, so the flag was
silently dropped and max_tokens: 5 was always injected for chat mode.
"""
mock_request = MagicMock()
mock_user_api_key_dict = MagicMock()
mock_user_api_key_dict.user_id = "test-user"
mock_user_api_key_dict.token = "test-token"
mock_prisma_client = MagicMock()
mock_router = MagicMock()
mock_router.get_deployment.return_value = None
mock_router.get_model_list.return_value = []
mock_can_user_make_model_call = AsyncMock()
mock_ahealth_check = AsyncMock(return_value={"status": "healthy"})
mock_run_with_timeout = AsyncMock(return_value={"status": "healthy"})
def mock_reject_os_environ(params):
return None
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
patch("litellm.proxy.proxy_server.llm_router", mock_router),
patch("litellm.proxy.proxy_server.premium_user", False),
patch(
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
mock_can_user_make_model_call,
),
patch(
"litellm.proxy.health_endpoints._health_endpoints.litellm.ahealth_check",
mock_ahealth_check,
),
patch(
"litellm.proxy.health_endpoints._health_endpoints.run_with_timeout",
mock_run_with_timeout,
),
patch(
"litellm.proxy.health_endpoints._health_endpoints._reject_os_environ_references",
mock_reject_os_environ,
),
):
await health_test_model_connection(
request=mock_request,
mode="chat",
litellm_params={
"model": "azure/gpt-4o",
"api_key": "fake-key",
"api_base": "https://fake.invalid",
},
model_info={"health_check_supports_max_tokens": False},
user_api_key_dict=mock_user_api_key_dict,
)
model_params = mock_ahealth_check.call_args.kwargs.get("model_params", {})
assert "max_tokens" not in model_params, (
"max_tokens must not be sent when health_check_supports_max_tokens is False; "
f"got model_params={model_params}"
)
@pytest.mark.asyncio
async def test_health_services_endpoint_datadog_llm_observability():
"""