Add timeout for managed identity capability discovery

docs/mtlspop_managed_identity.md

@@ -273,6 +273,32 @@ When attestation is configured:
 - The MAA JWT is embedded in the `/issuecredential` request body as `attestation_token`
 - The token cache partitions attested vs non-attested tokens (`att=1` vs `att=0`) — they are not interchangeable
 
+## Bounding capability discovery
+
+Credential chains can bound uncached managed identity capability discovery by supplying
+`ManagedIdentityCapabilitiesOptions.CapabilityDiscoveryTimeout`:
+
+```csharp
+var options = new ManagedIdentityCapabilitiesOptions
+{
+    CapabilityDiscoveryTimeout = TimeSpan.FromSeconds(2)
+};
+
+ManagedIdentityCapabilities capabilities =
+    await managedIdentityApplication.GetManagedIdentityCapabilitiesAsync(
+        options,
+        CancellationToken.None);
+```
+
+The timeout is a total discovery budget covering lock contention, IMDSv2 probing and
+retries, IMDSv1 fallback, compute metadata retrieval, and binding-strength detection.
+Omitting the timeout preserves the existing unlimited behavior and cancellation
+ordering. When a timeout is configured, an already-canceled caller token stops uncached
+discovery before environment detection. Caller cancellation continues to surface as
+cancellation, while expiration of the discovery budget throws `MsalServiceException`
+with error code `request_timeout` and a capability-discovery-specific message. A
+timed-out discovery result is not cached, so a later call can retry.
+
 ---
 
 ## Constraints

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

@@ -65,7 +65,8 @@ public async Task<HttpResponse> SendRequestAsync(
             Func<HttpRequestMessage, X509Certificate2, X509Chain, SslPolicyErrors, bool> validateServerCert,
             CancellationToken cancellationToken,
             IRetryPolicy retryPolicy,
-            int retryCount = 0)
+            int retryCount = 0,
+            CancellationToken retryDelayCancellationToken = default)
         {
             Exception timeoutException = null;
             HttpResponse response = null;
@@ -113,7 +114,13 @@ public async Task<HttpResponse> SendRequestAsync(
                 timeoutException = exception;
             }
 
-            while (!_disableInternalRetries && await retryPolicy.PauseForRetryAsync(response, timeoutException, retryCount, logger).ConfigureAwait(false))
+            while (!_disableInternalRetries &&
+                await retryPolicy.PauseForRetryAsync(
+                    response,
+                    timeoutException,
+                    retryCount,
+                    logger,
+                    retryDelayCancellationToken).ConfigureAwait(false))
             {
                 retryCount++;
 
@@ -128,7 +135,8 @@ public async Task<HttpResponse> SendRequestAsync(
                     validateServerCert,
                     cancellationToken,
                     retryPolicy,
-                    retryCount) // Pass the updated retry count
+                    retryCount,
+                    retryDelayCancellationToken) // Pass the updated retry count and delay budget
                     .ConfigureAwait(false);
             }
 

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

@@ -31,6 +31,7 @@ internal interface IHttpManager
         /// <param name="cancellationToken"></param>
         /// <param name="retryPolicy">Retry policy to be used for the request.</param>
         /// <param name="retryCount">Number of retries to be attempted in case of retriable status codes.</param>
+        /// <param name="retryDelayCancellationToken">Cancellation token observed only while waiting between retries.</param>
         /// <returns></returns>
         Task<HttpResponse> SendRequestAsync(
            Uri endpoint,
@@ -43,6 +44,7 @@ Task<HttpResponse> SendRequestAsync(
            Func<HttpRequestMessage, X509Certificate2, X509Chain, SslPolicyErrors, bool> validateServerCertificate,
            CancellationToken cancellationToken,
            IRetryPolicy retryPolicy,
-           int retryCount = 0);
+           int retryCount = 0,
+           CancellationToken retryDelayCancellationToken = default);
     }
 }

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

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.Core;
 
@@ -45,7 +46,12 @@ internal virtual Task DelayAsync(int milliseconds)
             return Task.Delay(milliseconds);
         }
 
-        public async Task<bool> PauseForRetryAsync(HttpResponse response, Exception exception, int retryCount, ILoggerAdapter logger)
+        public async Task<bool> PauseForRetryAsync(
+            HttpResponse response,
+            Exception exception,
+            int retryCount,
+            ILoggerAdapter logger,
+            CancellationToken retryDelayCancellationToken)
         {
             // Check if the status code is retriable and if the current retry count is less than max retries
             if (_retryCondition(response, exception) &&

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

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.Core;
 
@@ -19,7 +20,13 @@ internal interface IRetryPolicy
         /// <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>
+        /// <param name="retryDelayCancellationToken">The cancellation token to observe while waiting to retry.</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);
+        Task<bool> PauseForRetryAsync(
+            HttpResponse response,
+            Exception exception,
+            int retryCount,
+            ILoggerAdapter logger,
+            CancellationToken retryDelayCancellationToken);
     }
 }

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

@@ -3,6 +3,7 @@
 
 using System;
 using System.Net;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.Core;
 
@@ -28,17 +29,22 @@ internal class ImdsRetryPolicy : IRetryPolicy
             ImdsRetryPolicy.ExponentialDeltaBackoffMs
         );
 
-        internal virtual Task DelayAsync(int milliseconds)
+        internal virtual Task DelayAsync(int milliseconds, CancellationToken cancellationToken)
         {
-            return Task.Delay(milliseconds);
+            return Task.Delay(milliseconds, cancellationToken);
         }
 
         protected virtual bool ShouldRetry(HttpResponse response, Exception exception)
         {
             return HttpRetryConditions.Imds(response, exception);
         }
 
-        public async Task<bool> PauseForRetryAsync(HttpResponse response, Exception exception, int retryCount, ILoggerAdapter logger)
+        public async Task<bool> PauseForRetryAsync(
+            HttpResponse response,
+            Exception exception,
+            int retryCount,
+            ILoggerAdapter logger,
+            CancellationToken retryDelayCancellationToken)
         {
             int httpStatusCode = (int)response.StatusCode;
 
@@ -61,7 +67,7 @@ public async Task<bool> PauseForRetryAsync(HttpResponse response, Exception exce
                 logger.Warning($"Retrying request in {retryAfterDelay}ms (retry attempt: {retryCount + 1})");
 
                 // Pause execution for the calculated delay
-                await DelayAsync(retryAfterDelay).ConfigureAwait(false);
+                await DelayAsync(retryAfterDelay, retryDelayCancellationToken).ConfigureAwait(false);
 
                 return true;
             }

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

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.Core;
 
@@ -27,7 +28,12 @@ internal virtual Task DelayAsync(int milliseconds)
             return Task.Delay(milliseconds);
         }
 
