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

mcp_use/client.py

@@ -9,8 +9,8 @@ import json
 from typing import Any
 
 from .config import create_connector_from_config, load_config_file
+from .logging import logger
 from .session import MCPSession
-from .tools.converter import ModelProvider
 
 
 class MCPClient:
@@ -23,19 +23,16 @@ class MCPClient:
     def __init__(
         self,
         config: str | dict[str, Any] | None = None,
-        model_provider: str | ModelProvider = "openai",
     ) -> 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.
-            model_provider: The model provider to use for tool conversion.
         """
-        self.model_provider = model_provider
         self.config: dict[str, Any] = {}
         self.sessions: dict[str, MCPSession] = {}
-        self.active_session: str | None = None
+        self.active_sessions: list[str] = []
 
         # Load configuration if provided
         if config is not None:
@@ -87,9 +84,9 @@ class MCPClient:
         if "mcpServers" in self.config and name in self.config["mcpServers"]:
             del self.config["mcpServers"][name]
 
-            # If we removed the active session, set active_session to None
-            if name == self.active_session:
-                self.active_session = None
+            # If we removed an active session, remove it from active_sessions
+            if name in self.active_sessions:
+                self.active_sessions.remove(name)
 
     def get_server_names(self) -> list[str]:
         """Get the list of configured server names.
@@ -108,17 +105,11 @@ class MCPClient:
         with open(filepath, "w") as f:
             json.dump(self.config, f, indent=2)
 
-    async def create_session(
-        self,
-        server_name: str | None = None,
-        auto_initialize: bool = True,
-    ) -> MCPSession:
+    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.
-                        If None, uses the first available server.
-            auto_initialize: Whether to automatically initialize the session.
 
         Returns:
             The created MCPSession.
@@ -131,10 +122,6 @@ class MCPClient:
         if not servers:
             raise ValueError("No MCP servers defined in config")
 
-        # If server_name not specified, use the first one
-        if not server_name:
-            server_name = next(iter(servers.keys()))
-
         if server_name not in servers:
             raise ValueError(f"Server '{server_name}' not found in config")
 
@@ -142,85 +129,123 @@ class MCPClient:
         connector = create_connector_from_config(server_config)
 
         # Create the session
-        session = MCPSession(connector, self.model_provider)
-        self.sessions[server_name] = session
-
-        # Make this the active session
-        self.active_session = server_name
-
-        # Initialize if requested
+        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
 
-    def get_session(self, server_name: str | None = None) -> MCPSession:
+    async def create_all_sessions(
+        self,
+        auto_initialize: bool = True,
+    ) -> dict[str, MCPSession]:
+        """Create a session for the specified server.
+
+        Args:
+            auto_initialize: Whether to automatically initialize the session.
+
+        Returns:
+            The created MCPSession. If server_name is None, returns the first created session.
+
+        Raises:
+            ValueError: If no servers are configured or the specified server doesn't exist.
+        """
+        # Get server config
+        servers = self.config.get("mcpServers", {})
+        if not servers:
+            raise ValueError("No MCP servers defined in config")
+
+        # Create sessions for all servers
+        for name in servers:
+            session = await self.create_session(name, auto_initialize)
+            if auto_initialize:
+                await session.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 active session.
+                        If None, uses the first active session.
 
         Returns:
             The MCPSession for the specified server.
 
         Raises:
-            ValueError: If no active session exists or the specified session doesn't exist.
+            ValueError: If no active sessions exist or the specified session doesn't exist.
         """
-        if server_name is None:
-            if self.active_session is None:
-                raise ValueError("No active session")
-            server_name = self.active_session
-
         if server_name not in self.sessions:
             raise ValueError(f"No session exists for server '{server_name}'")
 
         return self.sessions[server_name]
 
-    async def close_session(self, server_name: str | None = None) -> None:
+    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}
+
+    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 active session.
+                        If None, uses the first active session.
 
         Raises:
-            ValueError: If no active session exists or the specified session doesn't exist.
+            ValueError: If no active sessions exist or the specified session doesn't exist.
         """
-        session = self.get_session(server_name)
-        await session.disconnect()
-
-        # Remove the session
-        if server_name is None:
-            server_name = self.active_session
-
-        if server_name in self.sessions:
+        # 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.info(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]
 
-        # If we closed the active session, set active_session to None
-        if server_name == self.active_session:
-            self.active_session = None
+            # Remove from active_sessions
+            if server_name in self.active_sessions:
+                self.active_sessions.remove(server_name)
 
     async def close_all_sessions(self) -> None:
-        """Close all active sessions."""
-        for server_name in list(self.sessions.keys()):
-            await self.close_session(server_name)
-
-    async def __aenter__(self) -> "MCPClient":
-        """Enter the async context manager.
-
-        Creates a session for the first available server if no sessions exist.
-
-        Returns:
-            The client instance.
-        """
-        if not self.sessions and self.config.get("mcpServers"):
-            await self.create_session()
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
-        """Exit the async context manager.
+        """Close all active sessions.
 
-        Closes all active sessions.
+        This method ensures all sessions are closed even if some fail.
         """
-        await self.close_all_sessions()
+        # 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.info(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.info("All sessions closed successfully")

mcp_use/config.py

@@ -5,28 +5,9 @@ This module provides functionality to load MCP configuration from JSON files.
 """
 
 import json
-import os
-import subprocess
 from typing import Any
 
 from .connectors import BaseConnector, HttpConnector, StdioConnector, WebSocketConnector
-from .session import MCPSession
-from .tools.converter import ModelProvider
-
-
-def _execute_command(command: str, args: list[str], env: dict[str, str] | None = None) -> None:
-    """Execute a command with given arguments and environment variables.
-
-    Args:
-        command: The command to execute
-        args: List of command arguments
-        env: Optional environment variables to set
-    """
-    full_env = os.environ.copy()
-    if env:
-        full_env.update(env)
-
-    subprocess.Popen([command] + args, env=full_env)
 
 
 def load_config_file(filepath: str) -> dict[str, Any]:
@@ -62,9 +43,9 @@ def create_connector_from_config(server_config: dict[str, Any]) -> BaseConnector
     # HTTP connector
     elif "url" in server_config:
         return HttpConnector(
-            url=server_config["url"],
+            base_url=server_config["url"],
             headers=server_config.get("headers", None),
-            auth=server_config.get("auth", None),
+            auth_token=server_config.get("auth_token", None),
         )
 
     # WebSocket connector
@@ -72,42 +53,7 @@ def create_connector_from_config(server_config: dict[str, Any]) -> BaseConnector
         return WebSocketConnector(
             url=server_config["ws_url"],
             headers=server_config.get("headers", None),
-            auth=server_config.get("auth", None),
+            auth_token=server_config.get("auth_token", None),
         )
 
     raise ValueError("Cannot determine connector type from config")
-
-
-def create_session_from_config(
-    filepath: str,
-    server_name: str | None = None,
-    model_provider: str | ModelProvider = "openai",
-) -> MCPSession:
-    """Create an MCPSession from a configuration file.
-
-    Args:
-        filepath: Path to the configuration file
-        server_name: Name of the server to use from config, uses first if None
-        model_provider: Model provider to use for tool conversion
-
-    Returns:
-        Configured MCPSession instance
-    """
-    config = load_config_file(filepath)
-
-    # Get server config
-    servers = config.get("mcpServers", {})
-    if not servers:
-        raise ValueError("No MCP servers defined in config")
-
-    # If server_name not specified, use the first one
-    if not server_name:
-        server_name = next(iter(servers.keys()))
-
-    if server_name not in servers:
-        raise ValueError(f"Server '{server_name}' not found in config")
-
-    server_config = servers[server_name]
-    connector = create_connector_from_config(server_config)
-
-    return MCPSession(connector, model_provider)

mcp_use/connectors/base.py

@@ -8,9 +8,11 @@ must implement.
 from abc import ABC, abstractmethod
 from typing import Any
 
-from mcp.types import CallToolResult
+from mcp import ClientSession
+from mcp.types import CallToolResult, Tool
 
-from mcp_use.types import Tool
+from ..logging import logger
+from ..task_managers import ConnectionManager
 
 
 class BaseConnector(ABC):
