scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 154

fix: give each team member an independent budget instead of sharing the team default

Browse source 
Previously, members added to a team without an explicit per-member budget were
all linked to the same `litellm_budgettable` row referenced by the team's
`metadata.team_member_budget_id`. Updating one member's budget via
`/team/member_update` mutated the shared row and silently changed every other
member's budget too.

Now both write paths produce a private, per-member budget:

- `add_new_member` clones the team's default budget into a fresh row when a
  member is added without `max_budget_in_team`/`allowed_models`. If no team
  default exists, the membership is created with no budget.
- `_upsert_budget_and_membership` detects when an existing membership still
  points at the team's default budget id and clones-on-write, relinking the
  membership to the new private budget before applying the update.
- `team_member_update` reads `team_member_budget_id` from team metadata and
  passes it through so the helper can make this distinction.

Adds unit tests for clone-on-write, in-place update of a private budget, and
the no-default-no-budget add path.

Made-with: Cursor
This commit is contained in:
shivam 2026-04-21 17:58:50 -07:00
parent 26fcbc93e5
commit 27a105bcf9
No known key found for this signature in database
5 changed files with 369 additions and 52 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

45
litellm/proxy/management_endpoints/common_utils.py
View file

@ -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:

13
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -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

70
litellm/proxy/management_helpers/utils.py
View file

@ -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 = (

101
tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py
View file

@ -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()

192
tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py
View file

@ -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():