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

google/adk/tools/mcp_tool/mcp_session_manager.py

@@ -12,8 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import asyncio
-from contextlib import asynccontextmanager
+
 from contextlib import AsyncExitStack
 import functools
 import logging
@@ -71,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)
@@ -105,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}'
@@ -117,45 +114,6 @@ def retry_on_closed_resource(async_reinit_func_name: str):
   return decorator
 
 
-@asynccontextmanager
-async def tracked_stdio_client(server, errlog, process=None):
-  """A wrapper around stdio_client that ensures proper process tracking and cleanup."""
-  our_process = process
-
-  # If no process was provided, create one
-  if our_process is None:
-    our_process = await asyncio.create_subprocess_exec(
-        server.command,
-        *server.args,
-        stdin=asyncio.subprocess.PIPE,
-        stdout=asyncio.subprocess.PIPE,
-        stderr=errlog,
-    )
-
-  # Use the original stdio_client, but ensure process cleanup
-  try:
-    async with stdio_client(server=server, errlog=errlog) as client:
-      yield client, our_process
-  finally:
-    # Ensure the process is properly terminated if it still exists
-    if our_process and our_process.returncode is None:
-      try:
-        logger.info(
-            f'Terminating process {our_process.pid} from tracked_stdio_client'
-        )
-        our_process.terminate()
-        try:
-          await asyncio.wait_for(our_process.wait(), timeout=3.0)
-        except asyncio.TimeoutError:
-          # Force kill if it doesn't terminate quickly
-          if our_process.returncode is None:
-            logger.warning(f'Forcing kill of process {our_process.pid}')
-            our_process.kill()
-      except ProcessLookupError:
-        # Process already gone, that's fine
-        logger.info(f'Process {our_process.pid} already terminated')
-
-
 class MCPSessionManager:
   """Manages MCP client sessions.
 
@@ -166,162 +124,78 @@ class MCPSessionManager:
   def __init__(
       self,
       connection_params: StdioServerParameters | SseServerParams,
-      exit_stack: AsyncExitStack,
       errlog: TextIO = sys.stderr,
   ):
     """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._process = None  # Track the subprocess
-    self._active_processes = set()  # Track all processes created
-    self._active_file_handles = set()  # Track file handles
+    # 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,
-  ) -> tuple[ClientSession, Optional[asyncio.subprocess.Process]]:
-    """Creates a new MCP session and tracks the associated process."""
-    session, process = await self._initialize_session(
-        connection_params=self._connection_params,
-        exit_stack=self._exit_stack,
-        errlog=self._errlog,
-    )
-    self._process = process  # Store reference to process
-
-    # Track the process
-    if process:
-      self._active_processes.add(process)
-
-    return session, process
-
-  @classmethod
-  async def _initialize_session(
-      cls,
-      *,
-      connection_params: StdioServerParameters | SseServerParams,
-      exit_stack: AsyncExitStack,
-      errlog: TextIO = sys.stderr,
-  ) -> tuple[ClientSession, Optional[asyncio.subprocess.Process]]:
-    """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.
+  async def create_session(self) -> ClientSession:
+    """Creates and initializes an MCP client session.
 
     Returns:
         ClientSession: The initialized MCP client session.
     """
-    process = None
-
-    if isinstance(connection_params, StdioServerParameters):
-      # For stdio connections, we need to track the subprocess
-      client, process = await cls._create_stdio_client(
-          server=connection_params,
-          errlog=errlog,
-          exit_stack=exit_stack,
-      )
-    elif isinstance(connection_params, SseServerParams):
-      # For SSE connections, create the client without a subprocess
-      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 the session with the client
-    transports = await exit_stack.enter_async_context(client)
-    session = await exit_stack.enter_async_context(ClientSession(*transports))
-    await session.initialize()
+    # Create a new exit stack for this session
+    self._exit_stack = AsyncExitStack()
 
-    return session, process
-
-  @staticmethod
-  async def _create_stdio_client(
-      server: StdioServerParameters,
-      errlog: TextIO,
-      exit_stack: AsyncExitStack,
-  ) -> tuple[Any, asyncio.subprocess.Process]:
-    """Create stdio client and return both the client and process.
-
-    This implementation adapts to how the MCP stdio_client is created.
-    The actual implementation may need to be adjusted based on the MCP library
-    structure.
-    """
-    # Create the subprocess directly so we can track it
-    process = await asyncio.create_subprocess_exec(
-        server.command,
-        *server.args,
-        stdin=asyncio.subprocess.PIPE,
-        stdout=asyncio.subprocess.PIPE,
-        stderr=errlog,
-    )
-
-    # Create the stdio client using the MCP library
     try:
-      # Method 1: Try using the existing process if stdio_client supports it
-      client = stdio_client(server=server, errlog=errlog, process=process)
-    except TypeError:
-      # Method 2: If the above doesn't work, let stdio_client create its own process
-      # and we'll need to terminate both processes later
-      logger.warning(
-          'Using stdio_client with its own process - may lead to duplicate'
-          ' processes'
-      )
-      client = stdio_client(server=server, errlog=errlog)
+      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}'
+        )
 
-    return client, process
+      transports = await self._exit_stack.enter_async_context(client)
+      session = await self._exit_stack.enter_async_context(
+          ClientSession(*transports)
+      )
+      await session.initialize()
 
-  async def _emergency_cleanup(self):
-    """Perform emergency cleanup of resources when normal cleanup fails."""
-    logger.info('Performing emergency cleanup of MCPSessionManager resources')
+      self._session = session
+      return session
 
-    # Clean up any tracked processes
-    for proc in list(self._active_processes):
-      try:
-        if proc and proc.returncode is None:
-          logger.info(f'Emergency termination of process {proc.pid}')
-          proc.terminate()
-          try:
-            await asyncio.wait_for(proc.wait(), timeout=1.0)
-          except asyncio.TimeoutError:
-            logger.warning(f"Process {proc.pid} didn't terminate, forcing kill")
-            proc.kill()
-        self._active_processes.remove(proc)
-      except Exception as e:
-        logger.error(f'Error during process cleanup: {e}')
+    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
 
-    # Clean up any tracked file handles
-    for handle in list(self._active_file_handles):
+  async def close(self):
+    """Closes the session and cleans up resources."""
+    if self._exit_stack:
       try:
-        if not handle.closed:
-          logger.info('Closing file handle')
-          handle.close()
-        self._active_file_handles.remove(handle)
+        await self._exit_stack.aclose()
       except Exception as e:
-        logger.error(f'Error closing file handle: {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,6 +12,8 @@
 # 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
@@ -23,7 +25,6 @@ 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
@@ -43,6 +44,8 @@ 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.
@@ -53,44 +56,40 @@ 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")
-    super().__init__(name=mcp_tool.name, description=mcp_tool.description or "")
+    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 = mcp_session
     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()
-
   @override
   def _get_declaration(self) -> FunctionDeclaration:
     """Gets the function declaration for the tool.
@@ -105,7 +104,6 @@ class MCPTool(BaseTool):
     )
     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()