fix(proxy/auth): handle tz-aware temp_budget_expiry

[devin-ai-integration]

Relevant issues

Linear ticket

Resolves LIT-4576

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have added meaningful tests
• My PR passes all CI/CD checks (e.g., lint, format, unit tests)
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review

Screenshots / Proof of Fix

/key/update documents temp_budget_expiry as ISO-8601 with timezone (e.g. "2026-01-20T00:00:00Z"). Setting it that way stored a tz-aware value and bricked the key on the next request. Reproduced end to end against a live proxy on localhost:4000 hitting a real Anthropic model (anthropic/claude-haiku-4-5)

Same script for both runs:

MASTER=sk-1234; BASE=http://localhost:4000

# 1) generate a key
KEY=$(curl -s $BASE/key/generate -H "Authorization: Bearer $MASTER" -H "Content-Type: application/json" \
  -d '{"models":["anthropic-haiku-4-5"],"max_budget":0.01}' | python3 -c "import sys,json;print(json.load(sys.stdin)['key'])")

# 2) update with the DOCUMENTED tz-aware format
curl -s $BASE/key/update -H "Authorization: Bearer $MASTER" -H "Content-Type: application/json" \
  -d "{\"key\":\"$KEY\",\"temp_budget_increase\":100,\"temp_budget_expiry\":\"2026-01-20T00:00:00Z\"}"
#   -> stored metadata: {"temp_budget_expiry": "2026-01-20T00:00:00+00:00", ...}  (tz-aware)

# 3) make a request with the key
curl -s -w "\nHTTP %{http_code}\n" $BASE/v1/chat/completions -H "Authorization: Bearer $KEY" -H "Content-Type: application/json" \
  -d '{"model":"anthropic-haiku-4-5","messages":[{"role":"user","content":"hi"}]}'

Before (commit fdf380d, pre-fix)

{"error":{"message":"Authentication Error, can't compare offset-naive and offset-aware datetimes","type":"auth_error","param":"None","code":"401"}}
HTTP 401

Proxy traceback:

File ".../litellm/proxy/auth/user_api_key_auth.py", line 2698, in _update_key_budget_with_temp_budget_increase
  temp_budget_increase = _get_temp_budget_increase(valid_token) or 0.0
File ".../litellm/proxy/auth/user_api_key_auth.py", line 2688, in _get_temp_budget_increase
  if expiry > datetime.now():
TypeError: can't compare offset-naive and offset-aware datetimes

After (commit 6261056, this PR)

{"id":"chatcmpl-...","model":"anthropic-haiku-4-5","object":"chat.completion","choices":[{"finish_reason":"stop","index":0,"message":{"content":"Hello! 👋 How can I help you today?","role":"assistant"}}],"usage":{"completion_tokens":16,"prompt_tokens":8,"total_tokens":24}}
HTTP 200

Type

🐛 Bug Fix

Changes

UpdateKeyRequest.temp_budget_expiry: Optional[datetime] accepts the Z suffix and Pydantic parses it to a tz-aware UTC datetime; prepare_metadata_fields stores v.isoformat(), which keeps the +00:00 offset. On the next request _get_temp_budget_increase ran datetime.fromisoformat(...) > datetime.now(), and the aware-vs-naive comparison raised

_get_temp_budget_increase in litellm/proxy/auth/user_api_key_auth.py now treats a naive parsed expiry as UTC and always compares against datetime.now(timezone.utc):

expiry = datetime.fromisoformat(valid_token_metadata["temp_budget_expiry"])
if expiry.tzinfo is None:
    expiry = expiry.replace(tzinfo=timezone.utc)
if expiry > datetime.now(timezone.utc):
    return valid_token_metadata["temp_budget_increase"]

The naive branch keeps the previous behavior for any legacy metadata written without a timezone

tests/proxy_unit_tests/test_proxy_utils.py adds test_get_temp_budget_increase_tz_aware_expiry covering both the future-expiry (returns the increase) and past-expiry (returns None) branches using the documented "...Z" input format. The test fails before this change with the same TypeError and passes after

Final Attestation

• The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

Link to Devin session: https://app.devin.ai/sessions/29199d8f443e46a79091b097e31176f7
Requested by: @shivamrawat1

[devin-ai-integration]

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

• Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
• Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

• Disable automatic comment, CI, and merge conflict monitoring

[greptile-apps]

Greptile Summary

This PR fixes a TypeError in _get_temp_budget_increase that occurred when temp_budget_expiry was stored as a tz-aware ISO string (e.g. "2026-01-20T00:00:00Z") — the documented format — and then compared against a naive datetime.now(). The fix normalizes any naive parsed expiry to UTC before comparing against datetime.now(timezone.utc).

• user_api_key_auth.py: Two lines added to _get_temp_budget_increase: if expiry.tzinfo is None, stamp it as UTC via replace(tzinfo=timezone.utc), then compare against the UTC-aware datetime.now(timezone.utc). timezone was already imported at the top of the file, so no new import is required.
• test_proxy_utils.py: New test test_get_temp_budget_increase_tz_aware_expiry exercises both the future-expiry (returns the budget increase) and past-expiry (returns None) branches using tz-aware strings, complementing the existing naive-datetime test.

Confidence Score: 5/5

Safe to merge — the change is small, targeted to the comparison in a single helper function, and the fix matches the documented format contract.

The two-line change is correct: it normalises a tz-naive parsed expiry to UTC and compares against datetime.now(timezone.utc), resolving the aware-vs-naive comparison error. timezone is already imported, no side effects touch the broader auth path, and the new test reproduces both the broken case and the expected outcome.

No files require special attention.

Important Files Changed

Filename	Overview
litellm/proxy/auth/user_api_key_auth.py	Adds tz-normalization to _get_temp_budget_increase: naive stored expiry is coerced to UTC before comparison with datetime.now(timezone.utc), fixing the TypeError raised when tz-aware ISO strings (e.g. 2026-01-20T00:00:00Z) were stored. timezone was already imported at line 16, so no new import needed.
tests/proxy_unit_tests/test_proxy_utils.py	Adds test_get_temp_budget_increase_tz_aware_expiry covering the future-expiry (returns increase) and past-expiry (returns None) cases using tz-aware ISO strings. The existing test_get_temp_budget_increase continues to cover the naive-datetime path.

_{Reviews (1): Last reviewed commit: "fix(proxy/auth): handle tz-aware temp_bu..." | Re-trigger Greptile}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[codspeed-hq]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing litellm_fix_temp_budget_expiry_tz (6261056) with litellm_internal_staging (66dea7d)}