@@ -19,43 +21,119 @@ class BaseConnector(ABC):
     This class defines the interface that all MCP connectors must implement.
     """
 
+    def __init__(self):
+        """Initialize base connector with common attributes."""
+        self.client: ClientSession | None = None
+        self._connection_manager: ConnectionManager | None = None
+        self._tools: list[Tool] | None = None
+        self._connected = False
+
     @abstractmethod
     async def connect(self) -> None:
         """Establish a connection to the MCP implementation."""
         pass
 
-    @abstractmethod
     async def disconnect(self) -> None:
         """Close the connection to the MCP implementation."""
-        pass
+        if not self._connected:
+            logger.debug("Not connected to MCP implementation")
+            return
+
+        logger.info("Disconnecting from MCP implementation")
+        await self._cleanup_resources()
+        self._connected = False
+        logger.info("Disconnected from MCP implementation")
+
+    async def _cleanup_resources(self) -> None:
+        """Clean up all resources associated with this connector."""
+        errors = []
+
+        # First close the client session
+        if self.client:
+            try:
+                logger.debug("Closing client session")
+                await self.client.__aexit__(None, None, None)
+            except Exception as e:
+                error_msg = f"Error closing client session: {e}"
+                logger.warning(error_msg)
+                errors.append(error_msg)
+            finally:
+                self.client = None
+
+        # Then stop the connection manager
+        if self._connection_manager:
+            try:
+                logger.debug("Stopping connection manager")
+                await self._connection_manager.stop()
+            except Exception as e:
+                error_msg = f"Error stopping connection manager: {e}"
+                logger.warning(error_msg)
+                errors.append(error_msg)
+            finally:
+                self._connection_manager = None
+
+        # Reset tools
+        self._tools = None
+
+        if errors:
+            logger.warning(f"Encountered {len(errors)} errors during resource cleanup")
 
-    @abstractmethod
     async def initialize(self) -> dict[str, Any]:
         """Initialize the MCP session and return session information."""
-        pass
+        if not self.client:
+            raise RuntimeError("MCP client is not connected")
+
+        logger.info("Initializing MCP session")
+
+        # Initialize the session
+        result = await self.client.initialize()
+
+        # Get available tools
+        tools_result = await self.client.list_tools()
+        self._tools = tools_result.tools
+
+        logger.info(f"MCP session initialized with {len(self._tools)} tools")
+
+        return result
 
     @property
-    @abstractmethod
     def tools(self) -> list[Tool]:
         """Get the list of available tools."""
-        pass
+        if not self._tools:
+            raise RuntimeError("MCP client is not initialized")
+        return self._tools
 
-    @abstractmethod
     async def call_tool(self, name: str, arguments: dict[str, Any]) -> CallToolResult:
         """Call an MCP tool with the given arguments."""
-        pass
+        if not self.client:
+            raise RuntimeError("MCP client is not connected")
+
+        logger.debug(f"Calling tool '{name}' with arguments: {arguments}")
+        result = await self.client.call_tool(name, arguments)
+        return result
 
-    @abstractmethod
     async def list_resources(self) -> list[dict[str, Any]]:
         """List all available resources from the MCP implementation."""
-        pass
+        if not self.client:
+            raise RuntimeError("MCP client is not connected")
+
+        logger.debug("Listing resources")
+        resources = await self.client.list_resources()
+        return resources
 
-    @abstractmethod
     async def read_resource(self, uri: str) -> tuple[bytes, str]:
         """Read a resource by URI."""
-        pass
+        if not self.client:
+            raise RuntimeError("MCP client is not connected")
+
+        logger.debug(f"Reading resource: {uri}")
+        resource = await self.client.read_resource(uri)
+        return resource.content, resource.mimeType
 
-    @abstractmethod
     async def request(self, method: str, params: dict[str, Any] | None = None) -> Any:
         """Send a raw request to the MCP implementation."""
-        pass
+        if not self.client:
+            raise RuntimeError("MCP client is not connected")
+
+        logger.debug(f"Sending request: {method} with params: {params}")
+        return await self.client.request({"method": method, "params": params or {}})

mcp_use/connectors/http.py

@@ -2,24 +2,30 @@
 HTTP connector for MCP implementations.
 
 This module provides a connector for communicating with MCP implementations
-through HTTP APIs.
+through HTTP APIs with SSE for transport.
 """
 
-from typing import Any
-
-import aiohttp
+from mcp import ClientSession
 
+from ..logging import logger
+from ..task_managers import SseConnectionManager
 from .base import BaseConnector
 
 
 class HttpConnector(BaseConnector):
-    """Connector for MCP implementations using HTTP transport.
+    """Connector for MCP implementations using HTTP transport with SSE.
 
-    This connector uses HTTP requests to communicate with remote MCP implementations.
+    This connector uses HTTP/SSE to communicate with remote MCP implementations,
+    using a connection manager to handle the proper lifecycle management.
     """
 
     def __init__(
-        self, base_url: str, auth_token: str | None = None, headers: dict[str, str] | None = None
+        self,
+        base_url: str,
+        auth_token: str | None = None,
+        headers: dict[str, str] | None = None,
+        timeout: float = 5,
+        sse_read_timeout: float = 60 * 5,
     ):
         """Initialize a new HTTP connector.
 
