From 9b72cf720b39b9e5daf16e596dc0ce99f280631a Mon Sep 17 00:00:00 2001 From: Ishkirat-Singh Date: Wed, 9 Sep 2026 16:03:35 +0530 Subject: [PATCH] fix(proxy): bind the probe mode to the configured deployment for non-admins `mode` selects which provider operation the health probe performs, so it belongs to the configuration match like every other request value: a non-admin may only probe with the deployment's configured mode, or with none and let it be detected. Also split the authorization helper into a gate, the key/user access check and the orchestrator, and move the response filtering out of the handler, keeping each function's complexity within the repository's ceiling. --- .../health_endpoints/_health_endpoints.py | 143 ++++++++++++------ .../health_endpoints/test_health_endpoints.py | 42 ++++- 2 files changed, 135 insertions(+), 50 deletions(-) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index ffb48a460b1..dafe2ad377e 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -2006,62 +2006,53 @@ async def health_liveliness_options(): _NON_ADMIN_TEST_CONNECTION_RESULT_KEYS: Final[frozenset[str]] = frozenset(("model", "error", "mode_error")) -async def _authorize_test_connection( +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) + if not outcome_only: + return cleaned + return {k: v for k, v in cleaned.items() if k in _NON_ADMIN_TEST_CONNECTION_RESULT_KEYS} + + +def _probe_is_configured_deployment( *, model_params: "Deployment", + configured_litellm_params: Mapping[str, object], + configured_mode: str | None, + 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()): + return False + return requested_mode is None or requested_mode == configured_mode + + +async def _assert_caller_can_call_model( + *, + model: str, user_api_key_dict: UserAPIKeyAuth, prisma_client: "PrismaClient", - premium_user: bool, llm_router: "Router | None", - configured_model_name: str | None, - configured_litellm_params: Mapping[str, object], - request_litellm_params: Mapping[str, object], -) -> 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: - 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. - """ +) -> 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, can_user_call_model, get_user_object, ) - from litellm.proxy.management_endpoints.model_management_endpoints import ( - ModelManagementAuthChecks, - ) from litellm.proxy.proxy_server import llm_model_list, user_api_key_cache - try: - await ModelManagementAuthChecks.can_user_make_model_call( - model_params=model_params, - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - premium_user=premium_user, - ) - return False - except HTTPException as management_denial: - if management_denial.status_code != 403: - raise - 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: await can_key_call_model( - model=configured_model_name, + model=model, llm_model_list=llm_model_list, valid_token=user_api_key_dict, llm_router=llm_router, @@ -2075,13 +2066,65 @@ async def _authorize_test_connection( ) except UserNotFoundError: user_object = None - await can_user_call_model( - model=configured_model_name, - llm_router=llm_router, - user_object=user_object, - ) + 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 + + +async def _authorize_test_connection( + *, + model_params: "Deployment", + user_api_key_dict: UserAPIKeyAuth, + prisma_client: "PrismaClient", + premium_user: bool, + llm_router: "Router | None", + configured_model_name: str | None, + configured_litellm_params: Mapping[str, object], + configured_mode: str | None, + 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. + """ + from litellm.proxy.management_endpoints.model_management_endpoints import ( + ModelManagementAuthChecks, + ) + + try: + await ModelManagementAuthChecks.can_user_make_model_call( + model_params=model_params, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + premium_user=premium_user, + ) + return False + except HTTPException as management_denial: + if ( + management_denial.status_code != 403 + or configured_model_name is None + or not _probe_is_configured_deployment( + model_params=model_params, + configured_litellm_params=configured_litellm_params, + configured_mode=configured_mode, + request_litellm_params=request_litellm_params, + requested_mode=requested_mode, + ) + ): + raise + await _assert_caller_can_call_model( + model=configured_model_name, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + llm_router=llm_router, + ) return True @@ -2278,7 +2321,9 @@ 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"), request_litellm_params=request_litellm_params, + requested_mode=mode or request_litellm_params.get("mode"), ) mode = mode or litellm_params.pop("mode", None) @@ -2293,9 +2338,9 @@ async def test_model_connection( ) # Clean the result for display - 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} + cleaned_result: Final = _test_connection_result_for_display( + {**litellm_params, **result}, outcome_only=admitted_as_caller + ) 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 543c4e2453d..c6a3aa6d5f0 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -4191,7 +4191,7 @@ def _configured_non_team_deployment(): api_base="https://configured.invalid/v1", api_key="CONFIGURED-API-KEY", ), - model_info=ModelInfo(id="non-team-deployment-id"), + model_info=ModelInfo(id="non-team-deployment-id", mode="chat"), ) @@ -4421,3 +4421,43 @@ async def test_test_model_connection_tolerates_missing_user_record_for_non_admin assert result["status"] == "success" assert user_check.await_args.kwargs["user_object"] is None + + +@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 + + deployment = Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="openai/gpt-4o", api_key="CONFIGURED-API-KEY"), + model_info=ModelInfo(id="non-team-deployment-id", mode="chat"), + ) + 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)) 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="image_generation", + litellm_params={"model": "openai/gpt-4o"}, + 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()