From 7b234ebfd6e0416e48b266dc288aa00c6e5fe5c6 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 12 Mar 2026 16:40:57 -0700 Subject: [PATCH] right error codes for mcp --- .../litellm_core_utils/llm_request_utils.py | 5 +- .../proxy/_experimental/mcp_server/server.py | 105 +++++-- litellm/proxy/auth/auth_exception_handler.py | 2 +- .../test_llm_request_utils.py | 39 +++ .../mcp_server/test_mcp_error_handling.py | 278 ++++++++++++++++++ 5 files changed, 406 insertions(+), 23 deletions(-) create mode 100644 tests/test_litellm/litellm_core_utils/test_llm_request_utils.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_error_handling.py diff --git a/litellm/litellm_core_utils/llm_request_utils.py b/litellm/litellm_core_utils/llm_request_utils.py index a2292d6e00f..2c0920bc76e 100644 --- a/litellm/litellm_core_utils/llm_request_utils.py +++ b/litellm/litellm_core_utils/llm_request_utils.py @@ -77,8 +77,7 @@ def get_proxy_server_request_headers(litellm_params: Optional[dict]) -> dict: if litellm_params is None: return {} - proxy_request_headers = ( - litellm_params.get("proxy_server_request", {}).get("headers", {}) or {} - ) + proxy_server_request = litellm_params.get("proxy_server_request") or {} + proxy_request_headers = proxy_server_request.get("headers", {}) or {} return proxy_request_headers diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index cd06de2a2df..92bbc54057a 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2192,7 +2192,9 @@ if MCP_AVAILABLE: def _get_mcp_servers_in_path(path: str) -> Optional[List[str]]: """ - Get the MCP servers from the path + Get the MCP servers from the path. + Handles both /mcp/ (full path) and / (child app receives + stripped path when mounted at /mcp). """ import re @@ -2234,6 +2236,12 @@ if MCP_AVAILABLE: mcp_servers_from_path = [server_name] else: mcp_servers_from_path = [servers_and_path] + else: + # Child app receives path like /undefined when mounted at /mcp + # Extract first path segment as server name + segments = path.strip("/").split("/") + if segments and segments[0] and "?" not in segments[0]: + mcp_servers_from_path = [segments[0]] return mcp_servers_from_path async def extract_mcp_auth_context(scope, path): @@ -2362,6 +2370,8 @@ if MCP_AVAILABLE: scope: Scope, receive: Receive, send: Send ) -> None: """Handle MCP requests through StreamableHTTP.""" + from litellm.proxy._types import ProxyException + try: path = scope.get("path", "") ( @@ -2382,12 +2392,17 @@ if MCP_AVAILABLE: verbose_logger.debug( f"MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}" ) - # https://datatracker.ietf.org/doc/html/rfc9728#name-www-authenticate-response + # Validate MCP servers exist and check OAuth requirements for server_name in mcp_servers or []: server = global_mcp_server_manager.get_mcp_server_by_name( server_name, client_ip=_client_ip ) - if server and server.auth_type == MCPAuth.oauth2 and not oauth2_headers: + if server is None: + raise HTTPException( + status_code=404, + detail=f"MCP server '{server_name}' not found", + ) + if server.auth_type == MCPAuth.oauth2 and not oauth2_headers: request = StarletteRequest(scope) base_url = get_request_base_url(request) @@ -2442,18 +2457,38 @@ if MCP_AVAILABLE: return await session_manager.handle_request(scope, receive, send) - except HTTPException: - # Re-raise HTTP exceptions to preserve status codes and details - raise + except HTTPException as e: + try: + detail = e.detail + if isinstance(detail, dict): + content = {"error": detail} + else: + content = {"error": {"message": str(detail)}} + error_response = JSONResponse( + status_code=e.status_code, + content=content, + ) + await error_response(scope, receive, send) + except Exception: + raise e + except ProxyException as e: + status_code = int(e.code) if e.code else 500 + verbose_logger.warning( + "MCP auth error (status=%s): %s", status_code, e.message + ) + try: + error_response = JSONResponse( + status_code=status_code, + content={"error": {"message": e.message, "type": e.type}}, + ) + await error_response(scope, receive, send) + except Exception: + raise 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 try: - from starlette.responses import JSONResponse - from starlette.status import HTTP_500_INTERNAL_SERVER_ERROR - error_response = JSONResponse( - status_code=HTTP_500_INTERNAL_SERVER_ERROR, + status_code=500, content={"error": "MCP request failed", "details": str(e)}, ) await error_response(scope, receive, send) @@ -2461,11 +2496,12 @@ if MCP_AVAILABLE: verbose_logger.exception( f"Failed to send error response: {response_error}" ) - # If we can't send a proper response, re-raise the original error raise e async def handle_sse_mcp(scope: Scope, receive: Receive, send: Send) -> None: """Handle MCP requests through SSE.""" + from litellm.proxy._types import ProxyException + try: path = scope.get("path", "") ( @@ -2486,6 +2522,16 @@ if MCP_AVAILABLE: verbose_logger.debug( f"MCP server auth headers: {list(mcp_server_auth_headers.keys()) if mcp_server_auth_headers else None}" ) + # Validate MCP servers exist (same as handle_streamable_http_mcp) + for server_name in mcp_servers or []: + server = global_mcp_server_manager.get_mcp_server_by_name( + server_name, client_ip=_sse_client_ip + ) + if server is None: + raise HTTPException( + status_code=404, + detail=f"MCP server '{server_name}' not found", + ) set_auth_context( user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, @@ -2501,16 +2547,38 @@ if MCP_AVAILABLE: await asyncio.sleep(0.1) await sse_session_manager.handle_request(scope, receive, send) + except HTTPException as e: + try: + detail = e.detail + if isinstance(detail, dict): + content = {"error": detail} + else: + content = {"error": {"message": str(detail)}} + error_response = JSONResponse( + status_code=e.status_code, + content=content, + ) + await error_response(scope, receive, send) + except Exception: + raise e + except ProxyException as e: + status_code = int(e.code) if e.code else 500 + verbose_logger.warning( + "MCP SSE auth error (status=%s): %s", status_code, e.message + ) + try: + error_response = JSONResponse( + status_code=status_code, + content={"error": {"message": e.message, "type": e.type}}, + ) + await error_response(scope, receive, send) + except Exception: + raise e except Exception as e: verbose_logger.exception(f"Error handling MCP request: {e}") - # Instead of re-raising, try to send a graceful error response try: - # Send a proper HTTP error response instead of letting the exception bubble up - from starlette.responses import JSONResponse - from starlette.status import HTTP_500_INTERNAL_SERVER_ERROR - error_response = JSONResponse( - status_code=HTTP_500_INTERNAL_SERVER_ERROR, + status_code=500, content={"error": "MCP request failed", "details": str(e)}, ) await error_response(scope, receive, send) @@ -2518,7 +2586,6 @@ if MCP_AVAILABLE: verbose_logger.exception( f"Failed to send error response: {response_error}" ) - # If we can't send a proper response, re-raise the original error raise e app = FastAPI( diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 9c306acd2c6..2768b6e0804 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -75,7 +75,7 @@ class UserAPIKeyAuthExceptionHandler: request=request, use_x_forwarded_for=general_settings.get("use_x_forwarded_for", False), ) - verbose_proxy_logger.exception( + verbose_proxy_logger.warning( "litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - {}\nRequester IP Address:{}".format( str(e), requester_ip, diff --git a/tests/test_litellm/litellm_core_utils/test_llm_request_utils.py b/tests/test_litellm/litellm_core_utils/test_llm_request_utils.py new file mode 100644 index 00000000000..0a2efe94063 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_llm_request_utils.py @@ -0,0 +1,39 @@ +""" +Tests for llm_request_utils helper functions. +""" + +import pytest + +from litellm.litellm_core_utils.llm_request_utils import ( + get_proxy_server_request_headers, +) + + +class TestGetProxyServerRequestHeaders: + """Tests for get_proxy_server_request_headers.""" + + def test_should_return_empty_dict_when_litellm_params_is_none(self): + assert get_proxy_server_request_headers(None) == {} + + def test_should_return_empty_dict_when_proxy_server_request_is_none(self): + """When proxy_server_request is explicitly None (e.g. MCP tool calls), + should return {} instead of raising 'NoneType has no attribute get'.""" + litellm_params = {"proxy_server_request": None} + assert get_proxy_server_request_headers(litellm_params) == {} + + def test_should_return_empty_dict_when_proxy_server_request_missing(self): + litellm_params = {} + assert get_proxy_server_request_headers(litellm_params) == {} + + def test_should_return_empty_dict_when_headers_is_none(self): + litellm_params = {"proxy_server_request": {"headers": None}} + assert get_proxy_server_request_headers(litellm_params) == {} + + def test_should_return_headers_when_present(self): + expected = {"Authorization": "Bearer sk-123", "Content-Type": "application/json"} + litellm_params = {"proxy_server_request": {"headers": expected}} + assert get_proxy_server_request_headers(litellm_params) == expected + + def test_should_return_empty_dict_when_proxy_server_request_has_no_headers_key(self): + litellm_params = {"proxy_server_request": {"body": "{}"}} + assert get_proxy_server_request_headers(litellm_params) == {} diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_error_handling.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_error_handling.py new file mode 100644 index 00000000000..20bce600a20 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_error_handling.py @@ -0,0 +1,278 @@ +""" +Tests for MCP endpoint error handling. + +Verifies that auth failures and other errors in MCP endpoints return proper +HTTP status codes (401, 403) instead of 500, and that ProxyException is +caught and converted to a JSON error response. +""" + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + + +class TestHandleStreamableHttpMcpErrorHandling: + """Tests that handle_streamable_http_mcp properly handles ProxyException.""" + + @pytest.mark.asyncio + async def test_should_return_401_on_auth_failure(self): + """Auth failures (ProxyException with code=401) should return 401, not 500.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + ) + except ImportError: + pytest.skip("MCP server not available") + + from litellm.proxy._types import ProxyErrorTypes, ProxyException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/undefined", + "headers": [(b"content-type", b"application/json")], + "query_string": b"", + "server": ("localhost", 8000), + "scheme": "http", + } + receive = AsyncMock() + send = AsyncMock() + + auth_error = ProxyException( + message="Malformed API Key passed in. Ensure Key has `Bearer ` prefix.", + type=ProxyErrorTypes.auth_error, + param="None", + code=401, + ) + + with patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + side_effect=auth_error, + ): + await handle_streamable_http_mcp(scope, receive, send) + + assert send.called + # First call is http.response.start with status code + start_call = send.call_args_list[0] + message = start_call[0][0] + assert message["type"] == "http.response.start" + assert message["status"] == 401 + + @pytest.mark.asyncio + async def test_should_return_403_on_forbidden(self): + """ProxyException with code=403 should return 403, not 500.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + ) + except ImportError: + pytest.skip("MCP server not available") + + from litellm.proxy._types import ProxyErrorTypes, ProxyException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/some_server", + "headers": [(b"content-type", b"application/json")], + "query_string": b"", + "server": ("localhost", 8000), + "scheme": "http", + } + receive = AsyncMock() + send = AsyncMock() + + forbidden_error = ProxyException( + message="Access denied", + type=ProxyErrorTypes.auth_error, + param="None", + code=403, + ) + + with patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + side_effect=forbidden_error, + ): + await handle_streamable_http_mcp(scope, receive, send) + + assert send.called + start_call = send.call_args_list[0] + message = start_call[0][0] + assert message["type"] == "http.response.start" + assert message["status"] == 403 + + @pytest.mark.asyncio + async def test_should_return_error_body_as_json(self): + """ProxyException responses should include error message and type in JSON body.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + ) + except ImportError: + pytest.skip("MCP server not available") + + import json + + from litellm.proxy._types import ProxyErrorTypes, ProxyException + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/undefined", + "headers": [(b"content-type", b"application/json")], + "query_string": b"", + "server": ("localhost", 8000), + "scheme": "http", + } + receive = AsyncMock() + send = AsyncMock() + + auth_error = ProxyException( + message="Malformed API Key", + type=ProxyErrorTypes.auth_error, + param="None", + code=401, + ) + + with patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + side_effect=auth_error, + ): + await handle_streamable_http_mcp(scope, receive, send) + + # Second call is http.response.body + body_call = send.call_args_list[1] + body_message = body_call[0][0] + assert body_message["type"] == "http.response.body" + body = json.loads(body_message["body"]) + assert body["error"]["message"] == "Malformed API Key" + assert body["error"]["type"] == ProxyErrorTypes.auth_error + + @pytest.mark.asyncio + async def test_should_return_404_for_nonexistent_mcp_server(self): + """Non-existent MCP server names should return 404, not 200.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + ) + except ImportError: + pytest.skip("MCP server not available") + + scope = { + "type": "http", + "method": "POST", + "path": "/undefined", + "headers": [ + (b"content-type", b"application/json"), + (b"accept", b"application/json, text/event-stream"), + (b"authorization", b"Bearer sk-1620"), + ], + "query_string": b"", + "server": ("localhost", 8000), + "scheme": "http", + } + receive = AsyncMock() + send = AsyncMock() + + mock_auth = MagicMock() + with patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(mock_auth, None, ["undefined"], {}, None, {}), + ), patch( + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager" + ) as mock_mgr: + mock_mgr.get_mcp_server_by_name.return_value = None + + await handle_streamable_http_mcp(scope, receive, send) + + assert send.called + start_call = send.call_args_list[0] + message = start_call[0][0] + assert message["type"] == "http.response.start" + assert message["status"] == 404 + + @pytest.mark.asyncio + async def test_should_still_500_on_unexpected_exceptions(self): + """Non-ProxyException errors should still result in a 500 response.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + handle_streamable_http_mcp, + ) + except ImportError: + pytest.skip("MCP server not available") + + scope = { + "type": "http", + "method": "POST", + "path": "/mcp/test", + "headers": [(b"content-type", b"application/json")], + "query_string": b"", + "server": ("localhost", 8000), + "scheme": "http", + } + receive = AsyncMock() + send = AsyncMock() + + with patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + side_effect=RuntimeError("unexpected crash"), + ): + await handle_streamable_http_mcp(scope, receive, send) + + assert send.called + start_call = send.call_args_list[0] + message = start_call[0][0] + assert message["type"] == "http.response.start" + assert message["status"] == 500 + + +class TestHandleSseMcpErrorHandling: + """Tests that handle_sse_mcp properly handles ProxyException.""" + + @pytest.mark.asyncio + async def test_should_return_401_on_auth_failure(self): + """Auth failures in SSE handler should return 401, not 500.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + handle_sse_mcp, + ) + except ImportError: + pytest.skip("MCP server not available") + + from litellm.proxy._types import ProxyErrorTypes, ProxyException + + scope = { + "type": "http", + "method": "GET", + "path": "/sse", + "headers": [(b"accept", b"text/event-stream")], + "query_string": b"", + "server": ("localhost", 8000), + "scheme": "http", + } + receive = AsyncMock() + send = AsyncMock() + + auth_error = ProxyException( + message="Malformed API Key passed in.", + type=ProxyErrorTypes.auth_error, + param="None", + code=401, + ) + + with patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + side_effect=auth_error, + ): + await handle_sse_mcp(scope, receive, send) + + assert send.called + start_call = send.call_args_list[0] + message = start_call[0][0] + assert message["type"] == "http.response.start" + assert message["status"] == 401