diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 726874bc82d..67dee47ca8f 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1865,31 +1865,6 @@ async def _validate_update_key_data( detail=f"User={data.user_id} is not allowed to update key={data.key} to belong to user={existing_key_row.user_id}", ) - # Validate team exists when non-admin changes team_id (LIT-1884) - if ( - data.team_id is not None - and not _is_proxy_admin - ): - try: - _team_obj = await get_team_object( - team_id=data.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - check_db_only=True, - ) - if _team_obj is None: - raise HTTPException( - status_code=400, - detail=f"Team not found for team_id={data.team_id}. Non-admin users cannot set keys to non-existent teams.", - ) - except HTTPException: - raise - except Exception: - raise HTTPException( - status_code=400, - detail=f"Team not found for team_id={data.team_id}. Non-admin users cannot set keys to non-existent teams.", - ) - common_key_access_checks( user_api_key_dict=user_api_key_dict, data=data, @@ -1929,6 +1904,13 @@ async def _validate_update_key_data( check_db_only=True, ) + # Validate team exists when non-admin sets a new team_id (LIT-1884) + if team_obj is None and data.team_id is not None and not _is_proxy_admin: + raise HTTPException( + status_code=400, + detail=f"Team not found for team_id={data.team_id}. Non-admin users cannot set keys to non-existent teams.", + ) + if team_obj is not None: await _check_team_key_limits( team_table=team_obj, 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 e1b4a5288fb..08e9b5028d0 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 @@ -7768,12 +7768,15 @@ class TestLIT1884KeyUpdateValidation: async def test_internal_user_cannot_set_invalid_team_id(self): """ Non-admin users should not be able to update a key to a non-existent team. + get_team_object raises HTTPException(404) when team doesn't exist in DB. """ data = UpdateKeyRequest(key="sk-test-key", team_id="nonexistent-team") existing_key_row = MagicMock() existing_key_row.user_id = "internal-user-123" existing_key_row.token = "hashed_token" existing_key_row.team_id = None + existing_key_row.organization_id = None + existing_key_row.project_id = None user_api_key_dict = UserAPIKeyAuth( user_id="internal-user-123", @@ -7782,7 +7785,10 @@ class TestLIT1884KeyUpdateValidation: with patch( "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", - AsyncMock(side_effect=Exception("Team not found")), + AsyncMock(side_effect=HTTPException( + status_code=404, + detail="Team doesn't exist in db. Team=nonexistent-team.", + )), ): with pytest.raises(HTTPException) as exc_info: await _validate_update_key_data( @@ -7794,8 +7800,8 @@ class TestLIT1884KeyUpdateValidation: prisma_client=AsyncMock(), user_api_key_cache=MagicMock(), ) - assert exc_info.value.status_code == 400 - assert "Team not found" in str(exc_info.value.detail) + assert exc_info.value.status_code == 404 + assert "Team doesn't exist" in str(exc_info.value.detail) @pytest.mark.asyncio async def test_admin_can_remove_user_id(self):