mcp-use 0.1.0__py3-none-any.whl → 1.0.1__py3-none-any.whl

mcp_use/connectors/stdio.py

@@ -5,13 +5,12 @@ This module provides a connector for communicating with MCP implementations
 through the standard input/output streams.
 """
 
-from typing import Any
+import sys
 
 from mcp import ClientSession, StdioServerParameters
-from mcp.client.stdio import stdio_client
-from mcp.types import Tool
 
 from ..logging import logger
+from ..task_managers import StdioConnectionManager
 from .base import BaseConnector
 
 
@@ -19,7 +18,8 @@ class StdioConnector(BaseConnector):
     """Connector for MCP implementations using stdio transport.
 
     This connector uses the stdio transport to communicate with MCP implementations
-    that are executed as child processes.
+    that are executed as child processes. It uses a connection manager to handle
+    the proper lifecycle management of the stdio client.
     """
 
     def __init__(
@@ -27,6 +27,7 @@ class StdioConnector(BaseConnector):
         command: str = "npx",
         args: list[str] | None = None,
         env: dict[str, str] | None = None,
+        errlog=sys.stderr,
     ):
         """Initialize a new stdio connector.
 
@@ -34,91 +35,44 @@ class StdioConnector(BaseConnector):
             command: The command to execute.
             args: Optional command line arguments.
             env: Optional environment variables.
+            errlog: Stream to write error output to.
         """
+        super().__init__()
         self.command = command
-        self.args = args
+        self.args = args or []  # Ensure args is never None
         self.env = env
-        self.client: ClientSession | None = None
-        self._stdio_ctx = None
-        self._tools: list[Tool] | None = None
+        self.errlog = errlog
 
     async def connect(self) -> None:
         """Establish a connection to the MCP implementation."""
-        server_params = StdioServerParameters(command=self.command, args=self.args, env=self.env)
-        self._stdio_ctx = stdio_client(server_params)
-        read_stream, write_stream = await self._stdio_ctx.__aenter__()
-        self.client = ClientSession(read_stream, write_stream, sampling_callback=None)
-        await self.client.__aenter__()
-
-    async def disconnect(self) -> None:
-        """Close the connection to the MCP implementation."""
+        if self._connected:
+            logger.debug("Already connected to MCP implementation")
+            return
+
+        logger.info(f"Connecting to MCP implementation: {self.command}")
         try:
-            if self.client:
-                await self.client.__aexit__(None, None, None)
-                self.client = None
-            if self._stdio_ctx:
-                await self._stdio_ctx.__aexit__(None, None, None)
-                self._stdio_ctx = None
+            # Create server parameters
+            server_params = StdioServerParameters(
+                command=self.command, args=self.args, env=self.env
+            )
+
+            # Create and start the connection manager
+            self._connection_manager = StdioConnectionManager(server_params, self.errlog)
+            read_stream, write_stream = await self._connection_manager.start()
+
+            # Create the client session
+            self.client = ClientSession(read_stream, write_stream, sampling_callback=None)
+            await self.client.__aenter__()
+
+            # Mark as connected
+            self._connected = True
+            logger.info(f"Successfully connected to MCP implementation: {self.command}")
+
         except Exception as e:
-            logger.warning(f"Warning: Error during stdio connector disconnect: {e}")
-        finally:
-            # Always clean up references even if there were errors
-            self.client = None
-            self._stdio_ctx = None
-            self._tools = None
-
-    async def __aenter__(self) -> "StdioConnector":
-        """Enter the async context manager."""
-        await self.connect()
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
-        """Exit the async context manager."""
-        await self.disconnect()
-
-    async def initialize(self) -> dict[str, Any]:
-        """Initialize the MCP session and return session information."""
-        if not self.client:
-            raise RuntimeError("MCP client is not connected")
-
-        # Initialize the session
-        result = await self.client.initialize()
-
-        # Get available tools
-        tools_result = await self.client.list_tools()
-        self._tools = tools_result.tools
-
-        return result
-
-    @property
-    def tools(self) -> list[Tool]:
-        """Get the list of available tools."""
-        if not self._tools:
-            raise RuntimeError("MCP client is not initialized")
-        return self._tools
-
-    async def call_tool(self, name: str, arguments: dict[str, Any]) -> Any:
-        """Call an MCP tool with the given arguments."""
-        if not self.client:
-            raise RuntimeError("MCP client is not connected")
-        return await self.client.call_tool(name, arguments)
-
-    async def list_resources(self) -> list[dict[str, Any]]:
-        """List all available resources from the MCP implementation."""
-        if not self.client:
-            raise RuntimeError("MCP client is not connected")
-        resources = await self.client.list_resources()
-        return resources
-
-    async def read_resource(self, uri: str) -> tuple[bytes, str]:
-        """Read a resource by URI."""
-        if not self.client:
-            raise RuntimeError("MCP client is not connected")
-        resource = await self.client.read_resource(uri)
-        return resource.content, resource.mimeType
-
-    async def request(self, method: str, params: dict[str, Any] | None = None) -> Any:
-        """Send a raw request to the MCP implementation."""
-        if not self.client:
-            raise RuntimeError("MCP client is not connected")
-        return await self.client.request({"method": method, "params": params or {}})
+            logger.error(f"Failed to connect to MCP implementation: {e}")
+
+            # Clean up any resources if connection failed
+            await self._cleanup_resources()
+
+            # Re-raise the original exception
+            raise

mcp_use/connectors/websocket.py

@@ -10,20 +10,26 @@ import json
 import uuid
 from typing import Any
 
-import websockets
+from mcp.types import Tool
 from websockets.client import WebSocketClientProtocol
 
+from ..logging import logger
+from ..task_managers import ConnectionManager, WebSocketConnectionManager
 from .base import BaseConnector
 
 
 class WebSocketConnector(BaseConnector):
     """Connector for MCP implementations using WebSocket transport.
 
