diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 07286b4fa80..b0ea6b41ac5 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -355,6 +355,7 @@ async def _upsert_budget_and_membership( tpm_limit: Optional[int] = None, rpm_limit: Optional[int] = None, allowed_models: Optional[List[str]] = None, + team_default_budget_id: Optional[str] = None, ): """ Helper function to Create/Update or Delete the budget within the team membership @@ -368,6 +369,11 @@ async def _upsert_budget_and_membership( tpm_limit: Tokens per minute limit for the team member rpm_limit: Requests per minute limit for the team member allowed_models: Per-member model scope. None = don't change. [] = remove restrictions. Non-empty list = enforce. + team_default_budget_id: The team's shared default member budget id (from + team metadata.team_member_budget_id), if any. When the membership's + existing_budget_id matches this, we clone-on-write so editing one + member's budget does not mutate the shared default (and therefore + every other member who still points at it). If max_budget, tpm_limit, rpm_limit, and allowed_models are all None, the user's budget is removed from the team membership. If any of these values exist, a budget is updated or created and linked to the team membership. @@ -385,7 +391,13 @@ async def _upsert_budget_and_membership( ) return - if existing_budget_id is not None: + is_shared_default = ( + existing_budget_id is not None + and team_default_budget_id is not None + and existing_budget_id == team_default_budget_id + ) + + if existing_budget_id is not None and not is_shared_default: # Update the existing budget in-place to preserve fields not being changed. # Only write fields that the caller explicitly provided (non-None). update_data: Dict[str, Any] = { @@ -405,11 +417,40 @@ async def _upsert_budget_and_membership( ) return - # No existing budget — create a new one and link it to the membership. + # Either there is no existing budget, OR the membership is still pointing + # at the team's shared default member budget. In both cases we create a + # NEW private budget for this user and (re)link the membership to it. create_data: Dict[str, Any] = { "created_by": user_api_key_dict.user_id or "", "updated_by": user_api_key_dict.user_id or "", } + + # If we're forking off the shared default, seed the new row with the + # default's values so fields the caller did not change carry over. + if is_shared_default: + default_budget_row = await tx.litellm_budgettable.find_unique( + where={"budget_id": existing_budget_id} + ) + if default_budget_row is not None: + default_budget_dict = default_budget_row.model_dump() + for field in ( + "max_budget", + "soft_budget", + "max_parallel_requests", + "tpm_limit", + "rpm_limit", + "model_max_budget", + "budget_duration", + "allowed_models", + ): + value = default_budget_dict.get(field) + if value is None: + continue + if isinstance(value, list) and len(value) == 0: + continue + create_data[field] = value + + # Caller-provided values take precedence over the cloned defaults. if max_budget is not None: create_data["max_budget"] = max_budget if tpm_limit is not None: diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 8e21b851857..bf912fba4f8 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1570,8 +1570,7 @@ async def update_team( # noqa: PLR0915 current_org_id = getattr(existing_team_row, "organization_id", None) if ( data.organization_id != current_org_id - and user_api_key_dict.user_role - != LitellmUserRoles.PROXY_ADMIN.value + and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value ): # Is the caller org_admin of the destination org? caller_memberships = ( @@ -2609,6 +2608,15 @@ async def team_member_update( identified_budget_id = tm.budget_id break + # If this membership still points at the team's shared default member + # budget, _upsert_budget_and_membership will clone-on-write so that the + # update only touches this user (not every member sharing the default). + team_default_budget_id: Optional[str] = None + if team_table.metadata is not None: + raw_default_budget_id = team_table.metadata.get("team_member_budget_id") + if isinstance(raw_default_budget_id, str): + team_default_budget_id = raw_default_budget_id + ### upsert new budget async with prisma_client.db.tx() as tx: await _upsert_budget_and_membership( @@ -2621,6 +2629,7 @@ async def team_member_update( tpm_limit=data.tpm_limit, rpm_limit=data.rpm_limit, allowed_models=data.allowed_models, + team_default_budget_id=team_default_budget_id, ) ### update team member role diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 3e42d392077..5cf53ae06f5 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -140,6 +140,62 @@ async def handle_budget_for_entity( return existing_budget_id +# Fields on LiteLLM_BudgetTable that represent the budget's *configuration* +# (i.e. the values an admin sets). We copy these when cloning a team's +# default member-budget into an individual member-budget so that the new +# row starts with the same limits as the default. +_CLONABLE_BUDGET_FIELDS: Tuple[str, ...] = ( + "max_budget", + "soft_budget", + "max_parallel_requests", + "tpm_limit", + "rpm_limit", + "model_max_budget", + "budget_duration", + "allowed_models", +) + + +async def _clone_team_default_budget_for_member( + prisma_client: PrismaClient, + default_team_budget_id: str, + user_api_key_dict: UserAPIKeyAuth, + litellm_proxy_admin_name: str, +) -> Optional[str]: + """ + Create a new budget row that copies the values from the team's default + member budget. Returns the new budget_id, or None if the default budget + no longer exists in the DB. + + Used when adding a new team member without an explicit per-member budget, + so the member starts with the team default's values but gets their own + private budget row (which can be edited independently). + """ + default_budget = await prisma_client.db.litellm_budgettable.find_unique( + where={"budget_id": default_team_budget_id} + ) + if default_budget is None: + return None + + default_budget_dict = default_budget.model_dump() + cloned_data: dict = { + "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + } + for field in _CLONABLE_BUDGET_FIELDS: + value = default_budget_dict.get(field) + if value is None: + continue + # Skip empty list defaults (e.g. allowed_models = []) so the cloned + # row matches the "no value set" shape rather than carrying a default. + if isinstance(value, list) and len(value) == 0: + continue + cloned_data[field] = value + + new_budget = await prisma_client.db.litellm_budgettable.create(data=cloned_data) + return new_budget.budget_id + + async def add_new_member( new_member: Member, max_budget_in_team: Optional[float], @@ -221,8 +277,20 @@ async def add_new_member( response = await prisma_client.db.litellm_budgettable.create(data=budget_data) _budget_id = response.budget_id + elif default_team_budget_id is not None: + # No per-member budget was provided, but the team has a default member + # budget. Clone the default budget into a new row for this user so that + # later edits to one member's budget do not bleed into other members. + # If the default no longer exists in the DB, fall back to no budget. + _budget_id = await _clone_team_default_budget_for_member( + prisma_client=prisma_client, + default_team_budget_id=default_team_budget_id, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ) else: - _budget_id = default_team_budget_id + # No per-member budget and no team default → member gets no budget. + _budget_id = None if _budget_id and returned_user is not None and returned_user.user_id is not None: _returned_team_membership = ( diff --git a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py index 8e511518892..f4bf0d7b2be 100644 --- a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py +++ b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py @@ -268,3 +268,104 @@ async def test_upsert_rpm_only_creates_new_budget(mock_tx, fake_user): }, }, ) + + +# TEST: clone-on-write when membership still points at the team's shared default budget +@pytest.mark.asyncio +async def test_upsert_clones_when_pointing_at_shared_default(mock_tx, fake_user): + """ + When a member's existing budget_id is the same row as the team's shared + default member budget, updating that member's budget must NOT mutate the + shared row. Instead we should create a new private budget for this member + (seeded with the default's values) and re-link the membership to it. + """ + shared_default_id = "team-default-budget-1" + + # Default budget row in the DB: $200 cap, daily reset, 500 tpm. + default_row = MagicMock() + default_row.model_dump.return_value = { + "budget_id": shared_default_id, + "max_budget": 200.0, + "soft_budget": None, + "max_parallel_requests": None, + "tpm_limit": 500, + "rpm_limit": None, + "model_max_budget": None, + "budget_duration": "1d", + "allowed_models": [], + } + mock_tx.litellm_budgettable.find_unique = AsyncMock(return_value=default_row) + + # Caller is changing only this member's max_budget. + await _upsert_budget_and_membership( + mock_tx, + team_id="team-shared", + user_id="user-shared", + max_budget=50.0, + existing_budget_id=shared_default_id, + user_api_key_dict=fake_user, + team_default_budget_id=shared_default_id, + ) + + # Must NOT touch the shared default row in place. + mock_tx.litellm_budgettable.update.assert_not_called() + + # Must create a new private budget seeded with the default's values, + # with the caller's max_budget overriding the cloned default. + mock_tx.litellm_budgettable.create.assert_awaited_once_with( + data={ + "created_by": fake_user.user_id, + "updated_by": fake_user.user_id, + "max_budget": 50.0, # caller wins + "tpm_limit": 500, # cloned from default + "budget_duration": "1d", # cloned from default + }, + include={"team_membership": True}, + ) + + # Membership must be re-linked to the new private budget. + new_budget_id = mock_tx.litellm_budgettable.create.return_value.budget_id + mock_tx.litellm_teammembership.upsert.assert_awaited_once_with( + where={"user_id_team_id": {"user_id": "user-shared", "team_id": "team-shared"}}, + data={ + "create": { + "user_id": "user-shared", + "team_id": "team-shared", + "litellm_budget_table": {"connect": {"budget_id": new_budget_id}}, + }, + "update": { + "litellm_budget_table": {"connect": {"budget_id": new_budget_id}}, + }, + }, + ) + + +# TEST: when team default exists but member already has their own budget, in-place update +@pytest.mark.asyncio +async def test_upsert_updates_in_place_when_member_has_private_budget( + mock_tx, fake_user +): + """ + If the member's budget_id is different from the team's shared default + (i.e. they already have a private budget), we should keep the current + in-place behavior and not allocate a new row. + """ + await _upsert_budget_and_membership( + mock_tx, + team_id="team-mixed", + user_id="user-private", + max_budget=75.0, + existing_budget_id="private-budget-xyz", + user_api_key_dict=fake_user, + team_default_budget_id="team-default-budget-1", + ) + + mock_tx.litellm_budgettable.update.assert_awaited_once_with( + where={"budget_id": "private-budget-xyz"}, + data={ + "max_budget": 75.0, + "updated_by": fake_user.user_id, + }, + ) + mock_tx.litellm_budgettable.create.assert_not_called() + mock_tx.litellm_teammembership.upsert.assert_not_called() diff --git a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py index c9828fc64f8..459072cf9d3 100644 --- a/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py @@ -20,14 +20,13 @@ from litellm.proxy.management_helpers.utils import add_new_member @pytest.mark.asyncio -async def test_add_new_member_uses_default_team_budget_id(): +async def test_add_new_member_clones_default_team_budget_id(): """ - Test that add_new_member uses the default_team_budget_id when max_budget_in_team is None. + Test that add_new_member CLONES the team's default member budget when + max_budget_in_team is None and a default_team_budget_id is provided. - This test verifies that: - 1. When max_budget_in_team is None - 2. And default_team_budget_id is provided - 3. The team membership is created with the default_team_budget_id + Cloning (rather than sharing the same budget row) is what lets admins later + edit one member's budget without mutating every other member's budget. """ from litellm.proxy._types import LitellmUserRoles @@ -35,17 +34,15 @@ async def test_add_new_member_uses_default_team_budget_id(): test_user_id = "test_user_123" test_team_id = "test_team_456" test_default_budget_id = "default_budget_789" + test_cloned_budget_id = "cloned_budget_xyz" test_admin_name = "test_admin" - # Create a Member object with user_id new_member = Member(user_id=test_user_id, role="user") - # Create UserAPIKeyAuth object user_api_key_dict = UserAPIKeyAuth( user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN ) - # Mock the prisma client mock_prisma_client = AsyncMock() # Mock the user table upsert operation @@ -60,56 +57,140 @@ async def test_add_new_member_uses_default_team_budget_id(): return_value=mock_user_response ) + # Mock the default budget row fetched for cloning. + mock_default_budget_row = MagicMock() + mock_default_budget_row.model_dump.return_value = { + "budget_id": test_default_budget_id, + "max_budget": 100.0, + "soft_budget": None, + "max_parallel_requests": None, + "tpm_limit": 1000, + "rpm_limit": None, + "model_max_budget": None, + "budget_duration": "1d", + "allowed_models": [], + } + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( + return_value=mock_default_budget_row + ) + + # Mock the cloned budget row that .create() returns. + mock_cloned_budget_row = MagicMock() + mock_cloned_budget_row.budget_id = test_cloned_budget_id + mock_prisma_client.db.litellm_budgettable.create = AsyncMock( + return_value=mock_cloned_budget_row + ) + # Mock the team membership creation mock_team_membership_response = MagicMock() mock_team_membership_response.model_dump.return_value = { "team_id": test_team_id, "user_id": test_user_id, - "budget_id": test_default_budget_id, + "budget_id": test_cloned_budget_id, "litellm_budget_table": None, } mock_prisma_client.db.litellm_teammembership.create = AsyncMock( return_value=mock_team_membership_response ) - # Call the function with max_budget_in_team=None and a default_team_budget_id result_user, result_team_membership = await add_new_member( new_member=new_member, - max_budget_in_team=None, # This is the key - no max budget specified + max_budget_in_team=None, prisma_client=mock_prisma_client, team_id=test_team_id, user_api_key_dict=user_api_key_dict, litellm_proxy_admin_name=test_admin_name, - default_team_budget_id=test_default_budget_id, # This should be used + default_team_budget_id=test_default_budget_id, ) - # Verify that the user was created/updated correctly assert result_user is not None assert result_user.user_id == test_user_id - # Verify that the team membership was created correctly + # Membership should be linked to the new cloned budget, not the shared default. assert result_team_membership is not None - assert result_team_membership.team_id == test_team_id - assert result_team_membership.user_id == test_user_id - assert result_team_membership.budget_id == test_default_budget_id + assert result_team_membership.budget_id == test_cloned_budget_id + assert result_team_membership.budget_id != test_default_budget_id - # Verify that the prisma client methods were called correctly mock_prisma_client.db.litellm_usertable.upsert.assert_called_once() mock_prisma_client.db.litellm_teammembership.create.assert_called_once() - # Verify that no budget table creation was called (since max_budget_in_team is None) - assert ( - not hasattr(mock_prisma_client.db, "litellm_budgettable") - or not mock_prisma_client.db.litellm_budgettable.create.called + # The clone must have happened: find_unique on the default, create for the clone. + mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( + where={"budget_id": test_default_budget_id} ) + mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + cloned_create_data = ( + mock_prisma_client.db.litellm_budgettable.create.call_args.kwargs["data"] + ) + # Cloned values from the default budget row + assert cloned_create_data["max_budget"] == 100.0 + assert cloned_create_data["tpm_limit"] == 1000 + assert cloned_create_data["budget_duration"] == "1d" + assert cloned_create_data["created_by"] == user_api_key_dict.user_id - # Verify the team membership was created with the correct budget_id team_membership_call_args = ( mock_prisma_client.db.litellm_teammembership.create.call_args ) - assert team_membership_call_args is not None create_data = team_membership_call_args.kwargs["data"] - assert create_data["budget_id"] == test_default_budget_id + assert create_data["budget_id"] == test_cloned_budget_id + + +@pytest.mark.asyncio +async def test_add_new_member_no_budget_when_no_default_and_no_max_budget(): + """ + Test that add_new_member links no budget to the team membership when + neither max_budget_in_team nor default_team_budget_id is provided. + + When the team has no default member budget, new members get nothing. + """ + from litellm.proxy._types import LitellmUserRoles + + test_user_id = "test_user_no_budget" + test_team_id = "test_team_no_budget" + test_admin_name = "test_admin" + + new_member = Member(user_id=test_user_id, role="user") + + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + mock_prisma_client = AsyncMock() + + mock_user_response = MagicMock() + mock_user_response.model_dump.return_value = { + "user_id": test_user_id, + "user_email": None, + "teams": [test_team_id], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + return_value=mock_user_response + ) + + # Even though we mock these, they must NOT be called on the no-budget path. + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock() + mock_prisma_client.db.litellm_budgettable.create = AsyncMock() + mock_prisma_client.db.litellm_teammembership.create = AsyncMock() + + result_user, result_team_membership = await add_new_member( + new_member=new_member, + max_budget_in_team=None, + prisma_client=mock_prisma_client, + team_id=test_team_id, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=test_admin_name, + default_team_budget_id=None, + ) + + assert result_user is not None + assert result_user.user_id == test_user_id + + # No budget id, so no team membership row is created. + assert result_team_membership is None + mock_prisma_client.db.litellm_budgettable.find_unique.assert_not_called() + mock_prisma_client.db.litellm_budgettable.create.assert_not_called() + mock_prisma_client.db.litellm_teammembership.create.assert_not_called() @pytest.mark.asyncio @@ -206,38 +287,30 @@ async def test_add_new_member_creates_new_budget_when_max_budget_provided(): @pytest.mark.asyncio -async def test_add_new_member_with_user_email(): +async def test_add_new_member_with_user_email_clones_default_budget(): """ - Test add_new_member with user_email instead of user_id and default budget. - - This test verifies that: - 1. When new_member has user_email instead of user_id - 2. And max_budget_in_team is None - 3. The default_team_budget_id is used correctly + Test add_new_member with user_email instead of user_id and a team default + budget. The default budget should be CLONED into a new private row for + this user, not shared with other members of the team. """ from litellm.proxy._types import LitellmUserRoles - # Setup test data test_user_email = "test@example.com" test_team_id = "test_team_456" test_default_budget_id = "default_budget_789" + test_cloned_budget_id = "cloned_budget_for_email_user" test_admin_name = "test_admin" - # Create a Member object with user_email new_member = Member(user_email=test_user_email, role="user") - # Create UserAPIKeyAuth object user_api_key_dict = UserAPIKeyAuth( user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN ) - # Mock the prisma client mock_prisma_client = AsyncMock() - # Mock get_data to return empty list (no existing user) mock_prisma_client.get_data = AsyncMock(return_value=[]) - # Mock insert_data for new user creation mock_user_response = MagicMock() mock_user_response.model_dump.return_value = { "user_id": "generated_user_id", @@ -247,19 +320,41 @@ async def test_add_new_member_with_user_email(): } mock_prisma_client.insert_data = AsyncMock(return_value=mock_user_response) - # Mock the team membership creation + # Default budget that will be cloned + mock_default_budget_row = MagicMock() + mock_default_budget_row.model_dump.return_value = { + "budget_id": test_default_budget_id, + "max_budget": 25.0, + "soft_budget": None, + "max_parallel_requests": None, + "tpm_limit": None, + "rpm_limit": None, + "model_max_budget": None, + "budget_duration": None, + "allowed_models": [], + } + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( + return_value=mock_default_budget_row + ) + + # Cloned budget result + mock_cloned_budget_row = MagicMock() + mock_cloned_budget_row.budget_id = test_cloned_budget_id + mock_prisma_client.db.litellm_budgettable.create = AsyncMock( + return_value=mock_cloned_budget_row + ) + mock_team_membership_response = MagicMock() mock_team_membership_response.model_dump.return_value = { "team_id": test_team_id, "user_id": "generated_user_id", - "budget_id": test_default_budget_id, + "budget_id": test_cloned_budget_id, "litellm_budget_table": None, } mock_prisma_client.db.litellm_teammembership.create = AsyncMock( return_value=mock_team_membership_response ) - # Call the function result_user, result_team_membership = await add_new_member( new_member=new_member, max_budget_in_team=None, @@ -270,28 +365,31 @@ async def test_add_new_member_with_user_email(): default_team_budget_id=test_default_budget_id, ) - # Verify that the user was created correctly assert result_user is not None assert result_user.user_email == test_user_email - # Verify that the team membership was created with the default budget_id + # Membership should point at the cloned (private) budget, not the shared default. assert result_team_membership is not None - assert result_team_membership.budget_id == test_default_budget_id + assert result_team_membership.budget_id == test_cloned_budget_id - # Verify that get_data was called to check for existing user mock_prisma_client.get_data.assert_called_once_with( key_val={"user_email": test_user_email}, table_name="user", query_type="find_all", ) - # Verify that insert_data was called to create new user mock_prisma_client.insert_data.assert_called_once() insert_call_args = mock_prisma_client.insert_data.call_args insert_data = insert_call_args.kwargs["data"] assert insert_data["user_email"] == test_user_email assert insert_data["teams"] == [test_team_id] + # Confirm the clone path ran + mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( + where={"budget_id": test_default_budget_id} + ) + mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + @pytest.mark.asyncio async def test_attach_object_permission_to_dict_with_object_permission_id():