From 348739a435af7d39c0215731ad86b36decda1b87 Mon Sep 17 00:00:00 2001 From: Akhand Pratap Shukla <145320759+A-P-Shukla@users.noreply.github.com> Date: Tue, 3 Nov 2026 12:37:39 +0530 Subject: [PATCH 1/2] fix(proxy): preserve team budget window reset on cap update --- .../management_endpoints/team_endpoints.py | 42 ++++++++++++++--- .../test_team_endpoints.py | 46 ++++++++++++++++++- 2 files changed, 80 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index f505406fe0e..63ac53c279e 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/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, @@ -2750,7 +2751,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 @@ -2978,7 +2979,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 @@ -2991,14 +2996,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. diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 2f4dad0ccd7..0a1773f2475 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -13213,6 +13213,47 @@ def test_set_budget_reset_at_sets_value_when_budget_duration_provided(): assert updated_kv["budget_reset_at"] is not None +@pytest.mark.asyncio +async def test_update_team_cap_change_preserves_open_budget_window_reset( + disable_audit_logging_for_mocked_team, monkeypatch +): + from datetime import datetime, timezone, tzinfo + + from litellm.models.team import BudgetLimitEntry + 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": "7d", "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 + assert written_window["reset_at"] == stored_reset_at.isoformat() + + @pytest.mark.asyncio async def test_clear_team_member_budget_duration_calls_update_budget(): """ @@ -15896,7 +15937,7 @@ 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, existing_metadata, existing_budget_limits=None): """Mock just enough of update_team to reach (or pass) the estimate gate.""" from unittest.mock import AsyncMock, MagicMock, patch @@ -15918,6 +15959,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) From c339941f50db1460c8f7ed6271b956fd1aad56dd Mon Sep 17 00:00:00 2001 From: Akhand Pratap Shukla <145320759+A-P-Shukla@users.noreply.github.com> Date: Wed, 4 Nov 2026 12:03:55 +0530 Subject: [PATCH 2/2] test(proxy): cover team budget window duration changes Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../test_team_endpoints.py | 21 ++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 0a1773f2475..57539fac2c3 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -19,6 +19,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,12 +13215,15 @@ def test_set_budget_reset_at_sets_value_when_budget_duration_provided(): @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, monkeypatch + 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.models.team import BudgetLimitEntry from litellm.proxy.management_endpoints.team_endpoints import update_team class FrozenDateTime(datetime): @@ -13238,7 +13242,7 @@ async def test_update_team_cap_change_preserves_open_budget_window_reset( await update_team( data=UpdateTeamRequest( team_id="test_team_id", - budget_limits=[{"budget_duration": "7d", "max_budget": 10}], + budget_limits=[{"budget_duration": updated_duration, "max_budget": 10}], ), http_request=MagicMock(spec=Request), user_api_key_dict=UserAPIKeyAuth( @@ -13251,7 +13255,10 @@ async def test_update_team_cap_change_preserves_open_budget_window_reset( written_window: Final = json.loads(prisma.db.litellm_teamtable.update.call_args.kwargs["data"]["budget_limits"])[0] assert written_window["max_budget"] == 10 - assert written_window["reset_at"] == stored_reset_at.isoformat() + 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 @@ -15937,7 +15944,11 @@ def test_team_output_token_estimate_admin_gate_matrix(label, request_body, exist ) -def _wire_update_team(stack, existing_metadata, existing_budget_limits=None): +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