From 787bee98162f3ad956402ce4377de8813d827ad3 Mon Sep 17 00:00:00 2001 From: Harshit Jain Date: Sat, 31 Jan 2026 17:49:53 +0530 Subject: [PATCH] fix: proxy failure cases, now log ip and user agent, key hash, name --- litellm/integrations/prometheus.py | 30 +--- litellm/proxy/auth/auth_checks.py | 8 +- litellm/proxy/auth/auth_exception_handler.py | 61 ++++++-- litellm/proxy/auth/auth_utils.py | 145 ++++++++++++------ litellm/proxy/auth/user_api_key_auth.py | 14 +- litellm/proxy/utils.py | 10 +- .../test_prometheus_invalid_key_filtering.py | 140 +++++++++++------ 7 files changed, 273 insertions(+), 135 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 2c897cb0692e..1fd67954ba72 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1478,35 +1478,19 @@ def _is_invalid_api_key_request( """ Determine if a request has an invalid API key based on status code and exception. - This method prevents invalid authentication attempts from being recorded in - Prometheus metrics. A 401 status code is the definitive indicator of authentication - failure. Additionally, we check exception messages for authentication error patterns - to catch cases where the exception hasn't been converted to a ProxyException yet. + Returns True only when we truly cannot record useful metrics (e.g. missing required + data). We no longer skip 401/invalid-key requests - all requests including + authentication failures and bad requests must be tracked for debugging, security + auditing, abuse detection, and capacity planning. Args: - status_code: HTTP status code (401 indicates authentication error) - exception: Exception object to check for auth-related error messages + status_code: HTTP status code + exception: Exception object (unused, kept for API compatibility) Returns: - True if the request has an invalid API key and metrics should be skipped, + True if metrics should be skipped (currently always False - track all requests), False otherwise """ - if status_code == 401: - return True - - # Handle cases where AssertionError is raised before conversion to ProxyException - if exception is not None: - exception_str = str(exception).lower() - auth_error_patterns = [ - "virtual key expected", - "expected to start with 'sk-'", - "authentication error", - "invalid api key", - "api key not valid", - ] - if any(pattern in exception_str for pattern in auth_error_patterns): - return True - return False def _should_skip_metrics_for_invalid_key( diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e0b056d450fb..81502b8698df 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -128,7 +128,7 @@ async def common_checks( message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}", type=ProxyErrorTypes.team_model_access_denied, param="model", - code=status.HTTP_401_UNAUTHORIZED, + code=status.HTTP_400_BAD_REQUEST, ) ## 2.1 If user can call model (if personal key) @@ -1983,7 +1983,7 @@ def _can_object_call_model( object_type=object_type ), param="model", - code=status.HTTP_401_UNAUTHORIZED, + code=status.HTTP_400_BAD_REQUEST, ) @@ -2084,7 +2084,7 @@ async def can_user_call_model( message=f"User not allowed to access model. No default model access, only team models allowed. Tried to access {model}", type=ProxyErrorTypes.key_model_access_denied, param="model", - code=status.HTTP_401_UNAUTHORIZED, + code=status.HTTP_400_BAD_REQUEST, ) return _can_object_call_model( @@ -2666,7 +2666,7 @@ def _can_object_call_vector_stores( object_type ), param="vector_store", - code=status.HTTP_401_UNAUTHORIZED, + code=status.HTTP_400_BAD_REQUEST, ) return True diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 9c306acd2c6a..90378bd8f3a6 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -9,7 +9,10 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth -from litellm.proxy.auth.auth_utils import _get_request_ip_address +from litellm.proxy.auth.auth_utils import ( + _get_request_ip_address, + add_client_context_to_request_data, +) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.types.services import ServiceTypes @@ -30,6 +33,7 @@ async def _handle_authentication_error( route: str, parent_otel_span: Optional[Span], api_key: str, + valid_token: Optional[UserAPIKeyAuth] = None, ) -> UserAPIKeyAuth: """ Handles Connection Errors when reading a Virtual Key from LiteLLM DB @@ -71,30 +75,63 @@ async def _handle_authentication_error( ) else: # raise the exception to the caller + use_x_forwarded_for = general_settings.get("use_x_forwarded_for", False) requester_ip = _get_request_ip_address( request=request, - use_x_forwarded_for=general_settings.get("use_x_forwarded_for", False), + use_x_forwarded_for=use_x_forwarded_for, + ) + user_agent = request.headers.get("user-agent", "") if request else "" + + # Ensure request_data has client context for callbacks (Prometheus, etc.) + add_client_context_to_request_data( + request=request, + request_data=request_data, + use_x_forwarded_for=use_x_forwarded_for, + ) + + key_name = ( + valid_token.key_alias or getattr(valid_token, "key_name", None) + if valid_token + else "" ) verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - {}\nRequester IP Address:{}".format( + "litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - {}\nRequester IP Address:{}\nUser-Agent:{}\nKey Hash:{}\nKey Name:{}".format( str(e), - requester_ip, + requester_ip or "", + user_agent or "", + api_key or "", + key_name, ), - extra={"requester_ip": requester_ip}, + extra={ + "requester_ip": requester_ip, + "user_agent": user_agent, + "key_hash": api_key, + "key_name": key_name, + }, ) - # Log this exception to OTEL, Datadog etc - user_api_key_dict = UserAPIKeyAuth( - parent_otel_span=parent_otel_span, - api_key=api_key, - request_route=route, - ) + # Log this exception to OTEL, Datadog etc - use valid_token when available (e.g. model access denied) + if valid_token is not None: + user_api_key_dict = valid_token + else: + user_api_key_dict = UserAPIKeyAuth( + parent_otel_span=parent_otel_span, + api_key=api_key, + request_route=route, + key_alias="", + ) # Allow callbacks to transform the error response + error_type = ProxyErrorTypes.auth_error + if isinstance(e, ProxyException) and hasattr(e, "type"): + try: + error_type = ProxyErrorTypes(e.type) + except (ValueError, TypeError): + pass transformed_exception = await proxy_logging_obj.post_call_failure_hook( request_data=request_data, original_exception=e, user_api_key_dict=user_api_key_dict, - error_type=ProxyErrorTypes.auth_error, + error_type=error_type, route=route, ) # Use transformed exception if callback returned one, otherwise use original diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 2bd84a1d98d9..4049dcad14f0 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -26,6 +26,28 @@ def _get_request_ip_address( return client_ip +def add_client_context_to_request_data( + request: Request, request_data: dict, use_x_forwarded_for: bool = False +) -> None: + """ + Add client_ip (requester_ip_address) and User-Agent to request_data metadata early. + Call this at the start of the request pipeline so failures have client context for logging. + """ + if "metadata" not in request_data: + request_data["metadata"] = {} + metadata = request_data["metadata"] + + requester_ip = _get_request_ip_address( + request=request, use_x_forwarded_for=use_x_forwarded_for + ) + metadata["requester_ip_address"] = requester_ip or "" + + user_agent = "" + if hasattr(request, "headers") and "user-agent" in request.headers: + user_agent = request.headers.get("user-agent", "") + metadata["user_agent"] = user_agent + + def _check_valid_ip( allowed_ips: Optional[List[str]], request: Request, @@ -314,17 +336,17 @@ def get_request_route(request: Request) -> str: def normalize_request_route(route: str) -> str: """ Normalize request routes by replacing dynamic path parameters with placeholders. - + This prevents high cardinality in Prometheus metrics by collapsing routes like: - /v1/responses/1234567890 -> /v1/responses/{response_id} - /v1/threads/thread_123 -> /v1/threads/{thread_id} - + Args: route: The request route path - + Returns: Normalized route with dynamic parameters replaced by placeholders - + Examples: >>> normalize_request_route("/v1/responses/abc123") '/v1/responses/{response_id}' @@ -337,58 +359,90 @@ def normalize_request_route(route: str) -> str: # Format: (regex_pattern, replacement_template) patterns = [ # Responses API - must come before generic patterns - (r'^(/(?:openai/)?v1/responses)/([^/]+)(/input_items)$', r'\1/{response_id}\3'), - (r'^(/(?:openai/)?v1/responses)/([^/]+)(/cancel)$', r'\1/{response_id}\3'), - (r'^(/(?:openai/)?v1/responses)/([^/]+)$', r'\1/{response_id}'), - (r'^(/responses)/([^/]+)(/input_items)$', r'\1/{response_id}\3'), - (r'^(/responses)/([^/]+)(/cancel)$', r'\1/{response_id}\3'), - (r'^(/responses)/([^/]+)$', r'\1/{response_id}'), - + (r"^(/(?:openai/)?v1/responses)/([^/]+)(/input_items)$", r"\1/{response_id}\3"), + (r"^(/(?:openai/)?v1/responses)/([^/]+)(/cancel)$", r"\1/{response_id}\3"), + (r"^(/(?:openai/)?v1/responses)/([^/]+)$", r"\1/{response_id}"), + (r"^(/responses)/([^/]+)(/input_items)$", r"\1/{response_id}\3"), + (r"^(/responses)/([^/]+)(/cancel)$", r"\1/{response_id}\3"), + (r"^(/responses)/([^/]+)$", r"\1/{response_id}"), # Threads API - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/steps)/([^/]+)$', r'\1/{thread_id}\3/{run_id}\5/{step_id}'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/steps)$', r'\1/{thread_id}\3/{run_id}\5'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/cancel)$', r'\1/{thread_id}\3/{run_id}\5'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/submit_tool_outputs)$', r'\1/{thread_id}\3/{run_id}\5'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)$', r'\1/{thread_id}\3/{run_id}'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/runs)$', r'\1/{thread_id}\3'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/messages)/([^/]+)$', r'\1/{thread_id}\3/{message_id}'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)(/messages)$', r'\1/{thread_id}\3'), - (r'^(/(?:openai/)?v1/threads)/([^/]+)$', r'\1/{thread_id}'), - + ( + r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/steps)/([^/]+)$", + r"\1/{thread_id}\3/{run_id}\5/{step_id}", + ), + ( + r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/steps)$", + r"\1/{thread_id}\3/{run_id}\5", + ), + ( + r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/cancel)$", + r"\1/{thread_id}\3/{run_id}\5", + ), + ( + r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)(/submit_tool_outputs)$", + r"\1/{thread_id}\3/{run_id}\5", + ), + ( + r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)/([^/]+)$", + r"\1/{thread_id}\3/{run_id}", + ), + (r"^(/(?:openai/)?v1/threads)/([^/]+)(/runs)$", r"\1/{thread_id}\3"), + ( + r"^(/(?:openai/)?v1/threads)/([^/]+)(/messages)/([^/]+)$", + r"\1/{thread_id}\3/{message_id}", + ), + (r"^(/(?:openai/)?v1/threads)/([^/]+)(/messages)$", r"\1/{thread_id}\3"), + (r"^(/(?:openai/)?v1/threads)/([^/]+)$", r"\1/{thread_id}"), # Vector Stores API - (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/files)/([^/]+)$', r'\1/{vector_store_id}\3/{file_id}'), - (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/files)$', r'\1/{vector_store_id}\3'), - (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/file_batches)/([^/]+)$', r'\1/{vector_store_id}\3/{batch_id}'), - (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)(/file_batches)$', r'\1/{vector_store_id}\3'), - (r'^(/(?:openai/)?v1/vector_stores)/([^/]+)$', r'\1/{vector_store_id}'), - + ( + r"^(/(?:openai/)?v1/vector_stores)/([^/]+)(/files)/([^/]+)$", + r"\1/{vector_store_id}\3/{file_id}", + ), + ( + r"^(/(?:openai/)?v1/vector_stores)/([^/]+)(/files)$", + r"\1/{vector_store_id}\3", + ), + ( + r"^(/(?:openai/)?v1/vector_stores)/([^/]+)(/file_batches)/([^/]+)$", + r"\1/{vector_store_id}\3/{batch_id}", + ), + ( + r"^(/(?:openai/)?v1/vector_stores)/([^/]+)(/file_batches)$", + r"\1/{vector_store_id}\3", + ), + (r"^(/(?:openai/)?v1/vector_stores)/([^/]+)$", r"\1/{vector_store_id}"), # Assistants API - (r'^(/(?:openai/)?v1/assistants)/([^/]+)$', r'\1/{assistant_id}'), - + (r"^(/(?:openai/)?v1/assistants)/([^/]+)$", r"\1/{assistant_id}"), # Files API - (r'^(/(?:openai/)?v1/files)/([^/]+)(/content)$', r'\1/{file_id}\3'), - (r'^(/(?:openai/)?v1/files)/([^/]+)$', r'\1/{file_id}'), - + (r"^(/(?:openai/)?v1/files)/([^/]+)(/content)$", r"\1/{file_id}\3"), + (r"^(/(?:openai/)?v1/files)/([^/]+)$", r"\1/{file_id}"), # Batches API - (r'^(/(?:openai/)?v1/batches)/([^/]+)(/cancel)$', r'\1/{batch_id}\3'), - (r'^(/(?:openai/)?v1/batches)/([^/]+)$', r'\1/{batch_id}'), - + (r"^(/(?:openai/)?v1/batches)/([^/]+)(/cancel)$", r"\1/{batch_id}\3"), + (r"^(/(?:openai/)?v1/batches)/([^/]+)$", r"\1/{batch_id}"), # Fine-tuning API - (r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/events)$', r'\1/{fine_tuning_job_id}\3'), - (r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/cancel)$', r'\1/{fine_tuning_job_id}\3'), - (r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/checkpoints)$', r'\1/{fine_tuning_job_id}\3'), - (r'^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)$', r'\1/{fine_tuning_job_id}'), - + ( + r"^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/events)$", + r"\1/{fine_tuning_job_id}\3", + ), + ( + r"^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/cancel)$", + r"\1/{fine_tuning_job_id}\3", + ), + ( + r"^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)(/checkpoints)$", + r"\1/{fine_tuning_job_id}\3", + ), + (r"^(/(?:openai/)?v1/fine_tuning/jobs)/([^/]+)$", r"\1/{fine_tuning_job_id}"), # Models API - (r'^(/(?:openai/)?v1/models)/([^/]+)$', r'\1/{model}'), + (r"^(/(?:openai/)?v1/models)/([^/]+)$", r"\1/{model}"), ] - + # Apply patterns in order for pattern, replacement in patterns: normalized = re.sub(pattern, replacement, route) if normalized != route: return normalized - + # Return original route if no pattern matched return route @@ -644,6 +698,7 @@ def get_customer_user_header_from_mapping(user_id_mapping) -> Optional[str]: return header_name return None + def _get_customer_id_from_standard_headers( request_headers: Optional[dict], ) -> Optional[str]: @@ -679,7 +734,9 @@ def get_end_user_id_from_request_body( from litellm.proxy.proxy_server import general_settings # Check 1: Standard customer ID headers (always checked, no configuration required) - customer_id = _get_customer_id_from_standard_headers(request_headers=request_headers) + customer_id = _get_customer_id_from_standard_headers( + request_headers=request_headers + ) if customer_id is not None: return customer_id diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 7290528cb5ad..92e947287eb2 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -42,6 +42,7 @@ from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler from litellm.proxy.auth.auth_utils import ( abbreviate_api_key, + add_client_context_to_request_data, get_end_user_id_from_request_body, get_model_from_request, get_request_route, @@ -247,7 +248,9 @@ async def get_global_proxy_spend( proxy_logging_obj: ProxyLogging, ) -> Optional[float]: global_proxy_spend = None - if litellm.max_budget > 0 and prisma_client is not None: # user set proxy max budget + if ( + litellm.max_budget > 0 and prisma_client is not None + ): # user set proxy max budget # Use event-driven coordination to prevent cache stampede cache_key = "{}:spend".format(litellm_proxy_admin_name) global_proxy_spend = await _fetch_global_spend_with_event_coordination( @@ -1219,6 +1222,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 route=route, parent_otel_span=parent_otel_span, api_key=api_key, + valid_token=valid_token, ) @@ -1246,6 +1250,14 @@ async def user_api_key_auth( request_data = populate_request_with_path_params( request_data=request_data, request=request ) + # Capture client context early so failures have IP/User-Agent for logging and metrics + from litellm.proxy.proxy_server import general_settings + + add_client_context_to_request_data( + request=request, + request_data=request_data, + use_x_forwarded_for=general_settings.get("use_x_forwarded_for", False), + ) route: str = get_request_route(request=request) ## CHECK IF ROUTE IS ALLOWED diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 6bbf0df74def..9877825f6716 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1633,8 +1633,16 @@ def _is_proxy_only_llm_api_error( if RouteChecks.is_llm_api_route(route) is not True: return False + proxy_only_error_types = { + ProxyErrorTypes.auth_error, + ProxyErrorTypes.key_model_access_denied, + ProxyErrorTypes.team_model_access_denied, + ProxyErrorTypes.user_model_access_denied, + ProxyErrorTypes.org_model_access_denied, + ProxyErrorTypes.token_not_found_in_db, + } return isinstance(original_exception, HTTPException) or ( - error_type == ProxyErrorTypes.auth_error + error_type in proxy_only_error_types ) async def _handle_logging_proxy_only_error( diff --git a/tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py b/tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py index ff433480d5e2..565b7f3ef7bd 100644 --- a/tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py +++ b/tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py @@ -1,8 +1,8 @@ """ -Unit tests for Prometheus invalid API key request filtering. +Unit tests for Prometheus request tracking. -Tests functionality that prevents invalid API key requests (401 status codes) -from being recorded in Prometheus metrics. +Tests that all requests including 401/invalid-key failures are tracked in metrics +for debugging, security auditing, abuse detection, and capacity planning. """ import os @@ -29,12 +29,14 @@ def prometheus_logger(): class ExceptionWithCode: """Exception-like object with 'code' attribute (ProxyException pattern).""" + def __init__(self, code): self.code = code class ExceptionWithStatusCode: """Exception-like object with 'status_code' attribute.""" + def __init__(self, status_code): self.status_code = status_code @@ -42,17 +44,25 @@ def __init__(self, status_code): class TestExtractStatusCode: """Test status code extraction from various sources.""" - @pytest.mark.parametrize("exception_class,code_value,expected", [ - (ExceptionWithCode, "401", 401), - (ExceptionWithStatusCode, 401, 401), - ]) - def test_extract_from_exception(self, prometheus_logger, exception_class, code_value, expected): + @pytest.mark.parametrize( + "exception_class,code_value,expected", + [ + (ExceptionWithCode, "401", 401), + (ExceptionWithStatusCode, 401, 401), + ], + ) + def test_extract_from_exception( + self, prometheus_logger, exception_class, code_value, expected + ): exception = exception_class(code_value) assert prometheus_logger._extract_status_code(exception=exception) == expected def test_extract_from_kwargs(self, prometheus_logger): exception = ExceptionWithCode("401") - assert prometheus_logger._extract_status_code(kwargs={"exception": exception}) == 401 + assert ( + prometheus_logger._extract_status_code(kwargs={"exception": exception}) + == 401 + ) def test_extract_from_enum_values(self, prometheus_logger): enum_values = Mock(status_code="401") @@ -60,45 +70,62 @@ def test_extract_from_enum_values(self, prometheus_logger): class TestInvalidAPIKeyDetection: - """Test invalid API key request detection logic.""" - - @pytest.mark.parametrize("status_code,expected", [ - (401, True), - (200, False), - (500, False), - (None, False), - ]) - def test_status_code_detection(self, prometheus_logger, status_code, expected): - assert prometheus_logger._is_invalid_api_key_request(status_code=status_code) == expected - - def test_auth_error_message_detection(self, prometheus_logger): - exception = AssertionError("LiteLLM Virtual Key expected. Received=invalid-key-12345, expected to start with 'sk-'.") - assert prometheus_logger._is_invalid_api_key_request(status_code=None, exception=exception) is True - - def test_non_auth_exception_not_detected(self, prometheus_logger): + """Test that we no longer skip metrics - all requests are tracked.""" + + @pytest.mark.parametrize("status_code", [401, 200, 500, None]) + def test_no_skip_for_any_status_code(self, prometheus_logger, status_code): + """All status codes are tracked - never skip metrics.""" + assert ( + prometheus_logger._is_invalid_api_key_request(status_code=status_code) + is False + ) + + def test_auth_error_message_not_skipped(self, prometheus_logger): + exception = AssertionError( + "LiteLLM Virtual Key expected. Received=invalid-key-12345, expected to start with 'sk-'." + ) + assert ( + prometheus_logger._is_invalid_api_key_request( + status_code=None, exception=exception + ) + is False + ) + + def test_non_auth_exception_not_skipped(self, prometheus_logger): exception = ValueError("Some other error") - assert prometheus_logger._is_invalid_api_key_request(status_code=None, exception=exception) is False + assert ( + prometheus_logger._is_invalid_api_key_request( + status_code=None, exception=exception + ) + is False + ) class TestSkipMetricsValidation: - """Test high-level validation method that orchestrates detection and extraction.""" + """Test that we never skip metrics - all requests are tracked.""" - def test_skip_for_401_exception(self, prometheus_logger): - """Test full flow: extraction -> detection -> skip decision.""" + def test_no_skip_for_401_exception(self, prometheus_logger): + """401 requests are now tracked for security auditing and abuse detection.""" exception = ExceptionWithCode("401") - assert prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception) is True + assert ( + prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception) + is False + ) - def test_skip_for_auth_error_message(self, prometheus_logger): - """Test full flow: exception message -> detection -> skip decision.""" + def test_no_skip_for_auth_error_message(self, prometheus_logger): + """Auth error messages are now tracked.""" exception = AssertionError("expected to start with 'sk-'") - assert prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception) is True + assert ( + prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception) + is False + ) def test_no_skip_for_valid_request(self, prometheus_logger): assert prometheus_logger._should_skip_metrics_for_invalid_key() is False class TestAsyncHooks: - """Test async hook methods skip metrics for invalid API keys.""" + """Test async hook methods record metrics for all requests including 401s.""" @pytest.fixture def mock_user_api_key(self): @@ -115,24 +142,33 @@ def mock_user_api_key(self): return user_key @pytest.mark.asyncio - async def test_post_call_failure_hook_skips_401(self, prometheus_logger, mock_user_api_key): + async def test_post_call_failure_hook_records_401( + self, prometheus_logger, mock_user_api_key + ): + """401 failures are now recorded for security auditing and abuse detection.""" exception = ExceptionWithCode("401") exception.__class__.__name__ = "ProxyException" - with patch.object(prometheus_logger, 'litellm_proxy_failed_requests_metric') as mock_failed, \ - patch.object(prometheus_logger, 'litellm_proxy_total_requests_metric') as mock_total: + with patch.object( + prometheus_logger, "litellm_proxy_failed_requests_metric" + ) as mock_failed, patch.object( + prometheus_logger, "litellm_proxy_total_requests_metric" + ) as mock_total: + mock_failed.labels.return_value = Mock(inc=Mock()) + mock_total.labels.return_value = Mock(inc=Mock()) await prometheus_logger.async_post_call_failure_hook( - request_data={"model": "test-model"}, + request_data={"model": "test-model", "metadata": {}}, original_exception=exception, - user_api_key_dict=mock_user_api_key + user_api_key_dict=mock_user_api_key, ) - mock_failed.labels.assert_not_called() - mock_total.labels.assert_not_called() + mock_failed.labels.assert_called_once() + mock_total.labels.assert_called_once() @pytest.mark.asyncio - async def test_log_failure_event_skips_401(self, prometheus_logger): + async def test_log_failure_event_records_401(self, prometheus_logger): + """401 failures in log_failure_event are now recorded.""" exception = ExceptionWithCode("401") kwargs = { "model": "test-model", @@ -140,6 +176,9 @@ async def test_log_failure_event_skips_401(self, prometheus_logger): "metadata": { "user_api_key_hash": "test-key", "user_api_key_user_id": "test-user", + "user_api_key_alias": None, + "user_api_key_team_id": None, + "user_api_key_team_alias": None, }, "model_group": "test-model", }, @@ -147,15 +186,16 @@ async def test_log_failure_event_skips_401(self, prometheus_logger): "litellm_params": {}, } - with patch.object(prometheus_logger, 'litellm_llm_api_failed_requests_metric') as mock_failed, \ - patch.object(prometheus_logger, 'set_llm_deployment_failure_metrics') as mock_deployment: + with patch.object( + prometheus_logger, "litellm_llm_api_failed_requests_metric" + ) as mock_failed, patch.object( + prometheus_logger, "set_llm_deployment_failure_metrics" + ) as mock_deployment: + mock_failed.labels.return_value = Mock(inc=Mock()) await prometheus_logger.async_log_failure_event( - kwargs=kwargs, - response_obj=None, - start_time=None, - end_time=None + kwargs=kwargs, response_obj=None, start_time=None, end_time=None ) - mock_failed.labels.assert_not_called() - mock_deployment.assert_not_called() + mock_failed.labels.assert_called_once() + mock_deployment.assert_called_once()