mcp-use 1.3.12__py3-none-any.whl → 1.4.0__py3-none-any.whl

mcp_use/client/auth/oauth_callback.py

@@ -0,0 +1,215 @@
+"""OAuth callback server implementation."""
+
+import asyncio
+import html
+from dataclasses import dataclass
+
+import anyio
+import uvicorn
+from starlette.applications import Starlette
+from starlette.requests import Request
+from starlette.responses import HTMLResponse
+from starlette.routing import Route
+
+from mcp_use.logging import logger
+
+
+@dataclass
+class CallbackResponse:
+    """Response data from OAuth callback."""
+
+    code: str | None = None  # Authorization code (success)
+    state: str | None = None  # CSRF protection token
+    error: str | None = None  # Errors code (if failed)
+    error_description: str | None = None
+    error_uri: str | None = None
+
+
+class OAuthCallbackServer:
+    """Local server to handle OAuth callback."""
+
+    def __init__(self, port: int):
+        """Initialize the callback server.
+
+        Args:
+            port: Port to listen on.
+        """
+        self.port = port
+        self.redirect_uri: str | None = None
+        # Thread safe way to pass callback data to the main OAuth flow
+        self.response_queue: asyncio.Queue[CallbackResponse] = asyncio.Queue(maxsize=1)
+        self.server: uvicorn.Server | None = None
+        self._shutdown_event = anyio.Event()
+
+    async def start(self) -> str:
+        """Start the callback server and return the redirect URI."""
+        app = self._create_app()
+
+        # Create the server
+        config = uvicorn.Config(
+            app,
+            host="127.0.0.1",
+            port=self.port,
+            log_level="error",  # Suppress uvicorn logs
+        )
+        self.server = uvicorn.Server(config)
+
+        # Start server in background
+        self._server_task = asyncio.create_task(self.server.serve())
+
+        # Wait a moment for server to start
+        await asyncio.sleep(0.1)
+
+        self.redirect_uri = f"http://localhost:{self.port}/callback"
+        return self.redirect_uri
+
+    async def wait_for_code(self, timeout: float = 300) -> CallbackResponse:
+        """Wait for the OAuth callback with a timeout (default 5 minutes)."""
+        try:
+            response = await asyncio.wait_for(self.response_queue.get(), timeout=timeout)
+            return response
+        except TimeoutError:
+            raise TimeoutError(f"OAuth callback not received within {timeout} seconds") from None
+        finally:
+            await self.shutdown()
+
+    async def shutdown(self):
+        """Shutdown the callback server."""
+        self._shutdown_event.set()
+        if self.server:
+            self.server.should_exit = True
+            if hasattr(self, "_server_task"):
+                try:
+                    await asyncio.wait_for(self._server_task, timeout=5.0)
+                except TimeoutError:
+                    self._server_task.cancel()
+
+    def _create_app(self) -> Starlette:
+        """Create the Starlette application."""
+
+        async def callback(request: Request) -> HTMLResponse:
+            """Handle the OAuth callback."""
+            params = request.query_params
+
+            # Extract OAuth parameters
+            response = CallbackResponse(
+                code=params.get("code"),
+                state=params.get("state"),
+                error=params.get("error"),
+                error_description=params.get("error_description"),
+                error_uri=params.get("error_uri"),
+            )
+
+            # Log the callback response
+            logger.debug(
+                f"OAuth callback received: error={response.error}, error_description={response.error_description}"
+            )
+            if response.code:
+                logger.debug("OAuth callback received authorization code")
+            else:
+                logger.error(f"OAuth callback error: {response.error} - {response.error_description}")
+
+            # Put response in queue
+            try:
+                self.response_queue.put_nowait(response)
+            except asyncio.QueueFull:
+                pass  # Ignore if queue is already full
+
+            # Return success page
+            if response.code:
+                html = self._success_html()
+            else:
+                html = self._error_html(response.error, response.error_description)
+
+            return HTMLResponse(content=html)
+
+        routes = [Route("/callback", callback)]
+        return Starlette(routes=routes)
+
+    def _success_html(self) -> str:
+        """HTML response for successful authorization."""
+        return """
+        <!DOCTYPE html>
+        <html>
+        <head>
+            <title>Authorization Successful</title>
+            <style>
+                body {
+                    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
+                    display: flex;
+                    justify-content: center;
+                    align-items: center;
+                    height: 100vh;
+                    margin: 0;
+                    background-color: #f5f5f5;
+                }
+                .container {
+                    text-align: center;
+                    padding: 2rem;
+                    background: white;
+                    border-radius: 8px;
+                    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+                }
+                h1 { color: #22c55e; margin-bottom: 0.5rem; }
+                p { color: #666; margin-top: 0.5rem; }
+                .icon { font-size: 48px; margin-bottom: 1rem; }
+            </style>
+        </head>
+        <body>
+            <div class="container">
+                <div class="icon">✅</div>
+                <h1>Authorization Successful!</h1>
+                <p>You can now close this window and return to your application.</p>
+            </div>
+            <script>
+                // Auto-close after 3 seconds
+                setTimeout(() => window.close(), 3000);
+            </script>
+        </body>
+        </html>
+        """
+
+    def _error_html(self, error: str | None, description: str | None) -> str:
+        """HTML response for authorization error."""
+        error_msg = html.escape(error or "Unknown error")
+        desc_msg = html.escape(description or "Authorization was not completed successfully.")
+
+        return f"""
+        <!DOCTYPE html>
+        <html>
+        <head>
+            <title>Authorization Error</title>
+            <style>
+                body {{
+                    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
+                    display: flex;
+                    justify-content: center;
+                    align-items: center;
+                    height: 100vh;
+                    margin: 0;
+                    background-color: #f5f5f5;
+                }}
+                .container {{
+                    text-align: center;
+                    padding: 2rem;
+                    background: white;
+                    border-radius: 8px;
+                    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+                    max-width: 500px;
+                }}
+                h1 {{ color: #ef4444; margin-bottom: 0.5rem; }}
+                .error {{ color: #dc2626; font-weight: 600; margin: 1rem 0; }}
+                .description {{ color: #666; margin-top: 0.5rem; }}
+                .icon {{ font-size: 48px; margin-bottom: 1rem; }}
+            </style>
+        </head>
+        <body>
+            <div class="container">
+                <div class="icon">❌</div>
+                <h1>Authorization Error</h1>
+                <p class="error">{error_msg}</p>
+                <p class="description">{desc_msg}</p>
+            </div>
+        </body>
+        </html>
+        """