-    This connector uses WebSockets to communicate with remote MCP implementations.
+    This connector uses WebSockets to communicate with remote MCP implementations,
+    using a connection manager to handle the proper lifecycle management.
     """
 
     def __init__(
-        self, url: str, auth_token: str | None = None, headers: dict[str, str] | None = None
+        self,
+        url: str,
+        auth_token: str | None = None,
+        headers: dict[str, str] | None = None,
     ):
         """Initialize a new WebSocket connector.
 
@@ -37,15 +43,43 @@ class WebSocketConnector(BaseConnector):
         self.headers = headers or {}
         if auth_token:
             self.headers["Authorization"] = f"Bearer {auth_token}"
+
         self.ws: WebSocketClientProtocol | None = None
+        self._connection_manager: ConnectionManager | None = None
+        self._receiver_task: asyncio.Task | None = None
         self.pending_requests: dict[str, asyncio.Future] = {}
+        self._tools: list[Tool] | None = None
+        self._connected = False
 
     async def connect(self) -> None:
         """Establish a connection to the MCP implementation."""
-        self.ws = await websockets.connect(self.url, extra_headers=self.headers)
+        if self._connected:
+            logger.debug("Already connected to MCP implementation")
+            return
+
+        logger.info(f"Connecting to MCP implementation via WebSocket: {self.url}")
+        try:
+            # Create and start the connection manager
+            self._connection_manager = WebSocketConnectionManager(self.url, self.headers)
+            self.ws = await self._connection_manager.start()
 
-        # Start the message receiver task
-        self._receiver_task = asyncio.create_task(self._receive_messages())
+            # Start the message receiver task
+            self._receiver_task = asyncio.create_task(
+                self._receive_messages(), name="websocket_receiver_task"
+            )
+
+            # Mark as connected
+            self._connected = True
+            logger.info(f"Successfully connected to MCP implementation via WebSocket: {self.url}")
+
+        except Exception as e:
+            logger.error(f"Failed to connect to MCP implementation via WebSocket: {e}")
+
+            # Clean up any resources if connection failed
+            await self._cleanup_resources()
+
+            # Re-raise the original exception
+            raise
 
     async def _receive_messages(self) -> None:
         """Continuously receive and process messages from the WebSocket."""
@@ -65,7 +99,12 @@ class WebSocketConnector(BaseConnector):
                         future.set_result(data["result"])
                     elif "error" in data:
                         future.set_exception(Exception(data["error"]))
