From f55d13ebba6981899f0ba6cebc0d6c1d30aaebd6 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 26 Jun 2026 22:59:30 +0300 Subject: [PATCH] fix(team): persist budget_duration on /team/member_add member budgets (#31443) /team/member_add could not set budget_duration on an individual member budget. add_new_member created the budget row with only max_budget and allowed_models, and TeamMemberAddRequest had no budget_duration field, so a member added with an explicit per-member budget while the team ran a recurring member budget got a lifetime cap instead of a recurring allowance. Thread budget_duration from TeamMemberAddRequest through _process_team_members into add_new_member, and pull the member-budget resolution into a helper that writes budget_duration plus a computed budget_reset_at. When only a budget_duration is supplied and the team has a default member budget, the default is cloned and its reset window overridden so the member keeps the default's max_budget rather than becoming uncapped; a duration with no team default creates a window-only budget. Invalid durations are rejected with a 400 before any DB write, symmetric with /team/member_update. The available-team self-join bypass only grants the ability to join, so reject per-member budget and model controls (max_budget_in_team, budget_duration, allowed_models) for non-admin self-join callers in _validate_team_member_add_permissions, before any DB write. Otherwise a self-joining non-admin could set their own cap, reset window, or model scope past the team default; admins, team admins, and org admins are unaffected and a clean self-join still inherits the team default budget. Resolves LIT-4052 --- litellm/proxy/_types.py | 4 + .../management_endpoints/team_endpoints.py | 24 +++ litellm/proxy/management_helpers/utils.py | 95 ++++++--- .../test_team_endpoints.py | 90 ++++++++ .../test_management_helpers_utils.py | 200 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 6 files changed, 390 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ed030c31b19..559720eb64e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3658,6 +3658,10 @@ class TeamMemberAddRequest(MemberAddRequest): default=None, description="Maximum budget allocated to this user within the team. If not set, user has unlimited budget within team limits", ) + budget_duration: Optional[str] = Field( + default=None, + description="Duration after which this team member's budget resets (e.g. '1h', '24h', '7d', '30d'). If not set, the budget never resets.", + ) allowed_models: Optional[List[str]] = Field( default=None, description="List of models this team member can access. If not set, inherits the team's default_team_member_models or all team models.", diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 249cfdae4eb..f533774ce18 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2225,6 +2225,26 @@ async def _validate_team_member_add_permissions( }, ) + # Available-team self-join grants only the ability to join; per-member + # budget and model controls stay admin-only. Reject them here so a + # self-joining non-admin cannot set their own cap, reset window, or model + # scope via the bypass. + if ( + data.max_budget_in_team is not None + or data.budget_duration is not None + or data.allowed_models is not None + ): + raise HTTPException( + status_code=403, + detail={ + "error": ( + "Available-team self-join cannot set per-member budget or " + "model controls (max_budget_in_team, budget_duration, " + "allowed_models); these are admin-only." + ) + }, + ) + # Available-team self-join: caller may add only themselves, only as a # standard user. Enforce that here so the bypass cannot be used as a # privilege-escalation or cross-user-injection primitive. @@ -2290,6 +2310,7 @@ async def _process_team_members( team_id=data.team_id, default_team_budget_id=default_team_budget_id, allowed_models=member_allowed_models, + budget_duration=data.budget_duration, ) except Exception as e: raise HTTPException( @@ -2315,6 +2336,7 @@ async def _process_team_members( team_id=data.team_id, default_team_budget_id=default_team_budget_id, allowed_models=member_allowed_models, + budget_duration=data.budget_duration, ) except Exception as e: raise HTTPException( @@ -2578,6 +2600,8 @@ async def team_member_add( except HTTPException as e: raise e + _validate_budget_duration(data.budget_duration) + prisma_client = cast(PrismaClient, prisma_client) existing_team_row = await get_team_object( diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 4ce8633ed57..0c5439e26e2 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -167,6 +167,7 @@ async def _clone_team_default_budget_for_member( default_team_budget_id: str, user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, + budget_duration_override: Optional[str] = None, ) -> Optional[str]: """ Create a new budget row that copies the values from the team's default @@ -176,6 +177,10 @@ async def _clone_team_default_budget_for_member( 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). + + ``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 + member's reset cadence without discarding the team default's max_budget. """ default_budget = await BudgetRepository(prisma_client).table.find_unique( where={"budget_id": default_team_budget_id} @@ -198,6 +203,9 @@ async def _clone_team_default_budget_for_member( continue cloned_data[field] = value + if budget_duration_override is not None: + cloned_data["budget_duration"] = budget_duration_override + # Start the member's budget window at clone time, not the pool's reset # timestamp — otherwise a member joining mid-cycle inherits a stale reset. if cloned_data.get("budget_duration"): @@ -209,6 +217,55 @@ async def _clone_team_default_budget_for_member( return new_budget.budget_id +async def _resolve_member_budget_id( + prisma_client: PrismaClient, + user_api_key_dict: UserAPIKeyAuth, + litellm_proxy_admin_name: str, + max_budget_in_team: Optional[float], + allowed_models: Optional[list[str]], + budget_duration: Optional[str], + default_team_budget_id: Optional[str], +) -> Optional[str]: + """ + 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. + """ + has_explicit_limit = 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: + return 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, + budget_duration_override=budget_duration, + ) + + if not has_explicit_limit and budget_duration is None: + return None + + budget_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, + } + if max_budget_in_team is not None: + budget_data["max_budget"] = max_budget_in_team + if allowed_models is not None: + budget_data["allowed_models"] = allowed_models + if budget_duration is not None: + budget_data["budget_duration"] = budget_duration + budget_data["budget_reset_at"] = get_budget_reset_time( + budget_duration=budget_duration + ) + response = await BudgetRepository(prisma_client).table.create(data=budget_data) + return response.budget_id + + async def add_new_member( new_member: Member, max_budget_in_team: Optional[float], @@ -218,6 +275,7 @@ async def add_new_member( litellm_proxy_admin_name: str, default_team_budget_id: Optional[str] = None, allowed_models: Optional[List[str]] = None, + budget_duration: Optional[str] = None, ) -> Tuple[LiteLLM_UserTable, Optional[LiteLLM_TeamMembership]]: """ Add a new member to a team @@ -278,34 +336,15 @@ async def add_new_member( }, ) - # Check if trying to set a budget or model scope for team member - if max_budget_in_team is not None or allowed_models is not None: - # create a new budget item for this member - budget_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, - } - if max_budget_in_team is not None: - budget_data["max_budget"] = max_budget_in_team - if allowed_models is not None: - budget_data["allowed_models"] = allowed_models - response = await BudgetRepository(prisma_client).table.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: - # No per-member budget and no team default → member gets no budget. - _budget_id = None + _budget_id = await _resolve_member_budget_id( + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + max_budget_in_team=max_budget_in_team, + allowed_models=allowed_models, + budget_duration=budget_duration, + default_team_budget_id=default_team_budget_id, + ) if _budget_id and returned_user is not None and returned_user.user_id is not None: _returned_team_membership = await TeamMembershipRepository( diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index a649bc7225e..ecce92e7324 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1270,6 +1270,95 @@ async def test_available_team_self_join_blocks_admin_role_in_member_list(): assert exc_info.value.status_code == 403 +@pytest.mark.asyncio +@pytest.mark.parametrize( + "budget_control", + [ + {"max_budget_in_team": 1000.0}, + {"budget_duration": "1h"}, + {"allowed_models": ["gpt-4o"]}, + ], +) +async def test_available_team_self_join_blocks_member_budget_controls(budget_control): + """A self-joining non-admin must not be able to set their own per-member + budget or model controls via the available-team bypass; only proxy/team/org + admins may. Without this guard a self-joiner could shorten their budget + reset window or widen their cap/model scope past the team default.""" + from litellm.proxy._types import Member, TeamMemberAddRequest + from litellm.proxy.management_endpoints.team_endpoints import ( + _validate_team_member_add_permissions, + ) + + user = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER) + team = MagicMock(spec=LiteLLM_TeamTable) + team.team_id = "public-team" + team.members_with_roles = [] + team.organization_id = None + + data = TeamMemberAddRequest( + team_id="public-team", + member=Member(role="user", user_id="alice"), + **budget_control, + ) + + with ( + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", + return_value=False, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_available_team", + return_value=True, + ), + pytest.raises(HTTPException) as exc_info, + ): + await _validate_team_member_add_permissions( + user_api_key_dict=user, + complete_team_data=team, + data=data, + ) + + assert exc_info.value.status_code == 403 + assert "admin-only" in str(exc_info.value.detail).lower() + + +@pytest.mark.asyncio +async def test_available_team_self_join_allows_no_budget_controls(): + """The clean self-join (no per-member budget/model controls) must still be + permitted, so the new guard does not break the legitimate join path.""" + from litellm.proxy._types import Member, TeamMemberAddRequest + from litellm.proxy.management_endpoints.team_endpoints import ( + _validate_team_member_add_permissions, + ) + + user = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER) + team = MagicMock(spec=LiteLLM_TeamTable) + team.team_id = "public-team" + team.members_with_roles = [] + team.organization_id = None + + data = TeamMemberAddRequest( + team_id="public-team", + member=Member(role="user", user_id="alice"), + ) + + with ( + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", + return_value=False, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._is_available_team", + return_value=True, + ), + ): + await _validate_team_member_add_permissions( + user_api_key_dict=user, + complete_team_data=team, + data=data, + ) + + @pytest.mark.asyncio async def test_update_team_member_permissions_blocks_non_admin_via_available_team( mock_db_client, @@ -1385,6 +1474,7 @@ async def test_process_team_members_single_member(): team_id="test-team-123", default_team_budget_id="budget-123", allowed_models=None, + budget_duration=None, ) 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 463aba6f744..dbbdca65cc6 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,6 +1,7 @@ import json import os import sys +from datetime import datetime, timezone from litellm._uuid import uuid from unittest.mock import AsyncMock, MagicMock @@ -283,6 +284,81 @@ async def test_add_new_member_clones_default_team_budget_id(): assert create_data["budget_id"] == test_cloned_budget_id +@pytest.mark.asyncio +async def test_add_new_member_budget_duration_only_clones_default_max_budget(): + """When only a budget_duration is given and the team has a default member + budget, the member must clone the default (keeping its max_budget) and just + override the reset window. Creating a fresh duration-only row instead would + silently drop the team default's cap, leaving the member uncapped.""" + from litellm.proxy._types import LitellmUserRoles + + new_member = Member(user_id="dur-clone-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": "dur-clone-user", + "user_email": None, + "teams": ["team-dc"], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + return_value=mock_user_response + ) + mock_default_budget_row = MagicMock() + mock_default_budget_row.model_dump.return_value = { + "budget_id": "default-dc", + "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_cloned_budget_row = MagicMock() + mock_cloned_budget_row.budget_id = "cloned-dc" + 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": "team-dc", + "user_id": "dur-clone-user", + "budget_id": "cloned-dc", + "litellm_budget_table": None, + } + mock_prisma_client.db.litellm_teammembership.create = AsyncMock( + return_value=mock_team_membership_response + ) + + await add_new_member( + new_member=new_member, + max_budget_in_team=None, + prisma_client=mock_prisma_client, + team_id="team-dc", + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="test_admin", + default_team_budget_id="default-dc", + budget_duration="7d", + ) + + mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + cloned_create_data = ( + mock_prisma_client.db.litellm_budgettable.create.call_args.kwargs["data"] + ) + assert cloned_create_data["max_budget"] == 100.0 # kept from the team default + assert cloned_create_data["budget_duration"] == "7d" # overridden by the caller + assert cloned_create_data["budget_reset_at"] > datetime.now(timezone.utc) + + @pytest.mark.asyncio async def test_add_new_member_no_budget_when_no_default_and_no_max_budget(): """ @@ -434,6 +510,130 @@ async def test_add_new_member_creates_new_budget_when_max_budget_provided(): assert create_data["budget_id"] == test_new_budget_id +@pytest.mark.asyncio +async def test_add_new_member_persists_budget_duration(): + """Regression for the member_add half of the recurring-member-budget gap: + a budget_duration passed to add_new_member must be written to the new + member budget along with a future budget_reset_at, so the per-member budget + recurs instead of acting as a lifetime cap.""" + from litellm.proxy._types import LitellmUserRoles + + new_member = Member(user_id="user-dur", 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": "user-dur", + "user_email": None, + "teams": ["team-dur"], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + return_value=mock_user_response + ) + mock_budget_response = MagicMock() + mock_budget_response.budget_id = "budget-dur" + mock_prisma_client.db.litellm_budgettable.create = AsyncMock( + return_value=mock_budget_response + ) + mock_team_membership_response = MagicMock() + mock_team_membership_response.model_dump.return_value = { + "team_id": "team-dur", + "user_id": "user-dur", + "budget_id": "budget-dur", + "litellm_budget_table": None, + } + mock_prisma_client.db.litellm_teammembership.create = AsyncMock( + return_value=mock_team_membership_response + ) + + await add_new_member( + new_member=new_member, + max_budget_in_team=10.0, + prisma_client=mock_prisma_client, + team_id="team-dur", + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="test_admin", + default_team_budget_id=None, + allowed_models=["gpt-4o-mini"], + budget_duration="30d", + ) + + mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + budget_data = mock_prisma_client.db.litellm_budgettable.create.call_args.kwargs[ + "data" + ] + assert budget_data["max_budget"] == 10.0 + assert budget_data["allowed_models"] == ["gpt-4o-mini"] + assert budget_data["budget_duration"] == "30d" + reset_at = budget_data["budget_reset_at"] + assert isinstance(reset_at, datetime) + assert reset_at.tzinfo is not None + assert reset_at > datetime.now(timezone.utc) + + +@pytest.mark.asyncio +async def test_add_new_member_persists_budget_duration_without_max_budget(): + """budget_duration alone must still create a member budget; otherwise an + explicit recurring window passed without a cap would be silently dropped.""" + from litellm.proxy._types import LitellmUserRoles + + new_member = Member(user_id="user-dur2", 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": "user-dur2", + "user_email": None, + "teams": ["team-dur2"], + "user_role": "internal_user", + } + mock_prisma_client.db.litellm_usertable.upsert = AsyncMock( + return_value=mock_user_response + ) + mock_budget_response = MagicMock() + mock_budget_response.budget_id = "budget-dur2" + mock_prisma_client.db.litellm_budgettable.create = AsyncMock( + return_value=mock_budget_response + ) + mock_team_membership_response = MagicMock() + mock_team_membership_response.model_dump.return_value = { + "team_id": "team-dur2", + "user_id": "user-dur2", + "budget_id": "budget-dur2", + "litellm_budget_table": None, + } + mock_prisma_client.db.litellm_teammembership.create = AsyncMock( + return_value=mock_team_membership_response + ) + + _, result_team_membership = await add_new_member( + new_member=new_member, + max_budget_in_team=None, + prisma_client=mock_prisma_client, + team_id="team-dur2", + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name="test_admin", + default_team_budget_id=None, + budget_duration="7d", + ) + + mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + budget_data = mock_prisma_client.db.litellm_budgettable.create.call_args.kwargs[ + "data" + ] + assert budget_data["budget_duration"] == "7d" + assert budget_data["budget_reset_at"] > datetime.now(timezone.utc) + assert result_team_membership is not None + assert result_team_membership.budget_id == "budget-dur2" + + @pytest.mark.asyncio async def test_add_new_member_with_user_email_clones_default_budget(): """ diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index cdc91371dfe..43be9ad6e39 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30842,6 +30842,11 @@ export interface components { * @description List of models this team member can access. If not set, inherits the team's default_team_member_models or all team models. */ allowed_models?: string[] | null; + /** + * Budget Duration + * @description Duration after which this team member's budget resets (e.g. '1h', '24h', '7d', '30d'). If not set, the budget never resets. + */ + budget_duration?: string | null; /** * Max Budget In Team * @description Maximum budget allocated to this user within the team. If not set, user has unlimited budget within team limits