@@ -27,100 +33,50 @@ class HttpConnector(BaseConnector):
             base_url: The base URL of the MCP HTTP API.
             auth_token: Optional authentication token.
             headers: Optional additional headers.
+            timeout: Timeout for HTTP operations in seconds.
+            sse_read_timeout: Timeout for SSE read operations in seconds.
         """
+        super().__init__()
         self.base_url = base_url.rstrip("/")
         self.auth_token = auth_token
         self.headers = headers or {}
         if auth_token:
             self.headers["Authorization"] = f"Bearer {auth_token}"
-        self.session: aiohttp.ClientSession | None = None
+        self.timeout = timeout
+        self.sse_read_timeout = sse_read_timeout
 
     async def connect(self) -> None:
         """Establish a connection to the MCP implementation."""
-        self.session = aiohttp.ClientSession(headers=self.headers)
-
-    async def disconnect(self) -> None:
-        """Close the connection to the MCP implementation."""
-        if self.session:
-            await self.session.close()
-            self.session = None
-
-    async def _request(self, method: str, endpoint: str, data: dict[str, Any] | None = None) -> Any:
-        """Send an HTTP request to the MCP API.
-
-        Args:
-            method: The HTTP method (GET, POST, etc.).
-            endpoint: The API endpoint path.
-            data: Optional request data.
-
-        Returns:
-            The parsed JSON response.
-        """
-        if not self.session:
-            raise RuntimeError("HTTP session is not connected")
-
-        url = f"{self.base_url}/{endpoint.lstrip('/')}"
-
-        if method.upper() == "GET" and data:
-            # For GET requests, convert data to query parameters
-            async with self.session.get(url, params=data) as response:
-                response.raise_for_status()
-                return await response.json()
-        else:
-            # For other methods, send data as JSON body
-            async with self.session.request(method, url, json=data) as response:
-                response.raise_for_status()
-                return await response.json()
-
-    async def initialize(self) -> dict[str, Any]:
-        """Initialize the MCP session and return session information."""
-        return await self._request("POST", "initialize")
-
-    async def list_tools(self) -> list[dict[str, Any]]:
-        """List all available tools from the MCP implementation."""
-        result = await self._request("GET", "tools")
-        return result.get("tools", [])
-
-    async def call_tool(self, name: str, arguments: dict[str, Any]) -> Any:
-        """Call an MCP tool with the given arguments."""
-        return await self._request("POST", f"tools/{name}", arguments)
-
-    async def list_resources(self) -> list[dict[str, Any]]:
-        """List all available resources from the MCP implementation."""
-        result = await self._request("GET", "resources")
-        return result
-
-    async def read_resource(self, uri: str) -> tuple[bytes, str]:
-        """Read a resource by URI."""
-        # For resources, we may need to handle binary data
-        if not self.session:
-            raise RuntimeError("HTTP session is not connected")
-
-        url = f"{self.base_url}/resources/read"
-
-        async with self.session.get(url, params={"uri": uri}) as response:
-            response.raise_for_status()
-
-            # Check if this is a JSON response or binary data
-            content_type = response.headers.get("Content-Type", "")
-            if "application/json" in content_type:
-                data = await response.json()
-                content = data.get("content", b"")
-                mime_type = data.get("mimeType", "")
-
-                # If content is base64 encoded, decode it
-                if isinstance(content, str):
-                    import base64
-
-                    content = base64.b64decode(content)
-
-                return content, mime_type
-            else:
-                # Assume binary response
-                content = await response.read()
-                return content, content_type
-
-    async def request(self, method: str, params: dict[str, Any] | None = None) -> Any:
-        """Send a raw request to the MCP implementation."""
-        # For custom methods, we'll use the RPC-style endpoint
-        return await self._request("POST", "rpc", {"method": method, "params": params or {}})
+        if self._connected:
+            logger.debug("Already connected to MCP implementation")
+            return
+
+        logger.info(f"Connecting to MCP implementation via HTTP/SSE: {self.base_url}")
+        try:
+            # 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.info(
+                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 via HTTP/SSE: {e}")
+
+            # Clean up any resources if connection failed
+            await self._cleanup_resources()
+
+            # Re-raise the original exception
+            raise