From 81518b7ca3ee25ddc1e221e8fb4ee0f62684e2ca Mon Sep 17 00:00:00 2001 From: daanhendrio <255322319+daanhendrio@users.noreply.github.com> Date: Fri, 10 Apr 2026 22:09:54 +0000 Subject: [PATCH] fix #25509 --- litellm/proxy/_types.py | 7 + .../management_endpoints/common_utils.py | 18 +- .../management_endpoints/team_endpoints.py | 21 ++ .../test_upsert_budget_membership.py | 79 +++++++ .../proxy/test_team_member_update.py | 210 +++++++++++++++++- 5 files changed, 332 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 96e221a9ac0..132235a87ce 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3775,6 +3775,12 @@ class TeamMemberUpdateRequest(TeamMemberDeleteRequest): rpm_limit: Optional[int] = Field( default=None, description="Requests per minute limit for this team member" ) + 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): @@ -3782,6 +3788,7 @@ class TeamMemberUpdateResponse(MemberUpdateResponse): max_budget_in_team: Optional[float] = None tpm_limit: Optional[int] = None rpm_limit: Optional[int] = 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 efc42d3355c..37e76b85a2c 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -1,6 +1,8 @@ +from datetime import datetime, timedelta from typing import TYPE_CHECKING, Any, Dict, Optional, Union from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.caching import DualCache from litellm.proxy._types import ( KeyRequestBase, @@ -354,6 +356,7 @@ async def _upsert_budget_and_membership( user_api_key_dict: UserAPIKeyAuth, tpm_limit: Optional[int] = None, rpm_limit: Optional[int] = None, + budget_duration: Optional[str] = None, ): """ Helper function to Create/Update or Delete the budget within the team membership @@ -366,11 +369,17 @@ async def _upsert_budget_and_membership( user_api_key_dict: User API Key dictionary containing user information tpm_limit: Tokens per minute limit for the team member rpm_limit: Requests per minute limit for the team member + budget_duration: Budget reset period for the team member (e.g. '30d', '1mo') - If max_budget, tpm_limit, and rpm_limit are all None, the user's budget is removed from the team membership. + If max_budget, tpm_limit, rpm_limit, 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 max_budget is None and tpm_limit is None and rpm_limit is None: + if ( + max_budget is None + and tpm_limit is None + and rpm_limit is None + and budget_duration is None + ): # disconnect the budget since all limits are None await tx.litellm_teammembership.update( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, @@ -389,6 +398,11 @@ async def _upsert_budget_and_membership( create_data["tpm_limit"] = tpm_limit if rpm_limit is not None: create_data["rpm_limit"] = rpm_limit + if budget_duration is not None: + create_data["budget_duration"] = budget_duration + create_data["budget_reset_at"] = datetime.utcnow() + timedelta( + seconds=duration_in_seconds(duration=budget_duration) + ) 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 e4e0b64af59..354a3ef8a0a 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2426,6 +2426,25 @@ async def team_member_update( identified_budget_id = tm.budget_id break + ### 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: + team_member_budget_id = (team_table.metadata or {}).get( + "team_member_budget_id" + ) + if team_member_budget_id: + _team_budget_row = await prisma_client.db.litellm_budgettable.find_unique( + where={"budget_id": team_member_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( @@ -2437,6 +2456,7 @@ async def team_member_update( user_api_key_dict=user_api_key_dict, tpm_limit=data.tpm_limit, rpm_limit=data.rpm_limit, + budget_duration=effective_budget_duration, ) ### update team member role @@ -2469,6 +2489,7 @@ async def team_member_update( max_budget_in_team=data.max_budget_in_team, tpm_limit=data.tpm_limit, rpm_limit=data.rpm_limit, + 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 a36dc7ff2e3..facbf3efb0f 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 @@ -275,6 +275,85 @@ async def test_upsert_rpm_limit_update_creates_new_budget(mock_tx, fake_user): ) +# 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_now = dt(2026, 1, 1, 0, 0, 0) + + with patch( + "litellm.proxy.management_endpoints.common_utils.datetime" + ) as mock_dt: + mock_dt.utcnow.return_value = fake_now + + 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 "budget_reset_at" in call_data + # 30 days from fake_now + from datetime import timedelta + assert call_data["budget_reset_at"] == fake_now + timedelta(days=30) + + # 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_now = dt(2026, 1, 1, 0, 0, 0) + + with patch( + "litellm.proxy.management_endpoints.common_utils.datetime" + ) as mock_dt: + mock_dt.utcnow.return_value = fake_now + + 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 "budget_reset_at" in call_data + + # 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): 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