feat(auth): implement regional access boundary support for standalone JWT and async service accounts by nbayati · Pull Request #17025 · googleapis/google-cloud-python

nbayati · 2026-05-11T16:36:44Z

This PR implements the following changes:

Add RAB support to async service account credential type, by providing async manager and fetching methods.
Update unit tests to accept both mtls and standard lookup endpoint urls.
Refactor before_request to use a _after_refresh hook so we don't have to override the method.
Add RAb support for self signed jwt flow through jwt.py
some small enhancements for test coverage and backward compatibility

…ement async refresh manager

…ndpoint.

…support Regional Access Boundary logic

gemini-code-assist

Code Review

This pull request implements asynchronous support for Regional Access Boundary (RAB) management, including background refresh tasks and mTLS endpoint support. Key changes include the addition of _AsyncRegionalAccessBoundaryRefreshManager, updates to JWT and service account credentials to handle RAB state during cloning and serialization, and comprehensive test coverage for these new flows. However, a critical issue was identified where the refresh method in google/auth/jwt.py was renamed to _perform_refresh_token, which will break token updates as the base class expects a refresh method. Additionally, a typo was found in a test assertion URL.

lsirac · 2026-05-12T21:02:56Z

+        except asyncio.TimeoutError:
+            return False, {}, False
+
+        response_body1 = await response.content()


This is an async_generator, not an awaitable method. I think this crashes as written. Should call await response.read() instead.

Let's also add a test that will catch this, and cover in try/catch block

I actually was able to run the code as-is and verify that it works. We have the same pattern at the top of this file in _token_endpoint_request_no_throw method.

Though I did dig a bit to see why it's working and this is my current understanding:
in google-auth the async Request transport adapter wraps the raw response in a custom _CombinedResponse class (defined in google/auth/transport/_aiohttp_requests.py).
The _CombinedResponse class explicitly implements content as an asynchronous method to mirror the synchronous requests interface. Because it is a coroutine method returning the decoded byte body, await response.content() works as properly.

I renamed the variable to response_bytes though, to make it a bit more intuitive.

Ty for validating. If this is only intended to support the legacy _aiohttp_requests transport then great, but we do then need to update docstrings (currently says google.auth.aio.transport.Request)

Done. Updated docstrings to document support for both standard (read()) and legacy (content()) async responses.

lsirac · 2026-05-12T21:08:05Z

+        try:
+            if timeout:
+                response = await asyncio.wait_for(
+                    request(method="GET", url=url, headers=headers), timeout=timeout


Should timeout be passed to the request here?

Good catch! I've added the timeout to the request.

lsirac · 2026-05-12T21:11:42Z

+                )
+            else:
+                response = await request(method="GET", url=url, headers=headers)
+        except asyncio.TimeoutError:


What happens if the request fails? I think we may need to catch additional errors here (timeout error, transport error..)

Good catch! I have updated the code to catch google.auth.exceptions.TransportError along with asyncio.TimeoutError and return (False, {}, False) on failure.

If a request fails due to a connection issue or timeout, the method now returns a non-retryable failure, which triggers the background refresh manager to enter the cooldown period. Any other unexpected exceptions will be caught and logged as warnings by the background worker task.

lsirac · 2026-05-12T21:15:19Z

@@ -369,6 +401,8 @@ def _copy_regional_access_boundary_manager(self, target):
        # but share the immutable data reference to avoid unnecessary initial lookups.
        new_manager = _regional_access_boundary_utils._RegionalAccessBoundaryManager()
        new_manager._data = self._rab_manager._data
+        # Preserve the type of refresh manager (sync or async)
+        new_manager.refresh_manager = self._rab_manager.refresh_manager.__class__()


I don't think this is safe across sync/async creds.

Using the source refresh-manager means we can have an async refresh manager on a sync cred which can later call asyncio.create_task() from a sync before_request.

_use_blocking_regional_access_boundary_lookup is not kept, so a credential configured for blocking RAB can become non-blocking after with_scopes(), with_quota_project(), etc (breaking gcloud)

I think the cred should keep its initialized manager type and we should only copy the RAB data/config.

You're right, I had missed that the initialization of the credential would take care of the manager and didn't have to copy/create it here! Updated this method to only copy over the RAB data, and also _use_blocking_regional_access_boundary_lookup.

