fix(budgets): compare budget_limits window reset_at as an absolute instant

This commit is contained in:
Devin AI 2026-07-28 15:45:12 +00:00
parent daf22ec871
commit fdf8567c7b
2 changed files with 52 additions and 2 deletions

View file

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

View file

@ -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()