mcp-use 1.2.12__py3-none-any.whl → 1.3.0__py3-none-any.whl

mcp_use/connectors/sandbox.py

@@ -0,0 +1,291 @@
+"""
+Sandbox connector for MCP implementations.
+
+This module provides a connector for communicating with MCP implementations
+that are executed inside a sandbox environment (currently using E2B).
+"""
+
+import asyncio
+import os
+import sys
+import time
+
+import aiohttp
+from mcp import ClientSession
+
+from ..logging import logger
+from ..task_managers import SseConnectionManager
+
+# Import E2B SDK components (optional dependency)
+try:
+    logger.debug("Attempting to import e2b_code_interpreter...")
+    from e2b_code_interpreter import (
+        CommandHandle,
+        Sandbox,
+    )
+
+    logger.debug("Successfully imported e2b_code_interpreter")
+except ImportError as e:
+    logger.debug(f"Failed to import e2b_code_interpreter: {e}")
+    CommandHandle = None
+    Sandbox = None
+
+from ..types.sandbox import SandboxOptions
+from .base import BaseConnector
+
+
+class SandboxConnector(BaseConnector):
+    """Connector for MCP implementations running in a sandbox environment.
+
+    This connector runs a user-defined stdio command within a sandbox environment,
+    currently implemented using E2B, potentially wrapped by a utility like 'supergateway'
+    to expose its stdio.
+    """
+
+    def __init__(
+        self,
+        command: str,
+        args: list[str],
+        env: dict[str, str] | None = None,
+        e2b_options: SandboxOptions | None = None,
+        timeout: float = 5,
+        sse_read_timeout: float = 60 * 5,
+    ):
+        """Initialize a new sandbox connector.
+
+        Args:
+            command: The user's MCP server command to execute in the sandbox.
+            args: Command line arguments for the user's MCP server command.
+            env: Environment variables for the user's MCP server command.
+            e2b_options: Configuration options for the E2B sandbox environment.
+                        See SandboxOptions for available options and defaults.
+            timeout: Timeout for the sandbox process in seconds.
+            sse_read_timeout: Timeout for the SSE connection in seconds.
+        """
+        super().__init__()
+        if Sandbox is None:
+            raise ImportError(
+                "E2B SDK (e2b-code-interpreter) not found. "
+                "Please install it with 'pip install mcp-use[e2b]' "
+                "(or 'pip install e2b-code-interpreter')."
+            )
+
+        self.user_command = command
+        self.user_args = args or []
+        self.user_env = env or {}
+
+        _e2b_options = e2b_options or {}
+
+        self.api_key = _e2b_options.get("api_key") or os.environ.get("E2B_API_KEY")
+        if not self.api_key:
+            raise ValueError(
+                "E2B API key is required. Provide it via 'e2b_options.api_key' "
+                "or the E2B_API_KEY environment variable."
+            )
+
+        self.sandbox_template_id = _e2b_options.get("sandbox_template_id", "base")
+        self.supergateway_cmd_parts = _e2b_options.get(
+            "supergateway_command", "npx -y supergateway"
+        )
+
+        self.sandbox: Sandbox | None = None
+        self.process: CommandHandle | None = None
+        self.client: ClientSession | None = None
+        self.errlog = sys.stderr
+        self.base_url: str | None = None
+        self._connected = False
+        self._connection_manager: SseConnectionManager | None = None
+
+        # SSE connection parameters
+        self.headers = {}
+        self.timeout = timeout
+        self.sse_read_timeout = sse_read_timeout
+
+        self.stdout_lines: list[str] = []
+        self.stderr_lines: list[str] = []
+        self._server_ready = asyncio.Event()
+
+    def _handle_stdout(self, data: str) -> None:
+        """Handle stdout data from the sandbox process."""
+        self.stdout_lines.append(data)
+        logger.debug(f"[SANDBOX STDOUT] {data}", end="", flush=True)
+
+    def _handle_stderr(self, data: str) -> None:
+        """Handle stderr data from the sandbox process."""
+        self.stderr_lines.append(data)
+        logger.debug(f"[SANDBOX STDERR] {data}", file=self.errlog, end="", flush=True)
+
+    async def wait_for_server_response(self, base_url: str, timeout: int = 30) -> bool:
+        """Wait for the server to respond to HTTP requests.
+        Args:
+            base_url: The base URL to check for server readiness
+            timeout: Maximum time to wait in seconds
+        Returns:
+            True if server is responding, raises TimeoutError otherwise
+        """
+        logger.info(f"Waiting for server at {base_url} to respond...")
+        sys.stdout.flush()
+
+        start_time = time.time()
+        ping_url = f"{base_url}/sse"
+
+        # Try to connect to the server
+        while time.time() - start_time < timeout:
+            try:
+                async with aiohttp.ClientSession() as session:
+                    try:
+                        # First try the endpoint
+                        async with session.get(ping_url, timeout=2) as response:
+                            if response.status == 200:
+                                elapsed = time.time() - start_time
+                                logger.info(
+                                    f"Server is ready! "
+                                    f"SSE endpoint responded with 200 after {elapsed:.1f}s"
+                                )
+                                return True
+                    except Exception:
+                        # If sse endpoint doesn't work, try the base URL
+                        async with session.get(base_url, timeout=2) as response:
+                            if response.status < 500:  # Accept any non-server error
+                                elapsed = time.time() - start_time
+                                logger.info(
+                                    f"Server is ready! Base URL responded with "
+                                    f"{response.status} after {elapsed:.1f}s"
+                                )
+                                return True
+            except Exception:
+                # Wait a bit before trying again
+                await asyncio.sleep(0.5)
+                continue
+
+            # If we get here, the request failed
+            await asyncio.sleep(0.5)
+
+            # Log status every 5 seconds
+            elapsed = time.time() - start_time
+            if int(elapsed) % 5 == 0:
+                logger.info(f"Still waiting for server to respond... ({elapsed:.1f}s elapsed)")
+                sys.stdout.flush()
+
+        # If we get here, we timed out
+        raise TimeoutError(f"Timeout waiting for server to respond (waited {timeout} seconds)")
+
+    async def connect(self):
+        """Connect to the sandbox and start the MCP server."""
+
+        if self._connected:
+            logger.debug("Already connected to MCP implementation")
+            return
+
+        logger.debug("Connecting to MCP implementation in sandbox")
+
+        try:
+            # Create and start the sandbox
+            self.sandbox = Sandbox(
+                template=self.sandbox_template_id,
+                api_key=self.api_key,
+            )
+
+            # Get the host for the sandbox
+            host = self.sandbox.get_host(3000)
+            self.base_url = f"https://{host}".rstrip("/")
+
+            # Append command with args
+            command = f"{self.user_command} {' '.join(self.user_args)}"
+
+            # Construct the full command with supergateway
+            full_command = f'{self.supergateway_cmd_parts} \
+                --base-url {self.base_url} \
+                --port 3000 \
+                --cors \
+                --stdio "{command}"'
+
+            logger.debug(f"Full command: {full_command}")
+
+            # Start the process in the sandbox with our stdout/stderr handlers
+            self.process: CommandHandle = self.sandbox.commands.run(
+                full_command,
+                envs=self.user_env,
+                timeout=1000 * 60 * 10,  # 10 minutes timeout
+                background=True,
+                on_stdout=self._handle_stdout,
+                on_stderr=self._handle_stderr,
+            )
+
+            # Wait for the server to be ready
+            await self.wait_for_server_response(self.base_url, timeout=30)
+            logger.debug("Initializing connection manager...")
+
+            # Create the SSE connection URL
+            sse_url = f"{self.base_url}/sse"
+
+            # Create and start the connection manager
+            self._connection_manager = SseConnectionManager(
+                sse_url, self.headers, self.timeout, self.sse_read_timeout
+            )
+            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.debug(
+                f"Successfully connected to MCP implementation via HTTP/SSE: {self.base_url}"
+            )
+
+        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()
+
+            raise e
+
+    async def _cleanup_resources(self) -> None:
+        """Clean up all resources associated with this connector, including the sandbox.
+        This method extends the base implementation to also terminate the sandbox instance
+        and clean up any processes running in the sandbox.
+        """
+        logger.debug("Cleaning up sandbox resources")
+
+        # Terminate any running process
+        if self.process:
+            try:
+                logger.debug("Terminating sandbox process")
+                self.process.kill()
+            except Exception as e:
+                logger.warning(f"Error terminating sandbox process: {e}")
+            finally:
+                self.process = None
+
+        # Close the sandbox
+        if self.sandbox:
+            try:
+                logger.debug("Closing sandbox instance")
+                self.sandbox.kill()
+                logger.debug("Sandbox instance closed successfully")
+            except Exception as e:
+                logger.warning(f"Error closing sandbox: {e}")
+            finally:
+                self.sandbox = None
+
+        # Then call the parent method to clean up the rest
+        await super()._cleanup_resources()
+
+        # Clear any collected output
+        self.stdout_lines = []
+        self.stderr_lines = []
+        self.base_url = None
+
+    async def disconnect(self) -> None:
+        """Close the connection to the MCP implementation."""
+        if not self._connected:
+            logger.debug("Not connected to MCP implementation")
+            return
+
+        logger.debug("Disconnecting from MCP implementation")
+        await self._cleanup_resources()
+        self._connected = False
+        logger.debug("Disconnected from MCP implementation")

