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

google/adk/tools/mcp_tool/mcp_toolset.py

@@ -12,93 +12,63 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 from contextlib import AsyncExitStack
+import logging
+import os
+import signal
 import sys
-from types import TracebackType
-from typing import List, Optional, TextIO, Tuple, Type
+from typing import List
+from typing import Optional
+from typing import TextIO
+from typing import Union
 
-from .mcp_session_manager import MCPSessionManager, SseServerParams, retry_on_closed_resource
+from typing_extensions import override
+
+from ...agents.readonly_context import ReadonlyContext
+from ..base_tool import BaseTool
+from ..base_toolset import BaseToolset
+from ..base_toolset import ToolPredicate
+from .mcp_session_manager import MCPSessionManager
+from .mcp_session_manager import retry_on_closed_resource
+from .mcp_session_manager import SseServerParams
 
 # 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, StdioServerParameters
+  from mcp import ClientSession
+  from mcp import StdioServerParameters
   from mcp.types import ListToolsResult
 except ImportError as e:
   import sys
 
   if sys.version_info < (3, 10):
     raise ImportError(
-        'MCP Tool requires Python 3.10 or above. Please upgrade your Python'
-        ' version.'
+        "MCP Tool requires Python 3.10 or above. Please upgrade your Python"
+        " version."
     ) from e
   else:
     raise e
 
 from .mcp_tool import MCPTool
 
+logger = logging.getLogger("google_adk." + __name__)
+
 
-class MCPToolset:
+class MCPToolset(BaseToolset):
   """Connects to a MCP Server, and retrieves MCP Tools into ADK Tools.
 
   Usage:
-  Example 1: (using from_server helper):
   ```
-  async def load_tools():
-    return await MCPToolset.from_server(
-      connection_params=StdioServerParameters(
-          command='npx',
-          args=["-y", "@modelcontextprotocol/server-filesystem"],
+  root_agent = LlmAgent(
+      tools=MCPToolset(
+          connection_params=StdioServerParameters(
+              command='npx',
+              args=["-y", "@modelcontextprotocol/server-filesystem"],
           )
-    )
-
-  # Use the tools in an LLM agent
-  tools, exit_stack = await load_tools()
-  agent = LlmAgent(
-      tools=tools
-  )
-  ...
-  await exit_stack.aclose()
-  ```
-
-  Example 2: (using `async with`):
-
-  ```
-  async def load_tools():
-    async with MCPToolset(
-      connection_params=SseServerParams(url="http://0.0.0.0:8090/sse")
-    ) as toolset:
-      tools = await toolset.load_tools()
-
-      agent = LlmAgent(
-          ...
-          tools=tools
       )
+  )
   ```
-
-  Example 3: (provide AsyncExitStack):
-  ```
-  async def load_tools():
-    async_exit_stack = AsyncExitStack()
-    toolset = MCPToolset(
-      connection_params=StdioServerParameters(...),
-    )
-    async_exit_stack.enter_async_context(toolset)
-    tools = await toolset.load_tools()
-    agent = LlmAgent(
-        ...
-        tools=tools
-    )
-    ...
-    await async_exit_stack.aclose()
-
-  ```
-
-  Attributes:
-    connection_params: The connection parameters to the MCP server. Can be
-      either `StdioServerParameters` or `SseServerParams`.
-    exit_stack: The async exit stack to manage the connection to the MCP server.
-    session: The MCP session being initialized with the connection.
   """
 
   def __init__(
@@ -106,161 +76,151 @@ class MCPToolset:
       *,
       connection_params: StdioServerParameters | SseServerParams,
       errlog: TextIO = sys.stderr,
-      exit_stack=AsyncExitStack(),
+      tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
   ):
     """Initializes the MCPToolset.
 
-    Usage:
-    Example 1: (using from_server helper):
-    ```
-    async def load_tools():
-      return await MCPToolset.from_server(
-        connection_params=StdioServerParameters(
-            command='npx',
-            args=["-y", "@modelcontextprotocol/server-filesystem"],
-            )
-      )
-
-    # Use the tools in an LLM agent
-    tools, exit_stack = await load_tools()
-    agent = LlmAgent(
-        tools=tools
-    )
-    ...
-    await exit_stack.aclose()
-    ```
-
-    Example 2: (using `async with`):
-
-    ```
-    async def load_tools():
-      async with MCPToolset(
-        connection_params=SseServerParams(url="http://0.0.0.0:8090/sse")
-      ) as toolset:
-        tools = await toolset.load_tools()
-
-        agent = LlmAgent(
-            ...
-            tools=tools
-        )
-    ```
-
-    Example 3: (provide AsyncExitStack):
-    ```
-    async def load_tools():
-      async_exit_stack = AsyncExitStack()
-      toolset = MCPToolset(
-        connection_params=StdioServerParameters(...),
-      )
-      async_exit_stack.enter_async_context(toolset)
-      tools = await toolset.load_tools()
-      agent = LlmAgent(
-          ...
-          tools=tools
-      )
-      ...
-      await async_exit_stack.aclose()
-
-    ```
-
     Args:
       connection_params: The connection parameters to the MCP server. Can be:
         `StdioServerParameters` for using local mcp server (e.g. using `npx` or
         `python3`); or `SseServerParams` for a local/remote SSE server.
+      errlog: (Optional) TextIO stream for error logging. Use only for
+        initializing a local stdio MCP session.
     """
