aip-agents-binary 0.0.0__py3-none-any.whl

aip_agents/mcp/client/base_mcp_client.py

@@ -0,0 +1,377 @@
+"""Base MCP Client for aip-agents with Session Persistence.
+
+This base class provides persistent session management for MCP connections and serves
+as the foundation for SDK-specific MCP clients (LangChain, Google ADK, etc.). It handles
+session pooling, initialization, tool discovery, and cleanup while leaving SDK-specific
+tool conversion to subclasses.
+
+Authors:
+    Putu Ravindra Wiguna (putu.r.wiguna@gdplabs.id)
+"""
+
+import asyncio
+from abc import ABC, abstractmethod
+from typing import Any
+
+from gllm_tools.mcp.client.client import MCPClient
+from gllm_tools.mcp.client.config import MCPConfiguration
+from gllm_tools.mcp.client.resource import MCPResource
+from mcp.types import CallToolResult, Tool
+
+from aip_agents.mcp.client.persistent_session import PersistentMCPSession
+from aip_agents.mcp.client.session_pool import MCPSessionPool
+from aip_agents.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class BaseMCPClient(MCPClient, ABC):
+    """Base MCP Client with persistent session management for aip-agents.
+
+    This class provides:
+    - Persistent session management across tool calls
+    - One-time tool registration and caching
+    - Automatic connection health monitoring and reconnection
+    - Centralized cleanup of all MCP resources
+    - Generic tool discovery from all configured MCP servers
+
+    Subclasses should implement SDK-specific tool conversion in _process_tool() method.
+    """
+
+    def __init__(self, servers: dict[str, MCPConfiguration]):
+        """Initialize the base MCP client with session pool.
+
+        Args:
+            servers (dict[str, MCPConfiguration]): Dictionary of MCP server configurations by server name
+        """
+        super().__init__(servers)
+        self.session_pool = MCPSessionPool()
+        self._tools_cache: list[Tool] = []
+        self._initialized = False
+        # Guard against concurrent initialization
+        self._init_lock: asyncio.Lock = asyncio.Lock()
+
+    async def initialize(self) -> None:
+        """Initialize all MCP sessions and cache tools once.
+
+        This method is idempotent and only performs initialization if not already done.
+        It establishes persistent connections to all configured MCP servers and caches
+        available tools for efficient access.
+
+        Raises:
+            Exception: If any session initialization fails
+        """
+        if self._initialized:
+            logger.debug("BaseMCPClient already initialized, skipping")
+            return
+
+        # Ensure only one concurrent initializer proceeds
+        async with self._init_lock:
+            if self._initialized:
+                logger.debug("BaseMCPClient already initialized (post-lock), skipping")
+                return
+
+            logger.info(f"Initializing BaseMCPClient with {len(self.servers)} MCP servers")
+
+            try:
+                # Initialize all persistent sessions concurrently
+                await self.session_pool.initialize_all_sessions(self.servers)
+
+                # Cache all available tools from all sessions
+                self._tools_cache = await self.session_pool.get_all_tools()
+
+                self._initialized = True
+                logger.info(
+                    f"BaseMCPClient initialization complete: "
+                    f"{len(self._tools_cache)} tools cached from {self.session_pool.session_count} sessions"
+                )
+            except asyncio.CancelledError:
+                logger.info("BaseMCPClient initialization cancelled; running cleanup")
+                await self.cleanup()
+                raise
+            except Exception as e:
+                logger.error(f"Failed to initialize BaseMCPClient: {e}", exc_info=True)
+                # Cleanup any partially initialized sessions on failure
+                await self.cleanup()
+                raise
+
+    @abstractmethod
+    async def get_tools(self, server: str | None = None) -> list[Any]:
+        """Get framework-specific tools from MCP servers.
+
+        This method must be implemented by subclasses to provide framework-specific
+        tool conversion (e.g., StructuredTool for LangChain, FunctionTool for Google ADK).
+
+        Args:
+            server (str | None): Optional server name to filter tools from a specific server.
+                    If None, returns tools from all configured servers.
+
+        Returns:
+            list[Any]: List of framework-specific tool objects.
+        """
+        pass  # pragma: no cover  # Abstract method - cannot be executed directly
+
+    async def get_raw_mcp_tools(self, server: str | None = None) -> list[Tool]:
+        """Get raw MCP tools - for subclasses to perform framework-specific conversions.
+
+        This method provides access to the cached raw MCP Tool objects.
+        Subclasses use this to convert to framework-specific tools.
+
+        Args:
+            server (str | None): Optional server name to filter tools from a specific server.
+                    If None, returns tools from all configured servers.
+
+        Returns:
+            list[Tool]: List of raw MCP Tool objects. Empty list if not initialized or no tools available.
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        if server:
+            if server not in self.servers:
+                logger.warning(f"Server '{server}' not found in configuration")
+                return []
+
+            try:
+                session = self.session_pool.get_session(server)
+                return await session.list_tools()
+            except KeyError:
+                logger.warning(f"No active session found for server '{server}'")
+                return []
+            except Exception as e:
+                logger.error(f"Failed to get tools from server '{server}': {e}")
+                return []
+        else:
+            return self._tools_cache.copy()
+
+    def get_tools_count(self, server: str | None = None) -> int:
+        """Get count of raw MCP tools without expensive copying.
+
+        This is an efficient way to get tool counts for logging/metrics
+        without the overhead of copying tool lists.
+
+        Args:
+            server (str | None): Optional server name to filter tools from a specific server.
+                    If None, returns count from all configured servers.
+
+        Returns:
+            int: Count of raw MCP tools available.
+        """
+        if not self._initialized:
+            return 0
+
+        if server:
+            if server not in self.servers:
+                return 0
+            try:
+                session = self.session_pool.get_session(server)
+                return session.get_tools_count()
+            except KeyError:
+                return 0
+            except Exception as e:
+                logger.error(f"Failed to get tools count from server '{server}': {e}")
+                return 0
+        else:
+            return len(self._tools_cache)
+
+    async def get_session(self, server_name: str) -> PersistentMCPSession:
+        """Get a persistent session for a specific server.
+
+        Args:
+            server_name (str): The name of the MCP server
+
+        Returns:
+            PersistentMCPSession: Persistent MCP session for the specified server
+
+        Raises:
+            KeyError: If the server is not configured or no active session exists
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        try:
+            return self.session_pool.get_session(server_name)
+        except KeyError:
+            logger.error(f"Server '{server_name}' not found in configuration")
+            raise
+        except Exception as e:
+            logger.error(f"Failed to get session for server '{server_name}': {e}")
+            raise
+
+    async def call_tool(self, server_name: str, tool_name: str, arguments: dict[str, Any]) -> CallToolResult:
+        """Execute a tool on a specific MCP server using persistent session.
+
+        Args:
+            server_name (str): The MCP server to execute the tool on
+            tool_name (str): The name of the tool to execute
+            arguments (dict[str, Any]): Arguments for the tool execution
+
+        Returns:
+            CallToolResult: The result of the tool execution
+
+        Raises:
+            KeyError: If the server doesn't exist
+            Exception: If tool execution fails
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        try:
+            session = self.session_pool.get_session(server_name)
+            return await session.call_tool(tool_name, arguments)
+        except KeyError:
+            logger.error(f"Server '{server_name}' not found in configuration")
+            raise
+        except Exception as e:
+            logger.error(f"Failed to call tool '{tool_name}' on server '{server_name}': {e}")
+            raise
+
+    async def read_resource(self, server_name: str, resource_uri: str) -> Any:
+        """Read an MCP resource from a specific server using persistent session.
+
+        Args:
+            server_name (str): The MCP server to read the resource from
+            resource_uri (str): The URI of the resource to read
+
+        Returns:
+            Any: The resource content
+
+        Raises:
+            KeyError: If the server doesn't exist
+            Exception: If resource reading fails
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        try:
+            session = self.session_pool.get_session(server_name)
+            return await session.read_resource(resource_uri)
+        except KeyError:
+            logger.error(f"Server '{server_name}' not found in configuration")
+            raise
+        except Exception as e:
+            logger.error(f"Failed to read resource '{resource_uri}' from server '{server_name}': {e}")
+            raise
+
+    async def cleanup(self) -> None:
+        """Clean up all MCP resources and close sessions.
+
+        This method properly closes all persistent sessions and cleans up resources.
+        It can be called multiple times safely.
+        """
+        should_cleanup_pool = self._initialized or self.session_pool.session_count > 0
+
+        if should_cleanup_pool:
+            logger.info("Cleaning up BaseMCPClient resources")
+            cleanup_error = None
+            try:
+                await self.session_pool.close_all_sessions()
+            except TimeoutError as e:
+                cleanup_error = e
+                logger.warning(f"Timeout during BaseMCPClient cleanup: {e}")
+            except Exception as e:
+                cleanup_error = e
+                logger.error(f"Unexpected error during BaseMCPClient cleanup: {e}", exc_info=True)
+                # Don't re-raise - cleanup should be best-effort
+            finally:
+                # Always clear cache and mark as not initialized, even if session cleanup failed
+                self._tools_cache.clear()
+                self._initialized = False
+                if cleanup_error:
+                    logger.info("BaseMCPClient cleanup completed (session cleanup had errors)")
+                else:
+                    logger.info("BaseMCPClient cleanup completed successfully")
+        else:
+            self._tools_cache.clear()
+            self._initialized = False
+            logger.debug("BaseMCPClient cleanup called with no active sessions, skipping pool close")
+
+    @property
+    def is_initialized(self) -> bool:
+        """Check if the client is fully initialized.
+
+        Returns:
+            bool: True if sessions are initialized and tools are cached, False otherwise
+        """
+        return self._initialized
+
+    @property
+    def active_sessions(self) -> dict[str, PersistentMCPSession]:
+        """Get all active persistent sessions.
+
+        Returns:
+            dict[str, PersistentMCPSession]: Dictionary of active sessions by server name
+        """
+        if not self._initialized:
+            return {}
+        return self.session_pool.get_all_active_sessions()
+
+    def get_session_count(self) -> int:
+        """Get the number of active MCP sessions.
+
+        Returns:
+            Number of active persistent sessions
+        """
+        return self.session_pool.session_count
+
+    def _get_annotated_server(self, resource: MCPResource) -> str | None:
+        """Extract server name from resource annotations if valid.
+
+        Args:
+            resource (MCPResource): The MCP resource.
+
+        Returns:
+            str | None: Annotated server name if valid and configured, else None.
+        """
+        if not self.servers:
+            return None
+
+        annotations = getattr(resource, "annotations", None)
+        if annotations and hasattr(annotations, "server") and annotations.server:
+            annotated = annotations.server
+            if annotated in self.servers:
+                return annotated
+            logger.warning(f"Resource annotation requested server '{annotated}' which is not configured; falling back")
+        return None
+
+    def _get_default_server(self) -> str:
+        """Get the default server name (first configured server).
+
+        Returns:
+            str: First server name.
+
+        Raises:
+            ValueError: If no servers configured.
+        """
+        if not self.servers:
+            raise ValueError("No MCP servers configured; cannot resolve server for resource")
+
+        server_list = list(self.servers.keys())
+        if len(server_list) > 1:
+            logger.warning(
+                "Multiple MCP servers configured but resource provided no annotation; "
+                f"defaulting to first server '{server_list[0]}'"
+            )
+        return server_list[0]
+
+    def _determine_server_name_for_resource(self, resource: MCPResource) -> str:
+        """Determine the server name for an MCP resource.
+
+        Args:
+            resource (MCPResource): The MCP resource to determine server for.
+
+        Returns:
+            str: Server name string.
+
+        Raises:
+            ValueError: If cannot determine server.
+        """
+        try:
+            annotated = self._get_annotated_server(resource)
+            if annotated:
+                return annotated
+            return self._get_default_server()
+        except Exception as e:
+            error_msg = f"Failed to determine server for resource {resource.uri}: {e}"
+            logger.error(error_msg)
+            raise ValueError(error_msg) from e

