From 5b5bbac769e548199393e54aacf346953c5c5528 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 01:16:21 +0000 Subject: [PATCH 1/2] fix(team): link new members to the shared team member budget so /team/update applies to them Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/management_helpers/utils.py | 23 +-- .../test_management_helpers_utils.py | 135 +++++++++--------- 2 files changed, 82 insertions(+), 76 deletions(-) diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index f3bd4b0f6dd..2e7458232cc 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -291,9 +291,9 @@ async def _clone_team_default_budget_for_member( 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). + Used when adding a new team member with a per-member ``budget_duration`` + but no other per-member limit, so the member keeps the team default's + values in their own private budget row while the reset window differs. ``budget_duration_override`` replaces the default's reset window for this member while keeping the default's other limits, so an admin can set a @@ -344,14 +344,21 @@ async def _resolve_member_budget_id( """ Resolve the budget a new team member should be linked to. - Explicit per-member limits create a fresh budget. Otherwise the team's - default member budget is cloned (with ``budget_duration`` overriding its - reset window while keeping its other limits). A lone ``budget_duration`` - with no team default creates a window-only budget. With nothing set the - member gets no budget. + Explicit per-member limits create a fresh budget. Otherwise the member is + linked to the team's shared default member budget, so later ``/team/update`` + changes reach them; ``/team/member_update`` clones that row on first write. + A lone ``budget_duration`` clones the default with the reset window + overridden, or creates a window-only budget when there is no team default. + With nothing set the member gets no budget. """ has_explicit_limit: Final = max_budget_in_team is not None or allowed_models is not None + if not has_explicit_limit and default_team_budget_id is not None and budget_duration is None: + default_budget: Final = await _budget_table(prisma_client, tx).find_unique( + where={"budget_id": default_team_budget_id} + ) + return default_team_budget_id if default_budget is not None else None + if not has_explicit_limit and default_team_budget_id is not None: return await _clone_team_default_budget_for_member( prisma_client=prisma_client, 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 a6b1fc32eda..00de5171aa7 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 @@ -164,13 +164,16 @@ async def test_management_otel_span_redacts_nested_submission_env_var_secrets( @pytest.mark.asyncio -async def test_add_new_member_clones_default_team_budget_id(): +async def test_add_new_member_links_default_team_budget_id(): """ - 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. + A member added without any per-member limit must be LINKED to the team's + shared default member budget, not given a private copy of it. - 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. + Linking is what makes a later ``/team/update team_member_budget=...`` + reach existing members: the auth check reads the budget row behind the + membership, so a private clone would freeze the member at the old cap. + Per-member isolation is handled by ``/team/member_update`` cloning the + shared row on first write. """ from litellm.proxy._types import LitellmUserRoles @@ -178,7 +181,6 @@ async def test_add_new_member_clones_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" new_member = Member(user_id=test_user_id, role="user") @@ -202,36 +204,19 @@ async def test_add_new_member_clones_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_default_budget_row.budget_id = test_default_budget_id 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_prisma_client.db.litellm_budgettable.create = AsyncMock() # 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_cloned_budget_id, + "budget_id": test_default_budget_id, "litellm_budget_table": None, } mock_prisma_client.db.litellm_teammembership.create = AsyncMock( @@ -251,33 +236,67 @@ async def test_add_new_member_clones_default_team_budget_id(): assert result_user is not None assert result_user.user_id == test_user_id - # Membership should be linked to the new cloned budget, not the shared default. + # Membership points at the shared default row itself. assert result_team_membership is not None - assert result_team_membership.budget_id == test_cloned_budget_id - assert result_team_membership.budget_id != test_default_budget_id + assert result_team_membership.budget_id == test_default_budget_id mock_prisma_client.db.litellm_usertable.upsert.assert_called_once() mock_prisma_client.db.litellm_teammembership.create.assert_called_once() - # The clone must have happened: find_unique on the default, create for the clone. + # The default is only checked for existence; no private budget row is created. 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 + mock_prisma_client.db.litellm_budgettable.create.assert_not_called() team_membership_call_args = ( mock_prisma_client.db.litellm_teammembership.create.call_args ) create_data = team_membership_call_args.kwargs["data"] - assert create_data["budget_id"] == test_cloned_budget_id + assert create_data["budget_id"] == test_default_budget_id + + +@pytest.mark.asyncio +async def test_add_new_member_no_budget_when_default_budget_row_is_missing(): + """If team metadata still names a default member budget whose row was + deleted, the member must get no budget rather than a dangling link that + the membership foreign key would reject.""" + from litellm.proxy._types import LitellmUserRoles + + new_member = Member(user_id="missing-default-user", 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": "missing-default-user", + "user_email": None, + "teams": ["team-md"], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user_response) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=mock_user_response + ) + mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_budgettable.create = AsyncMock() + mock_prisma_client.db.litellm_teammembership.create = AsyncMock() + + _, result_team_membership = await add_new_member( + new_member=new_member, + max_budget_in_team=None, + prisma_client=mock_prisma_client, + team_id="team-md", + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="test_admin", + default_team_budget_id="deleted-default", + ) + + assert result_team_membership is None + mock_prisma_client.db.litellm_budgettable.create.assert_not_called() + mock_prisma_client.db.litellm_teammembership.create.assert_not_called() @pytest.mark.asyncio @@ -636,18 +655,17 @@ async def test_add_new_member_persists_budget_duration_without_max_budget(): @pytest.mark.asyncio -async def test_add_new_member_with_user_email_clones_default_budget(): +async def test_add_new_member_with_user_email_links_default_budget(): """ 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. + budget. The membership must link the shared default row so team-level + budget updates keep applying to this member. """ from litellm.proxy._types import LitellmUserRoles 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" new_member = Member(user_email=test_user_email, role="user") @@ -669,35 +687,18 @@ async def test_add_new_member_with_user_email_clones_default_budget(): } mock_prisma_client.insert_data = AsyncMock(return_value=mock_user_response) - # 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_default_budget_row.budget_id = test_default_budget_id 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_prisma_client.db.litellm_budgettable.create = AsyncMock() 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_cloned_budget_id, + "budget_id": test_default_budget_id, "litellm_budget_table": None, } mock_prisma_client.db.litellm_teammembership.create = AsyncMock( @@ -717,9 +718,8 @@ async def test_add_new_member_with_user_email_clones_default_budget(): assert result_user is not None assert result_user.user_email == test_user_email - # 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_cloned_budget_id + assert result_team_membership.budget_id == test_default_budget_id mock_prisma_client.get_data.assert_called_once_with( key_val={"user_email": test_user_email}, @@ -733,11 +733,10 @@ async def test_add_new_member_with_user_email_clones_default_budget(): 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() + mock_prisma_client.db.litellm_budgettable.create.assert_not_called() @pytest.mark.asyncio From ab99be9dad023d7d5e25d3e9352f7954357d4963 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 01:43:21 +0000 Subject: [PATCH 2/2] test(team): cover team_member_budget propagation and per-member isolation end to end Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_management_helpers_utils.py | 172 +++++++++++++++--- 1 file changed, 148 insertions(+), 24 deletions(-) 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 00de5171aa7..7512a3dfae8 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 @@ -1,13 +1,17 @@ import json +from collections.abc import Mapping from datetime import datetime, timezone -from litellm._uuid import uuid -from unittest.mock import AsyncMock, MagicMock +from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch import pytest - +import litellm +from litellm._uuid import uuid from litellm.proxy._types import ( + LiteLLM_BudgetTable, LiteLLM_TeamMembership, + LiteLLM_TeamTable, LiteLLM_UserTable, Member, UserAPIKeyAuth, @@ -165,19 +169,8 @@ async def test_management_otel_span_redacts_nested_submission_env_var_secrets( @pytest.mark.asyncio async def test_add_new_member_links_default_team_budget_id(): - """ - A member added without any per-member limit must be LINKED to the team's - shared default member budget, not given a private copy of it. - - Linking is what makes a later ``/team/update team_member_budget=...`` - reach existing members: the auth check reads the budget row behind the - membership, so a private clone would freeze the member at the old cap. - Per-member isolation is handled by ``/team/member_update`` cloning the - shared row on first write. - """ from litellm.proxy._types import LitellmUserRoles - # Setup test data test_user_id = "test_user_123" test_team_id = "test_team_456" test_default_budget_id = "default_budget_789" @@ -236,14 +229,12 @@ async def test_add_new_member_links_default_team_budget_id(): assert result_user is not None assert result_user.user_id == test_user_id - # Membership points at the shared default row itself. assert result_team_membership is not None assert result_team_membership.budget_id == test_default_budget_id mock_prisma_client.db.litellm_usertable.upsert.assert_called_once() mock_prisma_client.db.litellm_teammembership.create.assert_called_once() - # The default is only checked for existence; no private budget row is created. mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( where={"budget_id": test_default_budget_id} ) @@ -258,9 +249,6 @@ async def test_add_new_member_links_default_team_budget_id(): @pytest.mark.asyncio async def test_add_new_member_no_budget_when_default_budget_row_is_missing(): - """If team metadata still names a default member budget whose row was - deleted, the member must get no budget rather than a dangling link that - the membership foreign key would reject.""" from litellm.proxy._types import LitellmUserRoles new_member = Member(user_id="missing-default-user", role="user") @@ -656,11 +644,6 @@ async def test_add_new_member_persists_budget_duration_without_max_budget(): @pytest.mark.asyncio async def test_add_new_member_with_user_email_links_default_budget(): - """ - Test add_new_member with user_email instead of user_id and a team default - budget. The membership must link the shared default row so team-level - budget updates keep applying to this member. - """ from litellm.proxy._types import LitellmUserRoles test_user_email = "test@example.com" @@ -739,6 +722,147 @@ async def test_add_new_member_with_user_email_links_default_budget(): mock_prisma_client.db.litellm_budgettable.create.assert_not_called() +class _FakeBudgetTable: + def __init__(self) -> None: + self.rows: dict[str, dict[str, object]] = {} + + def _record(self, budget_id: str) -> LiteLLM_BudgetTable: + row: Final = self.rows[budget_id] + return LiteLLM_BudgetTable(**{k: v for k, v in row.items() if k in LiteLLM_BudgetTable.model_fields}) + + async def create( + self, *, data: Mapping[str, object], include: Mapping[str, bool] | None = None + ) -> LiteLLM_BudgetTable: + budget_id: Final = str(data.get("budget_id") or uuid.uuid4()) + self.rows[budget_id] = {**data, "budget_id": budget_id} + return self._record(budget_id) + + async def find_unique(self, *, where: Mapping[str, str]) -> LiteLLM_BudgetTable | None: + return self._record(where["budget_id"]) if where["budget_id"] in self.rows else None + + async def update(self, *, where: Mapping[str, str], data: Mapping[str, object]) -> LiteLLM_BudgetTable: + self.rows[where["budget_id"]] = {**self.rows[where["budget_id"]], **data} + return self._record(where["budget_id"]) + + +class _FakeMembershipTable: + def __init__(self, budgets: _FakeBudgetTable) -> None: + self.budgets: Final = budgets + self.budget_ids: dict[tuple[str, str], str | None] = {} + + def membership(self, team_id: str, user_id: str) -> LiteLLM_TeamMembership: + budget_id: Final = self.budget_ids[(team_id, user_id)] + return LiteLLM_TeamMembership( + user_id=user_id, + team_id=team_id, + budget_id=budget_id, + litellm_budget_table=self.budgets._record(budget_id) if budget_id is not None else None, + ) + + async def create(self, *, data: Mapping[str, str], include: Mapping[str, bool]) -> LiteLLM_TeamMembership: + self.budget_ids[(data["team_id"], data["user_id"])] = data["budget_id"] + return self.membership(data["team_id"], data["user_id"]) + + async def upsert(self, *, where: Mapping[str, Mapping[str, str]], data: Mapping[str, Mapping[str, object]]) -> None: + key: Final = where["user_id_team_id"] + connect: Final = data["update"]["litellm_budget_table"] + assert isinstance(connect, dict) + self.budget_ids[(key["team_id"], key["user_id"])] = connect["connect"]["budget_id"] + + +class _FakeUserTable: + async def upsert(self, *, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]]) -> LiteLLM_UserTable: + return LiteLLM_UserTable(user_id=where["user_id"], teams=list(data["create"].get("teams", []))) + + async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: + return 1 + + +class _FakeDb: + def __init__(self) -> None: + self.litellm_budgettable: Final = _FakeBudgetTable() + self.litellm_teammembership: Final = _FakeMembershipTable(self.litellm_budgettable) + self.litellm_usertable: Final = _FakeUserTable() + + +@pytest.mark.asyncio +async def test_team_update_reaches_inherited_members_but_not_overridden_ones(): + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.auth.auth_checks import _check_team_member_budget + from litellm.proxy.management_endpoints.common_utils import _upsert_budget_and_membership + from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler + from litellm.proxy.utils import ProxyLogging + + db: Final = _FakeDb() + prisma_client: Final = MagicMock() + prisma_client.db = db + admin: Final = UserAPIKeyAuth(user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN) + team_id: Final = "team-shared-default" + default_budget: Final = await db.litellm_budgettable.create(data={"budget_id": "team-default", "max_budget": 100.0}) + team: Final = LiteLLM_TeamTable(team_id=team_id, metadata={"team_member_budget_id": default_budget.budget_id}) + + for user_id in ("inherits", "overridden"): + await add_new_member( + new_member=Member(user_id=user_id, role="user"), + max_budget_in_team=None, + prisma_client=prisma_client, + team_id=team_id, + user_api_key_dict=admin, + litellm_proxy_admin_name="admin", + default_team_budget_id=default_budget.budget_id, + ) + + await _upsert_budget_and_membership( + db, + team_id=team_id, + user_id="overridden", + existing_budget_id=default_budget.budget_id, + user_api_key_dict=admin, + budget_patch={"max_budget": 50.0}, + team_default_budget_id=default_budget.budget_id, + ) + assert db.litellm_teammembership.membership(team_id, "inherits").budget_id == default_budget.budget_id + assert db.litellm_teammembership.membership(team_id, "overridden").budget_id != default_budget.budget_id + assert db.litellm_budgettable.rows[default_budget.budget_id]["max_budget"] == 100.0 + + with patch( # test-quality-ok: update_budget reads this module global; no dependency injection seam exists + "litellm.proxy.proxy_server.prisma_client", prisma_client + ): + await TeamMemberBudgetHandler.upsert_team_member_budget_table( + team_table=team, + user_api_key_dict=admin, + updated_kv={}, + team_member_budget=1.0, + ) + + async def spend_from_membership(counter_key: str, fallback_spend: float, max_budget: float | None = None) -> float: + return fallback_spend + + async def check(user_id: str, spend: float) -> None: + membership: Final = db.litellm_teammembership.membership(team_id, user_id).model_copy(update={"spend": spend}) + with patch( # test-quality-ok: production auth reads this module global; no dependency injection seam exists + "litellm.proxy.proxy_server.get_current_spend", spend_from_membership + ): + await _check_team_member_budget( + team_object=team, + user_object=LiteLLM_UserTable(user_id=user_id), + valid_token=UserAPIKeyAuth(token="tok", user_id=user_id, team_id=team_id), + prisma_client=prisma_client, + user_api_key_cache=MagicMock(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + team_membership=membership, + team_membership_loaded=True, + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await check("inherits", spend=2.0) + assert exc_info.value.max_budget == 1.0 + await check("overridden", spend=2.0) + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await check("overridden", spend=60.0) + assert exc_info.value.max_budget == 50.0 + + @pytest.mark.asyncio async def test_attach_object_permission_to_dict_with_object_permission_id(): """