Fix for #5375 - Add a FormatResultAsync to IAuthenticationOperation

src/client/Microsoft.Identity.Client/AuthScheme/Bearer/BearerAuthenticationOperation.cs

@@ -2,13 +2,15 @@
 // Licensed under the MIT License.
 
 using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.Identity.Client.Cache.Items;
 using Microsoft.Identity.Client.Internal;
 using Microsoft.Identity.Client.Utils;
 
 namespace Microsoft.Identity.Client.AuthScheme.Bearer
 {
-    internal class BearerAuthenticationOperation : IAuthenticationOperation
+    internal class BearerAuthenticationOperation : IAuthenticationOperation2
     {
         internal const string BearerTokenType = "bearer";
 
@@ -25,6 +27,12 @@ public void FormatResult(AuthenticationResult authenticationResult)
             // no-op
         }
 
+        public Task FormatResultAsync(AuthenticationResult authenticationResult, CancellationToken cancellationToken = default)
+        {
+            // no-op, return completed task
+            return Task.CompletedTask;
+        }
+
         public IReadOnlyDictionary<string, string> GetTokenRequestParams()
         {
             // ESTS issues Bearer tokens by default, no need for any extra params

src/client/Microsoft.Identity.Client/AuthScheme/IAuthenticationOperation.cs

@@ -2,7 +2,6 @@
 // Licensed under the MIT License.
 
 using System.Collections.Generic;
-using Microsoft.Identity.Client.Cache.Items;
 
 namespace Microsoft.Identity.Client.AuthScheme
 {
@@ -54,7 +53,7 @@ public interface IAuthenticationOperation
         /// Creates the access token that goes into an Authorization HTTP header. 
         /// </summary>
         void FormatResult(AuthenticationResult authenticationResult);
-
+        
         /// <summary>
         /// Expected to match the token_type parameter returned by ESTS. Used to disambiguate
         /// between ATs of different types (e.g. Bearer and PoP) when loading from cache etc.

src/client/Microsoft.Identity.Client/AuthScheme/IAuthenticationOperation2.cs

@@ -0,0 +1,21 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.Identity.Client.AuthScheme
+{
+    /// <summary>
+    /// This is an extensibility API and should only be used by SDKs.
+    /// Enhanced version of IAuthenticationOperation that supports asynchronous token formatting.
+    /// Used to modify the experience depending on the type of token asked with async capabilities.
+    /// </summary>
+    public interface IAuthenticationOperation2 : IAuthenticationOperation
+    {
+        /// <summary>
+        /// Will be invoked instead of IAuthenticationOperation.FormatResult
+        /// </summary>
+        Task FormatResultAsync(AuthenticationResult authenticationResult, CancellationToken cancellationToken = default);
+    }
+}

src/client/Microsoft.Identity.Client/AuthScheme/PoP/MtlsPopAuthenticationOperation.cs

@@ -4,13 +4,15 @@
 using System.Collections.Generic;
 using System.Security.Cryptography;
 using System.Security.Cryptography.X509Certificates;
+using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.Identity.Client.Internal;
 using Microsoft.Identity.Client.OAuth2;
 using Microsoft.Identity.Client.Utils;
 
 namespace Microsoft.Identity.Client.AuthScheme.PoP
 {
-    internal class MtlsPopAuthenticationOperation : IAuthenticationOperation
+    internal class MtlsPopAuthenticationOperation : IAuthenticationOperation2
     {
         private readonly X509Certificate2 _mtlsCert;
 
@@ -36,11 +38,6 @@ public IReadOnlyDictionary<string, string> GetTokenRequestParams()
             };
         }
 
-        public void FormatResult(AuthenticationResult authenticationResult)
-        {
-            //no-op
-        }
-
         private static string ComputeX5tS256KeyId(X509Certificate2 certificate)
         {
             // Extract the raw bytes of the certificate’s public key.
@@ -55,5 +52,15 @@ private static string ComputeX5tS256KeyId(X509Certificate2 certificate)
                 return Base64UrlHelpers.Encode(hash);
             }
         }
+
+        public Task FormatResultAsync(AuthenticationResult authenticationResult, CancellationToken cancellationToken = default)
+        {
+            return Task.CompletedTask;
+        }
+
+        public void FormatResult(AuthenticationResult authenticationResult)
+        {
+            // no-op
+        }
     }
 }

