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

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()

mcp_use/task_managers/http.py

@@ -0,0 +1,62 @@
+"""
+HTTP connection management for MCP implementations.
+
+This module provides a connection manager for HTTP-based MCP connections.
+"""
+
+import aiohttp
+
+from ..logging import logger
+from .base import ConnectionManager
+
+
+class HttpConnectionManager(ConnectionManager[aiohttp.ClientSession]):
+    """Connection manager for HTTP-based MCP connections.
+
+    This class handles the lifecycle of HTTP client sessions, ensuring proper
+    connection establishment and cleanup.
+    """
+
+    def __init__(
+        self,
+        base_url: str,
+        headers: dict[str, str] | None = None,
+    ):
+        """Initialize a new HTTP connection manager.
+
+        Args:
+            base_url: The base URL for HTTP requests
+            headers: Optional headers to include in all HTTP requests
+        """
+        super().__init__()
+        self.base_url = base_url.rstrip("/")
+        self.headers = headers or {}
+
+    async def _establish_connection(self) -> aiohttp.ClientSession:
+        """Establish an HTTP client session.
+
+        Returns:
+            The established HTTP client session
+
+        Raises:
+            Exception: If session cannot be established
+        """
+        logger.debug(f"Creating HTTP client session for: {self.base_url}")
+        try:
+            session = aiohttp.ClientSession(headers=self.headers)
+            return session
+        except Exception as e:
+            logger.error(f"Failed to create HTTP client session: {e}")
+            raise
+
+    async def _close_connection(self, connection: aiohttp.ClientSession) -> None:
+        """Close the HTTP client session.
+
+        Args:
+            connection: The HTTP client session to close
+        """
+        try:
+            logger.debug("Closing HTTP client session")
+            await connection.close()
+        except Exception as e:
+            logger.warning(f"Error closing HTTP client session: {e}")

mcp_use/task_managers/stdio.py

@@ -0,0 +1,73 @@
+"""
+StdIO connection management for MCP implementations.
+
+This module provides a connection manager for stdio-based MCP connections
+that ensures proper task isolation and resource cleanup.
+"""
+
+import sys
+from typing import Any, TextIO
+
+from mcp import StdioServerParameters
+from mcp.client.stdio import stdio_client
+
+from ..logging import logger
+from .base import ConnectionManager
+
+
+class StdioConnectionManager(ConnectionManager[tuple[Any, Any]]):
+    """Connection manager for stdio-based MCP connections.
+
+    This class handles the proper task isolation for stdio_client context managers
+    to prevent the "cancel scope in different task" error. It runs the stdio_client
+    in a dedicated task and manages its lifecycle.
+    """
+
+    def __init__(
+        self,
+        server_params: StdioServerParameters,
+        errlog: TextIO = sys.stderr,
+    ):
+        """Initialize a new stdio connection manager.
+
+        Args:
+            server_params: The parameters for the stdio server
+            errlog: The error log stream
+        """
+        super().__init__()
+        self.server_params = server_params
+        self.errlog = errlog
+        self._stdio_ctx = None
+
+    async def _establish_connection(self) -> tuple[Any, Any]:
+        """Establish a stdio connection.
+
+        Returns:
+            A tuple of (read_stream, write_stream)
+
+        Raises:
+            Exception: If connection cannot be established.
+        """
+        # Create the context manager
+        self._stdio_ctx = stdio_client(self.server_params, self.errlog)
+
+        # Enter the context manager
+        read_stream, write_stream = await self._stdio_ctx.__aenter__()
+
+        # Return the streams
+        return (read_stream, write_stream)
+
+    async def _close_connection(self, connection: tuple[Any, Any]) -> None:
+        """Close the stdio connection.
+
+        Args:
+            connection: The connection to close (ignored, we use the context manager)
+        """
+        if self._stdio_ctx:
+            # Exit the context manager
+            try:
+                await self._stdio_ctx.__aexit__(None, None, None)
+            except Exception as e:
+                logger.warning(f"Error closing stdio context: {e}")
+            finally:
+                self._stdio_ctx = None

mcp_use/task_managers/websocket.py

@@ -0,0 +1,63 @@
+"""
+WebSocket connection management for MCP implementations.
+
+This module provides a connection manager for WebSocket-based MCP connections.
+"""
+
+import websockets
+from websockets.client import WebSocketClientProtocol
+
+from ..logging import logger
+from .base import ConnectionManager
+
+
+class WebSocketConnectionManager(ConnectionManager[WebSocketClientProtocol]):
+    """Connection manager for WebSocket-based MCP connections.
+
+    This class handles the lifecycle of WebSocket connections, ensuring proper
+    connection establishment and cleanup.
+    """
+
+    def __init__(
+        self,
+        url: str,
+        headers: dict[str, str] | None = None,
+    ):
+        """Initialize a new WebSocket connection manager.
+
+        Args:
+            url: The WebSocket URL to connect to
+            headers: Optional headers to include in the WebSocket connection
+        """
+        super().__init__()
+        self.url = url
+        self.headers = headers or {}
+
+    async def _establish_connection(self) -> WebSocketClientProtocol:
+        """Establish a WebSocket connection.
+
+        Returns:
+            The established WebSocket connection
+
+        Raises:
+            Exception: If connection cannot be established
+        """
+        logger.debug(f"Connecting to WebSocket: {self.url}")
+        try:
+            ws = await websockets.connect(self.url, extra_headers=self.headers)
+            return ws
+        except Exception as e:
+            logger.error(f"Failed to connect to WebSocket: {e}")
+            raise
+
+    async def _close_connection(self, connection: WebSocketClientProtocol) -> None:
+        """Close the WebSocket connection.
+
+        Args:
+            connection: The WebSocket connection to close
+        """
+        try:
+            logger.debug("Closing WebSocket connection")
+            await connection.close()
+        except Exception as e:
+            logger.warning(f"Error closing WebSocket connection: {e}")