feat: support mTLS in sync default CUJs through google-auth migration (except custom args and custom httpx.Client)

PiperOrigin-RevId: 884659552

Author: yinghsienwu

google/genai/_api_client.py

@@ -1,4 +1,4 @@
-# Copyright 2025 Google LLC
+# Copyright 2026 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -34,7 +34,7 @@
 import sys
 import threading
 import time
-from typing import Any, AsyncIterator, Iterator, Optional, Tuple, TYPE_CHECKING, Union
+from typing import Any, AsyncIterator, Iterator, Optional, TYPE_CHECKING, Tuple, Union
 from urllib.parse import urlparse
 from urllib.parse import urlunparse
 import warnings
@@ -44,9 +44,13 @@
 import google.auth
 import google.auth.credentials
 from google.auth.credentials import Credentials
+from google.auth.transport import mtls
+from google.auth.transport.requests import AuthorizedSession
 import httpx
 from pydantic import BaseModel
 from pydantic import ValidationError
+import requests
+from requests.structures import CaseInsensitiveDict
 import tenacity
 
 from . import _common
@@ -182,12 +186,6 @@ def join_url_path(base_url: str, path: str) -> str:
 def load_auth(*, project: Union[str, None]) -> Tuple[Credentials, str]:
   """Loads google auth credentials and project id."""
 
-  ## Set GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES to false
-  ## to disable bound token sharing. Tracking on
-  ## https://github.com/googleapis/python-genai/issues/1956
-  os.environ['GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES'] = (
-      'false'
-  )
   credentials, loaded_project_id = google.auth.default(  # type: ignore[no-untyped-call]
       scopes=['https://www.googleapis.com/auth/cloud-platform'],
   )
@@ -235,7 +233,12 @@ class HttpResponse:
 
   def __init__(
       self,
-      headers: Union[dict[str, str], httpx.Headers, 'CIMultiDictProxy[str]'],
+      headers: Union[
+          dict[str, str],
+          httpx.Headers,
+          'CIMultiDictProxy[str]',
+          CaseInsensitiveDict,
+      ],
       response_stream: Union[Any, str] = None,
       byte_stream: Union[Any, bytes] = None,
   ):
