From 2778646a2e86c5d2de6f7be75f03985d877d2ef2 Mon Sep 17 00:00:00 2001
From: Emile Ferreira <32413750+emileferreira@users.noreply.github.com>
Date: Mon, 17 Aug 2026 16:04:09 +0200
Subject: [PATCH] fix(retry): retry google-auth transport errors

---
 google/genai/_api_client.py               |  3 +-
 google/genai/tests/client/test_retries.py | 59 +++++++++++++++++++++--
 2 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/google/genai/_api_client.py b/google/genai/_api_client.py
index 3c4cbc247..e6a91395d 100644
--- a/google/genai/_api_client.py
+++ b/google/genai/_api_client.py
@@ -577,7 +577,8 @@ def retry_args(options: Optional[HttpRetryOptions]) -> _common.StringDict:
   retriable_codes = options.http_status_codes or _RETRY_HTTP_STATUS_CODES
   retry = tenacity.retry_if_exception(
       lambda e: (isinstance(e, errors.APIError) and e.code in retriable_codes)
-      or isinstance(e, _HTTPX_TRANSIENT_EXC),
+      or isinstance(e, _HTTPX_TRANSIENT_EXC)
+      or isinstance(e, auth_exceptions.TransportError),
   )
   wait = tenacity.wait_exponential_jitter(
       initial=options.initial_delay or _RETRY_INITIAL_DELAY,
diff --git a/google/genai/tests/client/test_retries.py b/google/genai/tests/client/test_retries.py
index d02711313..c4ea03639 100644
--- a/google/genai/tests/client/test_retries.py
+++ b/google/genai/tests/client/test_retries.py
@@ -35,6 +35,7 @@
   StaticCredentials = mock.MagicMock()
   AsyncAuthorizedSession = mock.MagicMock()
 
+from google.auth import exceptions as auth_exceptions
 from google.oauth2 import credentials
 import httpx
 import tenacity
@@ -187,10 +188,10 @@ def test_retry_args_enabled_with_custom_values_are_not_overridden():
       assert not retry.predicate(e)
 
 
-def test_retry_args_retries_httpx_transport_errors():
-  # httpx transport errors (timeouts, connect errors) bypass APIError but are
-  # transient infrastructure failures, so the predicate must still retry them
-  # when HttpRetryOptions is configured. See issue #2337.
+def test_retry_args_retries_transport_errors():
+  # Transport errors bypass APIError but are transient infrastructure failures,
+  # so the predicate must still retry them when HttpRetryOptions is configured.
+  # See issue #2337.
   args = api_client.retry_args(types.HttpRetryOptions())
   retry = args['retry']
 
@@ -198,9 +199,13 @@ def test_retry_args_retries_httpx_transport_errors():
   assert retry.predicate(httpx.ReadTimeout('read stalled'))
   assert retry.predicate(httpx.ConnectTimeout('connect stalled'))
   assert retry.predicate(httpx.ConnectError('connect refused'))
+  assert retry.predicate(
+      auth_exceptions.TransportError('token endpoint unavailable')
+  )
 
   # Unrelated transport errors are not retried.
   assert not retry.predicate(httpx.InvalidURL('bad url'))
+  assert not retry.predicate(auth_exceptions.RefreshError('invalid grant'))
   assert not retry.predicate(ValueError('not a transport error'))
 
 
@@ -505,6 +510,52 @@ async def run():
   asyncio.run(run())
 
 
+def test_async_retries_google_auth_transport_error():
+  api_client.has_aiohttp = False
+
+  class FlakyCredentials(credentials.Credentials):
+
+    def __init__(self):
+      super().__init__(token=None)
+      self.refresh_calls = 0
+
+    def refresh(self, request):
+      self.refresh_calls += 1
+      if self.refresh_calls == 1:
+        raise auth_exceptions.TransportError('token endpoint unavailable')
+      self.token = 'magic_token'
+
+  async def run():
+    mock_transport = mock.Mock(spec=httpx.AsyncBaseTransport)
+    mock_transport.handle_async_request.return_value = _httpx_response(200)
+    credential = FlakyCredentials()
+
+    client = api_client.BaseApiClient(
+        vertexai=True,
+        project='test_project',
+        location='global',
+        http_options=_transport_options(
+            http_options=types.HttpOptions(retry_options=_RETRY_OPTIONS),
+            async_transport=mock_transport,
+        ),
+    )
+
+    with mock.patch(
+        'google.auth.default',
+        return_value=(credential, 'test_project'),
+        autospec=True,
+    ):
+      response = await client.async_request(
+          http_method='GET', path='path', request_dict={}
+      )
+
+    assert credential.refresh_calls == 2
+    mock_transport.handle_async_request.assert_called_once()
+    assert response.headers['status-code'] == '200'
+
+  asyncio.run(run())
+
+
 def test_async_retries_failed_request_retries_successfully():
   api_client.has_aiohttp = False