lsirac · 2026-05-12T21:50:31Z

+        except ValueError:
+            response_data = response_body
+
+        if response.status == http_client.OK:


In google.auth.aio.transport.Response, the HTTP status code is exposed via the status_code property, not status. Passing a compliant google.auth.aio transport callable raises AttributeError: 'Response' object has no attribute 'status'. Please update the async lookup and grant methods to check .status_code.

Actually decided to only apply the fix to the RAB lookup, and instead open a bug to bring the grant methods up to date separately. This was can keep the blast radius of this PR smaller and limit it to only RAB changes without touching any token fetching flows.

Created #17139 to track the necessary changes for the token endpoints.

lsirac · 2026-05-12T21:52:34Z

@@ -288,3 +289,145 @@ async def refresh_grant(
        request, token_uri, body, can_retry=can_retry
    )
    return client._handle_refresh_grant_response(response_data, refresh_token)
+
+


is _jwt_async.py out of scope?

lsirac · 2026-05-12T21:53:12Z

        await credentials_async.Credentials.before_request(
            self, request, method, url, headers
        )
+        self._maybe_start_regional_access_boundary_refresh(request, url)
+        self._rab_manager.apply_headers(headers)


This may be redundant, why is this needed?

Thanks for flagging this. I actually realized we no longer need the before_request override, if we follow what we did with the sync RAB flow and using the _after_refresh.

With this hook in place, the base class now automatically orchestrates the token refresh, RAB lookup, and header application in the correct order.

lsirac · 2026-05-12T21:55:17Z

+                # A refresh is already in progress.
+                return
+
+            async def _worker():


Unlike the synchronous refresh manager which safely deepcopies the transport (copied_request = copy.deepcopy(request)), the async manager passes the exact same request instance directly into the background coroutine task. Because start_refresh is invoked inside before_request, the main application coroutine immediately proceeds to make its actual service API HTTP call using the exact same request transport while the background task is concurrently using it, risking HTTP state corruption or interleaved headers.

Additionally, spawning asyncio.create_task(_worker()) without tracking cancellation hooks upon client session closure can potentially cause dangling tasks that raise RuntimeError: Session is closed when executing against closed client sessions.

Both concerns are valid, but safe under the hood:

Deepcopying the Transport is impossible and not needed in async
The async transport object (e.g., aiohttp_requests.Request) contains an active aiohttp.ClientSession with open TCP sockets. Attempting to copy.deepcopy(request) will instantly raise a TypeError: cannot pickle 'ClientSession' object and crash the application at runtime. Unlike synchronous transports running on separate OS threads, async HTTP clients (like aiohttp or httpx) are natively designed to handle concurrent requests sharing the same connection session. All request-specific states (headers, payloads) are stored in localized coroutine call stack frames, preventing any HTTP state corruption or interleaving.

Session closure and dangling tasks are handled safely
The background worker is a single-shot asyncio task that executes exactly one lookup request and then immediately terminates and gets garbage-collected.
If the user's application closes the underlying client session while a background task is still running, the resulting RuntimeError: Session is closed is safely caught by the worker's generic except Exception as e: block. It logs a warning, fails open cleanly, and does not raise an unhandled exception or crash the application.

I think no code changes are required here, as the current design is fully protected.

lsirac · 2026-05-12T21:56:50Z

