diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index d065b062517..d1bbeeb4947 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -1215,7 +1215,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) + parsed_reset_at: Final = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00")) + reset_at: Final = ( + 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 new_value: Final = await ResetBudgetJob._window_carried_spend(window, counter_key, spend_counter_cache) @@ -1259,7 +1264,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 5d3afd95a55..070c590842f 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 @@ -944,6 +944,75 @@ 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: a `reset_at` serialized with a non-UTC offset (what + `litellm_settings.timezone: Asia/Shanghai` produces) must be compared in UTC. + Dropping the offset kept an already-expired window blocked for the length of + the offset, so the window never reset on time.""" + shanghai = timezone(timedelta(hours=8)) + expired = (datetime.now(shanghai) - timedelta(minutes=5)).isoformat() + + key_rows = [ + { + "token": "sk-shanghai-expired", + "budget_limits": [{"budget_duration": "24h", "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-shanghai-expired:window:24h", value=0.0 + ) + + +def test_reset_budget_windows_skips_unexpired_window_with_non_utc_offset(monkeypatch): + """The mirror of the guard above: a `+08:00` timestamp still in the future must + not reset early just because its wall-clock value already trails naive UTC now.""" + shanghai = timezone(timedelta(hours=8)) + future = (datetime.now(shanghai) + timedelta(minutes=5)).isoformat() + + key_rows = [ + { + "token": "sk-shanghai-future", + "budget_limits": [{"budget_duration": "24h", "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_keeps_window_span_at_one_duration(monkeypatch): + """The offset-dropping bug also widened the metered span: the window start is + derived as `reset_at - budget_duration`, so a reset that fires `offset` late + metered `24h + offset` of spend against a 24h budget, shrinking the effective + daily allowance. Assert the reset fires close enough to the stored `reset_at` + that the span stays one duration.""" + shanghai = timezone(timedelta(hours=8)) + reset_at = datetime.now(shanghai) - timedelta(minutes=1) + + key_rows = [ + { + "token": "sk-shanghai-span", + "budget_limits": [{"budget_duration": "24h", "reset_at": reset_at.isoformat()}], + } + ] + 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_awaited_once() + overshoot = datetime.now(timezone.utc) - reset_at + assert overshoot < timedelta(hours=1), f"reset fired {overshoot} after reset_at, widening the metered window" + + 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.