+
+                    logger.debug(f"Received response for request {request_id}")
+                else:
+                    logger.debug(f"Received message: {data}")
         except Exception as e:
+            logger.error(f"Error in WebSocket message receiver: {e}")
             # If the websocket connection was closed or errored,
             # reject all pending requests
             for future in self.pending_requests.values():
@@ -74,22 +113,63 @@ class WebSocketConnector(BaseConnector):
 
     async def disconnect(self) -> None:
         """Close the connection to the MCP implementation."""
-        if self._receiver_task:
-            self._receiver_task.cancel()
-            try:
-                await self._receiver_task
-            except asyncio.CancelledError:
-                pass
+        if not self._connected:
+            logger.debug("Not connected to MCP implementation")
+            return
+
+        logger.info("Disconnecting from MCP implementation")
+        await self._cleanup_resources()
+        self._connected = False
+        logger.info("Disconnected from MCP implementation")
 
-        if self.ws:
-            await self.ws.close()
-            self.ws = None
+    async def _cleanup_resources(self) -> None:
+        """Clean up all resources associated with this connector."""
+        errors = []
+
+        # First cancel the receiver task
+        if self._receiver_task and not self._receiver_task.done():
+            try:
+                logger.debug("Cancelling WebSocket receiver task")
+                self._receiver_task.cancel()
+                try:
+                    await self._receiver_task
+                except asyncio.CancelledError:
+                    logger.debug("WebSocket receiver task cancelled successfully")
+                except Exception as e:
+                    logger.warning(f"Error during WebSocket receiver task cancellation: {e}")
+            except Exception as e:
+                error_msg = f"Error cancelling WebSocket receiver task: {e}"
+                logger.warning(error_msg)
+                errors.append(error_msg)
+            finally:
+                self._receiver_task = None
 
         # Reject any pending requests
-        for future in self.pending_requests.values():
-            if not future.done():
-                future.set_exception(ConnectionError("WebSocket disconnected"))
-        self.pending_requests.clear()
+        if self.pending_requests:
+            logger.debug(f"Rejecting {len(self.pending_requests)} pending requests")
+            for future in self.pending_requests.values():
+                if not future.done():
+                    future.set_exception(ConnectionError("WebSocket disconnected"))
+            self.pending_requests.clear()
+
+        # Then stop the connection manager
+        if self._connection_manager:
+            try:
+                logger.debug("Stopping connection manager")
+                await self._connection_manager.stop()
+            except Exception as e:
+                error_msg = f"Error stopping connection manager: {e}"
+                logger.warning(error_msg)
+                errors.append(error_msg)
+            finally:
+                self._connection_manager = None
+                self.ws = None
+
+        # Reset tools
+        self._tools = None
+
+        if errors:
+            logger.warning(f"Encountered {len(errors)} errors during resource cleanup")
 
     async def _send_request(self, method: str, params: dict[str, Any] | None = None) -> Any:
         """Send a request and wait for a response."""
@@ -106,37 +186,60 @@ class WebSocketConnector(BaseConnector):
         # Send the request
         await self.ws.send(json.dumps({"id": request_id, "method": method, "params": params or {}}))
 
+        logger.debug(f"Sent request {request_id} method: {method}")
+
         # Wait for the response
         try:
             return await future
         except Exception as e:
             # Remove the request from pending requests
             self.pending_requests.pop(request_id, None)
-            raise e
+            logger.error(f"Error waiting for response to request {request_id}: {e}")
+            raise
 
     async def initialize(self) -> dict[str, Any]:
         """Initialize the MCP session and return session information."""
-        return await self._send_request("initialize")
+        logger.info("Initializing MCP session")
+        result = await self._send_request("initialize")
+
+        # Get available tools
+        tools_result = await self.list_tools()
+        self._tools = [Tool(**tool) for tool in tools_result]
+
+        logger.info(f"MCP session initialized with {len(self._tools)} tools")
+        return result
 
     async def list_tools(self) -> list[dict[str, Any]]:
         """List all available tools from the MCP implementation."""
+        logger.debug("Listing tools")
         result = await self._send_request("tools/list")
         return result.get("tools", [])
 
