feat: allowlisting MCP tools for generate_content sessions

google/genai/_adapters.py

@@ -14,7 +14,9 @@
 #
 
 import typing
+from typing import Union
 
+from ._mcp_utils import AllowedToolsMcpSession
 from ._mcp_utils import mcp_to_gemini_tools
 from .types import FunctionCall, Tool
 
@@ -28,8 +30,8 @@ class McpToGenAiToolAdapter:
 
   def __init__(
       self,
-      session: "mcp.ClientSession",  # type: ignore # noqa: F821
-      list_tools_result: "mcp_types.ListToolsResult",  # type: ignore
+      session: Union['ClientSession', AllowedToolsMcpSession],
+      list_tools_result: 'mcp_types.ListToolsResult',  # type: ignore
       is_agent_platform: bool = False,
   ) -> None:
     self._mcp_session = session

google/genai/_extra_utils.py

@@ -564,17 +564,15 @@ async def parse_config_for_mcp_sessions(
   parsed_config_copy = parsed_config.model_copy(update={'tools': None})
   if parsed_config.tools:
     parsed_config_copy.tools = []
-    if not _mcp_utils._is_mcp_loaded():
+    if not _mcp_utils._is_mcp_loaded() and not any(
+        isinstance(tool, _mcp_utils.AllowedToolsMcpSession)
+        for tool in parsed_config.tools
+    ):
       # No MCP tools possible if `mcp` isn't loaded; pass through unchanged.
       parsed_config_copy.tools.extend(parsed_config.tools)
     else:
-      try:
-        from mcp import ClientSession as _McpClientSession  # pylint: disable=g-import-not-at-top
-      except ImportError:
-        _McpClientSession = type('DummySession', (), {})  # type: ignore
-
       for tool in parsed_config.tools:
-        if isinstance(tool, _McpClientSession):
+        if _mcp_utils.is_mcp_client_session(tool):
           mcp_to_genai_tool_adapter = McpToGenAiToolAdapter(
               tool, await tool.list_tools(), is_agent_platform=is_agent_platform
           )

google/genai/_mcp_utils.py

@@ -20,7 +20,7 @@
 
 from importlib.metadata import PackageNotFoundError, version
 import typing
-from typing import Any
+from typing import Any, Optional, Sequence, Union
 
 import google.auth
 from google.auth.transport.requests import Request
@@ -29,6 +29,11 @@
 from . import types
 from ._api_client import _MULTI_REGIONAL_LOCATIONS
 
+if sys.version_info >= (3, 10):
+  from typing import TypeGuard
+else:
+  from typing_extensions import TypeGuard
+
 def _is_mcp_loaded() -> bool:
   return "mcp" in sys.modules
 
@@ -43,6 +48,98 @@ def _is_mcp_loaded() -> bool:
   streamable_http_client: Any = None
   create_mcp_http_client: Any = None
 
+
+class AllowedToolsMcpSession:
+  """Wraps an MCP ``ClientSession`` with an allowlist of tool names.
+
+  Use :func:`with_allowed_tools` to construct instances. Filtering happens in
+  ``list_tools`` so function declarations and the AFC adapter map stay aligned.
+  """
+
+  def __init__(self, session: Any, allowed_tools: Sequence[str]):
+    if not allowed_tools:
+      raise ValueError(
+          'allowed_tools must be a non-empty sequence of tool names.'
+      )
+    self._session = session
+    self._allowed_tools = set(allowed_tools)
+
+  @property
+  def allowed_tools(self) -> set[str]:
+    return set(self._allowed_tools)
+
+  @property
+  def session(self) -> Any:
+    return self._session
+
+  async def list_tools(self) -> Any:
+    """Returns session tools filtered to the allowlist."""
+    try:
+      from mcp import types as mcp_types
+    except ImportError as e:
+      raise ImportError(
+          'The mcp package is required to use with_allowed_tools.'
+      ) from e
+
+    result = await self._session.list_tools()
+    available = {tool.name for tool in result.tools}
+    missing = sorted(self._allowed_tools - available)
+    if missing:
+      raise ValueError(
+          'allowed_tools includes tool names not provided by the MCP session:'
+          f' {missing}. Available tools: {sorted(available)}.'
+      )
+    filtered = [
+        tool for tool in result.tools if tool.name in self._allowed_tools
+    ]
+    return mcp_types.ListToolsResult(tools=filtered)
+
+  async def call_tool(
+      self,
+      name: str,
+      arguments: Optional[dict[str, Any]] = None,
+  ) -> Any:
+    """Calls a tool on the underlying session if it is allowlisted."""
+    if name not in self._allowed_tools:
+      raise ValueError(
+          f'Tool {name!r} is not in allowed_tools'
+          f' {sorted(self._allowed_tools)}.'
+      )
+    return await self._session.call_tool(
+        name=name, arguments=arguments if arguments is not None else {}
+    )
+
+
+def with_allowed_tools(
+    session: Any, allowed_tools: Sequence[str]
+) -> AllowedToolsMcpSession:
+  """Returns an MCP session wrapper that exposes only ``allowed_tools``.
+
+  Example::
+
+    from google.genai import mcp as genai_mcp
+
+    config = types.GenerateContentConfig(
+        tools=[genai_mcp.with_allowed_tools(session, ['tool_a', 'tool_b'])],
+    )
+  """
+  return AllowedToolsMcpSession(session, allowed_tools)
+
+
+def is_mcp_client_session(
+    obj: Any,
+) -> TypeGuard[Union[McpClientSession, AllowedToolsMcpSession]]:
+  """Returns True if ``obj`` is an MCP ClientSession or allowlist wrapper."""
+  if isinstance(obj, AllowedToolsMcpSession):
+    return True
+  if not _is_mcp_loaded():
+    return False
+  try:
+    from mcp import ClientSession as _McpClientSession
+  except ImportError:
+    return False
+  return isinstance(obj, _McpClientSession)
+
 def mcp_to_gemini_tool(tool: McpTool) -> types.Tool:
   """Translates an MCP tool to a Google GenAI tool."""
   return types.Tool(
@@ -83,32 +180,26 @@ def mcp_to_gemini_tools(
 
 def has_mcp_tool_usage(tools: types.ToolListUnion) -> bool:
   """Checks whether the list of tools contains any MCP tools or sessions."""
+  for tool in tools:
+    if is_mcp_client_session(tool):
+      return True
   if not _is_mcp_loaded():
     return False
   try:
-    from mcp import ClientSession as _McpClientSession
     from mcp.types import Tool as _McpTool
   except ImportError:
-    _McpClientSession = type('DummySession', (), {})  # type: ignore
     _McpTool = type('DummyTool', (), {})  # type: ignore
 
   for tool in tools:
-    if isinstance(tool, _McpTool) or isinstance(tool, _McpClientSession):
+    if isinstance(tool, _McpTool):
       return True
   return False
 
 
 def has_mcp_session_usage(tools: types.ToolListUnion) -> bool:
   """Checks whether the list of tools contains any MCP sessions."""
-  if not _is_mcp_loaded():
-    return False
-  try:
-    from mcp import ClientSession as _McpClientSession
-  except ImportError:
-    _McpClientSession = type('DummySession', (), {})  # type: ignore
-
   for tool in tools:
-    if isinstance(tool, _McpClientSession):
+    if is_mcp_client_session(tool):
       return True
   return False
 

google/genai/live.py

@@ -1164,19 +1164,20 @@ async def _t_live_connect_config(
   parameter_model_copy = parameter_model.model_copy(update={'tools': None})
   if parameter_model.tools:
     parameter_model_copy.tools = []
-    if not _mcp_utils._is_mcp_loaded():
+    if not _mcp_utils._is_mcp_loaded() and not any(
+        isinstance(tool, _mcp_utils.AllowedToolsMcpSession)
+        for tool in parameter_model.tools
+    ):
       # No MCP tools possible if `mcp` isn't loaded; pass through unchanged.
       parameter_model_copy.tools.extend(parameter_model.tools)
     else:
       try:
-        from mcp import ClientSession as _McpClientSession  # pylint: disable=g-import-not-at-top
         from mcp.types import Tool as _McpTool  # pylint: disable=g-import-not-at-top
       except ImportError:
-        _McpClientSession = type('DummySession', (), {})  # type: ignore
         _McpTool = type('DummyTool', (), {})  # type: ignore
 
       for tool in parameter_model.tools:
-        if isinstance(tool, _McpClientSession):
+        if _mcp_utils.is_mcp_client_session(tool):
           mcp_to_genai_tool_adapter = McpToGenAiToolAdapter(
               tool, await tool.list_tools()
           )

google/genai/mcp.py

@@ -0,0 +1,24 @@
+# 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.
+#
+
+"""Public helpers for using MCP sessions with the Google GenAI SDK."""
+
+from ._mcp_utils import AllowedToolsMcpSession
+from ._mcp_utils import with_allowed_tools
+
+__all__ = [
+    'AllowedToolsMcpSession',
+    'with_allowed_tools',
+]

google/genai/tests/mcp/test_has_mcp_tool_usage.py

@@ -66,6 +66,26 @@ def __init__(self):
       MockMcpClientSession(),
   ]
   assert _mcp_utils.has_mcp_tool_usage(mcp_tools)
+  assert _mcp_utils.has_mcp_session_usage(mcp_tools)
+
+
+def test_allowed_tools_mcp_session():
+  """Allowlist wrappers are detected as MCP session usage."""
+  if McpClientSession is None:
+    return
+
+  class MockMcpClientSession(McpClientSession):
+
+    def __init__(self):
+      self._read_stream = None
+      self._write_stream = None
+
+  wrapped = _mcp_utils.with_allowed_tools(
+      MockMcpClientSession(), ['tool_a']
+  )
+  assert _mcp_utils.is_mcp_client_session(wrapped)
+  assert _mcp_utils.has_mcp_tool_usage([wrapped])
+  assert _mcp_utils.has_mcp_session_usage([wrapped])
 
 
 def test_no_mcp_tools():