mcp_use/connectors/utils.py

@@ -0,0 +1,13 @@
+from typing import Any
+
+
+def is_stdio_server(server_config: dict[str, Any]) -> bool:
+    """Check if the server configuration is for a stdio server.
+
+    Args:
+        server_config: The server configuration section
+
+    Returns:
+        True if the server is a stdio server, False otherwise
+    """
+    return "command" in server_config and "args" in server_config

mcp_use/session.py

@@ -7,6 +7,8 @@ which handles authentication, initialization, and tool discovery.
 
 from typing import Any
 
+from mcp.types import Tool
+
 from .connectors.base import BaseConnector
 
 
@@ -30,7 +32,7 @@ class MCPSession:
         """
         self.connector = connector
         self.session_info: dict[str, Any] | None = None
-        self.tools: list[dict[str, Any]] = []
+        self.tools: list[Tool] = []
         self.auto_connect = auto_connect
 
     async def __aenter__(self) -> "MCPSession":
@@ -73,9 +75,6 @@ class MCPSession:
         # Initialize the session
         self.session_info = await self.connector.initialize()
 
-        # Discover available tools
-        await self.discover_tools()
-
         return self.session_info
 
     @property
@@ -86,28 +85,3 @@ class MCPSession:
             True if the connector is connected, False otherwise.
         """
         return hasattr(self.connector, "client") and self.connector.client is not None
