diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a1650b4968..ddce7481ce 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -367,6 +367,16 @@ async def _common_key_generation_helper( # noqa: PLR0915 premium_user=premium_user, ) + if ( + data.metadata is not None + 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.", + ) + # check if user set default key/generate params on config.yaml if litellm.default_key_generate_params is not None: for elem in data: @@ -860,6 +870,15 @@ async def prepare_key_update_data( data_json: dict = data.model_dump(exclude_unset=True) data_json.pop("key", None) data_json.pop("new_key", None) + if ( + data.metadata is not None + and data.metadata.get("service_account_id") is not None + and (data.team_id or existing_key_row.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.", + ) non_default_values = {} for k, v in data_json.items(): if ( 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 5490999903..e0102f8cd7 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 @@ -11,9 +11,20 @@ sys.path.insert( from unittest.mock import AsyncMock, MagicMock -import pytest +from fastapi import HTTPException -from litellm.proxy.management_endpoints.key_management_endpoints import _list_key_helper +from litellm.proxy._types import ( + GenerateKeyRequest, + LiteLLM_VerificationToken, + LitellmUserRoles, + UpdateKeyRequest, +) +from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.proxy.management_endpoints.key_management_endpoints import ( + _common_key_generation_helper, + _list_key_helper, + prepare_key_update_data, +) from litellm.proxy.proxy_server import app client = TestClient(app) @@ -492,3 +503,74 @@ def test_get_new_token_with_invalid_key(): assert exc_info.value.status_code == 400 assert "New key must start with 'sk-'" in str(exc_info.value.detail) + + + +@pytest.mark.asyncio +async def test_generate_service_account_requires_team_id(): + with pytest.raises(HTTPException): + await _common_key_generation_helper( + data=GenerateKeyRequest( + metadata={"service_account_id": "sa"}, + team_id=None, + ), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1" + ), + litellm_changed_by=None, + team_table=None, + ) + + +@pytest.mark.asyncio +async def test_generate_service_account_works_with_team_id(): + from unittest.mock import patch + + # Mock the database and router dependencies from proxy_server + with patch('litellm.proxy.proxy_server.prisma_client') as mock_prisma, \ + patch('litellm.proxy.proxy_server.llm_router') as mock_router, \ + patch('litellm.proxy.proxy_server.premium_user', False), \ + patch('litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn') as mock_generate_key: + + # Configure mocks + mock_prisma.return_value = AsyncMock() + mock_router.return_value = None + # Mock the response from generate_key_helper_fn + mock_generate_key.return_value = { + "key": "sk-test-key", + "expires": None, + "user_id": "test-user", + "team_id": "IJ" + } + + # This should not raise an exception since team_id is provided + await _common_key_generation_helper( + data=GenerateKeyRequest( + metadata={"service_account_id": "sa"}, + team_id="IJ", + ), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1" + ), + litellm_changed_by=None, + team_table=None, + ) + + + +@pytest.mark.asyncio +async def test_update_service_account_requires_team_id(): + data = UpdateKeyRequest(key="sk-1", metadata={"service_account_id": "sa"}) + existing_key = LiteLLM_VerificationToken(token="hashed", team_id=None) + + with pytest.raises(HTTPException): + await prepare_key_update_data(data=data, existing_key_row=existing_key) + + +@pytest.mark.asyncio +async def test_update_service_account_works_with_team_id(): + data = UpdateKeyRequest(key="sk-1", metadata={"service_account_id": "sa"}, team_id="IJ") + existing_key = LiteLLM_VerificationToken(token="hashed") + + await prepare_key_update_data(data=data, existing_key_row=existing_key) +