google-adk 0.5.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

google/adk/tools/mcp_tool/mcp_session_manager.py

@@ -12,15 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 from contextlib import AsyncExitStack
 import functools
+import logging
 import sys
-from typing import Any, TextIO
+from typing import Any
+from typing import Optional
+from typing import TextIO
+
 import anyio
 from pydantic import BaseModel
 
 try:
-  from mcp import ClientSession, StdioServerParameters
+  from mcp import ClientSession
+  from mcp import StdioServerParameters
   from mcp.client.sse import sse_client
   from mcp.client.stdio import stdio_client
 except ImportError as e:
@@ -34,6 +40,8 @@ except ImportError as e:
   else:
     raise e
 
+logger = logging.getLogger('google_adk.' + __name__)
+
 
 class SseServerParams(BaseModel):
   """Parameters for the MCP SSE connection.
@@ -62,29 +70,27 @@ def retry_on_closed_resource(async_reinit_func_name: str):
 
   Usage:
   class MCPTool:
-    ...
-    async def create_session(self):
-      self.session = ...
+      ...
+      async def create_session(self):
+          self.session = ...
 
-    @retry_on_closed_resource('create_session')
-    async def use_session(self):
-      await self.session.call_tool()
+      @retry_on_closed_resource('create_session')
+      async def use_session(self):
+          await self.session.call_tool()
 
   Args:
-    async_reinit_func_name: The name of the async function to recreate session.
+      async_reinit_func_name: The name of the async function to recreate session.
 
   Returns:
-    The decorated function.
+      The decorated function.
   """
 
   def decorator(func):
-    @functools.wraps(
-        func
-    )  # Preserves original function metadata (name, docstring)
+    @functools.wraps(func)  # Preserves original function metadata
     async def wrapper(self, *args, **kwargs):
       try:
         return await func(self, *args, **kwargs)
-      except anyio.ClosedResourceError:
+      except anyio.ClosedResourceError as close_err:
         try:
           if hasattr(self, async_reinit_func_name) and callable(
               getattr(self, async_reinit_func_name)
@@ -96,7 +102,7 @@ def retry_on_closed_resource(async_reinit_func_name: str):
                 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}'
@@ -118,73 +124,78 @@ class MCPSessionManager:
   def __init__(
       self,
       connection_params: StdioServerParameters | SseServerParams,
-      exit_stack: AsyncExitStack,
       errlog: TextIO = sys.stderr,
-  ) -> ClientSession:
+  ):
     """Initializes the MCP session manager.
 
-    Example usage:
-    ```
-    mcp_session_manager = MCPSessionManager(
-        connection_params=connection_params,
-        exit_stack=exit_stack,
-    )
-    session = await mcp_session_manager.create_session()
-    ```
-
     Args:
         connection_params: Parameters for the MCP connection (Stdio or SSE).
-        exit_stack: AsyncExitStack to manage the session lifecycle.
         errlog: (Optional) TextIO stream for error logging. Use only for
           initializing a local stdio MCP session.
     """
