diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index fb600cee26b..3569e45bf30 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -45,6 +45,9 @@ class ResetBudgetJob: ## Reset Team Budget await self.reset_budget_for_litellm_teams() + ### RESET TEAM MEMBER BUDGET (based on team budget duration) ### + await self.reset_budget_for_team_members_by_team() + ### RESET ENDUSER (Customer) BUDGET and corresponding Budget duration ### await self.reset_budget_for_litellm_budget_table() @@ -481,6 +484,85 @@ class ResetBudgetJob: ) verbose_proxy_logger.exception("Failed to reset budget for teams: %s", e) + async def reset_budget_for_team_members_by_team(self): + """ + Resets the budget for team members based on their team's budget_duration. + + This fixes issue #19105 where team member budgets don't reset because they + don't have their own budget_duration set - they should inherit from the team. + """ + now = datetime.utcnow() + start_time = time.time() + teams_to_process: Optional[List[LiteLLM_TeamTable]] = None + try: + # Get all teams that need budget reset + teams_to_process = await self.prisma_client.get_data( + table_name="team", query_type="find_all", reset_at=now + ) + + if teams_to_process is not None and len(teams_to_process) > 0: + team_ids_to_reset = [ + team.team_id + for team in teams_to_process + if team.team_id is not None + ] + + verbose_proxy_logger.debug( + "Resetting team member budgets for teams: %s", team_ids_to_reset + ) + + # Reset spend for all team members in these teams + result = await self.prisma_client.db.litellm_teammembership.update_many( + where={ + "team_id": {"in": team_ids_to_reset} + }, + data={ + "spend": 0, + }, + ) + + verbose_proxy_logger.debug( + "Reset %s team member budgets", result + ) + + end_time = time.time() + asyncio.create_task( + self.proxy_logging_obj.service_logging_obj.async_service_success_hook( + service=ServiceTypes.RESET_BUDGET_JOB, + duration=end_time - start_time, + call_type="reset_budget_team_members_by_team", + start_time=start_time, + end_time=end_time, + event_metadata={ + "num_teams_found": len(teams_to_process) if teams_to_process else 0, + "teams_found": json.dumps( + teams_to_process, indent=4, default=str + ), + }, + ) + ) + except Exception as e: + end_time = time.time() + asyncio.create_task( + self.proxy_logging_obj.service_logging_obj.async_service_failure_hook( + service=ServiceTypes.RESET_BUDGET_JOB, + duration=end_time - start_time, + error=e, + call_type="reset_budget_team_members_by_team", + start_time=start_time, + end_time=end_time, + event_metadata={ + "num_teams_found": len(teams_to_process) if teams_to_process else 0, + "teams_found": json.dumps( + teams_to_process, indent=4, default=str + ), + }, + ) + ) + verbose_proxy_logger.exception( + "Failed to reset budget for team members by team: %s", e + ) + @staticmethod async def _reset_budget_common( item: Union[LiteLLM_TeamTable, LiteLLM_UserTable, LiteLLM_VerificationToken], diff --git a/reproduce_issue_19105.py b/reproduce_issue_19105.py new file mode 100644 index 00000000000..f1e9f9359b4 --- /dev/null +++ b/reproduce_issue_19105.py @@ -0,0 +1,178 @@ +""" +Reproduction script for issue #19105: Team member budgets not enforced + +This script demonstrates the bug where: +1. Team member budget (max_budget_in_team) is not enforced - requests go through even when exceeded +2. Team member spend doesn't reset daily while team spend does + +Setup: +- Create a team with max_budget and budget_duration +- Add a key to the team +- Add a user with max_budget_in_team +- Make requests to exceed the user's budget +- Verify requests still go through (BUG) +- Trigger budget reset and verify team member spend doesn't reset (BUG) +""" + +import asyncio +import os +import sys +from datetime import datetime, timedelta, timezone + +# Add litellm to path +sys.path.insert(0, os.path.abspath(os.path.dirname(__file__))) + +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob + + +async def reproduce_issue(): + print("=" * 80) + print("REPRODUCING ISSUE #19105: Team Member Budget Not Enforced") + print("=" * 80) + + # Initialize Prisma client + prisma_client = PrismaClient() + await prisma_client.connect() + + try: + # Step 1: Create a team with budget and budget_duration + print("\n1. Creating team with budget...") + team = await prisma_client.db.litellm_teamtable.create( + data={ + "team_alias": "test_team_19105", + "max_budget": 1.0, + "budget_duration": "1d", # Daily reset + "spend": 0.0, + } + ) + print(f" Created team: {team.team_id}, budget: ${team.max_budget}, duration: {team.budget_duration}") + + # Step 2: Create a key for the team + print("\n2. Creating key for team...") + key = await prisma_client.db.litellm_verificationtoken.create( + data={ + "token": f"sk-test-{team.team_id}", + "team_id": team.team_id, + } + ) + print(f" Created key: {key.token}") + + # Step 3: Create a budget table for team member + print("\n3. Creating budget table for team member...") + budget = await prisma_client.db.litellm_budgettable.create( + data={ + "max_budget": 0.01, # Very low budget to trigger limit + # Note: NOT setting budget_duration - this is the bug! + } + ) + print(f" Created budget: budget_id={budget.budget_id}, max_budget=${budget.max_budget}") + print(f" ⚠️ Budget duration: {budget.budget_duration} (None - this is the problem!)") + + # Step 4: Create a user (internal user) + print("\n4. Creating internal user...") + user = await prisma_client.db.litellm_usertable.create( + data={ + "user_id": "test_user_19105", + "user_email": "test@example.com", + } + ) + print(f" Created user: {user.user_id}") + + # Step 5: Create team membership linking user to team with budget + print("\n5. Creating team membership with max_budget_in_team...") + membership = await prisma_client.db.litellm_teammembership.create( + data={ + "team_id": team.team_id, + "user_id": user.user_id, + "budget_id": budget.budget_id, + "spend": 0.015, # Already over budget! + } + ) + print(f" Created membership: user={user.user_id}, team={team.team_id}") + print(f" Member spend: ${membership.spend}, Member budget: ${budget.max_budget}") + print(f" ⚠️ Spend ${membership.spend} > Budget ${budget.max_budget} - should be blocked!") + + # Step 6: Check if budget enforcement would work + print("\n6. Testing budget check logic...") + membership_with_budget = await prisma_client.db.litellm_teammembership.find_unique( + where={"user_id_team_id": {"user_id": user.user_id, "team_id": team.team_id}}, + include={"litellm_budget_table": True}, + ) + + if membership_with_budget and membership_with_budget.litellm_budget_table: + team_member_budget = membership_with_budget.litellm_budget_table.max_budget + team_member_spend = membership_with_budget.spend or 0.0 + + print(f" Spend: ${team_member_spend}, Budget: ${team_member_budget}") + if team_member_spend >= team_member_budget: + print(f" ✓ Budget check WOULD block request (spend >= budget)") + else: + print(f" ✗ Budget check would NOT block request") + + # Step 7: Test budget reset logic + print("\n7. Testing budget reset logic...") + print(f" Current team member spend: ${membership.spend}") + print(f" Budget reset_at: {budget.budget_reset_at}") + + # Create ProxyLogging object (needed for ResetBudgetJob) + proxy_logging_obj = ProxyLogging(user_api_key_cache=None) + reset_job = ResetBudgetJob( + proxy_logging_obj=proxy_logging_obj, + prisma_client=prisma_client, + ) + + # Manually set budget reset_at to past to trigger reset + now = datetime.now(timezone.utc) + past_time = now - timedelta(hours=1) + await prisma_client.db.litellm_budgettable.update( + where={"budget_id": budget.budget_id}, + data={"budget_reset_at": past_time}, + ) + print(f" Set budget reset_at to: {past_time} (in the past)") + + # Run budget reset + print(f"\n8. Running budget reset job...") + await reset_job.reset_budget() + + # Check if team member spend was reset + membership_after = await prisma_client.db.litellm_teammembership.find_unique( + where={"user_id_team_id": {"user_id": user.user_id, "team_id": team.team_id}}, + ) + + print(f" Team member spend after reset: ${membership_after.spend}") + if membership_after.spend == 0.0: + print(f" ✓ Team member spend WAS reset") + else: + print(f" ✗ BUG: Team member spend NOT reset (still ${membership_after.spend})") + print(f" This is because the budget table doesn't have a budget_duration set!") + + # Clean up + print("\n9. Cleaning up...") + await prisma_client.db.litellm_teammembership.delete( + where={"user_id_team_id": {"user_id": user.user_id, "team_id": team.team_id}} + ) + await prisma_client.db.litellm_usertable.delete(where={"user_id": user.user_id}) + await prisma_client.db.litellm_verificationtoken.delete(where={"token": key.token}) + await prisma_client.db.litellm_budgettable.delete(where={"budget_id": budget.budget_id}) + await prisma_client.db.litellm_teamtable.delete(where={"team_id": team.team_id}) + print(" Cleanup complete") + + print("\n" + "=" * 80) + print("CONCLUSION:") + print("=" * 80) + print("Team member budgets don't reset because:") + print("1. Team members use litellm_budget_table for budget limits") + print("2. Budget resets only happen if litellm_budget_table has budget_duration set") + print("3. When creating team members with max_budget_in_team, the budget_duration") + print(" is NOT automatically inherited from the team") + print("4. Without budget_duration, the budget never resets") + print("\nFIX: Team member budgets should inherit budget_duration from their team") + print("=" * 80) + + finally: + await prisma_client.disconnect() + + +if __name__ == "__main__": + asyncio.run(reproduce_issue()) diff --git a/tests/proxy_unit_tests/test_team_member_budget_reset.py b/tests/proxy_unit_tests/test_team_member_budget_reset.py new file mode 100644 index 00000000000..979f736113f --- /dev/null +++ b/tests/proxy_unit_tests/test_team_member_budget_reset.py @@ -0,0 +1,175 @@ +""" +Test for issue #19105: Team member budget reset + +Tests that team member budgets reset based on the team's budget_duration, +not just their individual budget table's duration. +""" + +import asyncio +import pytest +import sys +import os +from datetime import datetime, timedelta, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +# Add project root to path +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +import litellm +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob + + +@pytest.mark.asyncio +async def test_team_member_budget_resets_with_team(): + """ + Test that team member budgets reset when the team's budget resets, + even if the member's budget table doesn't have its own budget_duration. + + This tests the fix for issue #19105. + """ + # Mock Prisma client + prisma_client = MagicMock(spec=PrismaClient) + prisma_client.db = MagicMock() + + # Mock ProxyLogging + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.service_logging_obj = MagicMock() + proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock() + proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock() + + # Create test data + now = datetime.utcnow() + past_time = now - timedelta(hours=1) + + # Team with budget_duration that needs reset + team = MagicMock() + team.team_id = "team_123" + team.max_budget = 10.0 + team.budget_duration = "1d" + team.budget_reset_at = past_time # In the past, needs reset + team.spend = 5.0 + + # Mock get_data to return the team + async def mock_get_data(table_name, query_type, **kwargs): + if table_name == "team" and "reset_at" in kwargs: + return [team] + return None + + prisma_client.get_data = AsyncMock(side_effect=mock_get_data) + + # Mock update_many for team membership + update_result = MagicMock() + update_result.count = 2 # 2 team members updated + prisma_client.db.litellm_teammembership.update_many = AsyncMock( + return_value=update_result + ) + + # Create reset job + reset_job = ResetBudgetJob( + proxy_logging_obj=proxy_logging_obj, prisma_client=prisma_client + ) + + # Run the team member reset + await reset_job.reset_budget_for_team_members_by_team() + + # Verify that update_many was called with correct parameters + prisma_client.db.litellm_teammembership.update_many.assert_called_once() + call_args = prisma_client.db.litellm_teammembership.update_many.call_args + + # Check the where clause includes the team_id + assert call_args.kwargs["where"]["team_id"] == {"in": ["team_123"]} + + # Check that spend is reset to 0 + assert call_args.kwargs["data"]["spend"] == 0 + + # Verify success hook was called + proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_called_once() + + +@pytest.mark.asyncio +async def test_team_member_budget_reset_no_teams(): + """ + Test that the function handles the case where no teams need resetting. + """ + # Mock Prisma client + prisma_client = MagicMock(spec=PrismaClient) + prisma_client.db = MagicMock() + + # Mock ProxyLogging + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.service_logging_obj = MagicMock() + proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock() + + # Mock get_data to return empty list (no teams need reset) + prisma_client.get_data = AsyncMock(return_value=[]) + + # Mock update_many + prisma_client.db.litellm_teammembership.update_many = AsyncMock() + + # Create reset job + reset_job = ResetBudgetJob( + proxy_logging_obj=proxy_logging_obj, prisma_client=prisma_client + ) + + # Run the team member reset + await reset_job.reset_budget_for_team_members_by_team() + + # Verify that update_many was NOT called (no teams to process) + prisma_client.db.litellm_teammembership.update_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_team_member_budget_reset_multiple_teams(): + """ + Test that team member budgets are reset for multiple teams. + """ + # Mock Prisma client + prisma_client = MagicMock(spec=PrismaClient) + prisma_client.db = MagicMock() + + # Mock ProxyLogging + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.service_logging_obj = MagicMock() + proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock() + + # Create test data - multiple teams + now = datetime.utcnow() + past_time = now - timedelta(hours=1) + + teams = [] + for i in range(3): + team = MagicMock() + team.team_id = f"team_{i}" + team.max_budget = 10.0 + team.budget_duration = "1d" + team.budget_reset_at = past_time + teams.append(team) + + # Mock get_data to return multiple teams + prisma_client.get_data = AsyncMock(return_value=teams) + + # Mock update_many + update_result = MagicMock() + update_result.count = 5 # 5 team members updated across all teams + prisma_client.db.litellm_teammembership.update_many = AsyncMock( + return_value=update_result + ) + + # Create reset job + reset_job = ResetBudgetJob( + proxy_logging_obj=proxy_logging_obj, prisma_client=prisma_client + ) + + # Run the team member reset + await reset_job.reset_budget_for_team_members_by_team() + + # Verify that update_many was called with all team IDs + call_args = prisma_client.db.litellm_teammembership.update_many.call_args + assert set(call_args.kwargs["where"]["team_id"]["in"]) == { + "team_0", + "team_1", + "team_2", + }