fix #25509

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

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,

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

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

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