+    @property
+    def tools(self) -> list[Tool]:
+        """Get the list of available tools."""
+        if not self._tools:
+            raise RuntimeError("MCP client is not initialized")
+        return self._tools
+
     async def call_tool(self, name: str, arguments: dict[str, Any]) -> Any:
         """Call an MCP tool with the given arguments."""
+        logger.debug(f"Calling tool '{name}' with arguments: {arguments}")
         return await self._send_request("tools/call", {"name": name, "arguments": arguments})
 
     async def list_resources(self) -> list[dict[str, Any]]:
         """List all available resources from the MCP implementation."""
+        logger.debug("Listing resources")
         result = await self._send_request("resources/list")
         return result
 
     async def read_resource(self, uri: str) -> tuple[bytes, str]:
         """Read a resource by URI."""
+        logger.debug(f"Reading resource: {uri}")
         result = await self._send_request("resources/read", {"uri": uri})
         return result.get("content", b""), result.get("mimeType", "")
 
     async def request(self, method: str, params: dict[str, Any] | None = None) -> Any:
         """Send a raw request to the MCP implementation."""
+        logger.debug(f"Sending request: {method} with params: {params}")
         return await self._send_request(method, params)

mcp_use/session.py

@@ -8,7 +8,6 @@ which handles authentication, initialization, and tool discovery.
 from typing import Any
 
 from .connectors.base import BaseConnector
-from .tools.converter import ModelProvider, ToolConverter
 
 
 class MCPSession:
@@ -21,18 +20,15 @@ class MCPSession:
     def __init__(
         self,
         connector: BaseConnector,
-        model_provider: str | ModelProvider,
         auto_connect: bool = True,
     ) -> None:
         """Initialize a new MCP session.
 
         Args:
             connector: The connector to use for communicating with the MCP implementation.
-            model_provider: The model provider to convert tools for.
             auto_connect: Whether to automatically connect to the MCP implementation.
         """
         self.connector = connector
-        self.tool_converter = ToolConverter(model_provider)
         self.session_info: dict[str, Any] | None = None
         self.tools: list[dict[str, Any]] = []
         self.auto_connect = auto_connect
@@ -100,14 +96,6 @@ class MCPSession:
         self.tools = self.connector.tools
         return self.tools
 