src/client/Microsoft.Identity.Client/AuthScheme/PoP/PopAuthenticationOperation.cs

@@ -6,6 +6,8 @@
 using System.Globalization;
 using System.Security.Cryptography;
 using System.Text;
+using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.Identity.Client.AppConfig;
 using Microsoft.Identity.Client.Cache.Items;
 using Microsoft.Identity.Client.Internal;
@@ -21,7 +23,7 @@
 
 namespace Microsoft.Identity.Client.AuthScheme.PoP
 {
-    internal class PopAuthenticationOperation : IAuthenticationOperation
+    internal class PopAuthenticationOperation : IAuthenticationOperation2
     {
         private readonly PoPAuthenticationConfiguration _popAuthenticationConfiguration;
         private readonly IPoPCryptoProvider _popCryptoProvider;
@@ -86,6 +88,14 @@ public void FormatResult(AuthenticationResult authenticationResult)
             authenticationResult.AccessToken = popToken;
         }
 
+        public Task FormatResultAsync(AuthenticationResult authenticationResult, CancellationToken cancellationToken = default)
+        {
+            // For now, PoP token creation is synchronous, so we wrap the sync method
+            // Future enhancement could make crypto operations truly async
+            FormatResult(authenticationResult);
+            return Task.CompletedTask;
+        }
+
         private JObject CreateBody(string accessToken)
         {
             var publicKeyJwk = JToken.Parse(_popCryptoProvider.CannonicalPublicKeyJwk);

src/client/Microsoft.Identity.Client/AuthScheme/PoP/PopBrokerAuthenticationOperation.cs

@@ -5,6 +5,7 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.Cache.Items;
 using Microsoft.Identity.Client.Internal;
@@ -15,7 +16,7 @@ namespace Microsoft.Identity.Client.AuthScheme.PoP
     //Authentication Scheme used when MSAL Broker and pop are used together.
     //Tokens acquired from brokers will not be saved in the local ache and MSAL will not search the local cache during silent authentication.
     //This is because tokens are cached in the broker instead so MSAL will rely on the broker's cache for silent requests.
-    internal class PopBrokerAuthenticationOperation : IAuthenticationOperation
+    internal class PopBrokerAuthenticationOperation : IAuthenticationOperation2
     {
         public int TelemetryTokenType => TelemetryTokenTypeConstants.Pop;
 
@@ -25,9 +26,14 @@ internal class PopBrokerAuthenticationOperation : IAuthenticationOperation
 
         public string AccessTokenType => Constants.PoPTokenType;
 
+        public Task FormatResultAsync(AuthenticationResult authenticationResult, CancellationToken cancellationToken = default)
+        {
+            return Task.CompletedTask;
+        }
+
         public void FormatResult(AuthenticationResult authenticationResult)
         {
-            //no-op
+            // no-op
         }
 
         public IReadOnlyDictionary<string, string> GetTokenRequestParams()

src/client/Microsoft.Identity.Client/AuthScheme/SSHCertificates/SSHCertAuthenticationOperation.cs

@@ -3,13 +3,15 @@
 
 using System;
 using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.Identity.Client.Cache.Items;
 using Microsoft.Identity.Client.Internal;
 using Microsoft.Identity.Client.OAuth2;
 
 namespace Microsoft.Identity.Client.AuthScheme.SSHCertificates
 {
-    internal class SSHCertAuthenticationOperation : IAuthenticationOperation
+    internal class SSHCertAuthenticationOperation : IAuthenticationOperation2
     {
         internal const string SSHCertTokenType = "ssh-cert";
         private readonly string _jwk;
@@ -40,6 +42,11 @@ public SSHCertAuthenticationOperation(string keyId, string jwk)
 
         public string KeyId { get; }
 
+        public Task FormatResultAsync(AuthenticationResult authenticationResult, CancellationToken cancellationToken = default)
+        {
+            return Task.CompletedTask;
+        }
+
         public void FormatResult(AuthenticationResult authenticationResult)
         {
             // no-op

src/client/Microsoft.Identity.Client/AuthenticationResult.cs

@@ -6,6 +6,7 @@
 using System.ComponentModel;
 using System.Globalization;
 using System.Security.Claims;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.AuthScheme;
 using Microsoft.Identity.Client.Cache;
@@ -21,7 +22,7 @@ namespace Microsoft.Identity.Client
     /// </summary>
     public partial class AuthenticationResult
     {
-        private readonly IAuthenticationOperation _authenticationScheme;
+        private IAuthenticationOperation _authenticationScheme;
 
         /// <summary>
         /// Constructor meant to help application developers test their apps. Allows mocking of authentication flows.
@@ -126,19 +127,76 @@ public AuthenticationResult(
 
         }
 
-        internal AuthenticationResult(
+        /// <summary>
+        /// This method must be used by the product code to create an <see cref="AuthenticationResult"/> instance.
+        /// It calls IAuthenticationOperation.FormatResult or FormatResultAsync on the authentication scheme
+        /// </summary>
+        internal static async Task<AuthenticationResult> CreateAsync(
             MsalAccessTokenCacheItem msalAccessTokenCacheItem,
             MsalIdTokenCacheItem msalIdTokenCacheItem,
             IAuthenticationOperation authenticationScheme,
-            Guid correlationID,
+            Guid correlationId,
             TokenSource tokenSource,
             ApiEvent apiEvent,
             Account account,
             string spaAuthCode,
-            IReadOnlyDictionary<string, string> additionalResponseParameters)
+            IReadOnlyDictionary<string, string> additionalResponseParameters,
+            CancellationToken cancellationToken = default)
         {
-            _authenticationScheme = authenticationScheme ?? throw new ArgumentNullException(nameof(authenticationScheme));
+            if (authenticationScheme == null)
+            {
+                throw new ArgumentNullException(nameof(authenticationScheme));
+            }
+
+            // Create the AuthenticationResult without calling FormatResult in constructor
+            var result = new AuthenticationResult(
+                msalAccessTokenCacheItem,
+                msalIdTokenCacheItem,
+                correlationId,
+                tokenSource,
+                apiEvent,
+                account,
+                spaAuthCode,
+                additionalResponseParameters,
+                authenticationScheme);
+
+            // Apply token formatting (async if supported, sync otherwise)
+            var measuredResultDuration = await StopwatchService.MeasureCodeBlockAsync(async () =>
+            {
+                if (authenticationScheme is IAuthenticationOperation2 asyncAuthScheme)
+                {
+                    await asyncAuthScheme.FormatResultAsync(result, cancellationToken).ConfigureAwait(false);
+                }
+                else
+                {
+                    authenticationScheme.FormatResult(result);
+                }
+            }).ConfigureAwait(false);
+
+            // Update telemetry metadata
+            result.AuthenticationResultMetadata.DurationCreatingExtendedTokenInUs = measuredResultDuration.Microseconds;
+            result.AuthenticationResultMetadata.TelemetryTokenType = authenticationScheme.TelemetryTokenType;
+
+            return result;
+        }
+
+        //Default constructor for testing
+        internal AuthenticationResult() { }
 
+        /// <summary>
+        /// This is to help CreateAsync
+        /// </summary>
+        private AuthenticationResult(
+            MsalAccessTokenCacheItem msalAccessTokenCacheItem,
+            MsalIdTokenCacheItem msalIdTokenCacheItem,
+            Guid correlationID,
+            TokenSource tokenSource,
+            ApiEvent apiEvent,
+            Account account,
+            string spaAuthCode,
+            IReadOnlyDictionary<string, string> additionalResponseParameters,
+            IAuthenticationOperation authenticationScheme)
+        {
             string homeAccountId =
                 msalAccessTokenCacheItem?.HomeAccountId ??
                 msalIdTokenCacheItem?.HomeAccountId;
@@ -163,7 +221,7 @@ internal AuthenticationResult(
             TenantId = msalIdTokenCacheItem?.IdToken?.TenantId;
             IdToken = msalIdTokenCacheItem?.Secret;
             SpaAuthCode = spaAuthCode;
-
+            _authenticationScheme = authenticationScheme;
             CorrelationId = correlationID;
             ApiEvent = apiEvent;
             AuthenticationResultMetadata = new AuthenticationResultMetadata(tokenSource);
@@ -189,20 +247,8 @@ internal AuthenticationResult(
 
                 AccessToken = msalAccessTokenCacheItem.Secret;
             }
-
-            var measuredResultDuration = StopwatchService.MeasureCodeBlock(() =>
-            {
-                //Important: only call this at the end
-                authenticationScheme.FormatResult(this);
-            });
-
-            AuthenticationResultMetadata.DurationCreatingExtendedTokenInUs = measuredResultDuration.Microseconds;
-            AuthenticationResultMetadata.TelemetryTokenType = authenticationScheme.TelemetryTokenType;
         }
 
-        //Default constructor for testing
-        internal AuthenticationResult() { }
-
         /// <summary>
         /// Access Token that can be used as a bearer token to access protected web APIs
         /// </summary>

src/client/Microsoft.Identity.Client/Extensibility/ExternalBoundTokenScheme.cs

@@ -2,14 +2,16 @@
 // Licensed under the MIT License.
 
 using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.Identity.Client.AuthScheme;
 using Microsoft.Identity.Client.Cache.Items;
 using Microsoft.Identity.Client.Internal;
 using Microsoft.Identity.Client.Utils;
 
 namespace Microsoft.Identity.Client.Extensibility
 {
-    internal class ExternalBoundTokenScheme : IAuthenticationOperation
+    internal class ExternalBoundTokenScheme : IAuthenticationOperation2
     {
         private readonly string _keyId;
         private readonly string _tokenType;
@@ -28,6 +30,11 @@ public ExternalBoundTokenScheme(string keyId, string expectedTokenTypeFromEsts =
 
         public string AccessTokenType => _tokenType;
 
+        public Task FormatResultAsync(AuthenticationResult authenticationResult, CancellationToken cancellationToken = default)
+        {
+            return Task.CompletedTask;
+        }
+
         public void FormatResult(AuthenticationResult authenticationResult)
         {
             // no-op

src/client/Microsoft.Identity.Client/Extensibility/MsalAuthenticationExtension.cs

@@ -20,7 +20,7 @@ public class MsalAuthenticationExtension
         public Func<OnBeforeTokenRequestData, Task> OnBeforeTokenRequestHandler { get; set; }
 
         /// <summary>
-        /// Enables the developer to provide a custom authentication extension.
+        /// Important: IAuthenticationOperation2 exists, for asynchronous token formatting.
         /// </summary>
         public IAuthenticationOperation AuthenticationOperation { get; set; }
 

src/client/Microsoft.Identity.Client/Internal/Requests/ByRefreshTokenRequest.cs

@@ -37,7 +37,7 @@ protected override async Task<AuthenticationResult> ExecuteAsync(CancellationTok
                 throw new MsalServiceException(msalTokenResponse.Error, msalTokenResponse.ErrorDescription, null);
             }
 
-            return await CacheTokenResponseAndCreateAuthenticationResultAsync(msalTokenResponse).ConfigureAwait(false);
+            return await CacheTokenResponseAndCreateAuthenticationResultAsync(msalTokenResponse, cancellationToken).ConfigureAwait(false);
         }
 
         private static Dictionary<string, string> GetBodyParameters(string refreshTokenSecret)