[Identity] Add background token refresh to async creds

[pvaneck]

In many cases, a token's refresh_on time might have passed, but there is still ample time before expiration. In these cases, we should do a background refresh to prevent blocking on the refresh attempt. The pre-existing token can be used in the meantime.

[Copilot]

Pull request overview

This pull request adds background token refresh functionality to async credentials in the Azure Identity SDK. When a token's refresh_on time has passed but there's still ample time before expiration (at least 10 minutes), credentials will now refresh the token in the background instead of blocking the caller. The pre-existing token is returned immediately, providing better performance and responsiveness.

Changes:

• Added background token refresh logic to GetTokenMixin that creates async tasks for non-blocking token refresh when tokens are in the refresh window but still have sufficient validity
• Implemented task management and cleanup mechanisms to prevent resource leaks and ensure proper shutdown
• Updated all async credentials using GetTokenMixin to cancel background refresh tasks in their close() methods
• Added comprehensive test coverage for background refresh scenarios including edge cases, failures, cleanup, and multi-threading

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
sdk/identity/azure-identity/azure/identity/aio/_internal/get_token_mixin.py	Core implementation of background refresh logic including task creation, management, cleanup, and serialization support
sdk/identity/azure-identity/azure/identity/aio/_internal/managed_identity_base.py	Added background task cleanup call in close() method
sdk/identity/azure-identity/azure/identity/aio/_credentials/on_behalf_of.py	Added background task cleanup call in close() method
sdk/identity/azure-identity/azure/identity/aio/_credentials/imds.py	Added background task cleanup call in close() method
sdk/identity/azure-identity/azure/identity/aio/_credentials/client_secret.py	Added background task cleanup call in close() method
sdk/identity/azure-identity/azure/identity/aio/_credentials/client_assertion.py	Added background task cleanup call in close() method
sdk/identity/azure-identity/azure/identity/aio/_credentials/certificate.py	Added background task cleanup call in close() method
sdk/identity/azure-identity/azure/identity/aio/_credentials/authorization_code.py	Added background task cleanup call in close() method
sdk/identity/azure-identity/tests/test_get_token_mixin_async.py	Added comprehensive test coverage including background refresh behavior, deduplication, non-asyncio fallback, failure handling, cleanup, serialization, and multi-threading scenarios

Comments suppressed due to low confidence (5)

sdk/identity/azure-identity/azure/identity/aio/_internal/get_token_mixin.py:16

• The _BACKGROUND_REFRESH_MIN_VALIDITY_SECONDS constant (600 seconds = 10 minutes) should be documented to explain why this specific threshold was chosen. This helps future maintainers understand the trade-off between background refresh benefits and the risk of returning an expired token if the refresh fails and takes too long.

_BACKGROUND_REFRESH_MIN_VALIDITY_SECONDS = 600

sdk/identity/azure-identity/azure/identity/aio/_internal/get_token_mixin.py:97

• The __getstate__ method is implemented to ensure the credential can be pickled, but there's no corresponding __setstate__ method. When unpickling, the _background_refresh_tasks dictionary will be properly restored to an empty dict, but if there's any additional initialization logic needed for unpickled credentials, a __setstate__ method should be added for completeness and clarity.

    def __getstate__(self) -> Dict[str, Any]:
        state = self.__dict__.copy()
        state["_background_refresh_tasks"] = {}
        return state

sdk/identity/azure-identity/tests/test_get_token_mixin_async.py:7

• The pickle import is used in line 7 but the actual pickling/unpickling is never tested in test_pickle_with_active_background_refresh. The test only verifies getstate returns the correct state. Consider either removing the unused import or adding a full pickle/unpickle cycle test to ensure the credential can be properly serialized and deserialized.

import pickle

sdk/identity/azure-identity/tests/test_get_token_mixin_async.py:162

• Consider adding a test case that explicitly validates the boundary condition between background and inline refresh. Specifically, test a scenario where the token has exactly _BACKGROUND_REFRESH_MIN_VALIDITY_SECONDS + 1 seconds remaining to ensure background refresh is used, and another with _BACKGROUND_REFRESH_MIN_VALIDITY_SECONDS - 1 seconds to ensure inline refresh is used. This would make the threshold behavior more explicit and prevent regressions if the constant value changes.

@pytest.mark.parametrize("get_token_method", GET_TOKEN_METHODS)
async def test_background_refresh_does_not_block(get_token_method):
    """Background refresh should return the cached token immediately while refreshing in the background"""

    now = int(time.time())
    # Token has plenty of time before expiry but refresh_on has passed, triggering background refresh
    credential = MockCredential(cached_token=AccessTokenInfo(CACHED_TOKEN, now + 3600, refresh_on=now - 1))

    token = await getattr(credential, get_token_method)(SCOPE)

    # The cached token is returned immediately, before the background task completes
    assert token.token == CACHED_TOKEN
    # Wait for background refresh task to complete
    key = ((SCOPE,), None, None, False)
    assert key in credential._background_refresh_tasks
    await asyncio.gather(credential._background_refresh_tasks[key], return_exceptions=True)
    credential.request_token.assert_called_once()

sdk/identity/azure-identity/azure/identity/aio/_internal/get_token_mixin.py:24

• The _background_refresh_tasks dictionary is not protected by a lock and could experience race conditions if the same credential instance is accessed from multiple threads with different event loops (as shown in test_background_refresh_multithread_event_loops). While credential instances are typically not shared across threads in practice, consider documenting this limitation or adding thread-safety if cross-thread usage is intended to be supported.

        self._background_refresh_tasks: Dict[_BackgroundRefreshKey, "asyncio.Task[None]"] = {}