diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 7866a226ca4..bc51f8e6b19 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -91,7 +91,6 @@ from litellm.proxy.utils import ( ProxyLogging, normalize_route_for_root_path, ) -from litellm.repositories.table_repositories import TeamMembershipRepository from litellm.secret_managers.main import get_secret_bool from litellm.types.services import ServiceTypes @@ -1835,58 +1834,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: - if prisma_client is not None: - _cache_key: Final = f"{valid_token.team_id}_{valid_token.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 - _user_id: Final = valid_token.user_id - _team_id: Final = valid_token.team_id - - if _user_id is not None and _team_id is not None: - _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.dict()) - 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.max_budget - 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, - ) - if team_member_spend > team_member_budget: - raise litellm.BudgetExceededError( - current_cost=team_member_spend, - max_budget=team_member_budget, - entity_type=Litellm_EntityType.TEAM_MEMBER.value, - entity_id=f"{valid_token.user_id}:{valid_token.team_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_team_member_budget.py b/tests/test_litellm/proxy/auth/test_team_member_budget.py index 3dcd65ec46f..99dcb912383 100644 --- a/tests/test_litellm/proxy/auth/test_team_member_budget.py +++ b/tests/test_litellm/proxy/auth/test_team_member_budget.py @@ -437,3 +437,51 @@ async def test_team_member_budget_check_personal_key_not_team(): assert result is True mock_get_team_membership.assert_not_called() + +@pytest.mark.asyncio +async def test_team_member_budget_not_enforced_on_management_routes(): + """ + An exhausted team member budget must not lock the member out of the management API. + + Regression: the same check was implemented a second time inline in user_api_key_auth, + outside the route gating in common_checks, so /key/info answered 429 and the member + could not see which budget had stopped them. + """ + team_object = LiteLLM_TeamTable(team_id="test-team-1", team_alias="Test Team", spend=0.0, max_budget=None) + user_object = LiteLLM_UserTable(user_id="test-user-1", spend=0.0, max_budget=None) + valid_token = UserAPIKeyAuth( + token="test-token", + user_id="test-user-1", + team_id="test-team-1", + team_member_spend=99.0, + ) + team_membership = LiteLLM_TeamMembership( + user_id="test-user-1", + team_id="test-team-1", + spend=99.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), + ) + + with ( + patch( + "litellm.proxy.auth.auth_checks.get_team_membership", + new_callable=AsyncMock, + return_value=team_membership, + ), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + ): + for route in ("/key/info", "/key/test-token/budgets"): + await common_checks( + request_body={}, + team_object=team_object, + user_object=user_object, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route=route, + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=valid_token, + request=MagicMock(spec=Request), + ) 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 46e02b14114..1fd44e45d5e 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 @@ -6228,3 +6228,115 @@ async def test_unlicensed_jwt_auth_is_forbidden_not_unauthorized(): assert error.code == "403" assert "enterprise" in error.message.lower() + +class _MembershipCache: + """Serves the team-membership lookups both budget check implementations perform.""" + + def __init__(self, values): + self.values = values + + async def async_get_cache(self, key, **kwargs): + return self.values.get(key) + + async def async_set_cache(self, key, value, **kwargs): + self.values[key] = value + + def get_cache(self, key, **kwargs): + return self.values.get(key) + + def set_cache(self, key, value, **kwargs): + self.values[key] = value + + async def async_batch_get_cache(self, keys, **kwargs): + return [self.values.get(key) for key in keys] + + def delete_cache(self, key, **kwargs): + self.values.pop(key, None) + + async def async_delete_cache(self, key, **kwargs): + self.values.pop(key, None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["/key/info", "/key/some-key-hash/budgets"]) +async def test_exhausted_team_member_budget_does_not_block_management_routes(route): + """ + A member who has spent their in-team budget must still be able to read key management routes. + + Regression: the team member budget check was implemented twice. The copy inside common_checks + is correctly skipped for non-LLM routes, but a second inline copy in _user_api_key_auth_builder + ran on every route, so the member got a 429 from /key/info and could not find out which budget + had stopped them. LLM routes are still blocked by the surviving check in common_checks. + """ + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj + + api_key = "sk-team-member-out-of-budget" + valid_token = UserAPIKeyAuth( + api_key=api_key, + token="hashed-team-member-key", + user_id="member-user", + team_id="member-team", + user_role=LitellmUserRoles.INTERNAL_USER, + team_member_spend=99.0, + ) + membership = LiteLLM_TeamMembership( + user_id="member-user", + team_id="member-team", + spend=99.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0), + ) + cache = _MembershipCache( + { + "member-team_member-user": membership, + "team_membership:member-user:member-team": membership, + } + ) + + request = Request(scope={"type": "http"}) + request._url = URL(url=route) + + proxy_logging_obj = MagicMock() + proxy_logging_obj.internal_usage_cache.dual_cache.async_get_cache = AsyncMock(return_value=None) + proxy_logging_obj.internal_usage_cache.dual_cache.async_set_cache = AsyncMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + + with ( + patch( + "litellm.proxy.auth.resolvers.store.IdentityStore._resolve_key", + new_callable=AsyncMock, + return_value=valid_token, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + return_value=LiteLLM_TeamTableCachedObj(team_id="member-team", team_alias="Member Team"), + ), + patch( + "litellm.proxy.auth.user_api_key_auth.get_user_object", + new_callable=AsyncMock, + return_value=LiteLLM_UserTable(user_id="member-user", spend=0.0), + ), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), + patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj), + patch("litellm.proxy.proxy_server.master_key", "sk-master-key"), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.user_custom_auth", None), + patch("litellm.proxy.proxy_server.jwt_handler", None), + ): + result = await _user_api_key_auth_builder( + request=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={}, + ) + + assert result.user_id == "member-user" + assert result.team_id == "member-team"