From f7a5fb8467629bb85b11385ff843312548d526e1 Mon Sep 17 00:00:00 2001 From: Ishkirat-Singh Date: Wed, 9 Sep 2026 16:21:09 +0530 Subject: [PATCH] fix(proxy): accept an inferred probe mode and drop the added commentary A deployment without model_info.mode gets its probe mode from litellm.model_cost inside ahealth_check, so compare a non-admin's explicit mode against that same lookup instead of rejecting it. Use Mapping types and mark the two unavoidable mutable constructions, and remove the explanatory comments and docstrings per the repository's comment policy. --- .../health_endpoints/_health_endpoints.py | 47 +++++++-------- .../health_endpoints/test_health_endpoints.py | 59 +++++++++++-------- 2 files changed, 55 insertions(+), 51 deletions(-) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index dafe2ad377e..2b08daf03cf 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -2001,17 +2001,26 @@ async def health_liveliness_options(): return Response(headers=response_headers, status_code=200) -# What a caller admitted through the non-admin path gets to see of the probe: -# the model and the outcome, none of the deployment's routing configuration. _NON_ADMIN_TEST_CONNECTION_RESULT_KEYS: Final[frozenset[str]] = frozenset(("model", "error", "mode_error")) -def _test_connection_result_for_display(endpoint_data: dict, *, outcome_only: bool) -> dict: - """Clean the probe result for display; ``outcome_only`` hides the deployment's configuration.""" - cleaned: Final = _clean_endpoint_data(endpoint_data, details=True) +def _test_connection_result_for_display( + litellm_params: Mapping[str, object], result: Mapping[str, object], *, outcome_only: bool +) -> Mapping[str, object]: + cleaned: Final = _clean_endpoint_data({**litellm_params, **result}, details=True) if not outcome_only: return cleaned - return {k: v for k, v in cleaned.items() if k in _NON_ADMIN_TEST_CONNECTION_RESULT_KEYS} + return { # mutable-ok: fresh filtered copy handed to the caller + k: v for k, v in cleaned.items() if k in _NON_ADMIN_TEST_CONNECTION_RESULT_KEYS + } + + +def _configured_probe_mode(model_info: Mapping[str, object] | None, model: object) -> str | None: + configured: Final = model_info.get("mode") if model_info else None + if configured is not None: + return str(configured) + cost_entry: Final = litellm.model_cost.get(model) if isinstance(model, str) else None + return cost_entry.get("mode") if cost_entry else None def _probe_is_configured_deployment( @@ -2022,11 +2031,6 @@ def _probe_is_configured_deployment( request_litellm_params: Mapping[str, object], requested_mode: str | None, ) -> bool: - """Whether a probe targets a team-less configured deployment exactly as configured. - - A request value that differs from the configuration (model, provider, - endpoint, credentials, mode, ...) describes a different probe. - """ if getattr(model_params.model_info, "team_id", None) is not None: return False if any(configured_litellm_params.get(key) != value for key, value in request_litellm_params.items()): @@ -2041,7 +2045,6 @@ async def _assert_caller_can_call_model( prisma_client: "PrismaClient", llm_router: "Router | None", ) -> None: - """Raise 403 unless the caller's key and user may call ``model``.""" from litellm.proxy.auth.auth_checks import ( UserNotFoundError, can_key_call_model, @@ -2068,7 +2071,10 @@ async def _assert_caller_can_call_model( user_object = None await can_user_call_model(model=model, llm_router=llm_router, user_object=user_object) except ProxyException as e: - raise HTTPException(status_code=403, detail={"error": str(e.message)}) from e + raise HTTPException( + status_code=403, + detail={"error": str(e.message)}, # mutable-ok: same 403 payload shape as the rest of this endpoint + ) from e async def _authorize_test_connection( @@ -2084,16 +2090,7 @@ async def _authorize_test_connection( request_litellm_params: Mapping[str, object], requested_mode: str | None, ) -> bool: - """Decide whether the caller may probe this model. - - Proxy admins and team admins may probe any model they manage, as before. - Any other user may probe a configured deployment that has no team, exactly - as configured, when their key and user are allowed to call its model. - Anything else stays a management operation. - - Returns True when the caller was admitted through that non-admin path, so - the response can be limited to the outcome of the probe. - """ + """Returns True when a non-admin was admitted to probe a team-less deployment as configured.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelManagementAuthChecks, ) @@ -2321,7 +2318,7 @@ async def test_model_connection( llm_router=llm_router, configured_model_name=configured_model_name, configured_litellm_params=config_litellm_params, - configured_mode=(loaded_model_info or {}).get("mode"), + configured_mode=_configured_probe_mode(loaded_model_info, config_litellm_params.get("model")), request_litellm_params=request_litellm_params, requested_mode=mode or request_litellm_params.get("mode"), ) @@ -2339,7 +2336,7 @@ async def test_model_connection( # Clean the result for display cleaned_result: Final = _test_connection_result_for_display( - {**litellm_params, **result}, outcome_only=admitted_as_caller + litellm_params, result, outcome_only=admitted_as_caller ) return { 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 c6a3aa6d5f0..b9b05c37f10 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -4205,10 +4205,6 @@ def _internal_user(): @pytest.mark.asyncio async def test_test_model_connection_allows_internal_user_to_probe_configured_model_they_can_call(): - """ - An internal user may test a configured, non-team model they are allowed - to call, and the probe runs with the configured credentials. - """ mock_router = MagicMock() mock_router.get_deployment.return_value = _configured_non_team_deployment() @@ -4233,7 +4229,6 @@ async def test_test_model_connection_allows_internal_user_to_probe_configured_mo assert key_check.await_args.kwargs["model"] == "gpt-4o" assert user_check.await_args.kwargs["model"] == "gpt-4o" assert health_check.await_args.kwargs["model_params"]["api_key"] == "CONFIGURED-API-KEY" - # the caller learns the outcome, not the deployment's routing configuration assert "api_base" not in result["result"] assert "api_key" not in result["result"] @@ -4278,11 +4273,6 @@ async def test_test_model_connection_denies_internal_user_without_model_access() async def test_test_model_connection_keeps_connection_overrides_admin_only(): from fastapi import HTTPException - """ - A request that sets its own connection fields describes a different - endpoint than the configured one; probing that stays a management - operation, so an internal user is denied before any access check runs. - """ mock_router = MagicMock() mock_router.get_deployment.return_value = _configured_non_team_deployment() @@ -4312,11 +4302,6 @@ async def test_test_model_connection_keeps_connection_overrides_admin_only(): @pytest.mark.asyncio async def test_test_model_connection_keeps_model_overrides_admin_only(): - """ - The probe runs whatever `litellm_params.model` the request sends, so a - non-admin may only send the configured deployment's own model: the access - check and the probe must target the same model. - """ from fastapi import HTTPException mock_router = MagicMock() @@ -4345,11 +4330,6 @@ async def test_test_model_connection_keeps_model_overrides_admin_only(): @pytest.mark.asyncio async def test_test_model_connection_keeps_team_deployments_admin_only_for_non_admins(): - """ - A team deployment keeps its team-admin policy: a non-admin whose key can - call the same model name must not reach the non-admin path with the - team's configured credentials. - """ from fastapi import HTTPException from litellm.proxy._types import LiteLLM_TeamTable @@ -4393,7 +4373,6 @@ async def test_test_model_connection_keeps_team_deployments_admin_only_for_non_a @pytest.mark.asyncio async def test_test_model_connection_tolerates_missing_user_record_for_non_admin(): - """A key whose user record is gone is judged on the key's own model access.""" from litellm.proxy.auth.auth_checks import UserNotFoundError mock_router = MagicMock() @@ -4425,11 +4404,6 @@ async def test_test_model_connection_tolerates_missing_user_record_for_non_admin @pytest.mark.asyncio async def test_test_model_connection_keeps_mode_overrides_admin_only(): - """ - `mode` selects which provider operation the probe performs, so a non-admin - may only probe with the deployment's configured mode (or none, which - auto-detects it). - """ from fastapi import HTTPException from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4461,3 +4435,36 @@ async def test_test_model_connection_keeps_mode_overrides_admin_only(): assert exc_info.value.status_code == 403 key_check.assert_not_awaited() health_check.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_test_model_connection_accepts_mode_the_probe_would_infer(): + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + deployment = Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="gpt-4o", api_key="CONFIGURED-API-KEY"), + model_info=ModelInfo(id="non-team-deployment-id"), + ) + mock_router = MagicMock() + mock_router.get_deployment.return_value = deployment + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.llm_router", mock_router), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.auth.auth_checks.can_key_call_model", AsyncMock(return_value=True)), + patch("litellm.proxy.auth.auth_checks.get_user_object", AsyncMock(return_value=None)), + patch("litellm.proxy.auth.auth_checks.can_user_call_model", AsyncMock(return_value=True)), + patch("litellm.ahealth_check", AsyncMock(return_value={"status": "healthy"})) as health_check, + ): + result = await health_test_model_connection( + request=MagicMock(), + mode="chat", + litellm_params={"model": "gpt-4o"}, + model_info={"id": "non-team-deployment-id"}, + user_api_key_dict=_internal_user(), + ) + + assert result["status"] == "success" + assert health_check.await_args.kwargs["mode"] == "chat"