mcp_use/client/client.py

@@ -0,0 +1,356 @@
+"""
+Client for managing MCP servers and sessions.
+
+This module provides a high-level client that manages MCP servers, connectors,
+and sessions from configuration.
+"""
+
+import json
+import warnings
+from typing import Any
+
+from mcp.client.session import ElicitationFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
+
+from mcp_use.client.config import create_connector_from_config, load_config_file
+from mcp_use.client.connectors.sandbox import SandboxOptions
+from mcp_use.client.middleware import Middleware, default_logging_middleware
+from mcp_use.client.session import MCPSession
+from mcp_use.logging import logger
+from mcp_use.telemetry.telemetry import telemetry
+
+
+class MCPClient:
+    """Client for managing MCP servers and sessions.
+
+    This class provides a unified interface for working with MCP servers,
+    handling configuration, connector creation, and session management.
+    """
+
+    def __init__(
+        self,
+        config: str | dict[str, Any] | None = None,
+        allowed_servers: list[str] | None = None,
+        sandbox: bool = False,
+        sandbox_options: SandboxOptions | None = None,
+        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,
+    ) -> None:
+        """Initialize a new MCP client.
+
+        Args:
+            config: Either a dict containing configuration or a path to a JSON config file.
+                   If None, an empty configuration is used.
+            sandbox: Whether to use sandboxed execution mode for running MCP servers.
+            sandbox_options: Optional sandbox configuration options.
+            sampling_callback: Optional sampling callback function.
+        """
+        self.config: dict[str, Any] = {}
+        self.allowed_servers: list[str] = allowed_servers
+        self.sandbox = sandbox
+        self.sandbox_options = sandbox_options
+        self.sessions: dict[str, MCPSession] = {}
+        self.active_sessions: list[str] = []
+        self.sampling_callback = sampling_callback
+        self.elicitation_callback = elicitation_callback
+        self.message_handler = message_handler
+        self.logging_callback = logging_callback
+        # Add default logging middleware if no middleware provided, or prepend it to existing middleware
+        default_middleware = [default_logging_middleware]
+        if middleware:
+            self.middleware = default_middleware + middleware
+        else:
+            self.middleware = default_middleware
+        # Load configuration if provided
+        if config is not None:
+            if isinstance(config, str):
+                self.config = load_config_file(config)
+            else:
+                self.config = config
+
+    @classmethod
+    def from_dict(
+        cls,
+        config: dict[str, Any],
+        sandbox: bool = False,
+        sandbox_options: SandboxOptions | None = None,
+        sampling_callback: SamplingFnT | None = None,
+        elicitation_callback: ElicitationFnT | None = None,
+        message_handler: MessageHandlerFnT | None = None,
+        logging_callback: LoggingFnT | None = None,
+    ) -> "MCPClient":
+        """Create a MCPClient from a dictionary.
+
+        Args:
+            config: The configuration dictionary.
+            sandbox: Whether to use sandboxed execution mode for running MCP servers.
+            sandbox_options: Optional sandbox configuration options.
+            sampling_callback: Optional sampling callback function.
+            elicitation_callback: Optional elicitation callback function.
+        """
+        return cls(
+            config=config,
+            sandbox=sandbox,
+            sandbox_options=sandbox_options,
+            sampling_callback=sampling_callback,
+            elicitation_callback=elicitation_callback,
+            message_handler=message_handler,
+            logging_callback=logging_callback,
+        )
+
+    @classmethod
+    def from_config_file(
+        cls,
+        filepath: str,
+        sandbox: bool = False,
+        sandbox_options: SandboxOptions | None = None,
+        sampling_callback: SamplingFnT | None = None,
+        elicitation_callback: ElicitationFnT | None = None,
+        message_handler: MessageHandlerFnT | None = None,
+        logging_callback: LoggingFnT | None = None,
+    ) -> "MCPClient":
+        """Create a MCPClient from a configuration file.
+
+        Args:
+            filepath: The path to the configuration file.
+            sandbox: Whether to use sandboxed execution mode for running MCP servers.
+            sandbox_options: Optional sandbox configuration options.
+            sampling_callback: Optional sampling callback function.
+            elicitation_callback: Optional elicitation callback function.
+        """
+        return cls(
+            config=load_config_file(filepath),
+            sandbox=sandbox,
+            sandbox_options=sandbox_options,
+            sampling_callback=sampling_callback,
+            elicitation_callback=elicitation_callback,
+            message_handler=message_handler,
+            logging_callback=logging_callback,
+        )
+
+    @telemetry("client_add_server")
+    def add_server(
+        self,
+        name: str,
+        server_config: dict[str, Any],
+    ) -> None:
+        """Add a server configuration.
+
+        Args:
+            name: The name to identify this server.
+            server_config: The server configuration.
+        """
+        if "mcpServers" not in self.config:
+            self.config["mcpServers"] = {}
+
+        self.config["mcpServers"][name] = server_config
+
+    @telemetry("client_remove_server")
+    def remove_server(self, name: str) -> None:
+        """Remove a server configuration.
+
+        Args:
+            name: The name of the server to remove.
+        """
+        if "mcpServers" in self.config and name in self.config["mcpServers"]:
+            del self.config["mcpServers"][name]
+
+            # If we removed an active session, remove it from active_sessions
+            if name in self.active_sessions:
+                self.active_sessions.remove(name)
+
+    def add_middleware(self, middleware: Middleware) -> None:
+        """Add a middleware.
+
+        Args:
+            middleware: The middleware to add
+        """
+        if len(self.sessions) == 0 and middleware not in self.middleware:
+            self.middleware.append(middleware)
+            return
+
+        if middleware not in self.middleware:
+            self.middleware.append(middleware)
+            for session in self.sessions.values():
+                session.connector.middleware_manager.add_middleware(middleware)
+
+    def get_server_names(self) -> list[str]:
+        """Get the list of configured server names.
+
+        Returns:
+            List of server names.
+        """
+        return list(self.config.get("mcpServers", {}).keys())
+
+    def save_config(self, filepath: str) -> None:
+        """Save the current configuration to a file.
+
+        Args:
+            filepath: The path to save the configuration to.
+        """
+        with open(filepath, "w") as f:
+            json.dump(self.config, f, indent=2)
+
+    @telemetry("client_create_session")
+    async def create_session(self, server_name: str, auto_initialize: bool = True) -> MCPSession:
+        """Create a session for the specified server.
+
+        Args:
+            server_name: The name of the server to create a session for.
+            auto_initialize: Whether to automatically initialize the session.
+
+        Returns:
+            The created MCPSession.
+
+        Raises:
+            ValueError: If the specified server doesn't exist.
+        """
+        # Get server config
+        servers = self.config.get("mcpServers", {})
+        if not servers:
+            warnings.warn("No MCP servers defined in config", UserWarning, stacklevel=2)
+            return None
+
+        if server_name not in servers:
+            raise ValueError(f"Server '{server_name}' not found in config")
+
+        server_config = servers[server_name]
+
+        # Create connector with options and client-level auth
+        connector = create_connector_from_config(
+            server_config,
+            sandbox=self.sandbox,
+            sandbox_options=self.sandbox_options,
+            sampling_callback=self.sampling_callback,
+            elicitation_callback=self.elicitation_callback,
+            message_handler=self.message_handler,
+            logging_callback=self.logging_callback,
+            middleware=self.middleware,
+        )
+
+        # Create the session
+        session = MCPSession(connector)
+        if auto_initialize:
+            await session.initialize()
+        self.sessions[server_name] = session
+
+        # Add to active sessions
+        if server_name not in self.active_sessions:
+            self.active_sessions.append(server_name)
+
+        return session
+
+    @telemetry("client_create_all_sessions")
+    async def create_all_sessions(
+        self,
+        auto_initialize: bool = True,
+    ) -> dict[str, MCPSession]:
+        """Create sessions for all configured servers.
+
+        Args:
+            auto_initialize: Whether to automatically initialize the sessions.
+
+        Returns:
+            Dictionary mapping server names to their MCPSession instances.
+
+        Warns:
+            UserWarning: If no servers are configured.
+        """
+        # Get server config
+        servers = self.config.get("mcpServers", {})
+        if not servers:
+            warnings.warn("No MCP servers defined in config", UserWarning, stacklevel=2)
+            return {}
+
+        # Create sessions only for allowed servers if applicable else create for all servers
+        for name in servers:
+            if self.allowed_servers is None or name in self.allowed_servers:
+                await self.create_session(name, auto_initialize)
+
+        return self.sessions
+
+    def get_session(self, server_name: str) -> MCPSession:
+        """Get an existing session.
+
+        Args:
+            server_name: The name of the server to get the session for.
+                        If None, uses the first active session.
+
+        Returns:
+            The MCPSession for the specified server.
+
+        Raises:
+            ValueError: If no active sessions exist or the specified session doesn't exist.
+        """
+        if server_name not in self.sessions:
+            raise ValueError(f"No session exists for server '{server_name}'")
+
+        return self.sessions[server_name]
+
+    def get_all_active_sessions(self) -> dict[str, MCPSession]:
+        """Get all active sessions.
+
+        Returns:
+            Dictionary mapping server names to their MCPSession instances.
+        """
+        return {name: self.sessions[name] for name in self.active_sessions if name in self.sessions}
+
+    @telemetry("client_close_session")
+    async def close_session(self, server_name: str) -> None:
+        """Close a session.
+
+        Args:
+            server_name: The name of the server to close the session for.
+                        If None, uses the first active session.
+
+        Raises:
+            ValueError: If no active sessions exist or the specified session doesn't exist.
+        """
+        # Check if the session exists
+        if server_name not in self.sessions:
+            logger.warning(f"No session exists for server '{server_name}', nothing to close")
+            return
+
+        # Get the session
+        session = self.sessions[server_name]
+
+        try:
+            # Disconnect from the session
+            logger.debug(f"Closing session for server '{server_name}'")
+            await session.disconnect()
+        except Exception as e:
+            logger.error(f"Error closing session for server '{server_name}': {e}")
+        finally:
+            # Remove the session regardless of whether disconnect succeeded
+            del self.sessions[server_name]
+
+            # Remove from active_sessions
+            if server_name in self.active_sessions:
+                self.active_sessions.remove(server_name)
+
+    @telemetry("client_close_all_sessions")
+    async def close_all_sessions(self) -> None:
+        """Close all active sessions.
+
+        This method ensures all sessions are closed even if some fail.
+        """
+        # Get a list of all session names first to avoid modification during iteration
+        server_names = list(self.sessions.keys())
+        errors = []
+
+        for server_name in server_names:
+            try:
+                logger.debug(f"Closing session for server '{server_name}'")
+                await self.close_session(server_name)
+            except Exception as e:
+                error_msg = f"Failed to close session for server '{server_name}': {e}"
+                logger.error(error_msg)
+                errors.append(error_msg)
+
+        # Log summary if there were errors
+        if errors:
+            logger.error(f"Encountered {len(errors)} errors while closing sessions")
+        else:
+            logger.debug("All sessions closed successfully")

