From 2e5e78c35df00e48d457a332caa10788c38f2efd Mon Sep 17 00:00:00 2001
From: Sebin P Johnson <8169523+sebinpj@users.noreply.github.com>
Date: Sat, 30 May 2026 09:46:23 +0530
Subject: [PATCH] fix black formatting; add router-resolved path test

Black reformats the long ternary in _update_litellm_params_for_health_check
call to a multi-line form. Also add a second regression test covering the
deployment-resolved-by-id path where server-side model_info carries
health_check_supports_max_tokens: False.
---
 .../health_endpoints/_health_endpoints.py     |  6 +-
 .../health_endpoints/test_health_endpoints.py | 72 +++++++++++++++++++
 2 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py
index 14f273d3a51..ede9dbbb11e 100644
--- a/litellm/proxy/health_endpoints/_health_endpoints.py
+++ b/litellm/proxy/health_endpoints/_health_endpoints.py
@@ -1849,7 +1849,11 @@ async def test_model_connection(
         )
         # Include health_check_params if provided
         litellm_params = _update_litellm_params_for_health_check(
-            model_info=resolved_model_info if resolved_model_info is not None else (model_info or {}),
+            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)
diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py
index 0ba1f7f3a96..ba6477de5d0 100644
--- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py
+++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py
@@ -761,6 +761,78 @@ async def test_test_model_connection_respects_health_check_supports_max_tokens()
     )
 
 
+@pytest.mark.asyncio
+async def test_test_model_connection_uses_server_model_info_when_deployment_resolved():
+    """
+    When the router resolves a deployment by id, the server-side model_info
+    must be forwarded to the health-check helper, not the caller-supplied one.
+
+    This covers the router-resolved path for the same bug as
+    test_test_model_connection_respects_health_check_supports_max_tokens.
+    """
+    from litellm.types.router import Deployment, LiteLLM_Params
+
+    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_can_user_make_model_call = AsyncMock()
+    mock_ahealth_check = AsyncMock(return_value={"status": "healthy"})
+    mock_run_with_timeout = AsyncMock(return_value={"status": "healthy"})
+
+    server_deployment = Deployment(
+        model_name="my-model",
+        litellm_params=LiteLLM_Params(
+            model="azure/gpt-4o",
+            api_key="server-key",
+            api_base="https://server.invalid/v1",
+        ),
+        model_info={"id": "dep-1", "health_check_supports_max_tokens": False},
+    )
+    mock_router = MagicMock()
+    mock_router.get_deployment.return_value = server_deployment
+
+    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"},
+            model_info={"id": "dep-1"},
+            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 server model_info has "
+        "health_check_supports_max_tokens: False; "
+        f"got model_params={model_params}"
+    )
+
+
 @pytest.mark.asyncio
 async def test_health_services_endpoint_datadog_llm_observability():
     """