cloudbuilderspa
diff --git a/‎src/google/adk/tools/mcp_tool/mcp_session_manager.py
Lines changed: 147 additions & 90 deletions b/‎src/google/adk/tools/mcp_tool/mcp_session_manager.py
Lines changed: 147 additions & 90 deletions
@@ -18,9 +18,12 @@
 from contextlib import AsyncExitStack
 from datetime import timedelta
 import functools
+import hashlib
+import json
 import logging
 import sys
 from typing import Any
+from typing import Dict
 from typing import Optional
 from typing import TextIO
 from typing import Union
@@ -105,74 +108,39 @@ class StreamableHTTPConnectionParams(BaseModel):
   terminate_on_close: bool = True
 
 
-def retry_on_closed_resource(session_manager_field_name: str):
-  """Decorator to automatically reinitialize session and retry action.
+def retry_on_closed_resource(func):
+  """Decorator to automatically retry action when MCP session is closed.
 
-  When MCP session was closed, the decorator will automatically recreate the
-  session and retry the action with the same parameters.
-
-  Note:
-  1. session_manager_field_name is the name of the class member field that
-  contains the MCPSessionManager instance.
-  2. The session manager must have a reinitialize_session() async method.
-
-  Usage:
-  class MCPTool:
-      def __init__(self):
-          self._mcp_session_manager = MCPSessionManager(...)
-
-      @retry_on_closed_resource('_mcp_session_manager')
-      async def use_session(self):
-          session = await self._mcp_session_manager.create_session()
-          await session.call_tool()
+  When MCP session was closed, the decorator will automatically retry the
+  action once. The create_session method will handle creating a new session
+  if the old one was disconnected.
 
   Args:
-      session_manager_field_name: The name of the session manager field.
+      func: The function to decorate.
 
   Returns:
       The decorated function.
   """
 
-  def decorator(func):
-    @functools.wraps(func)  # Preserves original function metadata
-    async def wrapper(self, *args, **kwargs):
-      try:
-        return await func(self, *args, **kwargs)
-      except anyio.ClosedResourceError as close_err:
-        try:
-          if hasattr(self, session_manager_field_name):
-            session_manager = getattr(self, session_manager_field_name)
-            if hasattr(session_manager, 'reinitialize_session') and callable(
-                getattr(session_manager, 'reinitialize_session')
-            ):
-              await session_manager.reinitialize_session()
-            else:
-              raise ValueError(
-                  f'Session manager {session_manager_field_name} does not have'
-                  ' 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:
-          raise RuntimeError(
-              f'Error reinitializing: {reinit_err}'
-          ) from reinit_err
-        return await func(self, *args, **kwargs)
-
-    return wrapper
-
-  return decorator
+  @functools.wraps(func)  # Preserves original function metadata
+  async def wrapper(self, *args, **kwargs):
+    try:
+      return await func(self, *args, **kwargs)
+    except anyio.ClosedResourceError:
+      # Simply retry the function - create_session will handle
+      # detecting and replacing disconnected sessions
+      logger.info('Retrying %s due to closed resource', func.__name__)
+      return await func(self, *args, **kwargs)
+
+  return wrapper
 
 
 class MCPSessionManager:
   """Manages MCP client sessions.
 
   This class provides methods for creating and initializing MCP client sessions,
-  handling different connection parameters (Stdio and SSE).
+  handling different connection parameters (Stdio and SSE) and supporting
+  session pooling based on authentication headers.
   """
 
   def __init__(
@@ -209,30 +177,125 @@ def __init__(
     else:
       self._connection_params = connection_params
     self._errlog = errlog
-    # Each session manager maintains its own exit stack for proper cleanup
-    self._exit_stack: Optional[AsyncExitStack] = None
-    self._session: Optional[ClientSession] = None
+
+    # Session pool: maps session keys to (session, exit_stack) tuples
+    self._sessions: Dict[str, tuple[ClientSession, AsyncExitStack]] = {}
+
     # Lock to prevent race conditions in session creation
     self._session_lock = asyncio.Lock()
 
-  async def create_session(self) -> ClientSession:
+  def _generate_session_key(
+      self, merged_headers: Optional[Dict[str, str]] = None
+  ) -> str:
+    """Generates a session key based on connection params and merged headers.
+
+    For StdioConnectionParams, returns a constant key since headers are not
+    supported. For SSE and StreamableHTTP connections, generates a key based
+    on the provided merged headers.
+
+    Args:
+        merged_headers: Already merged headers (base + additional).
+
+    Returns:
+        A unique session key string.
+    """
+    if isinstance(self._connection_params, StdioConnectionParams):
+      # For stdio connections, headers are not supported, so use constant key
+      return 'stdio_session'
+
+    # For SSE and StreamableHTTP connections, use merged headers
+    if merged_headers:
+      headers_json = json.dumps(merged_headers, sort_keys=True)
+      headers_hash = hashlib.md5(headers_json.encode()).hexdigest()
+      return f'session_{headers_hash}'
+    else:
+      return 'session_no_headers'
+
+  def _merge_headers(
+      self, additional_headers: Optional[Dict[str, str]] = None
+  ) -> Optional[Dict[str, str]]:
+    """Merges base connection headers with additional headers.
+
+    Args:
+        additional_headers: Optional headers to merge with connection headers.
+
+    Returns:
+        Merged headers dictionary, or None if no headers are provided.
+    """
+    if isinstance(self._connection_params, StdioConnectionParams) or isinstance(
+        self._connection_params, StdioServerParameters
+    ):
+      # Stdio connections don't support headers
+      return None
+
+    base_headers = {}
+    if (
+        hasattr(self._connection_params, 'headers')
+        and self._connection_params.headers
+    ):
+      base_headers = self._connection_params.headers.copy()
+
+    if additional_headers:
+      base_headers.update(additional_headers)
+
+    return base_headers
+
+  def _is_session_disconnected(self, session: ClientSession) -> bool:
+    """Checks if a session is disconnected or closed.
+
+    Args:
+        session: The ClientSession to check.
+
+    Returns:
+        True if the session is disconnected, False otherwise.
+    """
+    return session._read_stream._closed or session._write_stream._closed
+
+  async def create_session(
+      self, headers: Optional[Dict[str, str]] = None
+  ) -> ClientSession:
     """Creates and initializes an MCP client session.
 