mcp_use/client/config.py

@@ -0,0 +1,106 @@
+"""
+Configuration loader for MCP session.
+
+This module provides functionality to load MCP configuration from JSON files.
+"""
+
+import json
+from typing import Any
+
+from mcp.client.session import ElicitationFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
+
+from mcp_use.client.connectors.base import BaseConnector
+from mcp_use.client.connectors.http import HttpConnector
+from mcp_use.client.connectors.sandbox import SandboxConnector, SandboxOptions
+from mcp_use.client.connectors.stdio import StdioConnector
+from mcp_use.client.connectors.utils import is_stdio_server
+from mcp_use.client.connectors.websocket import WebSocketConnector
+from mcp_use.client.middleware import Middleware
+
+
+def load_config_file(filepath: str) -> dict[str, Any]:
+    """Load a configuration file.
+
+    Args:
+        filepath: Path to the configuration file
+
+    Returns:
+        The parsed configuration
+    """
+    with open(filepath) as f:
+        return json.load(f)
+
+
+def create_connector_from_config(
+    server_config: dict[str, Any],
+    sandbox: bool = False,
+    sandbox_options: SandboxOptions | None = None,
+    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,
+) -> BaseConnector:
+    """Create a connector based on server configuration.
+    This function can be called with just the server_config parameter:
+    create_connector_from_config(server_config)
+    Args:
+        server_config: The server configuration section
+        sandbox: Whether to use sandboxed execution mode for running MCP servers.
+        sandbox_options: Optional sandbox configuration options.
+        sampling_callback: Optional sampling callback function.
+    Returns:
+        A configured connector instance
+    """
+
+    # Stdio connector (command-based)
+    if is_stdio_server(server_config) and not sandbox:
+        return StdioConnector(
+            command=server_config["command"],
+            args=server_config["args"],
+            env=server_config.get("env", None),
+            sampling_callback=sampling_callback,
+            elicitation_callback=elicitation_callback,
+            message_handler=message_handler,
+            logging_callback=logging_callback,
+            middleware=middleware,
+        )
+
+    # Sandboxed connector
+    elif is_stdio_server(server_config) and sandbox:
+        return SandboxConnector(
+            command=server_config["command"],
+            args=server_config["args"],
+            env=server_config.get("env", None),
+            e2b_options=sandbox_options,
+            sampling_callback=sampling_callback,
+            elicitation_callback=elicitation_callback,
+            message_handler=message_handler,
+            logging_callback=logging_callback,
+            middleware=middleware,
+        )
+
+    # HTTP connector
+    elif "url" in server_config:
+        return HttpConnector(
+            base_url=server_config["url"],
+            headers=server_config.get("headers", None),
+            auth=server_config.get("auth", {}),
+            timeout=server_config.get("timeout", 5),
+            sse_read_timeout=server_config.get("sse_read_timeout", 60 * 5),
+            sampling_callback=sampling_callback,
+            elicitation_callback=elicitation_callback,
+            message_handler=message_handler,
+            logging_callback=logging_callback,
+            middleware=middleware,
+        )
+
+    # WebSocket connector
+    elif "ws_url" in server_config:
+        return WebSocketConnector(
+            url=server_config["ws_url"],
+            headers=server_config.get("headers", None),
+            auth=server_config.get("auth", {}),
+        )
+
+    raise ValueError("Cannot determine connector type from config")

mcp_use/client/connectors/__init__.py

@@ -0,0 +1,20 @@
+"""
+Connectors for various MCP transports.
+
+This module provides interfaces for connecting to MCP implementations
+through different transport mechanisms.
+"""
+
+from .base import BaseConnector  # noqa: F401
+from .http import HttpConnector  # noqa: F401
+from .sandbox import SandboxConnector  # noqa: F401
+from .stdio import StdioConnector  # noqa: F401
+from .websocket import WebSocketConnector  # noqa: F401
+
+__all__ = [
+    "BaseConnector",
+    "StdioConnector",
+    "HttpConnector",
+    "WebSocketConnector",
+    "SandboxConnector",
+]