From b9dd397746ba6d6579e65ece9f5fe756c21ba786 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 02:09:00 +0000 Subject: [PATCH] fix(prometheus): count 401 auth failures in litellm_proxy_failed_requests_metric Invalid or unknown virtual keys were filtered out of the proxy failed and total request counters entirely. Count them with hashed_api_key unset so caller-chosen key strings cannot create unbounded label series, and normalize the request route on the auth failure path so dynamic path ids do not leak into the route label either. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/prometheus.py | 8 +- litellm/proxy/auth/auth_exception_handler.py | 3 +- .../test_prometheus_invalid_key_filtering.py | 90 +++++++++++++------ .../proxy/auth/test_auth_exception_handler.py | 28 ++++++ 4 files changed, 94 insertions(+), 35 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 69a38e83835..09be00f2b7b 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -2610,12 +2610,6 @@ class PrometheusLogger(CustomLogger): StandardLoggingPayloadSetup, ) - if self._should_skip_metrics_for_invalid_key( - user_api_key_dict=user_api_key_dict, - exception=original_exception, - ): - return - status_code: Final = self._extract_status_code(exception=original_exception) try: @@ -2633,7 +2627,7 @@ class PrometheusLogger(CustomLogger): end_user=user_api_key_dict.end_user_id, user=user_api_key_dict.user_id, user_email=user_api_key_dict.user_email, - hashed_api_key=user_api_key_dict.api_key, + hashed_api_key=None if status_code == 401 else user_api_key_dict.api_key, api_key_alias=user_api_key_dict.key_alias, team=user_api_key_dict.team_id, team_alias=user_api_key_dict.team_alias, diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index b36c8a038fc..ba4c095c00f 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -23,6 +23,7 @@ from litellm.proxy.auth.auth_utils import ( _get_request_ip_address, is_invalid_virtual_key_error, mark_invalid_virtual_key_error, + normalize_request_route, ) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.types.services import ServiceTypes @@ -172,7 +173,7 @@ class UserAPIKeyAuthExceptionHandler: # so the handler is side-effect-free for the caller's identity object. user_api_key_dict = resolved_identity.model_copy() if resolved_identity is not None else UserAPIKeyAuth() user_api_key_dict.parent_otel_span = parent_otel_span - user_api_key_dict.request_route = route + user_api_key_dict.request_route = normalize_request_route(route) user_api_key_dict.api_key = user_api_key_dict.api_key or UserAPIKeyAuth(api_key=api_key).api_key # Stamp identity onto the request's server span now, before the request 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 278a4ef1df6..9e8348e860a 100644 --- a/tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py +++ b/tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py @@ -1,18 +1,18 @@ """ Unit tests for Prometheus invalid API key request filtering. -Tests functionality that prevents invalid API key requests (401 status codes) -from being recorded in Prometheus metrics. +Tests the 401 detection helpers, that LLM-level metrics skip invalid API key +requests, and that the proxy-level failed request counter still records them. """ from unittest.mock import Mock, patch import pytest +from fastapi import HTTPException from prometheus_client import REGISTRY - from litellm.integrations.prometheus import PrometheusLogger -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth @pytest.fixture(scope="function") @@ -129,28 +129,29 @@ class TestSkipMetricsValidation: class TestAsyncHooks: - """Test async hook methods skip metrics for invalid API keys.""" - - @pytest.fixture - def mock_user_api_key(self): - """Create a mock UserAPIKeyAuth object.""" - user_key = Mock(spec=UserAPIKeyAuth) - user_key.api_key = "test-key" - user_key.end_user_id = None - user_key.user_id = None - user_key.user_email = None - user_key.key_alias = None - user_key.team_id = None - user_key.team_alias = None - user_key.request_route = "/test" - return user_key + """Test how async hook methods treat invalid API key requests.""" @pytest.mark.asyncio - async def test_post_call_failure_hook_skips_401( - self, prometheus_logger, mock_user_api_key + @pytest.mark.parametrize( + "exception", + [ + HTTPException( + status_code=401, + detail="LiteLLM Virtual Key expected. Received=nota****tall, expected to start with 'sk-'.", + ), + ProxyException( + message="Authentication Error, Invalid proxy server token passed.", + type=ProxyErrorTypes.token_not_found_in_db, + param="key", + code=401, + ), + ], + ) + async def test_post_call_failure_hook_counts_401_without_key_hash( + self, prometheus_logger, exception ): - exception = ExceptionWithCode("401") - exception.__class__.__name__ = "ProxyException" + unauthenticated = UserAPIKeyAuth(request_route="/v1/chat/completions") + unauthenticated.api_key = "notakeyatall" with ( patch.object( @@ -160,15 +161,50 @@ class TestAsyncHooks: prometheus_logger, "litellm_proxy_total_requests_metric" ) as mock_total, ): - await prometheus_logger.async_post_call_failure_hook( request_data={"model": "test-model"}, original_exception=exception, - user_api_key_dict=mock_user_api_key, + user_api_key_dict=unauthenticated, ) - mock_failed.labels.assert_not_called() - mock_total.labels.assert_not_called() + failed_labels = mock_failed.labels.call_args.kwargs + assert failed_labels["exception_status"] == "401" + assert failed_labels["hashed_api_key"] is None + assert failed_labels["route"] == "/v1/chat/completions" + mock_failed.labels.return_value.inc.assert_called_once() + assert mock_total.labels.call_args.kwargs["status_code"] == "401" + mock_total.labels.return_value.inc.assert_called_once() + + @pytest.mark.asyncio + async def test_post_call_failure_hook_keeps_resolved_identity_labels_for_401( + self, prometheus_logger + ): + expired_key = UserAPIKeyAuth( + api_key="sk-expired", + key_alias="expired-alias", + team_id="team-1", + ) + exception = ProxyException( + message="Authentication Error - Expired Key.", + type=ProxyErrorTypes.expired_key, + param="key", + code=401, + ) + + with patch.object( + prometheus_logger, "litellm_proxy_failed_requests_metric" + ) as mock_failed: + await prometheus_logger.async_post_call_failure_hook( + request_data={"model": "test-model"}, + original_exception=exception, + user_api_key_dict=expired_key, + ) + + failed_labels = mock_failed.labels.call_args.kwargs + assert failed_labels["exception_status"] == "401" + assert failed_labels["hashed_api_key"] is None + assert failed_labels["api_key_alias"] == "expired-alias" + assert failed_labels["team"] == "team-1" @pytest.mark.asyncio async def test_log_failure_event_skips_401(self, prometheus_logger): diff --git a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py index 21e0b83791f..08a9d0ebf01 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -487,6 +487,34 @@ async def test_route_passed_to_post_call_failure_hook(): assert call_args["user_api_key_dict"].request_route == test_route +@pytest.mark.asyncio +async def test_dynamic_route_normalized_on_auth_failure(): + handler = UserAPIKeyAuthExceptionHandler() + + with ( + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + ) as mock_post_call_failure_hook, + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.proxy_server.general_settings", {} + ), + pytest.raises(ProxyException), + ): + await handler._handle_authentication_error( + HTTPException(status_code=401, detail="Authentication Error, Invalid proxy server token passed"), + MagicMock(), + {}, + "/v1/responses/resp_attacker_controlled_id", + None, + "sk-doesnotexist", + ) + + hook_kwargs = mock_post_call_failure_hook.call_args.kwargs + assert hook_kwargs["route"] == "/v1/responses/resp_attacker_controlled_id" + assert hook_kwargs["user_api_key_dict"].request_route == "/v1/responses/{response_id}" + + @pytest.mark.asyncio async def test_resolved_identity_exported_on_auth_failure(): """Regression: when auth fails AFTER the key/team/user identity is resolved