From 92ecef29b5731b2842494d7ab102f44d99e3d8ce Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 24 Sep 2026 09:51:23 -0700 Subject: [PATCH] fix(auth): enforce and alert on team member budgets only in common_checks The builder re-checked the team member budget inline before common_checks ran the same check, so one request that crossed a team_member_max_budget_alert_emails threshold dispatched two alerts. Drop the inline check; common_checks is the single authorization point and already covers per-member rows, the team default member budget, zero-cost skips and the cross-pod spend counter. Its 422 message now uses the TeamMember=user:team form the builder and budget reservation already returned. --- litellm/proxy/auth/auth_checks.py | 8 +- litellm/proxy/auth/user_api_key_auth.py | 71 -------------- .../proxy/auth/test_user_api_key_auth.py | 95 +++++++++++-------- 3 files changed, 59 insertions(+), 115 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index dd658148b47..4c5dd4cde94 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5646,12 +5646,16 @@ async def _check_team_member_budget( ) if team_member_spend >= team_member_budget: + entity_id: Final = f"{valid_token.user_id}:{team_object.team_id}" raise litellm.BudgetExceededError( current_cost=team_member_spend, max_budget=team_member_budget, - message=f"Budget has been exceeded! User={valid_token.user_id} in Team={team_object.team_id} Current cost: {team_member_spend}, Max budget: {team_member_budget}", + message=( + f"Budget has been exceeded! TeamMember={entity_id} " + f"Current cost: {team_member_spend}, Max budget: {team_member_budget}" + ), entity_type=Litellm_EntityType.TEAM_MEMBER.value, - entity_id=f"{valid_token.user_id}:{team_object.team_id}", + entity_id=entity_id, ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 5d004974635..d1523861492 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -50,7 +50,6 @@ from litellm.proxy.auth.auth_checks import ( _get_user_role, _is_model_cost_zero, _is_user_proxy_admin, - _team_member_max_budget_alert_check, _virtual_key_max_budget_alert_check, _virtual_key_max_budget_check, _virtual_key_soft_budget_check, @@ -119,7 +118,6 @@ from litellm.proxy.common_utils.model_listing_utils import claude_code_requested from litellm.proxy.common_utils.realtime_utils import _realtime_request_body from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, - team_membership_auth_cache_key, ) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup @@ -134,7 +132,6 @@ from litellm.proxy.utils import ( ProxyLogging, normalize_route_for_root_path, ) -from litellm.repositories.table_repositories import TeamMembershipRepository from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.secret_managers.main import get_secret_bool from litellm.types.services import ServiceTypes @@ -2239,74 +2236,6 @@ async def _user_api_key_auth_builder( if skip_budget_checks: verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model) - # Check 3. Check if user is in their team budget - if not skip_budget_checks and valid_token.team_member_spend is not None: - _user_id: Final = valid_token.user_id - _team_id: Final = valid_token.team_id - if prisma_client is not None and _user_id is not None and _team_id is not None: - _cache_key: Final = team_membership_auth_cache_key(team_id=_team_id, user_id=_user_id) - - team_member_info = await user_api_key_cache.async_get_cache( - key=_cache_key, - model_type=LiteLLM_TeamMembership, - ) - if team_member_info is None: - # read from DB - _db_member: Final = await TeamMembershipRepository(prisma_client).table.find_first( - where={ - "user_id": _user_id, - "team_id": _team_id, - }, - include={"litellm_budget_table": True}, - ) - if _db_member is not None: - team_member_info = LiteLLM_TeamMembership(**_db_member.model_dump()) - await user_api_key_cache.async_set_cache( - key=_cache_key, - value=team_member_info, - model_type=LiteLLM_TeamMembership, - ttl=5, - ) - - if team_member_info is not None and team_member_info.litellm_budget_table is not None: - team_member_budget: Final = team_member_info.litellm_budget_table.effective_max_budget( - now=datetime.now(timezone.utc), - ) - if team_member_budget is not None and team_member_budget > 0: - # Read from cross-pod counter (Redis-first) if available - from litellm.proxy.proxy_server import get_current_spend - - team_member_spend = valid_token.team_member_spend - if valid_token.user_id is not None and valid_token.team_id is not None: - team_member_spend = await get_current_spend( - counter_key=f"spend:team_member:{valid_token.user_id}:{valid_token.team_id}", - fallback_spend=team_member_spend, - max_budget=team_member_budget, - ) - _team_member_max_budget_alert_check( - team_id=_team_id, - team_alias=valid_token.team_alias, - team_metadata=valid_token.team_metadata, - organization_id=valid_token.org_id, - user_id=_user_id, - user_email=user_obj.user_email if user_obj is not None else None, - proxy_logging_obj=proxy_logging_obj, - spend=team_member_spend, - max_budget=team_member_budget, - ) - if team_member_spend >= team_member_budget: - _entity_id: Final = f"{valid_token.user_id}:{valid_token.team_id}" - raise litellm.BudgetExceededError( - current_cost=team_member_spend, - max_budget=team_member_budget, - message=( - f"Budget has been exceeded! TeamMember={_entity_id} " - f"Current cost: {team_member_spend}, Max budget: {team_member_budget}" - ), - entity_type=Litellm_EntityType.TEAM_MEMBER.value, - entity_id=_entity_id, - ) - # Check 3. If token is expired if valid_token.expires is not None: current_time = datetime.now(timezone.utc) 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 8d50ba5062f..b70e3a55326 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 @@ -7855,6 +7855,30 @@ async def test_temp_budget_increase_applied_for_cached_key(): assert cached_after.max_budget == 2.0 +async def _authenticate_and_authorize(mock_request, api_key): + """Builder then the single common_checks gate, the same sequence user_api_key_auth runs.""" + from litellm.proxy.auth.user_api_key_auth import _authorize_authenticated_request + + request_data = {"model": "claude-sonnet-5", "messages": [{"role": "user", "content": "hi"}]} + auth_obj = await _user_api_key_auth_builder( + request=mock_request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data=request_data, + ) + recovered = await _authorize_authenticated_request( + user_api_key_auth_obj=auth_obj, + request=mock_request, + request_data=request_data, + route="/v1/messages", + api_key=f"Bearer {api_key}", + ) + return recovered or auth_obj + + @pytest.mark.asyncio @pytest.mark.parametrize( "team_member_spend, expect_blocked", @@ -7868,7 +7892,7 @@ async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spe """A team member counter sitting exactly at the cap (where a resized reservation lands it) must be rejected by the cached-key auth path like every other budget check.""" from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj - from litellm.proxy.common_utils.user_api_key_cache import team_membership_auth_cache_key + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key from litellm.proxy.utils import hash_token api_key = "sk-team-member-exact-cap" @@ -7898,7 +7922,7 @@ async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spe value=LiteLLM_UserTable(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER), ) await user_api_key_cache.async_set_cache( - key=team_membership_auth_cache_key(team_id=team_id, user_id=user_id), + key=team_membership_reservation_cache_key(team_id=team_id, user_id=user_id), value=LiteLLM_TeamMembership( user_id=user_id, team_id=team_id, @@ -7920,15 +7944,7 @@ async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spe proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) async def _auth(): - return await _user_api_key_auth_builder( - request=mock_request, - api_key=f"Bearer {api_key}", - azure_api_key_header="", - anthropic_api_key_header=None, - google_ai_studio_api_key_header=None, - azure_apim_header=None, - request_data={"model": "claude-sonnet-5", "messages": [{"role": "user", "content": "hi"}]}, - ) + return await _authenticate_and_authorize(mock_request, api_key) with ( patch( # test-quality-ok: the builder reads proxy settings from module globals, no injection seam @@ -7973,7 +7989,10 @@ async def test_cached_key_team_member_budget_emails_configured_thresholds( """The team's team_member_max_budget_alert_emails thresholds fire from the cached-key auth path, including on the request that trips the hard cap, and stay silent below the lowest threshold.""" from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj - from litellm.proxy.common_utils.user_api_key_cache import team_membership_auth_cache_key + from litellm.proxy.common_utils.user_api_key_cache import ( + team_membership_auth_cache_key, + team_membership_reservation_cache_key, + ) from litellm.proxy.utils import hash_token api_key = "sk-team-member-alert-thresholds" @@ -7998,7 +8017,11 @@ async def test_cached_key_team_member_budget_emails_configured_thresholds( ) await user_api_key_cache.async_set_cache( key=f"team_id:{team_id}", - value=LiteLLM_TeamTableCachedObj(team_id=team_id), + value=LiteLLM_TeamTableCachedObj( + team_id=team_id, + team_alias="platform", + metadata={"team_member_max_budget_alert_emails": alert_emails}, + ), ) await user_api_key_cache.async_set_cache( key=user_id, @@ -8006,16 +8029,20 @@ async def test_cached_key_team_member_budget_emails_configured_thresholds( user_id=user_id, user_email="member@example.com", user_role=LitellmUserRoles.INTERNAL_USER ), ) - await user_api_key_cache.async_set_cache( - key=team_membership_auth_cache_key(team_id=team_id, user_id=user_id), - value=LiteLLM_TeamMembership( - user_id=user_id, - team_id=team_id, - spend=team_member_spend, - budget_id="budget-alert-thresholds", - litellm_budget_table=LiteLLM_BudgetTable(max_budget=2.4), - ), + membership = LiteLLM_TeamMembership( + user_id=user_id, + team_id=team_id, + spend=team_member_spend, + budget_id="budget-alert-thresholds", + litellm_budget_table=LiteLLM_BudgetTable(max_budget=2.4), ) + # A live proxy holds the row under both keys, so any second team-member check in the + # auth flow would find it too and send a duplicate alert. + for membership_cache_key in ( + team_membership_reservation_cache_key(team_id=team_id, user_id=user_id), + team_membership_auth_cache_key(team_id=team_id, user_id=user_id), + ): + await user_api_key_cache.async_set_cache(key=membership_cache_key, value=membership) mock_request = MagicMock() mock_request.url.path = "/v1/messages" @@ -8030,15 +8057,7 @@ async def test_cached_key_team_member_budget_emails_configured_thresholds( proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock(return_value=None) async def _auth(): - return await _user_api_key_auth_builder( - request=mock_request, - api_key=f"Bearer {api_key}", - azure_api_key_header="", - anthropic_api_key_header=None, - google_ai_studio_api_key_header=None, - azure_apim_header=None, - request_data={"model": "claude-sonnet-5", "messages": [{"role": "user", "content": "hi"}]}, - ) + return await _authenticate_and_authorize(mock_request, api_key) with ( patch( # test-quality-ok: the builder reads proxy settings from module globals, no injection seam @@ -8089,7 +8108,7 @@ async def test_cached_key_team_member_budget_honours_temp_increase(expiry_offset """A member over their permanent cap is admitted while a temp_budget_increase is unexpired and blocked again once it expires, on the cached-key auth path.""" from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj - from litellm.proxy.common_utils.user_api_key_cache import team_membership_auth_cache_key + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key from litellm.proxy.utils import hash_token api_key = "sk-team-member-temp-budget" @@ -8119,7 +8138,7 @@ async def test_cached_key_team_member_budget_honours_temp_increase(expiry_offset value=LiteLLM_UserTable(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER), ) await user_api_key_cache.async_set_cache( - key=team_membership_auth_cache_key(team_id=team_id, user_id=user_id), + key=team_membership_reservation_cache_key(team_id=team_id, user_id=user_id), value=LiteLLM_TeamMembership( user_id=user_id, team_id=team_id, @@ -8145,15 +8164,7 @@ async def test_cached_key_team_member_budget_honours_temp_increase(expiry_offset proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) async def _auth(): - return await _user_api_key_auth_builder( - request=mock_request, - api_key=f"Bearer {api_key}", - azure_api_key_header="", - anthropic_api_key_header=None, - google_ai_studio_api_key_header=None, - azure_apim_header=None, - request_data={"model": "claude-sonnet-5", "messages": [{"role": "user", "content": "hi"}]}, - ) + return await _authenticate_and_authorize(mock_request, api_key) with ( patch( # test-quality-ok: the builder reads proxy settings from module globals, no injection seam