feat: Add mTLS support for DiscoveryEngineSearchTool

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 938274974

Author: xuanyang15

src/google/adk/tools/discovery_engine_search_tool.py

@@ -28,6 +28,7 @@
 from google.cloud import discoveryengine_v1beta as discoveryengine
 from google.genai import types
 
+from ..utils._mtls_utils import get_api_endpoint
 from .function_tool import FunctionTool
 
 logger = logging.getLogger('google_adk.' + __name__)
@@ -37,6 +38,7 @@
 )
 
 _DEFAULT_ENDPOINT = 'discoveryengine.googleapis.com'
+_DEFAULT_MTLS_ENDPOINT = 'discoveryengine.mtls.googleapis.com'
 _GLOBAL_LOCATION = 'global'
 _LOCATION_PATTERN = re.compile(
     r'/locations/([a-z0-9-]+)(?:/|$)', flags=re.IGNORECASE
@@ -85,6 +87,18 @@ def _resolve_location(resource_id: str, location: Optional[str]) -> str:
   return _GLOBAL_LOCATION
 
 
+def _get_api_endpoint(location: str) -> str:
+  """Returns API endpoint based on mTLS configuration and cert availability."""
+  default_template = '{location}-' + _DEFAULT_ENDPOINT
+  mtls_template = '{location}-' + _DEFAULT_MTLS_ENDPOINT
+
+  return get_api_endpoint(
+      location=location,
+      default_template=default_template,
+      mtls_template=mtls_template,
+  )
+
+
 def _build_client_options(
     resource_id: str,
     quota_project_id: Optional[str],
@@ -95,9 +109,7 @@ def _build_client_options(
   resolved_location = _resolve_location(resource_id, location)
 
   if resolved_location != _GLOBAL_LOCATION:
-    client_options_kwargs['api_endpoint'] = (
-        f'{resolved_location}-{_DEFAULT_ENDPOINT}'
-    )
+    client_options_kwargs['api_endpoint'] = _get_api_endpoint(resolved_location)
   if quota_project_id:
     client_options_kwargs['quota_project_id'] = quota_project_id
 

tests/unittests/tools/test_discovery_engine_search_tool.py

@@ -112,16 +112,19 @@ def test_init_with_data_store_specs_without_search_engine_id_raises_error(
           ),
       ],
   )
+  @mock.patch.object(discovery_engine_search_tool, "_get_api_endpoint")
   @mock.patch.object(discovery_engine_search_tool, "client_options")
   @mock.patch.object(discoveryengine, "SearchServiceClient")
   def test_init_with_regional_location_uses_regional_endpoint(
       self,
       mock_search_client,
       mock_client_options,
+      mock_get_api_endpoint,
       tool_kwargs,
       expected_endpoint,
   ):
     """Test initialization uses the expected regional API endpoint."""
+    mock_get_api_endpoint.return_value = expected_endpoint
     DiscoveryEngineSearchTool(**tool_kwargs)
 
     mock_client_options.ClientOptions.assert_called_once_with(
@@ -132,12 +135,14 @@ def test_init_with_regional_location_uses_regional_endpoint(
         client_options=mock_client_options.ClientOptions.return_value,
     )
 
+  @mock.patch.object(discovery_engine_search_tool, "_get_api_endpoint")
   @mock.patch.object(discovery_engine_search_tool, "client_options")
   @mock.patch.object(discoveryengine, "SearchServiceClient")
   def test_init_with_explicit_location_override_uses_input_location(
-      self, mock_search_client, mock_client_options
+      self, mock_search_client, mock_client_options, mock_get_api_endpoint
   ):
     """Test initialization uses explicit location when resource has none."""
+    mock_get_api_endpoint.return_value = "eu-discoveryengine.googleapis.com"
     DiscoveryEngineSearchTool(
         data_store_id="test_data_store",
         location="eu",
@@ -239,12 +244,14 @@ def test_init_with_global_location_keeps_default_endpoint(
         credentials="credentials", client_options=None
     )
 
+  @mock.patch.object(discovery_engine_search_tool, "_get_api_endpoint")
   @mock.patch.object(discovery_engine_search_tool, "client_options")
   @mock.patch.object(discoveryengine, "SearchServiceClient")
   def test_init_with_regional_location_and_quota_project_id(
-      self, mock_search_client, mock_client_options
+      self, mock_search_client, mock_client_options, mock_get_api_endpoint
   ):
     """Test initialization uses endpoint and quota project id together."""
+    mock_get_api_endpoint.return_value = "eu-discoveryengine.googleapis.com"
     mock_credentials = mock.MagicMock()
     mock_credentials.quota_project_id = "test-quota-project"