refactor(auth): derive temp budget increase without mutating the token (#34121)

* refactor(auth): derive temp budget bump without mutation, tz-aware auth datetimes

_update_key_budget_with_temp_budget_increase mutated max_budget in place, so correctness depended on every resolution path handing it a fresh copy of the cached token; one future re-cache of a live token would compound the bump per request. Return a model_copy instead so no caller can leak an increased budget into shared state.

Also fixes the three remaining DTZ005 naive datetime.now() calls in user_api_key_auth.py (auth span start, builder start_time, service-log end_time; all consumers convert to epoch or subtract same-pair datetimes) and ratchets the DTZ005 strict budget 244 -> 241.

* test: pin non-mutation of the temp budget helper input

Adversarial mutation-testing showed reverting the helper to in-place mutation still passed every test: the cache's copy-on-read layer masks the mutation in the integration test and the direct unit test only inspected the return value. Assert the input object is left untouched and the result is a distinct object so the purity guarantee itself is load-bearing.

litellm/proxy/auth/user_api_key_auth.py

@@ -1011,7 +1011,7 @@ def _ensure_parent_otel_span_on_request_state(request: Request) -> None:
         return
     if getattr(request.state, "parent_otel_span", None) is not None:
         return
-    start_time = datetime.now()
+    start_time = datetime.now(timezone.utc)
     try:
         request.state.litellm_received_at = start_time
     except Exception:
@@ -1061,7 +1061,7 @@ async def _user_api_key_auth_builder(
     # Prefer the receive-instant stamped by the early helper in
     # user_api_key_auth (before body parse) — overwriting it would shorten
     # the preprocessing-duration measurement by the body-parse window.
-    start_time = getattr(request.state, "litellm_received_at", None) or datetime.now()
+    start_time = getattr(request.state, "litellm_received_at", None) or datetime.now(timezone.utc)
     try:
         request.state.litellm_received_at = start_time
     except Exception:
@@ -2607,7 +2607,7 @@ async def _return_user_api_key_auth_obj(
     start_time: datetime,
     user_role: Optional[LitellmUserRoles] = None,
 ) -> UserAPIKeyAuth:
-    end_time = datetime.now()
+    end_time = datetime.now(timezone.utc)
 
     asyncio.create_task(
         user_api_key_service_logger_obj.async_service_success_hook(
@@ -2696,9 +2696,10 @@ def _update_key_budget_with_temp_budget_increase(
 ) -> UserAPIKeyAuth:
     if valid_token.max_budget is None:
         return valid_token
-    temp_budget_increase = _get_temp_budget_increase(valid_token) or 0.0
-    valid_token.max_budget = valid_token.max_budget + temp_budget_increase
-    return valid_token
+    temp_budget_increase = _get_temp_budget_increase(valid_token)
+    if not temp_budget_increase:
+        return valid_token
+    return valid_token.model_copy(update={"max_budget": valid_token.max_budget + temp_budget_increase})
 
 
 async def _lookup_end_user_and_apply_budget(

ruff-strict-budget.json

@@ -93,7 +93,7 @@
     "limit": 33
   },
   "DTZ005": {
-    "limit": 244
+    "limit": 241
   },
   "DTZ006": {
     "limit": 13

tests/proxy_unit_tests/test_proxy_utils.py

@@ -1780,7 +1780,10 @@ def test_update_key_budget_with_temp_budget_increase():
             "temp_budget_expiry": expiry_in_isoformat,
         },
     )
-    assert _update_key_budget_with_temp_budget_increase(valid_token).max_budget == 200
+    result = _update_key_budget_with_temp_budget_increase(valid_token)
+    assert result.max_budget == 200
+    assert result is not valid_token
+    assert valid_token.max_budget == 100
 
 
 @pytest.mark.asyncio

tests/test_litellm/proxy/auth/test_user_api_key_auth.py

@@ -4529,6 +4529,9 @@ async def test_temp_budget_increase_applied_for_cached_key():
     Seed the auth cache with a key whose spend (5.0) exceeds its original
     max_budget (2.0) but is under the effective budget (2.0 + 100.0). The cache-hit
     request must not raise and the resolved token must carry max_budget == 102.0.
+
+    Resolving twice must yield 102.0 both times and leave the cached object at the
+    original 2.0: the increase is derived per request, never compounded or persisted.
     """
     from datetime import datetime, timedelta
 
@@ -4574,14 +4577,22 @@ async def test_temp_budget_increase_applied_for_cached_key():
             new_callable=AsyncMock,
         ),
     ):
-        result = await _user_api_key_auth_builder(
-            request=mock_request,
-            api_key=f"Bearer {api_key}",
-            azure_api_key_header="",
-            anthropic_api_key_header=None,
-            google_ai_studio_api_key_header=None,
-            azure_apim_header=None,
-            request_data={"model": "gpt-4o-mini"},
+        results = tuple(
+            [
+                await _user_api_key_auth_builder(
+                    request=mock_request,
+                    api_key=f"Bearer {api_key}",
+                    azure_api_key_header="",
+                    anthropic_api_key_header=None,
+                    google_ai_studio_api_key_header=None,
+                    azure_apim_header=None,
+                    request_data={"model": "gpt-4o-mini"},
+                )
+                for _ in range(2)
+            ]
         )
 
-    assert result.max_budget == 102.0
+    assert all(result.max_budget == 102.0 for result in results)
+
+    cached_after = await user_api_key_cache.async_get_cache(key=hashed_token)
+    assert cached_after.max_budget == 2.0