aip-agents-binary 0.5.25b8__py3-none-any.whl → 0.6.0__py3-none-any.whl

aip_agents/mcp/client/persistent_session.py

@@ -8,6 +8,7 @@ Authors:
 """
 
 import asyncio
+from collections.abc import Awaitable, Callable
 from typing import Any
 
 from gllm_tools.mcp.client.config import MCPConfiguration
@@ -69,6 +70,13 @@ class PersistentMCPSession:
 
         self._initialized = False
         self._lock = asyncio.Lock()
+        self._owner_task: asyncio.Task | None = None
+        self._owner_ready: asyncio.Event = asyncio.Event()
+        self._owner_exception: Exception | None = None
+        self._timeout = float(config.get("timeout", 30.0))
+        self._request_queue: asyncio.Queue[
+            tuple[Callable[..., Awaitable[Any]], tuple[Any, ...], asyncio.Future, bool]
+        ] = asyncio.Queue()
 
     async def initialize(self) -> None:
         """Initialize session once and cache tools.
@@ -85,50 +93,28 @@ class PersistentMCPSession:
             # Double-check pattern
             if self._initialized:
                 return
+            if self._owner_task is None or self._owner_task.done():
+                self._owner_ready = asyncio.Event()
+                self._owner_exception = None
+                self._owner_task = asyncio.create_task(self._owner_loop())
 
-            try:
-                logger.info(f"Initializing persistent session for {self.server_name}")
-
-                # Start connection manager
-                read_stream, write_stream = await self.connection_manager.start()
-
-                # Create client session
-                self.client_session = ClientSession(read_stream, write_stream)
-                await self.client_session.__aenter__()
-
-                # MCP handshake
-                result = await self.client_session.initialize()
-                logger.debug(f"MCP handshake complete for {self.server_name}: {result.capabilities}")
-
-                # Discover and cache tools
-                if result.capabilities.tools:
-                    tools_result = await self.client_session.list_tools()
-                    self.tools = tools_result.tools if tools_result else []
-                    self._filtered_tools_cache = None  # Invalidate cache when tools change
-                    logger.info(f"Cached {len(self.tools)} tools for {self.server_name}")
-                else:
-                    logger.info(f"No tools available for {self.server_name}")
-
-                # Warn once per initialization if allowed_tools references unknown names
-                if self._allowed_tools_set:
-                    self._warn_on_unknown_allowed_tools(list(self._allowed_tools_set), self.tools)
-
-                # Discover resources (for future use)
-                if result.capabilities.resources:
-                    try:
-                        resources_result = await self.client_session.list_resources()
-                        if resources_result and resources_result.resources:
-                            logger.debug(f"Found {len(resources_result.resources)} resources for {self.server_name}")
-                    except Exception:
-                        logger.debug(f"Could not list resources for {self.server_name}, skipping")
-
-                self._initialized = True
-                logger.info(f"Session initialization complete for {self.server_name}")
-
-            except Exception as e:
-                logger.error(f"Failed to initialize session for {self.server_name}: {e}", exc_info=True)
-                await self._cleanup_on_error()
-                raise ConnectionError(f"Failed to initialize MCP session for {self.server_name}: {str(e)}") from e
+        try:
+            await asyncio.wait_for(self._owner_ready.wait(), timeout=self._timeout)
+        except asyncio.CancelledError:
+            if self._owner_task and not self._owner_task.done():
+                self._owner_task.cancel()
+            raise
+        except TimeoutError as e:
+            logger.error(f"Initialization timed out for {self.server_name} after {self._timeout}s")
+            if self._owner_task and not self._owner_task.done():
+                self._owner_task.cancel()
+            self._owner_exception = ConnectionError(
+                f"Initialization timed out for {self.server_name} after {self._timeout}s"
+            )
+            self._owner_ready.set()
+            raise self._owner_exception from e
+        if self._owner_exception:
+            raise self._owner_exception
 
     async def call_tool(self, name: str, arguments: dict[str, Any]) -> CallToolResult:
         """Call MCP tool using persistent session.
@@ -143,7 +129,19 @@ class PersistentMCPSession:
         Raises:
             Exception: If tool call fails
         """
