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

google/adk/tools/mcp_tool/mcp_toolset.py

@@ -12,19 +12,13 @@
 # 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 typing import List
 from typing import Optional
 from typing import TextIO
 from typing import Union
 
-from typing_extensions import override
-
 from ...agents.readonly_context import ReadonlyContext
 from ..base_tool import BaseTool
 from ..base_toolset import BaseToolset
@@ -36,7 +30,6 @@ 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
   from mcp import StdioServerParameters
   from mcp.types import ListToolsResult
 except ImportError as e:
@@ -58,16 +51,31 @@ logger = logging.getLogger("google_adk." + __name__)
 class MCPToolset(BaseToolset):
   """Connects to a MCP Server, and retrieves MCP Tools into ADK Tools.
 
+  This toolset manages the connection to an MCP server and provides tools
+  that can be used by an agent. It properly implements the BaseToolset
+  interface for easy integration with the agent framework.
+
   Usage:
-  ```
-  root_agent = LlmAgent(
-      tools=MCPToolset(
-          connection_params=StdioServerParameters(
-              command='npx',
-              args=["-y", "@modelcontextprotocol/server-filesystem"],
-          )
-      )
+  ```python
+  toolset = MCPToolset(
+      connection_params=StdioServerParameters(
+          command='npx',
+          args=["-y", "@modelcontextprotocol/server-filesystem"],
+      ),
+      tool_filter=['read_file', 'list_directory']  # Optional: filter specific tools
   )
+
+  # Use in an agent
+  agent = LlmAgent(
+      model='gemini-2.0-flash',
+      name='enterprise_assistant',
+      instruction='Help user accessing their file systems',
+      tools=[toolset],
+  )
+
+  # Cleanup is handled automatically by the agent framework
+  # But you can also manually close if needed:
+  # await toolset.close()
   ```
   """
 
@@ -75,152 +83,89 @@ class MCPToolset(BaseToolset):
       self,
       *,
       connection_params: StdioServerParameters | SseServerParams,
-      errlog: TextIO = sys.stderr,
       tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
+      errlog: TextIO = sys.stderr,
   ):
     """Initializes the MCPToolset.
 
     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.
+        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.
+        tool_filter: Optional filter to select specific tools. Can be either:
+            - A list of tool names to include
+            - A ToolPredicate function for custom filtering logic
+        errlog: TextIO stream for error logging.
     """
+    super().__init__(tool_filter=tool_filter)
 
     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(
+    # Create the session manager that will handle the MCP connection
+    self._mcp_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
-
-  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
-
-    # Step 3: Manual cleanup of the subprocess if we have its PID
-    if self._process_pid:
-      await self._ensure_process_terminated(self._process_pid)
-
-    # Step 4: Ask the session manager to do any additional cleanup it can
-    await self._session_manager._emergency_cleanup()
-
-  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)
-
-    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")
-  @override
+  @retry_on_closed_resource("_reinitialize_session")
   async def get_tools(
       self,
       readonly_context: Optional[ReadonlyContext] = None,
-  ) -> List[MCPTool]:
-    """Loads all tools from the MCP Server.
+  ) -> List[BaseTool]:
+    """Return all tools in the toolset based on the provided context.
+
+    Args:
+        readonly_context: Context used to filter tools available to the agent.
+            If None, all tools in the toolset are returned.
 
     Returns:
