From bf4069a80c46a3da804e1d078f9052e9b8ce449c Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Mon, 24 Aug 2026 11:24:21 -0700
Subject: [PATCH] fix(proxy): authorize health test-connection on final probe
 params

Move _update_litellm_params_for_health_check before can_user_make_model_call
so health_check_params cannot retarget the probe after the auth check. Type
the Pegasus test helper signature and drop the redundant test narrative.
---
 .../proxy/health_endpoints/_health_endpoints.py   | 11 ++++++-----
 .../proxy/test_health_check_max_tokens.py         | 15 +++------------
 2 files changed, 9 insertions(+), 17 deletions(-)

diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py
index 1feddca9328..c23cde6052e 100644
--- a/litellm/proxy/health_endpoints/_health_endpoints.py
+++ b/litellm/proxy/health_endpoints/_health_endpoints.py
@@ -1952,8 +1952,13 @@ async def test_model_connection(
             **request_litellm_params,
         }
 
-        ## Auth check
         resolved_model_info: Final = loaded_model_info if loaded_model_info is not None else model_info
+        litellm_params = _update_litellm_params_for_health_check(
+            model_info=resolved_model_info or {},
+            litellm_params=litellm_params,
+        )
+
+        ## Auth check, on the final probe params so health_check_params cannot retarget it afterwards
         await ModelManagementAuthChecks.can_user_make_model_call(
             model_params=Deployment(
                 model_name="test_model",
@@ -1964,10 +1969,6 @@ async def test_model_connection(
             prisma_client=prisma_client,
             premium_user=premium_user,
         )
-        litellm_params = _update_litellm_params_for_health_check(
-            model_info=resolved_model_info or {},
-            litellm_params=litellm_params,
-        )
         mode = mode or litellm_params.pop("mode", None)
 
         result: Final = await run_with_timeout(
diff --git a/tests/test_litellm/proxy/test_health_check_max_tokens.py b/tests/test_litellm/proxy/test_health_check_max_tokens.py
index e20b18c8813..162360328d5 100644
--- a/tests/test_litellm/proxy/test_health_check_max_tokens.py
+++ b/tests/test_litellm/proxy/test_health_check_max_tokens.py
@@ -549,17 +549,6 @@ async def test_run_model_health_check_skips_auto_router_deployment():
     assert result == {}
 
 
-# ---------------------------------------------------------------------------
-# model_info.health_check_params
-#
-# Some providers require a payload field litellm does not synthesize for a
-# probe. Bedrock TwelveLabs Pegasus rejects any Invoke body without a top-level
-# `mediaSource`, so every health check on such a deployment failed with
-# "Invalid JSON: $: required property 'mediaSource' not found". The config key
-# was accepted and then never read, so operators had no way to supply it.
-# ---------------------------------------------------------------------------
-
-
 def test_health_check_params_merge_into_probe_params():
     """health_check_params reach the probe request for the deployment that declares them."""
     media_source = {"s3Location": {"uri": "s3://my-bucket/clip.mp4"}}
@@ -640,7 +629,9 @@ def test_health_check_params_apply_to_non_chat_modes():
     assert "max_tokens" not in updated
 
 
-async def _pegasus_health_check_request_body(model_info: dict, monkeypatch) -> dict:
+async def _pegasus_health_check_request_body(
+    model_info: dict[str, object], monkeypatch: pytest.MonkeyPatch
+) -> dict[str, object]:
     monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
     litellm.in_memory_llm_clients_cache.flush_cache()