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

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

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"]
 

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)

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)

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)

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(

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