[Fix] Ensure team_id is a required field for generating service account keys (#14270)

* generate_service_account_key_fn

* fix validate_team_id_used_in_service_account_request

* fix types

* test_validate_team_id_used_in_service_account_request_requires_team_id

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

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:

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)
+