feat: Enable MCP Tool Auth (Experimental)

PiperOrigin-RevId: 773002759

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

@@ -18,9 +18,12 @@ import asyncio
 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)
+  @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
-
-  return decorator
+  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 @@ class MCPSessionManager:
     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
+      # Check if we have an existing session
+      if session_key in self._sessions:
+        session, exit_stack = self._sessions[session_key]
 
-      # Create a new exit stack for this session
-      self._exit_stack = AsyncExitStack()
+        # 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 @@ class MCPSessionManager:
         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 @@ class MCPSessionManager:
         ):
           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 @@ class MCPSessionManager:
               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 @@ class MCPSessionManager:
               )
           )
         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

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

@@ -14,10 +14,13 @@
 
 from __future__ import annotations
 
+import base64
+import json
 import logging
 from typing import Optional
 
 from google.genai.types import FunctionDeclaration
+from google.oauth2.credentials import Credentials
 from typing_extensions import override
 
 from .._gemini_schema_util import _to_gemini_schema
@@ -42,13 +45,15 @@ except ImportError as e:
 
 from ...auth.auth_credential import AuthCredential
 from ...auth.auth_schemes import AuthScheme
-from ..base_tool import BaseTool
+from ...auth.auth_tool import AuthConfig
+from ..base_authenticated_tool import BaseAuthenticatedTool
+#  import
 from ..tool_context import ToolContext
 
 logger = logging.getLogger("google_adk." + __name__)
 
 
-class MCPTool(BaseTool):
+class MCPTool(BaseAuthenticatedTool):
   """Turns an MCP Tool into an ADK Tool.
 
   Internally, the tool initializes from a MCP Tool, and uses the MCP Session to
@@ -77,19 +82,17 @@ class MCPTool(BaseTool):
     Raises:
         ValueError: If mcp_tool or mcp_session_manager is None.
     """
-    if mcp_tool is None:
-      raise ValueError("mcp_tool cannot be None")
-    if mcp_session_manager is None:
-      raise ValueError("mcp_session_manager cannot be None")
     super().__init__(
         name=mcp_tool.name,
         description=mcp_tool.description if mcp_tool.description else "",
+        auth_config=AuthConfig(
+            auth_scheme=auth_scheme, raw_auth_credential=auth_credential
+        )
+        if auth_scheme
+        else None,
     )
     self._mcp_tool = mcp_tool
     self._mcp_session_manager = mcp_session_manager
-    # TODO(cheliu): Support passing auth to MCP Server.
-    self._auth_scheme = auth_scheme
-    self._auth_credential = auth_credential
 
   @override
   def _get_declaration(self) -> FunctionDeclaration:
@@ -105,8 +108,11 @@ class MCPTool(BaseTool):
     )
     return function_decl
 
-  @retry_on_closed_resource("_mcp_session_manager")
-  async def run_async(self, *, args, tool_context: ToolContext):
+  @retry_on_closed_resource
+  @override
+  async def _run_async_impl(
+      self, *, args, tool_context: ToolContext, credential: AuthCredential
+  ):
     """Runs the tool asynchronously.
 
     Args:
@@ -116,8 +122,66 @@ class MCPTool(BaseTool):
     Returns:
         Any: The response from the tool.
     """
+    # Extract headers from credential for session pooling
+    headers = await self._get_headers(tool_context, credential)
+
     # Get the session from the session manager
-    session = await self._mcp_session_manager.create_session()
+    session = await self._mcp_session_manager.create_session(headers=headers)
 
     response = await session.call_tool(self.name, arguments=args)
     return response
+
+  async def _get_headers(
+      self, tool_context: ToolContext, credential: AuthCredential
+  ) -> Optional[dict[str, str]]:
+    headers = None
+    if credential:
+      if credential.oauth2:
+        headers = {"Authorization": f"Bearer {credential.oauth2.access_token}"}
+      elif credential.google_oauth2_json:
+        google_credential = Credentials.from_authorized_user_info(
+            json.loads(credential.google_oauth2_json)
+        )
+        headers = {"Authorization": f"Bearer {google_credential.token}"}
+      elif credential.http:
+        # Handle HTTP authentication schemes
+        if (
+            credential.http.scheme.lower() == "bearer"
+            and credential.http.credentials.token
+        ):
+          headers = {
+              "Authorization": f"Bearer {credential.http.credentials.token}"
+          }
+        elif credential.http.scheme.lower() == "basic":
+          # Handle basic auth
+          if (
+              credential.http.credentials.username
+              and credential.http.credentials.password
+          ):
+
+            credentials = f"{credential.http.credentials.username}:{credential.http.credentials.password}"
+            encoded_credentials = base64.b64encode(
+                credentials.encode()
+            ).decode()
+            headers = {"Authorization": f"Basic {encoded_credentials}"}
+        elif credential.http.credentials.token:
+          # Handle other HTTP schemes with token
+          headers = {
+              "Authorization": (
+                  f"{credential.http.scheme} {credential.http.credentials.token}"
+              )
+          }
+      elif credential.api_key:
+        # For API keys, we'll add them as headers since MCP typically uses header-based auth
+        # The specific header name would depend on the API, using a common default
+        # TODO Allow user to specify the header name for API keys.
+        headers = {"X-API-Key": credential.api_key}
+      elif credential.service_account:
+        # Service accounts should be exchanged for access tokens before reaching this point
+        # If we reach here, we can try to use google_oauth2_json or log a warning
+        logger.warning(
+            "Service account credentials should be exchanged for access"
+            " tokens before MCP session creation"
+        )
+
+    return headers

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

@@ -22,6 +22,8 @@ from typing import TextIO
 from typing import Union
 
 from ...agents.readonly_context import ReadonlyContext
+from ...auth.auth_credential import AuthCredential
+from ...auth.auth_schemes import AuthScheme
 from ..base_tool import BaseTool
 from ..base_toolset import BaseToolset
 from ..base_toolset import ToolPredicate
@@ -94,6 +96,8 @@ class MCPToolset(BaseToolset):
       ],
       tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
       errlog: TextIO = sys.stderr,
+      auth_scheme: Optional[AuthScheme] = None,
+      auth_credential: Optional[AuthCredential] = None,
   ):
     """Initializes the MCPToolset.
 
@@ -110,6 +114,8 @@ class MCPToolset(BaseToolset):
         list of tool names to include - A ToolPredicate function for custom
         filtering logic
       errlog: TextIO stream for error logging.
+      auth_scheme: The auth scheme of the tool for tool calling
+      auth_credential: The auth credential of the tool for tool calling
     """
     super().__init__(tool_filter=tool_filter)
 
@@ -124,8 +130,10 @@ class MCPToolset(BaseToolset):
         connection_params=self._connection_params,
         errlog=self._errlog,
     )
+    self._auth_scheme = auth_scheme
+    self._auth_credential = auth_credential
 
-  @retry_on_closed_resource("_mcp_session_manager")
+  @retry_on_closed_resource
   async def get_tools(
       self,
       readonly_context: Optional[ReadonlyContext] = None,
@@ -151,6 +159,8 @@ class MCPToolset(BaseToolset):
       mcp_tool = MCPTool(
           mcp_tool=tool,
           mcp_session_manager=self._mcp_session_manager,
+          auth_scheme=self._auth_scheme,
+          auth_credential=self._auth_credential,
       )
 
       if self._is_tool_selected(mcp_tool, readonly_context):