aip-agents-binary 0.5.21__py3-none-macosx_13_0_arm64.whl → 0.6.8__py3-none-macosx_13_0_arm64.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,47 +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:
-                    resources_result = await self.client_session.list_resources()
-                    if resources_result:
-                        logger.debug(f"Found {len(resources_result.resources)} resources for {self.server_name}")
-
-                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.
@@ -140,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))
@@ -170,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:
@@ -185,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)
@@ -199,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)
@@ -231,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.
@@ -252,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."""
@@ -303,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.
 
@@ -357,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

@@ -2,6 +2,7 @@ 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
@@ -90,6 +91,7 @@ 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.
         """
     @property
     def is_initialized(self) -> bool:
@@ -98,6 +100,13 @@ class PersistentMCPSession:
         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.
 

aip_agents/mcp/client/transports.py

@@ -12,10 +12,11 @@ from collections.abc import AsyncIterator
 from enum import StrEnum
 from typing import Any, Protocol
 
+import httpx
 from gllm_tools.mcp.client.config import MCPConfiguration
 from mcp.client.sse import sse_client
 from mcp.client.stdio import StdioServerParameters, stdio_client
-from mcp.client.streamable_http import streamablehttp_client
+from mcp.client.streamable_http import streamable_http_client
 
 from aip_agents.utils.logger import get_logger
 
@@ -45,6 +46,19 @@ DEFAULT_TIMEOUT: float = 30.0
 """Default connection timeout in seconds."""
 
 
+def _sanitize_headers(config: MCPConfiguration) -> dict[str, str]:
+    """Remove headers with None values to avoid invalid HTTP headers.
+
+    Args:
+        config (MCPConfiguration): Transport configuration containing optional headers.
+
+    Returns:
+        dict[str, str]: Filtered headers with None values removed.
+    """
+    headers = config.get("headers", {}) or {}
+    return {key: value for key, value in headers.items() if value is not None}
+
+
 class TransportType(StrEnum):
     """Enum for supported MCP transport types."""
 
@@ -115,7 +129,7 @@ class SSETransport(Transport):
 
         url = f"{base_url}/sse" if not base_url.endswith("/sse") else base_url
         timeout = self.config.get("timeout", DEFAULT_TIMEOUT)
-        headers = self.config.get("headers", {})
+        headers = _sanitize_headers(self.config)
         logger.debug(f"Attempting SSE connection to {url} with headers: {list(headers.keys())}")
         try:
             self.ctx = sse_client(url=url, timeout=timeout, sse_read_timeout=300.0, headers=headers)
@@ -129,6 +143,27 @@ class SSETransport(Transport):
 class HTTPTransport(Transport):
     """Streamable HTTP transport handler."""
 
+    def __init__(self, server_name: str, config: MCPConfiguration) -> None:
+        """Initialize the HTTP transport.
+
+        Args:
+            server_name (str): Name of the MCP server.
+            config (MCPConfiguration): Configuration for the transport.
+        """
+        super().__init__(server_name, config)
+        self._http_client: httpx.AsyncClient | None = None
+
+    async def close(self) -> None:
+        """Clean up the transport connection and any owned HTTP client."""
+        await super().close()
+        if self._http_client:
+            try:
+                await self._http_client.aclose()
+            except Exception as e:
+                logger.warning(f"Error during HTTP client cleanup for {self.server_name}: {e}")
+            finally:
+                self._http_client = None
+
     async def connect(self) -> tuple[AsyncIterator[bytes], AsyncIterator[bytes], TransportContext]:
         """Connect using streamable HTTP transport.
 
@@ -147,14 +182,27 @@ class HTTPTransport(Transport):
 
         url = f"{base_url}/mcp" if not base_url.endswith("/mcp") else base_url
         timeout = self.config.get("timeout", DEFAULT_TIMEOUT)
-        headers = self.config.get("headers", {})
+        headers = _sanitize_headers(self.config)
         logger.debug(f"Attempting streamable HTTP connection to {url} with headers: {list(headers.keys())}")
         try:
-            self.ctx = streamablehttp_client(url=url, timeout=timeout, headers=headers)
+            http_client = httpx.AsyncClient(
+                timeout=httpx.Timeout(timeout),
+                headers=headers,
+                follow_redirects=True,
+            )
+            self._http_client = http_client
+            self.ctx = streamable_http_client(url=url, http_client=http_client)
             read_stream, write_stream, _ = await self.ctx.__aenter__()
             logger.info(f"Connected to {self.server_name} via HTTP")
             return read_stream, write_stream, self.ctx
         except Exception as e:
+            if self._http_client:
+                try:
+                    await self._http_client.aclose()
+                except Exception as close_exc:
+                    logger.warning(f"Error during HTTP client cleanup for {self.server_name}: {close_exc}")
+                finally:
+                    self._http_client = None
             raise ConnectionError(f"HTTP connection failed for {self.server_name}: {str(e)}") from e
 
 

aip_agents/mcp/client/transports.pyi

@@ -77,6 +77,15 @@ class SSETransport(Transport):
 
 class HTTPTransport(Transport):
     """Streamable HTTP transport handler."""
+    def __init__(self, server_name: str, config: MCPConfiguration) -> None:
+        """Initialize the HTTP transport.
+
+        Args:
+            server_name (str): Name of the MCP server.
+            config (MCPConfiguration): Configuration for the transport.
+        """
+    async def close(self) -> None:
+        """Clean up the transport connection and any owned HTTP client."""
     ctx: Incomplete
     async def connect(self) -> tuple[AsyncIterator[bytes], AsyncIterator[bytes], TransportContext]:
         """Connect using streamable HTTP transport.