From 76c9eca25dec8bd87bc8e957bf381c0e6dea6b17 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 21 Jul 2026 14:02:40 -0700 Subject: [PATCH] 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 | 13 +++++---- ruff-strict-budget.json | 2 +- tests/proxy_unit_tests/test_proxy_utils.py | 5 +++- .../proxy/auth/test_user_api_key_auth.py | 29 +++++++++++++------ 4 files changed, 32 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 5b21a7265a0..4623c3d153c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/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( diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 448a0079674..d3d70ff5ff4 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -93,7 +93,7 @@ "limit": 33 }, "DTZ005": { - "limit": 244 + "limit": 241 }, "DTZ006": { "limit": 13 diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index d22b343d843..ee18c96c393 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/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 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 59b8228530a..d410f7ff72c 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 @@ -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