diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index f6174f62ab4..ffb48a460b1 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -8,7 +8,7 @@ import time import traceback from collections.abc import Iterable, Mapping from datetime import datetime, timedelta, timezone -from typing import Any, Final, Literal, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, cast import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, Response, status @@ -64,6 +64,10 @@ from litellm.proxy.middleware.in_flight_requests_middleware import ( ) from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager from litellm.router import Router + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + from litellm.types.router import Deployment from litellm.router_utils.clientside_credential_handler import ( _ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the router path clientside_credential_keys, @@ -1997,24 +2001,36 @@ 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")) + + async def _authorize_test_connection( *, - model_params: Any, + model_params: "Deployment", user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, + prisma_client: "PrismaClient", premium_user: bool, - llm_router: Any, + llm_router: "Router | None", configured_model_name: str | None, + configured_litellm_params: Mapping[str, object], request_litellm_params: Mapping[str, object], -) -> 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 model they are allowed to call, but - only as configured: a request that sets its own connection fields describes - a different endpoint, and probing that stays a management operation. + 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: + a request value that differs from the configuration (model, provider, + endpoint, credentials, ...) describes a different probe, and that stays a + management operation. Team deployments keep their team-admin policy. + + Returns True when the caller was admitted through that non-admin path, so + the response can be limited to the outcome of the probe. """ from litellm.proxy.auth.auth_checks import ( + UserNotFoundError, can_key_call_model, can_user_call_model, get_user_object, @@ -2031,11 +2047,16 @@ async def _authorize_test_connection( prisma_client=prisma_client, premium_user=premium_user, ) - return + return False except HTTPException as management_denial: if management_denial.status_code != 403: raise - if configured_model_name is None or any(field in request_litellm_params for field in _CONFIG_CONNECTION_FIELDS): + if configured_model_name is None or getattr(model_params.model_info, "team_id", None) is not None: + raise + if any( + key != "mode" and configured_litellm_params.get(key) != value + for key, value in request_litellm_params.items() + ): raise try: @@ -2045,12 +2066,15 @@ async def _authorize_test_connection( valid_token=user_api_key_dict, llm_router=llm_router, ) - user_object = await get_user_object( - user_id=user_api_key_dict.user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - ) + try: + user_object = await get_user_object( + user_id=user_api_key_dict.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + ) + except UserNotFoundError: + user_object = None await can_user_call_model( model=configured_model_name, llm_router=llm_router, @@ -2058,6 +2082,7 @@ async def _authorize_test_connection( ) except ProxyException as e: raise HTTPException(status_code=403, detail={"error": str(e.message)}) from e + return True @router.post( @@ -2241,7 +2266,7 @@ async def test_model_connection( ) ## Auth check, on the final probe params so health_check_params cannot retarget it afterwards - await _authorize_test_connection( + admitted_as_caller: Final = await _authorize_test_connection( model_params=Deployment( model_name="test_model", litellm_params=LiteLLM_Params(**litellm_params), @@ -2252,6 +2277,7 @@ async def test_model_connection( premium_user=premium_user, llm_router=llm_router, configured_model_name=configured_model_name, + configured_litellm_params=config_litellm_params, request_litellm_params=request_litellm_params, ) mode = mode or litellm_params.pop("mode", None) @@ -2267,7 +2293,9 @@ async def test_model_connection( ) # Clean the result for display - cleaned_result: Final = _clean_endpoint_data({**litellm_params, **result}, details=True) + cleaned_result = _clean_endpoint_data({**litellm_params, **result}, details=True) + if admitted_as_caller: + cleaned_result = {k: v for k, v in cleaned_result.items() if k in _NON_ADMIN_TEST_CONNECTION_RESULT_KEYS} return { "status": "error" if "error" in result else "success", 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 ecf4523fb45..543c4e2453d 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -4233,6 +4233,9 @@ 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"] @pytest.mark.asyncio @@ -4305,3 +4308,116 @@ async def test_test_model_connection_keeps_connection_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_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() + mock_router.get_deployment.return_value = _configured_non_team_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)) as key_check, + patch("litellm.ahealth_check", AsyncMock()) as health_check, + pytest.raises(HTTPException) as exc_info, + ): + await health_test_model_connection( + request=MagicMock(), + mode="chat", + litellm_params={"model": "openai/gpt-4o-mini"}, + model_info={"id": "non-team-deployment-id"}, + user_api_key_dict=_internal_user(), + ) + + 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_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 + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + team_deployment = Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="openai/gpt-4o", api_key="TEAM-A-API-KEY"), + model_info=ModelInfo(id="team-a-deployment-id", team_id="team-a"), + ) + mock_router = MagicMock() + mock_router.get_deployment.return_value = team_deployment + team_row = SimpleNamespace( + model_dump=lambda: LiteLLM_TeamTable(team_id="team-a", members_with_roles=[]).model_dump() + ) + + 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.management_endpoints.model_management_endpoints.TeamRepository") as MockTeamRepo, + patch("litellm.proxy.auth.auth_checks.can_key_call_model", AsyncMock(return_value=True)) as key_check, + patch("litellm.ahealth_check", AsyncMock()) as health_check, + ): + repo = MagicMock() + repo.table.find_unique = AsyncMock(return_value=team_row) + MockTeamRepo.return_value = repo + with pytest.raises(HTTPException) as exc_info: + await health_test_model_connection( + request=MagicMock(), + mode="chat", + litellm_params={"model": "openai/gpt-4o"}, + model_info={"id": "team-a-deployment-id"}, + user_api_key_dict=_internal_user(), + ) + + 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_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() + mock_router.get_deployment.return_value = _configured_non_team_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(side_effect=UserNotFoundError("gone")), + ), + patch("litellm.proxy.auth.auth_checks.can_user_call_model", AsyncMock(return_value=True)) as user_check, + patch("litellm.ahealth_check", AsyncMock(return_value={"status": "healthy"})), + ): + result = await health_test_model_connection( + request=MagicMock(), + mode="chat", + litellm_params={"model": "openai/gpt-4o"}, + model_info={"id": "non-team-deployment-id"}, + user_api_key_dict=_internal_user(), + ) + + assert result["status"] == "success" + assert user_check.await_args.kwargs["user_object"] is None