@@ -66,6 +67,12 @@ class Credentials(
        credentials = credentials.with_quota_project('myproject-123')
    """

+    def __init__(self, *args, **kwargs):


_service_account_async.Credentials lacks a __setstate__ override. When older pickled async credentials unpickle, they fall back to CredentialsWithRegionalAccessBoundary.__setstate__, which attaches a synchronous _RegionalAccessBoundaryRefreshManager. If a background RAB lookup triggers on this unpickled credential, a synchronous background thread will invoke async def _lookup_regional_access_boundary synchronously without awaiting it, causing a fatal thread crash (AttributeError: 'coroutine' object has no attribute 'get'). Please implement __setstate__ to ensure self._rab_manager.refresh_manager is always restored as an _AsyncRegionalAccessBoundaryRefreshManager().

Good catch. I've added __setstate__ and a unit test.

…y data and config

… timeout fixes

…e hook

…ager in service account credentials

…used on RAB

… pickling docs

…nc jwt

lsirac · 2026-05-16T02:42:57Z

-        new_manager._data = self._rab_manager._data
-        target._rab_manager = new_manager
+        """Copies the regional access boundary manager state to another instance."""
+        target._rab_manager._data = self._rab_manager._data


Tests are failing:

FAILED tests/test_external_account.py::TestCredentials::test_with_scopes_full_options_propagated - AttributeError: 'CredentialsImpl' object has no attribute '_rab_manager'
FAILED tests/test_external_account.py::TestCredentials::test_with_quota_project_full_options_propagated - AttributeError: 'CredentialsImpl' object has no attribute '_rab_manager'

Probably tests are configured incorrectly? I don't think it's possible for there to not be a rab manager?

lsirac · 2026-05-16T02:48:20Z

+            if hasattr(response, "read"):
+                response_bytes = await response.read()
+            else:
+                response_bytes = await response.content()


This can raise an error too, but currently not being caught.

lsirac · 2026-05-16T02:50:34Z

+                else response_bytes
+            )
+            response_data = json.loads(response_body)
+        except (UnicodeDecodeError, ValueError):


As written, a retryable error response without a JSON body would not be retried because we return here before checking the status.

lsirac · 2026-05-16T03:03:17Z

+        expected_url_standard = "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/{}/allowedLocations".format(
+            self.SERVICE_ACCOUNT_EMAIL
+        )
+        expected_url_mtls = "https://iamcredentials.mtls.googleapis.com/v1/projects/-/serviceAccounts/{}/allowedLocations".format(


Can you remind me what prompted the mTLS addition?

We should have test(s) that set that signal and makes sure the right endpoint is called on refresh etc.

nbayati added 8 commits May 7, 2026 17:12

feat: add async support for RAB to ServiceAccountCredentials and impl…

b33d818

…ement async refresh manager

Add unit tests for the async RAB implementation.

80529c6

fix async unit tests

73617c5

Update unit tests to accept both mtls and standard allowedLocations e…

9e43ecc

…ndpoint.

test: verify iam endpoint constant resolution in mTLS environments

4d5787b

refactor: introduce _after_refresh hook in Credentials base class to …

e3f8e90

…support Regional Access Boundary logic

add __setstate__ to the base RAB class for backward compatibility

95af8e5

Implement RAB support for jwt credentials

1b4270b

gemini-code-assist Bot reviewed May 11, 2026

View reviewed changes

Comment thread packages/google-auth/google/auth/jwt.py

Comment thread packages/google-auth/tests/test_external_account_authorized_user.py Outdated

nbayati changed the title ~~Add RAB support for async SA and jwt.py~~ feat(auth): implement regional access boundary support for standalone JWT and async service accounts May 11, 2026

fix lint errors

0478330

nbayati marked this pull request as ready for review May 11, 2026 20:40

nbayati requested review from a team as code owners May 11, 2026 20:40

fix: preserve refresh manager type when copying RAB manager

9650ac7

macastelaz self-assigned this May 12, 2026

macastelaz approved these changes May 12, 2026

View reviewed changes

Comment thread packages/google-auth/google/auth/_regional_access_boundary_utils.py Outdated

Comment thread packages/google-auth/google/auth/_regional_access_boundary_utils.py Outdated

lsirac requested changes May 12, 2026

View reviewed changes

nbayati and others added 9 commits May 14, 2026 10:03

refactor(auth): optimize RAB manager copy logic to only share boundar…

5b6e715

…y data and config

fix(auth): enhance client lookup robustness with defensive checks and…

2729e70

… timeout fixes

refactor(auth): centralize async RAB lifecycle via _after_refresh bas…

e1ee432

…e hook

feat: add pickling support for _AsyncRegionalAccessBoundaryRefreshMan…

39fbe10

…ager in service account credentials

revert changes to the _token_endpoint_request_no_throw to keep PR foc…

d20d7ba

…used on RAB

fix(auth): align async client with AIO transport spec and add unit tests

c1b946e

test(auth): assert closed session safety in async RAB refresh and fix…

75b3231

… pickling docs

docs(auth): clarify async RAB transport requirements in docstrings

a952527

feat(auth): support async blocking RAB lookups and add support to asy…

be67ab6

…nc jwt

lsirac requested changes May 16, 2026

View reviewed changes

Conversation

nbayati commented May 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

nbayati commented May 11, 2026 •

edited

Loading