mcp-use 1.3.11__py3-none-any.whl → 1.3.13__py3-none-any.whl

mcp_use/client/session.py

@@ -0,0 +1,162 @@
+"""
+Session manager for MCP connections.
+
+This module provides a session manager for MCP connections,
+which handles authentication, initialization, and tool discovery.
+"""
+
+from datetime import timedelta
+from typing import Any
+
+from mcp.types import CallToolResult, GetPromptResult, Prompt, ReadResourceResult, Resource, Tool
+from pydantic import AnyUrl
+
+from mcp_use.client.connectors.base import BaseConnector
+from mcp_use.telemetry.telemetry import telemetry
+
+
+class MCPSession:
+    """Session manager for MCP connections.
+
+    This class manages the lifecycle of an MCP connection, including
+    authentication, initialization, and tool discovery.
+    """
+
+    def __init__(
+        self,
+        connector: BaseConnector,
+        auto_connect: bool = True,
+    ) -> None:
+        """Initialize a new MCP session.
+
+        Args:
+            connector: The connector to use for communicating with the MCP implementation.
+            auto_connect: Whether to automatically connect to the MCP implementation.
+        """
+        self.connector = connector
+        self.session_info: dict[str, Any] | None = None
+        self.auto_connect = auto_connect
+
+    async def __aenter__(self) -> "MCPSession":
+        """Enter the async context manager.
+
+        Returns:
+            The session instance.
+        """
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit the async context manager.
+
+        Args:
+            exc_type: The exception type, if an exception was raised.
+            exc_val: The exception value, if an exception was raised.
+            exc_tb: The exception traceback, if an exception was raised.
+        """
+        await self.disconnect()
+
+    async def connect(self) -> None:
+        """Connect to the MCP implementation."""
+        await self.connector.connect()
+
+    async def disconnect(self) -> None:
+        """Disconnect from the MCP implementation."""
+        await self.connector.disconnect()
+
+    @telemetry("session_initialize")
+    async def initialize(self) -> dict[str, Any]:
+        """Initialize the MCP session and discover available tools.
+
+        Returns:
+            The session information returned by the MCP implementation.
+        """
+        # Make sure we're connected
+        if not self.is_connected and self.auto_connect:
+            await self.connect()
+
+        # Initialize the session
+        self.session_info = await self.connector.initialize()
+
+        return self.session_info
+
+    @property
+    def is_connected(self) -> bool:
+        """Check if the connector is connected.
+
+        Returns:
+            True if the connector is connected, False otherwise.
+        """
+        return self.connector.is_connected
+
+    # Convenience methods for MCP operations
+    @telemetry("session_call_tool")
+    async def call_tool(
+        self, name: str, arguments: dict[str, Any], read_timeout_seconds: timedelta | None = None
+    ) -> CallToolResult:
+        """Call an MCP tool.
+
+        Args:
+            name: The name of the tool to call.
+            arguments: The arguments to pass to the tool.
+            read_timeout_seconds: Optional timeout for the tool call.
+
+        Returns:
+            The result of the tool call.
+
+        Raises:
+            RuntimeError: If the connection is lost and cannot be reestablished.
+        """
+        return await self.connector.call_tool(name, arguments, read_timeout_seconds)
+
+    @telemetry("session_list_tools")
+    async def list_tools(self) -> list[Tool]:
+        """List all available tools from the MCP server.
+
+        Returns:
+            List of available tools.
+        """
+        return await self.connector.list_tools()
+
+    @telemetry("session_list_resources")
+    async def list_resources(self) -> list[Resource]:
+        """List all available resources from the MCP server.
+
+        Returns:
+            List of available resources.
+        """
+        return await self.connector.list_resources()
+
+    @telemetry("session_read_resource")
+    async def read_resource(self, uri: AnyUrl) -> ReadResourceResult:
+        """Read a resource by URI.
+
+        Args:
+            uri: The URI of the resource to read.
+
+        Returns:
+            The resource content.
+        """
+        return await self.connector.read_resource(uri)
+
+    @telemetry("session_list_prompts")
+    async def list_prompts(self) -> list[Prompt]:
+        """List all available prompts from the MCP server.
+
+        Returns:
+            List of available prompts.
+        """
+        return await self.connector.list_prompts()
+
+    @telemetry("session_get_prompt")
+    async def get_prompt(self, name: str, arguments: dict[str, Any] | None = None) -> GetPromptResult:
+        """Get a prompt by name.
+
+        Args:
+            name: The name of the prompt to get.
+            arguments: Optional arguments for the prompt.
+
+        Returns:
+            The prompt result with messages.
+        """
+        return await self.connector.get_prompt(name, arguments)

mcp_use/client/task_managers/__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 ConnectionManager
+from .sse import SseConnectionManager
+from .stdio import StdioConnectionManager
+from .streamable_http import StreamableHttpConnectionManager
+from .websocket import WebSocketConnectionManager
+
+__all__ = [
+    "ConnectionManager",
+    "StdioConnectionManager",
+    "WebSocketConnectionManager",
+    "SseConnectionManager",
+    "StreamableHttpConnectionManager",
+]

mcp_use/client/task_managers/base.py

@@ -0,0 +1,145 @@
+"""
+Connection management for MCP implementations.
+
+This module provides an abstract base class for different types of connection
+managers used in MCP connectors.
+"""
+
+import asyncio
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+from mcp_use.logging import logger
+
+# Type variable for connection types
+T = TypeVar("T")
+
+
+class ConnectionManager(Generic[T], ABC):
+    """Abstract base class for connection managers.
+
+    This class defines the interface for different types of connection managers
+    used with MCP connectors.
+    """
+
+    def __init__(self) -> None:
+        """Initialize a new connection manager."""
+        self._ready_event = asyncio.Event()
+        self._done_event = asyncio.Event()
+        self._stop_event = asyncio.Event()
+        self._exception: Exception | None = None
+        self._connection: T | None = None
+        self._task: asyncio.Task[None] | None = None
+
+    @abstractmethod
+    async def _establish_connection(self) -> T:
+        """Establish the connection.
+
+        This method should be implemented by subclasses to establish
+        the specific type of connection needed.
+
+        Returns:
+            The established connection.
+
+        Raises:
+            Exception: If connection cannot be established.
+        """
+        pass
+
+    @abstractmethod
+    async def _close_connection(self) -> None:
+        """Close the connection.
+
+        This method should be implemented by subclasses to close
+        the specific type of connection.
+
+        """
+        pass
+
+    async def start(self) -> T:
+        """Start the connection manager and establish a connection.
+
+        Returns:
+            The established connection.
+
+        Raises:
+            Exception: If connection cannot be established.
+        """
+        # Reset state
+        self._ready_event.clear()
+        self._done_event.clear()
+        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")
+
+        # Wait for the connection to be ready or fail
+        await self._ready_event.wait()
+
+        # If there was an exception, raise it
+        if self._exception:
+            raise self._exception
+
+        # Return the connection
+        if self._connection is None:
+            raise RuntimeError("Connection was not established")
+        return self._connection
+
+    async def stop(self) -> None:
+        """Stop the connection manager and close the connection."""
+        # Signal stop to the connection task instead of cancelling it, avoids
+        # propagating CancelledError to unrelated tasks.
+        if self._task and not self._task.done():
+            logger.debug(f"Signaling stop to {self.__class__.__name__} task")
+            self._stop_event.set()
+            # Wait for it to finish gracefully
+            await self._task
+
+        # Ensure cleanup completed
+        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.
+
+        This task establishes and maintains the connection until cancelled.
+        """
+        logger.debug(f"Starting {self.__class__.__name__} task")
+        try:
+            # Establish the connection
+            self._connection = await self._establish_connection()
+            logger.debug(f"{self.__class__.__name__} connected successfully")
+
+            # Signal that the connection is ready
+            self._ready_event.set()
+
+            # Wait until stop is requested
+            await self._stop_event.wait()
+
+        except Exception as e:
+            # Store the exception
+            self._exception = e
+            logger.error(f"Error in {self.__class__.__name__} task: {e}")
+
+            # Signal that the connection is ready (with error)
+            self._ready_event.set()
+
+        finally:
+            # Close the connection if it was established
+            if self._connection is not None:
+                try:
+                    await self._close_connection()
+                except Exception as e:
+                    logger.warning(f"Error closing connection in {self.__class__.__name__}: {e}")
+                self._connection = None
+
+            # Signal that the connection is done
+            self._done_event.set()

