feat: Add mTLS support to Google API tools

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

Author: xuanyang15

src/google/adk/integrations/parameter_manager/parameter_client.py

@@ -25,7 +25,7 @@
 from google.oauth2 import service_account
 
 from ... import version
-from ...utils import _mtls_utils
+from ...utils._mtls_utils import get_api_endpoint
 
 USER_AGENT = f"google-adk/{version.__version__}"
 
@@ -112,7 +112,7 @@ def __init__(
     client_options = None
     if location:
       client_options = {
-          "api_endpoint": _mtls_utils.get_api_endpoint(
+          "api_endpoint": get_api_endpoint(
               location,
               _DEFAULT_REGIONAL_ENDPOINT_TEMPLATE,
               _DEFAULT_MTLS_REGIONAL_ENDPOINT_TEMPLATE,

src/google/adk/tools/google_api_tool/google_api_toolset.py

@@ -14,22 +14,28 @@
 
 from __future__ import annotations
 
+import logging
 from typing import Dict
 from typing import List
 from typing import Optional
 from typing import Union
 
+import httpx
 from typing_extensions import override
 
 from ...agents.readonly_context import ReadonlyContext
 from ...auth.auth_credential import ServiceAccount
 from ...auth.auth_schemes import OpenIdConnectWithConfig
 from ...tools.base_toolset import BaseToolset
 from ...tools.base_toolset import ToolPredicate
+from ...utils._mtls_utils import MtlsClientCerts
+from ...utils._mtls_utils import use_client_cert_effective
 from ..openapi_tool import OpenAPIToolset
 from .google_api_tool import GoogleApiTool
 from .googleapi_to_openapi_converter import GoogleApiToOpenApiConverter
 
+logger = logging.getLogger('google_adk.' + __name__)
+
 
 class GoogleApiToolset(BaseToolset):
   """Google API Toolset contains tools for interacting with Google APIs.
@@ -75,6 +81,20 @@ def __init__(
     self._additional_headers = additional_headers
     self._additional_scopes = additional_scopes
     self._discovery_url = discovery_url
+
+    self._httpx_client_factory = None
+    use_client_cert = use_client_cert_effective()
+
+    if use_client_cert:
+      self._mtls_certs = MtlsClientCerts()
+      cert_path, key_path, passphrase = self._mtls_certs.get_certs()
+      if cert_path and key_path and passphrase:
+
+        def client_factory():
+          return httpx.AsyncClient(cert=(cert_path, key_path, passphrase))
+
+        self._httpx_client_factory = client_factory
+
     self._openapi_toolset = self._load_toolset_with_oidc_auth()
 
   @override
@@ -134,6 +154,7 @@ def _load_toolset_with_oidc_auth(self) -> OpenAPIToolset:
             grant_types_supported=['authorization_code'],
             scopes=scopes,
         ),
+        httpx_client_factory=self._httpx_client_factory,
     )
 
   def configure_auth(self, client_id: str, client_secret: str):
@@ -147,3 +168,5 @@ def configure_sa_auth(self, service_account: ServiceAccount):
   async def close(self):
     if self._openapi_toolset:
       await self._openapi_toolset.close()
+    if hasattr(self, '_mtls_certs') and self._mtls_certs:
+      self._mtls_certs.close()

src/google/adk/tools/google_api_tool/googleapi_to_openapi_converter.py

@@ -17,13 +17,18 @@
 import argparse
 import json
 import logging
+import socket
 from typing import Any
 from typing import Dict
 from typing import List
 
 # Google API client
 from googleapiclient.discovery import build
 from googleapiclient.errors import HttpError
+import httplib2
+
+from ...utils._mtls_utils import MtlsClientCerts
+from ...utils._mtls_utils import use_client_cert_effective
 
 # Configure logging
 logger = logging.getLogger("google_adk." + __name__)
@@ -54,6 +59,8 @@ def __init__(
         "paths": {},
         "components": {"schemas": {}, "securitySchemes": {}},
     }
+    self._use_client_cert = use_client_cert_effective()
+    self._mtls_certs = MtlsClientCerts() if self._use_client_cert else None
 
   def fetch_google_api_spec(self) -> None:
     """Fetches the Google API specification using discovery service."""
@@ -63,15 +70,42 @@ def fetch_google_api_spec(self) -> None:
           self._api_name,
           self._api_version,
       )
+
+      # Determine if we should use mTLS
+      # self._use_client_cert is already initialized in __init__
+
+      http_client = None
+      discovery_url = self._discovery_url
+
+      if self._use_client_cert and self._mtls_certs:
+        cert_path, key_path, passphrase = self._mtls_certs.get_certs()
+        if cert_path and key_path and passphrase:
+          # Set default HTTP timeout similar to googleapiclient.http.build_http()
+          http_timeout = socket.getdefaulttimeout() or 60
+          http_client = httplib2.Http(timeout=http_timeout)
+          try:
+            http_client.redirect_codes = http_client.redirect_codes - {308}
+          except AttributeError:
+            pass
+          http_client.add_certificate(key_path, cert_path, "", passphrase)
+
+          if not discovery_url:
+            discovery_url = "https://www.mtls.googleapis.com/discovery/v1/apis/{api}/{apiVersion}/rest"
+
       # Build a resource object for the specified API
-      if self._discovery_url:
+      if discovery_url:
         self._google_api_resource = build(
             self._api_name,
             self._api_version,
-            discoveryServiceUrl=self._discovery_url,
+            discoveryServiceUrl=discovery_url,
+            http=http_client,
         )
       else:
-        self._google_api_resource = build(self._api_name, self._api_version)
+        self._google_api_resource = build(
+            self._api_name,
+            self._api_version,
+            http=http_client,
+        )
 
       # Access the underlying API discovery document
       self._google_api_spec = self._google_api_resource._rootDesc
@@ -136,9 +170,13 @@ def _convert_info(self) -> None:
 
   def _convert_servers(self) -> None:
     """Convert server information."""
-    base_url = self._google_api_spec.get(
-        "rootUrl", ""
-    ) + self._google_api_spec.get("servicePath", "")
+    use_client_cert = getattr(self, "_use_client_cert", False)
+    if use_client_cert and "mtlsRootUrl" in self._google_api_spec:
+      root_url = self._google_api_spec["mtlsRootUrl"]
+    else:
+      root_url = self._google_api_spec.get("rootUrl", "")
+
+    base_url = root_url + self._google_api_spec.get("servicePath", "")
 
     # Remove trailing slash if present
     if base_url.endswith("/"):

src/google/adk/utils/_mtls_utils.py

@@ -19,6 +19,8 @@
 import enum
 import logging
 import os
+import tempfile
+import threading
 from typing import TYPE_CHECKING
 from urllib.parse import urlsplit
 from urllib.parse import urlunsplit
@@ -149,3 +151,69 @@ def configure_session_for_mtls(session: requests.Session) -> bool:
   if is_mtls:
     session.mount("https://", _MutualTlsAdapter(cert, key))
   return bool(is_mtls)
+
+
+class MtlsClientCerts:
+  """Manages the creation and lifecycle of client certificates for mTLS.
+
+  Extracts certificates to a temporary directory that is automatically cleaned up
+  when the instance is garbage collected.
+  """
+
+  def __init__(self) -> None:
+    self._tempdir: tempfile.TemporaryDirectory[str] | None = None
+    self.cert_path: str | None = None
+    self.key_path: str | None = None
+    self.passphrase: bytes | None = None
+    self._lock = threading.Lock()
+    self._initialized = False
+
+  def get_certs(self) -> tuple[str | None, str | None, bytes | None]:
+    """Extracts and returns the certificate paths and passphrase.
+
+    Returns:
+        A tuple of (cert_path, key_path, passphrase) if client certificates
+        are available, otherwise (None, None, None).
+    """
+    with self._lock:
+      if self._initialized:
+        return self.cert_path, self.key_path, self.passphrase
+
+      if not mtls.has_default_client_cert_source():
+        self._initialized = True
+        return None, None, None
+
+      self._tempdir = tempfile.TemporaryDirectory()
+      cert_path_tmp = os.path.join(self._tempdir.name, "cert.pem")
+      key_path_tmp = os.path.join(self._tempdir.name, "key.pem")
+
+      try:
+        cert_source = mtls.default_client_encrypted_cert_source(
+            cert_path_tmp, key_path_tmp
+        )
+        _, _, passphrase = cert_source()
+      except Exception as e:
+        # If extraction fails, we should fail loud.
+        self._tempdir.cleanup()
+        self._tempdir = None
+        raise RuntimeError(
+            f"Failed to extract default client certificates for mTLS: {e}"
+        ) from e
+
+      self.cert_path = cert_path_tmp
+      self.key_path = key_path_tmp
+      self.passphrase = passphrase
+      self._initialized = True
+
+      return self.cert_path, self.key_path, self.passphrase
+
+  def close(self) -> None:
+    """Manually cleans up the temporary directory."""
+    with self._lock:
+      if self._tempdir:
+        self._tempdir.cleanup()
+        self._tempdir = None
+      self.cert_path = None
+      self.key_path = None
+      self.passphrase = None
+      self._initialized = False

tests/unittests/integrations/parameter_manager/test_parameter_client.py

@@ -115,7 +115,7 @@ def test_init_with_auth_token(self, mock_pm_client_class):
       "google.adk.integrations.parameter_manager.parameter_client.default_service_credential"
   )
   @patch(
-      "google.adk.integrations.parameter_manager.parameter_client._mtls_utils.get_api_endpoint"
+      "google.adk.integrations.parameter_manager.parameter_client.get_api_endpoint"
   )
   def test_init_with_location(
       self,

tests/unittests/tools/google_api_tool/test_google_api_toolset.py

@@ -523,3 +523,44 @@ def test_init_with_tool_name_prefix(
     )
 
     assert tool_set.tool_name_prefix == tool_name_prefix
+
+  @mock.patch(
+      "google.adk.tools.google_api_tool.google_api_toolset.OpenAPIToolset"
+  )
+  @mock.patch(
+      "google.adk.tools.google_api_tool.google_api_toolset.GoogleApiToOpenApiConverter"
+  )
+  @mock.patch(
+      "google.adk.tools.google_api_tool.google_api_toolset.MtlsClientCerts"
+  )
+  @mock.patch(
+      "google.adk.tools.google_api_tool.google_api_toolset.use_client_cert_effective"
+  )
+  async def test_mtls_cleanup_on_close(
+      self,
+      mock_use_client_cert,
+      mock_mtls_certs_class,
+      mock_converter_class,
+      mock_openapi_toolset_class,
+  ):
+    """Test that mTLS temp files are cleaned up on close."""
+    mock_converter_class.return_value = mock.MagicMock()
+    mock_openapi_toolset_instance = mock.MagicMock()
+    mock_openapi_toolset_instance.close = mock.AsyncMock()
+    mock_openapi_toolset_class.return_value = mock_openapi_toolset_instance
+
+    mock_use_client_cert.return_value = True
+    mock_mtls_certs_instance = mock.MagicMock()
+    mock_mtls_certs_instance.get_certs.return_value = ("cert", "key", b"pass")
+    mock_mtls_certs_class.return_value = mock_mtls_certs_instance
+
+    tool_set = GoogleApiToolset(
+        api_name=TEST_API_NAME, api_version=TEST_API_VERSION
+    )
+
+    assert tool_set._httpx_client_factory is not None
+
+    await tool_set.close()
+
+    mock_openapi_toolset_instance.close.assert_called_once()
+    mock_mtls_certs_instance.close.assert_called_once()

tests/unittests/tools/google_api_tool/test_googleapi_to_openapi_converter.py

@@ -197,6 +197,14 @@ def calendar_api_spec():
   }
 
 
+@pytest.fixture(autouse=True)
+def disable_mtls_by_default(monkeypatch):
+  monkeypatch.setattr(
+      "google.auth.transport.mtls.should_use_client_cert",
+      lambda: False,
+  )
+
+
 @pytest.fixture
 def converter():
   """Fixture that provides a basic converter instance."""
@@ -279,7 +287,50 @@ def test_fetch_google_api_spec_with_discovery_url(
 
     assert converter._google_api_spec == calendar_api_spec
     mock_build.assert_called_once_with(
-        "calendar", "v3", discoveryServiceUrl=discovery_url
+        "calendar", "v3", discoveryServiceUrl=discovery_url, http=None
+    )
+
+  def test_fetch_google_api_spec_with_mtls(
+      self, monkeypatch, mock_api_resource, calendar_api_spec
+  ):
+    """Test fetching Google API specification with mTLS enabled."""
+    mock_build = MagicMock(return_value=mock_api_resource)
+    monkeypatch.setattr(
+        "google.adk.tools.google_api_tool.googleapi_to_openapi_converter.build",
+        mock_build,
+    )
+
+    # Enable mTLS
+    monkeypatch.setattr(
+        "google.auth.transport.mtls.should_use_client_cert",
+        lambda: True,
+    )
+    monkeypatch.setattr(
+        "google.auth.transport.mtls.has_default_client_cert_source",
+        lambda: True,
+    )
+
+    mock_cert_source = MagicMock(
+        return_value=("/path/to/cert", "/path/to/key", b"passphrase")
+    )
+    monkeypatch.setattr(
+        "google.auth.transport.mtls.default_client_encrypted_cert_source",
+        lambda c, k: mock_cert_source,
+    )
+
+    converter = GoogleApiToOpenApiConverter("calendar", "v3")
+    converter.fetch_google_api_spec()
+
+    assert converter._google_api_spec == calendar_api_spec
+
+    # Verify build was called with the http parameter set and mtls url
+    mock_build.assert_called_once()
+    _, kwargs = mock_build.call_args
+    assert "http" in kwargs
+    assert kwargs["http"] is not None
+    assert (
+        kwargs["discoveryServiceUrl"]
+        == "https://www.mtls.googleapis.com/discovery/v1/apis/{api}/{apiVersion}/rest"
     )
 
   def test_fetch_google_api_spec_error(self, monkeypatch, converter):