refactor: enhance mcp tool session management

1. remove unnecessary cached session instance in mcp toolset
2. move session reinitialization logic from mcp tool and mcp toolset to mcp session manager
3. add lock for the code block of session creation to avoid race conditions

PiperOrigin-RevId: 770949529

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

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import asyncio
 from contextlib import AsyncExitStack
 from datetime import timedelta
 import functools
@@ -34,7 +35,6 @@ try:
   from mcp.client.stdio import stdio_client
   from mcp.client.streamable_http import streamablehttp_client
 except ImportError as e:
-  import sys
 
   if sys.version_info < (3, 10):
     raise ImportError(
@@ -105,30 +105,29 @@ class StreamableHTTPConnectionParams(BaseModel):
   terminate_on_close: bool = True
 
 
-def retry_on_closed_resource(async_reinit_func_name: str):
+def retry_on_closed_resource(session_manager_field_name: str):
   """Decorator to automatically reinitialize session and retry action.
 
   When MCP session was closed, the decorator will automatically recreate the
   session and retry the action with the same parameters.
 
   Note:
-  1. async_reinit_func_name is the name of the class member function that
+  1. session_manager_field_name is the name of the class member field that
-  reinitializes the MCP session.
+  contains the MCPSessionManager instance.
-  2. Both the decorated function and the async_reinit_func_name must be async
+  2. The session manager must have a reinitialize_session() async method.
-  functions.
 
   Usage:
   class MCPTool:
-      ...
+      def __init__(self):
-      async def create_session(self):
+          self._mcp_session_manager = MCPSessionManager(...)
-          self.session = ...
 
-      @retry_on_closed_resource('create_session')
+      @retry_on_closed_resource('_mcp_session_manager')
       async def use_session(self):
-          await self.session.call_tool()
+          session = await self._mcp_session_manager.create_session()
+          await session.call_tool()
 
   Args:
-      async_reinit_func_name: The name of the async function to recreate session.
+      session_manager_field_name: The name of the session manager field.
 
   Returns:
       The decorated function.
@@ -141,15 +140,21 @@ def retry_on_closed_resource(async_reinit_func_name: str):
         return await func(self, *args, **kwargs)
       except anyio.ClosedResourceError as close_err:
         try:
-          if hasattr(self, async_reinit_func_name) and callable(
+          if hasattr(self, session_manager_field_name):
-              getattr(self, async_reinit_func_name)
+            session_manager = getattr(self, session_manager_field_name)
+            if hasattr(session_manager, 'reinitialize_session') and callable(
+                getattr(session_manager, 'reinitialize_session')
             ):
-            async_init_fn = getattr(self, async_reinit_func_name)
+              await session_manager.reinitialize_session()
-            await async_init_fn()
             else:
               raise ValueError(
-                f'Function {async_reinit_func_name} does not exist in decorated'
+                  f'Session manager {session_manager_field_name} does not have'
-                ' class. Please check the function name in'
+                  ' reinitialize_session method.'
+              ) from close_err
+          else:
+            raise ValueError(
+                f'Session manager field {session_manager_field_name} does not'
+                ' exist in decorated class. Please check the field name in'
                 ' retry_on_closed_resource decorator.'
             ) from close_err
         except Exception as reinit_err:
@@ -207,6 +212,8 @@ class MCPSessionManager:
     # Each session manager maintains its own exit stack for proper cleanup
     self._exit_stack: Optional[AsyncExitStack] = None
     self._session: Optional[ClientSession] = None
+    # Lock to prevent race conditions in session creation
+    self._session_lock = asyncio.Lock()
 
   async def create_session(self) -> ClientSession:
     """Creates and initializes an MCP client session.
@@ -214,6 +221,13 @@ class MCPSessionManager:
     Returns:
         ClientSession: The initialized MCP client session.
     """
+    # Fast path: if session already exists, return it without acquiring lock
+    if self._session is not None:
+      return self._session
+
+    # Use async lock to prevent race conditions
+    async with self._session_lock:
+      # Double-check: session might have been created while waiting for lock
       if self._session is not None:
         return self._session
 
@@ -233,7 +247,9 @@ class MCPSessionManager:
               timeout=self._connection_params.timeout,
               sse_read_timeout=self._connection_params.sse_read_timeout,
           )
-      elif isinstance(self._connection_params, StreamableHTTPConnectionParams):
+        elif isinstance(
+            self._connection_params, StreamableHTTPConnectionParams
+        ):
           client = streamablehttp_client(
               url=self._connection_params.url,
               headers=self._connection_params.headers,
@@ -280,18 +296,28 @@ class MCPSessionManager:
 
   async def close(self):
     """Closes the session and cleans up resources."""
+    if not self._exit_stack:
+      return
+    async with self._session_lock:
       if self._exit_stack:
         try:
           await self._exit_stack.aclose()
         except Exception as e:
           # Log the error but don't re-raise to avoid blocking shutdown
           print(
-            f'Warning: Error during MCP session cleanup: {e}', file=self._errlog
+              f'Warning: Error during MCP session cleanup: {e}',
+              file=self._errlog,
           )
         finally:
           self._exit_stack = None
           self._session = None
 
+  async def reinitialize_session(self):
+    """Reinitializes the session when connection is lost."""
+    # Close the old session and create a new one
+    await self.close()
+    await self.create_session()
+
 
 SseServerParams = SseConnectionParams
 

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

@@ -105,7 +105,7 @@ class MCPTool(BaseTool):
     )
     return function_decl
 
-  @retry_on_closed_resource("_reinitialize_session")
+  @retry_on_closed_resource("_mcp_session_manager")
   async def run_async(self, *, args, tool_context: ToolContext):
     """Runs the tool asynchronously.
 
@@ -122,9 +122,3 @@ class MCPTool(BaseTool):
     # TODO(cheliu): Support passing tool context to MCP Server.
     response = await session.call_tool(self.name, arguments=args)
     return response
-
-  async def _reinitialize_session(self):
-    """Reinitializes the session when connection is lost."""
-    # Close the old session and create a new one
-    await self._mcp_session_manager.close()
-    await self._mcp_session_manager.create_session()

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

@@ -28,10 +28,8 @@ from ..base_toolset import ToolPredicate
 from .mcp_session_manager import MCPSessionManager
 from .mcp_session_manager import retry_on_closed_resource
 from .mcp_session_manager import SseConnectionParams
-from .mcp_session_manager import SseServerParams
 from .mcp_session_manager import StdioConnectionParams
 from .mcp_session_manager import StreamableHTTPConnectionParams
-from .mcp_session_manager import StreamableHTTPServerParams
 
 # Attempt to import MCP Tool from the MCP library, and hints user to upgrade
 # their Python version to 3.10 if it fails.
@@ -127,9 +125,7 @@ class MCPToolset(BaseToolset):
         errlog=self._errlog,
     )
 
