diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 678137100a5..e891425274f 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -63,7 +63,11 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) -from litellm.proxy._types import SpecialMCPServerNames, UserAPIKeyAuth +from litellm.proxy._types import ( + ProxyException, + SpecialMCPServerNames, + UserAPIKeyAuth, +) from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.litellm_pre_call_utils import ( LiteLLMProxyRequestSetup, @@ -229,6 +233,28 @@ def _jsonrpc_text_has_top_level_method(text: str) -> bool: return False +def _proxy_exception_to_http_exception(exc: ProxyException) -> HTTPException: + """Map a ``ProxyException`` to an ``HTTPException`` that preserves its real + status code and headers. + + ``user_api_key_auth`` raises ``ProxyException`` (not ``HTTPException``) on + auth failures. The MCP ASGI handlers re-raise ``HTTPException`` to keep the + status and any ``WWW-Authenticate`` challenge, but a ``ProxyException`` would + otherwise fall through to their generic handler and be flattened to a 500 — + dropping the 401 + challenge an OAuth client needs to re-authenticate, so the + tool call surfaces as a cancelled/terminated session instead. + """ + try: + status_code = int(exc.code) + except (TypeError, ValueError): + status_code = 500 + return HTTPException( + status_code=status_code, + detail=exc.message, + headers=exc.headers or None, + ) + + if MCP_AVAILABLE: from mcp.server import Server from mcp.server.lowlevel.server import NotificationOptions @@ -4019,6 +4045,12 @@ if MCP_AVAILABLE: except HTTPException: # Re-raise HTTP exceptions to preserve status codes and details raise + except ProxyException as e: + # Auth failures from user_api_key_auth arrive as ProxyException, not + # HTTPException. Preserve the real status (e.g. 401 + WWW-Authenticate) + # so OAuth clients can re-authenticate instead of receiving a generic + # 500 that surfaces as a cancelled tool call. + raise _proxy_exception_to_http_exception(e) except Exception as e: verbose_logger.exception(f"Error handling MCP request: {e}") # Try to send a graceful error response for non-HTTP exceptions @@ -4136,6 +4168,12 @@ if MCP_AVAILABLE: # Re-raise HTTP exceptions to preserve status codes and details # (e.g. 401 + WWW-Authenticate challenges from OAuth pass-through). raise + except ProxyException as e: + # Auth failures from user_api_key_auth arrive as ProxyException, not + # HTTPException. Preserve the real status (e.g. 401 + WWW-Authenticate) + # so OAuth clients can re-authenticate instead of receiving a generic + # 500 that surfaces as a cancelled tool call. + raise _proxy_exception_to_http_exception(e) except Exception as e: verbose_logger.exception(f"Error handling MCP request: {e}") # Try to send a graceful error response for non-HTTP exceptions diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index c86ae966f21..f44552c2943 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -6293,3 +6293,153 @@ async def test_get_allowed_mcp_servers_from_mcp_server_names_empty_list_fails_cl ) assert result == [] + + +class TestProxyExceptionToHttpException: + """Auth failures reach the MCP ASGI handlers as ProxyException, not + HTTPException. The handlers must map them back to their real status and + headers; otherwise they fall through to the generic 500 handler, dropping + the 401 + WWW-Authenticate challenge an OAuth client needs to re-authenticate + and surfacing the tool call as a cancelled/terminated session. + """ + + def test_preserves_401_status_and_www_authenticate_header(self): + from litellm.proxy._experimental.mcp_server.server import ( + _proxy_exception_to_http_exception, + ) + from litellm.proxy._types import ProxyException + + exc = ProxyException( + message="Authentication Error, invalid token", + type="auth_error", + param="key", + code=401, + headers={"WWW-Authenticate": 'Bearer resource_metadata="/x"'}, + ) + + http_exc = _proxy_exception_to_http_exception(exc) + + assert http_exc.status_code == 401 + assert http_exc.detail == "Authentication Error, invalid token" + assert http_exc.headers["WWW-Authenticate"] == 'Bearer resource_metadata="/x"' + + def test_preserves_403_status(self): + from litellm.proxy._experimental.mcp_server.server import ( + _proxy_exception_to_http_exception, + ) + from litellm.proxy._types import ProxyException + + http_exc = _proxy_exception_to_http_exception( + ProxyException( + message="Forbidden", type="auth_error", param="key", code=403 + ) + ) + + assert http_exc.status_code == 403 + + def test_non_numeric_code_falls_back_to_500(self): + from litellm.proxy._experimental.mcp_server.server import ( + _proxy_exception_to_http_exception, + ) + from litellm.proxy._types import ProxyException + + # ProxyException normalises code to the string "None" when unset. + http_exc = _proxy_exception_to_http_exception( + ProxyException(message="boom", type="server_error", param=None, code=None) + ) + + assert http_exc.status_code == 500 + + +class TestStreamableHttpAuthErrorMapping: + """End-to-end guard for the handler wiring: a ProxyException from auth must + propagate as the real HTTPException (401 + WWW-Authenticate), not be + flattened to a generic 500 by the catch-all handler. + """ + + @pytest.mark.asyncio + async def test_streamable_http_propagates_proxy_exception_as_401(self): + from litellm.proxy._experimental.mcp_server import server as mcp_module + from litellm.proxy._types import ProxyException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/some_server", + "headers": [(b"x-litellm-api-key", b"sk-bad")], + } + + async def receive(): + return {"type": "http.request", "body": b"{}", "more_body": False} + + sent = [] + + async def send(message): + sent.append(message) + + auth_failure = ProxyException( + message="Authentication Error, invalid token", + type="auth_error", + param="key", + code=401, + headers={"WWW-Authenticate": "Bearer"}, + ) + + with patch.object( + mcp_module, + "extract_mcp_auth_context", + new=AsyncMock(side_effect=auth_failure), + ): + with pytest.raises(HTTPException) as exc_info: + await mcp_module.handle_streamable_http_mcp(scope, receive, send) + + assert exc_info.value.status_code == 401 + assert exc_info.value.headers["WWW-Authenticate"] == "Bearer" + # Must not have emitted a 500 body via the generic catch-all. + assert not any( + m.get("type") == "http.response.start" and m.get("status") == 500 + for m in sent + ) + + @pytest.mark.asyncio + async def test_sse_propagates_proxy_exception_as_401(self): + from litellm.proxy._experimental.mcp_server import server as mcp_module + from litellm.proxy._types import ProxyException + + scope = { + "type": "http", + "method": "GET", + "path": "/mcp/some_server", + "headers": [(b"x-litellm-api-key", b"sk-bad")], + } + + async def receive(): + return {"type": "http.request", "body": b"", "more_body": False} + + sent = [] + + async def send(message): + sent.append(message) + + auth_failure = ProxyException( + message="Authentication Error, invalid token", + type="auth_error", + param="key", + code=401, + headers={"WWW-Authenticate": "Bearer"}, + ) + + with patch.object( + mcp_module, + "extract_mcp_auth_context", + new=AsyncMock(side_effect=auth_failure), + ): + with pytest.raises(HTTPException) as exc_info: + await mcp_module.handle_sse_mcp(scope, receive, send) + + assert exc_info.value.status_code == 401 + assert exc_info.value.headers["WWW-Authenticate"] == "Bearer" + assert not any( + m.get("type") == "http.response.start" and m.get("status") == 500 + for m in sent + )