massgen 0.0.3__py3-none-any.whl → 0.1.0__py3-none-any.whl

massgen/mcp_tools/client.py

@@ -0,0 +1,894 @@
+# -*- coding: utf-8 -*-
+"""
+MCP client implementation for connecting to MCP servers. This module provides enhanced MCP client
+functionality to connect with MCP servers and integrate external tools into the MassGen workflow.
+"""
+import asyncio
+from dataclasses import dataclass
+from datetime import timedelta
+from enum import Enum
+from types import TracebackType
+from typing import Any, Awaitable, Callable, Dict, List, Optional, Union
+
+from mcp import ClientSession, StdioServerParameters
+from mcp import types as mcp_types
+from mcp.client.stdio import get_default_environment, stdio_client
+from mcp.client.streamable_http import streamablehttp_client
+
+from ..logger_config import logger
+from .circuit_breaker import MCPCircuitBreaker
+from .config_validator import MCPConfigValidator
+from .exceptions import (
+    MCPConnectionError,
+    MCPError,
+    MCPServerError,
+    MCPTimeoutError,
+    MCPValidationError,
+)
+from .security import (
+    prepare_command,
+    sanitize_tool_name,
+    substitute_env_variables,
+    validate_tool_arguments,
+)
+
+
+class ConnectionState(Enum):
+    """Connection state for MCP clients."""
+
+    DISCONNECTED = "disconnected"
+    CONNECTING = "connecting"
+    CONNECTED = "connected"
+    DISCONNECTING = "disconnecting"
+    FAILED = "failed"
+
+
+# Hook types reference: https://docs.anthropic.com/en/docs/claude-code/sdk/sdk-python#hook-types
+class HookType(Enum):
+    """Available hook types for MCP tool execution."""
+
+    PRE_TOOL_USE = "PreToolUse"
+
+
+def _ensure_timedelta(value: Union[int, float, timedelta], default_seconds: float) -> timedelta:
+    """
+    Ensure a value is converted to timedelta for consistent timeout handling.
+
+    Raises:
+        MCPValidationError: If value is invalid
+    """
+    if isinstance(value, timedelta):
+        if value.total_seconds() <= 0:
+            raise MCPValidationError(
+                f"Timeout must be positive, got {value.total_seconds()} seconds",
+                field="timeout",
+                value=value.total_seconds(),
+            )
+        return value
+    elif isinstance(value, (int, float)):
+        if value <= 0:
+            raise MCPValidationError(
+                f"Timeout must be positive, got {value} seconds",
+                field="timeout",
+                value=value,
+            )
+        return timedelta(seconds=value)
+    else:
+        logger.warning(f"Invalid timeout value {value}, using default {default_seconds}s")
+        return timedelta(seconds=default_seconds)
+
+
+@dataclass
+class _ServerClient:
+    """Internal container for per-server state."""
+
+    session: Optional[ClientSession] = None
+    manager_task: Optional[asyncio.Task] = None
+    connected_event: asyncio.Event = None
+    disconnect_event: asyncio.Event = None
+    connection_lock: asyncio.Lock = None
+    connection_state: ConnectionState = ConnectionState.DISCONNECTED
+    initialized: bool = False
+
+    def __post_init__(self):
+        if self.connected_event is None:
+            self.connected_event = asyncio.Event()
+        if self.disconnect_event is None:
+            self.disconnect_event = asyncio.Event()
+        if self.connection_lock is None:
+            self.connection_lock = asyncio.Lock()
+
+
+class MCPClient:
+    """
+    Unified MCP client for communicating with single or multiple MCP servers.
+    Provides improved security, error handling, and async context management.
+
+    Accepts a list of server configurations and automatically handles:
+    - Consistent tool naming: Always uses prefixed names (mcp__server__tool)
+    - Circuit breaker protection for all servers
+    - Parallel connection for multi-server scenarios
+    - Sequential connection for single-server scenarios
+    """
+
+    def __init__(
+        self,
+        server_configs: List[Dict[str, Any]],
+        *,
+        timeout_seconds: int = 30,
+        allowed_tools: Optional[List[str]] = None,
+        exclude_tools: Optional[List[str]] = None,
+        status_callback: Optional[Callable[[str, Dict[str, Any]], Awaitable[None]]] = None,
+        hooks: Optional[Dict[HookType, List[Callable[[str, Dict[str, Any]], Awaitable[bool]]]]] = None,
+    ):
+        """
+        Initialize MCP client.
+
+        Args:
+            server_configs: List of server configuration dicts (always a list, even for single server)
+            timeout_seconds: Timeout for operations in seconds
+            allowed_tools: Optional list of tool names to include (if None, includes all)
+            exclude_tools: Optional list of tool names to exclude (if None, excludes none)
+            status_callback: Optional async callback for status updates
+            hooks: Optional dict mapping hook types to lists of hook functions
+        """
+        # Validate all server configs
+        self._server_configs = [MCPConfigValidator.validate_server_config(config) for config in server_configs]
+
+        # Set name to first server's name for backward compatibility
+        self.name = self._server_configs[0]["name"]
+
+        self.timeout_seconds = timeout_seconds
+        self.allowed_tools = allowed_tools
+        self.exclude_tools = exclude_tools
+        self.status_callback = status_callback
+        self.hooks = hooks or {}
+
+        # Initialize circuit breaker for ALL scenarios
+        self._circuit_breaker = MCPCircuitBreaker()
+
+        # Per-server tracking
+        self._server_clients: Dict[str, _ServerClient] = {}
+        for config in self._server_configs:
+            self._server_clients[config["name"]] = _ServerClient()
+
+        # Unified registry for tools
+        self.tools: Dict[str, mcp_types.Tool] = {}
+        self._tool_to_server: Dict[str, str] = {}
+
+        # Connection management
+        self._initialized = False
+        self._cleanup_done = False
+        self._cleanup_lock = asyncio.Lock()
+        self._context_managed = False
+
+    @property
+    def session(self) -> Optional[ClientSession]:
+        """Return first server's session for backward compatibility."""
+        if self._server_configs:
+            first_server_name = self._server_configs[0]["name"]
+            server_client = self._server_clients.get(first_server_name)
+            if server_client:
+                return server_client.session
+        return None
+
+    def _get_server_session(self, server_name: str) -> ClientSession:
+        """Get session for server, raising error if not connected."""
+        server_client = self._server_clients.get(server_name)
+        if not server_client or not server_client.session:
+            raise MCPConnectionError(
+                f"Server '{server_name}' not connected",
+                server_name=server_name,
+            )
+        return server_client.session
+
+    async def connect(self) -> None:
+        """Connect to MCP server(s) and discover capabilities with circuit breaker integration."""
+        if self._initialized:
+            return
+
+        logger.info(f"Connecting to {len(self._server_configs)} MCP server(s)...")
+
+        # Send connecting status if callback is available
+        if self.status_callback:
+            await self.status_callback(
+                "connecting",
+                {
+                    "message": f"Connecting to {len(self._server_configs)} MCP server(s)",
+                    "server_count": len(self._server_configs),
+                },
+            )
+
+        if len(self._server_configs) > 1:
+            # Multi-server: connect in parallel
+            await self._connect_all_parallel()
+        else:
+            # Single-server: connect sequentially
+            await self._connect_single()
+
+        # Only mark as initialized if at least one server connected successfully
+        self._initialized = any(sc.initialized for sc in self._server_clients.values())
+
+        # Count successful and failed connections
+        successful_count = len([sc for sc in self._server_clients.values() if sc.initialized])
+        failed_count = len(self._server_configs) - successful_count
+
+        # Send connection summary status if callback is available
+        if self.status_callback:
+            await self.status_callback(
+                "connection_summary",
+                {
+                    "message": f"Connected to {successful_count}/{len(self._server_configs)} server(s)" + (f" ({failed_count} failed)" if failed_count > 0 else ""),
+                    "successful_count": successful_count,
+                    "failed_count": failed_count,
+                    "total_count": len(self._server_configs),
+                    "tools_count": len(self.tools),
+                },
+            )
+
+    async def _connect_server(self, server_name: str, config: Dict[str, Any]) -> bool:
+        """Connect to a single server with circuit breaker integration.
+
+        Returns:
+            True on success, False on failure
+        """
+        server_client = self._server_clients[server_name]
+
+        async with server_client.connection_lock:
+            # Check circuit breaker
+            if self._circuit_breaker.should_skip_server(server_name):
+                logger.warning(f"Skipping server {server_name} due to circuit breaker")
+                server_client.connection_state = ConnectionState.FAILED
+                return False
+
+            server_client.connection_state = ConnectionState.CONNECTING
+
+            try:
+                # Start background manager task
+                server_client.manager_task = asyncio.create_task(
+                    self._run_manager(server_name, config),
+                )
+
+                # Wait for connection
+                await asyncio.wait_for(server_client.connected_event.wait(), timeout=30.0)
+
+                if not server_client.initialized or server_client.connection_state != ConnectionState.CONNECTED:
+                    raise MCPConnectionError(f"Failed to connect to {server_name}")
+
+                # Record success
+                self._circuit_breaker.record_success(server_name)
+                logger.info(f"✅ MCP server '{server_name}' connected successfully!")
+                return True
+
+            except Exception as e:
+                self._circuit_breaker.record_failure(server_name)
+                server_client.connection_state = ConnectionState.FAILED
+                logger.error(f"Failed to connect to {server_name}: {e}")
+
+                # Cleanup manager task to prevent resource leak
+                if server_client.manager_task and not server_client.manager_task.done():
+                    server_client.disconnect_event.set()
+                    try:
+                        await asyncio.wait_for(server_client.manager_task, timeout=5.0)
+                    except asyncio.TimeoutError:
+                        logger.warning(f"Manager task for {server_name} didn't shutdown gracefully, cancelling")
+                        server_client.manager_task.cancel()
+                        try:
+                            await server_client.manager_task
+                        except asyncio.CancelledError:
+                            pass
+                    except Exception as cleanup_error:
+                        logger.error(f"Error cleaning up manager task for {server_name}: {cleanup_error}")
+                    finally:
+                        server_client.manager_task = None
+
+                return False
+
+    async def _connect_single(self) -> None:
+        """Connect to single server."""
+        config = self._server_configs[0]
+        server_name = config["name"]
+
+        success = await self._connect_server(server_name, config)
+        if not success:
+            raise MCPConnectionError(f"Failed to connect to {server_name}")
+
+    async def _connect_all_parallel(self) -> None:
+        """Connect to all servers in parallel."""
+        tasks = [self._connect_server(c["name"], c) for c in self._server_configs]
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+
+        # Log results
+        successful = sum(1 for r in results if r is True)
+        logger.info(f"Connected to {successful}/{len(self._server_configs)} servers")
+
+    def _create_transport_context(self, config: Dict[str, Any]):
+        """Create the appropriate transport context manager based on config."""
+        transport_type = config.get("type", "stdio")
+        server_name = config["name"]
+
+        if transport_type == "stdio":
+            command = config.get("command", [])
+            args = config.get("args", [])
+
+            logger.debug(f"Setting up stdio transport for {server_name}: command={command}, args={args}")
+
+            # Handle command preparation
+            if isinstance(command, str):
+                full_command = prepare_command(command)
+                if args:
+                    full_command.extend(args)
+            elif isinstance(command, list):
+                full_command = command + (args or [])
+            else:
+                full_command = args or []
+
+            if not full_command:
+                raise MCPConnectionError(f"No command specified for stdio transport in {server_name}")
+
+            # Merge provided env with system env
+            env = config.get("env", {})
+            if env:
+                env = {**get_default_environment(), **env}
+            else:
+                env = get_default_environment()
+
+            # Perform environment variable substitution for args
+            substituted_args = []
+            for arg in full_command[1:] if len(full_command) > 1 else []:
+                if isinstance(arg, str):
+                    try:
+                        substituted_args.append(substitute_env_variables(arg))
+                    except ValueError as e:
+                        raise MCPConnectionError(f"Environment variable substitution failed in args: {e}", server_name=server_name) from e
+                else:
+                    substituted_args.append(arg)
+
+            # Perform environment variable substitution for env dict
+            for key, value in list(env.items()):
+                if isinstance(value, str):
+                    try:
+                        env[key] = substitute_env_variables(value)
+                    except ValueError as e:
+                        raise MCPConnectionError(f"Environment variable substitution failed for {key}: {e}", server_name=server_name) from e
+
+            # Extract cwd if provided in config
+            cwd = config.get("cwd")
+
+            server_params = StdioServerParameters(
+                command=full_command[0],
+                args=substituted_args,
+                env=env,
+                cwd=cwd,
+            )
+
+            # Open errlog file to redirect MCP server stderr output
+            from ..logger_config import get_log_session_dir
+
+            log_dir = get_log_session_dir()
+            errlog_path = log_dir / f"mcp_{server_name}_stderr.log"
+            errlog_file = open(errlog_path, "w", encoding="utf-8")
+
+            # Store errlog file handle for cleanup
+            if not hasattr(self, "_errlog_files"):
+                self._errlog_files = {}
+            self._errlog_files[server_name] = errlog_file
+
+            return stdio_client(server_params, errlog=errlog_file)
+
+        elif transport_type == "streamable-http":
+            url = config["url"]
+            headers = config.get("headers", {})
+
+            # Perform environment variable substitution for headers
+            substituted_headers = {}
+            for key, value in headers.items():
+                if isinstance(value, str):
+                    try:
+                        substituted_headers[key] = substitute_env_variables(value)
+                    except ValueError as e:
+                        raise MCPConnectionError(f"Environment variable substitution failed in header {key}: {e}", server_name=server_name) from e
+                else:
+                    substituted_headers[key] = value
+
+            timeout_raw = config.get("timeout", self.timeout_seconds)
+            http_read_timeout_raw = config.get("http_read_timeout", 60 * 5)
+
+            timeout = _ensure_timedelta(timeout_raw, self.timeout_seconds)
+            http_read_timeout = _ensure_timedelta(http_read_timeout_raw, 60 * 5)
+
+            return streamablehttp_client(
+                url=url,
+                headers=substituted_headers,
+                timeout=timeout,
+                sse_read_timeout=http_read_timeout,
+            )
+        else:
+            raise MCPConnectionError(f"Unsupported transport type: {transport_type}")
+
+    async def _run_manager(self, server_name: str, config: Dict[str, Any]) -> None:
+        """Background task that owns the transport and session contexts for a server."""
+        server_client = self._server_clients[server_name]
+        connection_successful = False
+
+        try:
+            transport_ctx = self._create_transport_context(config)
+
+            async with transport_ctx as session_params:
+                read, write = session_params[0:2]
+
+                session_timeout_timedelta = _ensure_timedelta(self.timeout_seconds, 30.0)
+
+                async with ClientSession(read, write, read_timeout_seconds=session_timeout_timedelta) as session:
+                    # Initialize and expose session
+                    server_client.session = session
+                    await session.initialize()
+                    await self._discover_capabilities(server_name, config)
+                    server_client.initialized = True
+                    server_client.connection_state = ConnectionState.CONNECTED
+                    connection_successful = True
+                    server_client.connected_event.set()
+
+                    logger.info(f"✅ MCP server '{server_name}' connected successfully!")
+
+                    # Send connected status if callback is available
+                    if self.status_callback:
+                        await self.status_callback(
+                            "connected",
+                            {
+                                "server": server_name,
+                                "message": f"Server '{server_name}' ready",
+                            },
+                        )
+
+                    # Wait until disconnect is requested
+                    await server_client.disconnect_event.wait()
+
+        except Exception as e:
+            logger.error(f"MCP manager error for {server_name}: {e}", exc_info=True)
+
+            if self.status_callback:
+                await self.status_callback(
+                    "error",
+                    {
+                        "server": server_name,
+                        "message": f"Failed to connect to MCP server '{server_name}': {e}",
+                        "error": str(e),
+                    },
+                )
+
+            if not server_client.connected_event.is_set():
+                server_client.connected_event.set()
+        finally:
+            # Clear session state
+            server_client.initialized = False
+            server_client.session = None
+            if not connection_successful:
+                server_client.connection_state = ConnectionState.FAILED
+                if not server_client.connected_event.is_set():
+                    server_client.connected_event.set()
+            else:
+                server_client.connection_state = ConnectionState.DISCONNECTED
+
+    async def _discover_capabilities(self, server_name: str, config: Dict[str, Any]) -> None:
+        """Discover server capabilities (tools, resources, prompts) with name prefixing for multi-server."""
+        logger.debug(f"Discovering capabilities for {server_name}")
+
+        session = self._get_server_session(server_name)
+
+        try:
+            # Combine backend-level and per-server tool filtering
+            server_exclude = config.get("exclude_tools", [])
+            combined_exclude = list(set((self.exclude_tools or []) + server_exclude))
+
+            server_allowed = config.get("allowed_tools")
+            combined_allowed = server_allowed if server_allowed is not None else self.allowed_tools
+
+            # List tools
+            available_tools = await session.list_tools()
+            tools_list = getattr(available_tools, "tools", []) if available_tools else []
+
+            for tool in tools_list:
+                if combined_exclude and tool.name in combined_exclude:
+                    continue
+                if combined_allowed is None or tool.name in combined_allowed:
+                    # Always apply name prefixing for consistency
+                    prefixed_name = sanitize_tool_name(tool.name, server_name)
+
+                    self.tools[prefixed_name] = tool
+                    self._tool_to_server[prefixed_name] = server_name
+
+            logger.info(f"Discovered capabilities for {server_name}: " f"{len([t for t, s in self._tool_to_server.items() if s == server_name])} tools")
+
+        except Exception as e:
+            logger.error(f"Failed to discover server capabilities for {server_name}: {e}", exc_info=True)
+            raise MCPConnectionError(f"Failed to discover server capabilities: {e}") from e
+
+    async def disconnect(self) -> None:
+        """Disconnect from all MCP servers."""
+        if not self._initialized:
+            return
+
+        # Disconnect all servers (works for single or multiple)
+        tasks = [self._disconnect_one(name, client) for name, client in self._server_clients.items() if client.connection_state != ConnectionState.DISCONNECTED]
+
+        if tasks:
+            await asyncio.gather(*tasks, return_exceptions=True)
+
+        self._initialized = False
+
+    async def _disconnect_one(self, server_name: str, server_client: _ServerClient) -> None:
+        """Disconnect a single server."""
+        server_client.connection_state = ConnectionState.DISCONNECTING
+
+        if server_client.manager_task and not server_client.manager_task.done():
+            server_client.disconnect_event.set()
+            try:
+                await asyncio.wait_for(server_client.manager_task, timeout=5.0)
+            except asyncio.TimeoutError:
+                logger.warning(f"Manager task for {server_name} didn't shutdown gracefully, cancelling")
+                server_client.manager_task.cancel()
+                try:
+                    await server_client.manager_task
+                except asyncio.CancelledError:
+                    logger.debug(f"Manager task for {server_name} cancelled successfully")
+            except Exception as e:
+                logger.error(f"Error during manager task shutdown for {server_name}: {e}")
+            finally:
+                server_client.manager_task = None
+
+        server_client.initialized = False
+        server_client.connection_state = ConnectionState.DISCONNECTED
+
+    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
+        """
+        Call an MCP tool with validation and timeout handling.
+
+        Args:
+            tool_name: Name of the tool to call (always prefixed as mcp__server__toolname)
+            arguments: Tool arguments
+
+        Returns:
+            Tool execution result
+
+        Raises:
+            MCPError: If tool is not available
+            MCPConnectionError: If no active session
+            MCPValidationError: If arguments are invalid
+            MCPTimeoutError: If tool call times out
+            MCPServerError: If tool execution fails
+        """
+        if tool_name not in self.tools:
+            available_tools = list(self.tools.keys())
+            raise MCPError(
+                f"Tool '{tool_name}' not available",
+                context={"available_tools": available_tools, "total": len(available_tools)},
+            )
+
+        # Validate tool arguments
+        try:
+            validated_arguments = validate_tool_arguments(arguments)
+        except ValueError as e:
+            raise MCPValidationError(
+                f"Invalid tool arguments: {e}",
+                field="arguments",
+                value=arguments,
+                context={"tool_name": tool_name},
+            ) from e
+
+        # Execute pre-tool hooks
+        pre_tool_hooks = self.hooks.get(HookType.PRE_TOOL_USE, [])
+        for hook in pre_tool_hooks:
+            try:
+                allowed = await hook(tool_name, validated_arguments)
+                if not allowed:
+                    raise MCPValidationError(
+                        "Tool call blocked by pre-tool hook",
+                        field="tool_name",
+                        value=tool_name,
+                        context={"arguments": validated_arguments},
+                    )
+            except Exception as e:
+                if isinstance(e, MCPValidationError):
+                    raise
+                logger.warning(f"Pre-tool hook error for {tool_name}: {e}", exc_info=True)
+
+        # Extract server name from prefixed tool name (always prefixed)
+        server_name = self._tool_to_server.get(tool_name)
+        if not server_name:
+            raise MCPError(f"Tool '{tool_name}' not mapped to any server")
+
+        # Extract original tool name (remove prefix - always prefixed)
+        original_tool_name = tool_name[len(f"mcp__{server_name}__") :]
+
+        session = self._get_server_session(server_name)
+
+        logger.debug(f"Calling tool {original_tool_name} on {server_name} with arguments: {validated_arguments}")
+
+        # Send tool call start status if callback is available
+        if self.status_callback:
+            await self.status_callback(
+                "tool_call_start",
+                {
+                    "server": server_name,
+                    "tool": original_tool_name,
+                    "message": f"Calling tool '{original_tool_name}' on server '{server_name}'",
+                    "arguments": validated_arguments,
+                },
+            )
+
+        try:
+            # Add timeout to tool calls
+            result = await asyncio.wait_for(
+                session.call_tool(original_tool_name, validated_arguments),
+                timeout=self.timeout_seconds,
+            )
+            logger.debug(f"Tool {original_tool_name} completed successfully on {server_name}")
+
+            # Send tool call success status if callback is available
+            if self.status_callback:
+                await self.status_callback(
+                    "tool_call_success",
+                    {
+                        "server": server_name,
+                        "tool": original_tool_name,
+                        "message": f"Tool '{original_tool_name}' executed successfully",
+                    },
+                )
+
+            return result
+
+        except asyncio.TimeoutError:
+            if self.status_callback:
+                await self.status_callback(
+                    "tool_call_timeout",
+                    {
+                        "server": server_name,
+                        "tool": original_tool_name,
+                        "message": f"Tool '{original_tool_name}' timed out after {self.timeout_seconds} seconds",
+                        "timeout": self.timeout_seconds,
+                    },
+                )
+
+            # Record failure with circuit breaker
+            self._circuit_breaker.record_failure(server_name)
+
+            raise MCPTimeoutError(
+                f"Tool call timed out after {self.timeout_seconds} seconds",
+                timeout_seconds=self.timeout_seconds,
+                operation=f"call_tool({original_tool_name})",
+                context={"tool_name": original_tool_name, "server_name": server_name},
+            )
+        except Exception as e:
+            logger.error(f"Tool call failed for {original_tool_name} on {server_name}: {e}", exc_info=True)
+
+            # Record failure with circuit breaker
+            self._circuit_breaker.record_failure(server_name)
+
+            if self.status_callback:
+                await self.status_callback(
+                    "tool_call_error",
+                    {
+                        "server": server_name,
+                        "tool": original_tool_name,
+                        "message": f"Tool '{original_tool_name}' failed: {e}",
+                        "error": str(e),
+                    },
+                )
+
+            raise MCPServerError(
+                f"Tool call failed: {e}",
+                server_name=server_name,
+                context={"tool_name": original_tool_name, "arguments": validated_arguments},
+            ) from e
+
+    def get_available_tools(self) -> List[str]:
+        """Get list of available tool names."""
+        return list(self.tools.keys())
+
+    def is_connected(self) -> bool:
+        """Check if any servers are connected."""
+        return self._initialized and any(sc.initialized for sc in self._server_clients.values())
+
+    def get_server_names(self) -> List[str]:
+        """Get list of connected server names."""
+        return [name for name, sc in self._server_clients.items() if sc.initialized]
+
+    def get_active_sessions(self) -> List[ClientSession]:
+        """Return active MCP ClientSession objects for all connected servers."""
+        sessions = []
+        for server_client in self._server_clients.values():
+            if server_client.session is not None and server_client.initialized:
+                sessions.append(server_client.session)
+        return sessions
+
+    async def health_check_all(self) -> Dict[str, bool]:
+        """
+        Perform health check on all connected MCP servers.
+
+        Returns:
+            Dictionary mapping server names to health status
+        """
+        health_status = {}
+
+        for server_name, server_client in self._server_clients.items():
+            if not server_client.initialized or not server_client.session:
+                health_status[server_name] = False
+                continue
+
+            try:
+                await server_client.session.list_tools()
+                health_status[server_name] = True
+            except Exception as e:
+                logger.warning(f"Health check failed for {server_name}: {e}")
+                health_status[server_name] = False
+
+        return health_status
+
+    async def health_check(self) -> bool:
+        """
+        Perform a health check on all servers.
+
+        Returns:
+            True if all connected servers are healthy, False otherwise
+        """
+        health_status = await self.health_check_all()
+        return all(health_status.values()) if health_status else False
+
+    async def _reconnect_failed_servers(self, max_retries: int = 3) -> Dict[str, bool]:
+        """
+        Attempt to reconnect any failed servers with circuit breaker integration.
+
+        Args:
+            max_retries: Maximum number of reconnection attempts per server
+
+        Returns:
+            Dictionary mapping server names to reconnection success status
+        """
+        health_status = await self.health_check_all()
+        reconnect_results = {}
+
+        for server_name, is_healthy in health_status.items():
+            if not is_healthy:
+                # Check circuit breaker before reconnecting
+                if self._circuit_breaker.should_skip_server(server_name):
+                    logger.warning(f"Skipping reconnection for {server_name} due to circuit breaker")
+                    reconnect_results[server_name] = False
+                    continue
+
+                logger.info(f"Attempting to reconnect failed server: {server_name}")
+
+                # Find the config for this server
+                config = next((c for c in self._server_configs if c["name"] == server_name), None)
+                if not config:
+                    reconnect_results[server_name] = False
+                    continue
+
+                success = False
+                for attempt in range(max_retries):
+                    try:
+                        if attempt > 0:
+                            await asyncio.sleep(1.0 * (2**attempt))  # Exponential backoff
+
+                        # Disconnect first
+                        server_client = self._server_clients[server_name]
+                        await self._disconnect_one(server_name, server_client)
+
+                        # Reconnect
+                        server_client.connected_event = asyncio.Event()
+                        server_client.disconnect_event = asyncio.Event()
+                        server_client.manager_task = asyncio.create_task(
+                            self._run_manager(server_name, config),
+                        )
+                        await asyncio.wait_for(server_client.connected_event.wait(), timeout=30.0)
+
+                        if server_client.initialized:
+                            self._circuit_breaker.record_success(server_name)
+                            success = True
+                            logger.info(f"Successfully reconnected server: {server_name}")
+                            break
+                    except Exception as e:
+                        logger.warning(f"Reconnection attempt {attempt + 1} failed for {server_name}: {e}")
+                        self._circuit_breaker.record_failure(server_name)
+
+                reconnect_results[server_name] = success
+            else:
+                reconnect_results[server_name] = True
+
+        return reconnect_results
+
+    async def reconnect(self, max_retries: int = 3) -> bool:
+        """
+        Attempt to reconnect all servers with circuit breaker integration.
+
+        Args:
+            max_retries: Maximum number of reconnection attempts
+                Uses exponential backoff between retries: 2s, 4s, 8s, 16s...
+
+        Returns:
+            True if all reconnections successful, False otherwise
+        """
+        results = await self._reconnect_failed_servers(max_retries)
+        return all(results.values()) if results else False
+
+    async def _cleanup(self) -> None:
+        """Comprehensive cleanup of all resources."""
+        async with self._cleanup_lock:
+            if self._cleanup_done:
+                return
+
+            logger.debug("Starting cleanup for MCPClient")
+
+            try:
+                # Disconnect all servers
+                await self.disconnect()
+
+                # Close errlog files
+                if hasattr(self, "_errlog_files"):
+                    for server_name, errlog_file in self._errlog_files.items():
+                        try:
+                            errlog_file.close()
+                        except Exception as e:
+                            logger.debug(f"Error closing errlog file for {server_name}: {e}")
+                    self._errlog_files.clear()
+
+                # Clear all references
+                self.tools.clear()
+                self._tool_to_server.clear()
+
+                self._cleanup_done = True
+                logger.debug("Cleanup completed for MCPClient")
+
+            except Exception as e:
+                logger.error(f"Error during cleanup: {e}")
+                raise
+
+    async def __aenter__(self) -> "MCPClient":
+        """Async context manager entry."""
+        self._context_managed = True
+        await self.connect()
+        return self
+
+    async def __aexit__(
+        self,
+        _exc_type: Optional[type],
+        _exc_val: Optional[BaseException],
+        _exc_tb: Optional[TracebackType],
+    ) -> None:
+        """Async context manager exit."""
+        try:
+            await self._cleanup()
+        except Exception as e:
+            logger.error(f"Error during context manager cleanup: {e}")
+        finally:
+            self._context_managed = False
+
+    @classmethod
+    async def create_and_connect(
+        cls,
+        server_configs: List[Dict[str, Any]],
+        *,
+        timeout_seconds: int = 30,
+        allowed_tools: Optional[List[str]] = None,
+        exclude_tools: Optional[List[str]] = None,
+    ) -> "MCPClient":
+        """
+        Create and connect MCP client in one step.
+
+        Args:
+            server_configs: List of server configuration dictionaries
+            timeout_seconds: Timeout for operations in seconds
+            allowed_tools: Optional list of tool names to include
+            exclude_tools: Optional list of tool names to exclude
+
+        Returns:
+            Connected MCPClient instance
+        """
+        client = cls(
+            server_configs,
+            timeout_seconds=timeout_seconds,
+            allowed_tools=allowed_tools,
+            exclude_tools=exclude_tools,
+        )
+        await client.connect()
+        return client