From 7cd265306e336107ac360806fa46017b648170bc Mon Sep 17 00:00:00 2001 From: Aarav Mittal Date: Mon, 3 Aug 2026 18:28:06 -0700 Subject: [PATCH] fix: avoid pickling MCP sessions when copying GenerateContentConfig Async generate_content deep-copied config objects and crashed when tools included an MCP ClientSession holding an asyncio.Future. Clear tools before deep copy and reattach by reference, matching the existing MCP parse helpers. Fixes #2669. --- google/genai/_extra_utils.py | 43 +++- google/genai/models.py | 32 +-- .../mcp/test_copy_generate_content_config.py | 211 ++++++++++++++++++ 3 files changed, 257 insertions(+), 29 deletions(-) create mode 100644 google/genai/tests/mcp/test_copy_generate_content_config.py diff --git a/google/genai/_extra_utils.py b/google/genai/_extra_utils.py index 8e1445fdd..c7e8c311a 100644 --- a/google/genai/_extra_utils.py +++ b/google/genai/_extra_utils.py @@ -65,6 +65,36 @@ def _create_generate_content_config_model( return config +def copy_generate_content_config( + config: Optional[types.GenerateContentConfigOrDict], + *, + deep: bool = True, +) -> Optional[types.GenerateContentConfig]: + """Copies GenerateContentConfig without pickling unpickleable tools. + + MCP ``ClientSession`` tools (and other live objects such as callables) cannot + be deep-copied via pickle. Pydantic applies ``deep`` before ``update``, so + tools must be cleared with a shallow copy first, then optionally deep-copied, + then reattached by reference. See + https://github.com/googleapis/python-genai/issues/2669. + """ + if not config: + return None + config_model = _create_generate_content_config_model(config) + tools = config_model.tools + # Clear tools before any deep copy; model_copy(deep=True, update=...) still + # deep-copies tools before applying update. + config_without_tools = config_model.model_copy(update={'tools': None}) + config_copy = ( + config_without_tools.model_copy(deep=True) + if deep + else config_without_tools + ) + if tools is not None: + config_copy.tools = list(tools) + return config_copy + + def _get_gcs_uri( src: Union[str, types.BatchJobSourceOrDict] ) -> Optional[str]: @@ -527,14 +557,13 @@ def parse_config_for_mcp_usage( config: Optional[types.GenerateContentConfigOrDict] = None, ) -> Optional[types.GenerateContentConfig]: """Returns a parsed config with an appended MCP header if MCP tools or sessions are used.""" - if not config: + # Shallow-copy so callers' configs are not mutated; tools may be unpickleable. + config_model_copy = copy_generate_content_config(config, deep=False) + if not config_model_copy: return None - config_model = _create_generate_content_config_model(config) - # Create a copy of the config model with the tools field cleared since some - # tools may not be pickleable. - config_model_copy = config_model.model_copy(update={'tools': None}) - config_model_copy.tools = config_model.tools - if config_model.tools and _mcp_utils.has_mcp_tool_usage(config_model.tools): + if config_model_copy.tools and _mcp_utils.has_mcp_tool_usage( + config_model_copy.tools + ): if config_model_copy.http_options is None: config_model_copy.http_options = types.HttpOptions(headers={}) if config_model_copy.http_options.headers is None: diff --git a/google/genai/models.py b/google/genai/models.py index d77346acf..0326b1ea8 100644 --- a/google/genai/models.py +++ b/google/genai/models.py @@ -6590,8 +6590,8 @@ def generate_content( response = types.GenerateContentResponse() i = 0 while remaining_remote_calls_afc > 0: - parsed_config_to_call = ( - parsed_config.model_copy(deep=True) if parsed_config else None + parsed_config_to_call = _extra_utils.copy_generate_content_config( + parsed_config ) function_map = _extra_utils.get_function_map(parsed_config) if function_map: @@ -6767,8 +6767,8 @@ def generate_content_stream( func_response_parts = None i = 0 while remaining_remote_calls_afc > 0: - parsed_config_to_call = ( - parsed_config.model_copy(deep=True) if parsed_config else None + parsed_config_to_call = _extra_utils.copy_generate_content_config( + parsed_config ) function_map = _extra_utils.get_function_map(parsed_config) if function_map: @@ -8695,12 +8695,8 @@ async def generate_content( ) ) - if not config: - parsed_config = None - elif isinstance(config, dict): - parsed_config = types.GenerateContentConfig(**config) - else: - parsed_config = config.model_copy(deep=True) + # Deep-copy config without pickling MCP sessions / other live tools (#2669). + parsed_config = _extra_utils.copy_generate_content_config(config) # Use AsyncExitStack to keep MCP connections alive across the entire AFC loop async with contextlib.AsyncExitStack() as stack: @@ -8805,9 +8801,7 @@ async def generate_content( is_caller_method_async=True, ) final_parsed_config_to_call = ( - final_parsed_config.model_copy(deep=True) - if final_parsed_config - else None + _extra_utils.copy_generate_content_config(final_parsed_config) ) if function_map: final_parsed_config_to_call = _extra_utils.get_usage_header( @@ -8930,12 +8924,8 @@ async def generate_content_stream( # The image shows a flat lay arrangement of freshly baked blueberry # scones. """ - if not config: - parsed_config = None - elif isinstance(config, dict): - parsed_config = types.GenerateContentConfig(**config) - else: - parsed_config = config.model_copy(deep=True) + # Deep-copy config without pickling MCP sessions / other live tools (#2669). + parsed_config = _extra_utils.copy_generate_content_config(config) incompatible_tools_indexes = ( _extra_utils.find_afc_incompatible_tool_indexes( @@ -9061,9 +9051,7 @@ async def stream_generator(): # type: ignore[no-untyped-def] ) final_parsed_config_to_call = ( - final_parsed_config.model_copy(deep=True) - if final_parsed_config - else None + _extra_utils.copy_generate_content_config(final_parsed_config) ) if function_map: final_parsed_config_to_call = _extra_utils.get_usage_header( diff --git a/google/genai/tests/mcp/test_copy_generate_content_config.py b/google/genai/tests/mcp/test_copy_generate_content_config.py new file mode 100644 index 000000000..310a1a154 --- /dev/null +++ b/google/genai/tests/mcp/test_copy_generate_content_config.py @@ -0,0 +1,211 @@ +# 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. +# + +"""Tests for copy_generate_content_config (issue #2669).""" + +import _asyncio +from types import SimpleNamespace +from unittest import mock + +import pytest + +from ... import _extra_utils +from ... import _transformers as t +from ... import models as models_module +from ... import types + +try: + from mcp import types as mcp_types + from mcp import ClientSession as McpClientSession +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.' + ) from e + else: + raise e + + +def _mock_async_models(): + api_client = SimpleNamespace(vertexai=False) + return models_module.AsyncModels(api_client) + + +class _UnpickleableMcpClientSession(McpClientSession): + + def __init__(self): + self._read_stream = None + self._write_stream = None + self._future = _asyncio.Future() + + 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'}}, + }, + ), + ] + ) + + +def test_copy_none_and_empty_dict(): + assert _extra_utils.copy_generate_content_config(None) is None + assert _extra_utils.copy_generate_content_config({}) is None + + +def test_copy_preserves_unpickleable_mcp_session(): + """Deep copy must not pickle MCP sessions that hold asyncio.Future.""" + + class MockMcpClientSession(McpClientSession): + + def __init__(self): + self._read_stream = None + self._write_stream = None + self._future = _asyncio.Future() + + session = MockMcpClientSession() + config = types.GenerateContentConfig( + temperature=0.5, + tools=[session], + ) + + with pytest.raises(TypeError, match='pickle'): + config.model_copy(deep=True) + + copied = _extra_utils.copy_generate_content_config(config) + assert copied is not config + assert copied.temperature == 0.5 + assert copied.tools is not None + assert len(copied.tools) == 1 + assert copied.tools[0] is session + # Caller config is unchanged. + assert config.tools is not None + assert config.tools[0] is session + + +def test_copy_from_dict_keeps_tool_identity(): + class MockMcpClientSession(McpClientSession): + + def __init__(self): + self._read_stream = None + self._write_stream = None + self._future = _asyncio.Future() + + session = MockMcpClientSession() + config = { + 'temperature': 0.25, + 'tools': [session], + } + copied = _extra_utils.copy_generate_content_config(config) + assert copied is not None + assert copied.temperature == 0.25 + assert copied.tools is not None + assert copied.tools[0] is session + + +def test_copy_deep_copies_other_nested_fields(): + config = types.GenerateContentConfig( + http_options=types.HttpOptions(headers={'x-test': '1'}), + ) + copied = _extra_utils.copy_generate_content_config(config) + assert copied is not None + assert copied.http_options is not None + assert copied.http_options is not config.http_options + assert copied.http_options.headers is not None + assert config.http_options is not None + assert config.http_options.headers is not None + assert copied.http_options.headers is not config.http_options.headers + copied.http_options.headers['x-test'] = '2' + assert config.http_options.headers['x-test'] == '1' + + +@pytest.mark.asyncio +async def test_async_generate_content_accepts_unpickleable_mcp_config_object(): + """Regression for #2669 on AsyncModels.generate_content.""" + async_models = _mock_async_models() + + async def fake_generate_content(self, *, model, contents, config): + del self, model, contents, config + return types.GenerateContentResponse( + candidates=[ + types.Candidate( + content=types.Content( + role='model', + parts=[types.Part(text='sunny')], + ) + ) + ] + ) + + with mock.patch.object( + models_module.AsyncModels, + '_generate_content', + fake_generate_content, + ): + response = await async_models.generate_content( + model='gemini-2.5-flash', + contents=t.t_contents('What is the weather in Boston?'), + config=types.GenerateContentConfig( + tools=[_UnpickleableMcpClientSession()] + ), + ) + assert response.text == 'sunny' + + +@pytest.mark.asyncio +async def test_async_generate_content_stream_accepts_unpickleable_mcp_config_object(): + """Regression for #2669 on AsyncModels.generate_content_stream.""" + async_models = _mock_async_models() + + async def fake_generate_content_stream(self, *, model, contents, config): + del self, model, contents, config + + async def _gen(): + yield types.GenerateContentResponse( + candidates=[ + types.Candidate( + content=types.Content( + role='model', + parts=[types.Part(text='sunny')], + ) + ) + ] + ) + + return _gen() + + with mock.patch.object( + models_module.AsyncModels, + '_generate_content_stream', + fake_generate_content_stream, + ): + stream = await async_models.generate_content_stream( + model='gemini-2.5-flash', + contents=t.t_contents('What is the weather in Boston?'), + config=types.GenerateContentConfig( + tools=[_UnpickleableMcpClientSession()] + ), + ) + chunks = [chunk async for chunk in stream] + assert chunks + assert chunks[0].text == 'sunny'