-    self._session = None
+  @retry_on_closed_resource("_mcp_session_manager")
-
-  @retry_on_closed_resource("_reinitialize_session")
   async def get_tools(
       self,
       readonly_context: Optional[ReadonlyContext] = None,
@@ -144,11 +140,10 @@ class MCPToolset(BaseToolset):
         List[BaseTool]: A list of tools available under the specified context.
     """
     # Get session from session manager
-    if not self._session:
+    session = await self._mcp_session_manager.create_session()
-      self._session = await self._mcp_session_manager.create_session()
 
     # Fetch available tools from the MCP server
-    tools_response: ListToolsResult = await self._session.list_tools()
+    tools_response: ListToolsResult = await session.list_tools()
 
     # Apply filtering based on context and tool_filter
     tools = []
@@ -162,14 +157,6 @@ class MCPToolset(BaseToolset):
         tools.append(mcp_tool)
     return tools
 
-  async def _reinitialize_session(self):
-    """Reinitializes the session when connection is lost."""
-    # Close the old session and clear cache
-    await self._mcp_session_manager.close()
-    self._session = await self._mcp_session_manager.create_session()
-
-    # Tools will be reloaded on next get_tools call
-
   async def close(self) -> None:
     """Performs cleanup and releases resources held by the toolset.
 
@@ -182,7 +169,3 @@ class MCPToolset(BaseToolset):
     except Exception as e:
       # Log the error but don't re-raise to avoid blocking shutdown
       print(f"Warning: Error during MCPToolset cleanup: {e}", file=self._errlog)
-    finally:
-      # Clear cached tools
-      self._tools_cache = None
-      self._tools_loaded = False