diff --git a/google/genai/_adapters.py b/google/genai/_adapters.py index 73105f6e3..a0ca51622 100644 --- a/google/genai/_adapters.py +++ b/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 diff --git a/google/genai/_extra_utils.py b/google/genai/_extra_utils.py index 8e1445fdd..68ac50b7a 100644 --- a/google/genai/_extra_utils.py +++ b/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 ) diff --git a/google/genai/_mcp_utils.py b/google/genai/_mcp_utils.py index 74b24b363..cee724df7 100644 --- a/google/genai/_mcp_utils.py +++ b/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 diff --git a/google/genai/live.py b/google/genai/live.py index b9cf0c33e..374a7b5aa 100644 --- a/google/genai/live.py +++ b/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() ) diff --git a/google/genai/mcp.py b/google/genai/mcp.py new file mode 100644 index 000000000..c19f89ac3 --- /dev/null +++ b/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', +] diff --git a/google/genai/tests/mcp/test_has_mcp_tool_usage.py b/google/genai/tests/mcp/test_has_mcp_tool_usage.py index 12d8c903a..dadf2e595 100644 --- a/google/genai/tests/mcp/test_has_mcp_tool_usage.py +++ b/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(): diff --git a/google/genai/tests/mcp/test_parse_config_for_mcp_sessions.py b/google/genai/tests/mcp/test_parse_config_for_mcp_sessions.py index 89693bcb8..0ff16858c 100644 --- a/google/genai/tests/mcp/test_parse_config_for_mcp_sessions.py +++ b/google/genai/tests/mcp/test_parse_config_for_mcp_sessions.py @@ -16,6 +16,7 @@ import _asyncio import pytest from ... import _extra_utils +from ... import mcp as genai_mcp from ... import types from ..._adapters import McpToGenAiToolAdapter @@ -34,6 +35,41 @@ raise e +class MockMcpClientSession(McpClientSession): + """Shared mock session with two weather tools.""" + + def __init__(self): + self._read_stream = None + self._write_stream = None + self.call_tool_names = [] + + async def list_tools(self): + return mcp_types.ListToolsResult( + tools=[ + mcp_types.Tool( + name='get_weather', + description='Get the weather in a city.', + inputSchema={ + 'type': 'object', + 'properties': {'location': {'type': 'string'}}, + }, + ), + mcp_types.Tool( + name='get_weather_2', + description='Different tool to get the weather.', + inputSchema={ + 'type': 'object', + 'properties': {'location': {'type': 'string'}}, + }, + ), + ] + ) + + async def call_tool(self, name, arguments=None): + self.call_tool_names.append(name) + return mcp_types.CallToolResult(content=[]) + + @pytest.mark.asyncio async def test_parse_empty_config_dict(): """Test conversion of empty GenerateContentConfigDict to parsed config.""" @@ -59,35 +95,6 @@ async def test_parse_empty_config_object(): @pytest.mark.asyncio async def test_parse_config_object_with_tools(): """Test conversion of GenerateContentConfig with tools to parsed config.""" - - class MockMcpClientSession(McpClientSession): - - def __init__(self): - self._read_stream = None - self._write_stream = None - - async def list_tools(self): - return mcp_types.ListToolsResult( - tools=[ - mcp_types.Tool( - name='get_weather', - description='Get the weather in a city.', - inputSchema={ - 'type': 'object', - 'properties': {'location': {'type': 'string'}}, - }, - ), - mcp_types.Tool( - name='get_weather_2', - description='Different tool to get the weather.', - inputSchema={ - 'type': 'object', - 'properties': {'location': {'type': 'string'}}, - }, - ), - ] - ) - mock_session_instance = MockMcpClientSession() config = types.GenerateContentConfig(tools=[mock_session_instance]) parsed_config, mcp_to_genai_tool_adapters = ( @@ -109,6 +116,55 @@ async def list_tools(self): ) +@pytest.mark.asyncio +async def test_parse_config_with_allowed_tools_filters_adapters_and_fds(): + """Allowlist keeps only named tools in FDs and AFC adapter keys.""" + mock_session = MockMcpClientSession() + wrapped = genai_mcp.with_allowed_tools(mock_session, ['get_weather']) + config = types.GenerateContentConfig(tools=[wrapped]) + parsed_config, mcp_to_genai_tool_adapters = ( + await _extra_utils.parse_config_for_mcp_sessions(config) + ) + assert mcp_to_genai_tool_adapters.keys() == {'get_weather'} + assert isinstance( + mcp_to_genai_tool_adapters['get_weather'], McpToGenAiToolAdapter + ) + names = [] + for tool in parsed_config.tools or []: + if tool.function_declarations: + for fd in tool.function_declarations: + names.append(fd.name) + assert names == ['get_weather'] + + +@pytest.mark.asyncio +async def test_parse_config_with_unknown_allowed_tool_raises(): + mock_session = MockMcpClientSession() + wrapped = genai_mcp.with_allowed_tools( + mock_session, ['get_weather', 'not_a_real_tool'] + ) + config = types.GenerateContentConfig(tools=[wrapped]) + with pytest.raises(ValueError, match='not_a_real_tool'): + await _extra_utils.parse_config_for_mcp_sessions(config) + + +def test_empty_allowed_tools_raises(): + mock_session = MockMcpClientSession() + with pytest.raises(ValueError, match='non-empty'): + genai_mcp.with_allowed_tools(mock_session, []) + + +@pytest.mark.asyncio +async def test_allowed_tools_call_tool_rejects_disallowed_name(): + mock_session = MockMcpClientSession() + wrapped = genai_mcp.with_allowed_tools(mock_session, ['get_weather']) + with pytest.raises(ValueError, match='get_weather_2'): + await wrapped.call_tool(name='get_weather_2', arguments={}) + assert mock_session.call_tool_names == [] + await wrapped.call_tool(name='get_weather', arguments={'location': 'SF'}) + assert mock_session.call_tool_names == ['get_weather'] + + @pytest.mark.asyncio async def test_parse_config_object_with_tools_complex_type(): """Test conversion of GenerateContentConfig with tools to parsed config.""" diff --git a/google/genai/types.py b/google/genai/types.py index 359c00d45..61bae6bd7 100644 --- a/google/genai/types.py +++ b/google/genai/types.py @@ -5335,9 +5335,21 @@ def _validate_tool_list(v: object, handler: Any) -> Any: if typing.TYPE_CHECKING: - ToolUnion = Union[Tool, Callable[..., Any], mcp_types.Tool, McpClientSession] + from ._mcp_utils import AllowedToolsMcpSession as _AllowedToolsMcpSession + + ToolUnion = Union[ + Tool, + Callable[..., Any], + mcp_types.Tool, + McpClientSession, + _AllowedToolsMcpSession, + ] ToolUnionDict = Union[ - ToolDict, Callable[..., Any], mcp_types.Tool, McpClientSession + ToolDict, + Callable[..., Any], + mcp_types.Tool, + McpClientSession, + _AllowedToolsMcpSession, ] ToolListUnion = list[ToolUnion] ToolListUnionDict = list[ToolUnionDict]