-      A list of MCPTools imported from the MCP Server.
+        List[BaseTool]: A list of tools available under the specified context.
     """
+    # Get session from session manager
     if not self._session:
-      await self._initialize()
+      self._session = await self._mcp_session_manager.create_session()
+
+    # Fetch available tools from the MCP server
     tools_response: ListToolsResult = await self._session.list_tools()
+
+    # Apply filtering based on context and tool_filter
     tools = []
     for tool in tools_response.tools:
       mcp_tool = MCPTool(
           mcp_tool=tool,
-          mcp_session=self._session,
-          mcp_session_manager=self._session_manager,
+          mcp_session_manager=self._mcp_session_manager,
       )
 
-      if self._is_selected(mcp_tool, readonly_context):
+      if self._is_tool_selected(mcp_tool, readonly_context):
         tools.append(mcp_tool)
     return tools
+
+  async def _reinitialize_session(self):
+    """Reinitializes the session when connection is lost."""
+    # Close the old session and clear cache
+    await self._mcp_session_manager.close()
+    self._session = await self._mcp_session_manager.create_session()
+
+    # Tools will be reloaded on next get_tools call
+
+  async def close(self) -> None:
+    """Performs cleanup and releases resources held by the toolset.
+
+    This method closes the MCP session and cleans up all associated resources.
+    It's designed to be safe to call multiple times and handles cleanup errors
+    gracefully to avoid blocking application shutdown.
+    """
+    try:
+      await self._mcp_session_manager.close()
+    except Exception as e:
+      # Log the error but don't re-raise to avoid blocking shutdown
+      print(f"Warning: Error during MCPToolset cleanup: {e}", file=self._errlog)
+    finally:
+      # Clear cached tools
+      self._tools_cache = None
+      self._tools_loaded = False

google/adk/tools/openapi_tool/auth/credential_exchangers/base_credential_exchanger.py

@@ -15,9 +15,7 @@
 import abc
 from typing import Optional
 
-from .....auth.auth_credential import (
-    AuthCredential,
-)
+from .....auth.auth_credential import AuthCredential
 from .....auth.auth_schemes import AuthScheme
 
 

google/adk/tools/openapi_tool/auth/credential_exchangers/service_account_exchanger.py

@@ -21,14 +21,13 @@ from google.auth.transport.requests import Request
 from google.oauth2 import service_account
 import google.oauth2.credentials
 
-from .....auth.auth_credential import (
-    AuthCredential,
-    AuthCredentialTypes,
-    HttpAuth,
-    HttpCredentials,
-)
+from .....auth.auth_credential import AuthCredential
+from .....auth.auth_credential import AuthCredentialTypes
+from .....auth.auth_credential import HttpAuth
+from .....auth.auth_credential import HttpCredentials
 from .....auth.auth_schemes import AuthScheme
-from .base_credential_exchanger import AuthCredentialMissingError, BaseAuthCredentialExchanger
+from .base_credential_exchanger import AuthCredentialMissingError
+from .base_credential_exchanger import BaseAuthCredentialExchanger
 
 
 class ServiceAccountCredentialExchanger(BaseAuthCredentialExchanger):

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

@@ -14,7 +14,11 @@
 
 import keyword
 import re
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Union
 
 from fastapi.openapi.models import Response
 from fastapi.openapi.models import Schema

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

@@ -12,10 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .openapi_spec_parser import OpenApiSpecParser, OperationEndpoint, ParsedOperation
+from .openapi_spec_parser import OpenApiSpecParser
+from .openapi_spec_parser import OperationEndpoint
+from .openapi_spec_parser import ParsedOperation
 from .openapi_toolset import OpenAPIToolset
 from .operation_parser import OperationParser
-from .rest_api_tool import AuthPreparationState, RestApiTool, snake_to_lower_camel, to_gemini_schema
+from .rest_api_tool import AuthPreparationState
+from .rest_api_tool import RestApiTool
+from .rest_api_tool import snake_to_lower_camel
+from .rest_api_tool import to_gemini_schema
 from .tool_auth_handler import ToolAuthHandler
 
 __all__ = [

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

@@ -103,12 +103,12 @@ class OpenAPIToolset(BaseToolset):
       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.
     """
+    super().__init__(tool_filter=tool_filter)
     if not spec_dict:
       spec_dict = self._load_spec(spec_str, spec_str_type)
     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
@@ -129,12 +129,7 @@ class OpenAPIToolset(BaseToolset):
     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
-        )
+        if self._is_tool_selected(tool, readonly_context)
     ]
 
   def get_tool(self, tool_name: str) -> Optional[RestApiTool]:

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

