From ca140b60d3b882a09036c44bdaf5de08ac95cc38 Mon Sep 17 00:00:00 2001 From: Akhilesh Arora Date: Thu, 21 May 2026 18:00:01 +0200 Subject: [PATCH] fix(proxy/auth): handle tz-aware temp_budget_expiry /key/update accepts ISO 8601 with timezone (e.g. "2026-01-20T00:00:00Z"). Pydantic parses it to a tz-aware datetime; prepare_metadata_fields stores v.isoformat() which preserves the +00:00 offset. On the next request _get_temp_budget_increase compared the parsed value against a naive datetime.now() and raised TypeError, bricking the key. Normalize a naive expiry to UTC and compare against datetime.now(timezone.utc). --- litellm/proxy/auth/user_api_key_auth.py | 4 ++- .../proxy/auth/test_user_api_key_auth.py | 32 +++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 6974860a22a..018f605d105 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2318,7 +2318,9 @@ def _get_temp_budget_increase(valid_token: UserAPIKeyAuth): and "temp_budget_expiry" in valid_token_metadata ): expiry = datetime.fromisoformat(valid_token_metadata["temp_budget_expiry"]) - if expiry > datetime.now(): + if expiry.tzinfo is None: + expiry = expiry.replace(tzinfo=timezone.utc) + if expiry > datetime.now(timezone.utc): return valid_token_metadata["temp_budget_increase"] return None diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index defd3bbcdcd..e24e6219831 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -3457,3 +3457,35 @@ async def test_user_api_key_auth_does_not_overwrite_end_user_id_set_by_builder() finally: for k, v in originals.items(): setattr(_proxy_server_mod, k, v) + + +def test_get_temp_budget_increase_tz_aware_expiry(): + """ + /key/update accepts ISO 8601 with timezone (e.g. "2026-01-20T00:00:00Z"); + Pydantic parses it to a tz-aware datetime and prepare_metadata_fields + stores .isoformat() -> "2026-01-20T00:00:00+00:00". The comparison must + not raise TypeError on the next request. Also covers the naive-input + path (legacy metadata written without timezone). + """ + from datetime import datetime, timedelta, timezone + + from litellm.proxy.auth.user_api_key_auth import _get_temp_budget_increase + + future_aware = (datetime.now(timezone.utc) + timedelta(days=1)).isoformat() + past_aware = (datetime.now(timezone.utc) - timedelta(days=1)).isoformat() + future_naive = (datetime.now() + timedelta(days=1)).isoformat() + + for expiry, expected in [ + (future_aware, 100), # tz-aware future + (past_aware, None), # tz-aware past + (future_naive, 100), # naive future, exercises the `tzinfo is None` branch + ]: + token = UserAPIKeyAuth( + max_budget=100, + spend=0, + metadata={ + "temp_budget_increase": 100, + "temp_budget_expiry": expiry, + }, + ) + assert _get_temp_budget_increase(token) == expected, expiry