google-adk 1.2.1__py3-none-any.whl → 1.4.0__py3-none-any.whl

google/adk/tools/mcp_tool/mcp_session_manager.py

@@ -14,12 +14,16 @@
 
 from __future__ import annotations
 
+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
@@ -34,7 +38,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(
@@ -47,102 +50,106 @@ except ImportError as e:
 logger = logging.getLogger('google_adk.' + __name__)
 
 
-class SseServerParams(BaseModel):
+class StdioConnectionParams(BaseModel):
+  """Parameters for the MCP Stdio connection.
+
+  Attributes:
+      server_params: Parameters for the MCP Stdio server.
+      timeout: Timeout in seconds for establishing the connection to the MCP
+        stdio server.
+  """
+
+  server_params: StdioServerParameters
+  timeout: float = 5.0
+
+
+class SseConnectionParams(BaseModel):
   """Parameters for the MCP SSE connection.
 
   See MCP SSE Client documentation for more details.
   https://github.com/modelcontextprotocol/python-sdk/blob/main/src/mcp/client/sse.py
+
+  Attributes:
+      url: URL for the MCP SSE server.
+      headers: Headers for the MCP SSE connection.
+      timeout: Timeout in seconds for establishing the connection to the MCP SSE
+        server.
+      sse_read_timeout: Timeout in seconds for reading data from the MCP SSE
+        server.
   """
 
   url: str
   headers: dict[str, Any] | None = None
-  timeout: float = 5
-  sse_read_timeout: float = 60 * 5
+  timeout: float = 5.0
+  sse_read_timeout: float = 60 * 5.0
 
 
-class StreamableHTTPServerParams(BaseModel):
+class StreamableHTTPConnectionParams(BaseModel):
   """Parameters for the MCP SSE connection.
 
   See MCP SSE Client documentation for more details.
   https://github.com/modelcontextprotocol/python-sdk/blob/main/src/mcp/client/streamable_http.py
+
+  Attributes:
+      url: URL for the MCP Streamable HTTP server.
+      headers: Headers for the MCP Streamable HTTP connection.
+      timeout: Timeout in seconds for establishing the connection to the MCP
+        Streamable HTTP server.
+      sse_read_timeout: Timeout in seconds for reading data from the MCP
+        Streamable HTTP server.
+      terminate_on_close: Whether to terminate the MCP Streamable HTTP server
+        when the connection is closed.
   """
 
   url: str
   headers: dict[str, Any] | None = None
-  timeout: float = 5
-  sse_read_timeout: float = 60 * 5
+  timeout: float = 5.0
+  sse_read_timeout: float = 60 * 5.0
   terminate_on_close: bool = True
 
 
-def retry_on_closed_resource(async_reinit_func_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
-  reinitializes the MCP session.
-  2. Both the decorated function and the async_reinit_func_name must be async
-  functions.
+def retry_on_closed_resource(func):
+  """Decorator to automatically retry action when MCP session is closed.
 
-  Usage:
-  class MCPTool:
-      ...
-      async def create_session(self):
-          self.session = ...
-
-      @retry_on_closed_resource('create_session')
-      async def use_session(self):
-          await self.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:
-      async_reinit_func_name: The name of the async function to recreate session.
+      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, async_reinit_func_name) and callable(
-              getattr(self, async_reinit_func_name)
-          ):
-            async_init_fn = getattr(self, async_reinit_func_name)
-            await async_init_fn()
-          else:
-            raise ValueError(
-                f'Function {async_reinit_func_name} does not exist in decorated'
-                ' class. Please check the function 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__(
       self,
       connection_params: Union[
-          StdioServerParameters, SseServerParams, StreamableHTTPServerParams
+          StdioServerParameters,
+          StdioConnectionParams,
+          SseConnectionParams,
+          StreamableHTTPConnectionParams,
       ],
       errlog: TextIO = sys.stderr,
   ):
@@ -155,105 +162,220 @@ class MCPSessionManager:
         errlog: (Optional) TextIO stream for error logging. Use only for
           initializing a local stdio MCP session.
     """
-    self._connection_params = connection_params
+    if isinstance(connection_params, StdioServerParameters):
+      # So far timeout is not configurable. Given MCP is still evolving, we
+      # would expect stdio_client to evolve to accept timeout parameter like
+      # other client.
+      logger.warning(
+          'StdioServerParameters is not recommended. Please use'
+          ' StdioConnectionParams.'
+      )
+      self._connection_params = StdioConnectionParams(
+          server_params=connection_params,
+          timeout=5,
+      )
+    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
 
-  async def create_session(self) -> ClientSession:
+    # 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()
+
+  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.
     """
-    if self._session is not None:
-      return self._session
-
-    # Create a new exit stack for this session
-    self._exit_stack = AsyncExitStack()
+    # 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:
+      # 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, StdioServerParameters):
-        # So far timeout is not configurable. Given MCP is still evolving, we
-        # would expect stdio_client to evolve to accept timeout parameter like
-        # other client.
-        client = stdio_client(
-            server=self._connection_params, errlog=self._errlog
-        )
-      elif isinstance(self._connection_params, SseServerParams):
-        client = sse_client(
-            url=self._connection_params.url,
-            headers=self._connection_params.headers,
-            timeout=self._connection_params.timeout,
-            sse_read_timeout=self._connection_params.sse_read_timeout,
-        )
-      elif isinstance(self._connection_params, StreamableHTTPServerParams):
-        client = streamablehttp_client(
-            url=self._connection_params.url,
-            headers=self._connection_params.headers,
-            timeout=timedelta(seconds=self._connection_params.timeout),
-            sse_read_timeout=timedelta(
-                seconds=self._connection_params.sse_read_timeout
-            ),
-            terminate_on_close=self._connection_params.terminate_on_close,
-        )
-      else:
-        raise ValueError(
-            'Unable to initialize connection. Connection should be'
-            ' StdioServerParameters or SseServerParams, but got'
-            f' {self._connection_params}'
-        )
-
-      transports = await self._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.
-      # The StdioServerParameters does not provide a timeout parameter for the
-      # session, so we need to set a default timeout for it. Other clients
-      # (SseServerParams and StreamableHTTPServerParams) already provide a
-      # timeout parameter in their configuration.
-      if isinstance(self._connection_params, StdioServerParameters):
-        # Default timeout for MCP session is 5 seconds, same as SseServerParams
-        # and StreamableHTTPServerParams.
-        # TODO :
-        #   1. make timeout configurable
-        #   2. Add StdioConnectionParams to include StdioServerParameters as a
-        #      field and rename other two params to XXXXConnetionParams. Ohter
-        #      two params are actually connection params, while stdio is
-        #      special, stdio_client takes the resposibility of starting the
-        #      server and working as a client.
-        session = await self._exit_stack.enter_async_context(
-            ClientSession(
-                *transports[:2],
-                read_timeout_seconds=timedelta(seconds=5),
-            )
-        )
-      else:
-        session = await self._exit_stack.enter_async_context(
-            ClientSession(*transports[:2])
-        )
-      await session.initialize()
-
-      self._session = session
-      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
-      raise
+      try:
+        if isinstance(self._connection_params, StdioConnectionParams):
+          client = stdio_client(
+              server=self._connection_params.server_params,
+              errlog=self._errlog,
+          )
+        elif isinstance(self._connection_params, SseConnectionParams):
+          client = sse_client(
+              url=self._connection_params.url,
+              headers=merged_headers,
+              timeout=self._connection_params.timeout,
+              sse_read_timeout=self._connection_params.sse_read_timeout,
+          )
+        elif isinstance(
+            self._connection_params, StreamableHTTPConnectionParams
+        ):
+          client = streamablehttp_client(
+              url=self._connection_params.url,
+              headers=merged_headers,
+              timeout=timedelta(seconds=self._connection_params.timeout),
+              sse_read_timeout=timedelta(
+                  seconds=self._connection_params.sse_read_timeout
+              ),
+              terminate_on_close=self._connection_params.terminate_on_close,
+          )
+        else:
+          raise ValueError(
+              'Unable to initialize connection. Connection should be'
+              ' StdioServerParameters or SseServerParams, but got'
+              f' {self._connection_params}'
+          )
+
+        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 exit_stack.enter_async_context(
+              ClientSession(
+                  *transports[:2],
+                  read_timeout_seconds=timedelta(
+                      seconds=self._connection_params.timeout
+                  ),
+              )
+          )
+        else:
+          session = await exit_stack.enter_async_context(
+              ClientSession(*transports[:2])
+          )
+        await session.initialize()
+
+        # 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 exit_stack:
+          await exit_stack.aclose()
+        raise
 
   async def close(self):
-    """Closes the session and cleans up resources."""
-    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
-        )
-      finally:
-        self._exit_stack = None
-        self._session = None
+    """Closes all sessions and cleans up resources."""
+    async with self._session_lock:
+      for session_key in list(self._sessions.keys()):
+        _, exit_stack = self._sessions[session_key]
+        try:
+          await exit_stack.aclose()
+        except Exception as e:
+          # Log the error but don't re-raise to avoid blocking shutdown
+          print(
+              'Warning: Error during MCP session cleanup for'
+              f' {session_key}: {e}',
+              file=self._errlog,
+          )
+        finally:
+          del self._sessions[session_key]
+
+
+SseServerParams = SseConnectionParams
+
+StreamableHTTPServerParams = StreamableHTTPConnectionParams

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,14 +45,16 @@ 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):
-  """Turns a MCP Tool into a Vertex Agent Framework Tool.
+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
   call the tool.
@@ -63,9 +68,9 @@ class MCPTool(BaseTool):
       auth_scheme: Optional[AuthScheme] = None,
       auth_credential: Optional[AuthCredential] = None,
   ):
-    """Initializes a MCPTool.
+    """Initializes an MCPTool.
 
-    This tool wraps a MCP Tool interface and uses a session manager to
+    This tool wraps an MCP Tool interface and uses a session manager to
     communicate with the MCP server.
 
     Args:
@@ -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,26 +108,80 @@ class MCPTool(BaseTool):
     )
     return function_decl
 
-  @retry_on_closed_resource("_reinitialize_session")
-  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:
         args: The arguments as a dict to pass to the tool.
-        tool_context: The tool context from upper level ADK agent.
+        tool_context: The tool context of the current invocation.
 
     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)
 
-    # 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()
+  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