mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
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.
This commit is contained in:
Yucheng Zhu
parent
1a1d6ce804
commit
6646bb4319
7 changed files with 243 additions and 271 deletions
|
|
@ -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())
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue