Reworked retry policy functionality & Created IMDS retry policy

src/client/Microsoft.Identity.Client/AppConfig/ApplicationConfiguration.cs

@@ -11,6 +11,7 @@
 using Microsoft.Identity.Client.Core;
 using Microsoft.Identity.Client.Extensibility;
 using Microsoft.Identity.Client.Http;
+using Microsoft.Identity.Client.Http.Retry;
 using Microsoft.Identity.Client.Instance;
 using Microsoft.Identity.Client.Instance.Discovery;
 using Microsoft.Identity.Client.Internal.Broker;
@@ -124,6 +125,8 @@ public string ClientVersion
 
         public Func<AppTokenProviderParameters, Task<AppTokenProviderResult>> AppTokenProvider;
 
+        internal IRetryPolicyFactory RetryPolicyFactory { get; set; }
+
         #region ClientCredentials
 
         // Indicates if claims or assertions are used within the configuration
@@ -207,6 +210,5 @@ public X509Certificate2 ClientCredentialCertificate
         public IDeviceAuthManager DeviceAuthManagerForTest { get; set; }        
         public bool IsInstanceDiscoveryEnabled { get; internal set; } = true;
         #endregion
-
     }
 }

src/client/Microsoft.Identity.Client/AppConfig/BaseAbstractApplicationBuilder.cs

@@ -14,6 +14,8 @@
 using Microsoft.Identity.Client.Utils;
 using Microsoft.IdentityModel.Abstractions;
 using Microsoft.Identity.Client.Internal;
+using Microsoft.Identity.Client.Http.Retry;
+
 #if SUPPORTS_SYSTEM_TEXT_JSON
 using System.Text.Json;
 #else
@@ -31,6 +33,12 @@ public abstract class BaseAbstractApplicationBuilder<T>
         internal BaseAbstractApplicationBuilder(ApplicationConfiguration configuration)
         {
             Config = configuration;
+
+            // Ensure the default retry policy factory is set if the test factory was not provided
+            if (Config.RetryPolicyFactory == null)
+            {
+                Config.RetryPolicyFactory = new RetryPolicyFactory();
+            }
         }
 
         internal ApplicationConfiguration Config { get; }
@@ -227,6 +235,17 @@ public T WithClientVersion(string clientVersion)
             return this as T;
         }
 
+        /// <summary>
+        /// Internal only: Allows tests to inject a custom retry policy factory.
+        /// </summary>
+        /// <param name="factory">The retry policy factory to use.</param>
+        /// <returns>The builder for chaining.</returns>
+        internal T WithRetryPolicyFactory(IRetryPolicyFactory factory)
+        {
+            Config.RetryPolicyFactory = factory;
+            return (T)this;
+        }
+
         internal virtual ApplicationConfiguration BuildConfiguration()
         {
             ResolveAuthority();

src/client/Microsoft.Identity.Client/Http/HttpManager.cs

@@ -13,6 +13,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.Core;
+using Microsoft.Identity.Client.Http.Retry;
 
 namespace Microsoft.Identity.Client.Http
 {
@@ -110,10 +111,11 @@ public async Task<HttpResponse> SendRequestAsync(
                 logger.Error("The HTTP request failed. " + exception.Message);
                 timeoutException = exception;
             }
-
-            while (!_disableInternalRetries && retryPolicy.PauseForRetry(response, timeoutException, retryCount))
+            
+            while (!_disableInternalRetries && await retryPolicy.PauseForRetryAsync(response, timeoutException, retryCount, logger).ConfigureAwait(false))
             {
-                logger.Warning($"Retry condition met. Retry count: {retryCount++} after waiting {retryPolicy.DelayInMilliseconds}ms.");
+                retryCount++;
+
                 return await SendRequestAsync(
                     endpoint,
                     headers,

src/client/Microsoft.Identity.Client/Http/IHttpManager.cs

@@ -9,6 +9,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.Core;
+using Microsoft.Identity.Client.Http.Retry;
 
 namespace Microsoft.Identity.Client.Http
 {

src/client/Microsoft.Identity.Client/Http/Retry/DefaultRetryPolicy.cs

@@ -0,0 +1,75 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Threading.Tasks;
+using Microsoft.Identity.Client.Core;
+
+namespace Microsoft.Identity.Client.Http.Retry
+{
+    class DefaultRetryPolicy : IRetryPolicy
+    {
+        // referenced in unit tests
+        public const int DefaultStsMaxRetries = 1;
+        public const int DefaultManagedIdentityMaxRetries = 3;
+
+        private const int DefaultStsRetryDelayMs = 1000;
+        private const int DefaultManagedIdentityRetryDelayMs = 1000;
+
+        public readonly int _defaultRetryDelayMs;
+        private readonly int _maxRetries;
+        private readonly Func<HttpResponse, Exception, bool> _retryCondition;
+        private readonly LinearRetryStrategy _linearRetryStrategy = new LinearRetryStrategy();
+
+        public DefaultRetryPolicy(RequestType requestType)
+        {
+            switch (requestType)
+            {
+                case RequestType.ManagedIdentityDefault:
+                    _defaultRetryDelayMs = DefaultManagedIdentityRetryDelayMs;
+                    _maxRetries = DefaultManagedIdentityMaxRetries;
+                    _retryCondition = HttpRetryConditions.DefaultManagedIdentity;
+                    break;
+                case RequestType.STS:
+                    _defaultRetryDelayMs = DefaultStsRetryDelayMs;
+                    _maxRetries = DefaultStsMaxRetries;
+                    _retryCondition = HttpRetryConditions.Sts;
+                    break;
+                default:
+                    throw new ArgumentOutOfRangeException(nameof(requestType), requestType, "Unknown request type");
+            }
+        }
+
+        internal virtual Task DelayAsync(int milliseconds)
+        {
+            return Task.Delay(milliseconds);
+        }
+
+        public async Task<bool> PauseForRetryAsync(HttpResponse response, Exception exception, int retryCount, ILoggerAdapter logger)
+        {
+            // Check if the status code is retriable and if the current retry count is less than max retries
+            if (_retryCondition(response, exception) &&
+                retryCount < _maxRetries)
+            {
+                // Use HeadersAsDictionary to check for "Retry-After" header
+                string retryAfter = string.Empty;
+                if (response?.HeadersAsDictionary != null)
+                {
+                    response.HeadersAsDictionary.TryGetValue("Retry-After", out retryAfter);
+                }
+
+                int retryAfterDelay = _linearRetryStrategy.CalculateDelay(retryAfter, _defaultRetryDelayMs);
+
+                logger.Warning($"Retrying request in {retryAfterDelay}ms (retry attempt: {retryCount + 1})");
+
+                // Pause execution for the calculated delay
+                await DelayAsync(retryAfterDelay).ConfigureAwait(false);
+
+                return true;
+            }
+
+            // If the status code is not retriable or max retries have been reached, do not retry
+            return false;
+        }
+    }
+}

src/client/Microsoft.Identity.Client/Http/Retry/ExponentialRetryStrategy.cs

@@ -0,0 +1,55 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+
+namespace Microsoft.Identity.Client.Http.Retry
+{
+    internal class ExponentialRetryStrategy
+    {
+        // Minimum backoff time in milliseconds
+        private int _minExponentialBackoff;
+        // Maximum backoff time in milliseconds
+        private int _maxExponentialBackoff;
+        // Maximum backoff time in milliseconds
+        private int _exponentialDeltaBackoff;
+
+        public ExponentialRetryStrategy(int minExponentialBackoff, int maxExponentialBackoff, int exponentialDeltaBackoff)
+        {
+            _minExponentialBackoff = minExponentialBackoff;
+            _maxExponentialBackoff = maxExponentialBackoff;
+            _exponentialDeltaBackoff = exponentialDeltaBackoff;
+        }
+
+        /// <summary>
+        /// Calculates the exponential delay based on the current retry attempt.
+        /// </summary>
+        /// <param name="currentRetry">The current retry attempt number.</param>
+        /// <returns>The calculated exponential delay in milliseconds.</returns>
+        /// <remarks>
+        /// The delay is calculated using the formula:
+        /// - If <paramref name="currentRetry"/> is 0, it returns the minimum backoff time.
+        /// - Otherwise, it calculates the delay as the minimum of:
+        ///   - (2^(currentRetry - 1)) * deltaBackoff
+        ///   - maxBackoff
+        /// This ensures that the delay increases exponentially with each retry attempt,
+        /// but does not exceed the maximum backoff time.
+        /// </remarks>
+        public int CalculateDelay(int currentRetry)
+        {
+            // Attempt 1
+            if (currentRetry == 0)
+            {
+                return _minExponentialBackoff;
+            }
+
+            // Attempt 2+
+            int exponentialDelay = Math.Min(
+                (int)(Math.Pow(2, currentRetry - 1) * _exponentialDeltaBackoff),
+                _maxExponentialBackoff
+            );
+
+            return exponentialDelay;
+        }
+    }
+}

src/client/Microsoft.Identity.Client/Http/HttpRetryCondition.cs → src/client/Microsoft.Identity.Client/Http/Retry/HttpRetryCondition.cs

@@ -4,15 +4,15 @@
 using System;
 using System.Threading.Tasks;
 
-namespace Microsoft.Identity.Client.Http
+namespace Microsoft.Identity.Client.Http.Retry
 {
     internal static class HttpRetryConditions
     {
         /// <summary>
         /// Retry policy specific to managed identity flow.
-        /// Avoid changing this, as it's breaking change.
+        /// Avoid changing this, as it's a breaking change.
         /// </summary>
-        public static bool ManagedIdentity(HttpResponse response, Exception exception)
+        public static bool DefaultManagedIdentity(HttpResponse response, Exception exception)
         {
             if (exception != null)
             {
@@ -21,12 +21,32 @@ public static bool ManagedIdentity(HttpResponse response, Exception exception)
 
             return (int)response.StatusCode switch
             {
-                //Not Found
+                // Not Found, Request Timeout, Too Many Requests, Server Error, Service Unavailable, Gateway Timeout
                 404 or 408 or 429 or 500 or 503 or 504 => true,
                 _ => false,
             };
         }
 
+        /// <summary>
+        /// Retry policy specific to IMDS Managed Identity.
+        /// </summary>
+        public static bool Imds(HttpResponse response, Exception exception)
+        {
+            if (exception != null)
+            {
+                return exception is TaskCanceledException ? true : false;
+            }
+
+            return (int)response.StatusCode switch
+            {
+                // Not Found, Request Timeout, Gone, Too Many Requests
+                404 or 408 or 410 or 429 => true,
+                // Server Error range
+                >= 500 and <= 599 => true,
+                _ => false,
+            };
+        }
+
         /// <summary>
         /// Retry condition for /token and /authorize endpoints
         /// </summary>

src/client/Microsoft.Identity.Client/Http/Retry/IRetryPolicy.cs

@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Threading.Tasks;
+using Microsoft.Identity.Client.Core;
+
+namespace Microsoft.Identity.Client.Http.Retry
+{
+    // Interface for implementing retry logic for HTTP requests. 
+    // Determines if a retry should occur and handles pause logic between retries.
+    internal interface IRetryPolicy
+    {
+        /// <summary>
+        /// Determines whether a retry should be attempted for a given HTTP response or exception,
+        /// and performs any necessary pause or delay logic before the next retry attempt.
+        /// </summary>
+        /// <param name="response">The HTTP response received from the request.</param>
+        /// <param name="exception">The exception encountered during the request.</param>
+        /// <param name="retryCount">The current retry attempt count.</param>
+        /// <param name="logger">The logger used for diagnostic and informational messages.</param>
+        /// <returns>A task that returns true if a retry should be performed; otherwise, false.</returns>
+        Task<bool> PauseForRetryAsync(HttpResponse response, Exception exception, int retryCount, ILoggerAdapter logger);
+    }
+}

src/client/Microsoft.Identity.Client/Http/Retry/IRetryPolicyFactory.cs

@@ -0,0 +1,10 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+namespace Microsoft.Identity.Client.Http.Retry
+{
+    internal interface IRetryPolicyFactory
+    {
+        IRetryPolicy GetRetryPolicy(RequestType requestType);
+    }
+}