mcp_use/client/task_managers/sse.py

@@ -0,0 +1,84 @@
+"""
+SSE connection management for MCP implementations.
+
+This module provides a connection manager for SSE-based MCP connections
+that ensures proper task isolation and resource cleanup.
+"""
+
+from typing import Any
+
+import httpx
+from mcp.client.sse import sse_client
+
+from mcp_use.client.task_managers.base import ConnectionManager
+from mcp_use.logging import logger
+
+
+class SseConnectionManager(ConnectionManager[tuple[Any, Any]]):
+    """Connection manager for SSE-based MCP connections.
+
+    This class handles the proper task isolation for sse_client context managers
+    to prevent the "cancel scope in different task" error. It runs the sse_client
+    in a dedicated task and manages its lifecycle.
+    """
+
+    def __init__(
+        self,
+        url: str,
+        headers: dict[str, str] | None = None,
+        timeout: float = 5,
+        sse_read_timeout: float = 60 * 5,
+        auth: httpx.Auth | None = None,
+    ):
+        """Initialize a new SSE connection manager.
+
+        Args:
+            url: The SSE endpoint URL
+            headers: Optional HTTP headers
+            timeout: Timeout for HTTP operations in seconds
+            sse_read_timeout: Timeout for SSE read operations in seconds
+            auth: Optional httpx.Auth instance for authentication
+        """
+        super().__init__()
+        self.url = url
+        self.headers = headers or {}
+        self.timeout = timeout
+        self.sse_read_timeout = sse_read_timeout
+        self.auth = auth
+        self._sse_ctx = None
+
+    async def _establish_connection(self) -> tuple[Any, Any]:
+        """Establish an SSE connection.
+
+        Returns:
+            A tuple of (read_stream, write_stream)
+
+        Raises:
+            Exception: If connection cannot be established.
+        """
+        # Create the context manager
+        self._sse_ctx = sse_client(
+            url=self.url,
+            headers=self.headers,
+            timeout=self.timeout,
+            sse_read_timeout=self.sse_read_timeout,
+            auth=self.auth,
+        )
+
+        # Enter the context manager
+        read_stream, write_stream = await self._sse_ctx.__aenter__()
+
+        # Return the streams
+        return (read_stream, write_stream)
+
+    async def _close_connection(self) -> None:
+        """Close the SSE connection."""
+
+        if self._sse_ctx:
+            # Exit the context manager
+            try:
+                await self._sse_ctx.__aexit__(None, None, None)
+            except Exception as e:
+                logger.warning(f"Error closing SSE context: {e}")
+            finally:
+                self._sse_ctx = None

