diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 92c920ca594..a286b61b2f3 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3860,6 +3860,12 @@ class TeamMemberUpdateRequest(TeamMemberDeleteRequest): default=None, description="List of models this team member can access. Pass an empty list to remove per-member model restrictions.", ) + budget_duration: Optional[str] = Field( + default=None, + description="Budget reset period for this team member (e.g. '30d', '1mo'). " + "Pass null to explicitly set a lifetime cap. " + "If omitted, inherits the team's team_member_budget_duration.", + ) class TeamMemberUpdateResponse(MemberUpdateResponse): @@ -3868,6 +3874,7 @@ class TeamMemberUpdateResponse(MemberUpdateResponse): tpm_limit: Optional[int] = None rpm_limit: Optional[int] = None allowed_models: Optional[List[str]] = None + budget_duration: Optional[str] = None class TeamModelAddRequest(BaseModel): diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index b0ea6b41ac5..fe7b85258d0 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy._types import ( KeyRequestBase, LiteLLM_ManagementEndpoint_MetadataFields, @@ -356,6 +357,7 @@ async def _upsert_budget_and_membership( rpm_limit: Optional[int] = None, allowed_models: Optional[List[str]] = None, team_default_budget_id: Optional[str] = None, + budget_duration: Optional[str] = None, ): """ Helper function to Create/Update or Delete the budget within the team membership @@ -374,8 +376,9 @@ async def _upsert_budget_and_membership( 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). + budget_duration: Budget reset period for the team member (e.g. '30d', '1mo') - If max_budget, tpm_limit, rpm_limit, and allowed_models are all None, the user's budget is removed from the team membership. + If max_budget, tpm_limit, rpm_limit, allowed_models, and budget_duration 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. """ if ( @@ -383,6 +386,7 @@ async def _upsert_budget_and_membership( and tpm_limit is None and rpm_limit is None and allowed_models is None + and budget_duration is None ): # disconnect the budget since all limits are None await tx.litellm_teammembership.update( @@ -411,6 +415,11 @@ async def _upsert_budget_and_membership( update_data["rpm_limit"] = rpm_limit if allowed_models is not None: update_data["allowed_models"] = allowed_models + if budget_duration is not None: + update_data["budget_duration"] = budget_duration + update_data["budget_reset_at"] = get_budget_reset_time( + budget_duration=budget_duration + ) await tx.litellm_budgettable.update( where={"budget_id": existing_budget_id}, data=update_data, @@ -459,6 +468,12 @@ async def _upsert_budget_and_membership( create_data["rpm_limit"] = rpm_limit if allowed_models is not None: create_data["allowed_models"] = allowed_models + if budget_duration is not None: + create_data["budget_duration"] = budget_duration + + bd = create_data.get("budget_duration") + if bd is not None: + create_data["budget_reset_at"] = get_budget_reset_time(budget_duration=bd) new_budget = await tx.litellm_budgettable.create( data=create_data, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index f254fea3e7f..01e71e5c4ee 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2714,15 +2714,28 @@ 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 + ### resolve effective budget_duration + # - Explicit value (including null) takes precedence + # - If omitted, inherit the team's configured team_member_budget_duration + if "budget_duration" in data.model_fields_set: + effective_budget_duration = data.budget_duration + else: + if team_default_budget_id: + _team_budget_row = await prisma_client.db.litellm_budgettable.find_unique( + where={"budget_id": team_default_budget_id} + ) + effective_budget_duration = ( + _team_budget_row.budget_duration if _team_budget_row else None + ) + else: + effective_budget_duration = None + ### upsert new budget async with prisma_client.db.tx() as tx: await _upsert_budget_and_membership( @@ -2736,6 +2749,7 @@ async def team_member_update( rpm_limit=data.rpm_limit, allowed_models=data.allowed_models, team_default_budget_id=team_default_budget_id, + budget_duration=effective_budget_duration, ) ### update team member role @@ -2769,6 +2783,7 @@ async def team_member_update( tpm_limit=data.tpm_limit, rpm_limit=data.rpm_limit, allowed_models=data.allowed_models, + budget_duration=effective_budget_duration, ) 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 f4bf0d7b2be..bc5d6562f7e 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 @@ -225,6 +225,80 @@ async def test_upsert_rpm_limit_update_creates_new_budget(mock_tx, fake_user): mock_tx.litellm_teammembership.upsert.assert_not_called() +# TEST: budget_duration is threaded through to the new budget row +@pytest.mark.asyncio +async def test_upsert_with_budget_duration(mock_tx, fake_user): + """ + When budget_duration is passed, it (and a derived budget_reset_at) should + appear in the create_data sent to Prisma. + + Regression test for: https://github.com/BerriAI/litellm/issues/25509 + """ + from unittest.mock import patch + from datetime import datetime as dt + + fake_reset_at = dt(2026, 1, 31, 0, 0, 0) + + with patch( + "litellm.proxy.management_endpoints.common_utils.get_budget_reset_time", + return_value=fake_reset_at, + ): + await _upsert_budget_and_membership( + mock_tx, + team_id="team-dur", + user_id="user-dur", + max_budget=10.0, + existing_budget_id=None, + user_api_key_dict=fake_user, + budget_duration="30d", + ) + + call_data = mock_tx.litellm_budgettable.create.call_args.kwargs["data"] + assert call_data["budget_duration"] == "30d" + assert call_data["budget_reset_at"] == fake_reset_at + + # membership upsert should still happen + mock_tx.litellm_teammembership.upsert.assert_awaited_once() + + +# TEST: budget_duration alone (no max_budget/limits) creates a budget rather than disconnecting +@pytest.mark.asyncio +async def test_upsert_budget_duration_only_creates_budget(mock_tx, fake_user): + """ + When only budget_duration is provided (no max_budget, tpm_limit, rpm_limit), + a new budget should be created rather than disconnecting. + """ + from unittest.mock import patch + from datetime import datetime as dt + + fake_reset_at = dt(2026, 1, 8, 0, 0, 0) + + with patch( + "litellm.proxy.management_endpoints.common_utils.get_budget_reset_time", + return_value=fake_reset_at, + ): + await _upsert_budget_and_membership( + mock_tx, + team_id="team-dur-only", + user_id="user-dur-only", + max_budget=None, + existing_budget_id=None, + user_api_key_dict=fake_user, + budget_duration="7d", + ) + + # Should NOT disconnect + mock_tx.litellm_teammembership.update.assert_not_called() + + # Should create a budget with budget_duration set + call_data = mock_tx.litellm_budgettable.create.call_args.kwargs["data"] + assert call_data["budget_duration"] == "7d" + assert call_data["budget_reset_at"] == fake_reset_at + + # Should upsert membership + mock_tx.litellm_teammembership.upsert.assert_awaited_once() + + # TEST: create new budget with only rpm_limit (no max_budget) @pytest.mark.asyncio async def test_upsert_rpm_only_creates_new_budget(mock_tx, fake_user): @@ -279,6 +353,9 @@ async def test_upsert_clones_when_pointing_at_shared_default(mock_tx, fake_user) 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. """ + from unittest.mock import patch + from datetime import datetime as dt + shared_default_id = "team-default-budget-1" # Default budget row in the DB: $200 cap, daily reset, 500 tpm. @@ -296,16 +373,21 @@ async def test_upsert_clones_when_pointing_at_shared_default(mock_tx, fake_user) } 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, - ) + fake_reset_at = dt(2026, 2, 1, 12, 0, 0) + with patch( + "litellm.proxy.management_endpoints.common_utils.get_budget_reset_time", + return_value=fake_reset_at, + ): + # 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() @@ -319,6 +401,7 @@ async def test_upsert_clones_when_pointing_at_shared_default(mock_tx, fake_user) "max_budget": 50.0, # caller wins "tpm_limit": 500, # cloned from default "budget_duration": "1d", # cloned from default + "budget_reset_at": fake_reset_at, }, include={"team_membership": True}, ) diff --git a/tests/test_litellm/proxy/test_team_member_update.py b/tests/test_litellm/proxy/test_team_member_update.py index 6561ec9e7fd..3cd1aae5fca 100644 --- a/tests/test_litellm/proxy/test_team_member_update.py +++ b/tests/test_litellm/proxy/test_team_member_update.py @@ -1,12 +1,220 @@ +import types +from unittest.mock import AsyncMock, MagicMock, patch + import pytest from fastapi import HTTPException from starlette.requests import Request import litellm.proxy.proxy_server as proxy_server -from litellm.proxy._types import TeamMemberUpdateRequest +from litellm.proxy._types import ( + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LitellmUserRoles, + Member, + TeamMemberUpdateRequest, + UserAPIKeyAuth, +) from litellm.proxy.management_endpoints.team_endpoints import team_member_update +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_request(): + scope = {"type": "http", "method": "POST", "path": "/team/member_update"} + return Request(scope) + + +def _admin_auth(): + return UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin-user", + ) + + +def _make_team_table(metadata=None): + return LiteLLM_TeamTable( + team_id="team-1", + members_with_roles=[Member(user_id="user-A", role="user")], + metadata=metadata or {}, + ) + + +def _team_info_response(team_table, budget_id=None): + membership = LiteLLM_TeamMembership( + user_id="user-A", + team_id="team-1", + budget_id=budget_id, + litellm_budget_table=None, + ) + return {"team_info": team_table, "team_memberships": [membership]} + + +def _mock_prisma(team_table, team_budget_duration=None): + """Build a minimal mock prisma client for team_member_update tests.""" + mock_existing_team = MagicMock() + mock_existing_team.model_dump.return_value = team_table.model_dump() + + mock_budget_row = MagicMock() + mock_budget_row.budget_duration = team_budget_duration + + # tx context manager + mock_tx = MagicMock() + mock_tx.__aenter__ = AsyncMock(return_value=mock_tx) + mock_tx.__aexit__ = AsyncMock(return_value=False) + + prisma = MagicMock() + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) + prisma.db.litellm_budgettable.find_unique = AsyncMock(return_value=mock_budget_row) + prisma.db.tx = MagicMock(return_value=mock_tx) + return prisma + + +# --------------------------------------------------------------------------- +# Tests for budget_duration resolution logic +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_team_member_update_explicit_budget_duration(): + """ + When budget_duration is explicitly provided in the request, it should be + passed straight to _upsert_budget_and_membership, ignoring any team setting. + """ + team_table = _make_team_table(metadata={"team_member_budget_id": "team-bud-1"}) + prisma = _mock_prisma(team_table, team_budget_duration="90d") + + data = TeamMemberUpdateRequest( + team_id="team-1", + user_id="user-A", + max_budget_in_team=5.0, + budget_duration="30d", + ) + + with patch("litellm.proxy.proxy_server.prisma_client", prisma), \ + patch("litellm.proxy.proxy_server.premium_user", True), \ + patch( + "litellm.proxy.management_endpoints.team_endpoints.team_info", + new=AsyncMock(return_value=_team_info_response(team_table)), + ), \ + patch( + "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership", + new=AsyncMock(), + ) as mock_upsert: + await team_member_update(data, _make_request(), _admin_auth()) + + mock_upsert.assert_awaited_once() + assert mock_upsert.call_args.kwargs["budget_duration"] == "30d" + # Should NOT have fetched the team-level budget row + prisma.db.litellm_budgettable.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_member_update_explicit_null_budget_duration(): + """ + When budget_duration is explicitly set to null in the request, None should + be used (lifetime cap), NOT the team's configured duration. + """ + team_table = _make_team_table(metadata={"team_member_budget_id": "team-bud-1"}) + prisma = _mock_prisma(team_table, team_budget_duration="30d") + + # Simulate the client sending {"budget_duration": null} by including the + # field in the model_fields_set while keeping the value None. + data = TeamMemberUpdateRequest.model_validate( + {"team_id": "team-1", "user_id": "user-A", "max_budget_in_team": 5.0, "budget_duration": None} + ) + assert "budget_duration" in data.model_fields_set # sanity check + + with patch("litellm.proxy.proxy_server.prisma_client", prisma), \ + patch("litellm.proxy.proxy_server.premium_user", True), \ + patch( + "litellm.proxy.management_endpoints.team_endpoints.team_info", + new=AsyncMock(return_value=_team_info_response(team_table)), + ), \ + patch( + "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership", + new=AsyncMock(), + ) as mock_upsert: + await team_member_update(data, _make_request(), _admin_auth()) + + mock_upsert.assert_awaited_once() + assert mock_upsert.call_args.kwargs["budget_duration"] is None + # Should NOT have fetched the team-level budget row + prisma.db.litellm_budgettable.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_member_update_inherits_team_budget_duration(): + """ + When budget_duration is omitted from the request and the team has a + team_member_budget_id, the team budget's duration should be inherited. + """ + team_table = _make_team_table(metadata={"team_member_budget_id": "team-bud-1"}) + prisma = _mock_prisma(team_table, team_budget_duration="30d") + + # budget_duration NOT included in request at all + data = TeamMemberUpdateRequest( + team_id="team-1", + user_id="user-A", + max_budget_in_team=5.0, + ) + assert "budget_duration" not in data.model_fields_set # sanity check + + with patch("litellm.proxy.proxy_server.prisma_client", prisma), \ + patch("litellm.proxy.proxy_server.premium_user", True), \ + patch( + "litellm.proxy.management_endpoints.team_endpoints.team_info", + new=AsyncMock(return_value=_team_info_response(team_table)), + ), \ + patch( + "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership", + new=AsyncMock(), + ) as mock_upsert: + await team_member_update(data, _make_request(), _admin_auth()) + + mock_upsert.assert_awaited_once() + assert mock_upsert.call_args.kwargs["budget_duration"] == "30d" + prisma.db.litellm_budgettable.find_unique.assert_awaited_once_with( + where={"budget_id": "team-bud-1"} + ) + + +@pytest.mark.asyncio +async def test_team_member_update_no_team_budget_duration_defaults_to_none(): + """ + When budget_duration is omitted and the team has no team_member_budget_id, + budget_duration should default to None. + """ + team_table = _make_team_table(metadata={}) # no team_member_budget_id + prisma = _mock_prisma(team_table) + + data = TeamMemberUpdateRequest( + team_id="team-1", + user_id="user-A", + max_budget_in_team=5.0, + ) + + with patch("litellm.proxy.proxy_server.prisma_client", prisma), \ + patch("litellm.proxy.proxy_server.premium_user", True), \ + patch( + "litellm.proxy.management_endpoints.team_endpoints.team_info", + new=AsyncMock(return_value=_team_info_response(team_table)), + ), \ + patch( + "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership", + new=AsyncMock(), + ) as mock_upsert: + await team_member_update(data, _make_request(), _admin_auth()) + + mock_upsert.assert_awaited_once() + assert mock_upsert.call_args.kwargs["budget_duration"] is None + prisma.db.litellm_budgettable.find_unique.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# Role / premium-user guard tests +# --------------------------------------------------------------------------- + @pytest.mark.asyncio async def test_ateam_member_update_admin_requires_premium(monkeypatch): # Arrange: patch prisma_client and premium_user