fix #25509

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

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,

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

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

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