From fdf8567c7b558ddbc27699745b7f88f85c7fb841 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 28 Jul 2026 15:45:12 +0000 Subject: [PATCH] fix(budgets): compare budget_limits window reset_at as an absolute instant --- .../proxy/common_utils/reset_budget_job.py | 6 ++- .../common_utils/test_reset_budget_job.py | 48 +++++++++++++++++++ 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 1210192416f..5ce84484afb 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -695,7 +695,9 @@ class ResetBudgetJob: reset_at_str = window.get("reset_at") if not reset_at_str: return False - reset_at = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00")).replace(tzinfo=None) + reset_at = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00")) + if reset_at.tzinfo is None: + reset_at = reset_at.replace(tzinfo=timezone.utc) if reset_at > now: return False spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.0) @@ -717,7 +719,7 @@ class ResetBudgetJob: from litellm.proxy.proxy_server import spend_counter_cache - now = datetime.utcnow() + now = datetime.now(timezone.utc) # Note on raw SQL: prisma-client-python does not support null-filtering # on `Json?` columns (no DbNull/JsonNull sentinel — see 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 f04d6f3cf5a..6e4555f42ee 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 @@ -1153,6 +1153,54 @@ def test_reset_budget_windows_resets_expired_key_window(monkeypatch): spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-expired:window:1d", value=0.0) +def test_reset_budget_windows_resets_tz_aware_expired_key_window(monkeypatch): + """Regression for #34896: a window whose `reset_at` carries a non-UTC offset + (e.g. Asia/Tokyo, +09:00) must be compared as an absolute instant, not by + dropping the offset. A reset_at that has already passed in UTC but reads + ~9h in the future once the offset is stripped must still reset now.""" + now_utc = datetime.now(timezone.utc) + tokyo = timezone(timedelta(hours=9)) + expired_tokyo = (now_utc - timedelta(minutes=5)).astimezone(tokyo).isoformat() + assert expired_tokyo.endswith("+09:00") + + key_rows = [ + { + "token": "sk-tokyo", + "budget_limits": [{"budget_duration": "1d", "reset_at": expired_tokyo}], + } + ] + job, prisma_client, spend_counter_cache = _make_reset_budget_windows_job( + monkeypatch, key_rows=key_rows, team_rows=[] + ) + + asyncio.run(job.reset_budget_windows()) + + prisma_client.db.litellm_verificationtoken.update.assert_awaited_once() + spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-tokyo:window:1d", value=0.0) + + +def test_reset_budget_windows_skips_tz_aware_unexpired_key_window(monkeypatch): + """The offset-aware comparison must not fire early either: a reset_at that is + genuinely in the future but expressed in a negative offset (US/Pacific, + -07:00) must be left untouched.""" + now_utc = datetime.now(timezone.utc) + pacific = timezone(timedelta(hours=-7)) + future_pacific = (now_utc + timedelta(hours=1)).astimezone(pacific).isoformat() + assert future_pacific.endswith("-07:00") + + key_rows = [ + { + "token": "sk-pacific", + "budget_limits": [{"budget_duration": "1d", "reset_at": future_pacific}], + } + ] + job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[]) + + asyncio.run(job.reset_budget_windows()) + + prisma_client.db.litellm_verificationtoken.update.assert_not_awaited() + + def test_reset_budget_windows_skips_unexpired_key_window(monkeypatch): """If `reset_at` is in the future, no write should happen for that key.""" now = datetime.utcnow()