-        await self.ensure_connected()
+        return await self._run_in_owner(self._call_tool_impl, name, arguments)
+
+    async def _call_tool_impl(self, name: str, arguments: dict[str, Any]) -> CallToolResult:
+        """Call MCP tool using the owner task.
+
+        Args:
+            name: Tool name.
+            arguments: Tool arguments.
+
+        Returns:
+            CallToolResult: Tool call result.
+        """
+        await self._ensure_connected_impl()
 
         if self._allowed_tools_set and name not in self._allowed_tools_set:
             allowed_display = ", ".join(sorted(self._allowed_tools_set))
@@ -173,10 +171,9 @@ class PersistentMCPSession:
         Raises:
             Exception: If resource reading fails
         """
-        await self.ensure_connected()
-        return await self._execute_read_resource(uri)
+        return await self._run_in_owner(self._execute_read_resource_impl, uri)
 
-    async def _execute_read_resource(self, uri: str) -> Any:
+    async def _execute_read_resource_impl(self, uri: str) -> Any:
         """Execute the reading of an MCP resource.
 
         Args:
@@ -188,6 +185,7 @@ class PersistentMCPSession:
         Raises:
             Exception: If resource reading fails
         """
+        await self._ensure_connected_impl()
         try:
             logger.debug(f"Reading resource '{uri}' on {self.server_name}")
             result = await self.client_session.read_resource(uri)
@@ -202,7 +200,15 @@ class PersistentMCPSession:
         Returns:
             list[Tool]: a copy of list of available tools, filtered to only allowed tools if configured
         """
-        await self.ensure_connected()
+        return await self._run_in_owner(self._list_tools_impl)
+
+    async def _list_tools_impl(self) -> list[Tool]:
+        """Return the cached tools list from the owner task.
+
+        Returns:
+            list[Tool]: Filtered tool list if allowed tools are configured, otherwise all tools.
+        """
+        await self._ensure_connected_impl()
 
         if not self._allowed_tools_set:
             return list(self.tools)
@@ -234,9 +240,13 @@ class PersistentMCPSession:
         Raises:
             Exception: If reconnection fails
         """
+        await self._run_in_owner(self._ensure_connected_impl)
+
+    async def _ensure_connected_impl(self) -> None:
+        """Ensure the session is connected, reconnecting if needed."""
         if not self._initialized or not self.connection_manager.is_connected:
             logger.info(f"Reconnecting session for {self.server_name}")
-            await self.initialize()
+            await self._initialize_impl()
 
     def _handle_connection_error(self, e: Exception, operation: str) -> None:
         """Handle connection-related errors with logging and reconnection marking.
@@ -255,29 +265,44 @@ class PersistentMCPSession:
         """Disconnect session gracefully.
 
         This method cleans up all resources and connections.
+        Always succeeds, even if the session was already in an error state.
         """
         logger.info(f"Disconnecting session for {self.server_name}")
 
-        async with self._lock:
+        if self._owner_task is None or self._owner_task.done():
+            await self._disconnect_impl()
+            return
+
+        try:
+            await self._run_in_owner(self._disconnect_impl, shutdown=True, ensure_initialized=False)
+        except ConnectionError:
+            # Owner task already failed; just clean up directly
+            logger.debug(f"Owner task already failed for {self.server_name}, cleaning up directly")
+            await self._disconnect_impl()
+        finally:
+            await self._await_owner_shutdown()
+
+    async def _await_owner_shutdown(self) -> None:
+        """Wait for the owner task to exit, cancelling on timeout."""
+        if not self._owner_task:
+            return
+
+        owner_task = self._owner_task
+        try:
+            await asyncio.wait_for(owner_task, timeout=self._timeout)
+        except TimeoutError:
+            logger.warning(f"Owner task for {self.server_name} did not exit within {self._timeout}s, cancelling")
+            owner_task.cancel()
             try:
-                # Close client session
-                if self.client_session:
-                    try:
-                        await self.client_session.__aexit__(None, None, None)
-                    except Exception as e:
-                        logger.warning(f"Error closing client session for {self.server_name}: {e}")
-                    self.client_session = None
-
-                # Stop connection manager
-                await self.connection_manager.stop()
-
-            except Exception as e:
-                logger.error(f"Error during disconnect for {self.server_name}: {e}")
-            finally:
-                self._initialized = False
-                self.tools.clear()
-                self._filtered_tools_cache = None  # Clear cache on disconnect
-                logger.info(f"Session disconnected for {self.server_name}")
+                await owner_task
+            except (asyncio.CancelledError, Exception):
+                pass
+        except asyncio.CancelledError:
+            pass
+        except Exception:
+            pass
+        finally:
+            self._owner_task = None
 
     async def _cleanup_on_error(self) -> None:
         """Internal cleanup method for error scenarios."""
@@ -306,6 +331,17 @@ class PersistentMCPSession:
         """
         return self._initialized and self.connection_manager.is_connected
 
+    @property
+    def allowed_tools(self) -> list[str] | None:
+        """Return the configured allowed tools, sorted if present.
+
+        Returns:
+            Sorted list of allowed tool names, or None if unrestricted.
+        """
+        if not self._allowed_tools_set:
+            return None
+        return sorted(self._allowed_tools_set)
+
     def update_allowed_tools(self, allowed_tools: list[str] | None) -> bool:
         """Update the list of allowed tools for this session.
 
@@ -360,3 +396,217 @@ class PersistentMCPSession:
                 logger.warning(
                     f"[{self.server_name}] Tool '{tool_name}' not found in available tools but specified in allowed_tools"
                 )
