From 6646bb4319367fdcba28a78cf96b19aa20dc97e4 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Sat, 29 Aug 2026 14:06:43 -0700 Subject: [PATCH] fix(logging): let invalid-key warnings inherit the proxy logger thresholds The dedicated stdout child logger no longer carries its own handler, level, or propagate flag, so the malformed virtual key warning obeys exactly the LITELLM_LOG threshold, handlers, and root propagation configured for "LiteLLM Proxy" in plain, JSON, and debug modes. The auth log line is written before the failure callbacks run, as before this change, so a callback that raises can no longer lose it, and it is classified on the original failure. The ProxyException conversion moves into a helper. Test patches carry test-quality reasons and the websocket regression test lives in the mapped test file. --- litellm/_logging.py | 19 +- litellm/proxy/auth/auth_exception_handler.py | 97 ++++--- litellm/proxy/auth/auth_utils.py | 2 +- .../test_user_api_key_auth.py | 38 --- .../proxy/auth/test_auth_exception_handler.py | 262 ++++++++---------- .../proxy/auth/test_user_api_key_auth.py | 43 +++ tests/test_litellm/test_logging.py | 53 +++- 7 files changed, 243 insertions(+), 271 deletions(-) diff --git a/litellm/_logging.py b/litellm/_logging.py index c2e540dd714..0ecab9b754d 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -510,23 +510,16 @@ else: handler.setFormatter(formatter) -verbose_proxy_logger: Final = logging.getLogger("LiteLLM Proxy") +verbose_proxy_logger = logging.getLogger("LiteLLM Proxy") +# Malformed virtual key rejections log through this child; LevelRoutingStreamHandler +# writes its WARNING records to stdout. It has no handler or level of its own. verbose_proxy_stdout_logger: Final = verbose_proxy_logger.getChild("stdout") -verbose_router_logger: Final = logging.getLogger("LiteLLM Router") -verbose_logger: Final = logging.getLogger("LiteLLM") - -verbose_proxy_stdout_handler: Final = LevelRoutingStreamHandler() -verbose_proxy_stdout_handler.setLevel(logging.WARNING) -verbose_proxy_stdout_handler.setFormatter(handler.formatter) -verbose_proxy_stdout_handler.addFilter(_secret_filter) -verbose_proxy_stdout_handler.addFilter(_correlation_filter) +verbose_router_logger = logging.getLogger("LiteLLM Router") +verbose_logger = logging.getLogger("LiteLLM") # Add the handler to the loggers verbose_router_logger.addHandler(handler) verbose_proxy_logger.addHandler(handler) -verbose_proxy_stdout_logger.setLevel(logging.WARNING) -verbose_proxy_stdout_logger.addHandler(verbose_proxy_stdout_handler) -verbose_proxy_stdout_logger.propagate = False verbose_logger.addHandler(handler) # Filters attached to the logger, not the handler, survive callers swapping in their own @@ -592,7 +585,6 @@ ALL_LOGGERS: Final = [ verbose_logger, verbose_router_logger, verbose_proxy_logger, - verbose_proxy_stdout_logger, ] @@ -700,7 +692,6 @@ def _turn_on_json(): handler: Final = LevelRoutingStreamHandler() handler.setFormatter(JsonFormatter()) _initialize_loggers_with_handler(handler) - verbose_proxy_stdout_logger.setLevel(logging.WARNING) # Set up exception handlers _setup_json_exception_handlers(JsonFormatter()) diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index ddb46c8821d..464cb0a4caf 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -54,6 +54,41 @@ def _with_requester_ip_address(request_data: dict[str, object], requester_ip: st return {**request_data, key: {**base, "requester_ip_address": requester_ip}} # mutable-ok: logging needs dicts +def _as_proxy_exception(e: Exception) -> ProxyException: + """Convert an authentication failure into the ProxyException the client receives.""" + if isinstance(e, litellm.BudgetExceededError): + return ProxyException( + message=e.message, + type=ProxyErrorTypes.budget_exceeded, + param=None, + code=getattr(e, "status_code", status.HTTP_429_TOO_MANY_REQUESTS), + ) + if isinstance(e, HTTPException): + return ProxyException( + message=getattr(e, "detail", f"Authentication Error({e})"), + type=ProxyErrorTypes.auth_error, + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_401_UNAUTHORIZED), + ) + if isinstance(e, ProxyException): + return e + if PrismaDBExceptionHandler.is_database_service_unavailable_error(e): + return ProxyException( + message=( + "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly." + ), + type=ProxyErrorTypes.no_db_connection, + param="None", + code=status.HTTP_503_SERVICE_UNAVAILABLE, + ) + return ProxyException( + message="Authentication Error, " + str(e), + type=ProxyErrorTypes.auth_error, + param=getattr(e, "param", "None"), + code=status.HTTP_401_UNAUTHORIZED, + ) + + class UserAPIKeyAuthExceptionHandler: @staticmethod async def _handle_authentication_error( @@ -115,8 +150,17 @@ class UserAPIKeyAuthExceptionHandler: request=request, use_x_forwarded_for=general_settings.get("use_x_forwarded_for") is True, ) - original_exception: Final = e is_invalid_virtual_key: Final = is_invalid_virtual_key_error(e) + is_quiet_log: Final = is_invalid_virtual_key and not litellm.log_client_error_tracebacks + logger: Final = verbose_proxy_stdout_logger if is_quiet_log else verbose_proxy_logger + logger.log( + logging.WARNING if is_quiet_log else logging.ERROR, + "litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - %s\nRequester IP Address:%s", + e, + requester_ip, + exc_info=True if litellm.log_client_error_tracebacks or not is_expected_client_error(e) else None, + extra={"requester_ip": requester_ip}, + ) # Log this exception to OTEL, Datadog etc. Reuse the identity resolved # before the failure (team alias/id, metadata, user) so the failed span @@ -154,7 +198,7 @@ class UserAPIKeyAuthExceptionHandler: # Allow callbacks to transform the error response transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook( request_data=_with_requester_ip_address(request_data, requester_ip), - original_exception=original_exception, + original_exception=e, user_api_key_dict=user_api_key_dict, error_type=ProxyErrorTypes.auth_error, route=route, @@ -163,51 +207,4 @@ class UserAPIKeyAuthExceptionHandler: if transformed_exception is not None: e = transformed_exception - proxy_exception: Final - if isinstance(e, litellm.BudgetExceededError): - proxy_exception = ProxyException( - message=e.message, - type=ProxyErrorTypes.budget_exceeded, - param=None, - code=getattr(e, "status_code", status.HTTP_429_TOO_MANY_REQUESTS), - ) - elif isinstance(e, HTTPException): - proxy_exception = ProxyException( - message=getattr(e, "detail", f"Authentication Error({e})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_401_UNAUTHORIZED), - ) - elif isinstance(e, ProxyException): - proxy_exception = e - elif PrismaDBExceptionHandler.is_database_service_unavailable_error(e): - proxy_exception = ProxyException( - message=( - "Service Unavailable, the authentication database is " - "temporarily unreachable. Please retry shortly." - ), - type=ProxyErrorTypes.no_db_connection, - param="None", - code=status.HTTP_503_SERVICE_UNAVAILABLE, - ) - else: - proxy_exception = ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_401_UNAUTHORIZED, - ) - final_exception: Final = mark_invalid_virtual_key_error(proxy_exception, is_invalid_virtual_key) - is_quiet_log: Final = ( - is_invalid_virtual_key_error(final_exception) and not litellm.log_client_error_tracebacks - ) - logger: Final = verbose_proxy_stdout_logger if is_quiet_log else verbose_proxy_logger - logger.log( - logging.WARNING if is_quiet_log else logging.ERROR, - "litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - %s\nRequester IP Address:%s", - final_exception, - requester_ip, - exc_info=not (is_expected_client_error(original_exception) and not litellm.log_client_error_tracebacks), - extra={"requester_ip": requester_ip}, - ) - raise final_exception + raise mark_invalid_virtual_key_error(_as_proxy_exception(e), is_invalid_virtual_key) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index ccb8e350892..0b29a3e3fdf 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -76,7 +76,7 @@ def mark_invalid_virtual_key_error(exception: ProxyException, is_invalid_virtual param=exception.param, code=exception.code, headers=exception.headers.copy(), - openai_code=exception.openai_code, + openai_code=None if exception.openai_code is None else str(exception.openai_code), provider_specific_fields=exception.provider_specific_fields, ) setattr(marked_exception, _INVALID_VIRTUAL_KEY_ERROR_MARKER, True) diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index 66c3d62f5c4..0cdf3500d50 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -874,44 +874,6 @@ async def test_user_api_key_auth_websocket(): assert mock_user_api_key_auth.call_args.kwargs["api_key"] == "Bearer some_api_key" -@pytest.mark.asyncio -async def test_user_api_key_auth_websocket_skips_duplicate_invalid_key_log(): - from litellm.proxy._types import ProxyException - from litellm.proxy.auth.auth_utils import mark_invalid_virtual_key_error - from litellm.proxy.auth.user_api_key_auth import WebSocketException, user_api_key_auth_websocket - - mock_websocket = MagicMock(spec=WebSocket) - mock_websocket.query_params = {"model": "some_model"} - mock_websocket.headers = {"authorization": "Bearer undefined"} - mock_websocket.scope = {"headers": [(b"authorization", b"Bearer undefined")]} - mock_websocket.url = URL(url="/v1/responses") - transformed_invalid_virtual_key = mark_invalid_virtual_key_error( - ProxyException( - message="Please authenticate again", - type="auth_error", - param="None", - code=status.HTTP_401_UNAUTHORIZED, - ), - True, - ) - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", - side_effect=transformed_invalid_virtual_key, - autospec=True, - ), - patch("litellm.proxy.auth.user_api_key_auth.verbose_proxy_logger.exception") as exception_log, - pytest.raises(WebSocketException) as exc_info, - ): - await user_api_key_auth_websocket(mock_websocket) - - assert exc_info.value.code == status.WS_1008_POLICY_VIOLATION - assert exc_info.value.reason == "" - exception_log.assert_not_called() - mock_websocket.close.assert_not_called() - - @pytest.mark.asyncio async def test_user_api_key_auth_websocket_carries_asgi_path(): """ 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 cf34e35f65a..b20c7f3cdc7 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -1,6 +1,7 @@ import asyncio import json import logging +from contextlib import contextmanager from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -117,22 +118,12 @@ async def test_handle_authentication_error_permanent_fault_gets_no_fallback_iden "prisma_error", [ DataError(data={"user_facing_error": {"meta": {"table": "test_table"}}}), - UniqueViolationError( - data={"user_facing_error": {"meta": {"table": "test_table"}}} - ), - ForeignKeyViolationError( - data={"user_facing_error": {"meta": {"table": "test_table"}}} - ), - MissingRequiredValueError( - data={"user_facing_error": {"meta": {"table": "test_table"}}} - ), + UniqueViolationError(data={"user_facing_error": {"meta": {"table": "test_table"}}}), + ForeignKeyViolationError(data={"user_facing_error": {"meta": {"table": "test_table"}}}), + MissingRequiredValueError(data={"user_facing_error": {"meta": {"table": "test_table"}}}), RawQueryError(data={"user_facing_error": {"meta": {"table": "test_table"}}}), - TableNotFoundError( - data={"user_facing_error": {"meta": {"table": "test_table"}}} - ), - RecordNotFoundError( - data={"user_facing_error": {"meta": {"table": "test_table"}}} - ), + TableNotFoundError(data={"user_facing_error": {"meta": {"table": "test_table"}}}), + RecordNotFoundError(data={"user_facing_error": {"meta": {"table": "test_table"}}}), ], ) async def test_handle_authentication_error_data_layer_errors_do_not_fall_back( @@ -781,46 +772,29 @@ async def test_handle_authentication_error_traceback_only_for_unexpected_errors( 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, - return_value=None, - ), - patch( # test-quality-ok: handler reads proxy_server globals at call time - "litellm.proxy.auth.auth_exception_handler.seed_request_identity", - ), + _auth_failure_boundaries(return_value=None), patch( # test-quality-ok: handler reads this global flag at call time "litellm.proxy.auth.auth_exception_handler.litellm.log_client_error_tracebacks", log_client_error_tracebacks, ), - patch( # test-quality-ok: handler reads proxy_server globals at call time - "litellm.proxy.proxy_server.general_settings", - {"allow_requests_on_db_unavailable": False}, - ), ): - verbose_proxy_logger.propagate = True - verbose_proxy_stdout_logger.propagate = True try: - try: - raise auth_error - except (HTTPException, ProxyException, ValueError) as caught: - with caplog.at_level(logging.DEBUG), pytest.raises(ProxyException): - await handler._handle_authentication_error( - caught, - MagicMock(), - {}, - "/v1/chat/completions", - None, - "sk-bad-key", - ) - finally: - verbose_proxy_logger.propagate = False - verbose_proxy_stdout_logger.propagate = False + raise auth_error + except (HTTPException, ProxyException, ValueError) as caught: + with caplog.at_level(logging.DEBUG), pytest.raises(ProxyException): + await handler._handle_authentication_error( + caught, + MagicMock(), + {}, + "/v1/chat/completions", + None, + "sk-bad-key", + ) - records = [r for r in caplog.records if "user_api_key_auth(): Exception occured" in r.getMessage()] + records = _auth_failure_records(caplog) assert len(records) == 1 assert records[0].levelno == expected_level - assert bool(records[0].exc_info) is expect_traceback + assert (records[0].exc_info is not None and records[0].exc_info[1] is auth_error) is expect_traceback assert records[0].name == ( verbose_proxy_stdout_logger.name if expected_level == logging.WARNING and not log_client_error_tracebacks @@ -828,145 +802,127 @@ async def test_handle_authentication_error_traceback_only_for_unexpected_errors( ) -@pytest.mark.asyncio -@pytest.mark.parametrize( - "transformed_exception", - [ - HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Please authenticate again"), - ProxyException( - message="Please authenticate again", - type=ProxyErrorTypes.auth_error, - param=None, - code=status.HTTP_401_UNAUTHORIZED, +@contextmanager +def _auth_failure_boundaries(**failure_hook_kwargs): + """Stub the proxy globals the handler reads and let its records reach caplog.""" + 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, + **failure_hook_kwargs, ), - ], -) -async def test_handle_authentication_error_preserves_invalid_virtual_key_marker_after_callback_transform( - caplog, - transformed_exception, -): - handler = UserAPIKeyAuthExceptionHandler() - original_exception = HTTPException( + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.auth.auth_exception_handler.seed_request_identity" + ), + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.proxy_server.general_settings", {"allow_requests_on_db_unavailable": False} + ), + ): + verbose_proxy_logger.propagate = True + yield + + +def _auth_failure_records(caplog) -> list[logging.LogRecord]: + return [r for r in caplog.records if "user_api_key_auth(): Exception occured" in r.getMessage()] + + +def _invalid_virtual_key_error() -> HTTPException: + return HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="LiteLLM Virtual Key expected. Received=unde****ined, expected to start with 'sk-'.", ) - with ( - patch( - "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", - new_callable=AsyncMock, - return_value=transformed_exception, - ), - patch("litellm.proxy.auth.auth_exception_handler.seed_request_identity"), - patch("litellm.proxy.proxy_server.general_settings", {"allow_requests_on_db_unavailable": False}), - ): - verbose_proxy_stdout_logger.propagate = True - try: - with pytest.raises(ProxyException) as exc_info: - await handler._handle_authentication_error( - original_exception, - MagicMock(), - {}, - "/v1/chat/completions", - None, - "undefined", - ) - finally: - verbose_proxy_stdout_logger.propagate = False - records = [r for r in caplog.records if "user_api_key_auth(): Exception occured" in r.getMessage()] +@pytest.mark.asyncio +async def test_handle_authentication_error_preserves_invalid_virtual_key_marker_after_callback_transform(caplog): + handler = UserAPIKeyAuthExceptionHandler() + transformed_exception = HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Please authenticate again") + + with _auth_failure_boundaries(return_value=transformed_exception), pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error( + _invalid_virtual_key_error(), MagicMock(), {}, "/v1/chat/completions", None, "undefined" + ) + + records = _auth_failure_records(caplog) assert len(records) == 1 assert records[0].name == verbose_proxy_stdout_logger.name assert records[0].levelno == logging.WARNING + assert records[0].exc_info is None + assert "LiteLLM Virtual Key expected" in records[0].getMessage() assert exc_info.value.message == "Please authenticate again" + assert exc_info.value.code == str(status.HTTP_401_UNAUTHORIZED) assert getattr(exc_info.value, "_litellm_invalid_virtual_key_error") is True - if isinstance(transformed_exception, ProxyException): - assert not hasattr(transformed_exception, "_litellm_invalid_virtual_key_error") @pytest.mark.asyncio -async def test_handle_authentication_error_keeps_unexpected_source_traceback_after_callback_4xx( - caplog, -): +async def test_handle_authentication_error_keeps_unexpected_source_traceback_after_callback_4xx(caplog): handler = UserAPIKeyAuthExceptionHandler() + transformed_exception = HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Please authenticate again") - with ( - patch( - "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", - new_callable=AsyncMock, - return_value=HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail="Please authenticate again", - ), - ), - patch("litellm.proxy.auth.auth_exception_handler.seed_request_identity"), - patch("litellm.proxy.proxy_server.general_settings", {"allow_requests_on_db_unavailable": False}), - ): - verbose_proxy_logger.propagate = True - try: - with pytest.raises(ProxyException) as exc_info: - await handler._handle_authentication_error( - ValueError("unexpected internal error"), - MagicMock(), - {}, - "/v1/chat/completions", - None, - "sk-bad-key", - ) - finally: - verbose_proxy_logger.propagate = False + try: + raise ValueError("unexpected internal error") + except ValueError as caught: + source_error = caught + with _auth_failure_boundaries(return_value=transformed_exception), pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error( + source_error, MagicMock(), {}, "/v1/chat/completions", None, "sk-bad-key" + ) - records = [r for r in caplog.records if "user_api_key_auth(): Exception occured" in r.getMessage()] + records = _auth_failure_records(caplog) assert len(records) == 1 assert records[0].name == verbose_proxy_logger.name assert records[0].levelno == logging.ERROR assert records[0].exc_info is not None - assert "Please authenticate again" in records[0].getMessage() + assert records[0].exc_info[1] is source_error + assert "unexpected internal error" in records[0].getMessage() + assert "Please authenticate again" not in records[0].getMessage() + assert exc_info.value.message == "Please authenticate again" assert exc_info.value.code == str(status.HTTP_401_UNAUTHORIZED) + assert not hasattr(exc_info.value, "_litellm_invalid_virtual_key_error") @pytest.mark.asyncio -async def test_handle_authentication_error_does_not_preserve_invalid_virtual_key_marker_for_callback_503( - caplog, -): +async def test_handle_authentication_error_does_not_preserve_invalid_virtual_key_marker_for_callback_503(caplog): handler = UserAPIKeyAuthExceptionHandler() - original_exception = HTTPException( - status_code=status.HTTP_401_UNAUTHORIZED, - detail="LiteLLM Virtual Key expected. Received=unde****ined, expected to start with 'sk-'.", - ) transformed_exception = HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Authentication service temporarily unavailable", ) - with ( - patch( - "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", - new_callable=AsyncMock, - return_value=transformed_exception, - ), - patch("litellm.proxy.auth.auth_exception_handler.seed_request_identity"), - patch("litellm.proxy.proxy_server.general_settings", {"allow_requests_on_db_unavailable": False}), - ): - verbose_proxy_logger.propagate = True - try: - with pytest.raises(ProxyException) as exc_info: - await handler._handle_authentication_error( - original_exception, - MagicMock(), - {}, - "/v1/chat/completions", - None, - "undefined", - ) - finally: - verbose_proxy_logger.propagate = False + with _auth_failure_boundaries(return_value=transformed_exception), pytest.raises(ProxyException) as exc_info: + await handler._handle_authentication_error( + _invalid_virtual_key_error(), MagicMock(), {}, "/v1/chat/completions", None, "undefined" + ) - records = [r for r in caplog.records if "user_api_key_auth(): Exception occured" in r.getMessage()] + records = _auth_failure_records(caplog) assert len(records) == 1 - assert records[0].name == verbose_proxy_logger.name - assert records[0].levelno == logging.ERROR - assert records[0].exc_info is not None - assert "Authentication service temporarily unavailable" in records[0].getMessage() + assert records[0].name == verbose_proxy_stdout_logger.name + assert records[0].levelno == logging.WARNING + assert records[0].exc_info is None + assert "LiteLLM Virtual Key expected" in records[0].getMessage() + assert exc_info.value.message == "Authentication service temporarily unavailable" assert exc_info.value.code == str(status.HTTP_503_SERVICE_UNAVAILABLE) assert not hasattr(exc_info.value, "_litellm_invalid_virtual_key_error") + + +@pytest.mark.asyncio +async def test_handle_authentication_error_logs_before_the_failure_hook_can_raise(caplog): + handler = UserAPIKeyAuthExceptionHandler() + hook_error = TypeError("function_setup() got multiple values for keyword argument 'start_time'") + + with _auth_failure_boundaries(side_effect=hook_error), pytest.raises(TypeError) as exc_info: + await handler._handle_authentication_error( + _invalid_virtual_key_error(), + MagicMock(), + {"start_time": "client-controlled"}, + "/v1/chat/completions", + None, + "undefined", + ) + + assert exc_info.value is hook_error + records = _auth_failure_records(caplog) + assert len(records) == 1 + assert records[0].name == verbose_proxy_stdout_logger.name + assert records[0].levelno == logging.WARNING + assert "LiteLLM Virtual Key expected" in records[0].getMessage() 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 d44f96d95bf..584234df84c 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 @@ -6872,3 +6872,46 @@ class TestLitellmReceivedAtStamping: assert result == earlier assert request.state.litellm_received_at == earlier + + +@pytest.mark.asyncio +async def test_user_api_key_auth_websocket_skips_duplicate_invalid_key_log(): + from fastapi import WebSocket + from starlette.datastructures import URL + + from litellm.proxy._types import ProxyException + from litellm.proxy.auth.auth_utils import mark_invalid_virtual_key_error + from litellm.proxy.auth.user_api_key_auth import WebSocketException, user_api_key_auth_websocket + + mock_websocket = MagicMock(spec=WebSocket) + mock_websocket.query_params = {"model": "some_model"} + mock_websocket.headers = {"authorization": "Bearer undefined"} + mock_websocket.scope = {"headers": [(b"authorization", b"Bearer undefined")]} + mock_websocket.url = URL(url="/v1/responses") + transformed_invalid_virtual_key = mark_invalid_virtual_key_error( + ProxyException( + message="Please authenticate again", + type="auth_error", + param="None", + code=status.HTTP_401_UNAUTHORIZED, + ), + True, + ) + + with ( + patch( # test-quality-ok: auth error test injects this runtime boundary + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", + side_effect=transformed_invalid_virtual_key, + autospec=True, + ), + patch( # test-quality-ok: auth error test injects this runtime boundary + "litellm.proxy.auth.user_api_key_auth.verbose_proxy_logger.exception" + ) as exception_log, + pytest.raises(WebSocketException) as exc_info, + ): + await user_api_key_auth_websocket(mock_websocket) + + assert exc_info.value.code == status.WS_1008_POLICY_VIOLATION + assert exc_info.value.reason == "" + exception_log.assert_not_called() + mock_websocket.close.assert_not_called() diff --git a/tests/test_litellm/test_logging.py b/tests/test_litellm/test_logging.py index 680dd2015c4..fe9fb5c8ef6 100644 --- a/tests/test_litellm/test_logging.py +++ b/tests/test_litellm/test_logging.py @@ -928,13 +928,24 @@ def test_records_below_warning_and_invalid_key_warnings_go_to_stdout(capsys): def test_verbose_loggers_route_records_by_level(): - for lg in (verbose_logger, verbose_router_logger, verbose_proxy_logger, verbose_proxy_stdout_logger): + for lg in (verbose_logger, verbose_router_logger, verbose_proxy_logger): assert any(isinstance(h, LevelRoutingStreamHandler) for h in lg.handlers), lg.name - assert verbose_proxy_stdout_logger.level == logging.WARNING - assert verbose_proxy_stdout_logger.handlers[0].level == logging.WARNING -def test_invalid_virtual_key_record_does_not_propagate_to_root_handler(): +def test_invalid_key_logger_inherits_proxy_logger_configuration(): + assert verbose_proxy_stdout_logger.parent is verbose_proxy_logger + assert verbose_proxy_stdout_logger.handlers == [] + assert verbose_proxy_stdout_logger.level == logging.NOTSET + assert verbose_proxy_stdout_logger.propagate is True + assert verbose_proxy_stdout_logger.getEffectiveLevel() == verbose_proxy_logger.getEffectiveLevel() + assert verbose_proxy_stdout_logger not in ALL_LOGGERS + + +def test_invalid_virtual_key_record_propagates_to_root_handler_like_ordinary_records(capsys): + proxy_handlers = list(verbose_proxy_logger.handlers) + original_levels = [h.level for h in proxy_handlers] + for h in proxy_handlers: + h.setLevel(logging.WARNING) root_logger = logging.getLogger() root_stream = StringIO() root_handler = logging.StreamHandler(root_stream) @@ -945,22 +956,34 @@ def test_invalid_virtual_key_record_does_not_propagate_to_root_handler(): verbose_proxy_stdout_logger.warning("invalid virtual key") finally: root_logger.removeHandler(root_handler) + for h, level in zip(proxy_handlers, original_levels, strict=True): + h.setLevel(level) - assert root_stream.getvalue() == "" + out, err = capsys.readouterr() + assert out.count("invalid virtual key") == 1 + assert err == "" + assert "ROOT WARNING invalid virtual key" in root_stream.getvalue() -def test_turn_on_json_preserves_invalid_key_warning_visibility(monkeypatch, capfd): - monkeypatch.setenv("LITELLM_LOG", "ERROR") - _turn_on_json() +def test_error_threshold_suppresses_invalid_key_warning_like_ordinary_warnings(capsys): + proxy_handlers = list(verbose_proxy_logger.handlers) + original_levels = [h.level for h in proxy_handlers] + for h in proxy_handlers: + h.setLevel(logging.ERROR) - verbose_proxy_stdout_logger.warning("invalid virtual key") + try: + verbose_proxy_stdout_logger.warning("invalid virtual key") + verbose_proxy_logger.warning("ordinary proxy warning") + verbose_proxy_logger.error("ordinary proxy error") + finally: + for h, level in zip(proxy_handlers, original_levels, strict=True): + h.setLevel(level) - out, err = capfd.readouterr() - assert [raw for raw in err.splitlines() if raw.strip()] == [] - records = [json.loads(raw) for raw in out.splitlines() if raw.strip()] - assert len(records) == 1 - assert records[0]["component"] == verbose_proxy_stdout_logger.name - assert records[0]["level"] == "WARNING" + out, err = capsys.readouterr() + assert out == "" + assert "invalid virtual key" not in err + assert "ordinary proxy warning" not in err + assert "ordinary proxy error" in err def test_ordinary_proxy_records_still_propagate_to_root_handler():