mcp-use 0.0.3__py3-none-any.whl → 0.0.4__py3-none-any.whl

mcp_use/connectors/http.py

@@ -8,18 +8,25 @@ through HTTP APIs.
 from typing import Any
 
 import aiohttp
+from mcp.types import Tool
 
+from ..logging import logger
+from ..task_managers import ConnectionManager, HttpConnectionManager
 from .base import BaseConnector
 
 
 class HttpConnector(BaseConnector):
     """Connector for MCP implementations using HTTP transport.
 
-    This connector uses HTTP requests to communicate with remote MCP implementations.
+    This connector uses HTTP requests to communicate with remote MCP implementations,
+    using a connection manager to handle the proper lifecycle management.
     """
 
     def __init__(
-        self, base_url: str, auth_token: str | None = None, headers: dict[str, str] | None = None
+        self,
+        base_url: str,
+        auth_token: str | None = None,
+        headers: dict[str, str] | None = None,
     ):
         """Initialize a new HTTP connector.
 
@@ -33,17 +40,70 @@ class HttpConnector(BaseConnector):
         self.headers = headers or {}
         if auth_token:
             self.headers["Authorization"] = f"Bearer {auth_token}"
+
         self.session: aiohttp.ClientSession | None = None
+        self._connection_manager: ConnectionManager | None = None
+        self._tools: list[Tool] | None = None
+        self._connected = False
 
     async def connect(self) -> None:
         """Establish a connection to the MCP implementation."""
-        self.session = aiohttp.ClientSession(headers=self.headers)
+        if self._connected:
+            logger.debug("Already connected to MCP implementation")
+            return
+
+        logger.info(f"Connecting to MCP implementation via HTTP: {self.base_url}")
+        try:
+            # Create and start the connection manager
+            self._connection_manager = HttpConnectionManager(self.base_url, self.headers)
+            self.session = await self._connection_manager.start()
+
+            # Mark as connected
+            self._connected = True
+            logger.info(f"Successfully connected to MCP implementation via HTTP: {self.base_url}")
+
+        except Exception as e:
+            logger.error(f"Failed to connect to MCP implementation via HTTP: {e}")
+
+            # Clean up any resources if connection failed
+            await self._cleanup_resources()
+
+            # Re-raise the original exception
+            raise
 
     async def disconnect(self) -> None:
         """Close the connection to the MCP implementation."""
-        if self.session:
-            await self.session.close()
-            self.session = None
+        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")
+
+    async def _cleanup_resources(self) -> None:
+        """Clean up all resources associated with this connector."""
+        errors = []
+
+        # 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.session = None
+
+        # Reset tools
+        self._tools = None
+
+        if errors:
+            logger.warning(f"Encountered {len(errors)} errors during resource cleanup")
 
     async def _request(self, method: str, endpoint: str, data: dict[str, Any] | None = None) -> Any:
         """Send an HTTP request to the MCP API.
@@ -60,6 +120,7 @@ class HttpConnector(BaseConnector):
             raise RuntimeError("HTTP session is not connected")
 
         url = f"{self.base_url}/{endpoint.lstrip('/')}"
+        logger.debug(f"Sending {method} request to {url}")
 
         if method.upper() == "GET" and data:
             # For GET requests, convert data to query parameters
@@ -74,24 +135,46 @@ class HttpConnector(BaseConnector):
 
     async def initialize(self) -> dict[str, Any]:
         """Initialize the MCP session and return session information."""