-        public async Task<bool> PauseForRetryAsync(HttpResponse response, Exception exception, int retryCount, ILoggerAdapter logger)
+        public async Task<bool> PauseForRetryAsync(
+            HttpResponse response,
+            Exception exception,
+            int retryCount,
+            ILoggerAdapter logger,
+            CancellationToken retryDelayCancellationToken)
         {
             // Check if the status code is retriable and if the current retry count is less than max retries
             if (HttpRetryConditions.RegionDiscovery(response, exception) &&

src/client/Microsoft.Identity.Client/Instance/Discovery/NetworkMetadataProvider.cs

@@ -142,7 +142,12 @@ private static string GetTenant(Uri uri)
 
         private sealed class InstanceDiscoveryRetryPolicy : IRetryPolicy
         {
-            public Task<bool> PauseForRetryAsync(HttpResponse response, Exception exception, int retryCount, ILoggerAdapter logger)
+            public Task<bool> PauseForRetryAsync(
+                HttpResponse response,
+                Exception exception,
+                int retryCount,
+                ILoggerAdapter logger,
+                CancellationToken retryDelayCancellationToken)
             {
                 return Task.FromResult(false);
             }

src/client/Microsoft.Identity.Client/ManagedIdentity/ImdsComputeMetadataManager.cs

@@ -26,7 +26,8 @@ internal static class ImdsComputeMetadataManager
         internal static async Task<ComputeMetadataResponse> GetComputeMetadataAsync(
             IHttpManager httpManager,
             ILoggerAdapter logger,
-            CancellationToken cancellationToken)
+            CancellationToken cancellationToken,
+            CancellationToken retryDelayCancellationToken)
         {
             var headers = new Dictionary<string, string>
             {
@@ -53,7 +54,8 @@ internal static async Task<ComputeMetadataResponse> GetComputeMetadataAsync(
                     mtlsCertificate: null,
                     validateServerCertificate: null,
                     cancellationToken: cancellationToken,
-                    retryPolicy: new ImdsRetryPolicy())
+                    retryPolicy: new ImdsRetryPolicy(),
+                    retryDelayCancellationToken: retryDelayCancellationToken)
                     .ConfigureAwait(false);
 
                 if (response is null || response.StatusCode != HttpStatusCode.OK)

src/client/Microsoft.Identity.Client/ManagedIdentity/ImdsManagedIdentitySource.cs

@@ -252,7 +252,8 @@ public static string ImdsQueryParamsHelper(
         public static async Task<(bool success, string failureReason)> ProbeImdsEndpointAsync(
             RequestContext requestContext,
             ImdsVersion imdsVersion,
-            CancellationToken cancellationToken)
+            CancellationToken cancellationToken,
+            CancellationToken retryDelayCancellationToken)
         {
             string apiVersionQueryParam;
             string imdsApiVersion;
@@ -308,7 +309,8 @@ public static string ImdsQueryParamsHelper(
                     mtlsCertificate: null,
                     validateServerCertificate: null,
                     cancellationToken: cancellationToken,
-                    retryPolicy: retryPolicy)
+                    retryPolicy: retryPolicy,
+                    retryDelayCancellationToken: retryDelayCancellationToken)
                 .ConfigureAwait(false);
             }
             catch (Exception ex) when (ex is not OperationCanceledException)

src/client/Microsoft.Identity.Client/ManagedIdentity/ManagedIdentityCapabilities.cs

@@ -10,7 +10,9 @@ namespace Microsoft.Identity.Client.ManagedIdentity
     /// detected source and the strength with which tokens can be bound to a key.
     /// </summary>
     /// <remarks>
-    /// This type is returned by <see cref="ManagedIdentityApplication.GetManagedIdentityCapabilitiesAsync"/>.
+    /// This type is returned by
+    /// <see cref="ManagedIdentityApplication.GetManagedIdentityCapabilitiesAsync(System.Threading.CancellationToken)"/> and
+    /// <see cref="ManagedIdentityApplication.GetManagedIdentityCapabilitiesAsync(ManagedIdentityCapabilitiesOptions, System.Threading.CancellationToken)"/>.
     /// It is useful for credential chains such as <c>DefaultAzureCredential</c> to decide whether
     /// managed identity is available and what binding strength the host supports.
     /// </remarks>

src/client/Microsoft.Identity.Client/ManagedIdentity/ManagedIdentityCapabilitiesOptions.cs

@@ -0,0 +1,26 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+
+namespace Microsoft.Identity.Client.ManagedIdentity
+{
+    /// <summary>
+    /// Configures managed identity capability discovery.
+    /// </summary>
+    public sealed class ManagedIdentityCapabilitiesOptions
+    {
+        /// <summary>
+        /// Gets or sets the total time allowed for uncached managed identity capability discovery.
+        /// </summary>
+        /// <remarks>
+        /// The timeout covers discovery lock contention, IMDS probes and retries, fallback,
+        /// compute metadata retrieval, and binding-strength detection. A <c>null</c> value
+        /// preserves the existing cancellation ordering and waits without a discovery timeout.
+        /// When a timeout is configured, an already-canceled caller token is observed before
+        /// uncached discovery begins. Cancellation is cooperative, so non-interruptible platform
+        /// work may finish before the timeout is observed.
+        /// </remarks>
+        public TimeSpan? CapabilityDiscoveryTimeout { get; set; }
+    }
+}