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

mcp_use/client/connectors/stdio.py

@@ -0,0 +1,109 @@
+"""
+StdIO connector for MCP implementations.
+
+This module provides a connector for communicating with MCP implementations
+through the standard input/output streams.
+"""
+
+import sys
+
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.session import ElicitationFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
+
+from mcp_use.client.connectors.base import BaseConnector
+from mcp_use.client.middleware import CallbackClientSession, Middleware
+from mcp_use.client.task_managers import StdioConnectionManager
+from mcp_use.logging import logger
+
+
+class StdioConnector(BaseConnector):
+    """Connector for MCP implementations using stdio transport.
+
+    This connector uses the stdio transport to communicate with MCP implementations
+    that are executed as child processes. It uses a connection manager to handle
+    the proper lifecycle management of the stdio client.
+    """
+
+    def __init__(
+        self,
+        command: str = "npx",
+        args: list[str] | None = None,
+        env: dict[str, str] | None = None,
+        errlog=sys.stderr,
+        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,
+    ):
+        """Initialize a new stdio connector.
+
+        Args:
+            command: The command to execute.
+            args: Optional command line arguments.
+            env: Optional environment variables.
+            errlog: Stream to write error output to.
+            sampling_callback: Optional callback to sample the client.
+            elicitation_callback: Optional callback to elicit the client.
+        """
+        super().__init__(
+            sampling_callback=sampling_callback,
+            elicitation_callback=elicitation_callback,
+            message_handler=message_handler,
+            logging_callback=logging_callback,
+            middleware=middleware,
+        )
+        self.command = command
+        self.args = args or []  # Ensure args is never None
+        self.env = env
+        self.errlog = errlog
+
+    async def connect(self) -> None:
+        """Establish a connection to the MCP implementation."""
+        if self._connected:
+            logger.debug("Already connected to MCP implementation")
+            return
+
+        logger.debug(f"Connecting to MCP implementation: {self.command}")
+        try:
+            # Create server parameters
+            server_params = StdioServerParameters(command=self.command, args=self.args, env=self.env)
+
+            # Create and start the connection manager
+            self._connection_manager = StdioConnectionManager(server_params, self.errlog)
+            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,
+                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: {self.command}")
+
+        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()
+
+            # Re-raise the original exception
+            raise
+
+    @property
+    def public_identifier(self) -> str:
+        """Get the identifier for the connector."""
+        return f"stdio:{self.command} {' '.join(self.args)}"

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

