From e8af22e4d01d1438bd05488a781436508285f05a Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Thu, 27 Aug 2026 10:02:37 -0700 Subject: [PATCH] fix(proxy): preserve timezone offsets when resetting multi-window budgets _reset_expired_window dropped the parsed offset with .replace(tzinfo=None), so a "+08:00" boundary was read as wall-clock and compared against a UTC now, delaying every non-UTC reset by the offset. compute_budget_reset_at writes those offsets itself whenever the configured budget timezone is not UTC. Keep the offset via _as_utc, which only fills in a tzinfo for legacy naive rows, and take `now` as an aware UTC datetime so the comparison is offset-aware on both sides. Rebased onto the current staging tip, where the surrounding job has been restructured: the two edits are re-applied to the current file rather than replayed, and the other three utcnow() sites belong to the key/user/team chunk resets this PR deliberately does not touch. Signed-off-by: Vineeth Sai --- .../proxy/common_utils/reset_budget_job.py | 9 +- .../common_utils/test_reset_budget_job.py | 111 ++++++++++++++++++ 2 files changed, 118 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 4ebcc549cdd..91038f7c714 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -1104,7 +1104,12 @@ class ResetBudgetJob: reset_at_str: Final = window.get("reset_at") if not reset_at_str: return False - reset_at: Final = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00")).replace(tzinfo=None) + # Keep the parsed offset. Dropping it left the wall-clock reading of a + # "+08:00" boundary compared against a UTC now, which delays the reset by + # the offset; compute_budget_reset_at writes those offsets itself whenever + # the configured budget timezone is not UTC. _as_utc only fills in a + # tzinfo for legacy rows written naive, which always meant UTC. + reset_at: Final = _as_utc(datetime.fromisoformat(reset_at_str.replace("Z", "+00:00"))) if reset_at > now: return False spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.0) @@ -1126,7 +1131,7 @@ class ResetBudgetJob: from litellm.proxy.proxy_server import spend_counter_cache - now: Final = datetime.utcnow() + now: Final = datetime.now(timezone.utc) for source in _WINDOW_SOURCES: try: await self._reset_windows_for(source=source, now=now, spend_counter_cache=spend_counter_cache) diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index aa844304604..9672f6f420b 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -12,6 +12,7 @@ import prisma import pytest +from litellm.caching.dual_cache import DualCache from litellm.proxy._types import LiteLLM_VerificationToken from litellm.proxy.common_utils import reset_budget_job as reset_budget_job_module from litellm.constants import ( @@ -2588,3 +2589,113 @@ def test_ambiguous_commit_replay_does_not_erase_newly_accrued_spend( assert client.key_spend == expected_spend assert client.commit_attempts == expected_commits assert client.reconnect_reasons == expected_reconnects + + +class TestBudgetWindowResetAtOffsets: + """ + Regression tests for https://github.com/BerriAI/litellm/issues/37454 + + `compute_budget_reset_at` writes `reset_at` in the configured budget + timezone, so a proxy on Asia/Shanghai stores `...T00:00:00+08:00`. + `_reset_expired_window` dropped that offset with `.replace(tzinfo=None)` and + compared the wall-clock reading against a UTC `now`, so the window stayed + enforced for another 8 hours after its own boundary. + """ + + @staticmethod + def _window(reset_at: str, budget_duration: str = "24h") -> dict: + return {"budget_duration": budget_duration, "max_budget": 100, "reset_at": reset_at} + + @staticmethod + async def _expired(window: dict, now: datetime, tz: str = "UTC") -> bool: + return await ResetBudgetJob._reset_expired_window( + window, + "spend:key:tok-1:window:" + window["budget_duration"], + DualCache(), + now, + BudgetResetSettings(timezone=tz, reset_time_of_day=dt_time(0, 0)), + ) + + @pytest.mark.asyncio + async def test_offset_window_is_not_reset_a_minute_early(self): + window = self._window("2026-08-20T00:00:00+08:00") + + assert ( + await self._expired(window, datetime(2026, 8, 19, 15, 59, tzinfo=timezone.utc), tz="Asia/Shanghai") is False + ) + assert window["reset_at"] == "2026-08-20T00:00:00+08:00" + + @pytest.mark.asyncio + async def test_offset_window_is_reset_exactly_at_its_own_boundary(self): + """2026-08-19 16:00 UTC IS 2026-08-20 00:00+08:00, so the window is due. + Reading it as naive 2026-08-20 00:00 delayed this by the 8 hour offset.""" + window = self._window("2026-08-20T00:00:00+08:00") + + assert ( + await self._expired(window, datetime(2026, 8, 19, 16, 0, tzinfo=timezone.utc), tz="Asia/Shanghai") is True + ) + assert window["reset_at"] != "2026-08-20T00:00:00+08:00" + + @pytest.mark.asyncio + async def test_negative_offset_window_is_not_reset_early(self): + """The offset cuts the other way too: a UTC-5 boundary is due later, not + sooner, and dropping the offset would fire it 5 hours early.""" + window = self._window("2026-08-19T00:00:00-05:00") # 2026-08-19 05:00 UTC + + assert ( + await self._expired(window, datetime(2026, 8, 19, 4, 59, tzinfo=timezone.utc), tz="America/New_York") + is False + ) + assert ( + await self._expired(window, datetime(2026, 8, 19, 5, 0, tzinfo=timezone.utc), tz="America/New_York") is True + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("reset_at", ["2026-08-19T12:00:00Z", "2026-08-19T12:00:00+00:00"]) + async def test_utc_windows_are_unchanged(self, reset_at): + """The shapes that already worked must keep working: this is the control.""" + assert await self._expired(self._window(reset_at), datetime(2026, 8, 19, 11, 59, tzinfo=timezone.utc)) is False + assert await self._expired(self._window(reset_at), datetime(2026, 8, 19, 12, 0, tzinfo=timezone.utc)) is True + + @pytest.mark.asyncio + async def test_legacy_naive_reset_at_is_read_as_utc(self): + """Rows written before offsets were stored carry no tzinfo; they always + meant UTC, and `_as_utc` keeps reading them that way.""" + window = self._window("2026-08-19T12:00:00") + + assert await self._expired(window, datetime(2026, 8, 19, 11, 59, tzinfo=timezone.utc)) is False + assert ( + await self._expired(self._window("2026-08-19T12:00:00"), datetime(2026, 8, 19, 12, 0, tzinfo=timezone.utc)) + is True + ) + + @pytest.mark.asyncio + async def test_only_the_expired_window_is_reset(self): + """A key with 24h and 30d windows must lose only the 24h counter.""" + cache = DualCache() + settings = BudgetResetSettings(timezone="Asia/Shanghai", reset_time_of_day=dt_time(0, 0)) + now = datetime(2026, 8, 19, 16, 0, tzinfo=timezone.utc) + due = self._window("2026-08-20T00:00:00+08:00", budget_duration="24h") + not_due = self._window("2026-09-01T00:00:00+08:00", budget_duration="30d") + + results = [ + await ResetBudgetJob._reset_expired_window( + window, f"spend:key:tok-1:window:{window['budget_duration']}", cache, now, settings + ) + for window in (due, not_due) + ] + + assert results == [True, False] + assert cache.in_memory_cache.get_cache("spend:key:tok-1:window:24h") == 0.0 + assert cache.in_memory_cache.get_cache("spend:key:tok-1:window:30d") is None + assert not_due["reset_at"] == "2026-09-01T00:00:00+08:00" + + @pytest.mark.asyncio + async def test_the_written_reset_at_keeps_the_configured_offset(self): + """The offset this fix reads is one the job writes itself, so a reset + window round-trips through the same parser next tick.""" + window = self._window("2026-08-20T00:00:00+08:00") + + await self._expired(window, datetime(2026, 8, 19, 16, 0, tzinfo=timezone.utc), tz="Asia/Shanghai") + + assert datetime.fromisoformat(window["reset_at"]).tzinfo is not None