Merge a3f8bfb61d into 784754ca48

litellm/proxy/management_endpoints/team_endpoints.py

@@ -49,6 +49,7 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.proxy._types import (
     UI_TEAM_ID,
     BlockTeamRequest,
+    BudgetLimitEntry,
     BudgetNewRequest,
     CommonProxyErrors,
     DeleteTeamRequest,
@@ -2755,7 +2756,7 @@ async def update_team(
             )
 
         # Check budget_duration and budget_reset_at
-        _set_budget_reset_at(data, updated_kv)
+        _set_budget_reset_at(data, updated_kv, existing_team.budget_limits)
 
         _team_member_fields_in_request: Final = {
             field
@@ -2983,7 +2984,11 @@ async def patch_team(
         raise handle_exception_on_proxy(e)
 
 
-def _set_budget_reset_at(data: UpdateTeamRequest, updated_kv: dict) -> None:
+def _set_budget_reset_at(
+    data: UpdateTeamRequest,
+    updated_kv: dict,
+    existing_budget_limits: Sequence[BudgetLimitEntry] | None = None,
+) -> None:
     """Set budget_reset_at in updated_kv if budget_duration is provided."""
     if data.budget_duration is not None:
         from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
@@ -2996,14 +3001,37 @@ def _set_budget_reset_at(data: UpdateTeamRequest, updated_kv: dict) -> None:
     if data.budget_limits is not None and len(data.budget_limits) > 0:
         from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
 
-        initialized_windows: Final = []
-        for window in data.budget_limits:
-            w = window if isinstance(window, dict) else window.model_dump()
-            w["reset_at"] = get_budget_reset_time(budget_duration=w["budget_duration"]).isoformat()
-            initialized_windows.append(w)
+        initialized_windows: Final = [
+            _team_window_with_reset_at(window, existing_budget_limits or ()) for window in data.budget_limits
+        ]
         updated_kv["budget_limits"] = json.dumps(initialized_windows)
 
 
+def _team_window_with_reset_at(
+    window: BudgetLimitEntry, existing_budget_limits: Sequence[BudgetLimitEntry]
+) -> dict[str, object]:
+    from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
+
+    now: Final = datetime.now(timezone.utc)
+    open_reset_at: Final = next(
+        (
+            stored_window.reset_at
+            for stored_window in existing_budget_limits
+            if stored_window.budget_duration == window.budget_duration
+            and stored_window.reset_at is not None
+            and (
+                stored_window.reset_at
+                if stored_window.reset_at.tzinfo is not None
+                else stored_window.reset_at.replace(tzinfo=timezone.utc)
+            )
+            > now
+        ),
+        None,
+    )
+    reset_at: Final = open_reset_at or get_budget_reset_time(budget_duration=window.budget_duration)
+    return {**window.model_dump(), "reset_at": reset_at.isoformat()}
+
+
 async def handle_update_object_permission(data_json: dict, existing_team_row: _ObjectPermissionRow) -> dict:
     """
     Handle the update of object permission for a team.

tests/unit/proxy/management_endpoints/test_team_endpoints.py

@@ -20,6 +20,7 @@ from litellm._uuid import uuid
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY
 from litellm.proxy._types import (
+    BudgetLimitEntry,
     LiteLLM_BudgetTable,
     LiteLLM_BudgetTableFull,
     LiteLLM_ModelTable,
@@ -13214,6 +13215,53 @@ def test_set_budget_reset_at_sets_value_when_budget_duration_provided():
     assert updated_kv["budget_reset_at"] is not None
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("updated_duration", "preserves_reset_at"), (("7d", True), ("14d", False)))
+async def test_update_team_cap_change_preserves_open_budget_window_reset(
+    disable_audit_logging_for_mocked_team: None,
+    monkeypatch: pytest.MonkeyPatch,
+    updated_duration: str,
+    preserves_reset_at: bool,
+):
+    from datetime import datetime, timezone, tzinfo
+
+    from litellm.proxy.management_endpoints.team_endpoints import update_team
+
+    class FrozenDateTime(datetime):
+        @classmethod
+        def now(cls, tz: tzinfo | None = None) -> datetime:
+            return datetime(2026, 1, 1, tzinfo=timezone.utc)
+
+    stored_reset_at = datetime(2026, 1, 8, tzinfo=timezone.utc)
+    monkeypatch.setattr("litellm.proxy.management_endpoints.team_endpoints.datetime", FrozenDateTime)
+    with ExitStack() as stack:
+        prisma = _wire_update_team(
+            stack,
+            {},
+            [BudgetLimitEntry(budget_duration="7d", max_budget=5, reset_at=stored_reset_at)],
+        )
+        await update_team(
+            data=UpdateTeamRequest(
+                team_id="test_team_id",
+                budget_limits=[{"budget_duration": updated_duration, "max_budget": 10}],
+            ),
+            http_request=MagicMock(spec=Request),
+            user_api_key_dict=UserAPIKeyAuth(
+                user_role=LitellmUserRoles.PROXY_ADMIN,
+                api_key=MASTER_KEY,
+                user_id="proxy-admin",
+            ),
+        )
+
+    written_window: Final = json.loads(prisma.db.litellm_teamtable.update.call_args.kwargs["data"]["budget_limits"])[0]
+
+    assert written_window["max_budget"] == 10
+    if preserves_reset_at:
+        assert written_window["reset_at"] == stored_reset_at.isoformat()
+    else:
+        assert written_window["reset_at"] != stored_reset_at.isoformat()
+
+
 @pytest.mark.asyncio
 async def test_clear_team_member_budget_duration_calls_update_budget():
     """
@@ -15897,7 +15945,11 @@ def test_team_output_token_estimate_admin_gate_matrix(label, request_body, exist
     )
 
 
-def _wire_update_team(stack, existing_metadata):
+def _wire_update_team(
+    stack: ExitStack,
+    existing_metadata: object,
+    existing_budget_limits: Sequence[BudgetLimitEntry] | None = None,
+) -> MagicMock:
     """Mock just enough of update_team to reach (or pass) the estimate gate."""
     from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -15919,6 +15971,9 @@ def _wire_update_team(stack, existing_metadata):
         "team_id": "test_team_id",
         "team_alias": "test_team",
         "metadata": existing_metadata,
+        "budget_limits": (
+            [window.model_dump() for window in existing_budget_limits] if existing_budget_limits is not None else None
+        ),
         "members_with_roles": [{"user_id": "team-admin", "role": "admin"}],
     }
     mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team)