mcp_use/client/task_managers/stdio.py

@@ -0,0 +1,69 @@
+"""
+StdIO connection management for MCP implementations.
+
+This module provides a connection manager for stdio-based MCP connections
+that ensures proper task isolation and resource cleanup.
+"""
+
+import sys
+from typing import Any, TextIO
+
+from mcp import StdioServerParameters
+from mcp.client.stdio import stdio_client
+
+from mcp_use.client.task_managers.base import ConnectionManager
+from mcp_use.logging import logger
+
+
+class StdioConnectionManager(ConnectionManager[tuple[Any, Any]]):
+    """Connection manager for stdio-based MCP connections.
+
+    This class handles the proper task isolation for stdio_client context managers
+    to prevent the "cancel scope in different task" error. It runs the stdio_client
+    in a dedicated task and manages its lifecycle.
+    """
+
+    def __init__(
+        self,
+        server_params: StdioServerParameters,
+        errlog: TextIO = sys.stderr,
+    ):
+        """Initialize a new stdio connection manager.
+
+        Args:
+            server_params: The parameters for the stdio server
+            errlog: The error log stream
+        """
+        super().__init__()
+        self.server_params = server_params
+        self.errlog = errlog
+        self._stdio_ctx = None
+
+    async def _establish_connection(self) -> tuple[Any, Any]:
+        """Establish a stdio connection.
+
+        Returns:
+            A tuple of (read_stream, write_stream)
+
+        Raises:
+            Exception: If connection cannot be established.
+        """
+        # Create the context manager
+        self._stdio_ctx = stdio_client(self.server_params, self.errlog)
+
+        # Enter the context manager
+        read_stream, write_stream = await self._stdio_ctx.__aenter__()
+
+        # Return the streams
+        return (read_stream, write_stream)
+
+    async def _close_connection(self) -> None:
+        """Close the stdio connection."""
+        if self._stdio_ctx:
+            # Exit the context manager
+            try:
+                await self._stdio_ctx.__aexit__(None, None, None)
+            except Exception as e:
+                logger.warning(f"Error closing stdio context: {e}")
+            finally:
+                self._stdio_ctx = None

mcp_use/client/task_managers/streamable_http.py

@@ -0,0 +1,86 @@
+"""
+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
+
+import httpx
+from mcp.client.streamable_http import streamablehttp_client
+
+from mcp_use.client.task_managers.base import ConnectionManager
+from mcp_use.logging import logger
+
+
+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,
+        auth: httpx.Auth | None = None,
+    ):
+        """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
+            auth: Optional httpx.Auth instance for authentication
+        """
+        super().__init__()
+        self.url = url
+        self.headers = headers or {}
+        self.timeout = timedelta(seconds=timeout)
+        self.read_timeout = timedelta(seconds=read_timeout)
+        self.auth = auth
+        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,
+            auth=self.auth,
+        )
+
+        # 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/client/task_managers/websocket.py

@@ -0,0 +1,68 @@
+"""
+WebSocket connection management for MCP implementations.
+
+This module provides a connection manager for WebSocket-based MCP connections.
+"""
+
+from typing import Any
+
+from mcp.client.websocket import websocket_client
+
+from mcp_use.client.task_managers.base import ConnectionManager
+from mcp_use.logging import logger
+
+
+class WebSocketConnectionManager(ConnectionManager[tuple[Any, Any]]):
+    """Connection manager for WebSocket-based MCP connections.
+
+    This class handles the lifecycle of WebSocket connections, ensuring proper
+    connection establishment and cleanup.
+    """
+
+    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.
+
+        Returns:
+            The established WebSocket connection
+
+        Raises:
+            Exception: If connection cannot be established
+        """
+        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
+        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