@@ -245,6 +248,8 @@ def __init__(
       self.headers = {
           key: ', '.join(headers.get_list(key)) for key in headers.keys()
       }
+    elif isinstance(headers, CaseInsensitiveDict):
+      self.headers = {key: value for key, value in headers.items()}
     elif type(headers).__name__ == 'CIMultiDictProxy':
       self.headers = {
           key: ', '.join(headers.getall(key)) for key in headers.keys()
@@ -321,15 +326,22 @@ def _copy_to_dict(self, response_payload: dict[str, object]) -> None:
 
   def _iter_response_stream(self) -> Iterator[str]:
     """Iterates over chunks retrieved from the API."""
-    if not isinstance(self.response_stream, httpx.Response):
+    if not (
+        isinstance(self.response_stream, httpx.Response)
+        or isinstance(self.response_stream, requests.Response)
+    ):
       raise TypeError(
           'Expected self.response_stream to be an httpx.Response object, '
           f'but got {type(self.response_stream).__name__}.'
       )
 
     chunk = ''
     balance = 0
-    for line in self.response_stream.iter_lines():
+    if isinstance(self.response_stream, httpx.Response):
+      response_stream = self.response_stream.iter_lines()
+    else:
+      response_stream = self.response_stream.iter_lines(decode_unicode=True)
+    for line in response_stream:
       if not line:
         continue
 
@@ -593,7 +605,10 @@ def __init__(
     elif http_options and _common.is_duck_type_of(http_options, HttpOptions):
       validated_http_options = http_options
 
-    if validated_http_options.base_url_resource_scope and not validated_http_options.base_url:
+    if (
+        validated_http_options.base_url_resource_scope
+        and not validated_http_options.base_url
+    ):
       # base_url_resource_scope is only valid when base_url is set.
       raise ValueError(
           'base_url must be set when base_url_resource_scope is set.'
@@ -729,8 +744,11 @@ def __init__(
         self._http_options
     )
     self._async_httpx_client_args = async_client_args
+    self.authorized_session: Optional[AuthorizedSession] = None
 
-    if self._http_options.httpx_client:
+    if self._use_google_auth_sync():
+      self._httpx_client = None
+    elif self._http_options.httpx_client:
       self._httpx_client = self._http_options.httpx_client
     else:
       self._httpx_client = SyncHttpxClient(**client_args)
@@ -759,7 +777,14 @@ def __init__(
     self._retry = tenacity.Retrying(**retry_kwargs)
     self._async_retry = tenacity.AsyncRetrying(**retry_kwargs)
 
-  async def _get_aiohttp_session(self) -> 'aiohttp.ClientSession':
+  def _use_google_auth_sync(self) -> Union[bool, None]:
+    return self.vertexai and not (
+        self._http_options.httpx_client or self._http_options.client_args
+    )
+
+  async def _get_aiohttp_session(
+      self,
+  ) -> 'aiohttp.ClientSession':
     """Returns the aiohttp client session."""
     if (
         self._aiohttp_session is None
@@ -1003,6 +1028,11 @@ def _websocket_base_url(self) -> str:
 
   def _access_token(self) -> str:
     """Retrieves the access token for the credentials."""
+    # Set GOOGLE_API_USE_CLIENT_CERTIFICATE to true to enable bound token sharing.
+    os.environ['GOOGLE_API_USE_CLIENT_CERTIFICATE'] = 'true'
+    os.environ['GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES'] = (
+        'true'
+    )
     with self._sync_auth_lock:
       if not self._credentials:
         self._credentials, project = load_auth(project=self.project)
@@ -1041,6 +1071,12 @@ async def _get_async_auth_lock(self) -> asyncio.Lock:
 
   async def _async_access_token(self) -> Union[str, Any]:
     """Retrieves the access token for the credentials asynchronously."""
+    # Set GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES to false
+    # to disable bound token sharing. Tracking on
+    # https://github.com/googleapis/python-genai/issues/1956
+    os.environ['GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES'] = (
+        'false'
+    )
     if not self._credentials:
       async_auth_lock = await self._get_async_auth_lock()
       async with async_auth_lock:
@@ -1190,31 +1226,44 @@ def _request_once(
         else:
           data = http_request.data
 
-    if stream:
-      httpx_request = self._httpx_client.build_request(
-          method=http_request.method,
-          url=http_request.url,
-          content=data,
+    if self._use_google_auth_sync():
+      url = str(http_request.url)
+      if self.authorized_session is None:
+        self.authorized_session = AuthorizedSession(  # type: ignore[no-untyped-call]
+            self._credentials,
+            max_refresh_attempts=1,
+        )
+        # Application default SSL credentials will be used to configure mtls
+        # channel.
+        self.authorized_session.configure_mtls_channel()  # type: ignore[no-untyped-call]
+        if self.authorized_session._is_mtls and 'googleapis.com' in url:
+          if 'sandbox' in url:
+            url = url.replace(
+                'sandbox.googleapis.com', 'mtls.sandbox.googleapis.com'
+            )
+          else:
+            url = url.replace('googleapis.com', 'mtls.googleapis.com')
+      response = self.authorized_session.request(  # type: ignore[no-untyped-call]
+          method=http_request.method.upper(),
+          url=url,
+          data=data,
           headers=http_request.headers,
           timeout=http_request.timeout,
-      )
-      response = self._httpx_client.send(httpx_request, stream=stream)
-      errors.APIError.raise_for_response(response)
-      return HttpResponse(
-          response.headers, response if stream else [response.text]
+          stream=stream,
       )
     else:
-      response = self._httpx_client.request(
+      httpx_request = self._httpx_client.build_request(  # type: ignore[union-attr]
           method=http_request.method,
           url=http_request.url,
-          headers=http_request.headers,
           content=data,
+          headers=http_request.headers,
           timeout=http_request.timeout,
       )
-      errors.APIError.raise_for_response(response)
-      return HttpResponse(
-          response.headers, response if stream else [response.text]
-      )
+      response = self._httpx_client.send(httpx_request, stream=stream)  # type: ignore[union-attr]
+    errors.APIError.raise_for_response(response)
+    return HttpResponse(
+        response.headers, response if stream else [response.text]
+    )
 
   def _request(
       self,
@@ -1590,7 +1639,7 @@ def _upload_fd(
       populate_server_timeout_header(upload_headers, timeout_in_seconds)
       retry_count = 0
       while retry_count < MAX_RETRY_COUNT:
-        response = self._httpx_client.request(
+        response = self._httpx_client.request(  # type: ignore[union-attr]
             method='POST',
             url=upload_url,
             headers=upload_headers,
@@ -1642,7 +1691,7 @@ def download_file(
       else:
         data = http_request.data
 
-    response = self._httpx_client.request(
+    response = self._httpx_client.request(  # type: ignore[union-attr]
         method=http_request.method,
         url=http_request.url,
         headers=http_request.headers,
@@ -1956,8 +2005,10 @@ def close(self) -> None:
     """Closes the API client."""
     # Let users close the custom client explicitly by themselves. Otherwise,
     # close the client when the object is garbage collected.
-    if not self._http_options.httpx_client:
+    if not self._http_options.httpx_client and self._httpx_client:
       self._httpx_client.close()
+    if self.authorized_session:
+      self.authorized_session.close()  # type: ignore[no-untyped-call]
 
   async def aclose(self) -> None:
     """Closes the API async client."""
@@ -1986,6 +2037,7 @@ def __del__(self) -> None:
     except Exception:  # pylint: disable=broad-except
       pass
 
+
 def get_token_from_credentials(
     client: 'BaseApiClient',
     credentials: google.auth.credentials.Credentials
@@ -1998,6 +2050,7 @@ def get_token_from_credentials(
     raise RuntimeError('Could not resolve API token from the environment')
   return credentials.token  # type: ignore[no-any-return]
 
+
 async def async_get_token_from_credentials(
     client: 'BaseApiClient',
     credentials: google.auth.credentials.Credentials

google/genai/errors.py

@@ -18,7 +18,7 @@
 from typing import Any, Callable, Optional, TYPE_CHECKING, Union
 import httpx
 import json
-import websockets
+import requests
 from . import _common
 
 
@@ -30,7 +30,11 @@
 class APIError(Exception):
   """General errors raised by the GenAI API."""
   code: int
-  response: Union['ReplayResponse', httpx.Response, 'aiohttp.ClientResponse']
+  response: Union[
+      requests.Response,
+      'ReplayResponse',
+      httpx.Response,
+  ]
 
   status: Optional[str] = None
   message: Optional[str] = None
@@ -40,7 +44,11 @@ def __init__(
       code: int,
       response_json: Any,
       response: Optional[
-          Union['ReplayResponse', httpx.Response, 'aiohttp.ClientResponse']
+          Union[
+              requests.Response,
+              'ReplayResponse',
+              httpx.Response,
+          ]
       ] = None,
   ):
     if isinstance(response_json, list) and len(response_json) == 1:
@@ -112,7 +120,7 @@ def _to_replay_record(self) -> _common.StringDict:
 
   @classmethod
   def raise_for_response(
-      cls, response: Union['ReplayResponse', httpx.Response]
+      cls, response: Union['ReplayResponse', httpx.Response, requests.Response]
   ) -> None:
     """Raises an error with detailed error message if the response has an error status."""
     if response.status_code == 200:
@@ -128,6 +136,16 @@ def raise_for_response(
             'message': message,
             'status': response.reason_phrase,
         }
+    elif isinstance(response, requests.Response):
+      try:
+        # do not do any extra muanipulation on the response.
+        # return the raw response json as is.
+        response_json = response.json()
+      except requests.exceptions.JSONDecodeError:
+        response_json = {
+            'message': response.text,
+            'status': response.reason,
+        }
     else:
       response_json = response.body_segments[0].get('error', {})
 
@@ -139,7 +157,11 @@ def raise_error(
       status_code: int,
       response_json: Any,
       response: Optional[
-          Union['ReplayResponse', httpx.Response, 'aiohttp.ClientResponse']
+          Union[
+              'ReplayResponse',
+              httpx.Response,
+              requests.Response,
+          ]
       ],
   ) -> None:
     """Raises an appropriate APIError subclass based on the status code.

google/genai/tests/client/test_client_close.py

@@ -43,6 +43,7 @@ def test_close_httpx_client():
       vertexai=True,
       project='test_project',
       location='global',
+      http_options=api_client.HttpOptions(client_args={'max_redirects': 10}),
   )
   client.close()
   assert client._api_client._httpx_client.is_closed
@@ -55,6 +56,7 @@ def test_httpx_client_context_manager():
       vertexai=True,
       project='test_project',
       location='global',
+      http_options=api_client.HttpOptions(client_args={'max_redirects': 10}),
   ) as client:
     pass
     assert not client._api_client._httpx_client.is_closed

google/genai/tests/client/test_client_initialization.py

@@ -20,6 +20,7 @@
 import concurrent.futures
 import logging
 import os
+import requests
 import ssl
 from unittest import mock
 
@@ -1332,13 +1333,16 @@ def refresh_side_effect(request):
   mock_creds.refresh = mock_refresh
 
   # Mock the actual request to avoid network calls
-  mock_httpx_response = httpx.Response(
-      status_code=200,
-      headers={},
-      text='{"candidates": [{"content": {"parts": [{"text": "response"}]}}]}',
+  mock_http_response = requests.Response()
+  mock_http_response.status_code = 200
+  mock_http_response.headers = {}
+  mock_http_response._content = (
+      b'{"candidates": [{"content": {"parts": [{"text": "response"}]}}]}'
+  )
+  mock_request = mock.Mock(return_value=mock_http_response)
+  monkeypatch.setattr(
+     google.auth.transport.requests.AuthorizedSession, "request", mock_request
   )
-  mock_request = mock.Mock(return_value=mock_httpx_response)
-  monkeypatch.setattr(api_client.SyncHttpxClient, "request", mock_request)
 
   client = Client(
       vertexai=True, project="fake_project_id", location="fake-location"