-
-    async def discover_tools(self) -> list[dict[str, Any]]:
-        """Discover available tools from the MCP implementation.
-
-        Returns:
-            The list of available tools in MCP format.
-        """
-        self.tools = self.connector.tools
-        return self.tools
-
-    async def call_tool(self, name: str, arguments: dict[str, Any]) -> Any:
-        """Call an MCP tool with the given arguments.
-
-        Args:
-            name: The name of the tool to call.
-            arguments: The arguments to pass to the tool.
-
-        Returns:
-            The result of the tool call.
-        """
-        # Make sure we're connected
-        if not self.is_connected and self.auto_connect:
-            await self.connect()
-
-        return await self.connector.call_tool(name, arguments)

mcp_use/task_managers/base.py

@@ -46,14 +46,12 @@ class ConnectionManager(Generic[T], ABC):
         pass
 
     @abstractmethod
-    async def _close_connection(self, connection: T) -> None:
+    async def _close_connection(self) -> 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
 
@@ -139,10 +137,10 @@ class ConnectionManager(Generic[T], ABC):
             self._ready_event.set()
 
         finally:
-            # Close the connection if it was established
+            # Close the connection if it was establishedSUPABASE_URL
             if self._connection is not None:
                 try:
-                    await self._close_connection(self._connection)
+                    await self._close_connection()
                 except Exception as e:
                     logger.warning(f"Error closing connection in {self.__class__.__name__}: {e}")
                 self._connection = None

mcp_use/task_managers/sse.py

