fix(proxy): convert budget window reset_at to UTC before expiry comparison

This commit is contained in:
Devin AI 2026-07-28 09:10:48 +00:00
parent daf22ec871
commit 24d88a25b9
2 changed files with 51 additions and 2 deletions

View file

@ -695,7 +695,12 @@ 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)
parsed_reset_at = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00"))
reset_at = (
parsed_reset_at.replace(tzinfo=timezone.utc)
if parsed_reset_at.tzinfo is None
else parsed_reset_at.astimezone(timezone.utc)
)
if reset_at > now:
return False
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.0)
@ -717,7 +722,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

@ -1196,6 +1196,50 @@ def test_reset_budget_windows_resets_expired_team_window(monkeypatch):
spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:team:team-expired:window:30d", value=0.0)
def test_reset_budget_windows_resets_expired_window_with_non_utc_offset(monkeypatch):
"""Regression for #34896: `reset_at` written with a non-UTC offset (e.g. +09:00
when `litellm_settings.timezone: Asia/Tokyo`) must be converted to UTC before
the expiry comparison. Previously the offset was dropped, so a window stayed
blocked for the length of the offset."""
jst = timezone(timedelta(hours=9))
expired = (datetime.now(jst) - timedelta(minutes=5)).isoformat()
key_rows = [
{
"token": "sk-jst-expired",
"budget_limits": [{"budget_duration": "30d", "reset_at": expired}],
}
]
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-jst-expired:window:30d", value=0.0)
def test_reset_budget_windows_skips_unexpired_window_with_non_utc_offset(monkeypatch):
"""The mirror of the #34896 guard: a `+09:00` timestamp that is still in the
future in real time must not be reset early just because its wall-clock value
is already behind naive UTC now."""
jst = timezone(timedelta(hours=9))
future = (datetime.now(jst) + timedelta(minutes=5)).isoformat()
key_rows = [
{
"token": "sk-jst-future",
"budget_limits": [{"budget_duration": "30d", "reset_at": future}],
}
]
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_handles_string_budget_limits(monkeypatch):
"""Defensive: if `query_raw` returns `budget_limits` as a JSON-encoded
string (driver-dependent), the code still parses and resets it.