Litellm individual team member budgets

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:

litellm/proxy/management_endpoints/team_endpoints.py

@@ -2608,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(
@@ -2620,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

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

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