diff --git a/litellm/_logging.py b/litellm/_logging.py index fbb35b72be2..a748edb4636 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -264,13 +264,16 @@ def _plain_log_format(stdout: TextIO | None, stderr: TextIO | None) -> str: class LevelRoutingStreamHandler(logging.StreamHandler): - """Writes records below WARNING to stdout and WARNING and above to stderr. + """Writes records below WARNING and selected WARNING records to stdout. Collectors that derive severity from the stream report every stderr line as an error. """ def emit(self, record: logging.LogRecord) -> None: - preferred: Final = sys.stdout if record.levelno < logging.WARNING else sys.stderr + is_stdout_record: Final = record.levelno < logging.WARNING or ( + record.levelno == logging.WARNING and record.__dict__.get("route_to_stdout") is True + ) + preferred: Final = sys.stdout if is_stdout_record else sys.stderr if preferred is None or getattr(preferred, "closed", False): self.stream = sys.stderr # rebind-ok: fall back to the pre-fix stream rather than raising per record else: @@ -357,6 +360,8 @@ def _get_standard_record_attrs() -> frozenset: _STANDARD_RECORD_ATTRS: Final = _get_standard_record_attrs() +# Stream-routing controls the local handler only; it is not an event attribute. +_JSON_EXCLUDED_LOG_RECORD_ATTRS: Final = frozenset(("route_to_stdout",)) # CorrelationContextFilter is the only legitimate source for these two JSON fields; # see JsonFormatter.format() for why they're excluded from the generic message-content @@ -397,7 +402,11 @@ class JsonFormatter(Formatter): # Include extra attributes passed via logger.debug("msg", extra={...}) for key, value in record.__dict__.items(): - if key not in _STANDARD_RECORD_ATTRS and key not in json_record: + if ( + key not in _STANDARD_RECORD_ATTRS + and key not in _JSON_EXCLUDED_LOG_RECORD_ATTRS + and key not in json_record + ): json_record[key] = value # trace_id/session_id are reserved: CorrelationContextFilter is the only diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index c2aebfd270d..c445342c382 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -2,6 +2,7 @@ Handles Authentication Errors """ +import logging from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final @@ -115,27 +116,20 @@ class UserAPIKeyAuthExceptionHandler: and e.status_code == status.HTTP_401_UNAUTHORIZED and "LiteLLM Virtual Key expected" in str(e.detail) ) - if is_invalid_virtual_key and not litellm.log_client_error_tracebacks: - verbose_proxy_logger.warning( - "litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - %s\nRequester IP Address:%s", - e, - requester_ip, - extra={"requester_ip": requester_ip}, - ) - elif is_expected_client_error(e) and not litellm.log_client_error_tracebacks: - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - %s\nRequester IP Address:%s", - e, - requester_ip, - extra={"requester_ip": requester_ip}, - ) - else: - verbose_proxy_logger.exception( - "litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - %s\nRequester IP Address:%s", - e, - requester_ip, - extra={"requester_ip": requester_ip}, - ) + log_level: Final = ( + logging.WARNING if is_invalid_virtual_key and not litellm.log_client_error_tracebacks else logging.ERROR + ) + verbose_proxy_logger.log( + log_level, + "litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - %s\nRequester IP Address:%s", + e, + requester_ip, + exc_info=not (is_expected_client_error(e) and not litellm.log_client_error_tracebacks), + extra={ + "requester_ip": requester_ip, + "route_to_stdout": is_invalid_virtual_key and not litellm.log_client_error_tracebacks, + }, + ) # Log this exception to OTEL, Datadog etc. Reuse the identity resolved # before the failure (team alias/id, metadata, user) so the failed span diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index e92d090a2fb..0d9efec1ac0 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -19,6 +19,7 @@ import fastapi import orjson from fastapi import HTTPException, Request, WebSocket, status from fastapi.security.api_key import APIKeyHeader +from starlette.exceptions import WebSocketException import litellm from litellm._logging import verbose_logger, verbose_proxy_logger @@ -32,6 +33,7 @@ from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm.integrations.otel.runtime import phase_span, seed_request_identity from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value +from litellm.litellm_core_utils.realtime_errors import websocket_close_reason from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, @@ -539,6 +541,16 @@ async def user_api_key_auth_websocket(websocket: WebSocket): try: return await user_api_key_auth(request=request, api_key=f"Bearer {api_key}") except Exception as e: + is_invalid_virtual_key: Final = ( + isinstance(e, ProxyException) + and e.code == str(status.HTTP_401_UNAUTHORIZED) + and "LiteLLM Virtual Key expected" in e.message + ) + if is_invalid_virtual_key: + raise WebSocketException( + code=status.WS_1008_POLICY_VIOLATION, + reason=websocket_close_reason(str(e), fallback="Invalid API key"), + ) verbose_proxy_logger.exception(e) await websocket.close(code=status.WS_1008_POLICY_VIOLATION) raise HTTPException(status_code=403, detail=str(e)) 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 0cdf3500d50..f1a00da49ac 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -874,6 +874,107 @@ 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_does_not_relog_invalid_virtual_key(): + from litellm.proxy._types import ProxyException + 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") + invalid_virtual_key = ProxyException( + message="LiteLLM Virtual Key expected", + type="auth_error", + param="None", + code=status.HTTP_401_UNAUTHORIZED, + ) + + with ( + patch( # test-quality-ok: delegates auth; this test isolates the WebSocket adapter branch + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", + side_effect=invalid_virtual_key, + autospec=True, + ), + patch( # test-quality-ok: assert non-target errors retain the adapter's legacy logger call + "litellm.proxy.auth.user_api_key_auth.verbose_proxy_logger.exception" + ) as exception_log, + pytest.raises(WebSocketException, match="LiteLLM Virtual Key expected") as exc_info, + ): + await user_api_key_auth_websocket(mock_websocket) + + assert exc_info.value.code == status.WS_1008_POLICY_VIOLATION + exception_log.assert_not_called() + mock_websocket.close.assert_not_called() + + +@pytest.mark.asyncio +async def test_user_api_key_auth_websocket_caps_invalid_virtual_key_close_reason(): + from litellm.litellm_core_utils.realtime_errors import WEBSOCKET_CLOSE_REASON_MAX_BYTES + from litellm.proxy._types import ProxyException + 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") + invalid_virtual_key = ProxyException( + message="LiteLLM Virtual Key expected" + "x" * WEBSOCKET_CLOSE_REASON_MAX_BYTES, + type="auth_error", + param="None", + code=status.HTTP_401_UNAUTHORIZED, + ) + + with ( + patch( # test-quality-ok: delegates auth; this test isolates the WebSocket adapter branch + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", + side_effect=invalid_virtual_key, + autospec=True, + ), + pytest.raises(WebSocketException) as exc_info, + ): + await user_api_key_auth_websocket(mock_websocket) + + assert len(exc_info.value.reason.encode()) <= WEBSOCKET_CLOSE_REASON_MAX_BYTES + + +@pytest.mark.asyncio +async def test_user_api_key_auth_websocket_preserves_non_target_error_handling(): + from litellm.proxy._types import ProxyException + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_websocket + + mock_websocket = MagicMock(spec=WebSocket) + mock_websocket.query_params = {"model": "some_model"} + mock_websocket.headers = {"authorization": "Bearer sk-unknown"} + mock_websocket.scope = {"headers": [(b"authorization", b"Bearer sk-unknown")]} + mock_websocket.url = URL(url="/v1/responses") + unknown_key = ProxyException( + message="Authentication Error, key not found in db", + type="auth_error", + param="None", + code=status.HTTP_401_UNAUTHORIZED, + ) + + with ( + patch( # test-quality-ok: delegates auth; this test isolates the WebSocket adapter branch + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", + side_effect=unknown_key, + autospec=True, + ), + patch( # test-quality-ok: assert non-target errors retain the adapter's legacy logger call + "litellm.proxy.auth.user_api_key_auth.verbose_proxy_logger.exception" + ) as exception_log, + pytest.raises(HTTPException, match="key not found in db") as exc_info, + ): + await user_api_key_auth_websocket(mock_websocket) + + assert exc_info.value.status_code == status.HTTP_403_FORBIDDEN + exception_log.assert_called_once_with(unknown_key) + mock_websocket.close.assert_awaited_once_with(code=status.WS_1008_POLICY_VIOLATION) + + @pytest.mark.asyncio async def test_user_api_key_auth_websocket_carries_asgi_path(): """ @@ -894,7 +995,9 @@ async def test_user_api_key_auth_websocket_carries_asgi_path(): } mock_websocket.url = URL(url="/v1/realtime") - with patch("litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True) as mock_user_api_key_auth: + with patch( # test-quality-ok: delegates auth; this test inspects the synthetic ASGI request + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True + ) as mock_user_api_key_auth: await user_api_key_auth_websocket(mock_websocket) request_arg = mock_user_api_key_auth.call_args.kwargs["request"] 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 dbc587a6d3f..0f220b61a3e 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -786,7 +786,7 @@ async def test_handle_authentication_error_traceback_only_for_unexpected_errors( patch( # test-quality-ok: handler reads proxy_server globals at call time "litellm.proxy.auth.auth_exception_handler.seed_request_identity", ), - patch( + 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, ), @@ -815,4 +815,7 @@ async def test_handle_authentication_error_traceback_only_for_unexpected_errors( records = [r for r in caplog.records if "user_api_key_auth(): Exception occured" in r.getMessage()] assert len(records) == 1 assert records[0].levelno == expected_level - assert (records[0].exc_info is not None) is expect_traceback + assert bool(records[0].exc_info) is expect_traceback + assert getattr(records[0], "route_to_stdout", False) is ( + expected_level == logging.WARNING and not log_client_error_tracebacks + ) diff --git a/tests/test_litellm/test_logging.py b/tests/test_litellm/test_logging.py index 087a1c8b3ad..b69a27426e4 100644 --- a/tests/test_litellm/test_logging.py +++ b/tests/test_litellm/test_logging.py @@ -128,6 +128,25 @@ def test_json_formatter_includes_extra_attributes(): assert obj["authorization"] == "Bearer sk-***" +def test_json_formatter_excludes_stream_routing_controls(): + """Stream routing controls must not become JSON event attributes.""" + formatter = JsonFormatter() + record = logging.LogRecord( + name="LiteLLM Proxy", + level=logging.WARNING, + pathname="", + lineno=0, + msg="invalid virtual key", + args=(), + exc_info=None, + ) + record.route_to_stdout = True + + obj = json.loads(formatter.format(record)) + + assert "route_to_stdout" not in obj + + def test_json_formatter_plain_message_unchanged(): """ Test that non-JSON messages are passed through as-is in the message field. @@ -845,7 +864,29 @@ class _FakeStream: return self._tty -def test_records_below_warning_go_to_stdout_and_the_rest_to_stderr(capsys): +def test_selected_warning_routes_as_json_to_stdout(capsys): + logger = logging.getLogger("test_json_level_routing") + logger.handlers.clear() + logger.propagate = False + logger.setLevel(logging.DEBUG) + handler = LevelRoutingStreamHandler() + handler.setFormatter(JsonFormatter()) + logger.addHandler(handler) + + try: + logger.warning("invalid virtual key", extra={"route_to_stdout": True}) + finally: + logger.handlers.clear() + + out, err = capsys.readouterr() + assert err == "" + log_record = json.loads(out) + assert log_record["message"] == "invalid virtual key" + assert log_record["level"] == "WARNING" + assert "route_to_stdout" not in log_record + + +def test_records_below_warning_and_selected_warnings_go_to_stdout(capsys): logger = logging.getLogger("test_level_routing") logger.handlers.clear() logger.propagate = False @@ -858,14 +899,16 @@ def test_records_below_warning_go_to_stdout_and_the_rest_to_stderr(capsys): logger.debug("d") logger.info("i") logger.warning("w") + logger.warning("marked-warning", extra={"route_to_stdout": True}) logger.error("e") + logger.error("marked-error", extra={"route_to_stdout": True}) logger.critical("c") finally: logger.handlers.clear() out, err = capsys.readouterr() - assert out.splitlines() == ["DEBUG d", "INFO i"] - assert err.splitlines() == ["WARNING w", "ERROR e", "CRITICAL c"] + assert out.splitlines() == ["DEBUG d", "INFO i", "WARNING marked-warning"] + assert err.splitlines() == ["WARNING w", "ERROR e", "ERROR marked-error", "CRITICAL c"] def test_verbose_loggers_route_records_by_level():