diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0595c44d69..66bd597755 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2,7 +2,16 @@ import enum import json import uuid from datetime import datetime -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Literal, + Optional, + Union, +) import httpx from pydantic import ( @@ -778,7 +787,6 @@ class GenerateKeyRequest(KeyRequestBase): description="Type of key that determines default allowed routes.", ) - class GenerateKeyResponse(KeyRequestBase): key: str # type: ignore key_name: Optional[str] = None diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 3868c9df69..8a3507e239 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -346,6 +346,35 @@ def handle_key_type(data: GenerateKeyRequest, data_json: dict) -> dict: data_json["allowed_routes"] = ["info_routes"] return data_json +async def validate_team_id_used_in_service_account_request( + team_id: Optional[str], + prisma_client: Optional[PrismaClient], +): + """ + Validate team_id is used in the request body for generating a service account key + """ + if team_id is None: + raise HTTPException( + status_code=400, + detail="team_id is required for service account keys. Please specify `team_id` in the request body.", + ) + + if prisma_client is None: + raise HTTPException( + status_code=400, + detail="prisma_client is required for service account keys. Please specify `prisma_client` in the request body.", + ) + + # check if team_id exists in the database + team = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id}, + ) + if team is None: + raise HTTPException( + status_code=400, + detail="team_id does not exist in the database. Please specify a valid `team_id` in the request body.", + ) + return True async def _common_key_generation_helper( # noqa: PLR0915 data: GenerateKeyRequest, @@ -372,9 +401,9 @@ async def _common_key_generation_helper( # noqa: PLR0915 and data.metadata.get("service_account_id") is not None and data.team_id is None ): - raise HTTPException( - status_code=400, - detail="team_id is required for service account keys. Please specify `team_id` in the request body.", + await validate_team_id_used_in_service_account_request( + team_id=data.team_id, + prisma_client=prisma_client, ) # check if user set default key/generate params on config.yaml @@ -756,6 +785,11 @@ async def generate_service_account_key_fn( user_custom_key_generate, ) + await validate_team_id_used_in_service_account_request( + team_id=data.team_id, + prisma_client=prisma_client, + ) + verbose_proxy_logger.debug("entered /key/generate") if user_custom_key_generate is not None: diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 893e5767ec..3a597adef0 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -576,3 +576,154 @@ async def test_update_service_account_works_with_team_id(): await prepare_key_update_data(data=data, existing_key_row=existing_key) + +@pytest.mark.asyncio +async def test_validate_team_id_used_in_service_account_request_requires_team_id(): + """ + Test that validate_team_id_used_in_service_account_request raises HTTPException + when team_id is None for service account key generation. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + validate_team_id_used_in_service_account_request, + ) + + mock_prisma_client = AsyncMock() + + # Test that HTTPException is raised when team_id is None + with pytest.raises(HTTPException) as exc_info: + await validate_team_id_used_in_service_account_request( + team_id=None, + prisma_client=mock_prisma_client, + ) + + assert exc_info.value.status_code == 400 + assert "team_id is required for service account keys" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_validate_team_id_used_in_service_account_request_requires_prisma_client(): + """ + Test that validate_team_id_used_in_service_account_request raises HTTPException + when prisma_client is None for service account key generation. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + validate_team_id_used_in_service_account_request, + ) + + # Test that HTTPException is raised when prisma_client is None + with pytest.raises(HTTPException) as exc_info: + await validate_team_id_used_in_service_account_request( + team_id="test-team-id", + prisma_client=None, + ) + + assert exc_info.value.status_code == 400 + assert "prisma_client is required for service account keys" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_validate_team_id_used_in_service_account_request_checks_team_exists(): + """ + Test that validate_team_id_used_in_service_account_request validates that + the team_id exists in the database for service account key generation. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + validate_team_id_used_in_service_account_request, + ) + + mock_prisma_client = AsyncMock() + + # Mock the database query to return None (team doesn't exist) + mock_find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique + + # Test that HTTPException is raised when team doesn't exist in DB + with pytest.raises(HTTPException) as exc_info: + await validate_team_id_used_in_service_account_request( + team_id="non-existent-team-id", + prisma_client=mock_prisma_client, + ) + + assert exc_info.value.status_code == 400 + assert "team_id does not exist in the database" in str(exc_info.value.detail) + + # Verify the database was queried with the correct parameters + mock_find_unique.assert_called_once_with( + where={"team_id": "non-existent-team-id"} + ) + + +@pytest.mark.asyncio +async def test_validate_team_id_used_in_service_account_request_success(): + """ + Test that validate_team_id_used_in_service_account_request returns True + when team_id exists in the database for service account key generation. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + validate_team_id_used_in_service_account_request, + ) + + mock_prisma_client = AsyncMock() + + # Mock the database query to return a team object (team exists) + mock_team = {"team_id": "existing-team-id", "team_name": "Test Team"} + mock_find_unique = AsyncMock(return_value=mock_team) + mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique + + # Test that function returns True when team exists + result = await validate_team_id_used_in_service_account_request( + team_id="existing-team-id", + prisma_client=mock_prisma_client, + ) + + assert result is True + + # Verify the database was queried with the correct parameters + mock_find_unique.assert_called_once_with( + where={"team_id": "existing-team-id"} + ) + + +@pytest.mark.asyncio +async def test_generate_service_account_key_endpoint_validation(): + """ + Test that the /key/service-account/generate endpoint properly validates + team_id requirement and team existence in database. + """ + from unittest.mock import patch + + from litellm.proxy.management_endpoints.key_management_endpoints import ( + generate_service_account_key_fn, + ) + + # Test case 1: Missing team_id + with pytest.raises(HTTPException) as exc_info: + await generate_service_account_key_fn( + data=GenerateKeyRequest(team_id=None), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1" + ), + litellm_changed_by=None, + ) + + assert exc_info.value.status_code == 400 + assert "team_id is required for service account keys" in str(exc_info.value.detail) + + # Test case 2: Team doesn't exist in database + with patch('litellm.proxy.proxy_server.prisma_client') as mock_prisma: + # Mock team not found + mock_find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_teamtable.find_unique = mock_find_unique + + with pytest.raises(HTTPException) as exc_info: + await generate_service_account_key_fn( + data=GenerateKeyRequest(team_id="non-existent-team"), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1" + ), + litellm_changed_by=None, + ) + + assert exc_info.value.status_code == 400 + assert "team_id does not exist in the database" in str(exc_info.value.detail) +