diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index c20991b8d43..2b3ec231ac8 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -1807,6 +1807,7 @@ async def test_model_connection( # Look up model configuration from router if model name is provided # This gets the litellm_params from proxy config (with resolved env vars) config_litellm_params: dict = {} + loaded_model_info: Optional[dict] = None if llm_router is not None: # Prefer disambiguation by deployment id (`model_info.id`) when # the caller supplies it. This is required when multiple @@ -1825,6 +1826,7 @@ async def test_model_connection( if deployment_by_id is not None: config_litellm_params = deployment_by_id.litellm_params.model_dump(exclude_none=True) + loaded_model_info = deployment_by_id.model_info.model_dump(exclude_none=True) elif model_name: # Fall back to model_name lookup for callers (e.g. the # "Add Model" wizard, or curl) that don't supply an id. @@ -1846,6 +1848,7 @@ async def test_model_connection( # config. These already have resolved environment # variables from proxy config. config_litellm_params = dict(deployments[0].get("litellm_params", {})) + loaded_model_info = dict(deployments[0].get("model_info") or {}) except Exception as e: verbose_proxy_logger.debug( f"Could not find model {model_name} in router: {e}. Proceeding with request params only." @@ -1856,11 +1859,12 @@ async def test_model_connection( litellm_params = {**config_litellm_params, **request_litellm_params} ## Auth check + auth_model_info = loaded_model_info if loaded_model_info is not None else model_info await ModelManagementAuthChecks.can_user_make_model_call( model_params=Deployment( model_name="test_model", litellm_params=LiteLLM_Params(**litellm_params), - model_info=model_info, + model_info=auth_model_info, ), user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, 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 a04ad5598df..917bedcb93f 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -696,6 +696,309 @@ async def test_test_model_connection_falls_back_to_deployments_zero_without_id() assert model_params.get("api_key") == "fake-key-A" +@pytest.mark.asyncio +async def test_test_model_connection_uses_loaded_deployment_team_id(): + """ + /health/test_connection must authorize using the team_id of the + deployment it actually loaded (by model_info.id), not the team_id + supplied in the request body. Requesting team A's deployment while + authenticated as an admin of team B must be denied. + """ + from fastapi import HTTPException + + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.management_endpoints.model_management_endpoints import ( + ModelManagementAuthChecks, + ) + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + mock_request = MagicMock() + + requester_team_id = "team-b" + deployment_owner_team_id = "team-a" + deployment_id = "team-a-deployment-id" + + requester_user_api_key_dict = UserAPIKeyAuth( + token="requester-token", + user_id="team-b-admin-user", + team_id=requester_team_id, + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + mock_prisma_client = MagicMock() + + other_team_deployment = Deployment( + model_name="team-a-model", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://team-a-api.invalid/v1", + api_key="TEAM-A-API-KEY", + ), + model_info=ModelInfo(id=deployment_id, team_id=deployment_owner_team_id), + ) + + mock_router = MagicMock() + mock_router.get_deployment.return_value = other_team_deployment + + async def fake_find_unique(*, where): + team_id = where["team_id"] + if team_id == requester_team_id: + return SimpleNamespace( + model_dump=lambda: LiteLLM_TeamTable( + team_id=requester_team_id, + members_with_roles=[ + { + "user_id": "team-b-admin-user", + "role": "admin", + } + ], + ).model_dump() + ) + if team_id == deployment_owner_team_id: + return SimpleNamespace( + model_dump=lambda: LiteLLM_TeamTable( + team_id=deployment_owner_team_id, + members_with_roles=[], + ).model_dump() + ) + return None + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), + patch("litellm.proxy.proxy_server.llm_router", mock_router), + patch("litellm.proxy.proxy_server.premium_user", True), + patch.object( + ModelManagementAuthChecks, + "can_user_make_model_call", + wraps=ModelManagementAuthChecks.can_user_make_model_call, + ) as spy_auth_check, + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.TeamRepository" + ) as MockTeamRepo, + ): + mock_team_repo_instance = MagicMock() + mock_team_repo_instance.table.find_unique = AsyncMock( + side_effect=fake_find_unique + ) + MockTeamRepo.return_value = mock_team_repo_instance + + with pytest.raises(HTTPException) as exc_info: + await health_test_model_connection( + request=mock_request, + mode="chat", + litellm_params={ + "model": "openai/gpt-4o", + "api_base": "https://swapped-base.invalid/v1", + }, + model_info={ + "id": deployment_id, + "team_id": requester_team_id, + }, + user_api_key_dict=requester_user_api_key_dict, + ) + + assert exc_info.value.status_code == 403 + assert spy_auth_check.called + passed_model_params = spy_auth_check.call_args.kwargs["model_params"] + assert passed_model_params.model_info.team_id == deployment_owner_team_id, ( + "Auth check must run against the loaded deployment's team_id " + f"({deployment_owner_team_id!r}); got " + f"{passed_model_params.model_info.team_id!r}." + ) + + +@pytest.mark.asyncio +async def test_test_model_connection_uses_loaded_deployment_team_id_via_model_name_fallback(): + """ + Companion to the id-lookup case: when the caller provides only a model + name (no `model_info.id`) and that name resolves via the router's + `model_name` fallback to a deployment owned by a different team, the + auth check must still run against the loaded deployment's `team_id`, + not the caller-supplied one in the request body. + """ + from fastapi import HTTPException + + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.management_endpoints.model_management_endpoints import ( + ModelManagementAuthChecks, + ) + + mock_request = MagicMock() + + requester_team_id = "team-b-2" + deployment_owner_team_id = "team-a-2" + + requester_user_api_key_dict = UserAPIKeyAuth( + token="requester-token-2", + user_id="team-b-admin-user-2", + team_id=requester_team_id, + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + mock_prisma_client = MagicMock() + + other_team_deployment_dict = { + "model_name": "shared-model-name", + "litellm_params": { + "model": "openai/gpt-4o", + "api_base": "https://team-a-api-2.invalid/v1", + "api_key": "TEAM-A-API-KEY-2", + }, + "model_info": { + "id": "team-a-deployment-id-2", + "team_id": deployment_owner_team_id, + }, + } + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [other_team_deployment_dict] + + async def fake_find_unique(*, where): + return SimpleNamespace( + model_dump=lambda: LiteLLM_TeamTable( + team_id=where["team_id"], + members_with_roles=( + [{"user_id": "team-b-admin-user-2", "role": "admin"}] + if where["team_id"] == requester_team_id + else [] + ), + ).model_dump() + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), + patch("litellm.proxy.proxy_server.llm_router", mock_router), + patch("litellm.proxy.proxy_server.premium_user", True), + patch.object( + ModelManagementAuthChecks, + "can_user_make_model_call", + wraps=ModelManagementAuthChecks.can_user_make_model_call, + ) as spy_auth_check, + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.TeamRepository" + ) as MockTeamRepo, + ): + mock_team_repo_instance = MagicMock() + mock_team_repo_instance.table.find_unique = AsyncMock( + side_effect=fake_find_unique + ) + MockTeamRepo.return_value = mock_team_repo_instance + + with pytest.raises(HTTPException) as exc_info: + await health_test_model_connection( + request=mock_request, + mode="chat", + litellm_params={ + "model": "shared-model-name", + "api_base": "https://swapped-base-2.invalid/v1", + }, + model_info={"team_id": requester_team_id}, + user_api_key_dict=requester_user_api_key_dict, + ) + + assert exc_info.value.status_code == 403 + + passed_model_params = spy_auth_check.call_args.kwargs["model_params"] + assert passed_model_params.model_info.team_id == deployment_owner_team_id + + +@pytest.mark.asyncio +async def test_test_model_connection_authorized_team_admin_passes_real_auth(): + """ + Positive-path companion to the deny tests above. When the caller is a + genuine admin of the team that owns the loaded deployment, the real + (unmocked) auth check must pass and the endpoint must reach the outbound + health probe. Guards against a regression that swaps the auth `team_id` + for something deny-all on the legit path. + """ + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.management_endpoints.model_management_endpoints import ( + ModelManagementAuthChecks, + ) + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + mock_request = MagicMock() + + owner_team_id = "team-owner" + owner_admin_user_id = "team-owner-admin" + owned_deployment_id = "owned-deployment-id" + + owner_admin_api_key_dict = UserAPIKeyAuth( + token="owner-admin-token", + user_id=owner_admin_user_id, + team_id=owner_team_id, + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + mock_prisma_client = MagicMock() + + owned_deployment = Deployment( + model_name="owner-model", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o-mini", + api_base="https://owner-real-api.invalid/v1", + api_key="owner-team-api-key", + ), + model_info=ModelInfo(id=owned_deployment_id, team_id=owner_team_id), + ) + + mock_router = MagicMock() + mock_router.get_deployment.return_value = owned_deployment + + async def fake_find_unique(*, where): + if where["team_id"] == owner_team_id: + return SimpleNamespace( + model_dump=lambda: LiteLLM_TeamTable( + team_id=owner_team_id, + members_with_roles=[ + {"user_id": owner_admin_user_id, "role": "admin"} + ], + ).model_dump() + ) + return None + + health_result = {"status": "healthy", "response_time_ms": 50} + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), + patch("litellm.proxy.proxy_server.llm_router", mock_router), + patch("litellm.proxy.proxy_server.premium_user", True), + patch.object( + ModelManagementAuthChecks, + "can_user_make_model_call", + wraps=ModelManagementAuthChecks.can_user_make_model_call, + ) as spy_auth_check, + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.TeamRepository" + ) as MockTeamRepo, + patch( + "litellm.proxy.health_endpoints._health_endpoints.litellm.ahealth_check", + AsyncMock(return_value=health_result), + ), + patch( + "litellm.proxy.health_endpoints._health_endpoints.run_with_timeout", + AsyncMock(return_value=health_result), + ), + ): + mock_team_repo_instance = MagicMock() + mock_team_repo_instance.table.find_unique = AsyncMock( + side_effect=fake_find_unique + ) + MockTeamRepo.return_value = mock_team_repo_instance + + result = await health_test_model_connection( + request=mock_request, + mode="chat", + litellm_params={"model": "openai/gpt-4o-mini"}, + model_info={"id": owned_deployment_id, "team_id": owner_team_id}, + user_api_key_dict=owner_admin_api_key_dict, + ) + + assert result["status"] == "success" + passed_model_params = spy_auth_check.call_args.kwargs["model_params"] + assert passed_model_params.model_info.team_id == owner_team_id + + @pytest.mark.asyncio @pytest.mark.parametrize( "status,error_message",