@@ -14,7 +14,11 @@
 
 import inspect
 from textwrap import dedent
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Union
 
 from fastapi.encoders import jsonable_encoder
 from fastapi.openapi.models import Operation

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

@@ -41,6 +41,16 @@ from .openapi_spec_parser import ParsedOperation
 from .operation_parser import OperationParser
 from .tool_auth_handler import ToolAuthHandler
 
+# Not supported by the Gemini API
+_OPENAPI_SCHEMA_IGNORE_FIELDS = (
+    "title",
+    "default",
+    "format",
+    "additional_properties",
+    "ref",
+    "def",
+)
+
 
 def snake_to_lower_camel(snake_case_string: str):
   """Converts a snake_case string to a lower_camel_case string.
@@ -121,7 +131,7 @@ def to_gemini_schema(openapi_schema: Optional[Dict[str, Any]] = None) -> Schema:
     snake_case_key = to_snake_case(key)
     # Check if the snake_case_key exists in the Schema model's fields.
     if snake_case_key in Schema.model_fields:
-      if snake_case_key in ["title", "default", "format"]:
+      if snake_case_key in _OPENAPI_SCHEMA_IGNORE_FIELDS:
         # Ignore these fields as Gemini backend doesn't recognize them, and will
         # throw exception if they appear in the schema.
         # Format: properties[expiration].format: only 'enum' and 'date-time' are

google/adk/tools/toolbox_toolset.py

@@ -12,8 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Any
+from typing import Callable
 from typing import List
+from typing import Mapping
 from typing import Optional
+from typing import Union
 
 import toolbox_core as toolbox
 from typing_extensions import override
@@ -40,12 +44,24 @@ class ToolboxToolset(BaseToolset):
       server_url: str,
       toolset_name: Optional[str] = None,
       tool_names: Optional[List[str]] = None,
+      auth_token_getters: Optional[dict[str, Callable[[], str]]] = None,
+      bound_params: Optional[
+          Mapping[str, Union[Callable[[], Any], Any]]
+      ] = None,
   ):
     """Args:
 
       server_url: The URL of the toolbox server.
       toolset_name: The name of the toolbox toolset to load.
       tool_names: The names of the tools to load.