+
+    async def _owner_loop(self) -> None:
+        """Run the owner task loop and process queued requests.
+
+        Returns:
+            None
+        """
+        shutdown_requested = False
+        try:
+            shutdown_requested = await self._initialize_owner()
+            if shutdown_requested:
+                return
+
+            while True:
+                func, args, future, shutdown = await self._request_queue.get()
+                if await self._process_owner_request(func, args, future):
+                    continue
+                if shutdown:
+                    shutdown_requested = True
+                    break
+        finally:
+            # Drain and cancel any pending requests to avoid hanging callers
+            await self._drain_pending_requests()
+            if not shutdown_requested and (
+                self._initialized or self.client_session or self.connection_manager.transport_type is not None
+            ):
+                await self._cleanup_on_error()
+            self._owner_task = None
+
+    async def _drain_pending_requests(self) -> None:
+        """Cancel all pending requests in the queue.
+
+        This prevents callers from hanging when the owner loop exits unexpectedly.
+        """
+        error = ConnectionError(f"Session for {self.server_name} is shutting down")
+        while not self._request_queue.empty():
+            try:
+                _, _, future, _ = self._request_queue.get_nowait()
+                if not future.done():
+                    future.set_exception(error)
+            except asyncio.QueueEmpty:
+                break
+
+    async def _initialize_owner(self) -> bool:
+        """Initialize the owner task and signal readiness.
+
+        Returns:
+            bool: True when initialization fails and the loop should stop.
+        """
+        try:
+            await self._initialize_impl()
+        except Exception as e:
+            self._owner_exception = e
+            self._owner_ready.set()
+            return True
+
+        self._owner_ready.set()
+        return False
+
+    async def _process_owner_request(
+        self,
+        func: Callable[..., Awaitable[Any]],
+        args: tuple[Any, ...],
+        future: asyncio.Future,
+    ) -> bool:
+        """Process a single queued request.
+
+        Args:
+            func: Coroutine function to execute.
+            args: Positional arguments for the function.
+            future: Future to resolve with the result or exception.
+
+        Returns:
+            bool: True if the request was skipped due to cancellation.
+        """
+        if future.cancelled():
+            return True
+
+        try:
+            result = await func(*args)
+        except asyncio.CancelledError as e:
+            # Owner task was cancelled - resolve future to prevent hanging caller
+            if not future.cancelled():
+                future.set_exception(e)
+            raise
+        except Exception as e:
+            if not future.cancelled():
+                future.set_exception(e)
+        else:
+            if not future.cancelled():
+                future.set_result(result)
+        return False
+
+    async def _run_in_owner(
+        self,
+        func: Callable[..., Awaitable[Any]],
+        *args: Any,
+        shutdown: bool = False,
+        ensure_initialized: bool = True,
+    ) -> Any:
+        """Execute a coroutine on the owner task.
+
+        Args:
+            func: Coroutine function to execute.
+            *args: Positional arguments to pass to func.
+            shutdown: Whether this request should shut down the owner loop.
+            ensure_initialized: Whether to initialize the owner task if needed.
+
+        Returns:
+            Any: The result of the coroutine call.
+
+        Raises:
+            ConnectionError: If the owner task died or session is shutting down.
+        """
+        if ensure_initialized:
+            await self.initialize()
+        else:
+            # For non-init calls (like disconnect), check if owner is alive
+            if self._owner_task is None or self._owner_task.done():
+                return await func(*args)
+            try:
+                await asyncio.wait_for(self._owner_ready.wait(), timeout=self._timeout)
+            except TimeoutError as e:
+                # Owner task is stuck, cancel it and raise
+                if self._owner_task and not self._owner_task.done():
+                    self._owner_task.cancel()
+                raise ConnectionError(
+                    f"Session for {self.server_name} initialization timed out after {self._timeout}s"
+                ) from e
+            if self._owner_exception:
+                # Propagate the error instead of silently returning None
+                raise ConnectionError(
+                    f"Session for {self.server_name} failed: {self._owner_exception}"
+                ) from self._owner_exception
+
+        # Check if owner task died after initialization (race condition guard)
+        if self._owner_task is None or self._owner_task.done():
+            raise ConnectionError(f"Session for {self.server_name} is no longer active")
+
+        loop = asyncio.get_running_loop()
+        future: asyncio.Future = loop.create_future()
+        await self._request_queue.put((func, args, future, shutdown))
+        return await future
+
+    async def _initialize_impl(self) -> None:
+        """Initialize the underlying MCP session on the owner task."""
+        if self._initialized:
+            return
+
+        try:
+            logger.info(f"Initializing persistent session for {self.server_name}")
+
+            # Start connection manager
+            read_stream, write_stream = await self.connection_manager.start()
+
+            # Create client session
+            self.client_session = ClientSession(read_stream, write_stream)
+            await self.client_session.__aenter__()
+
+            # MCP handshake
+            result = await self.client_session.initialize()
+            logger.debug(f"MCP handshake complete for {self.server_name}: {result.capabilities}")
+
+            # Discover and cache tools
+            if result.capabilities.tools:
+                tools_result = await self.client_session.list_tools()
+                self.tools = tools_result.tools if tools_result else []
+                self._filtered_tools_cache = None  # Invalidate cache when tools change
+                logger.info(f"Cached {len(self.tools)} tools for {self.server_name}")
+            else:
+                logger.info(f"No tools available for {self.server_name}")
+
+            # Warn once per initialization if allowed_tools references unknown names
+            if self._allowed_tools_set:
+                self._warn_on_unknown_allowed_tools(list(self._allowed_tools_set), self.tools)
+
+            # Discover resources (for future use)
+            if result.capabilities.resources:
+                try:
+                    resources_result = await self.client_session.list_resources()
+                    if resources_result and resources_result.resources:
+                        logger.debug(f"Found {len(resources_result.resources)} resources for {self.server_name}")
+                except Exception:
+                    logger.debug(f"Could not list resources for {self.server_name}, skipping")
+
+            self._initialized = True
+            logger.info(f"Session initialization complete for {self.server_name}")
+
+        except Exception as e:
+            logger.error(f"Failed to initialize session for {self.server_name}: {e}", exc_info=True)
+            await self._cleanup_on_error()
+            raise ConnectionError(f"Failed to initialize MCP session for {self.server_name}: {str(e)}") from e
+
+    async def _disconnect_impl(self) -> None:
+        """Disconnect the underlying MCP session on the owner task."""
+        try:
+            # Close client session
+            if self.client_session:
+                try:
+                    await self.client_session.__aexit__(None, None, None)
+                except Exception as e:
+                    logger.warning(f"Error closing client session for {self.server_name}: {e}")
+                self.client_session = None
+
+            # Stop connection manager
+            await self.connection_manager.stop()
+
+        except Exception as e:
+            logger.error(f"Error during disconnect for {self.server_name}: {e}")
+        finally:
+            self._initialized = False
+            self.tools.clear()
+            self._filtered_tools_cache = None  # Clear cache on disconnect
+            logger.info(f"Session disconnected for {self.server_name}")

