From e992271bd42eab8faa126fb7d7b3600ae93a5dcd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 01:21:41 -0800 Subject: [PATCH] fix(proxy): keep a resent budget reset time only once it is due A save that resends the same budget_duration keeps the stored reset time only when that time has already passed, so the scheduled sweep still resets the period it is due for. A row that is not due yet realigns its reset time on save, the same as before. --- litellm/proxy/common_utils/timezone_utils.py | 10 +++++++--- .../internal_user_endpoints.py | 14 +++++++++----- .../key_management_endpoints.py | 4 ++-- .../proxy/management_endpoints/team_endpoints.py | 4 ++-- .../test_internal_user_endpoints.py | 14 ++++++++------ .../test_key_management_endpoints.py | 4 +++- .../management_endpoints/test_team_endpoints.py | 4 +++- 7 files changed, 34 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/common_utils/timezone_utils.py b/litellm/proxy/common_utils/timezone_utils.py index 6ef4a427346..7a40931f50d 100644 --- a/litellm/proxy/common_utils/timezone_utils.py +++ b/litellm/proxy/common_utils/timezone_utils.py @@ -1,4 +1,5 @@ from datetime import datetime, time, timezone +from typing import Final from pydantic import ConfigDict @@ -81,10 +82,13 @@ def get_budget_reset_time(budget_duration: str) -> datetime: return compute_budget_reset_at(budget_duration, get_budget_reset_settings()) -def budget_period_unchanged( - requested_duration: object, stored_duration: str | None, stored_reset_at: datetime | None +def keeps_due_budget_reset( + requested_duration: str, stored_duration: str | None, stored_reset_at: datetime | None ) -> bool: - return isinstance(requested_duration, str) and requested_duration == stored_duration and stored_reset_at is not None + if stored_reset_at is None or requested_duration != stored_duration: + return False + stored_utc: Final = stored_reset_at if stored_reset_at.tzinfo else stored_reset_at.replace(tzinfo=timezone.utc) + return stored_utc <= datetime.now(timezone.utc) def _is_persistable_budget_duration(budget_duration: str) -> bool: diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 92604477379..246d2e3531d 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -48,7 +48,7 @@ from litellm.proxy.auth.password_policy import ( from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.common_utils.openai_error_payload import openai_shaped_proxy_exception, proxy_exception_for -from litellm.proxy.common_utils.timezone_utils import budget_period_unchanged +from litellm.proxy.common_utils.timezone_utils import keeps_due_budget_reset from litellm.proxy.common_utils.user_api_key_cache import ( AUTH_OBJECTS_TARGET, object_permission_cache_key, @@ -1565,10 +1565,14 @@ async def _update_single_user_helper( if existing_user_row is not None: existing_user_row = LiteLLM_UserTable.model_validate(existing_user_row.model_dump(exclude_none=True)) - if "budget_reset_at" in non_default_values and budget_period_unchanged( - non_default_values["budget_duration"], - existing_user_row.budget_duration, - existing_user_row.budget_reset_at, + if ( + "budget_reset_at" in non_default_values + and isinstance(non_default_values["budget_duration"], str) + and keeps_due_budget_reset( + non_default_values["budget_duration"], + existing_user_row.budget_duration, + existing_user_row.budget_reset_at, + ) ): del non_default_values["budget_reset_at"] diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 40bfbac0ea6..99f427e8f7e 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -78,7 +78,7 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( ) from litellm.proxy.common_utils.openai_error_payload import openai_error_param, proxy_exception_for from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys -from litellm.proxy.common_utils.timezone_utils import budget_period_unchanged, get_budget_reset_time +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time, keeps_due_budget_reset from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache, is_user_key_cache_key from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage @@ -3076,7 +3076,7 @@ async def prepare_key_update_data( non_default_values["budget_reset_at"] = None elif isinstance(budget_duration, str) and len(budget_duration) > 0: non_default_values["budget_duration"] = budget_duration - if not budget_period_unchanged( + if not keeps_due_budget_reset( budget_duration, existing_key_row.budget_duration, existing_key_row.budget_reset_at ): non_default_values["budget_reset_at"] = get_budget_reset_time(budget_duration=budget_duration) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 64f6c1b799e..c35acd23f58 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -134,7 +134,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for -from litellm.proxy.common_utils.timezone_utils import budget_period_unchanged, get_budget_reset_time +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time, keeps_due_budget_reset from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, team_membership_auth_cache_key, @@ -2986,7 +2986,7 @@ async def patch_team( def _set_budget_reset_at(data: UpdateTeamRequest, updated_kv: dict, existing_team: LiteLLM_TeamTable) -> None: """Set budget_reset_at in updated_kv if budget_duration is provided.""" - if data.budget_duration is not None and not budget_period_unchanged( + if data.budget_duration is not None and not keeps_due_budget_reset( data.budget_duration, existing_team.budget_duration, existing_team.budget_reset_at ): updated_kv["budget_reset_at"] = get_budget_reset_time(budget_duration=data.budget_duration) diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index d1050f2fdca..d9cb056503a 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2415,18 +2415,20 @@ async def test_user_status_update_refreshes_cached_user( @pytest.mark.asyncio @pytest.mark.parametrize( - ("request_fields", "role_default_duration", "moves_reset"), + ("request_fields", "role_default_duration", "stored_reset_at", "moves_reset"), [ - ({"budget_duration": "30d"}, None, False), - ({"budget_duration": "7d"}, None, True), - ({"user_role": LitellmUserRoles.INTERNAL_USER}, "30d", False), + ({"budget_duration": "30d"}, None, datetime(2020, 1, 1, tzinfo=timezone.utc), False), + ({"budget_duration": "30d"}, None, datetime(2099, 1, 1, tzinfo=timezone.utc), True), + ({"budget_duration": "7d"}, None, datetime(2020, 1, 1, tzinfo=timezone.utc), True), + ({"user_role": LitellmUserRoles.INTERNAL_USER}, "30d", datetime(2020, 1, 1, tzinfo=timezone.utc), False), ], ) -async def test_user_update_moves_the_budget_period_only_when_the_duration_changes( +async def test_user_update_keeps_a_due_budget_reset_when_the_duration_is_resent( mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, request_fields: dict[str, object], role_default_duration: str | None, + stored_reset_at: datetime, moves_reset: bool, ) -> None: import litellm @@ -2437,7 +2439,7 @@ async def test_user_update_moves_the_budget_period_only_when_the_duration_change saved_user: Final = LiteLLM_UserTable( user_id="user-spruce", budget_duration="30d", - budget_reset_at=datetime(2020, 1, 1, tzinfo=timezone.utc), + budget_reset_at=stored_reset_at, ) prisma_client: Final = mocker.MagicMock() prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user) diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 25cf59f5568..75d30417648 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -2110,11 +2110,13 @@ async def test_update_key_end_user_budget_id_folds_into_metadata_and_survives_om ("stored_duration", "stored_reset_at", "moves_reset"), [ ("30d", datetime(2020, 1, 1, tzinfo=timezone.utc), False), + ("30d", datetime(2020, 1, 1), False), + ("30d", datetime(2099, 1, 1, tzinfo=timezone.utc), True), ("7d", datetime(2020, 1, 1, tzinfo=timezone.utc), True), ("30d", None, True), ], ) -async def test_update_key_moves_the_budget_period_only_when_the_duration_changes( +async def test_update_key_keeps_a_due_budget_reset_when_the_duration_is_resent( stored_duration: str, stored_reset_at: datetime | None, moves_reset: bool ) -> None: existing_key: Final = LiteLLM_VerificationToken( diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 6983c3cf19c..c7b13238971 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -13218,11 +13218,13 @@ def test_set_budget_reset_at_sets_value_when_budget_duration_provided(): ("stored_duration", "stored_reset_at", "moves_reset"), [ ("30d", datetime(2020, 1, 1, tzinfo=timezone.utc), False), + ("30d", datetime(2020, 1, 1), False), + ("30d", datetime(2099, 1, 1, tzinfo=timezone.utc), True), ("7d", datetime(2020, 1, 1, tzinfo=timezone.utc), True), ("30d", None, True), ], ) -def test_set_budget_reset_at_moves_the_period_only_when_the_duration_changes( +def test_set_budget_reset_at_keeps_a_due_reset_when_the_duration_is_resent( stored_duration: str, stored_reset_at: datetime | None, moves_reset: bool ) -> None: from litellm.proxy._types import UpdateTeamRequest