-    self.connection_params = connection_params
-    self.exit_stack = exit_stack
-    self.errlog = errlog
+    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:
-    return await MCPSessionManager.initialize_session(
-        connection_params=self.connection_params,
-        exit_stack=self.exit_stack,
-        errlog=self.errlog,
-    )
-
-  @classmethod
-  async def initialize_session(
-      cls,
-      *,
-      connection_params: StdioServerParameters | SseServerParams,
-      exit_stack: AsyncExitStack,
-      errlog: TextIO = sys.stderr,
-  ) -> ClientSession:
-    """Initializes an MCP client session.
-
-    Args:
-        connection_params: Parameters for the MCP connection (Stdio or SSE).
-        exit_stack: AsyncExitStack to manage the session lifecycle.
-        errlog: (Optional) TextIO stream for error logging. Use only for
-          initializing a local stdio MCP session.
+    """Creates and initializes an MCP client session.
 
     Returns:
         ClientSession: The initialized MCP client session.
     """
-    if isinstance(connection_params, StdioServerParameters):
-      client = stdio_client(server=connection_params, errlog=errlog)
-    elif isinstance(connection_params, SseServerParams):
-      client = sse_client(
-          url=connection_params.url,
-          headers=connection_params.headers,
-          timeout=connection_params.timeout,
-          sse_read_timeout=connection_params.sse_read_timeout,
-      )
-    else:
-      raise ValueError(
-          'Unable to initialize connection. Connection should be'
-          ' StdioServerParameters or SseServerParams, but got'
-          f' {connection_params}'
+    if self._session is not None:
+      return self._session
+
+    # Create a new exit stack for this session
+    self._exit_stack = AsyncExitStack()
+
+    try:
+      if isinstance(self._connection_params, StdioServerParameters):
+        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,
+        )
+      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)
+      session = await self._exit_stack.enter_async_context(
+          ClientSession(*transports)
       )
+      await session.initialize()
 
-    transports = await exit_stack.enter_async_context(client)
-    session = await exit_stack.enter_async_context(ClientSession(*transports))
-    await session.initialize()
-    return session
+      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
+
+  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

google/adk/tools/mcp_tool/mcp_tool.py

@@ -12,17 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
+import logging
 from typing import Optional
 
 from google.genai.types import FunctionDeclaration
 from typing_extensions import override
 
-from .mcp_session_manager import MCPSessionManager, retry_on_closed_resource
+from .mcp_session_manager import MCPSessionManager
+from .mcp_session_manager import retry_on_closed_resource
 
 # Attempt to import MCP Tool from the MCP library, and hints user to upgrade
 # their Python version to 3.10 if it fails.
 try:
-  from mcp import ClientSession
   from mcp.types import Tool as McpBaseTool
 except ImportError as e:
   import sys
@@ -36,12 +38,14 @@ except ImportError as e:
     raise e
 
 
-from ..base_tool import BaseTool
 from ...auth.auth_credential import AuthCredential
 from ...auth.auth_schemes import AuthScheme
+from ..base_tool import BaseTool
 from ..openapi_tool.openapi_spec_parser.rest_api_tool import to_gemini_schema
 from ..tool_context import ToolContext
 
+logger = logging.getLogger("google_adk." + __name__)
+
 
 class MCPTool(BaseTool):
   """Turns a MCP Tool into a Vertex Agent Framework Tool.
@@ -52,44 +56,39 @@ class MCPTool(BaseTool):
 
   def __init__(
       self,
+      *,
       mcp_tool: McpBaseTool,
-      mcp_session: ClientSession,
       mcp_session_manager: MCPSessionManager,
       auth_scheme: Optional[AuthScheme] = None,
-      auth_credential: Optional[AuthCredential] | None = None,
+      auth_credential: Optional[AuthCredential] = None,
   ):
     """Initializes a MCPTool.
 
-    This tool wraps a MCP Tool interface and an active MCP Session. It invokes
-    the MCP Tool through executing the tool from remote MCP Session.
-
-    Example:
-        tool = MCPTool(mcp_tool=mcp_tool, mcp_session=mcp_session)
+    This tool wraps a MCP Tool interface and uses a session manager to
+    communicate with the MCP server.
 
     Args:
         mcp_tool: The MCP tool to wrap.
-        mcp_session: The MCP session to use to call the tool.
+        mcp_session_manager: The MCP session manager to use for communication.
         auth_scheme: The authentication scheme to use.
         auth_credential: The authentication credential to use.
 
     Raises:
-        ValueError: If mcp_tool or mcp_session is None.
+        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 is None:
-      raise ValueError("mcp_session cannot be None")
-    self.name = mcp_tool.name
-    self.description = mcp_tool.description if mcp_tool.description else ""
-    self.mcp_tool = mcp_tool
-    self.mcp_session = mcp_session
-    self.mcp_session_manager = mcp_session_manager
+    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 "",
+    )
+    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
-
-  async def _reinitialize_session(self):
-    self.mcp_session = await self.mcp_session_manager.create_session()
+    self._auth_scheme = auth_scheme
+    self._auth_credential = auth_credential
 
   @override
   def _get_declaration(self) -> FunctionDeclaration:
@@ -98,14 +97,13 @@ class MCPTool(BaseTool):
     Returns:
         FunctionDeclaration: The Gemini function declaration for the tool.
     """
-    schema_dict = self.mcp_tool.inputSchema
+    schema_dict = self._mcp_tool.inputSchema
     parameters = to_gemini_schema(schema_dict)
     function_decl = FunctionDeclaration(
         name=self.name, description=self.description, parameters=parameters
     )
     return function_decl
 
-  @override
   @retry_on_closed_resource("_reinitialize_session")
   async def run_async(self, *, args, tool_context: ToolContext):
     """Runs the tool asynchronously.
@@ -117,10 +115,15 @@ class MCPTool(BaseTool):
     Returns:
         Any: The response from the tool.
     """
+    # Get the session from the session manager
+    session = await self._mcp_session_manager.create_session()
+
     # TODO(cheliu): Support passing tool context to MCP Server.
-    try:
-      response = await self.mcp_session.call_tool(self.name, arguments=args)
-      return response
-    except Exception as e:
-      print(e)
-      raise e
+    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()