feat(mcp): gracefully handle tool execution errors and transport crashes

Previously, if an MCP tool returned a JSON-RPC error (e.g. 403 Forbidden) or if the underlying transport connection crashed, the resulting exceptions (McpError and ConnectionError) would bubble up and crash the entire ADK runner.

This change introduces robust error boundaries for MCP tools:
  - `McpTool.run_async()` now catches `McpError` and general exceptions, returning them as structured error dictionaries `{"error": ...}` to the LLM agent so the conversation can continue gracefully.
  - `SessionContext` races tool calls against the background session task so transport crashes surface immediately instead of hanging.
  - Fixes an AnyIO cancellation scope bug ("Attempted to exit cancel scope in a different task") by removing redundant `asyncio.wait_for` wrappers around exit stack context entry.
  - Connection errors trigger automatic retries via `@retry_on_errors` before finally surfacing the failure to the agent.
Fixes #4901, #4162

PiperOrigin-RevId: 902369269

Author: google-genai-bot

src/google/adk/tools/load_mcp_resource_tool.py

@@ -29,8 +29,7 @@
 from .base_tool import BaseTool
 
 if TYPE_CHECKING:
-  from mcp_toolset import McpToolset
-
+  from .mcp_tool.mcp_toolset import McpToolset
   from .tool_context import ToolContext
 
 logger = logging.getLogger("google_adk." + __name__)
@@ -39,7 +38,7 @@
 class LoadMcpResourceTool(BaseTool):
   """A tool that loads the MCP resources and adds them to the session."""
 