+    This method will check if an existing session for the given headers
+    is still connected. If it's disconnected, it will be cleaned up and
+    a new session will be created.
+
+    Args:
+        headers: Optional headers to include in the session. These will be
+                merged with any existing connection headers. Only applicable
+                for SSE and StreamableHTTP connections.
+
     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
+    # Merge headers once at the beginning
+    merged_headers = self._merge_headers(headers)
+
+    # Generate session key using merged headers
+    session_key = self._generate_session_key(merged_headers)
 
     # 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
-
-      # Create a new exit stack for this session
-      self._exit_stack = AsyncExitStack()
+      # Check if we have an existing session
+      if session_key in self._sessions:
+        session, exit_stack = self._sessions[session_key]
+
+        # Check if the existing session is still connected
+        if not self._is_session_disconnected(session):
+          # Session is still good, return it
+          return session
+        else:
+          # Session is disconnected, clean it up
+          logger.info('Cleaning up disconnected session: %s', session_key)
+          try:
+            await exit_stack.aclose()
+          except Exception as e:
+            logger.warning('Error during disconnected session cleanup: %s', e)
+          finally:
+            del self._sessions[session_key]
+
+      # Create a new session (either first time or replacing disconnected one)
+      exit_stack = AsyncExitStack()
 
       try:
         if isinstance(self._connection_params, StdioConnectionParams):
@@ -243,7 +306,7 @@ async def create_session(self) -> ClientSession:
         elif isinstance(self._connection_params, SseConnectionParams):
           client = sse_client(
               url=self._connection_params.url,
-              headers=self._connection_params.headers,
+              headers=merged_headers,
               timeout=self._connection_params.timeout,
               sse_read_timeout=self._connection_params.sse_read_timeout,
           )
@@ -252,7 +315,7 @@ async def create_session(self) -> ClientSession:
         ):
           client = streamablehttp_client(
               url=self._connection_params.url,
-              headers=self._connection_params.headers,
+              headers=merged_headers,
               timeout=timedelta(seconds=self._connection_params.timeout),
               sse_read_timeout=timedelta(
                   seconds=self._connection_params.sse_read_timeout
@@ -266,11 +329,11 @@ async def create_session(self) -> ClientSession:
               f' {self._connection_params}'
           )
 
-        transports = await self._exit_stack.enter_async_context(client)
+        transports = await exit_stack.enter_async_context(client)
         # The streamable http client returns a GetSessionCallback in addition to the read/write MemoryObjectStreams
         # needed to build the ClientSession, we limit then to the two first values to be compatible with all clients.
         if isinstance(self._connection_params, StdioConnectionParams):
-          session = await self._exit_stack.enter_async_context(
+          session = await exit_stack.enter_async_context(
               ClientSession(
                   *transports[:2],
                   read_timeout_seconds=timedelta(
@@ -279,44 +342,38 @@ async def create_session(self) -> ClientSession:
               )
           )
         else:
-          session = await self._exit_stack.enter_async_context(
+          session = await exit_stack.enter_async_context(
               ClientSession(*transports[:2])
           )
         await session.initialize()
 
-        self._session = session
+        # Store session and exit stack in the pool
+        self._sessions[session_key] = (session, exit_stack)
+        logger.debug('Created new session: %s', session_key)
         return session
 
       except Exception:
         # If session creation fails, clean up the exit stack
-        if self._exit_stack:
-          await self._exit_stack.aclose()
-          self._exit_stack = None
+        if exit_stack:
+          await exit_stack.aclose()
         raise
 
   async def close(self):
-    """Closes the session and cleans up resources."""
-    if not self._exit_stack:
-      return
+    """Closes all sessions and cleans up resources."""
     async with self._session_lock:
-      if self._exit_stack:
+      for session_key in list(self._sessions.keys()):
+        _, exit_stack = self._sessions[session_key]
         try:
-          await self._exit_stack.aclose()
+          await 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}',
+              'Warning: Error during MCP session cleanup for'
+              f' {session_key}: {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()
+          del self._sessions[session_key]
 
 
 SseServerParams = SseConnectionParams