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