fix: avoid pickling MCP sessions when copying GenerateContentConfig

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:

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(

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'