added websocket support

Author: eavanvalkenburg

python/semantic_kernel/connectors/mcp.py

@@ -11,6 +11,7 @@
 from mcp.client.session import ClientSession
 from mcp.client.sse import sse_client
 from mcp.client.stdio import StdioServerParameters, stdio_client
+from mcp.client.websocket import websocket_client
 from mcp.types import CallToolResult, EmbeddedResource, Prompt, PromptMessage, TextResourceContents, Tool
 from mcp.types import (
     ImageContent as MCPImageContent,
@@ -348,3 +349,44 @@ def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
         if self._client_kwargs:
             args.update(self._client_kwargs)
         return sse_client(**args)
+
+
+class MCPWebsocketPlugin(MCPPluginBase):
+    """MCP websocket server configuration."""
+
+    def __init__(
+        self,
+        name: str,
+        url: str,
+        session: ClientSession | None = None,
+        description: str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize the MCP websocket plugin.
+
+                The arguments are used to create a websocket client.
+        see mcp.client.websocket.websocket_client for more details.
+
+        Any extra arguments passed to the constructor will be passed to the
+        websocket client constructor.
+
+        Args:
+            name: The name of the plugin.
+            url: The URL of the MCP server.
+            session: The session to use for the MCP connection.
+            description: The description of the plugin.
+            kwargs: Any extra arguments to pass to the websocket client.
+
+        """
+        super().__init__(name, description, session)
+        self.url = url
+        self._client_kwargs = kwargs
+
+    def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]:
+        """Get an MCP websocket client."""
+        args: dict[str, Any] = {
+            "url": self.url,
+        }
+        if self._client_kwargs:
+            args.update(self._client_kwargs)
+        return websocket_client(**args)

python/tests/unit/connectors/mcp/test_mcp.py

@@ -5,10 +5,7 @@
 import pytest
 from mcp import ClientSession, ListToolsResult, StdioServerParameters, Tool
 
-from semantic_kernel.connectors.mcp import (
-    MCPSsePlugin,
-    MCPStdioPlugin,
-)
+from semantic_kernel.connectors.mcp import MCPSsePlugin, MCPStdioPlugin, MCPWebsocketPlugin
 from semantic_kernel.exceptions import KernelPluginInvalidConfigurationError
 
 if TYPE_CHECKING:
@@ -116,6 +113,37 @@ async def test_with_kwargs_stdio(mock_session, mock_client, list_tool_calls, ker
         assert len(loaded_plugin.functions["func2"].parameters) == 0
 
 
+@patch("semantic_kernel.connectors.mcp.websocket_client")
+@patch("semantic_kernel.connectors.mcp.ClientSession")
+async def test_with_kwargs_websocket(mock_session, mock_client, list_tool_calls, kernel: "Kernel"):
+    mock_read = MagicMock()
+    mock_write = MagicMock()
+
+    mock_generator = MagicMock()
+    # Make the mock_stdio_client return an AsyncMock for the context manager
+    mock_generator.__aenter__.return_value = (mock_read, mock_write)
+    mock_generator.__aexit__.return_value = (mock_read, mock_write)
+
+    # Make the mock_stdio_client return an AsyncMock for the context manager
+    mock_client.return_value = mock_generator
+    mock_session.return_value.__aenter__.return_value.list_tools.return_value = list_tool_calls
+    async with MCPWebsocketPlugin(
+        name="TestMCPPlugin",
+        description="Test MCP Plugin",
+        url="http://localhost:8080/websocket",
+    ) as plugin:
+        mock_client.assert_called_once_with(url="http://localhost:8080/websocket")
+        loaded_plugin = kernel.add_plugin(plugin)
+        assert loaded_plugin is not None
+        assert loaded_plugin.name == "TestMCPPlugin"
+        assert loaded_plugin.description == "Test MCP Plugin"
+        assert loaded_plugin.functions.get("func1") is not None
+        assert loaded_plugin.functions["func1"].parameters[0].name == "name"
+        assert loaded_plugin.functions["func1"].parameters[0].is_required
+        assert loaded_plugin.functions.get("func2") is not None
+        assert len(loaded_plugin.functions["func2"].parameters) == 0
+
+
 @patch("semantic_kernel.connectors.mcp.sse_client")
 @patch("semantic_kernel.connectors.mcp.ClientSession")
 async def test_with_kwargs_sse(mock_session, mock_client, list_tool_calls, kernel: "Kernel"):