mcp-use__mcp-use/libraries/python/mcp_use/client/connectors/sandbox.py

"""
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 httpx
from mcp import ClientSession
from mcp.client.session import ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
from mcp.types import Root

from mcp_use.client.middleware import CallbackClientSession, Middleware
from mcp_use.client.task_managers import SseConnectionManager
from mcp_use.logging import logger

# 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 typing import NotRequired, TypedDict

from mcp_use.client.connectors.base import BaseConnector


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'"""


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,
        sampling_callback: SamplingFnT | None = None,
        elicitation_callback: ElicitationFnT | None = None,
        message_handler: MessageHandlerFnT | None = None,
        logging_callback: LoggingFnT | None = None,
        middleware: list[Middleware] | None = None,
        roots: list[Root] | None = None,
        list_roots_callback: ListRootsFnT | None = None,
    ):
        """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.
            sampling_callback: Optional sampling callback.
            elicitation_callback: Optional elicitation callback.
            message_handler: Optional callback to handle messages.
            logging_callback: Optional callback to handle log messages.
            middleware: Optional list of middleware.
            roots: Optional initial list of roots to advertise to the server.
            list_roots_callback: Optional custom callback to handle roots/list requests.
        """
        super().__init__(
            sampling_callback=sampling_callback,
            elicitation_callback=elicitation_callback,
            message_handler=message_handler,
            logging_callback=logging_callback,
            middleware=middleware,
            roots=roots,
            list_roots_callback=list_roots_callback,
        )
        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_session: 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 httpx.AsyncClient(timeout=2) as client:
                    try:
                        # First try the endpoint
                        response = await client.get(ping_url)
                        if response.status_code == 200:
                            elapsed = time.time() - start_time
                            logger.info(f"Server is ready! SSE endpoint responded with 200 after {elapsed:.1f}s")
                            return True
                    except Exception:
                        # If sse endpoint doesn't work, try the base URL
                        response = await client.get(base_url)
                        if response.status_code < 500:  # Accept any non-server error
                            elapsed = time.time() - start_time
                            logger.info(
                                f"Server is ready! Base URL responded with {response.status_code} 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
            raw_client_session = ClientSession(
                read_stream,
                write_stream,
                sampling_callback=self.sampling_callback,
                elicitation_callback=self.elicitation_callback,
                list_roots_callback=self.list_roots_callback,
                message_handler=self._internal_message_handler,
                logging_callback=self.logging_callback,
                client_info=self.client_info,
            )
            await raw_client_session.__aenter__()

            # Wrap with middleware
            self.client_session = CallbackClientSession(
                raw_client_session, self.public_identifier, self.middleware_manager
            )

            # 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."""
        args_str = " ".join(self.user_args) if self.user_args else ""
        return f"sandbox:{self.user_command} {args_str}".strip()