@@ -66,12 +66,9 @@ class SseConnectionManager(ConnectionManager[tuple[Any, Any]]):
         # Return the streams
         return (read_stream, write_stream)
 
-    async def _close_connection(self, connection: tuple[Any, Any]) -> None:
-        """Close the SSE connection.
+    async def _close_connection(self) -> None:
+        """Close the SSE connection."""
 
-        Args:
-            connection: The connection to close (ignored, we use the context manager)
-        """
         if self._sse_ctx:
             # Exit the context manager
             try:

mcp_use/task_managers/stdio.py

@@ -57,12 +57,8 @@ class StdioConnectionManager(ConnectionManager[tuple[Any, Any]]):
         # 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)
-        """
+    async def _close_connection(self) -> None:
+        """Close the stdio connection."""
         if self._stdio_ctx:
             # Exit the context manager
             try:

mcp_use/task_managers/websocket.py

@@ -4,14 +4,15 @@ WebSocket connection management for MCP implementations.
 This module provides a connection manager for WebSocket-based MCP connections.
 """
 
-import websockets
-from websockets.client import ClientConnection
+from typing import Any
+
+from mcp.client.websocket import websocket_client
 
 from ..logging import logger
 from .base import ConnectionManager
 
 
-class WebSocketConnectionManager(ConnectionManager[ClientConnection]):
+class WebSocketConnectionManager(ConnectionManager[tuple[Any, Any]]):
     """Connection manager for WebSocket-based MCP connections.
 
     This class handles the lifecycle of WebSocket connections, ensuring proper
@@ -21,19 +22,16 @@ class WebSocketConnectionManager(ConnectionManager[ClientConnection]):
     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) -> ClientConnection:
+    async def _establish_connection(self) -> tuple[Any, Any]:
         """Establish a WebSocket connection.
 
         Returns:
@@ -43,21 +41,23 @@ class WebSocketConnectionManager(ConnectionManager[ClientConnection]):
             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: ClientConnection) -> 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}")
+        # Create the context manager
+        self._ws_ctx = websocket_client(self.url)
+
+        # Enter the context manager
+        read_stream, write_stream = await self._ws_ctx.__aenter__()
+
+        # Return the streams
+        return (read_stream, write_stream)
+
+    async def _close_connection(self) -> None:
+        """Close the WebSocket connection."""
+        if self._ws_ctx:
+            # Exit the context manager
+            try:
+                logger.debug("Closing WebSocket connection")
+                await self._ws_ctx.__aexit__(None, None, None)
+            except Exception as e:
+                logger.warning(f"Error closing WebSocket connection: {e}")
+            finally:
+                self._ws_ctx = None

mcp_use/types/clientoptions.py

@@ -0,0 +1,23 @@
+"""
+Options for MCP client configuration.
+
+This module provides data classes and type definitions for configuring the MCP client.
+"""
+
+from typing import NotRequired, TypedDict
+
+from .sandbox import SandboxOptions
+
+
+class ClientOptions(TypedDict):
+    """Options for configuring the MCP client.
+
+    This class encapsulates all configuration options for the MCPClient,
+    making it easier to extend the API without breaking backward compatibility.
+    """
+
+    is_sandboxed: NotRequired[bool]
+    """Whether to use sandboxed execution mode for running MCP servers."""
+
+    sandbox_options: NotRequired[SandboxOptions]
+    """Options for sandbox configuration when is_sandboxed=True."""

mcp_use/types/sandbox.py

@@ -0,0 +1,23 @@
+"""Type definitions for sandbox-related configurations."""
+
+from typing import NotRequired, TypedDict
+
+
+class SandboxOptions(TypedDict):
+    """Configuration options for sandbox execution.
+
+    This type defines the configuration options available when running
+    MCP servers in a sandboxed environment (e.g., using E2B).
+    """
+
+    api_key: str
+    """Direct API key for sandbox provider (e.g., E2B API key).
+    If not provided, will use E2B_API_KEY environment variable."""
+
+    sandbox_template_id: NotRequired[str]
+    """Template ID for the sandbox environment.
+    Default: 'base'"""
+
+    supergateway_command: NotRequired[str]
+    """Command to run supergateway.
+    Default: 'npx -y supergateway'"""