-    if not connection_params:
-      raise ValueError('Missing connection params in MCPToolset.')
-    self.connection_params = connection_params
-    self.errlog = errlog
-    self.exit_stack = exit_stack
 
-    self.session_manager = MCPSessionManager(
-        connection_params=self.connection_params,
-        exit_stack=self.exit_stack,
-        errlog=self.errlog,
+    if not connection_params:
+      raise ValueError("Missing connection params in MCPToolset.")
+    self._connection_params = connection_params
+    self._errlog = errlog
+    self._exit_stack = AsyncExitStack()
+    self._creator_task_id = None
+    self._process_pid = None  # Store the subprocess PID
+
+    self._session_manager = MCPSessionManager(
+        connection_params=self._connection_params,
+        exit_stack=self._exit_stack,
+        errlog=self._errlog,
     )
+    self._session = None
+    self.tool_filter = tool_filter
+    self._initialized = False
 
-  @classmethod
-  async def from_server(
-      cls,
-      *,
-      connection_params: StdioServerParameters | SseServerParams,
-      async_exit_stack: Optional[AsyncExitStack] = None,
-      errlog: TextIO = sys.stderr,
-  ) -> Tuple[List[MCPTool], AsyncExitStack]:
-    """Retrieve all tools from the MCP connection.
-
-    Usage:
-    ```
-    async def load_tools():
-      tools, exit_stack = await MCPToolset.from_server(
-        connection_params=StdioServerParameters(
-            command='npx',
-            args=["-y", "@modelcontextprotocol/server-filesystem"],
-        )
-      )
-    ```
-
-    Args:
-      connection_params: The connection parameters to the MCP server.
-      async_exit_stack: The async exit stack to use. If not provided, a new
-        AsyncExitStack will be created.
+  async def _initialize(self) -> ClientSession:
+    """Connects to the MCP Server and initializes the ClientSession."""
+    # Store the current task ID when initializing
+    self._creator_task_id = id(asyncio.current_task())
+    self._session, process = await self._session_manager.create_session()
+    # Store the process PID if available
+    if process and hasattr(process, "pid"):
+      self._process_pid = process.pid
+    self._initialized = True
+    return self._session
+
+  def _is_selected(
+      self, tool: BaseTool, readonly_context: Optional[ReadonlyContext]
+  ) -> bool:
+    """Checks if a tool should be selected based on the tool filter."""
+    if self.tool_filter is None:
+      return True
+    if isinstance(self.tool_filter, ToolPredicate):
+      return self.tool_filter(tool, readonly_context)
+    if isinstance(self.tool_filter, list):
+      return tool.name in self.tool_filter
+    return False
+
+  @override
+  async def close(self):
+    """Safely closes the connection to MCP Server with guaranteed resource cleanup."""
+    if not self._initialized:
+      return  # Nothing to close
+
+    logger.info("Closing MCP Toolset")
+
+    # Step 1: Try graceful shutdown of the session if it exists
+    if self._session:
+      try:
+        logger.info("Attempting graceful session shutdown")
+        await self._session.shutdown()
+      except Exception as e:
+        logger.warning(f"Session shutdown error (continuing cleanup): {e}")
+
+    # Step 2: Try to close the exit stack
+    try:
+      logger.info("Closing AsyncExitStack")
+      await self._exit_stack.aclose()
+      # If we get here, the exit stack closed successfully
+      logger.info("AsyncExitStack closed successfully")
+      return
+    except RuntimeError as e:
+      if "Attempted to exit cancel scope in a different task" in str(e):
+        logger.warning("Task mismatch during shutdown - using fallback cleanup")
+        # Continue to manual cleanup
+      else:
+        logger.error(f"Unexpected RuntimeError: {e}")
+        # Continue to manual cleanup
+    except Exception as e:
+      logger.error(f"Error during exit stack closure: {e}")
+      # Continue to manual cleanup
 
-    Returns:
-      A tuple of the list of MCPTools and the AsyncExitStack.
-      - tools: The list of MCPTools.
-      - async_exit_stack: The AsyncExitStack used to manage the connection to
-        the MCP server. Use `await async_exit_stack.aclose()` to close the
-        connection when server shuts down.
-    """
-    async_exit_stack = async_exit_stack or AsyncExitStack()
-    toolset = cls(
-        connection_params=connection_params,
-        exit_stack=async_exit_stack,
-        errlog=errlog,
-    )
+    # Step 3: Manual cleanup of the subprocess if we have its PID
+    if self._process_pid:
+      await self._ensure_process_terminated(self._process_pid)
 
-    await async_exit_stack.enter_async_context(toolset)
-    tools = await toolset.load_tools()
-    return (tools, async_exit_stack)
+    # Step 4: Ask the session manager to do any additional cleanup it can
+    await self._session_manager._emergency_cleanup()
 
-  async def _initialize(self) -> ClientSession:
-    """Connects to the MCP Server and initializes the ClientSession."""
-    self.session = await self.session_manager.create_session()
-    return self.session
+  async def _ensure_process_terminated(self, pid):
+    """Ensure a process is terminated using its PID."""
+    try:
+      # Check if process exists
+      os.kill(pid, 0)  # This just checks if the process exists
+
+      logger.info(f"Terminating process with PID {pid}")
+      # First try SIGTERM for graceful shutdown
+      os.kill(pid, signal.SIGTERM)
+
+      # Give it a moment to terminate
+      for _ in range(30):  # wait up to 3 seconds
+        await asyncio.sleep(0.1)
+        try:
+          os.kill(pid, 0)  # Process still exists
+        except ProcessLookupError:
+          logger.info(f"Process {pid} terminated successfully")
+          return
+
+      # If we get here, process didn't terminate gracefully
+      logger.warning(
+          f"Process {pid} didn't terminate gracefully, using SIGKILL"
+      )
+      os.kill(pid, signal.SIGKILL)
 
-  async def _exit(self):
-    """Closes the connection to MCP Server."""
-    await self.exit_stack.aclose()
+    except ProcessLookupError:
+      logger.info(f"Process {pid} already terminated")
+    except Exception as e:
+      logger.error(f"Error terminating process {pid}: {e}")
 
-  @retry_on_closed_resource('_initialize')
-  async def load_tools(self) -> List[MCPTool]:
+  @retry_on_closed_resource("_initialize")
+  @override
+  async def get_tools(
+      self,
+      readonly_context: Optional[ReadonlyContext] = None,
+  ) -> List[MCPTool]:
     """Loads all tools from the MCP Server.
 
     Returns:
       A list of MCPTools imported from the MCP Server.
     """
-    tools_response: ListToolsResult = await self.session.list_tools()
-    return [
-        MCPTool(
-            mcp_tool=tool,
-            mcp_session=self.session,
-            mcp_session_manager=self.session_manager,
-        )
-        for tool in tools_response.tools
-    ]
-
-  async def __aenter__(self):
-    try:
+    if not self._session:
       await self._initialize()
-      return self
-    except Exception as e:
-      raise e
+    tools_response: ListToolsResult = await self._session.list_tools()
+    tools = []
+    for tool in tools_response.tools:
+      mcp_tool = MCPTool(
+          mcp_tool=tool,
+          mcp_session=self._session,
+          mcp_session_manager=self._session_manager,
+      )
 
-  async def __aexit__(
-      self,
-      exc_type: Optional[Type[BaseException]],
-      exc: Optional[BaseException],
-      tb: Optional[TracebackType],
-  ) -> None:
-    await self._exit()
+      if self._is_selected(mcp_tool, readonly_context):
+        tools.append(mcp_tool)
+    return tools

google/adk/tools/openapi_tool/common/common.py

@@ -14,11 +14,7 @@
 
 import keyword
 import re
-from typing import Any
-from typing import Dict
-from typing import List
-from typing import Optional
-from typing import Union
+from typing import Any, Dict, List, Optional, Union
 
 from fastapi.openapi.models import Response
 from fastapi.openapi.models import Schema
@@ -100,6 +96,7 @@ class ApiParameter(BaseModel):
   py_name: Optional[str] = ''
   type_value: type[Any] = Field(default=None, init_var=False)
   type_hint: str = Field(default=None, init_var=False)
+  required: bool = False
 
   def model_post_init(self, _: Any):
     self.py_name = (

google/adk/tools/openapi_tool/openapi_spec_parser/openapi_toolset.py

@@ -20,18 +20,23 @@ from typing import Final
 from typing import List
 from typing import Literal
 from typing import Optional
+from typing import Union
 
+from typing_extensions import override
 import yaml
 
+from ....agents.readonly_context import ReadonlyContext
 from ....auth.auth_credential import AuthCredential
 from ....auth.auth_schemes import AuthScheme
+from ...base_toolset import BaseToolset
+from ...base_toolset import ToolPredicate
 from .openapi_spec_parser import OpenApiSpecParser
 from .rest_api_tool import RestApiTool
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger("google_adk." + __name__)
 
 
-class OpenAPIToolset:
+class OpenAPIToolset(BaseToolset):
   """Class for parsing OpenAPI spec into a list of RestApiTool.
 
   Usage:
@@ -61,6 +66,7 @@ class OpenAPIToolset:
       spec_str_type: Literal["json", "yaml"] = "json",
       auth_scheme: Optional[AuthScheme] = None,
       auth_credential: Optional[AuthCredential] = None,
+      tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
   ):
     """Initializes the OpenAPIToolset.
 
@@ -94,31 +100,46 @@ class OpenAPIToolset:
       auth_credential: The auth credential to use for all tools. Use
         AuthCredential or use helpers in
         `google.adk.tools.openapi_tool.auth.auth_helpers`
+      tool_filter: The filter used to filter the tools in the toolset. It can be
+        either a tool predicate or a list of tool names of the tools to expose.
     """
     if not spec_dict:
       spec_dict = self._load_spec(spec_str, spec_str_type)
-    self.tools: Final[List[RestApiTool]] = list(self._parse(spec_dict))
+    self._tools: Final[List[RestApiTool]] = list(self._parse(spec_dict))
     if auth_scheme or auth_credential:
       self._configure_auth_all(auth_scheme, auth_credential)
+    self.tool_filter = tool_filter
 
   def _configure_auth_all(
       self, auth_scheme: AuthScheme, auth_credential: AuthCredential
   ):
     """Configure auth scheme and credential for all tools."""
 
-    for tool in self.tools:
+    for tool in self._tools:
       if auth_scheme:
         tool.configure_auth_scheme(auth_scheme)
       if auth_credential:
         tool.configure_auth_credential(auth_credential)
 
-  def get_tools(self) -> List[RestApiTool]:
+  @override
+  async def get_tools(
+      self, readonly_context: Optional[ReadonlyContext] = None
+  ) -> List[RestApiTool]:
     """Get all tools in the toolset."""
-    return self.tools
+    return [
+        tool
+        for tool in self._tools
+        if self.tool_filter is None
+        or (
+            self.tool_filter(tool, readonly_context)
+            if isinstance(self.tool_filter, ToolPredicate)
+            else tool.name in self.tool_filter
+        )
+    ]
 
   def get_tool(self, tool_name: str) -> Optional[RestApiTool]:
     """Get a tool by name."""
-    matching_tool = filter(lambda t: t.name == tool_name, self.tools)
+    matching_tool = filter(lambda t: t.name == tool_name, self._tools)
     return next(matching_tool, None)
 
   def _load_spec(
@@ -142,3 +163,7 @@ class OpenAPIToolset:
       logger.info("Parsed tool: %s", tool.name)
       tools.append(tool)
     return tools
+
+  @override
+  async def close(self):
+    pass