Make Access token cache public (#49712)

g2vinay · web-flow · commit ade4b70f02fa · 2026-07-20T16:09:31.000-07:00
diff --git a/eng/versioning/version_client.txt b/eng/versioning/version_client.txt
@@ -570,6 +570,7 @@ io.clientcore:optional-dependency-tests;1.0.0-beta.1;1.0.0-beta.1
 # In the pom, the version update tag after the version should name the unreleased package and the dependency version:
 # <!-- {x-version-update;unreleased_com.azure:azure-core;dependency} -->
 
+unreleased_com.azure:azure-core;1.59.0-beta.1
 unreleased_com.azure.v2:azure-core;2.0.0-beta.1
 unreleased_com.azure.v2:azure-identity;2.0.0-beta.1
 unreleased_com.azure.v2:azure-data-appconfiguration;2.0.0-beta.1
diff --git a/sdk/core/azure-core/CHANGELOG.md b/sdk/core/azure-core/CHANGELOG.md
@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- Promoted `AccessTokenCache` to a public API in the `com.azure.core.credential` package. This class provides a thread-safe, proactively refreshing token cache that wraps a `TokenCredential`, supporting both synchronous and asynchronous token retrieval.
+
 ### Breaking Changes
 
 ### Bugs Fixed
diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/credential/AccessTokenCache.java b/sdk/core/azure-core/src/main/java/com/azure/core/credential/AccessTokenCache.java
@@ -1,11 +1,9 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-package com.azure.core.implementation;
+package com.azure.core.credential;
 
-import com.azure.core.credential.AccessToken;
-import com.azure.core.credential.TokenCredential;
-import com.azure.core.credential.TokenRequestContext;
+import com.azure.core.implementation.AccessTokenCacheInfo;
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.core.util.logging.LogLevel;
 import com.azure.core.util.logging.LoggingEventBuilder;
@@ -25,7 +23,57 @@
 import java.util.function.Supplier;
 
 /**
- * A token cache that supports caching a token and refreshing it.
+ * <p>
+ * {@code AccessTokenCache} is a thread-safe token cache that wraps a {@link TokenCredential} and manages proactive
+ * token refresh. It supports both asynchronous and synchronous token retrieval via
+ * {@link #getToken(TokenRequestContext, boolean)} and {@link #getTokenSync(TokenRequestContext, boolean)}.
+ * </p>
+ *
+ * <p>
+ * The cache maintains a single cached {@link AccessToken} per instance and proactively refreshes it before expiry
+ * (by default 5 minutes before the expiry time, or at the {@code refreshAt} time if provided by the credential).
+ * If a refresh fails while a non-expired token is still available, the cached token continues to be returned until
+ * it expires.
+ * </p>
+ *
+ * <p>
+ * When the {@code refreshOnContextChange} flag is {@code true}, the cache compares the incoming
+ * {@link TokenRequestContext} (scopes, tenant ID, claims) against the context used to acquire the current cached
+ * token. A mismatch causes an immediate token refresh regardless of expiry. This is the mechanism used to support
+ * Continuous Access Evaluation (CAE) claims challenges.
+ * </p>
+ *
+ * <p>
+ * <strong>Note:</strong> Each instance caches exactly one {@link AccessToken} associated with one
+ * {@link TokenRequestContext}. Do not share a single {@code AccessTokenCache} instance across calls that require
+ * different scopes or tenants simultaneously, as each new context will evict the previously cached token.
+ * </p>
+ *
+ * <p>
+ * This class is thread-safe. Multiple threads may call {@link #getToken(TokenRequestContext, boolean)} and
+ * {@link #getTokenSync(TokenRequestContext, boolean)} concurrently.
+ * </p>
+ *
+ * <p>
+ * <strong>Sample: Wrapping a TokenCredential with AccessTokenCache</strong>
+ * </p>
+ *
+ * <!-- src_embed com.azure.core.credential.accessTokenCache -->
+ * <pre>
+ * TokenCredential credential = new BasicAuthenticationCredential&#40;&quot;username&quot;, &quot;password&quot;&#41;;
+ * AccessTokenCache tokenCache = new AccessTokenCache&#40;credential&#41;;
+ * TokenRequestContext requestContext = new TokenRequestContext&#40;&#41;.addScopes&#40;&quot;https:&#47;&#47;management.azure.com&#47;.default&quot;&#41;;
+ * &#47;&#47; Async usage
+ * Mono&lt;AccessToken&gt; tokenMono = tokenCache.getToken&#40;requestContext, false&#41;;
+ * &#47;&#47; Sync usage
+ * AccessToken token = tokenCache.getTokenSync&#40;requestContext, false&#41;;
+ * </pre>
+ * <!-- end com.azure.core.credential.accessTokenCache -->
+ *
+ * @see TokenCredential
+ * @see AccessToken
+ * @see TokenRequestContext
+ * @see SimpleTokenCache
  */
 public final class AccessTokenCache {
     // The delay after a refresh to attempt another token refresh
@@ -48,12 +96,12 @@ public final class AccessTokenCache {
     private final Lock lock;
 
     /**
-     * Creates an instance of RefreshableTokenCredential with default scheme "Bearer".
+     * Creates an instance of {@code AccessTokenCache} that wraps the given {@link TokenCredential}.
      *
      * @param tokenCredential the token credential to be used to acquire the token.
      */
     public AccessTokenCache(TokenCredential tokenCredential) {
-        Objects.requireNonNull(tokenCredential, "The token credential cannot be null");
+        Objects.requireNonNull(tokenCredential, "'tokenCredential' cannot be null.");
         this.wip = new AtomicReference<>();
         this.tokenCredential = tokenCredential;
         this.cacheInfo = new AtomicReference<>(new AccessTokenCacheInfo(null, OffsetDateTime.now()));
@@ -70,11 +118,15 @@ public AccessTokenCache(TokenCredential tokenCredential) {
      * Asynchronously get a token from either the cache or replenish the cache with a new token.
      *
      * @param tokenRequestContext The request context for token acquisition.
-     * @param checkToForceFetchToken The flag indicating whether to force fetch a new token or not.
-     * @return The Publisher that emits an AccessToken
+     * @param refreshOnContextChange When {@code true}, compares the incoming {@link TokenRequestContext} against the
+     *     one used to acquire the current cached token. If the scopes, tenant ID, or claims differ, a fresh token is
+     *     fetched immediately regardless of expiry. Pass {@code false} to always use the cached token when it is
+     *     still valid.
+     * @return a {@link Mono} that emits the cached or newly acquired {@link AccessToken}.
+     * @throws IllegalArgumentException if {@code tokenRequestContext} is {@code null}.
      */
-    public Mono<AccessToken> getToken(TokenRequestContext tokenRequestContext, boolean checkToForceFetchToken) {
-        return Mono.defer(retrieveToken(tokenRequestContext, checkToForceFetchToken))
+    public Mono<AccessToken> getToken(TokenRequestContext tokenRequestContext, boolean refreshOnContextChange) {
+        return Mono.defer(retrieveToken(tokenRequestContext, refreshOnContextChange))
             // Keep resubscribing as long as Mono.defer [token acquisition] emits empty().
             .repeatWhenEmpty((Flux<Long> longFlux) -> longFlux
                 .concatMap(ignored -> Flux.just(true).delayElements(Duration.ofMillis(500))));
@@ -84,25 +136,29 @@ public Mono<AccessToken> getToken(TokenRequestContext tokenRequestContext, boole
      * Synchronously get a token from either the cache or replenish the cache with a new token.
      *
      * @param tokenRequestContext The request context for token acquisition.
-     * @param checkToForceFetchToken The flag indicating whether to force fetch a new token or not.
-     * @return The Publisher that emits an AccessToken
+     * @param refreshOnContextChange When {@code true}, compares the incoming {@link TokenRequestContext} against the
+     *     one used to acquire the current cached token. If the scopes, tenant ID, or claims differ, a fresh token is
+     *     fetched immediately regardless of expiry. Pass {@code false} to always use the cached token when it is
+     *     still valid.
+     * @return the cached or newly acquired {@link AccessToken}.
+     * @throws IllegalArgumentException if {@code tokenRequestContext} is {@code null}.
      */
-    public AccessToken getTokenSync(TokenRequestContext tokenRequestContext, boolean checkToForceFetchToken) {
+    public AccessToken getTokenSync(TokenRequestContext tokenRequestContext, boolean refreshOnContextChange) {
         lock.lock();
         try {
-            return retrieveTokenSync(tokenRequestContext, checkToForceFetchToken).get();
+            return retrieveTokenSync(tokenRequestContext, refreshOnContextChange).get();
         } finally {
             lock.unlock();
         }
     }
 
     private Supplier<Mono<? extends AccessToken>> retrieveToken(TokenRequestContext tokenRequestContext,
-        boolean checkToForceFetchToken) {
+        boolean refreshOnContextChange) {
         return () -> {
             try {
                 if (tokenRequestContext == null) {
-                    return Mono.error(LOGGER.logExceptionAsError(
-                        new IllegalArgumentException("The token request context input cannot be null.")));
+                    return Mono.error(LOGGER
+                        .logExceptionAsError(new IllegalArgumentException("'tokenRequestContext' cannot be null.")));
                 }
 
                 AccessTokenCacheInfo cache = this.cacheInfo.get();
@@ -115,9 +171,9 @@ private Supplier<Mono<? extends AccessToken>> retrieveToken(TokenRequestContext
                     Mono<AccessToken> fallback;
 
                     // Check if the incoming token request context is different from the cached one. A different
-                    // token request context, requires to fetch a new token as the cached one won't work for the
+                    // token request context requires fetching a new token as the cached one won't work for the
                     // passed in token request context.
-                    boolean forceRefresh = (checkToForceFetchToken && checkIfForceRefreshRequired(tokenRequestContext))
+                    boolean forceRefresh = (refreshOnContextChange && checkIfForceRefreshRequired(tokenRequestContext))
                         || this.tokenRequestContext == null;
 
                     if (forceRefresh) {
@@ -153,12 +209,12 @@ private Supplier<Mono<? extends AccessToken>> retrieveToken(TokenRequestContext
                             .flatMap(processTokenRefreshResult(sinksOne, now, fallback))
                             .doOnError(sinksOne::tryEmitError),
                         w -> w.set(null));
-                } else if (cachedToken != null && !cachedToken.isExpired() && !checkToForceFetchToken) {
+                } else if (cachedToken != null && !cachedToken.isExpired() && !refreshOnContextChange) {
                     // another thread might be refreshing the token proactively, but the current token is still valid
                     return Mono.just(cachedToken);
                 } else {
-                    // if a force refresh is possible, then exit and retry.
-                    if (checkToForceFetchToken) {
+                    // if a context-change refresh is pending, exit and retry.
+                    if (refreshOnContextChange) {
                         return Mono.empty();
                     }
                     // another thread is definitely refreshing the expired token
@@ -178,11 +234,10 @@ private Supplier<Mono<? extends AccessToken>> retrieveToken(TokenRequestContext
     }
 
     private Supplier<AccessToken> retrieveTokenSync(TokenRequestContext tokenRequestContext,
-        boolean checkToForceFetchToken) {
+        boolean refreshOnContextChange) {
         return () -> {
             if (tokenRequestContext == null) {
-                throw LOGGER.logExceptionAsError(
-                    new IllegalArgumentException("The token request context input cannot be null."));
+                throw LOGGER.logExceptionAsError(new IllegalArgumentException("'tokenRequestContext' cannot be null."));
             }
             AccessTokenCacheInfo cache = this.cacheInfo.get();
             AccessToken cachedToken = cache.getCachedAccessToken();
@@ -192,9 +247,9 @@ private Supplier<AccessToken> retrieveTokenSync(TokenRequestContext tokenRequest
             AccessToken fallback;
 
             // Check if the incoming token request context is different from the cached one. A different
-            // token request context, requires to fetch a new token as the cached one won't work for the
+            // token request context requires fetching a new token as the cached one won't work for the
             // passed in token request context.
-            boolean forceRefresh = (checkToForceFetchToken && checkIfForceRefreshRequired(tokenRequestContext))
+            boolean forceRefresh = (refreshOnContextChange && checkIfForceRefreshRequired(tokenRequestContext))
                 || this.tokenRequestContext == null;
 
             if (forceRefresh) {
@@ -245,7 +300,10 @@ private Supplier<AccessToken> retrieveTokenSync(TokenRequestContext tokenRequest
                 if (fallback != null) {
                     return fallback;
                 }
-                throw error;
+                if (error instanceof RuntimeException) {
+                    throw LOGGER.logExceptionAsError((RuntimeException) error);
+                }
+                throw LOGGER.logExceptionAsError(new RuntimeException(error));
             }
         };
     }
diff --git a/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/BearerTokenAuthenticationPolicy.java b/sdk/core/azure-core/src/main/java/com/azure/core/http/policy/BearerTokenAuthenticationPolicy.java
@@ -12,7 +12,7 @@
 import com.azure.core.http.HttpPipelineNextPolicy;
 import com.azure.core.http.HttpPipelineNextSyncPolicy;
 import com.azure.core.http.HttpResponse;
-import com.azure.core.implementation.AccessTokenCache;
+import com.azure.core.credential.AccessTokenCache;
 import com.azure.core.implementation.http.policy.AuthorizationChallengeParser;
 import com.azure.core.util.CoreUtils;
 import com.azure.core.util.logging.ClientLogger;
@@ -208,16 +208,16 @@ public void setAuthorizationHeaderSync(HttpPipelineCallContext context, TokenReq
     }
 
     private Mono<Void> setAuthorizationHeaderHelper(HttpPipelineCallContext context,
-        TokenRequestContext tokenRequestContext, boolean checkToForceFetchToken) {
-        return cache.getToken(tokenRequestContext, checkToForceFetchToken).flatMap(token -> {
+        TokenRequestContext tokenRequestContext, boolean refreshOnContextChange) {
+        return cache.getToken(tokenRequestContext, refreshOnContextChange).flatMap(token -> {
             setAuthorizationHeader(context.getHttpRequest().getHeaders(), token.getToken());
             return Mono.empty();
         });
     }
 
     private void setAuthorizationHeaderHelperSync(HttpPipelineCallContext context,
-        TokenRequestContext tokenRequestContext, boolean checkToForceFetchToken) {
-        AccessToken token = cache.getTokenSync(tokenRequestContext, checkToForceFetchToken);
+        TokenRequestContext tokenRequestContext, boolean refreshOnContextChange) {
+        AccessToken token = cache.getTokenSync(tokenRequestContext, refreshOnContextChange);
         setAuthorizationHeader(context.getHttpRequest().getHeaders(), token.getToken());
     }
 
diff --git a/sdk/core/azure-core/src/samples/java/com/azure/core/credential/AccessTokenCacheJavadocCodeSnippets.java b/sdk/core/azure-core/src/samples/java/com/azure/core/credential/AccessTokenCacheJavadocCodeSnippets.java
@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.core.credential;
+
+import reactor.core.publisher.Mono;
+
+/**
+ * Codesnippets for {@link AccessTokenCache}.
+ */
+public class AccessTokenCacheJavadocCodeSnippets {
+
+    public void accessTokenCacheSnippet() {
+        // BEGIN: com.azure.core.credential.accessTokenCache
+        TokenCredential credential = new BasicAuthenticationCredential("username", "password");
+        AccessTokenCache tokenCache = new AccessTokenCache(credential);
+        TokenRequestContext requestContext = new TokenRequestContext().addScopes("https://management.azure.com/.default");
+        // Async usage
+        Mono<AccessToken> tokenMono = tokenCache.getToken(requestContext, false);
+        // Sync usage
+        AccessToken token = tokenCache.getTokenSync(requestContext, false);
+        // END: com.azure.core.credential.accessTokenCache
+    }
+
+}
diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/credential/AccessTokenCacheTests.java b/sdk/core/azure-core/src/test/java/com/azure/core/credential/AccessTokenCacheTests.java
@@ -1,11 +1,8 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-package com.azure.core.implementation;
+package com.azure.core.credential;
 
-import com.azure.core.credential.AccessToken;
-import com.azure.core.credential.TokenCredential;
-import com.azure.core.credential.TokenRequestContext;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
@@ -83,15 +80,15 @@ public void testTenantIdChangeTriggersForceRefresh() {
 
         assertEquals(1, refreshCount.get());
 
-        // Second request with same tenant should not trigger refresh (checkToForceFetchToken=false)
+        // Second request with same tenant should not trigger refresh (refreshOnContextChange=false)
         StepVerifier.create(cache.getToken(context1, false)).assertNext(token -> {
             assertNotNull(token);
             assertEquals("token-1", token.getToken()); // Same token
         }).verifyComplete();
 
         assertEquals(1, refreshCount.get()); // No additional refresh
 
-        // Third request with different tenant should trigger force refresh with checkToForceFetchToken=true
+        // Third request with different tenant should trigger force refresh with refreshOnContextChange=true
         TokenRequestContext context2 = new TokenRequestContext().addScopes(SCOPE).setTenantId(TENANT_ID_2);
 
         StepVerifier.create(cache.getToken(context2, true)).assertNext(token -> {
diff --git a/sdk/core/azure-core/src/test/java/com/azure/core/credential/TokenCacheTests.java b/sdk/core/azure-core/src/test/java/com/azure/core/credential/TokenCacheTests.java
@@ -3,7 +3,6 @@
 
 package com.azure.core.credential;
 
-import com.azure.core.implementation.AccessTokenCache;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
diff --git a/sdk/identity/azure-identity-broker/src/test/java/com/azure/identity/broker/shr/resources/PopTokenAuthenticationPolicy.java b/sdk/identity/azure-identity-broker/src/test/java/com/azure/identity/broker/shr/resources/PopTokenAuthenticationPolicy.java
@@ -14,7 +14,7 @@
 import com.azure.core.http.HttpPipelineNextSyncPolicy;
 import com.azure.core.http.HttpResponse;
 import com.azure.core.http.policy.HttpPipelinePolicy;
-import com.azure.core.implementation.AccessTokenCache;
+import com.azure.core.credential.AccessTokenCache;
 import com.azure.core.util.CoreUtils;
 import com.azure.core.util.logging.ClientLogger;
 import reactor.core.publisher.Mono;
@@ -169,7 +169,7 @@ public HttpResponse processSync(HttpPipelineCallContext context, HttpPipelineNex
         }
     }
 
-    private Mono<Void> setAuthorizationHeaderHelper(HttpPipelineCallContext context, boolean checkToForceFetchToken) {
+    private Mono<Void> setAuthorizationHeaderHelper(HttpPipelineCallContext context, boolean refreshOnContextChange) {
         if (!"https".equals(context.getHttpRequest().getUrl().getProtocol())) {
             throw LOGGER.logExceptionAsError(new RuntimeException(
                 "Proof of possession token authentication is not permitted for non TLS-protected (HTTPS) endpoints."));
@@ -181,15 +181,15 @@ private Mono<Void> setAuthorizationHeaderHelper(HttpPipelineCallContext context,
                 .setRequestUrl(context.getHttpRequest().getUrl()));
 
         if (!CoreUtils.isNullOrEmpty(popNonce)) {
-            return this.cache.getToken(popTokenRequestContext, checkToForceFetchToken).flatMap((token) -> {
+            return this.cache.getToken(popTokenRequestContext, refreshOnContextChange).flatMap((token) -> {
                 setAuthorizationHeader(context.getHttpRequest().getHeaders(), token.getToken());
                 return Mono.empty();
             });
         }
         return Mono.empty();
     }
 
-    private void setAuthorizationHeaderHelperSync(HttpPipelineCallContext context, boolean checkToForceFetchToken) {
+    private void setAuthorizationHeaderHelperSync(HttpPipelineCallContext context, boolean refreshOnContextChange) {
         if (!"https".equals(context.getHttpRequest().getUrl().getProtocol())) {
             throw LOGGER.logExceptionAsError(new RuntimeException(
                 "Proof of possession token authentication is not permitted for non TLS-protected (HTTPS) endpoints."));
@@ -199,7 +199,7 @@ private void setAuthorizationHeaderHelperSync(HttpPipelineCallContext context, b
                 .setRequestMethod(context.getHttpRequest().getHttpMethod())
                 .setRequestUrl(context.getHttpRequest().getUrl()));
 
-        AccessToken token = this.cache.getTokenSync(popTokenRequestContext, checkToForceFetchToken);
+        AccessToken token = this.cache.getTokenSync(popTokenRequestContext, refreshOnContextChange);
         setAuthorizationHeader(context.getHttpRequest().getHeaders(), token.getToken());
     }
 
diff --git a/sdk/identity/azure-identity-perf/pom.xml b/sdk/identity/azure-identity-perf/pom.xml
diff --git a/sdk/identity/azure-identity/pom.xml b/sdk/identity/azure-identity/pom.xml
diff --git a/sdk/identity/ci.yml b/sdk/identity/ci.yml
