mcp-use 1.2.13__py3-none-any.whl → 1.3.1__py3-none-any.whl

mcp_use/connectors/sandbox.py

@@ -0,0 +1,296 @@
+"""
+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 'sandbox_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")
+
+    @property
+    def public_identifier(self) -> str:
+        """Get the identifier for the connector."""
+        return {"type": "sandbox", "command": self.user_command, "args": self.user_args}

mcp_use/connectors/stdio.py

@@ -61,8 +61,8 @@ class StdioConnector(BaseConnector):
             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__()
+            self.client_session = ClientSession(read_stream, write_stream, sampling_callback=None)
+            await self.client_session.__aenter__()
 
             # Mark as connected
             self._connected = True
@@ -76,3 +76,8 @@ class StdioConnector(BaseConnector):
 
             # Re-raise the original exception
             raise
+
+    @property
+    def public_identifier(self) -> str:
+        """Get the identifier for the connector."""
+        return {"type": "stdio", "command&args": f"{self.command} {' '.join(self.args)}"}

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/connectors/websocket.py

@@ -11,7 +11,7 @@ import uuid
 from typing import Any
 
 from mcp.types import Tool
-from websockets.client import WebSocketClientProtocol
+from websockets import ClientConnection
 
 from ..logging import logger
 from ..task_managers import ConnectionManager, WebSocketConnectionManager
@@ -44,7 +44,7 @@ class WebSocketConnector(BaseConnector):
         if auth_token:
             self.headers["Authorization"] = f"Bearer {auth_token}"
 
-        self.ws: WebSocketClientProtocol | None = None
+        self.ws: ClientConnection | None = None
         self._connection_manager: ConnectionManager | None = None
         self._receiver_task: asyncio.Task | None = None
         self.pending_requests: dict[str, asyncio.Future] = {}
@@ -243,3 +243,8 @@ class WebSocketConnector(BaseConnector):
         """Send a raw request to the MCP implementation."""
         logger.debug(f"Sending request: {method} with params: {params}")
         return await self._send_request(method, params)
+
+    @property
+    def public_identifier(self) -> str:
+        """Get the identifier for the connector."""
+        return {"type": "websocket", "url": self.url}

mcp_use/session.py

@@ -7,8 +7,6 @@ which handles authentication, initialization, and tool discovery.
 
 from typing import Any
 
-from mcp.types import Tool
-
 from .connectors.base import BaseConnector
 
 
@@ -32,7 +30,6 @@ class MCPSession:
         """
         self.connector = connector
         self.session_info: dict[str, Any] | None = None
-        self.tools: list[Tool] = []
         self.auto_connect = auto_connect
 
     async def __aenter__(self) -> "MCPSession":
@@ -84,4 +81,4 @@ class MCPSession:
         Returns:
             True if the connector is connected, False otherwise.
         """
-        return hasattr(self.connector, "client") and self.connector.client is not None
+        return self.connector.is_connected

mcp_use/task_managers/__init__.py

@@ -8,12 +8,13 @@ through different transport mechanisms.
 from .base import ConnectionManager
 from .sse import SseConnectionManager
 from .stdio import StdioConnectionManager
+from .streamable_http import StreamableHttpConnectionManager
 from .websocket import WebSocketConnectionManager
 
 __all__ = [
     "ConnectionManager",
-    "HttpConnectionManager",
     "StdioConnectionManager",
     "WebSocketConnectionManager",
     "SseConnectionManager",
+    "StreamableHttpConnectionManager",
 ]

mcp_use/task_managers/base.py

@@ -70,9 +70,7 @@ class ConnectionManager(Generic[T], ABC):
         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"
-        )
+        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()
@@ -105,6 +103,14 @@ class ConnectionManager(Generic[T], ABC):
         await self._done_event.wait()
         logger.debug(f"{self.__class__.__name__} task completed")
 
+    def get_streams(self) -> T | None:
+        """Get the current connection streams.
+
+        Returns:
+            The current connection (typically a tuple of read_stream, write_stream) or None if not connected.
+        """
+        return self._connection
+
     async def _connection_task(self) -> None:
         """Run the connection task.
 
@@ -137,7 +143,7 @@ class ConnectionManager(Generic[T], ABC):
             self._ready_event.set()
 
         finally:
-            # Close the connection if it was establishedSUPABASE_URL
+            # Close the connection if it was established
             if self._connection is not None:
                 try:
                     await self._close_connection()

mcp_use/task_managers/streamable_http.py

