Integrate authentication record in SharedTokenCacheCredential

sdk/identity/azure-identity/src/main/java/com/azure/identity/DeviceCodeCredentialBuilder.java

@@ -56,9 +56,9 @@ public DeviceCodeCredentialBuilder enablePersistentCache() {
     /**
      * Sets the {@link AuthenticationRecord} captured from a previous authentication.
      *
-     * @param authenticationRecord the authentication record to ser.
+     * @param authenticationRecord the authentication record to be configured.
      *
-     * @return An updated instance of this builder with if the shared token cache enabled specified.
+     * @return An updated instance of this builder with the configured authentication record.
      */
     public DeviceCodeCredentialBuilder authenticationRecord(AuthenticationRecord authenticationRecord) {
         this.identityClientOptions.setAuthenticationRecord(authenticationRecord);

sdk/identity/azure-identity/src/main/java/com/azure/identity/SharedTokenCacheCredential.java

@@ -11,6 +11,7 @@
 import com.azure.identity.implementation.IdentityClient;
 import com.azure.identity.implementation.IdentityClientBuilder;
 import com.azure.identity.implementation.IdentityClientOptions;
+import com.azure.identity.implementation.MsalAuthenticationAccount;
 import com.azure.identity.implementation.MsalToken;
 import com.azure.identity.implementation.util.LoggingUtil;
 import reactor.core.publisher.Mono;
@@ -26,7 +27,8 @@ public class SharedTokenCacheCredential implements TokenCredential {
     private final String username;
     private final String clientId;
     private final String tenantId;
-    private final AtomicReference<MsalToken> cachedToken;
+    private final AtomicReference<MsalAuthenticationAccount> cachedToken;
+
 
     private final IdentityClient identityClient;
     private final ClientLogger logger = new ClientLogger(SharedTokenCacheCredential.class);
@@ -65,6 +67,9 @@ public class SharedTokenCacheCredential implements TokenCredential {
                 .identityClientOptions(identityClientOptions)
                 .build();
         this.cachedToken = new AtomicReference<>();
+        if (identityClientOptions.getAuthenticationRecord() != null) {
+            cachedToken.set(new MsalAuthenticationAccount(identityClientOptions.getAuthenticationRecord()));
+        }
         LoggingUtil.logAvailableEnvironmentVariables(logger, configuration);
     }
 
@@ -75,18 +80,23 @@ public class SharedTokenCacheCredential implements TokenCredential {
     public Mono<AccessToken> getToken(TokenRequestContext request) {
         return Mono.defer(() -> {
             if (cachedToken.get() != null) {
-                return identityClient.authenticateWithPublicClientCache(request, cachedToken.get().getAccount())
+                return identityClient.authenticateWithPublicClientCache(request, cachedToken.get())
                     .onErrorResume(t -> Mono.empty());
             } else {
                 return Mono.empty();
             }
         }).switchIfEmpty(
             Mono.defer(() -> identityClient.authenticateWithSharedTokenCache(request, username)))
-            .map(msalToken -> {
-                cachedToken.set(msalToken);
-                return (AccessToken) msalToken;
-            })
+            .map(this::updateCache)
             .doOnNext(token -> LoggingUtil.logTokenSuccess(logger, request))
             .doOnError(error -> LoggingUtil.logTokenError(logger, request, error));
     }
+
+    private AccessToken updateCache(MsalToken msalToken) {
+        cachedToken.set(
+            new MsalAuthenticationAccount(
+                new AuthenticationRecord(msalToken.getAuthenticationResult(),
+                    identityClient.getTenantId(), identityClient.getClientId())));
+        return msalToken;
+    }
 }

sdk/identity/azure-identity/src/main/java/com/azure/identity/SharedTokenCacheCredentialBuilder.java

@@ -34,6 +34,18 @@ public SharedTokenCacheCredentialBuilder allowUnencryptedCache() {
         return this;
     }
 
+    /**
+     * Sets the {@link AuthenticationRecord} captured from a previous authentication.
+     *
+     * @param authenticationRecord the authentication record to be configured.
+     *
+     * @return An updated instance of this builder with the configured authentication record.
+     */
+    public SharedTokenCacheCredentialBuilder authenticationRecord(AuthenticationRecord authenticationRecord) {
+        this.identityClientOptions.setAuthenticationRecord(authenticationRecord);
+        return this;
+    }
+
     /**
      * Creates a new {@link SharedTokenCacheCredentialBuilder} with the current configurations.
      *