diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b012d25ee10..4214e3ad870 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3900,7 +3900,6 @@ async def _check_team_member_budget( and team_object.team_id is not None and valid_token is not None and valid_token.user_id is not None - and valid_token.project_id is None ): team_membership: Final = await get_team_membership( user_id=valid_token.user_id, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 938e557ae22..42da94e8ac4 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1049,7 +1049,7 @@ async def _team_member_budget_check_for_key( prisma_client: PrismaClient | None, user_api_key_cache: UserApiKeyCache, ) -> None: - if valid_token.team_member_spend is None or valid_token.project_id is not None: + if valid_token.team_member_spend is None: return if prisma_client is None: return diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 39e3c9273af..ead2dfb7837 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -540,9 +540,6 @@ async def _get_team_member_budget_counter( if team_object is None or team_object.team_id is None or user_object is None or valid_token.user_id is None: return None - if valid_token.project_id is not None: - return None - membership_cache_key: Final = f"team_membership:{valid_token.user_id}:{team_object.team_id}" cached_team_membership: Final = await user_api_key_cache.async_get_cache(key=membership_cache_key) team_membership: LiteLLM_TeamMembership | None = None diff --git a/tests/test_litellm/proxy/auth/test_team_member_budget.py b/tests/test_litellm/proxy/auth/test_team_member_budget.py index e4ec62810fe..5eaabe97d44 100644 --- a/tests/test_litellm/proxy/auth/test_team_member_budget.py +++ b/tests/test_litellm/proxy/auth/test_team_member_budget.py @@ -439,11 +439,11 @@ async def test_team_member_budget_check_personal_key_not_team(): @pytest.mark.asyncio -async def test_team_member_budget_check_skipped_for_project_scoped_key(): +async def test_team_member_budget_check_enforced_for_project_scoped_key(): """ - Regression for project-scoped keys being blocked by the team member budget: - when a key carries a project_id, the project budget governs and the team - member budget check must not raise, even if the member is over budget. + When a key carries a project_id, the project budget is enforced in addition + to, not instead of, the team member budget: an over-budget member must + still be blocked on a project-scoped key. """ request_body = { "model": "gpt-3.5-turbo", @@ -494,18 +494,21 @@ async def test_team_member_budget_check_skipped_for_project_scoped_key(): patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), patch("litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache), ): - result = await common_checks( - request_body=request_body, - team_object=team_object, - user_object=user_object, - end_user_object=None, - global_proxy_spend=None, - general_settings={}, - route="/chat/completions", - llm_router=None, - proxy_logging_obj=mock_proxy_logging_obj, - valid_token=valid_token, - request=mock_request, - ) + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await common_checks( + request_body=request_body, + team_object=team_object, + user_object=user_object, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=mock_proxy_logging_obj, + valid_token=valid_token, + request=mock_request, + ) - assert result is True + assert "Budget has been exceeded" in str(exc_info.value) + assert "test-user-1" in str(exc_info.value) + assert "test-team-1" in str(exc_info.value) diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 482ba03d2d6..a9db89fe232 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -5762,8 +5762,8 @@ async def test_temp_budget_increase_applied_for_cached_key(): class TestTeamMemberBudgetCheckForKey: """Tests for _team_member_budget_check_for_key: virtual keys with a - team member budget must be blocked once over budget, except project-scoped - keys, which are governed by the project budget instead.""" + team member budget must be blocked once over budget, including + project-scoped keys, which enforce the project budget in addition.""" def _valid_token(self, project_id=None) -> UserAPIKeyAuth: return UserAPIKeyAuth( @@ -5806,24 +5806,22 @@ class TestTeamMemberBudgetCheckForKey: ) @pytest.mark.asyncio - async def test_skips_for_project_scoped_key(self): + async def test_raises_for_project_scoped_key(self): from litellm.proxy.auth.user_api_key_auth import ( _team_member_budget_check_for_key, ) - cache = self._cache_with_membership() with patch( "litellm.proxy.proxy_server.get_current_spend", new_callable=AsyncMock, return_value=0.03, ): - await _team_member_budget_check_for_key( - valid_token=self._valid_token(project_id="project-1"), - prisma_client=MagicMock(), - user_api_key_cache=cache, - ) - - cache.async_get_cache.assert_not_called() + with pytest.raises(litellm.BudgetExceededError): + await _team_member_budget_check_for_key( + valid_token=self._valid_token(project_id="project-1"), + prisma_client=MagicMock(), + user_api_key_cache=self._cache_with_membership(), + ) async def _proxy_exception_for_key( diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 400b5864392..9a1f8277da3 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2586,9 +2586,9 @@ async def test_streaming_slow_path_processes_and_yields_chunk(spend_counter_stat @pytest.mark.asyncio -async def test_team_member_budget_counter_skipped_for_project_scoped_key(): - """Project-scoped keys are governed by the project budget, so no team - member budget counter should be reserved for them.""" +async def test_team_member_budget_counter_reserved_for_project_scoped_key(): + """Project-scoped keys enforce the project budget in addition to the team + member budget, so both counters must be reserved for them.""" from litellm.proxy.spend_tracking.budget_reservation import ( _get_team_member_budget_counter, ) @@ -2625,7 +2625,8 @@ async def test_team_member_budget_counter_skipped_for_project_scoped_key(): user_object=user_object, user_api_key_cache=cache, ) - assert counter is None + assert counter is not None + assert counter.counter_key == "spend:team_member:member-user:member-team" @pytest.mark.asyncio