aip_agents/mcp/client/persistent_session.pyi

@@ -0,0 +1,122 @@
+from _typeshed import Incomplete
+from aip_agents.mcp.client.connection_manager import MCPConnectionManager as MCPConnectionManager
+from aip_agents.mcp.utils.config_validator import validate_allowed_tools_list as validate_allowed_tools_list
+from aip_agents.utils.logger import get_logger as get_logger
+from collections.abc import Awaitable as Awaitable
+from gllm_tools.mcp.client.config import MCPConfiguration
+from mcp import ClientSession
+from mcp.types import CallToolResult, Tool as Tool
+from typing import Any
+
+logger: Incomplete
+
+class PersistentMCPSession:
+    """Persistent MCP session that reuses connections.
+
+    This session wrapper manages the connection lifecycle and caches tools
+    to avoid repeated initialization overhead. It provides automatic reconnection
+    and thread-safe operations.
+
+    Tool Filtering:
+        When allowed_tools is configured, tools are filtered inline during list_tools()
+        and permission checked in call_tool() using set lookup.
+    """
+    server_name: Incomplete
+    config: Incomplete
+    connection_manager: Incomplete
+    client_session: ClientSession | None
+    tools: list[Tool]
+    def __init__(self, server_name: str, config: MCPConfiguration, allowed_tools: list[str] | None = None) -> None:
+        """Initialize persistent session.
+
+        Args:
+            server_name: Name of the MCP server
+            config: MCP server configuration
+            allowed_tools: Optional list of tool names to allow. None or empty means all tools allowed.
+        """
+    async def initialize(self) -> None:
+        """Initialize session once and cache tools.
+
+        This method is idempotent and can be called multiple times safely.
+
+        Raises:
+            Exception: If session initialization fails
+        """
+    async def call_tool(self, name: str, arguments: dict[str, Any]) -> CallToolResult:
+        """Call MCP tool using persistent session.
+
+        Args:
+            name (str): Tool name
+            arguments (dict[str, Any]): Tool arguments
+
+        Returns:
+            CallToolResult: Tool call result
+
+        Raises:
+            Exception: If tool call fails
+        """
+    async def read_resource(self, uri: str) -> Any:
+        """Read an MCP resource using persistent session.
+
+        Args:
+            uri (str): The URI of the resource to read
+
+        Returns:
+            Any: The resource content
+
+        Raises:
+            Exception: If resource reading fails
+        """
+    async def list_tools(self) -> list[Tool]:
+        """Get cached tools list with allowed tools filtering applied.
+
+        Returns:
+            list[Tool]: a copy of list of available tools, filtered to only allowed tools if configured
+        """
+    def get_tools_count(self) -> int:
+        """Get count of allowed tools.
+
+        Returns:
+            Count of allowed tools
+        """
+    async def ensure_connected(self) -> None:
+        """Ensure connection is healthy, reconnect if needed.
+
+        This method provides automatic reconnection capability.
+
+        Raises:
+            Exception: If reconnection fails
+        """
+    async def disconnect(self) -> None:
+        """Disconnect session gracefully.
+
+        This method cleans up all resources and connections.
+        Always succeeds, even if the session was already in an error state.
+        """
+    @property
+    def is_initialized(self) -> bool:
+        """Check if session is initialized.
+
+        Returns:
+            bool: True if initialized and connected, False otherwise
+        """
+    @property
+    def allowed_tools(self) -> list[str] | None:
+        """Return the configured allowed tools, sorted if present.
+
+        Returns:
+            Sorted list of allowed tool names, or None if unrestricted.
+        """
+    def update_allowed_tools(self, allowed_tools: list[str] | None) -> bool:
+        """Update the list of allowed tools for this session.
+
+        Args:
+            allowed_tools: New list of allowed tool names or None for no restriction.
+                None and empty list both mean 'no restrictions, allow all tools'.
+
+        Returns:
+            bool: True if the configuration changed, False otherwise.
+
+        Raises:
+            ValueError: If allowed_tools contains invalid entries.
+        """

