diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 54b9b26123e..79740841e9e 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, @@ -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. diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 697d9810284..8b3df17bed3 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/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)