-  def __init__(self, mcp_toolset: McpToolset):
+  def __init__(self, mcp_toolset: McpToolset) -> None:
     super().__init__(
         name="load_mcp_resource",
         description="""Loads resources from the MCP server.

src/google/adk/tools/mcp_tool/__init__.py

@@ -17,17 +17,19 @@
 try:
   from .conversion_utils import adk_to_mcp_tool_type
   from .conversion_utils import gemini_to_json_schema
-  from .mcp_session_manager import SseConnectionParams
-  from .mcp_session_manager import StdioConnectionParams
-  from .mcp_session_manager import StreamableHTTPConnectionParams
-  from .mcp_tool import MCPTool
-  from .mcp_tool import McpTool
-  from .mcp_toolset import MCPToolset
-  from .mcp_toolset import McpToolset
+  from .mcp_session_manager import MCPSessionManager as MCPSessionManager
+  from .mcp_session_manager import SseConnectionParams as SseConnectionParams
+  from .mcp_session_manager import StdioConnectionParams as StdioConnectionParams
+  from .mcp_session_manager import StreamableHTTPConnectionParams as StreamableHTTPConnectionParams
+  from .mcp_tool import MCPTool as MCPTool
+  from .mcp_tool import McpTool as McpTool
+  from .mcp_toolset import MCPToolset as MCPToolset
+  from .mcp_toolset import McpToolset as McpToolset
 
   __all__.extend([
       'adk_to_mcp_tool_type',
       'gemini_to_json_schema',
+      'MCPSessionManager',
       'McpTool',
       'MCPTool',
       'McpToolset',

src/google/adk/tools/mcp_tool/mcp_session_manager.py

@@ -16,7 +16,9 @@
 
 import asyncio
 from collections import deque
+import concurrent.futures
 from contextlib import AsyncExitStack
+from dataclasses import dataclass
 from datetime import timedelta
 import functools
 import hashlib
@@ -25,13 +27,20 @@
 import sys
 import threading
 from typing import Any
+from typing import Callable
+from typing import cast
 from typing import Dict
 from typing import Optional
 from typing import Protocol
 from typing import runtime_checkable
 from typing import TextIO
+from typing import TYPE_CHECKING
+from typing import TypeVar
 from typing import Union
 
+if TYPE_CHECKING:
+  from .session_context import SessionContext
+
 from mcp import ClientSession
 from mcp import SamplingCapability
 from mcp import StdioServerParameters
@@ -44,8 +53,6 @@
 from pydantic import BaseModel
 from pydantic import ConfigDict
 
-from .session_context import SessionContext
-
 logger = logging.getLogger('google_adk.' + __name__)
 
 
@@ -146,7 +153,10 @@ class StreamableHTTPConnectionParams(BaseModel):
   httpx_client_factory: CheckableMcpHttpClientFactory = create_mcp_http_client
 
 
-def retry_on_errors(func):
+_F = TypeVar('_F', bound=Callable[..., Any])
+
+
+def retry_on_errors(func: _F) -> _F:
   """Decorator to automatically retry action when MCP session errors occur.
 
   When MCP session errors occur, the decorator will automatically retry the
@@ -165,7 +175,7 @@ def retry_on_errors(func):
   """
 
   @functools.wraps(func)  # Preserves original function metadata
-  async def wrapper(self, *args, **kwargs):
+  async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
     try:
       return await func(self, *args, **kwargs)
     except Exception as e:
@@ -182,7 +192,17 @@ async def wrapper(self, *args, **kwargs):
       logger.info('Retrying %s due to error: %s', func.__name__, e)
       return await func(self, *args, **kwargs)
 
-  return wrapper
+  return cast(_F, wrapper)
+
+
+@dataclass
+class _SessionEntry:
+  """A dataclass to hold session information."""
+
+  session: ClientSession
+  exit_stack: AsyncExitStack
+  loop: asyncio.AbstractEventLoop
+  context: SessionContext
 
 
 class MCPSessionManager:
@@ -205,7 +225,7 @@ def __init__(
       *,
       sampling_callback: Optional[SamplingFnT] = None,
       sampling_capabilities: Optional[SamplingCapability] = None,
-  ):
+  ) -> None:
     """Initializes the MCP session manager.
 
     Args:
@@ -237,10 +257,8 @@ def __init__(
       self._connection_params = connection_params
     self._errlog = errlog
 
-    # Session pool: maps session keys to (session, exit_stack, loop) tuples
-    self._sessions: Dict[
-        str, tuple[ClientSession, AsyncExitStack, asyncio.AbstractEventLoop]
-    ] = {}
+    # Session pool: maps session keys to _SessionEntry objects
+    self._sessions: Dict[str, _SessionEntry] = {}
 
     # Map of event loops to their respective locks to prevent race conditions
     # across different event loops in session creation.
@@ -312,35 +330,66 @@ def _merge_headers(
 
     return base_headers
 
-  def _is_session_disconnected(self, session: ClientSession) -> bool:
+  def _is_session_disconnected(
+      self,
+      entry: _SessionEntry,
+  ) -> bool:
     """Checks if a session is disconnected or closed.
 
     Args:
-        session: The ClientSession to check.
+        entry: The _SessionEntry to check.
 
     Returns:
         True if the session is disconnected, False otherwise.
     """
-    return session._read_stream._closed or session._write_stream._closed
+    if (
+        entry.session._read_stream._closed
+        or entry.session._write_stream._closed
+    ):
+      return True
+    if entry.context is not None and not entry.context._is_task_alive:  # pylint: disable=protected-access
+      return True
+    return False
+
+  def _get_session_context(
+      self, headers: Optional[Dict[str, str]] = None
+  ) -> Optional['SessionContext']:
+    """Returns the SessionContext for the session matching the given headers.
+
+    Note: This method reads from the session pool without acquiring
+    ``_session_lock``. This is safe because it is called immediately after
+    ``create_session()`` (which populates the entry under the lock) within
+    the same task, and dict reads are atomic in CPython.
+
+    Args:
+        headers: Optional headers used to identify the session.
+
+    Returns:
+        The SessionContext if a matching session exists, None otherwise.
+    """
+    merged_headers = self._merge_headers(headers)
+    session_key = self._generate_session_key(merged_headers)
+    entry = self._sessions.get(session_key)
+    if entry is not None:
+      return entry.context
+    return None
 
   async def _cleanup_session(
       self,
       session_key: str,
-      exit_stack: AsyncExitStack,
-      stored_loop: asyncio.AbstractEventLoop,
-  ):
+      entry: _SessionEntry,
+  ) -> None:
     """Cleans up a session, handling different event loops safely.
 
     Args:
         session_key: The session key to clean up.
-        exit_stack: The AsyncExitStack managing the session resources.
-        stored_loop: The event loop on which the session was created.
+        entry: The _SessionEntry managing the session resources.
     """
     current_loop = asyncio.get_running_loop()
     try:
-      if stored_loop is current_loop:
-        await exit_stack.aclose()
-      elif stored_loop.is_closed():
+      if entry.loop is current_loop:
+        await entry.exit_stack.aclose()
+      elif entry.loop.is_closed():
         logger.warning(
             f'Error cleaning up session {session_key}: original event loop'
             ' is closed, resources may be leaked.'
@@ -353,11 +402,11 @@ async def _cleanup_session(
             ' event loop.'
         )
         future = asyncio.run_coroutine_threadsafe(
-            exit_stack.aclose(), stored_loop
+            entry.exit_stack.aclose(), entry.loop
         )
 
         # Attach a callback so errors don't go unnoticed
-        def cleanup_done(f: asyncio.Future):
+        def cleanup_done(f: 'concurrent.futures.Future[Any]') -> None:
           try:
             if f.exception():
               logger.warning(
@@ -379,7 +428,9 @@ def cleanup_done(f: asyncio.Future):
       if session_key in self._sessions:
         del self._sessions[session_key]
 
-  def _create_client(self, merged_headers: Optional[Dict[str, str]] = None):
+  def _create_client(
+      self, merged_headers: Optional[Dict[str, str]] = None
+  ) -> Any:
     """Creates an MCP client based on the connection parameters.
 
     Args:
@@ -451,22 +502,22 @@ async def create_session(
     async with self._session_lock:
       # Check if we have an existing session
       if session_key in self._sessions:
-        session, exit_stack, stored_loop = self._sessions[session_key]
+        entry = self._sessions[session_key]
 
         # Check if the existing session is still connected and bound to the current loop
         current_loop = asyncio.get_running_loop()
-        if stored_loop is current_loop and not self._is_session_disconnected(
-            session
+        if entry.loop is current_loop and not self._is_session_disconnected(
+            entry
         ):
           # Session is still good, return it
-          return session
+          return entry.session
         else:
           # Session is disconnected or from a different loop, clean it up
           logger.info(
               'Cleaning up session (disconnected or different loop): %s',
               session_key,
           )
-          await self._cleanup_session(session_key, exit_stack, stored_loop)
+          await self._cleanup_session(session_key, entry)
 
       # Create a new session (either first time or replacing disconnected one)
       exit_stack = AsyncExitStack()
@@ -482,28 +533,30 @@ async def create_session(
       )
 
       try:
+        from .session_context import SessionContext
+
         client = self._create_client(merged_headers)
         is_stdio = isinstance(self._connection_params, StdioConnectionParams)
 
+        session_context = SessionContext(
+            client=client,
+            timeout=timeout_in_seconds,
+            sse_read_timeout=sse_read_timeout_in_seconds,
+            is_stdio=is_stdio,
+            sampling_callback=self._sampling_callback,
+            sampling_capabilities=self._sampling_capabilities,
+        )
         session = await asyncio.wait_for(
-            exit_stack.enter_async_context(
-                SessionContext(
-                    client=client,
-                    timeout=timeout_in_seconds,
-                    sse_read_timeout=sse_read_timeout_in_seconds,
-                    is_stdio=is_stdio,
-                    sampling_callback=self._sampling_callback,
-                    sampling_capabilities=self._sampling_capabilities,
-                )
-            ),
+            exit_stack.enter_async_context(session_context),
             timeout=timeout_in_seconds,
         )
 
-        # Store session, exit stack, and loop in the pool
-        self._sessions[session_key] = (
-            session,
-            exit_stack,
-            asyncio.get_running_loop(),
+        # Store session, exit stack, loop, and context in the pool
+        self._sessions[session_key] = _SessionEntry(
+            session=session,
+            exit_stack=exit_stack,
+            loop=asyncio.get_running_loop(),
+            context=session_context,
         )
         logger.debug('Created new session: %s', session_key)
         return session
@@ -519,7 +572,7 @@ async def create_session(
             )
         raise ConnectionError(f'Failed to create MCP session: {e}') from e
 
-  def __getstate__(self):
+  def __getstate__(self) -> Dict[str, Any]:
     """Custom pickling to exclude non-picklable runtime objects."""
     state = self.__dict__.copy()
     # Remove unpicklable entries or those that shouldn't persist across pickle
@@ -532,7 +585,7 @@ def __getstate__(self):
 
     return state
 
-  def __setstate__(self, state):
+  def __setstate__(self, state: Dict[str, Any]) -> None:
     """Custom unpickling to restore state."""
     self.__dict__.update(state)
     # Re-initialize members that were not pickled
@@ -543,12 +596,12 @@ def __setstate__(self, state):
     if not hasattr(self, '_errlog') or self._errlog is None:
       self._errlog = sys.stderr
 
-  async def close(self):
+  async def close(self) -> None:
     """Closes all sessions and cleans up resources."""
     async with self._session_lock:
       for session_key in list(self._sessions.keys()):
-        _, exit_stack, stored_loop = self._sessions[session_key]
-        await self._cleanup_session(session_key, exit_stack, stored_loop)
+        entry = self._sessions[session_key]
+        await self._cleanup_session(session_key, entry)
 
 
 SseServerParams = SseConnectionParams