@@ -0,0 +1,81 @@
+"""
+Streamable HTTP connection management for MCP implementations.
+
+This module provides a connection manager for streamable HTTP-based MCP connections
+that ensures proper task isolation and resource cleanup.
+"""
+
+from datetime import timedelta
+from typing import Any
+
+from mcp.client.streamable_http import streamablehttp_client
+
+from ..logging import logger
+from .base import ConnectionManager
+
+
+class StreamableHttpConnectionManager(ConnectionManager[tuple[Any, Any]]):
+    """Connection manager for streamable HTTP-based MCP connections.
+
+    This class handles the proper task isolation for HTTP streaming connections
+    to prevent the "cancel scope in different task" error. It runs the http_stream_client
+    in a dedicated task and manages its lifecycle.
+    """
+
+    def __init__(
+        self,
+        url: str,
+        headers: dict[str, str] | None = None,
+        timeout: float = 5,
+        read_timeout: float = 60 * 5,
+    ):
+        """Initialize a new streamable HTTP connection manager.
+
+        Args:
+            url: The HTTP endpoint URL
+            headers: Optional HTTP headers
+            timeout: Timeout for HTTP operations in seconds
+            read_timeout: Timeout for HTTP read operations in seconds
+        """
+        super().__init__()
+        self.url = url
+        self.headers = headers or {}
+        self.timeout = timedelta(seconds=timeout)
+        self.read_timeout = timedelta(seconds=read_timeout)
+        self._http_ctx = None
+
+    async def _establish_connection(self) -> tuple[Any, Any]:
+        """Establish a streamable HTTP connection.
+
+        Returns:
+            A tuple of (read_stream, write_stream)
+
+        Raises:
+            Exception: If connection cannot be established.
+        """
+        # Create the context manager
+        self._http_ctx = streamablehttp_client(
+            url=self.url,
+            headers=self.headers,
+            timeout=self.timeout,
+            sse_read_timeout=self.read_timeout,
+        )
+
+        # Enter the context manager. Ignoring the session id callback
+        read_stream, write_stream, _ = await self._http_ctx.__aenter__()
+
+        # Return the streams
+        return (read_stream, write_stream)
+
+    async def _close_connection(self) -> None:
+        """Close the streamable HTTP connection."""
+
+        if self._http_ctx:
+            # Exit the context manager
+            try:
+                await self._http_ctx.__aexit__(None, None, None)
+            except Exception as e:
+                # Only log if it's not a normal connection termination
+                logger.debug(f"Streamable HTTP context cleanup: {e}")
+            finally:
+                self._http_ctx = None

mcp_use/task_managers/websocket.py

@@ -22,14 +22,17 @@ class WebSocketConnectionManager(ConnectionManager[tuple[Any, Any]]):
     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 HTTP headers
         """
         super().__init__()
         self.url = url
+        self.headers = headers or {}
 
     async def _establish_connection(self) -> tuple[Any, Any]:
         """Establish a WebSocket connection.
@@ -42,6 +45,8 @@ class WebSocketConnectionManager(ConnectionManager[tuple[Any, Any]]):
         """
         logger.debug(f"Connecting to WebSocket: {self.url}")
         # Create the context manager
+        # Note: The current MCP websocket_client implementation doesn't support headers
+        # If headers need to be passed, this would need to be updated when MCP supports it
         self._ws_ctx = websocket_client(self.url)
 
         # Enter the context manager

mcp_use/telemetry/events.py

@@ -0,0 +1,93 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+
+
+class BaseTelemetryEvent(ABC):
+    """Base class for all telemetry events"""
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Event name for tracking"""
+        pass
+
+    @property
+    @abstractmethod
+    def properties(self) -> dict[str, Any]:
+        """Event properties to send with the event"""
+        pass
+
+
+@dataclass
+class MCPAgentExecutionEvent(BaseTelemetryEvent):
+    """Comprehensive event for tracking complete MCP agent execution"""
+
+    # Execution method and context
+    execution_method: str  # "run" or "astream"
+    query: str  # The actual user query
+    success: bool
+
+    # Agent configuration
+    model_provider: str
+    model_name: str
+    server_count: int
+    server_identifiers: list[dict[str, str]]
+    total_tools_available: int
+    tools_available_names: list[str]
+    max_steps_configured: int
+    memory_enabled: bool
+    use_server_manager: bool
+
+    # Execution PARAMETERS
+    max_steps_used: int | None
+    manage_connector: bool
+    external_history_used: bool
+
+    # Execution results
+    steps_taken: int | None = None
+    tools_used_count: int | None = None
+    tools_used_names: list[str] | None = None
+    response: str | None = None  # The actual response
+    execution_time_ms: int | None = None
+    error_type: str | None = None
+
+    # Context
+    conversation_history_length: int | None = None
+
+    @property
+    def name(self) -> str:
+        return "mcp_agent_execution"
+
+    @property
+    def properties(self) -> dict[str, Any]:
+        return {
+            # Core execution info
+            "execution_method": self.execution_method,
+            "query": self.query,
+            "query_length": len(self.query),
+            "success": self.success,
+            # Agent configuration
+            "model_provider": self.model_provider,
+            "model_name": self.model_name,
+            "server_count": self.server_count,
+            "server_identifiers": self.server_identifiers,
+            "total_tools_available": self.total_tools_available,
+            "tools_available_names": self.tools_available_names,
+            "max_steps_configured": self.max_steps_configured,
+            "memory_enabled": self.memory_enabled,
+            "use_server_manager": self.use_server_manager,
+            # Execution parameters (always include, even if None)
+            "max_steps_used": self.max_steps_used,
+            "manage_connector": self.manage_connector,
+            "external_history_used": self.external_history_used,
+            # Execution results (always include, even if None)
+            "steps_taken": self.steps_taken,
+            "tools_used_count": self.tools_used_count,
+            "tools_used_names": self.tools_used_names,
+            "response": self.response,
+            "response_length": len(self.response) if self.response else None,
+            "execution_time_ms": self.execution_time_ms,
+            "error_type": self.error_type,
+            "conversation_history_length": self.conversation_history_length,
+        }