-        return await self._request("POST", "initialize")
+        logger.info("Initializing MCP session")
+
+        # Initialize the session
+        result = await self._request("POST", "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._request("GET", "tools")
         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._request("POST", f"tools/{name}", 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._request("GET", "resources")
         return result
 
     async def read_resource(self, uri: str) -> tuple[bytes, str]:
         """Read a resource by URI."""
+        logger.debug(f"Reading resource: {uri}")
+
         # For resources, we may need to handle binary data
         if not self.session:
             raise RuntimeError("HTTP session is not connected")
@@ -122,5 +205,6 @@ class HttpConnector(BaseConnector):
 
     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}")
         # For custom methods, we'll use the RPC-style endpoint
         return await self._request("POST", "rpc", {"method": method, "params": params or {}})

mcp_use/connectors/stdio.py

@@ -5,13 +5,14 @@ This module provides a connector for communicating with MCP implementations
 through the standard input/output streams.
 """
 
+import sys
 from typing import Any
 
 from mcp import ClientSession, StdioServerParameters
-from mcp.client.stdio import stdio_client
-from mcp.types import Tool
+from mcp.types import CallToolResult, Tool
 
 from ..logging import logger
+from ..task_managers import ConnectionManager, StdioConnectionManager
 from .base import BaseConnector
 
 
@@ -19,7 +20,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 +29,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,53 +37,103 @@ class StdioConnector(BaseConnector):
             command: The command to execute.
             args: Optional command line arguments.
             env: Optional environment variables.
+            errlog: Stream to write error output to.
         """
         self.command = command
-        self.args = args
+        self.args = args or []  # Ensure args is never None
         self.env = env
+        self.errlog = errlog
         self.client: ClientSession | None = None
-        self._stdio_ctx = None
+        self._connection_manager: ConnectionManager | None = None
         self._tools: list[Tool] | None = None
+        self._connected = False
 
     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__()
+        if self._connected:
+            logger.debug("Already connected to MCP implementation")
+            return
+
+        logger.info(f"Connecting to MCP implementation: {self.command}")
+        try:
+            # 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.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
 
     async def disconnect(self) -> None:
         """Close the connection to the MCP implementation."""
-        try:
-            if self.client:
+        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")
+
+    async def _cleanup_resources(self) -> None:
+        """Clean up all resources associated with this connector."""
+        errors = []
+
+        # First close the client session
+        if self.client:
+            try:
+                logger.debug("Closing client session")
                 await self.client.__aexit__(None, None, None)
+            except Exception as e:
+                error_msg = f"Error closing client session: {e}"
+                logger.warning(error_msg)
+                errors.append(error_msg)
+            finally:
                 self.client = None
-            if self._stdio_ctx:
-                await self._stdio_ctx.__aexit__(None, None, None)
-                self._stdio_ctx = None
-        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()
+
+        # 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
+
+        # Reset tools
+        self._tools = None
+
+        if errors:
+            logger.warning(f"Encountered {len(errors)} errors during resource cleanup")
 
     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")
 
+        logger.info("Initializing MCP session")
+
         # Initialize the session
         result = await self.client.initialize()
 
@@ -88,6 +141,8 @@ class StdioConnector(BaseConnector):
         tools_result = await self.client.list_tools()
         self._tools = tools_result.tools
 
+        logger.info(f"MCP session initialized with {len(self._tools)} tools")
+
         return result
 
     @property
@@ -97,16 +152,21 @@ class StdioConnector(BaseConnector):
             raise RuntimeError("MCP client is not initialized")
         return self._tools
 
-    async def call_tool(self, name: str, arguments: dict[str, Any]) -> Any:
+    async def call_tool(self, name: str, arguments: dict[str, Any]) -> CallToolResult:
         """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)
+
+        logger.debug(f"Calling tool '{name}' with arguments: {arguments}")
+        result = await self.client.call_tool(name, arguments)
+        return result
 
     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")
+
+        logger.debug("Listing resources")
         resources = await self.client.list_resources()
         return resources
 
@@ -114,6 +174,8 @@ class StdioConnector(BaseConnector):
         """Read a resource by URI."""
         if not self.client:
             raise RuntimeError("MCP client is not connected")
+
+        logger.debug(f"Reading resource: {uri}")
         resource = await self.client.read_resource(uri)
         return resource.content, resource.mimeType
 
@@ -121,4 +183,6 @@ class StdioConnector(BaseConnector):
         """Send a raw request to the MCP implementation."""
         if not self.client:
             raise RuntimeError("MCP client is not connected")
+
+        logger.debug(f"Sending request: {method} with params: {params}")
         return await self.client.request({"method": method, "params": params or {}})

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/task_managers/__init__.py

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