Merge e8af22e4d0 into 982d3a5476

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)

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