From ef92f5a3bad86edbca860125069a28392c76df52 Mon Sep 17 00:00:00 2001 From: joshua Date: Sat, 10 Oct 2026 20:20:13 +0000 Subject: [PATCH 1/8] fix(mcp): log expected OAuth cold-start 401 at debug Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/auth/user_api_key_auth_mcp.py | 130 ++++----- litellm/proxy/auth/auth_exception_handler.py | 30 ++- litellm/proxy/auth/auth_utils.py | 17 ++ litellm/proxy/auth/user_api_key_auth.py | 3 +- tests/integration/mcp/test_mcp_oauth_flows.py | 40 ++- .../auth/test_user_api_key_auth_mcp.py | 247 ++++++++++++++++++ .../proxy/auth/test_auth_exception_handler.py | 162 ++++++++++++ 7 files changed, 562 insertions(+), 67 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index bb2fe2d59b4..2ae4be2c0af 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -87,6 +87,7 @@ if TYPE_CHECKING: _EMPTY_TOOLSET_GRANTS: Final[Mapping[str, Sequence[str]]] = MappingProxyType({}) +_GENERAL_SETTINGS_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) OPTIONAL_STRING_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) @@ -285,30 +286,31 @@ def _gateway_dcr_challenge_target( return targets[0] -def _is_gateway_dcr_challenge_scope( - route: str, +def _cold_start_fallback( + request: Request, + request_route: str, mcp_servers: list[str] | None, mcp_auth_header: str | None, mcp_server_auth_headers: dict[str, dict[str, str]] | None, - exc: Exception, - client_ip: str | None, -) -> bool: - """True when an unauthenticated MCP request should receive the RFC 9728 401 - challenge that advertises the gateway as the authorization server. - - Fires only for a genuine 401 with no client-supplied MCP auth headers (those mean - the caller is not a cold-start DCR client), on the scopes the gateway's keyless - flow serves: the aggregate ``/mcp`` endpoint, an ``x-mcp-servers``-scoped request - (the resource the client configured is still ``/mcp``), or a per-server path whose - single target advertises gateway-owned sign-in. Every other named target keeps - its existing behavior, failing closed to the original admission error.""" - if not _is_litellm_auth_admission_error(exc): - return False + bearer_presented: bool, +) -> Literal["defer_to_route_challenge", "gateway_challenge"] | None: + """Classify the structural cold-start response without inspecting an error.""" if _has_client_supplied_mcp_auth(mcp_auth_header, mcp_server_auth_headers): - return False - if len(MCPRequestHandler.extract_target_server_names_from_path(route)) == 0: - return True - return _gateway_dcr_challenge_target(route, mcp_servers, client_ip) is not None + return None + + mcp_servers_from_path: Final = _parse_mcp_server_names_from_path(request_route, mcp_servers) + client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) + if mcp_servers_from_path is not None and ( + _is_mcp_passthrough_cold_start(mcp_servers_from_path, client_ip=client_ip) + or (not bearer_presented and _is_legacy_delegate_cold_start(mcp_servers_from_path, client_ip=client_ip)) + ): + return "defer_to_route_challenge" + if ( + len(MCPRequestHandler.extract_target_server_names_from_path(request_route)) == 0 + or _gateway_dcr_challenge_target(request_route, mcp_servers, client_ip) is not None + ): + return "gateway_challenge" + return None def _gateway_dcr_challenge( @@ -353,48 +355,41 @@ def _admission_failure_fallback( exc: Exception, bearer_presented: bool, ) -> UserAPIKeyAuth: - """Map a failed LiteLLM admission to its anonymous fallback or challenge. + """Apply the structural cold-start response to an admission error.""" + if not _is_litellm_auth_admission_error(exc): + raise exc - Two fallbacks exist, both gated on a genuine 401 with no client-supplied - MCP auth headers. The pass-through cold start (RFC 9728 / MCP - Authorization spec discovery return) admits anonymously so the route's - 401 emitter can produce the per-server challenge. The aggregate - gateway-DCR scope converts the failure into the gateway's own - resource_metadata challenge, with the RFC 6750 ``invalid_token`` error - code when the caller DID present a bearer (an expired gateway session - must re-authorize, not retry a dead token). Anything else re-raises the - original admission error unchanged.""" - mcp_servers_from_path: Final = _parse_mcp_server_names_from_path(request_route, mcp_servers) - if ( - mcp_servers_from_path is not None - and not _has_client_supplied_mcp_auth(mcp_auth_header, mcp_server_auth_headers) - and _is_litellm_auth_admission_error(exc) - and ( - _is_mcp_passthrough_cold_start( - mcp_servers_from_path, - client_ip=IPAddressUtils.get_mcp_client_ip(request), - ) - or ( - not bearer_presented - and _is_legacy_delegate_cold_start( - mcp_servers_from_path, - client_ip=IPAddressUtils.get_mcp_client_ip(request), - ) - ) - ) + match _cold_start_fallback( + request, + request_route, + mcp_servers, + mcp_auth_header, + mcp_server_auth_headers, + bearer_presented, ): - verbose_logger.debug("MCP pass-through cold start: deferring admission to route 401 emitter") - return UserAPIKeyAuth() - if _is_gateway_dcr_challenge_scope( - route=request_route, - mcp_servers=mcp_servers, - mcp_auth_header=mcp_auth_header, - mcp_server_auth_headers=mcp_server_auth_headers, - exc=exc, - client_ip=IPAddressUtils.get_mcp_client_ip(request), - ): - raise _gateway_dcr_challenge(request, request_route, mcp_servers, invalid_token=bearer_presented) from exc - raise exc + case "defer_to_route_challenge": + verbose_logger.debug("MCP pass-through cold start: deferring admission to route 401 emitter") + return UserAPIKeyAuth() + case "gateway_challenge": + raise _gateway_dcr_challenge( + request, + request_route, + mcp_servers, + invalid_token=bearer_presented, + ) from exc + case None: + raise exc + + +def _has_litellm_credential_header(headers: Headers) -> bool: + if "x-litellm-api-key" in headers or "Authorization" in headers: + return True + + from litellm.proxy.proxy_server import general_settings # pyright: ignore[reportUnknownVariableType] # legacy map + + settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(general_settings) + custom_header_name: Final = OPTIONAL_STRING_ADAPTER.validate_python(settings.get("litellm_key_header_name")) + return custom_header_name is not None and custom_header_name in headers @dataclass(frozen=True, slots=True) @@ -629,6 +624,21 @@ class MCPRequestHandler: bearer_presented=True, ) else: + from litellm.proxy.auth.auth_utils import mark_expected_credential_challenge + + if not _has_litellm_credential_header(headers) and ( + _cold_start_fallback( + request=request, + request_route=request_route, + mcp_servers=mcp_servers, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + bearer_presented=False, + ) + is not None + ): + mark_expected_credential_challenge(request) + try: validated_user_api_key_auth = await user_api_key_auth(api_key=litellm_api_key, request=request) except (HTTPException, ProxyException) as exc: diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 0266131fc33..ae64688676a 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -23,6 +23,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_utils import ( # noqa: F401 # legacy module exports _get_request_ip_address, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export get_request_ip_address, + is_expected_credential_challenge, is_invalid_virtual_key_error, mark_invalid_virtual_key_error, normalize_request_route, @@ -193,14 +194,31 @@ class UserAPIKeyAuthExceptionHandler: log_extra: Final = {"requester_ip": requester_ip} 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 + is_expected_challenge: Final = is_expected_credential_challenge(e, request) + log_level: Final = ( + logging.DEBUG if is_expected_challenge else logging.WARNING if is_quiet_log else logging.ERROR + ) + logger: Final = ( + verbose_proxy_logger + if is_expected_challenge + else verbose_proxy_stdout_logger + if is_quiet_log + else verbose_proxy_logger + ) + exc_info: Final = ( + None + if is_expected_challenge + else True + if litellm.log_client_error_tracebacks or not is_expected_client_error(e) + else None + ) logger.log( - logging.WARNING if is_quiet_log else logging.ERROR, + log_level, "litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - %s\nRequester IP Address:%s%s", e, requester_ip, _identity_log_suffix(resolved_identity), - exc_info=True if litellm.log_client_error_tracebacks or not is_expected_client_error(e) else None, + exc_info=exc_info, extra=log_extra, ) @@ -250,8 +268,10 @@ class UserAPIKeyAuthExceptionHandler: e = transformed_exception final_exception: Final = mark_invalid_virtual_key_error(_as_proxy_exception(e), is_invalid_virtual_key) - # If a quiet-logged malformed-key transform yields non-401, escalate to ERROR - if is_quiet_log and str(final_exception.code) != str(status.HTTP_401_UNAUTHORIZED): + # Escalate callback-transformed quiet logs and expected challenges with non-401 codes + if (is_quiet_log or is_expected_challenge) and str(final_exception.code) != str( + status.HTTP_401_UNAUTHORIZED + ): verbose_proxy_logger.error( "litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - %s\nRequester IP Address:%s", final_exception, diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index e4fe88690f5..dd9cc65285f 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -41,6 +41,23 @@ from litellm.types.router import reject_server_owned_wif_params as _reject_serve from litellm.types.utils import CustomPricingLiteLLMParams, oauth_token_exchange_litellm_params +class MissingCredentialError(Exception): + pass + + +EXPECTED_CREDENTIAL_CHALLENGE_STATE_KEY: Final = "litellm_expected_credential_challenge" + + +def mark_expected_credential_challenge(request: Request) -> None: + setattr(request.state, EXPECTED_CREDENTIAL_CHALLENGE_STATE_KEY, True) + + +def is_expected_credential_challenge(exception: BaseException, request: Request) -> bool: + return isinstance(exception, MissingCredentialError) and ( + getattr(request.state, EXPECTED_CREDENTIAL_CHALLENGE_STATE_KEY, False) is True + ) + + def is_invalid_virtual_key_error(exception: BaseException | None) -> bool: """True when an authentication error rejects a malformed virtual key. diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index d6a51bc2fc9..62b7c1daf55 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -93,6 +93,7 @@ from litellm.proxy.auth.auth_object_prefetch import ( prefetch_identity_keys, ) from litellm.proxy.auth.auth_utils import ( + MissingCredentialError, abbreviate_api_key, fallback_target_model_name, get_end_user_id_from_request_body, @@ -1966,7 +1967,7 @@ async def user_api_key_auth_builder( raise Exception("No api key passed in.") elif api_key == "": # missing 'Bearer ' prefix - raise Exception("Malformed API Key passed in. Ensure Key has `Bearer ` prefix.") + raise MissingCredentialError("Malformed API Key passed in. Ensure Key has `Bearer ` prefix.") if route == "/user/auth": if general_settings.get("allow_user_auth", False) is True: diff --git a/tests/integration/mcp/test_mcp_oauth_flows.py b/tests/integration/mcp/test_mcp_oauth_flows.py index e95cdb902e8..d09af78a8a5 100644 --- a/tests/integration/mcp/test_mcp_oauth_flows.py +++ b/tests/integration/mcp/test_mcp_oauth_flows.py @@ -31,7 +31,7 @@ from integration._support.mcp import ( ) from integration._support.mcp_grants import create_toolset from integration._support.oauth_server import AuthorizationServer, oauth_server -from integration._support.process import owned_proxy +from integration._support.process import owned_proxy, owned_proxy_process ADD: Final = {"a": 2, "b": 3} CLIENT_REDIRECT: Final = "http://127.0.0.1:9/cb" @@ -357,6 +357,44 @@ def test_per_user_authorization_code_with_pkce_binds_the_token_to_the_authorizin ) +def test_anonymous_oauth_mcp_cold_start_does_not_log_malformed_key_at_error( + gateway: Gateway, + tmp_path: Path, +) -> None: + with mcp_peer() as peer, oauth_server() as auth, owned_proxy_process(gateway, tmp_path, {}) as owned: + with owned.gateway.scenario() as scenario: + alias: Final = "cold" + uuid.uuid4().hex[:8] + _register_oauth( + scenario, + peer, + auth, + alias, + auth_type="oauth2", + oauth2_flow="authorization_code", + credentials={"client_id": "cold-client", "client_secret": "cold-secret"}, + ) + baseline_log: Final = owned.log.read_text() + response: Final = owned.gateway.client.post( + f"/{alias}/mcp", + headers=ACCEPT, + json=INITIALIZE, + ) + + assert response.status_code == 401, response.text + assert "www-authenticate" in response.headers, response.headers + assert "resource_metadata=" in response.headers["www-authenticate"] + request_log: Final = eventually( + owned.log.read_text, + lambda value: f"POST /{alias}/mcp" in value[len(baseline_log) :], + ) + error_lines: Final = tuple( + line + for line in request_log[len(baseline_log) :].splitlines() + if "Malformed API Key" in line and "ERROR" in line + ) + assert error_lines == () + + def test_authorization_request_without_pkce_is_refused_before_reaching_the_authorization_server( gateway: Gateway, ) -> None: diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 63756ee108c..4ad43d30136 100644 --- a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -9,6 +9,7 @@ import pytest from fastapi import HTTPException from fastapi.testclient import TestClient from starlette.datastructures import Headers +from starlette.requests import Request from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, @@ -23,9 +24,255 @@ from litellm.proxy._types import ( SpecialMCPServerNames, UserAPIKeyAuth, ) +from litellm.types.mcp import MCPAuth +from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.agents import AgentCaller +@pytest.mark.asyncio +@pytest.mark.parametrize( + "path,extra_headers,auth_type,custom_key_header_name,hide_target,expected", + [ + pytest.param("/oauth/mcp", (), MCPAuth.oauth2, None, False, True, id="gateway-oauth-server"), + pytest.param("/mcp", (), None, None, False, True, id="aggregate-mcp"), + pytest.param( + "/oauth/mcp", + ((b"authorization", b""),), + MCPAuth.oauth2, + None, + False, + False, + id="empty-authorization", + ), + pytest.param( + "/oauth/mcp", + ((b"authorization", b"Basic malformed"),), + MCPAuth.oauth2, + None, + False, + False, + id="malformed-authorization", + ), + pytest.param( + "/oauth/mcp", + ((b"x-litellm-api-key", b""),), + MCPAuth.oauth2, + None, + False, + False, + id="empty-litellm-key", + ), + pytest.param( + "/oauth/mcp", + ((b"x-custom-litellm-key", b""),), + MCPAuth.oauth2, + "x-custom-litellm-key", + False, + False, + id="custom-key-header", + ), + pytest.param("/keyed/mcp", (), MCPAuth.api_key, None, False, False, id="non-oauth-target"), + pytest.param("/missing/mcp", (), None, None, False, False, id="unresolved-target"), + pytest.param("/hidden/mcp", (), MCPAuth.oauth2, None, True, False, id="target-hidden-by-client-ip"), + ], +) +async def test_process_mcp_request_marks_only_expected_credential_challenges( + path: str, + extra_headers: tuple[tuple[bytes, bytes], ...], + auth_type: MCPAuth | None, + custom_key_header_name: str | None, + hide_target: bool, + expected: bool, +) -> None: + from litellm.proxy.auth.auth_utils import EXPECTED_CREDENTIAL_CHALLENGE_STATE_KEY + + server: Final[MCPServer | None] = ( + MCPServer( + server_id="server-id", + name="oauth" if auth_type == MCPAuth.oauth2 else "keyed", + server_name="oauth" if auth_type == MCPAuth.oauth2 else "keyed", + url="https://upstream.example/mcp", + transport="http", + auth_type=auth_type, + extra_headers=["x-api-key"] if auth_type == MCPAuth.api_key else None, + ) + if auth_type is not None + else None + ) + scope: Final = { + "type": "http", + "method": "POST", + "path": path, + "headers": [(b"host", b"testserver"), *extra_headers], + } + with ( + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=HTTPException(status_code=401, detail="Authentication required"), + ) as mock_auth, + patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_manager, + patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.IPAddressUtils.get_mcp_client_ip", + return_value="203.0.113.10" if hide_target else None, + ), + patch( + "litellm.proxy.proxy_server.general_settings", + {"litellm_key_header_name": custom_key_header_name}, + ), + ): + mock_manager.get_mcp_server_by_name.return_value = None if hide_target else server + with pytest.raises(HTTPException): + await MCPRequestHandler.process_mcp_request(scope) + + auth_request: Final = mock_auth.call_args.kwargs["request"] + assert getattr(auth_request.state, EXPECTED_CREDENTIAL_CHALLENGE_STATE_KEY, False) is expected + + +@pytest.mark.parametrize( + "path,auth_type,delegate_auth_to_upstream,pass_through,bearer_presented,mcp_auth_header," + "mcp_server_auth_headers,mcp_servers,expected", + [ + pytest.param( + "/mcp/legacy", + MCPAuth.oauth2, + True, + False, + False, + None, + None, + None, + "defer_to_route_challenge", + id="legacy-delegate-without-bearer", + ), + pytest.param( + "/mcp/legacy", + MCPAuth.oauth2, + True, + False, + True, + None, + None, + None, + None, + id="legacy-delegate-with-bearer", + ), + pytest.param( + "/mcp/passthrough", + None, + False, + True, + True, + None, + None, + None, + "defer_to_route_challenge", + id="pass-through-with-bearer", + ), + pytest.param( + "/mcp/oauth", + MCPAuth.oauth2, + False, + False, + False, + None, + None, + None, + "gateway_challenge", + id="gateway-dcr-per-server", + ), + pytest.param( + "/mcp", + None, + False, + False, + False, + None, + None, + None, + "gateway_challenge", + id="aggregate-mcp", + ), + pytest.param( + "/mcp", + None, + False, + False, + False, + None, + None, + ["oauth"], + "gateway_challenge", + id="aggregate-mcp-scoped-by-header", + ), + pytest.param( + "/mcp/oauth", + MCPAuth.oauth2, + False, + False, + False, + "Bearer upstream", + None, + None, + None, + id="legacy-mcp-auth-header", + ), + pytest.param( + "/mcp/oauth", + MCPAuth.oauth2, + False, + False, + False, + None, + {"oauth": {"Authorization": "Bearer upstream"}}, + None, + None, + id="server-specific-mcp-auth-header", + ), + ], +) +def test_cold_start_fallback_classifies_structural_challenges( + path: str, + auth_type: MCPAuth | None, + delegate_auth_to_upstream: bool, + pass_through: bool, + bearer_presented: bool, + mcp_auth_header: str | None, + mcp_server_auth_headers: dict[str, dict[str, str]] | None, + mcp_servers: list[str] | None, + expected: Literal["defer_to_route_challenge", "gateway_challenge"] | None, +) -> None: + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import _cold_start_fallback + + server: Final[MCPServer | MagicMock | None] = ( + MagicMock(is_oauth_passthrough=True) + if pass_through + else MCPServer( + server_id="server-id", + name=path.rsplit("/", 1)[-1], + server_name=path.rsplit("/", 1)[-1], + url="https://upstream.example/mcp", + transport="http", + auth_type=auth_type, + delegate_auth_to_upstream=delegate_auth_to_upstream, + ) + if auth_type is not None + else None + ) + request: Final = Request({"type": "http", "method": "POST", "path": path, "headers": []}) + with patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as manager: + manager.get_mcp_server_by_name.return_value = server + result: Final = _cold_start_fallback( + request=request, + request_route=path, + mcp_servers=mcp_servers, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + bearer_presented=bearer_presented, + ) + + assert result == expected + + @pytest.mark.parametrize("headers,expected", [ ({"Authorization": "Bearer caller.jwt.signature"}, "caller.jwt.signature"), ({"AUTHORIZATION": "bEaReR caller.jwt.signature"}, "caller.jwt.signature"), diff --git a/tests/unit/proxy/auth/test_auth_exception_handler.py b/tests/unit/proxy/auth/test_auth_exception_handler.py index dd6e1f57fbc..13a88ecc5fa 100644 --- a/tests/unit/proxy/auth/test_auth_exception_handler.py +++ b/tests/unit/proxy/auth/test_auth_exception_handler.py @@ -1,5 +1,7 @@ import asyncio import json +import logging +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -1147,6 +1149,166 @@ async def test_handle_authentication_error_keeps_internal_message_on_model_acces assert [r for r in caplog.records if r.levelname == "WARNING" and "internal-models" in r.getMessage()] == [] +async def _capture_authentication_exception( + exception: BaseException, + request: Request, + log_client_error_tracebacks: bool, + callback_exception: HTTPException | None = None, +) -> ProxyException: + handler: Final = UserAPIKeyAuthExceptionHandler() + with ( + patch.object(litellm, "log_client_error_tracebacks", log_client_error_tracebacks), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + return_value=callback_exception, + ), + patch("litellm.proxy.auth.auth_exception_handler.seed_request_identity"), + patch("litellm.proxy.proxy_server.general_settings", {"allow_requests_on_db_unavailable": False}), + ): + try: + await handler.handle_authentication_error(exception, request, {}, "/mcp", None, "") + except ProxyException as raised: + return raised + pytest.fail("expected the authentication handler to raise a ProxyException") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tracebacks_enabled", [True, False]) +async def test_expected_missing_credential_logs_debug_without_traceback_and_preserves_response( + tracebacks_enabled: bool, + caplog: pytest.LogCaptureFixture, +) -> None: + from litellm.proxy.auth.auth_utils import MissingCredentialError, mark_expected_credential_challenge + + flagged_request: Final = Request({"type": "http", "method": "POST", "path": "/mcp", "headers": []}) + mark_expected_credential_challenge(flagged_request) + unflagged_request: Final = Request({"type": "http", "method": "POST", "path": "/mcp", "headers": []}) + message: Final = "Malformed API Key passed in. Ensure Key has `Bearer ` prefix." + previous_propagate: Final = verbose_proxy_logger.propagate + verbose_proxy_logger.propagate = True + try: + with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): + flagged_response: Final = await _capture_authentication_exception( + MissingCredentialError(message), + flagged_request, + tracebacks_enabled, + ) + flagged_records: Final = tuple( + record for record in caplog.records if "user_api_key_auth(): Exception occured" in record.getMessage() + ) + assert tuple(record.levelname for record in flagged_records) == ("DEBUG",) + assert flagged_records[0].exc_info is None + + caplog.clear() + with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): + unflagged_response: Final = await _capture_authentication_exception( + MissingCredentialError(message), + unflagged_request, + tracebacks_enabled, + ) + unflagged_records: Final = tuple( + record for record in caplog.records if "user_api_key_auth(): Exception occured" in record.getMessage() + ) + assert tuple(record.levelname for record in unflagged_records) == ("ERROR",) + assert unflagged_records[0].exc_info is not None + assert vars(flagged_response) == vars(unflagged_response) + assert flagged_response.to_dict() == unflagged_response.to_dict() + finally: + verbose_proxy_logger.propagate = previous_propagate + + +@pytest.mark.asyncio +async def test_unflagged_missing_credential_keeps_error_traceback(caplog: pytest.LogCaptureFixture) -> None: + from litellm.proxy.auth.auth_utils import MissingCredentialError + + previous_propagate: Final = verbose_proxy_logger.propagate + verbose_proxy_logger.propagate = True + try: + with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): + await _capture_authentication_exception( + MissingCredentialError("Malformed API Key passed in. Ensure Key has `Bearer ` prefix."), + Request({"type": "http", "method": "POST", "path": "/mcp", "headers": []}), + False, + ) + finally: + verbose_proxy_logger.propagate = previous_propagate + + records: Final = tuple( + record for record in caplog.records if "user_api_key_auth(): Exception occured" in record.getMessage() + ) + assert tuple(record.levelname for record in records) == ("ERROR",) + assert records[0].exc_info is not None + + +@pytest.mark.asyncio +async def test_expected_missing_credential_callback_403_is_escalated_and_unchanged( + caplog: pytest.LogCaptureFixture, +) -> None: + from litellm.proxy.auth.auth_utils import MissingCredentialError, mark_expected_credential_challenge + + flagged_request: Final = Request({"type": "http", "method": "POST", "path": "/mcp", "headers": []}) + mark_expected_credential_challenge(flagged_request) + unflagged_request: Final = Request({"type": "http", "method": "POST", "path": "/mcp", "headers": []}) + callback_exception: Final = HTTPException(status_code=403, detail="Forbidden") + message: Final = "Malformed API Key passed in. Ensure Key has `Bearer ` prefix." + previous_propagate: Final = verbose_proxy_logger.propagate + verbose_proxy_logger.propagate = True + try: + with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): + flagged_response: Final = await _capture_authentication_exception( + MissingCredentialError(message), + flagged_request, + False, + callback_exception, + ) + flagged_records: Final = tuple( + record for record in caplog.records if "user_api_key_auth(): Exception occured" in record.getMessage() + ) + assert tuple(record.levelname for record in flagged_records) == ("DEBUG", "ERROR") + assert flagged_response.code == str(status.HTTP_403_FORBIDDEN) + + caplog.clear() + with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): + unflagged_response: Final = await _capture_authentication_exception( + MissingCredentialError(message), + unflagged_request, + False, + callback_exception, + ) + assert vars(flagged_response) == vars(unflagged_response) + assert flagged_response.to_dict() == unflagged_response.to_dict() + finally: + verbose_proxy_logger.propagate = previous_propagate + + +@pytest.mark.asyncio +async def test_expected_credential_marker_does_not_quiet_policy_rejection( + caplog: pytest.LogCaptureFixture, +) -> None: + from litellm.proxy.auth.auth_utils import mark_expected_credential_challenge + + request: Final = Request({"type": "http", "method": "POST", "path": "/mcp", "headers": []}) + mark_expected_credential_challenge(request) + previous_propagate: Final = verbose_proxy_logger.propagate + verbose_proxy_logger.propagate = True + try: + with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): + response: Final = await _capture_authentication_exception( + HTTPException(status_code=403, detail="Policy rejected request"), + request, + False, + ) + finally: + verbose_proxy_logger.propagate = previous_propagate + + records: Final = tuple( + record for record in caplog.records if "user_api_key_auth(): Exception occured" in record.getMessage() + ) + assert tuple(record.levelname for record in records) == ("ERROR",) + assert response.code == str(status.HTTP_403_FORBIDDEN) + + def test_as_proxy_exception_keeps_jwt_scope_denial_message_shape(): detail = {"error": _DENIED_CLIENT_MESSAGE} denial = ModelAccessDeniedHTTPException( From 6aba6f110cf2e5ea286a434d6bd90ceb1eb695a4 Mon Sep 17 00:00:00 2001 From: joshua Date: Sat, 10 Oct 2026 20:32:37 +0000 Subject: [PATCH 2/8] refactor(mcp): tidy cold-start credential check and tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/auth/user_api_key_auth_mcp.py | 11 +- litellm/proxy/auth/auth_exception_handler.py | 6 +- .../proxy/auth/test_auth_exception_handler.py | 144 +++++++----------- 3 files changed, 64 insertions(+), 97 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 2ae4be2c0af..28a07bacf6a 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -87,7 +87,6 @@ if TYPE_CHECKING: _EMPTY_TOOLSET_GRANTS: Final[Mapping[str, Sequence[str]]] = MappingProxyType({}) -_GENERAL_SETTINGS_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) OPTIONAL_STRING_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) @@ -382,13 +381,17 @@ def _admission_failure_fallback( def _has_litellm_credential_header(headers: Headers) -> bool: - if "x-litellm-api-key" in headers or "Authorization" in headers: + if ( + MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY in headers + or MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_SECONDARY in headers + ): return True from litellm.proxy.proxy_server import general_settings # pyright: ignore[reportUnknownVariableType] # legacy map - settings: Final = _GENERAL_SETTINGS_ADAPTER.validate_python(general_settings) - custom_header_name: Final = OPTIONAL_STRING_ADAPTER.validate_python(settings.get("litellm_key_header_name")) + custom_header_name: Final = OPTIONAL_STRING_ADAPTER.validate_python( + general_settings.get("litellm_key_header_name") # pyright: ignore[reportUnknownMemberType] # legacy settings + ) return custom_header_name is not None and custom_header_name in headers diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index ae64688676a..5b7cb10f583 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -199,11 +199,7 @@ class UserAPIKeyAuthExceptionHandler: logging.DEBUG if is_expected_challenge else logging.WARNING if is_quiet_log else logging.ERROR ) logger: Final = ( - verbose_proxy_logger - if is_expected_challenge - else verbose_proxy_stdout_logger - if is_quiet_log - else verbose_proxy_logger + verbose_proxy_stdout_logger if is_quiet_log and not is_expected_challenge else verbose_proxy_logger ) exc_info: Final = ( None diff --git a/tests/unit/proxy/auth/test_auth_exception_handler.py b/tests/unit/proxy/auth/test_auth_exception_handler.py index 13a88ecc5fa..9c2a982a3ac 100644 --- a/tests/unit/proxy/auth/test_auth_exception_handler.py +++ b/tests/unit/proxy/auth/test_auth_exception_handler.py @@ -1178,6 +1178,7 @@ async def _capture_authentication_exception( async def test_expected_missing_credential_logs_debug_without_traceback_and_preserves_response( tracebacks_enabled: bool, caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, ) -> None: from litellm.proxy.auth.auth_utils import MissingCredentialError, mark_expected_credential_challenge @@ -1185,65 +1186,39 @@ async def test_expected_missing_credential_logs_debug_without_traceback_and_pres mark_expected_credential_challenge(flagged_request) unflagged_request: Final = Request({"type": "http", "method": "POST", "path": "/mcp", "headers": []}) message: Final = "Malformed API Key passed in. Ensure Key has `Bearer ` prefix." - previous_propagate: Final = verbose_proxy_logger.propagate - verbose_proxy_logger.propagate = True - try: - with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): - flagged_response: Final = await _capture_authentication_exception( - MissingCredentialError(message), - flagged_request, - tracebacks_enabled, - ) - flagged_records: Final = tuple( - record for record in caplog.records if "user_api_key_auth(): Exception occured" in record.getMessage() + monkeypatch.setattr(verbose_proxy_logger, "propagate", True) + with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): + flagged_response: Final = await _capture_authentication_exception( + MissingCredentialError(message), + flagged_request, + tracebacks_enabled, ) - assert tuple(record.levelname for record in flagged_records) == ("DEBUG",) - assert flagged_records[0].exc_info is None - - caplog.clear() - with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): - unflagged_response: Final = await _capture_authentication_exception( - MissingCredentialError(message), - unflagged_request, - tracebacks_enabled, - ) - unflagged_records: Final = tuple( - record for record in caplog.records if "user_api_key_auth(): Exception occured" in record.getMessage() - ) - assert tuple(record.levelname for record in unflagged_records) == ("ERROR",) - assert unflagged_records[0].exc_info is not None - assert vars(flagged_response) == vars(unflagged_response) - assert flagged_response.to_dict() == unflagged_response.to_dict() - finally: - verbose_proxy_logger.propagate = previous_propagate - - -@pytest.mark.asyncio -async def test_unflagged_missing_credential_keeps_error_traceback(caplog: pytest.LogCaptureFixture) -> None: - from litellm.proxy.auth.auth_utils import MissingCredentialError - - previous_propagate: Final = verbose_proxy_logger.propagate - verbose_proxy_logger.propagate = True - try: - with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): - await _capture_authentication_exception( - MissingCredentialError("Malformed API Key passed in. Ensure Key has `Bearer ` prefix."), - Request({"type": "http", "method": "POST", "path": "/mcp", "headers": []}), - False, - ) - finally: - verbose_proxy_logger.propagate = previous_propagate - - records: Final = tuple( + flagged_records: Final = tuple( record for record in caplog.records if "user_api_key_auth(): Exception occured" in record.getMessage() ) - assert tuple(record.levelname for record in records) == ("ERROR",) - assert records[0].exc_info is not None + assert tuple(record.levelname for record in flagged_records) == ("DEBUG",) + assert flagged_records[0].exc_info is None + + caplog.clear() + with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): + unflagged_response: Final = await _capture_authentication_exception( + MissingCredentialError(message), + unflagged_request, + tracebacks_enabled, + ) + unflagged_records: Final = tuple( + record for record in caplog.records if "user_api_key_auth(): Exception occured" in record.getMessage() + ) + assert tuple(record.levelname for record in unflagged_records) == ("ERROR",) + assert unflagged_records[0].exc_info is not None + assert vars(flagged_response) == vars(unflagged_response) + assert flagged_response.to_dict() == unflagged_response.to_dict() @pytest.mark.asyncio async def test_expected_missing_credential_callback_403_is_escalated_and_unchanged( caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, ) -> None: from litellm.proxy.auth.auth_utils import MissingCredentialError, mark_expected_credential_challenge @@ -1252,55 +1227,48 @@ async def test_expected_missing_credential_callback_403_is_escalated_and_unchang unflagged_request: Final = Request({"type": "http", "method": "POST", "path": "/mcp", "headers": []}) callback_exception: Final = HTTPException(status_code=403, detail="Forbidden") message: Final = "Malformed API Key passed in. Ensure Key has `Bearer ` prefix." - previous_propagate: Final = verbose_proxy_logger.propagate - verbose_proxy_logger.propagate = True - try: - with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): - flagged_response: Final = await _capture_authentication_exception( - MissingCredentialError(message), - flagged_request, - False, - callback_exception, - ) - flagged_records: Final = tuple( - record for record in caplog.records if "user_api_key_auth(): Exception occured" in record.getMessage() + monkeypatch.setattr(verbose_proxy_logger, "propagate", True) + with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): + flagged_response: Final = await _capture_authentication_exception( + MissingCredentialError(message), + flagged_request, + False, + callback_exception, ) - assert tuple(record.levelname for record in flagged_records) == ("DEBUG", "ERROR") - assert flagged_response.code == str(status.HTTP_403_FORBIDDEN) + flagged_records: Final = tuple( + record for record in caplog.records if "user_api_key_auth(): Exception occured" in record.getMessage() + ) + assert tuple(record.levelname for record in flagged_records) == ("DEBUG", "ERROR") + assert flagged_response.code == str(status.HTTP_403_FORBIDDEN) - caplog.clear() - with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): - unflagged_response: Final = await _capture_authentication_exception( - MissingCredentialError(message), - unflagged_request, - False, - callback_exception, - ) - assert vars(flagged_response) == vars(unflagged_response) - assert flagged_response.to_dict() == unflagged_response.to_dict() - finally: - verbose_proxy_logger.propagate = previous_propagate + caplog.clear() + with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): + unflagged_response: Final = await _capture_authentication_exception( + MissingCredentialError(message), + unflagged_request, + False, + callback_exception, + ) + assert vars(flagged_response) == vars(unflagged_response) + assert flagged_response.to_dict() == unflagged_response.to_dict() @pytest.mark.asyncio async def test_expected_credential_marker_does_not_quiet_policy_rejection( caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, ) -> None: from litellm.proxy.auth.auth_utils import mark_expected_credential_challenge request: Final = Request({"type": "http", "method": "POST", "path": "/mcp", "headers": []}) mark_expected_credential_challenge(request) - previous_propagate: Final = verbose_proxy_logger.propagate - verbose_proxy_logger.propagate = True - try: - with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): - response: Final = await _capture_authentication_exception( - HTTPException(status_code=403, detail="Policy rejected request"), - request, - False, - ) - finally: - verbose_proxy_logger.propagate = previous_propagate + monkeypatch.setattr(verbose_proxy_logger, "propagate", True) + with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): + response: Final = await _capture_authentication_exception( + HTTPException(status_code=403, detail="Policy rejected request"), + request, + False, + ) records: Final = tuple( record for record in caplog.records if "user_api_key_auth(): Exception occured" in record.getMessage() From 24d1bfe3390d19ede2852ded1b008d87498d9e6c Mon Sep 17 00:00:00 2001 From: joshua Date: Sat, 10 Oct 2026 21:21:15 +0000 Subject: [PATCH 3/8] ci: include MCP auth tests in coverage shard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-unit.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 7f01b38ddc8..33b1098b2e7 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -800,6 +800,7 @@ jobs: - shard: mcp-oauth test-path: |- + tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py From 9ee07c5cfc4e3e465c189307b432694cf2657829 Mon Sep 17 00:00:00 2001 From: joshua Date: Sat, 10 Oct 2026 21:34:27 +0000 Subject: [PATCH 4/8] ci: narrow MCP OAuth coverage shard tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-unit.yml | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 33b1098b2e7..cade9627618 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -800,7 +800,12 @@ jobs: - shard: mcp-oauth test-path: |- - tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py + tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py::test_process_mcp_request_marks_only_expected_credential_challenges + tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py::test_cold_start_fallback_classifies_structural_challenges + tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py::TestMCPPassthroughColdStartAdmission::test_cold_start_propagates_non_401_http_error + tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py::TestMCPPassthroughColdStartAdmission::test_cold_start_propagates_non_auth_proxy_exception + tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py::TestMCPPassthroughColdStartAdmission::test_cold_start_allows_401_for_path_passthrough_target + tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py::TestMCPPassthroughColdStartAdmission::test_cold_start_allows_proxy_exception_401_for_path_target tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py From 2f36119a32f284698559dd82b824d252c5193889 Mon Sep 17 00:00:00 2001 From: joshua Date: Sat, 10 Oct 2026 21:40:16 +0000 Subject: [PATCH 5/8] Revert "ci: narrow MCP OAuth coverage shard tests" This reverts commit 9ee07c5cfc4e3e465c189307b432694cf2657829. --- .github/workflows/test-unit.yml | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index cade9627618..33b1098b2e7 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -800,12 +800,7 @@ jobs: - shard: mcp-oauth test-path: |- - tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py::test_process_mcp_request_marks_only_expected_credential_challenges - tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py::test_cold_start_fallback_classifies_structural_challenges - tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py::TestMCPPassthroughColdStartAdmission::test_cold_start_propagates_non_401_http_error - tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py::TestMCPPassthroughColdStartAdmission::test_cold_start_propagates_non_auth_proxy_exception - tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py::TestMCPPassthroughColdStartAdmission::test_cold_start_allows_401_for_path_passthrough_target - tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py::TestMCPPassthroughColdStartAdmission::test_cold_start_allows_proxy_exception_401_for_path_target + tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py From f6b492a11142c00f1e9d40a3cfb2bc5b47bb0f7c Mon Sep 17 00:00:00 2001 From: joshua Date: Sat, 10 Oct 2026 21:40:17 +0000 Subject: [PATCH 6/8] Revert "ci: include MCP auth tests in coverage shard" This reverts commit 24d1bfe3390d19ede2852ded1b008d87498d9e6c. --- .github/workflows/test-unit.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 33b1098b2e7..7f01b38ddc8 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -800,7 +800,6 @@ jobs: - shard: mcp-oauth test-path: |- - tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py From f4c8845db0534989ea53032b2e1de902a7e5df85 Mon Sep 17 00:00:00 2001 From: joshua Date: Sat, 10 Oct 2026 21:59:47 +0000 Subject: [PATCH 7/8] test(mcp): cover cold-start log levels end to end without own-code doubles Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/mcp/test_mcp_oauth_flows.py | 103 +++++++++--- .../auth/test_user_api_key_auth_mcp.py | 153 ++++-------------- 2 files changed, 117 insertions(+), 139 deletions(-) diff --git a/tests/integration/mcp/test_mcp_oauth_flows.py b/tests/integration/mcp/test_mcp_oauth_flows.py index d09af78a8a5..fc957b5b443 100644 --- a/tests/integration/mcp/test_mcp_oauth_flows.py +++ b/tests/integration/mcp/test_mcp_oauth_flows.py @@ -8,7 +8,7 @@ import uuid from collections.abc import Iterator from dataclasses import dataclass from pathlib import Path -from typing import Final +from typing import Final, Literal from urllib.parse import parse_qs, urlsplit import httpx @@ -357,42 +357,107 @@ def test_per_user_authorization_code_with_pkce_binds_the_token_to_the_authorizin ) -def test_anonymous_oauth_mcp_cold_start_does_not_log_malformed_key_at_error( +@pytest.mark.parametrize( + "path_template,server_auth_type,auth_headers,expected_status,expect_www_authenticate," + "expected_challenge_fragment,expected_error_lines", + [ + pytest.param( + "/{alias}/mcp", + "oauth2", + (), + 401, + True, + "resource_metadata=", + 0, + id="per-server-oauth-no-credentials", + ), + pytest.param("/mcp", "oauth2", (), 401, True, None, 0, id="aggregate-oauth-no-credentials"), + pytest.param( + "/{alias}/mcp", + "oauth2", + (("Authorization", ""),), + 401, + True, + None, + 1, + id="oauth-empty-authorization", + ), + pytest.param( + "/{alias}/mcp", + "api_key", + (), + 401, + False, + None, + 1, + id="non-oauth-api-key", + ), + ], +) +def test_mcp_cold_start_malformed_key_log_level( gateway: Gateway, tmp_path: Path, + path_template: str, + server_auth_type: Literal["oauth2", "api_key"], + auth_headers: tuple[tuple[str, str], ...], + expected_status: int, + expect_www_authenticate: bool, + expected_challenge_fragment: str | None, + expected_error_lines: int, ) -> None: with mcp_peer() as peer, oauth_server() as auth, owned_proxy_process(gateway, tmp_path, {}) as owned: with owned.gateway.scenario() as scenario: alias: Final = "cold" + uuid.uuid4().hex[:8] - _register_oauth( - scenario, - peer, - auth, - alias, - auth_type="oauth2", - oauth2_flow="authorization_code", - credentials={"client_id": "cold-client", "client_secret": "cold-secret"}, - ) + if server_auth_type == "oauth2": + _register_oauth( + scenario, + peer, + auth, + alias, + auth_type="oauth2", + oauth2_flow="authorization_code", + credentials={"client_id": "cold-client", "client_secret": "cold-secret"}, + ) + else: + register_mcp( + scenario, + peer, + alias, + auth_type="api_key", + credentials={"auth_value": "static-" + uuid.uuid4().hex}, + extra_headers=["x-api-key"], + ) baseline_log: Final = owned.log.read_text() - response: Final = owned.gateway.client.post( - f"/{alias}/mcp", - headers=ACCEPT, + request_path: Final = path_template.format(alias=alias) + request_headers: Final = [*ACCEPT.items(), *auth_headers] + request: Final = owned.gateway.client.build_request( + "POST", + request_path, + headers=request_headers, json=INITIALIZE, ) + expected_authorization: Final = tuple( + value for name, value in auth_headers if name.lower() == "authorization" + ) + assert tuple(request.headers.get_list("authorization")) == expected_authorization + response: Final = owned.gateway.client.send(request) - assert response.status_code == 401, response.text - assert "www-authenticate" in response.headers, response.headers - assert "resource_metadata=" in response.headers["www-authenticate"] + assert response.status_code == expected_status, response.text + challenge_header: Final = response.headers.get("www-authenticate") + assert (challenge_header is not None) is expect_www_authenticate + assert expected_challenge_fragment is None or ( + challenge_header is not None and expected_challenge_fragment in challenge_header + ) request_log: Final = eventually( owned.log.read_text, - lambda value: f"POST /{alias}/mcp" in value[len(baseline_log) :], + lambda value: f"POST {request_path}" in value[len(baseline_log) :], ) error_lines: Final = tuple( line for line in request_log[len(baseline_log) :].splitlines() if "Malformed API Key" in line and "ERROR" in line ) - assert error_lines == () + assert len(error_lines) == expected_error_lines, error_lines def test_authorization_request_without_pkce_is_refused_before_reaching_the_authorization_server( diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 4ad43d30136..0c26fc5b62e 100644 --- a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -1,6 +1,7 @@ import contextlib import json import os +import uuid from datetime import datetime, timedelta, timezone from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -18,6 +19,7 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( is_mcp_admitted_user_subject, ) from litellm.proxy._types import ( + LiteLLM_MCPServerTable, LiteLLM_ObjectPermissionTable, ProxyException, SpecialHeaders, @@ -30,104 +32,6 @@ from litellm.types.agents import AgentCaller @pytest.mark.asyncio -@pytest.mark.parametrize( - "path,extra_headers,auth_type,custom_key_header_name,hide_target,expected", - [ - pytest.param("/oauth/mcp", (), MCPAuth.oauth2, None, False, True, id="gateway-oauth-server"), - pytest.param("/mcp", (), None, None, False, True, id="aggregate-mcp"), - pytest.param( - "/oauth/mcp", - ((b"authorization", b""),), - MCPAuth.oauth2, - None, - False, - False, - id="empty-authorization", - ), - pytest.param( - "/oauth/mcp", - ((b"authorization", b"Basic malformed"),), - MCPAuth.oauth2, - None, - False, - False, - id="malformed-authorization", - ), - pytest.param( - "/oauth/mcp", - ((b"x-litellm-api-key", b""),), - MCPAuth.oauth2, - None, - False, - False, - id="empty-litellm-key", - ), - pytest.param( - "/oauth/mcp", - ((b"x-custom-litellm-key", b""),), - MCPAuth.oauth2, - "x-custom-litellm-key", - False, - False, - id="custom-key-header", - ), - pytest.param("/keyed/mcp", (), MCPAuth.api_key, None, False, False, id="non-oauth-target"), - pytest.param("/missing/mcp", (), None, None, False, False, id="unresolved-target"), - pytest.param("/hidden/mcp", (), MCPAuth.oauth2, None, True, False, id="target-hidden-by-client-ip"), - ], -) -async def test_process_mcp_request_marks_only_expected_credential_challenges( - path: str, - extra_headers: tuple[tuple[bytes, bytes], ...], - auth_type: MCPAuth | None, - custom_key_header_name: str | None, - hide_target: bool, - expected: bool, -) -> None: - from litellm.proxy.auth.auth_utils import EXPECTED_CREDENTIAL_CHALLENGE_STATE_KEY - - server: Final[MCPServer | None] = ( - MCPServer( - server_id="server-id", - name="oauth" if auth_type == MCPAuth.oauth2 else "keyed", - server_name="oauth" if auth_type == MCPAuth.oauth2 else "keyed", - url="https://upstream.example/mcp", - transport="http", - auth_type=auth_type, - extra_headers=["x-api-key"] if auth_type == MCPAuth.api_key else None, - ) - if auth_type is not None - else None - ) - scope: Final = { - "type": "http", - "method": "POST", - "path": path, - "headers": [(b"host", b"testserver"), *extra_headers], - } - with ( - patch( - "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", - side_effect=HTTPException(status_code=401, detail="Authentication required"), - ) as mock_auth, - patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as mock_manager, - patch( - "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.IPAddressUtils.get_mcp_client_ip", - return_value="203.0.113.10" if hide_target else None, - ), - patch( - "litellm.proxy.proxy_server.general_settings", - {"litellm_key_header_name": custom_key_header_name}, - ), - ): - mock_manager.get_mcp_server_by_name.return_value = None if hide_target else server - with pytest.raises(HTTPException): - await MCPRequestHandler.process_mcp_request(scope) - - auth_request: Final = mock_auth.call_args.kwargs["request"] - assert getattr(auth_request.state, EXPECTED_CREDENTIAL_CHALLENGE_STATE_KEY, False) is expected - - @pytest.mark.parametrize( "path,auth_type,delegate_auth_to_upstream,pass_through,bearer_presented,mcp_auth_header," "mcp_server_auth_headers,mcp_servers,expected", @@ -230,7 +134,7 @@ async def test_process_mcp_request_marks_only_expected_credential_challenges( ), ], ) -def test_cold_start_fallback_classifies_structural_challenges( +async def test_cold_start_fallback_classifies_structural_challenges( path: str, auth_type: MCPAuth | None, delegate_auth_to_upstream: bool, @@ -242,35 +146,44 @@ def test_cold_start_fallback_classifies_structural_challenges( expected: Literal["defer_to_route_challenge", "gateway_challenge"] | None, ) -> None: from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import _cold_start_fallback + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager - server: Final[MCPServer | MagicMock | None] = ( - MagicMock(is_oauth_passthrough=True) - if pass_through - else MCPServer( - server_id="server-id", - name=path.rsplit("/", 1)[-1], - server_name=path.rsplit("/", 1)[-1], - url="https://upstream.example/mcp", - transport="http", - auth_type=auth_type, - delegate_auth_to_upstream=delegate_auth_to_upstream, - ) - if auth_type is not None - else None + server_alias: Final = f"cold-start-{uuid.uuid4().hex}" + request_path: Final = f"/mcp/{server_alias}" if path != "/mcp" else path + request_mcp_servers: Final[list[str] | None] = [server_alias] if mcp_servers is not None else None + server_auth_type: Final = auth_type or (MCPAuth.oauth2 if mcp_servers is not None else MCPAuth.none) + server_row: Final = LiteLLM_MCPServerTable( + server_id=f"cold-start-{uuid.uuid4().hex}", + server_name=server_alias, + alias=server_alias, + url=None, + transport="http", + auth_type=server_auth_type, + oauth2_flow="authorization_code" if server_auth_type == MCPAuth.oauth2 else None, + delegate_auth_to_upstream=delegate_auth_to_upstream, + dcr_bridge=server_auth_type == MCPAuth.oauth2 and not delegate_auth_to_upstream, + oauth_passthrough=pass_through, + extra_headers=["Authorization"] if pass_through else [], + created_at=datetime(2024, 1, 1, tzinfo=timezone.utc), + updated_at=datetime(2024, 1, 1, tzinfo=timezone.utc), ) - request: Final = Request({"type": "http", "method": "POST", "path": path, "headers": []}) - with patch("litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager") as manager: - manager.get_mcp_server_by_name.return_value = server + request: Final = Request({"type": "http", "method": "POST", "path": request_path, "headers": []}) + try: + await global_mcp_server_manager.add_server(server_row) + registered_server: Final[MCPServer | None] = global_mcp_server_manager.get_mcp_server_by_name(server_alias) + assert isinstance(registered_server, MCPServer) + assert registered_server.is_oauth_passthrough is pass_through result: Final = _cold_start_fallback( request=request, - request_route=path, - mcp_servers=mcp_servers, + request_route=request_path, + mcp_servers=request_mcp_servers, mcp_auth_header=mcp_auth_header, mcp_server_auth_headers=mcp_server_auth_headers, bearer_presented=bearer_presented, ) - - assert result == expected + assert result == expected + finally: + global_mcp_server_manager.remove_server(server_row) @pytest.mark.parametrize("headers,expected", [ From 75f9e1cafa929b4c97063519a5c11ed872198f46 Mon Sep 17 00:00:00 2001 From: joshua Date: Sat, 10 Oct 2026 23:42:12 +0000 Subject: [PATCH 8/8] test(auth): drive failure-hook tests through real callbacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/auth/test_auth_exception_handler.py | 51 +++++++++++++++---- 1 file changed, 41 insertions(+), 10 deletions(-) diff --git a/tests/unit/proxy/auth/test_auth_exception_handler.py b/tests/unit/proxy/auth/test_auth_exception_handler.py index 9c2a982a3ac..f8539fe814c 100644 --- a/tests/unit/proxy/auth/test_auth_exception_handler.py +++ b/tests/unit/proxy/auth/test_auth_exception_handler.py @@ -1,6 +1,7 @@ import asyncio import json import logging +from collections.abc import Callable from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -31,6 +32,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import INVALID_VIRTUAL_KEY_ERROR_MARKER from litellm.exceptions import BudgetExceededError +from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ( ModelAccessDeniedProxyException, ProxyErrorTypes, @@ -1149,21 +1151,40 @@ async def test_handle_authentication_error_keeps_internal_message_on_model_acces assert [r for r in caplog.records if r.levelname == "WARNING" and "internal-models" in r.getMessage()] == [] +class _AuthenticationFailureCallback(CustomLogger): + def __init__(self, exception: HTTPException | None) -> None: + self._exception: Final[HTTPException | None] = exception + + async def async_post_call_failure_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> HTTPException | None: + return self._exception + + +@pytest.fixture +def register_authentication_failure_callback( + monkeypatch: pytest.MonkeyPatch, +) -> Callable[[HTTPException | None], _AuthenticationFailureCallback]: + def register(exception: HTTPException | None) -> _AuthenticationFailureCallback: + callback: Final = _AuthenticationFailureCallback(exception) + monkeypatch.setattr(litellm, "callbacks", [callback]) + return callback + + return register + + async def _capture_authentication_exception( exception: BaseException, request: Request, log_client_error_tracebacks: bool, - callback_exception: HTTPException | None = None, ) -> ProxyException: handler: Final = UserAPIKeyAuthExceptionHandler() with ( patch.object(litellm, "log_client_error_tracebacks", log_client_error_tracebacks), - patch( - "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", - new_callable=AsyncMock, - return_value=callback_exception, - ), - patch("litellm.proxy.auth.auth_exception_handler.seed_request_identity"), patch("litellm.proxy.proxy_server.general_settings", {"allow_requests_on_db_unavailable": False}), ): try: @@ -1179,9 +1200,13 @@ async def test_expected_missing_credential_logs_debug_without_traceback_and_pres tracebacks_enabled: bool, caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch, + register_authentication_failure_callback: Callable[ + [HTTPException | None], _AuthenticationFailureCallback + ], ) -> None: from litellm.proxy.auth.auth_utils import MissingCredentialError, mark_expected_credential_challenge + register_authentication_failure_callback(None) flagged_request: Final = Request({"type": "http", "method": "POST", "path": "/mcp", "headers": []}) mark_expected_credential_challenge(flagged_request) unflagged_request: Final = Request({"type": "http", "method": "POST", "path": "/mcp", "headers": []}) @@ -1219,13 +1244,17 @@ async def test_expected_missing_credential_logs_debug_without_traceback_and_pres async def test_expected_missing_credential_callback_403_is_escalated_and_unchanged( caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch, + register_authentication_failure_callback: Callable[ + [HTTPException | None], _AuthenticationFailureCallback + ], ) -> None: from litellm.proxy.auth.auth_utils import MissingCredentialError, mark_expected_credential_challenge + callback_exception: Final = HTTPException(status_code=403, detail="Forbidden") + register_authentication_failure_callback(callback_exception) flagged_request: Final = Request({"type": "http", "method": "POST", "path": "/mcp", "headers": []}) mark_expected_credential_challenge(flagged_request) unflagged_request: Final = Request({"type": "http", "method": "POST", "path": "/mcp", "headers": []}) - callback_exception: Final = HTTPException(status_code=403, detail="Forbidden") message: Final = "Malformed API Key passed in. Ensure Key has `Bearer ` prefix." monkeypatch.setattr(verbose_proxy_logger, "propagate", True) with caplog.at_level(logging.DEBUG, logger=verbose_proxy_logger.name): @@ -1233,7 +1262,6 @@ async def test_expected_missing_credential_callback_403_is_escalated_and_unchang MissingCredentialError(message), flagged_request, False, - callback_exception, ) flagged_records: Final = tuple( record for record in caplog.records if "user_api_key_auth(): Exception occured" in record.getMessage() @@ -1247,7 +1275,6 @@ async def test_expected_missing_credential_callback_403_is_escalated_and_unchang MissingCredentialError(message), unflagged_request, False, - callback_exception, ) assert vars(flagged_response) == vars(unflagged_response) assert flagged_response.to_dict() == unflagged_response.to_dict() @@ -1257,9 +1284,13 @@ async def test_expected_missing_credential_callback_403_is_escalated_and_unchang async def test_expected_credential_marker_does_not_quiet_policy_rejection( caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch, + register_authentication_failure_callback: Callable[ + [HTTPException | None], _AuthenticationFailureCallback + ], ) -> None: from litellm.proxy.auth.auth_utils import mark_expected_credential_challenge + register_authentication_failure_callback(None) request: Final = Request({"type": "http", "method": "POST", "path": "/mcp", "headers": []}) mark_expected_credential_challenge(request) monkeypatch.setattr(verbose_proxy_logger, "propagate", True)