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"