first vommit

2026-07-04 14:31:50 +02:00
commit a04e68bc95
@@ -0,0 +1,17 @@
+# 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 the Google GenAI SDK's _mcp_utils module."""
@@ -0,0 +1,89 @@
+# 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 typing
+from typing import Any
+
+from ... import _mcp_utils
+from ... import types
+
+if typing.TYPE_CHECKING:
+  from mcp.types import Tool as McpTool
+  from mcp import ClientSession as McpClientSession
+else:
+  McpTool: typing.Type = Any
+  McpClientSession: typing.Type = Any
+  try:
+    from mcp.types import Tool as McpTool
+    from mcp import ClientSession as McpClientSession
+  except ImportError:
+    McpTool = None
+    McpClientSession = None
+
+
+def test_mcp_tools():
+  """Test whether the list of tools contains any MCP tools."""
+  if McpTool is None:
+    return
+  mcp_tools = [
+      McpTool(
+          name='tool',
+          description='tool-description',
+          inputSchema={
+              'type': 'OBJECT',
+              'properties': {
+                  'key1': {'type': 'STRING'},
+                  'key2': {'type': 'NUMBER'},
+              },
+          },
+      ),
+  ]
+  assert _mcp_utils.has_mcp_tool_usage(mcp_tools)
+
+
+def test_mcp_client_session():
+  """Test whether the list of tools contains any MCP tools."""
+
+  class MockMcpClientSession(McpClientSession):
+
+    def __init__(self):
+      self._read_stream = None
+      self._write_stream = None
+
+  mcp_tools = [
+      MockMcpClientSession(),
+  ]
+  assert _mcp_utils.has_mcp_tool_usage(mcp_tools)
+
+
+def test_no_mcp_tools():
+  if McpClientSession is None:
+    return
+  """Test whether the list of tools contains any MCP tools."""
+  gemini_tools = [
+      types.Tool(
+          function_declarations=[
+              types.FunctionDeclaration(
+                  name='tool',
+                  description='tool-description',
+                  parameters=types.Schema(
+                      type='OBJECT',
+                      properties={},
+                  ),
+              ),
+          ],
+      ),
+  ]
+  assert not _mcp_utils.has_mcp_tool_usage(gemini_tools)
@@ -0,0 +1,371 @@
+# 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 contextlib
+from unittest import mock
+
+import pytest
+
+from ... import _mcp_utils
+from ... import types
+from ..._api_client import BaseApiClient
+
+
+try:
+  from mcp import types as mcp_types
+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 test_empty_mcp_tools_list():
+  """Test conversion of empty MCP tools list to Gemini tools list."""
+  result = _mcp_utils.mcp_to_gemini_tools([])
+
+  assert result == []
+
+
+def test_unknown_field_conversion():
+  """Test conversion of MCP tools with unknown fields to Gemini tools."""
+  mcp_tools = [
+      mcp_types.Tool(
+          name='tool',
+          description='tool-description',
+          inputSchema={
+              'type': 'object',
+              'properties': {},
+              'unknown_field': 'unknownField',
+              'unknown_object': {},
+          },
+      ),
+  ]
+  result = _mcp_utils.mcp_to_gemini_tools(mcp_tools)
+  assert result == [
+      types.Tool(
+          function_declarations=[
+              types.FunctionDeclaration(
+                  name='tool',
+                  description='tool-description',
+                  parameters=types.Schema(
+                      type='OBJECT',
+                      properties={},
+                  ),
+              ),
+          ],
+      ),
+  ]
+
+
+def test_items_conversion():
+  """Test conversion of MCP tools with items to Gemini tools."""
+  mcp_tools = [
+      mcp_types.Tool(
+          name='tool',
+          description='tool-description',
+          inputSchema={
+              'type': 'array',
+              'items': {
+                  'type': 'object',
+                  'properties': {
+                      'key1': {
+                          'type': 'string',
+                      },
+                      'key2': {
+                          'type': 'number',
+                      },
+                  },
+              },
+          },
+      ),
+  ]
+  result = _mcp_utils.mcp_to_gemini_tools(mcp_tools)
+  assert result == [
+      types.Tool(
+          function_declarations=[
+              types.FunctionDeclaration(
+                  name='tool',
+                  description='tool-description',
+                  parameters=types.Schema(
+                      type='ARRAY',
+                      items=types.Schema(
+                          type='OBJECT',
+                          properties={
+                              'key1': types.Schema(type='STRING'),
+                              'key2': types.Schema(type='NUMBER'),
+                          },
+                      ),
+                  ),
+              ),
+          ],
+      ),
+  ]
+
+
+def test_any_of_conversion():
+  """Test conversion of MCP tools with any_of to Gemini tools."""
+  mcp_tools = [
+      mcp_types.Tool(
+          name='tool',
+          description='tool-description',
+          inputSchema={
+              'type': 'object',
+              'any_of': [
+                  {
+                      'type': 'string',
+                  },
+                  {
+                      'type': 'number',
+                  },
+              ],
+          },
+      ),
+  ]
+  result = _mcp_utils.mcp_to_gemini_tools(mcp_tools)
+  assert result == [
+      types.Tool(
+          function_declarations=[
+              types.FunctionDeclaration(
+                  name='tool',
+                  description='tool-description',
+                  parameters=types.Schema(
+                      type='OBJECT',
+                      any_of=[
+                          types.Schema(type='STRING'),
+                          types.Schema(type='NUMBER'),
+                      ],
+                  ),
+              ),
+          ],
+      ),
+  ]
+
+
+def test_properties_conversion():
+  """Test conversion of MCP tools with properties to Gemini tools."""
+  mcp_tools = [
+      mcp_types.Tool(
+          name='tool',
+          description='tool-description',
+          inputSchema={
+              'type': 'object',
+              'properties': {
+                  'key1': {
+                      'type': 'string',
+                  },
+                  'key2': {
+                      'type': 'number',
+                  },
+              },
+          },
+      ),
+  ]
+  result = _mcp_utils.mcp_to_gemini_tools(mcp_tools)
+  assert result == [
+      types.Tool(
+          function_declarations=[
+              types.FunctionDeclaration(
+                  name='tool',
+                  description='tool-description',
+                  parameters=types.Schema(
+                      type='OBJECT',
+                      properties={
+                          'key1': types.Schema(type='STRING'),
+                          'key2': types.Schema(type='NUMBER'),
+                      },
+                  ),
+              ),
+          ],
+      ),
+  ]
+
+
+def test_defs_conversion():
+    """Test conversion of MCP tools with shared definitions ($defs)."""
+    mcp_tools = [
+        mcp_types.Tool(
+            name='create_endpoint',
+            description='Creates an endpoint',
+            inputSchema={
+                'type': 'object',
+                'properties': {
+                    'machine_spec': {'$ref': '#/$defs/MachineSpec'}
+                },
+                '$defs': {
+                    'MachineSpec': {
+                        'type': 'object',
+                        'properties': {'machine_type': {'type': 'string'}}
+                    }
+                }
+            },
+        ),
+    ]
+
+    result = _mcp_utils.mcp_to_gemini_tools(mcp_tools, is_agent_platform=True)
+
+    schema = result[0].function_declarations[0].parameters_json_schema
+    assert '$defs' in schema
+    assert 'MachineSpec' in schema['$defs']
+
+
+def test_create_endpoint_one_of_conversion():
+    """Test oneOf translation for Vertex create_endpoint resource selection."""
+
+    mcp_tools = [
+        mcp_types.Tool(
+            name='create_endpoint',
+            description='Creates a Vertex AI Endpoint resource.',
+            inputSchema={
+                'type': 'object',
+                'properties': {
+                    'endpoint': {
+                        'type': 'object',
+                        'oneOf': [
+                            {'title': 'dedicated_resources', 'type': 'object'},
+                            {'title': 'automatic_resources', 'type': 'object'},
+                        ],
+                    }
+                },
+            },
+        ),
+    ]
+    result = _mcp_utils.mcp_to_gemini_tools(mcp_tools, is_agent_platform=True)
+    schema = result[0].function_declarations[0].parameters_json_schema
+    endpoint_schema = schema['properties']['endpoint']
+
+    assert 'oneOf' in endpoint_schema
+    assert len(endpoint_schema['oneOf']) == 2
+
+
+def test_update_endpoint_labels_conversion():
+    """Test additionalProperties translation for Vertex resource labels."""
+
+    mcp_tools = [
+        mcp_types.Tool(
+            name='update_endpoint',
+            description='Updates a Vertex AI Endpoint resource.',
+            inputSchema={
+                'type': 'object',
+                'properties': {
+                    'endpoint': {
+                        'type': 'object',
+                        'properties': {
+                            'labels': {
+                                'type': 'object',
+                                'additionalProperties': {'type': 'string'}
+                            }
+                        }
+                    }
+                }
+            },
+        ),
+    ]
+    result = _mcp_utils.mcp_to_gemini_tools(mcp_tools, is_agent_platform=True)
+    schema = result[0].function_declarations[0].parameters_json_schema
+    labels_schema = schema['properties']['endpoint']['properties']['labels']
+
+    assert 'additionalProperties' in labels_schema
+
+
+def test_agent_platform_preserves_unknown_fields():
+    """Test that Agent Platform translation passes all schema fields directly
+    to the backend.
+    """
+    mcp_tools = [
+        mcp_types.Tool(
+            name='tool',
+            description='tool-description',
+            inputSchema={
+                'type': 'object',
+                'properties': {},
+                # A new unknown field
+                'some_new_future_field': 'value',
+            },
+        ),
+    ]
+
+    result = _mcp_utils.mcp_to_gemini_tools(mcp_tools, is_agent_platform=True)
+    schema = result[0].function_declarations[0].parameters_json_schema
+
+    # Verify the entire schema is passed through intact, including the unknown field
+    assert 'some_new_future_field' in schema
+    assert schema['some_new_future_field'] == 'value'
+
+
+@pytest.mark.asyncio
+@mock.patch('httpx.AsyncClient')
+@mock.patch.object(_mcp_utils, 'streamable_http_client')
+@mock.patch.object(_mcp_utils, 'McpClientSession')
+@mock.patch('google.auth.default')
+async def test_connect_agent_platform_mcp_url_and_headers(
+    mock_auth_default, mock_session_cls, mock_streamable, mock_create_http
+):
+    """Tests that _mcp_utils._connect_agent_platform_mcp builds the correct
+    regional URL and injects auth headers.
+    """
+
+    mock_creds = mock.Mock()
+    mock_creds.token = 'fake-oauth-token'
+    mock_auth_default.return_value = (mock_creds, 'fake-project')
+
+    @contextlib.asynccontextmanager
+    async def mock_streamable_ctx(*args, **kwargs):
+      yield (mock.Mock(), mock.Mock(), mock.Mock())
+
+    mock_streamable.side_effect = mock_streamable_ctx
+
+    @contextlib.asynccontextmanager
+    async def mock_session_ctx(*args, **kwargs):
+      session_instance = mock.AsyncMock()
+      yield session_instance
+
+    mock_session_cls.side_effect = mock_session_ctx
+
+    mock_http_client_instance = mock.AsyncMock()
+    mock_http_client_instance.__aenter__.return_value = (
+        mock_http_client_instance
+    )
+    mock_http_client_instance.__aexit__.return_value = None
+    mock_create_http.return_value = mock_http_client_instance
+
+    api_client = BaseApiClient(
+      vertexai=True,
+      project='test-project-123',
+      location='europe-west4'
+    )
+
+    async with _mcp_utils._connect_agent_platform_mcp(
+        api_client, 'endpoints'
+    ) as session:
+      session.initialize.assert_awaited_once()
+
+    mock_streamable.assert_called_once()
+    assert (
+        mock_streamable.call_args.kwargs['url']
+        == 'https://europe-west4-aiplatform.googleapis.com/mcp/endpoints'
+    )
+
+    mock_create_http.assert_called_once()
+    called_headers = mock_create_http.call_args.kwargs['headers']
+
+    assert called_headers['Authorization'] == 'Bearer fake-oauth-token'
+    assert called_headers['X-Goog-User-Project'] == 'test-project-123'
+    assert 'mcp_used' in called_headers.get('x-goog-api-client', '')
@@ -0,0 +1,201 @@
+# 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 _asyncio
+import pytest
+from ... import _extra_utils
+from ... import types
+from ..._adapters import McpToGenAiToolAdapter
+
+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
+
+
+@pytest.mark.asyncio
+async def test_parse_empty_config_dict():
+  """Test conversion of empty GenerateContentConfigDict to parsed config."""
+  config = {}
+  parsed_config, mcp_to_genai_tool_adapters = (
+      await _extra_utils.parse_config_for_mcp_sessions(config)
+  )
+  assert parsed_config == None
+  assert not mcp_to_genai_tool_adapters
+
+
+@pytest.mark.asyncio
+async def test_parse_empty_config_object():
+  """Test conversion of empty GenerateContentConfig to parsed config."""
+  config = types.GenerateContentConfig()
+  parsed_config, mcp_to_genai_tool_adapters = (
+      await _extra_utils.parse_config_for_mcp_sessions(config)
+  )
+  assert config is not parsed_config  # config is not modified
+  assert not mcp_to_genai_tool_adapters
+
+
+@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 = (
+      await _extra_utils.parse_config_for_mcp_sessions(config)
+  )
+  assert len(config.tools) == 1
+  assert config.tools[0] is mock_session_instance
+  assert config is not parsed_config  # config is not modified
+  assert len(mcp_to_genai_tool_adapters) == 2
+  assert mcp_to_genai_tool_adapters.keys() == {
+      'get_weather',
+      'get_weather_2',
+  }
+  assert isinstance(
+      mcp_to_genai_tool_adapters['get_weather'], McpToGenAiToolAdapter
+  )
+  assert isinstance(
+      mcp_to_genai_tool_adapters['get_weather_2'], McpToGenAiToolAdapter
+  )
+
+
+@pytest.mark.asyncio
+async def test_parse_config_object_with_tools_complex_type():
+  """Test conversion of GenerateContentConfig with tools to parsed config."""
+
+  class MockMcpClientSession(McpClientSession):
+
+    def __init__(self):
+      self._read_stream = None
+      self._write_stream = None
+      self._future = _asyncio.Future()  # This object cannot be pickled.
+
+    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 = (
+      await _extra_utils.parse_config_for_mcp_sessions(config)
+  )
+  assert len(config.tools) == 1
+  assert config.tools[0] is mock_session_instance
+  assert config is not parsed_config  # config is not modified
+  assert len(mcp_to_genai_tool_adapters) == 2
+  assert mcp_to_genai_tool_adapters.keys() == {
+      'get_weather',
+      'get_weather_2',
+  }
+  assert isinstance(
+      mcp_to_genai_tool_adapters['get_weather'], McpToGenAiToolAdapter
+  )
+  assert isinstance(
+      mcp_to_genai_tool_adapters['get_weather_2'], McpToGenAiToolAdapter
+  )
+
+
+@pytest.mark.asyncio
+async def test_parse_config_object_with_non_mcp_tools():
+  """Test conversion of GenerateContentConfig with regular tools to parsed config."""
+
+  config = types.GenerateContentConfig(
+      tools=[
+          types.Tool(
+              function_declarations=[
+                  {'name': 'tool-1', 'description': 'tool-1-description'}
+              ]
+          ),
+          types.Tool(
+              function_declarations=[
+                  {'name': 'tool-2', 'description': 'tool-2-description'}
+              ]
+          ),
+      ]
+  )
+  parsed_config, mcp_to_genai_tool_adapters = (
+      await _extra_utils.parse_config_for_mcp_sessions(config)
+  )
+  assert len(config.tools) == 2
+  assert config is not parsed_config  # config is not modified
+  assert mcp_to_genai_tool_adapters == {}
+  assert parsed_config.tools == [
+      types.Tool(
+          function_declarations=[
+              {'name': 'tool-1', 'description': 'tool-1-description'}
+          ]
+      ),
+      types.Tool(
+          function_declarations=[
+              {'name': 'tool-2', 'description': 'tool-2-description'}
+          ]
+      ),
+  ]
@@ -0,0 +1,130 @@
+# 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
+import pytest
+from ... import _extra_utils
+from ... import types
+from ..._adapters import McpToGenAiToolAdapter
+
+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 test_parse_empty_config_dict():
+  """Test conversion of empty GenerateContentConfigDict to parsed config."""
+  config = {}
+  parsed_config = _extra_utils.parse_config_for_mcp_usage(config)
+  assert parsed_config == None
+
+
+def test_parse_empty_config_object():
+  """Test conversion of empty GenerateContentConfig to parsed config."""
+  config = types.GenerateContentConfig()
+  parsed_config = _extra_utils.parse_config_for_mcp_usage(config)
+  assert config is not parsed_config  # config is not modified
+
+
+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 = _extra_utils.parse_config_for_mcp_usage(config)
+  assert config.http_options is None
+  assert config is not parsed_config  # config is not modified
+  assert re.match(
+      r'mcp_used/\d+\.\d+\.\d+',
+      parsed_config.http_options.headers['x-goog-api-client'],
+  )
+
+
+def test_parse_config_object_with_tools_and_existing_headers():
+  """Test conversion of GenerateContentConfig with tools and existing headers to parsed config."""
+
+  config = types.GenerateContentConfig(
+      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'}},
+              },
+          ),
+      ],
+      http_options=types.HttpOptions(
+          headers={
+              'x-goog-api-client': 'google-genai-sdk/1.0.0 gl-python/1.0.0'
+          }
+      ),
+  )
+  parsed_config = _extra_utils.parse_config_for_mcp_usage(config)
+  assert not re.match(
+      r'mcp_used/\d+\.\d+\.\d+',
+      config.http_options.headers['x-goog-api-client'],
+  )
+  assert config is not parsed_config  # config is not modified
+  assert re.match(
+      r'google-genai-sdk/1.0.0 gl-python/1.0.0 mcp_used/\d+\.\d+\.\d+',
+      parsed_config.http_options.headers['x-goog-api-client'],
+  )
@@ -0,0 +1,72 @@
+# 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.
+#
+
+from importlib.metadata import version
+import re
+import typing
+from typing import Any
+
+from ... import _mcp_utils
+from ... import types
+
+_is_mcp_imported = False
+if typing.TYPE_CHECKING:
+  import mcp
+
+  _is_mcp_imported = True
+else:
+  try:
+    import mcp
+
+    _is_mcp_imported = True
+  except ImportError:
+    _is_mcp_imported = False
+
+
+def test_set_mcp_usage_header_from_empty_dict():
+  if not _is_mcp_imported:
+    return
+  """Test whether the MCP usage header is set correctly from an empty dict."""
+  headers = {}
+  _mcp_utils.set_mcp_usage_header(headers)
+  assert re.match(r'mcp_used/\d+\.\d+\.\d+', headers['x-goog-api-client'])
+
+
+def test_set_mcp_usage_header_with_existing_header():
+  if not _is_mcp_imported:
+    return
+  """Test whether the MCP usage header is set correctly from an existing header."""
+  headers = {'x-goog-api-client': 'google-genai-sdk/1.0.0 gl-python/1.0.0'}
+  _mcp_utils.set_mcp_usage_header(headers)
+  assert re.match(
+      r'google-genai-sdk/1.0.0 gl-python/1.0.0 mcp_used/\d+\.\d+\.\d+',
+      headers['x-goog-api-client'],
+  )
+
+
+def test_set_mcp_usage_header_with_existing_mcp_header():
+  if not _is_mcp_imported:
+    return
+  """Test whether the MCP usage header is set correctly from an existing MCP header."""
+  headers = {
+      'x-goog-api-client': (
+          'google-genai-sdk/1.0.0 gl-python/1.0.0 mcp_used/1.0.0'
+      )
+  }
+  _mcp_utils.set_mcp_usage_header(headers)
+  assert re.match(
+      r'google-genai-sdk/1.0.0 gl-python/1.0.0 mcp_used/\d+\.\d+\.\d+',
+      headers['x-goog-api-client'],
+  )