aip_agents/mcp/client/connection_manager.py

@@ -0,0 +1,231 @@
+"""MCP Connection Manager for persistent connection lifecycle management.
+
+This module implements the connection manager pattern inspired by mcp-use library
+to avoid cancel scope issues and provide persistent connections.
+
+Authors:
+    Putu Ravindra Wiguna (putu.r.wiguna@gdplabs.id)
+"""
+
+import asyncio
+import inspect
+from typing import Any
+
+from gllm_tools.mcp.client.config import MCPConfiguration
+
+from aip_agents.mcp.client.errors import MCPServerStartupError
+from aip_agents.mcp.client.transports import TransportType, create_transport
+from aip_agents.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class MCPConnectionManager:
+    """Manages MCP connection lifecycle following mcp-use patterns.
+
+    This connection manager handles the transport connection lifecycle in a background
+    task to avoid cancel scope issues during cleanup. It supports automatic transport
+    negotiation (HTTP -> SSE fallback) and graceful shutdown. Invalid explicit transports
+    are normalized via aliases (e.g., 'streamable_http' -> 'http') or fall back to
+    auto-detection with a warning.
+    """
+
+    TRANSPORT_ALIASES = {
+        "http": TransportType.HTTP,
+        "streamable-http": TransportType.HTTP,
+        "sse": TransportType.SSE,
+        "stdio": TransportType.STDIO,
+    }
+
+    def __init__(self, server_name: str, config: MCPConfiguration):
+        """Initialize connection manager.
+
+        Args:
+            server_name (str): Name of the MCP server
+            config (MCPConfiguration): MCP server configuration
+        """
+        self.server_name = server_name
+        self.config = config
+        self._task = None
+        self._connection = None
+        self._transport = None
+        self._stop_event = asyncio.Event()
+        self._ready_event = asyncio.Event()
+        self._done_event = asyncio.Event()
+        self._exception = None
+        self.transport_type = None
+        # Configurable retry settings (MCP-specific, defaults to reasonable values)
+        self.max_retries = config.get("max_retries", 3)
+        self.initial_retry_delay = config.get("initial_retry_delay", 1.0)
+
+    async def start(self) -> tuple[Any, Any]:
+        """Start connection in background task.
+
+        For HTTP/SSE transports, establishes connection directly to avoid anyio context issues.
+        For stdio transport, uses background task to manage subprocess lifecycle.
+
+        Returns:
+            tuple[Any, Any]: Tuple of (read_stream, write_stream) for ClientSession
+
+        Raises:
+            Exception: If connection establishment fails
+        """
+        logger.debug(f"Starting connection manager for {self.server_name}")
+
+        # Determine transport type first
+        self.transport_type = self._get_transport_type()
+
+        # For HTTP/SSE: connect directly (no background task needed)
+        # This avoids anyio.BrokenResourceError when streams cross task boundaries
+        if self.transport_type in (TransportType.HTTP, TransportType.SSE):
+            await self._establish_connection()
+            return self._connection
+
+        # For stdio: use background task to manage subprocess
+        self._task = asyncio.create_task(self._connection_task())
+        await self._ready_event.wait()
+
+        if self._exception:
+            raise self._exception
+
+        return self._connection
+
+    async def stop(self) -> None:
+        """Stop connection gracefully."""
+        logger.debug(f"Stopping connection manager for {self.server_name}")
+
+        # For HTTP/SSE (no background task), just close transport
+        if self.transport_type in (TransportType.HTTP, TransportType.SSE):
+            if self._transport:
+                try:
+                    close_result = self._transport.close()
+                    if inspect.isawaitable(close_result):
+                        await close_result
+                except Exception as exc:
+                    logger.warning(f"Failed to close transport cleanly for {self.server_name}: {exc}")
+            self._connection = None
+            return
+
+        # For stdio (with background task), wait for task to finish
+        if self._task and not self._task.done():
+            self._stop_event.set()
+            try:
+                await asyncio.wait_for(self._task, timeout=5.0)
+            except TimeoutError:
+                logger.warning(f"Connection manager for {self.server_name} did not stop gracefully")
+                self._task.cancel()
+        await self._done_event.wait()
+
+    @property
+    def is_connected(self) -> bool:
+        """Check if connection is active.
+
+        Returns:
+            bool: True if connected, False otherwise
+        """
+        # For HTTP/SSE (no background task), just check if connection exists
+        if self.transport_type in (TransportType.HTTP, TransportType.SSE):
+            return self._connection is not None
+
+        # For stdio (with background task), check task status too
+        return (
+            self._connection is not None
+            and self._task is not None
+            and not self._task.done()
+            and not self._stop_event.is_set()
+        )
+
+    def _auto_detect_transport_type(self) -> TransportType:
+        """Auto-detect transport type based on configuration.
+
+        Returns:
+            TransportType: Detected transport type
+        """
+        if "command" in self.config:
+            return TransportType.STDIO
+        elif "url" in self.config:
+            url = self.config["url"]
+            if url.endswith("/sse"):
+                return TransportType.SSE
+            else:
+                return TransportType.HTTP
+        else:
+            return TransportType.STDIO
+
+    def _get_transport_type(self) -> TransportType:
+        """Determine the transport type to use, prioritizing explicit config with aliases.
+
+        Returns:
+            TransportType: Transport type enum
+
+        Notes:
+            Invalid explicit transports trigger a warning and fallback to auto-detection; no exception raised.
+        """
+        explicit_transport = self.config.get("transport", "").lower().replace("_", "-")
+        if explicit_transport in self.TRANSPORT_ALIASES:
+            return self.TRANSPORT_ALIASES[explicit_transport]
+
+        if explicit_transport:
+            logger.warning(f"Unknown explicit transport '{explicit_transport}'. Falling back to auto-detection.")
+
+        return self._auto_detect_transport_type()
+
+    async def _establish_connection(self) -> None:
+        """Establish connection based on transport preference with fallback.
+
+        Uses configurable retries with exponential backoff for transient failures.
+
+        Raises:
+            ConnectionError: If all connection attempts fail
+        """
+        # transport_type may already be set by start() for HTTP/SSE
+        if not self.transport_type:
+            self.transport_type = self._get_transport_type()
+        details = f"URL: {self.config.get('url', 'N/A')}, Command: {self.config.get('command', 'N/A')}"
+        logger.info(f"Establishing connection to {self.server_name} via {self.transport_type} ({details})")
+
+        retry_delay = self.initial_retry_delay
+        for attempt in range(self.max_retries):
+            self._transport = create_transport(self.server_name, self.config, self.transport_type)
+            try:
+                read_stream, write_stream, _ = await self._transport.connect()
+                self._connection = (read_stream, write_stream)
+                logger.info(f"Connection established on attempt {attempt + 1}")
+                return
+            except ValueError:
+                # Config validation errors should not be retried
+                raise
+            except MCPServerStartupError:
+                raise
+            except Exception as e:
+                logger.warning(f"Connection attempt {attempt + 1} failed: {e}")
+                if attempt < self.max_retries - 1:
+                    await asyncio.sleep(retry_delay)
+                    retry_delay *= 2  # Exponential backoff
+                else:
+                    raise ConnectionError(
+                        f"Failed to establish connection to {self.server_name} "
+                        f"after {self.max_retries} attempts: {str(e)}"
+                    ) from e
+
+    async def _connection_task(self) -> None:
+        """Background task that manages the connection lifecycle."""
+        try:
+            await self._establish_connection()
+            self._ready_event.set()
+            await self._stop_event.wait()
+        except Exception as e:
+            logger.error(f"Connection failed for {self.server_name}: {e}", exc_info=True)
+            self._exception = e
+            self._ready_event.set()
+        finally:
+            if self._transport:
+                try:
+                    close_result = self._transport.close()
+                    if inspect.isawaitable(close_result):
+                        await close_result
+                except Exception as exc:  # pragma: no cover - defensive
+                    logger.warning("Failed to close transport cleanly for %s: %s", self.server_name, exc)
+            self._connection = None
+            self._done_event.set()
+            logger.debug(f"Connection manager cleanup complete for {self.server_name}")