diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index c4ecce74838..e7f308167bd 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -2366,6 +2366,22 @@ if MCP_AVAILABLE: ] return False + async def _send_error_response( + scope: Scope, + receive: Receive, + send: Send, + status_code: int, + content: dict, + headers: Optional[dict] = None, + ) -> None: + """Send a JSON error response. Shared by handle_streamable_http_mcp and handle_sse_mcp.""" + response = JSONResponse( + status_code=status_code, + content=content, + headers=headers or {}, + ) + await response(scope, receive, send) + async def handle_streamable_http_mcp( scope: Scope, receive: Receive, send: Send ) -> None: @@ -2460,17 +2476,15 @@ if MCP_AVAILABLE: except HTTPException as e: try: detail = e.detail - if isinstance(detail, dict): - content = {"error": detail} - else: - content = {"error": {"message": str(detail)}} - headers = dict(e.headers) if e.headers else {} - error_response = JSONResponse( - status_code=e.status_code, - content=content, - headers=headers, + content = ( + {"error": detail} + if isinstance(detail, dict) + else {"error": {"message": str(detail)}} + ) + headers = dict(e.headers) if e.headers else {} + await _send_error_response( + scope, receive, send, e.status_code, content, headers ) - await error_response(scope, receive, send) except Exception: raise e except ProxyException as e: @@ -2484,21 +2498,25 @@ if MCP_AVAILABLE: "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 _send_error_response( + scope, + receive, + send, + status_code, + {"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: - error_response = JSONResponse( - status_code=500, - content={"error": "MCP request failed", "details": str(e)}, + await _send_error_response( + scope, + receive, + send, + 500, + {"error": "MCP request failed", "details": str(e)}, ) - await error_response(scope, receive, send) except Exception as response_error: verbose_logger.exception( f"Failed to send error response: {response_error}" @@ -2571,17 +2589,15 @@ if MCP_AVAILABLE: except HTTPException as e: try: detail = e.detail - if isinstance(detail, dict): - content = {"error": detail} - else: - content = {"error": {"message": str(detail)}} - headers = dict(e.headers) if e.headers else {} - error_response = JSONResponse( - status_code=e.status_code, - content=content, - headers=headers, + content = ( + {"error": detail} + if isinstance(detail, dict) + else {"error": {"message": str(detail)}} + ) + headers = dict(e.headers) if e.headers else {} + await _send_error_response( + scope, receive, send, e.status_code, content, headers ) - await error_response(scope, receive, send) except Exception: raise e except ProxyException as e: @@ -2595,21 +2611,25 @@ if MCP_AVAILABLE: "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 _send_error_response( + scope, + receive, + send, + status_code, + {"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: - error_response = JSONResponse( - status_code=500, - content={"error": "MCP request failed", "details": str(e)}, + await _send_error_response( + scope, + receive, + send, + 500, + {"error": "MCP request failed", "details": str(e)}, ) - await error_response(scope, receive, send) except Exception as response_error: verbose_logger.exception( f"Failed to send error response: {response_error}" 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 index 20bce600a20..a3e90f8a5c5 100644 --- 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 @@ -276,3 +276,78 @@ class TestHandleSseMcpErrorHandling: message = start_call[0][0] assert message["type"] == "http.response.start" assert message["status"] == 401 + + @pytest.mark.asyncio + async def test_should_return_404_for_nonexistent_mcp_server(self): + """Non-existent MCP server names in SSE handler should return 404, not 200.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + handle_sse_mcp, + ) + except ImportError: + pytest.skip("MCP server not available") + + 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() + + 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_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"] == 404 + + @pytest.mark.asyncio + async def test_should_still_500_on_unexpected_exceptions(self): + """Non-ProxyException errors in SSE handler should still result in a 500 response.""" + try: + from litellm.proxy._experimental.mcp_server.server import ( + handle_sse_mcp, + ) + except ImportError: + pytest.skip("MCP server not available") + + 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() + + with patch( + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + side_effect=RuntimeError("unexpected crash"), + ): + 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"] == 500