aip-agents-binary 0.5.25b9__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

@@ -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
 
@@ -142,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.
 
@@ -163,11 +185,20 @@ class HTTPTransport(Transport):
         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)
+            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.

aip_agents/tools/browser_use/browser_use_tool.py

@@ -93,6 +93,7 @@ References:
 """
 
 import asyncio
+import copy
 import json
 from collections.abc import Callable
 from typing import Any, Literal
@@ -556,13 +557,20 @@ class BrowserUseTool(BaseTool):
             iframe_event = yield_iframe_activity(streaming_state.debug_url, "Receive streaming URL")
             self._log_stream_event("iframe_start", iframe_event)
             yield iframe_event
+            last_done_event: dict | None = None
             async for event in self._stream_agent_with_markers(agent, streaming_state, recorder):
                 self._log_stream_event("agent_event", event)
+                tool_info = event.get("tool_info") if isinstance(event, dict) else None
+                tool_calls = tool_info.get("tool_calls", []) if isinstance(tool_info, dict) else []
+                if any(call.get("name") == "done" for call in tool_calls if isinstance(call, dict)):
+                    last_done_event = event
                 yield event
             recording_event = self._recording_event(streaming_state)
             if recording_event:
                 self._log_stream_event("recording_event", recording_event)
                 yield recording_event
+                if last_done_event:
+                    yield copy.deepcopy(last_done_event)
         finally:
             await self._release_session(session, client)
 

aip_agents/tools/browser_use/streaming.py

@@ -223,6 +223,8 @@ def create_step_response(
     if is_done:
         final_tool = _get_done_tool_for_final_response(agent, tool_calls_dict)
         final_output = final_tool.get("output") or TASK_COMPLETED_MESSAGE
+        if not any(call.get("name") == "done" for call in tool_calls_dict):
+            tool_calls_dict.append(final_tool)
         tool_info = {
             "name": PRIMARY_TOOL_NAME,
             "args": final_tool.get("args", {}),

aip_agents/utils/langgraph/tool_managers/delegation_tool_manager.py

@@ -134,7 +134,19 @@ class DelegationToolManager(BaseLangGraphToolManager):
             Returns:
                 The result from the delegated agent, including artifacts if any.
             """
-            writer: StreamWriter = get_stream_writer()
+            try:
+                writer: StreamWriter = get_stream_writer()
+            except Exception as exc:
+                logger.warning(
+                    "DelegationToolManager: Stream writer unavailable; delegation streaming disabled.",
+                    extra={"error": str(exc), "error_type": type(exc).__name__},
+                )
+
+                def _noop_writer(_: Any) -> None:
+                    """No-op writer for non-graph execution contexts."""
+                    return None
+
+                writer = _noop_writer
 
             try:
                 # Check delegation depth limit before executing
@@ -295,6 +307,19 @@ class DelegationToolManager(BaseLangGraphToolManager):
             cfg_conf = cfg.get("configurable") if isinstance(cfg, dict) else None
             if isinstance(cfg_conf, dict):
                 parent_step_id = cfg_conf.get("parent_step_id")
+                if not parent_step_id:
+                    metadata = cfg.get("metadata") or {}
+                    tool_call_id = (
+                        cfg_conf.get("tool_call_id")
+                        or metadata.get("tool_call_id")
+                        or metadata.get("id")
+                        or (metadata.get("tool_call") or {}).get("id")
+                    )
+                    if tool_call_id and self.parent_agent is not None:
+                        thread_key = getattr(self.parent_agent, "thread_id_key", "thread_id")
+                        thread_id = cfg_conf.get(thread_key)
+                        parent_map = self.parent_agent._tool_parent_map_by_thread.get(str(thread_id), {})
+                        parent_step_id = parent_map.get(str(tool_call_id))
             _DELEGATION_PARENT_STEP_ID_CVAR.set(parent_step_id)
         except Exception:
             _DELEGATION_PARENT_STEP_ID_CVAR.set(None)