aip_agents/mcp/client/session_pool.pyi

@@ -0,0 +1,101 @@
+from _typeshed import Incomplete
+from aip_agents.mcp.client.persistent_session import PersistentMCPSession as PersistentMCPSession
+from aip_agents.mcp.utils.config_validator import validate_allowed_tools_config as validate_allowed_tools_config, validate_allowed_tools_list as validate_allowed_tools_list
+from aip_agents.utils.logger import get_logger as get_logger
+from gllm_tools.mcp.client.config import MCPConfiguration
+from mcp.types import Tool as Tool
+
+logger: Incomplete
+
+class MCPSessionPool:
+    """Manages pool of persistent MCP sessions.
+
+    This pool provides centralized management of MCP sessions, including
+    initialization, tool collection, and resource cleanup. Sessions are
+    reused across the agent lifecycle.
+    """
+    sessions: dict[str, PersistentMCPSession]
+    def __init__(self) -> None:
+        """Initialize empty session pool."""
+    async def get_or_create_session(self, server_name: str, config: MCPConfiguration, allowed_tools: list[str] | None = None) -> PersistentMCPSession:
+        """Get existing session or create new one.
+
+        Args:
+            server_name (str): Name of the MCP server
+            config (MCPConfiguration): MCP server configuration
+            allowed_tools (list[str] | None): Optional list of tool names to allow. None means all tools allowed.
+
+        Returns:
+            PersistentMCPSession: Persistent MCP session
+
+        Raises:
+            Exception: If session creation fails
+        """
+    async def initialize_all_sessions(self, server_configs: dict[str, MCPConfiguration]) -> None:
+        """Initialize all sessions and cache tools.
+
+        This method initializes all configured MCP servers concurrently
+        for better performance.
+
+        Args:
+            server_configs (dict[str, MCPConfiguration]): Dictionary of server configurations
+
+        Raises:
+            Exception: If any session initialization fails
+        """
+    def get_all_active_sessions(self) -> dict[str, PersistentMCPSession]:
+        """Get all active sessions.
+
+        Returns:
+            dict[str, PersistentMCPSession]: Dictionary of active sessions by server name
+        """
+    @property
+    def active_sessions(self) -> list[str]:
+        """Get list of active session names.
+
+        Returns:
+            list[str]: List of active session names
+        """
+    async def get_all_tools(self) -> list[Tool]:
+        """Get all cached tools from all active sessions.
+
+        Returns:
+            list[Tool]: List of all available tools across all sessions
+        """
+    async def close_session(self, server_name: str) -> None:
+        """Close specific session.
+
+        Args:
+            server_name (str): Name of the server session to close
+        """
+    async def close_all_sessions(self) -> None:
+        """Close all sessions gracefully.
+
+        This method ensures all resources are cleaned up properly.
+        """
+    @property
+    def is_initialized(self) -> bool:
+        """Check if session pool is initialized.
+
+        Returns:
+            bool: True if initialized, False otherwise
+        """
+    @property
+    def session_count(self) -> int:
+        """Get number of active sessions.
+
+        Returns:
+            int: Number of active sessions
+        """
+    def get_session(self, server_name: str) -> PersistentMCPSession:
+        """Get specific session by name.
+
+        Args:
+            server_name (str): Name of the server
+
+        Returns:
+            PersistentMCPSession: The requested session
+
+        Raises:
+            KeyError: If session doesn't exist
+        """