-    def get_tools_for_llm(self) -> list[dict[str, Any]]:
-        """Get the tools in the format required by the LLM.
-
-        Returns:
-            The list of tools in the LLM-specific format.
-        """
-        return self.tool_converter.convert_tools(self.tools)
-
     async def call_tool(self, name: str, arguments: dict[str, Any]) -> Any:
         """Call an MCP tool with the given arguments.
 
@@ -123,46 +111,3 @@ class MCPSession:
             await self.connect()
 
         return await self.connector.call_tool(name, arguments)
-
-    async def process_tool_calls(self, response: dict[str, Any]) -> list[dict[str, Any]]:
-        """Process tool calls from an LLM response.
-
-        This method parses tool calls from the LLM response, executes them,
-        and returns the results.
-
-        Args:
-            response: The response from the LLM.
-
-        Returns:
-            A list of tool call results, each containing 'name', 'arguments', and 'result' keys.
-        """
-        # Parse tool calls from the response
-        tool_calls = self.tool_converter.parse_tool_calls(response)
-
-        # Execute each tool call
-        results = []
-        for tool_call in tool_calls:
-            name = tool_call["name"]
-            arguments = tool_call["arguments"]
-
-            try:
-                result = await self.call_tool(name, arguments)
-                results.append(
-                    {
-                        "name": name,
-                        "arguments": arguments,
-                        "result": result,
-                        "error": None,
-                    }
-                )
-            except Exception as e:
-                results.append(
-                    {
-                        "name": name,
-                        "arguments": arguments,
-                        "result": None,
-                        "error": str(e),
-                    }
-                )
-
-        return results

mcp_use/task_managers/__init__.py

@@ -0,0 +1,19 @@
+"""
+Connectors for various MCP transports.
+
+This module provides interfaces for connecting to MCP implementations
+through different transport mechanisms.
+"""
+
+from .base import ConnectionManager
+from .sse import SseConnectionManager
+from .stdio import StdioConnectionManager
+from .websocket import WebSocketConnectionManager
+
+__all__ = [
+    "ConnectionManager",
+    "HttpConnectionManager",
+    "StdioConnectionManager",
+    "WebSocketConnectionManager",
+    "SseConnectionManager",
+]

mcp_use/task_managers/base.py

@@ -0,0 +1,151 @@
+"""
+Connection management for MCP implementations.
+
+This module provides an abstract base class for different types of connection
+managers used in MCP connectors.
+"""
+
+import asyncio
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+from ..logging import logger
+
+# Type variable for connection types
+T = TypeVar("T")
+
+
+class ConnectionManager(Generic[T], ABC):
+    """Abstract base class for connection managers.
+
+    This class defines the interface for different types of connection managers
+    used with MCP connectors.
+    """
+
+    def __init__(self):
+        """Initialize a new connection manager."""
+        self._ready_event = asyncio.Event()
+        self._done_event = asyncio.Event()
+        self._exception: Exception | None = None
+        self._connection: T | None = None
+        self._task: asyncio.Task | None = None
+
+    @abstractmethod
+    async def _establish_connection(self) -> T:
+        """Establish the connection.
+
+        This method should be implemented by subclasses to establish
+        the specific type of connection needed.
+
+        Returns:
+            The established connection.
+
+        Raises:
+            Exception: If connection cannot be established.
+        """
+        pass
+
+    @abstractmethod
+    async def _close_connection(self, connection: T) -> None:
+        """Close the connection.
+
+        This method should be implemented by subclasses to close
+        the specific type of connection.
+
+        Args:
+            connection: The connection to close.
+        """
+        pass
+
+    async def start(self) -> T:
+        """Start the connection manager and establish a connection.
+
+        Returns:
+            The established connection.
+
+        Raises:
+            Exception: If connection cannot be established.
+        """
+        # Reset state
+        self._ready_event.clear()
+        self._done_event.clear()
+        self._exception = None
+
+        # Create a task to establish and maintain the connection
+        self._task = asyncio.create_task(
+            self._connection_task(), name=f"{self.__class__.__name__}_task"
+        )
+
+        # Wait for the connection to be ready or fail
+        await self._ready_event.wait()
+
+        # If there was an exception, raise it
+        if self._exception:
+            raise self._exception
+
+        # Return the connection
+        if self._connection is None:
+            raise RuntimeError("Connection was not established")
+        return self._connection
+
+    async def stop(self) -> None:
+        """Stop the connection manager and close the connection."""
+        if self._task and not self._task.done():
+            # Cancel the task
+            logger.debug(f"Cancelling {self.__class__.__name__} task")
+            self._task.cancel()
+
+            # Wait for it to complete
+            try:
+                await self._task
+            except asyncio.CancelledError:
+                logger.debug(f"{self.__class__.__name__} task cancelled successfully")
+            except Exception as e:
+                logger.warning(f"Error stopping {self.__class__.__name__} task: {e}")
+
+        # Wait for the connection to be done
+        await self._done_event.wait()
+        logger.debug(f"{self.__class__.__name__} task completed")
+
+    async def _connection_task(self) -> None:
+        """Run the connection task.
+
+        This task establishes and maintains the connection until cancelled.
+        """
+        logger.debug(f"Starting {self.__class__.__name__} task")
+        try:
+            # Establish the connection
+            self._connection = await self._establish_connection()
+            logger.debug(f"{self.__class__.__name__} connected successfully")
+
+            # Signal that the connection is ready
+            self._ready_event.set()
+
+            # Wait indefinitely until cancelled
+            try:
+                # This keeps the connection open until cancelled
+                await asyncio.Event().wait()
+            except asyncio.CancelledError:
+                # Expected when stopping
+                logger.debug(f"{self.__class__.__name__} task received cancellation")
+                pass
+
+        except Exception as e:
+            # Store the exception
+            self._exception = e
+            logger.error(f"Error in {self.__class__.__name__} task: {e}")
+
+            # Signal that the connection is ready (with error)
+            self._ready_event.set()
+
+        finally:
+            # Close the connection if it was established
+            if self._connection is not None:
+                try:
+                    await self._close_connection(self._connection)
+                except Exception as e:
+                    logger.warning(f"Error closing connection in {self.__class__.__name__}: {e}")
+                self._connection = None
+
+            # Signal that the connection is done
+            self._done_event.set()