diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9f947e7fa4c..8dfaef6420a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2254,6 +2254,7 @@ class UserAPIKeyAuth( user_max_budget: Optional[float] = None request_route: Optional[str] = None user: Optional[Any] = None # Expanded user object when expand=user is used + auth_time_ms: Optional[float] = None # Time spent in authentication (milliseconds) model_config = ConfigDict(arbitrary_types_allowed=True) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a153c6e51cc..6fb869816c6 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -9,6 +9,7 @@ Returns a UserAPIKeyAuth object if the API key is valid import asyncio import secrets +import time from datetime import datetime, timezone from typing import List, Optional, Tuple, cast @@ -414,6 +415,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 parent_otel_span: Optional[Span] = None start_time = datetime.now() + auth_start_time = time.perf_counter() # For accurate timing measurement route: str = get_request_route(request=request) valid_token: Optional[UserAPIKeyAuth] = None custom_auth_api_key: bool = False @@ -856,6 +858,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 }, route=route, start_time=start_time, + auth_start_time=auth_start_time, ) asyncio.create_task( _cache_key_object( @@ -1242,6 +1245,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 valid_token_dict=valid_token_dict, route=route, start_time=start_time, + auth_start_time=auth_start_time, ) except Exception as e: return await UserAPIKeyAuthExceptionHandler._handle_authentication_error( @@ -1312,9 +1316,16 @@ async def _return_user_api_key_auth_obj( valid_token_dict: dict, route: str, start_time: datetime, + auth_start_time: Optional[float] = None, user_role: Optional[LitellmUserRoles] = None, ) -> UserAPIKeyAuth: end_time = datetime.now() + + # Calculate auth time using perf_counter for accurate timing + auth_time_ms = None + if auth_start_time is not None: + auth_end_time = time.perf_counter() + auth_time_ms = (auth_end_time - auth_start_time) * 1000 asyncio.create_task( user_api_key_service_logger_obj.async_service_success_hook( @@ -1335,6 +1346,7 @@ async def _return_user_api_key_auth_obj( "api_key": api_key, "parent_otel_span": parent_otel_span, "user_role": retrieved_user_role, + "auth_time_ms": auth_time_ms, **valid_token_dict, } if user_obj is not None: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 9be78264e85..dd86a10151b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -621,6 +621,10 @@ class LiteLLMProxyRequestSetup: ) # Add the full UserAPIKeyAuth object for MCP server access control data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict + # Extract auth_time_ms from UserAPIKeyAuth object if available + auth_time_ms = getattr(user_api_key_dict, "auth_time_ms", None) + if auth_time_ms is not None: + data[_metadata_variable_name]["auth_time_ms"] = auth_time_ms return data @staticmethod diff --git a/tests/test_litellm/proxy/auth/test_auth_time_tracking.py b/tests/test_litellm/proxy/auth/test_auth_time_tracking.py new file mode 100644 index 00000000000..c2be7545633 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_auth_time_tracking.py @@ -0,0 +1,149 @@ +""" +Tests for auth time tracking. + +Verifies that auth_time_ms is correctly calculated and stored in request metadata. +""" + +import asyncio +import os +import sys +import time +from datetime import datetime +from unittest.mock import AsyncMock, Mock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.proxy.auth.user_api_key_auth import _return_user_api_key_auth_obj, _user_api_key_auth_builder +from litellm.proxy._types import UserAPIKeyAuth + + +class TestAuthTimeTracking: + """Test suite for auth time tracking.""" + + def test_auth_time_is_calculated_and_stored(self): + """ + Test that auth_time_ms is calculated and stored in UserAPIKeyAuth object + when authentication completes successfully. + """ + # Create mock user object with proper attribute values + # Use spec to ensure only defined attributes exist + mock_user_obj = Mock(spec=['tpm_limit', 'rpm_limit', 'user_email', 'spend', 'max_budget']) + mock_user_obj.tpm_limit = 1000 + mock_user_obj.rpm_limit = 100 + mock_user_obj.user_email = "test@example.com" + # Set actual None values for optional fields + mock_user_obj.spend = None + mock_user_obj.max_budget = None + + # Create valid_token_dict + valid_token_dict = { + "api_key": "sk-test-key", + "user_id": "user-123", + "team_id": "team-456", + } + + start_time = datetime.now() + auth_start_time = time.perf_counter() + # Simulate some auth processing time + time.sleep(0.005) # 5ms delay + + # Patch helper functions to avoid Mock attribute issues + with patch('litellm.proxy.auth.user_api_key_auth._get_user_role', return_value=None), \ + patch('litellm.proxy.auth.user_api_key_auth._is_user_proxy_admin', return_value=False): + # Call _return_user_api_key_auth_obj using asyncio.run() + result = asyncio.run(_return_user_api_key_auth_obj( + user_obj=mock_user_obj, + api_key="sk-test-key", + parent_otel_span=None, + valid_token_dict=valid_token_dict, + route="/chat/completions", + start_time=start_time, + auth_start_time=auth_start_time, + user_role=None, + )) + + # Verify auth_time_ms is stored in the UserAPIKeyAuth object + assert hasattr(result, "auth_time_ms"), "auth_time_ms should be an attribute of UserAPIKeyAuth" + auth_time = result.auth_time_ms + + # Verify it's a numeric value + assert isinstance(auth_time, (int, float)), f"auth_time_ms should be numeric, got {type(auth_time)}" + + # Verify it's positive (should be at least 5ms due to our delay) + assert auth_time > 0, f"auth_time_ms should be positive, got {auth_time}" + assert auth_time >= 4.0, f"auth_time_ms should be at least 4ms, got {auth_time}" + + def test_auth_time_is_zero_for_instant_auth(self): + """ + Test that auth_time_ms is still calculated even when + authentication is very fast (near-zero time). + """ + # Create mock user object with proper attribute values + mock_user_obj = Mock(spec=['tpm_limit', 'rpm_limit', 'user_email', 'spend', 'max_budget']) + mock_user_obj.tpm_limit = 1000 + mock_user_obj.rpm_limit = 100 + mock_user_obj.user_email = "test@example.com" + mock_user_obj.spend = None + mock_user_obj.max_budget = None + + valid_token_dict = { + "api_key": "sk-test-key", + "user_id": "user-123", + } + + start_time = datetime.now() + auth_start_time = time.perf_counter() + + # Patch helper functions to avoid Mock attribute issues + with patch('litellm.proxy.auth.user_api_key_auth._get_user_role', return_value=None), \ + patch('litellm.proxy.auth.user_api_key_auth._is_user_proxy_admin', return_value=False): + # Call immediately (no delay) + result = asyncio.run(_return_user_api_key_auth_obj( + user_obj=mock_user_obj, + api_key="sk-test-key", + parent_otel_span=None, + valid_token_dict=valid_token_dict, + route="/chat/completions", + start_time=start_time, + auth_start_time=auth_start_time, + user_role=None, + )) + + # Verify timing is still tracked even for fast operations + assert hasattr(result, "auth_time_ms"), "auth_time_ms should be an attribute" + auth_time = result.auth_time_ms + assert isinstance(auth_time, (int, float)) + # Should be >= 0 (can be very small but should be tracked) + assert auth_time >= 0, f"auth_time_ms should be >= 0, got {auth_time}" + + def test_auth_time_stored_in_metadata(self): + """ + Test that auth_time_ms from UserAPIKeyAuth is extracted and stored + in request metadata by add_user_api_key_auth_to_request_metadata(). + """ + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + # Create a UserAPIKeyAuth object with auth_time_ms + user_api_key_auth = UserAPIKeyAuth( + api_key="sk-test-key", + user_id="user-123", + auth_time_ms=12.5, # Simulated auth time + ) + + # Create data dict with metadata + data = { + "metadata": {}, + } + + # Call add_user_api_key_auth_to_request_metadata + result_data = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( + data=data, + user_api_key_dict=user_api_key_auth, + _metadata_variable_name="metadata", + ) + + # Verify auth_time_ms is in metadata + assert "auth_time_ms" in result_data["metadata"], "auth_time_ms should be in metadata" + assert result_data["metadata"]["auth_time_ms"] == 12.5, f"Expected 12.5, got {result_data['metadata']['auth_time_ms']}" diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 5504aee3cac..93c042a6f59 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -913,6 +913,86 @@ def test_get_logging_payload_extracts_response_translation_time_from_model_call_ ), f"Expected 12.7, got {overhead_breakdown.get('response_translation_time_ms')}" +@patch("litellm.proxy.proxy_server.master_key", None) +@patch("litellm.proxy.proxy_server.general_settings", {}) +def test_get_logging_payload_extracts_auth_time_from_metadata(): + """ + Integration test: Verify that auth_time_ms stored in request metadata + gets extracted and included in overhead_breakdown. + + This tests the full flow: + 1. auth_time_ms is stored in request metadata (by add_user_api_key_auth_to_request_metadata) + 2. get_logging_payload extracts it from metadata + 3. It's included in overhead_breakdown in the metadata + """ + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObj, + ) + + # Create a logging_obj + logging_obj = LiteLLMLoggingObj( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "test"}], + stream=False, + call_type="completion", + start_time=None, + litellm_call_id="test-call-id", + function_id="test-function-id", + ) + + # Create kwargs with auth_time_ms in metadata (simulating what add_user_api_key_auth_to_request_metadata does) + kwargs = { + "model": "gpt-3.5-turbo", + "logging_obj": logging_obj, + "litellm_params": { + "metadata": { + "user_api_key": "sk-test-key", + "auth_time_ms": 8.3, # Simulated auth time from metadata + } + }, + "call_type": "completion", + } + + response_obj = { + "id": "test-response-789", + "choices": [{"message": {"content": "Hello!"}}], + "usage": { + "total_tokens": 100, + "prompt_tokens": 50, + "completion_tokens": 50, + }, + } + + start_time = datetime.datetime.now(timezone.utc) + end_time = datetime.datetime.now(timezone.utc) + + payload = get_logging_payload( + kwargs=kwargs, + response_obj=response_obj, + start_time=start_time, + end_time=end_time, + ) + + # Parse the metadata JSON string + metadata_json = payload.get("metadata") + assert metadata_json is not None, "metadata should not be None" + + metadata = json.loads(metadata_json) + + # Verify overhead_breakdown exists and contains auth_time_ms + overhead_breakdown = metadata.get("overhead_breakdown") + assert overhead_breakdown is not None, "overhead_breakdown should be present" + assert isinstance(overhead_breakdown, dict), "overhead_breakdown should be a dict" + + assert ( + "auth_time_ms" in overhead_breakdown + ), "auth_time_ms should be in overhead_breakdown" + + assert ( + overhead_breakdown["auth_time_ms"] == 8.3 + ), f"Expected 8.3, got {overhead_breakdown.get('auth_time_ms')}" + + @patch("litellm.proxy.proxy_server.master_key", None) @patch("litellm.proxy.proxy_server.general_settings", {}) def test_get_logging_payload_extracts_overhead_breakdown_from_multiple_sources():