+      auth_token_getters: A mapping of authentication service names to
+        callables that return the corresponding authentication token. see:
+        https://github.com/googleapis/mcp-toolbox-sdk-python/tree/main/packages/toolbox-core#authenticating-tools
+        for details.
+      bound_params: A mapping of parameter names to bind to specific values or
+        callables that are called to produce values as needed. see:
+        https://github.com/googleapis/mcp-toolbox-sdk-python/tree/main/packages/toolbox-core#binding-parameter-values
+        for details.
     The resulting ToolboxToolset will contain both tools loaded by tool_names
     and toolset_name.
     """
@@ -53,9 +69,11 @@ class ToolboxToolset(BaseToolset):
       raise ValueError("tool_names and toolset_name cannot both be None")
     super().__init__()
     self._server_url = server_url
-    self._toolbox_client = toolbox.ToolboxSyncClient(server_url)
+    self._toolbox_client = toolbox.ToolboxClient(server_url)
     self._toolset_name = toolset_name
     self._tool_names = tool_names
+    self._auth_token_getters = auth_token_getters or {}
+    self._bound_params = bound_params or {}
 
   @override
   async def get_tools(
@@ -65,11 +83,21 @@ class ToolboxToolset(BaseToolset):
     if self._toolset_name:
       tools.extend([
           FunctionTool(tool)
-          for tool in self._toolbox_client.load_toolset(self._toolset_name)
+          for tool in await self._toolbox_client.load_toolset(
+              self._toolset_name,
+              auth_token_getters=self._auth_token_getters,
+              bound_params=self._bound_params,
+          )
       ])
     if self._tool_names:
       tools.extend([
-          FunctionTool(self._toolbox_client.load_tool(tool_name))
+          FunctionTool(
+              await self._toolbox_client.load_tool(
+                  tool_name,
+                  auth_token_getters=self._auth_token_getters,
+                  bound_params=self._bound_params,
+              )
+          )
           for tool_name in self._tool_names
       ])
     return tools

google/adk/utils/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

google/adk/utils/instructions_utils.py

@@ -0,0 +1,131 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+
+from ..agents.readonly_context import ReadonlyContext
+from ..sessions.state import State
+
+__all__ = [
+    'inject_session_state',
+]
+
+
+async def inject_session_state(
+    template: str,
+    readonly_context: ReadonlyContext,
+) -> str:
+  """Populates values in the instruction template, e.g. state, artifact, etc.
+
+  This method is intended to be used in InstructionProvider based instruction
+  and global_instruction which are called with readonly_context.
+
+  e.g.
+  ```
+  ...
+  from google.adk.utils import instructions_utils
+
+  async def build_instruction(
+      readonly_context: ReadonlyContext,
+  ) -> str:
+    return await instructions_utils.inject_session_state(
+        'You can inject a state variable like {var_name} or an artifact '
+        '{artifact.file_name} into the instruction template.',
+        readonly_context,
+    )
+
+  agent = Agent(
+      model="gemini-2.0-flash",
+      name="agent",
+      instruction=build_instruction,
+  )
+  ```
+
+  Args:
+    template: The instruction template.
+    readonly_context: The read-only context
+
+  Returns:
+    The instruction template with values populated.
+  """
+
+  invocation_context = readonly_context._invocation_context
+
+  async def _async_sub(pattern, repl_async_fn, string) -> str:
+    result = []
+    last_end = 0
+    for match in re.finditer(pattern, string):
+      result.append(string[last_end : match.start()])
+      replacement = await repl_async_fn(match)
+      result.append(replacement)
+      last_end = match.end()
+    result.append(string[last_end:])
+    return ''.join(result)
+
+  async def _replace_match(match) -> str:
+    var_name = match.group().lstrip('{').rstrip('}').strip()
+    optional = False
+    if var_name.endswith('?'):
+      optional = True
+      var_name = var_name.removesuffix('?')
+    if var_name.startswith('artifact.'):
+      var_name = var_name.removeprefix('artifact.')
+      if invocation_context.artifact_service is None:
+        raise ValueError('Artifact service is not initialized.')
+      artifact = await invocation_context.artifact_service.load_artifact(
+          app_name=invocation_context.session.app_name,
+          user_id=invocation_context.session.user_id,
+          session_id=invocation_context.session.id,
+          filename=var_name,
+      )
+      if not var_name:
+        raise KeyError(f'Artifact {var_name} not found.')
+      return str(artifact)
+    else:
+      if not _is_valid_state_name(var_name):
+        return match.group()
+      if var_name in invocation_context.session.state:
+        return str(invocation_context.session.state[var_name])
+      else:
+        if optional:
+          return ''
+        else:
+          raise KeyError(f'Context variable not found: `{var_name}`.')
+
+  return await _async_sub(r'{+[^{}]*}+', _replace_match, template)
+
+
+def _is_valid_state_name(var_name):
+  """Checks if the variable name is a valid state name.
+
+  Valid state is either:
+    - Valid identifier
+    - <Valid prefix>:<Valid identifier>
+  All the others will just return as it is.
+
+  Args:
+    var_name: The variable name to check.
+
+  Returns:
+    True if the variable name is a valid state name, False otherwise.
+  """
+  parts = var_name.split(':')
+  if len(parts) == 1:
+    return var_name.isidentifier()
+
+  if len(parts) == 2:
+    prefixes = [State.APP_PREFIX, State.USER_PREFIX, State.TEMP_PREFIX]
+    if (parts[0] + ':') in prefixes:
+      return parts[1].isidentifier()
+  return False

google/adk/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 
 # version: date+base_cl
-__version__ = "1.0.0"
+__version__ = "1.1.0"