@@ -0,0 +1,257 @@
+"""
+WebSocket connector for MCP implementations.
+
+This module provides a connector for communicating with MCP implementations
+through WebSocket connections.
+"""
+
+import asyncio
+import json
+import uuid
+from typing import Any
+
+import httpx
+from mcp.types import Tool
+from websockets import ClientConnection
+
+from mcp_use.client.connectors.base import BaseConnector
+from mcp_use.client.task_managers import ConnectionManager, WebSocketConnectionManager
+from mcp_use.logging import logger
+
+
+class WebSocketConnector(BaseConnector):
+    """Connector for MCP implementations using WebSocket transport.
+
+    This connector uses WebSockets to communicate with remote MCP implementations,
+    using a connection manager to handle the proper lifecycle management.
+    """
+
+    def __init__(
+        self,
+        url: str,
+        headers: dict[str, str] | None = None,
+        auth: str | dict[str, Any] | httpx.Auth | None = None,
+    ):
+        """Initialize a new WebSocket connector.
+
+        Args:
+            url: The WebSocket URL to connect to.
+            headers: Optional additional headers.
+            auth: Authentication method - can be:
+                - A string token: Use Bearer token authentication
+                - A dict: Not supported for WebSocket (will log warning)
+                - An httpx.Auth object: Not supported for WebSocket (will log warning)
+        """
+        self.url = url
+        self.headers = headers or {}
+
+        # Handle authentication - WebSocket only supports bearer tokens
+        # An auth field it's not needed
+        if auth is not None:
+            if isinstance(auth, str):
+                self.headers["Authorization"] = f"Bearer {auth}"
+            else:
+                logger.warning("WebSocket connector only supports bearer token authentication")
+
+        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] = {}
+        self._tools: list[Tool] | None = None
+        self._connected = False
+
+    async def connect(self) -> None:
+        """Establish a connection to the MCP implementation."""
+        if self._connected:
+            logger.debug("Already connected to MCP implementation")
+            return
+
+        logger.debug(f"Connecting to MCP implementation via WebSocket: {self.url}")
+        try:
+            # Create and start the connection manager
+            self._connection_manager = WebSocketConnectionManager(self.url, self.headers)
+            self.ws = await self._connection_manager.start()
+
+            # Start the message receiver task
+            self._receiver_task = asyncio.create_task(self._receive_messages(), name="websocket_receiver_task")
+
+            # Mark as connected
+            self._connected = True
+            logger.debug(f"Successfully connected to MCP implementation via WebSocket: {self.url}")
+
+        except Exception as e:
+            logger.error(f"Failed to connect to MCP implementation via WebSocket: {e}")
+
+            # Clean up any resources if connection failed
+            await self._cleanup_resources()
+
+            # Re-raise the original exception
+            raise
+
+    async def _receive_messages(self) -> None:
+        """Continuously receive and process messages from the WebSocket."""
+        if not self.ws:
+            raise RuntimeError("WebSocket is not connected")
+
+        try:
+            async for message in self.ws:
+                # Parse the message
+                data = json.loads(message)
+
+                # Check if this is a response to a pending request
+                request_id = data.get("id")
+                if request_id and request_id in self.pending_requests:
+                    future = self.pending_requests.pop(request_id)
+                    if "result" in data:
+                        future.set_result(data["result"])
+                    elif "error" in data:
+                        future.set_exception(Exception(data["error"]))
+
+                    logger.debug(f"Received response for request {request_id}")
+                else:
+                    logger.debug(f"Received message: {data}")
+        except Exception as e:
+            logger.error(f"Error in WebSocket message receiver: {e}")
+            # If the websocket connection was closed or errored,
+            # reject all pending requests
+            for future in self.pending_requests.values():
+                if not future.done():
+                    future.set_exception(e)
+
+    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")
+
+    async def _cleanup_resources(self) -> None:
+        """Clean up all resources associated with this connector."""
+        errors = []
+
+        # First cancel the receiver task
+        if self._receiver_task and not self._receiver_task.done():
+            try:
+                logger.debug("Cancelling WebSocket receiver task")
+                self._receiver_task.cancel()
+                try:
+                    await self._receiver_task
+                except asyncio.CancelledError:
+                    logger.debug("WebSocket receiver task cancelled successfully")
+                except Exception as e:
+                    logger.warning(f"Error during WebSocket receiver task cancellation: {e}")
+            except Exception as e:
+                error_msg = f"Error cancelling WebSocket receiver task: {e}"
+                logger.warning(error_msg)
+                errors.append(error_msg)
+            finally:
+                self._receiver_task = None
+
+        # Reject any pending requests
+        if self.pending_requests:
+            logger.debug(f"Rejecting {len(self.pending_requests)} pending requests")
+            for future in self.pending_requests.values():
+                if not future.done():
+                    future.set_exception(ConnectionError("WebSocket disconnected"))
+            self.pending_requests.clear()
+
+        # 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
+                self.ws = None
+
+        # Reset tools
+        self._tools = None
+
+        if errors:
+            logger.warning(f"Encountered {len(errors)} errors during resource cleanup")
+
+    async def _send_request(self, method: str, params: dict[str, Any] | None = None) -> Any:
+        """Send a request and wait for a response."""
+        if not self.ws:
+            raise RuntimeError("WebSocket is not connected")
+
+        # Create a request ID
+        request_id = str(uuid.uuid4())
+
+        # Create a future to receive the response
+        future = asyncio.Future()
+        self.pending_requests[request_id] = future
+
+        # Send the request
+        await self.ws.send(json.dumps({"id": request_id, "method": method, "params": params or {}}))
+
+        logger.debug(f"Sent request {request_id} method: {method}")
+
+        # Wait for the response
+        try:
+            return await future
+        except Exception as e:
+            # Remove the request from pending requests
+            self.pending_requests.pop(request_id, None)
+            logger.error(f"Error waiting for response to request {request_id}: {e}")
+            raise
+
+    async def initialize(self) -> dict[str, Any]:
+        """Initialize the MCP session and return session information."""
+        logger.debug("Initializing MCP session")
+        result = await self._send_request("initialize")
+
+        # Get available tools
+        tools_result = await self.list_tools()
+        self._tools = [Tool(**tool) for tool in tools_result]
+
+        logger.debug(f"MCP session initialized with {len(self._tools)} tools")
+        return result
+
+    async def list_tools(self) -> list[dict[str, Any]]:
+        """List all available tools from the MCP implementation."""
+        logger.debug("Listing tools")
+        result = await self._send_request("tools/list")
+        return result.get("tools", [])
+
+    @property
+    def tools(self) -> list[Tool]:
+        """Get the list of available tools."""
+        if not self._tools:
+            raise RuntimeError("MCP client is not initialized")
+        return self._tools
+
+    async def call_tool(self, name: str, arguments: dict[str, Any]) -> Any:
+        """Call an MCP tool with the given arguments."""
+        logger.debug(f"Calling tool '{name}' with arguments: {arguments}")
+        return await self._send_request("tools/call", {"name": name, "arguments": arguments})
+
+    async def list_resources(self) -> list[dict[str, Any]]:
+        """List all available resources from the MCP implementation."""
+        logger.debug("Listing resources")
+        result = await self._send_request("resources/list")
+        return result
+
+    async def read_resource(self, uri: str) -> tuple[bytes, str]:
+        """Read a resource by URI."""
+        logger.debug(f"Reading resource: {uri}")
+        result = await self._send_request("resources/read", {"uri": uri})
+        return result.get("content", b""), result.get("mimeType", "")
+
+    async def request(self, method: str, params: dict[str, Any] | None = None) -> Any:
+        """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 f"websocket:{self.url}"

mcp_use/client/exceptions.py

@@ -0,0 +1,31 @@
+"""MCP-use exceptions."""
+
+
+class MCPError(Exception):
+    """Base exception for MCP-use."""
+
+    pass
+
+
+class OAuthDiscoveryError(MCPError):
+    """OAuth discovery auth metadata error"""
+
+    pass
+
+
+class OAuthAuthenticationError(MCPError):
+    """OAuth authentication-related errors"""
+
+    pass
+
+
+class ConnectionError(MCPError):
+    """Connection-related errors."""
+
+    pass
+
+
+class ConfigurationError(MCPError):
+    """Configuration-related errors."""
+
+    pass

mcp_use/client/middleware/__init__.py

@@ -0,0 +1,50 @@
+"""
+Middleware package for MCP request interception and processing.
+
+This package provides a flexible middleware system for intercepting MCP requests
+and responses, enabling logging, metrics, caching, and custom processing.
+
+The middleware system follows an Express.js-style pattern where middleware functions
+receive a request context and a call_next function, allowing them to process both
+incoming requests and outgoing responses.
+"""
+
+# Core middleware implementation
+# Default logging middleware
+from .logging import default_logging_middleware
+
+# Metrics middleware classes
+from .metrics import (
+    CombinedAnalyticsMiddleware,
+    ErrorTrackingMiddleware,
+    MetricsMiddleware,
+    PerformanceMetricsMiddleware,
+)
+
+# Protocol types for type-safe middleware
+from .middleware import (
+    CallbackClientSession,
+    MCPResponseContext,
+    Middleware,
+    MiddlewareContext,
+    MiddlewareManager,
+    NextFunctionT,
+)
+
+__all__ = [
+    # Core types and classes
+    "MiddlewareContext",
+    "MCPResponseContext",
+    "Middleware",
+    "MiddlewareManager",
+    "CallbackClientSession",
+    # Protocol types
+    "NextFunctionT",
+    # Default logging middleware
+    "default_logging_middleware",
+    # Metrics middleware
+    "MetricsMiddleware",
+    "PerformanceMetricsMiddleware",
+    "ErrorTrackingMiddleware",
+    "CombinedAnalyticsMiddleware",
+]

mcp_use/client/middleware/logging.py

@@ -0,0 +1,31 @@
+"""
+Default logging middleware for MCP requests.
+
+Simple debug logging for all MCP requests and responses.
+"""
+
+import time
+from typing import Any
+
+from mcp_use.client.middleware.middleware import Middleware, MiddlewareContext, NextFunctionT
+from mcp_use.logging import logger
+
+
+class LoggingMiddleware(Middleware):
+    """Default logging middleware that logs all MCP requests and responses with logger.debug."""
+
+    async def on_request(self, context: MiddlewareContext[Any], call_next: NextFunctionT) -> Any:
+        """Logs all MCP requests and responses with logger.debug."""
+        logger.debug(f"[{context.id}] {context.connection_id} -> {context.method}")
+        try:
+            result = await call_next(context)
+            duration = time.time() - context.timestamp
+            logger.debug(f"[{context.id}] {context.connection_id} <- {context.method} ({duration:.3f}s)")
+            return result
+        except Exception as e:
+            duration = time.time() - context.timestamp
+            logger.debug(f"[{context.id}] {context.connection_id} <- {context.method} FAILED ({duration:.3f}s): {e}")
+            raise
+
+
+default_logging_middleware = LoggingMiddleware()