Implemented GitHub feedback

Robbie-Microsoft · Robbie-Microsoft · commit b5eeaf3a98b5 · 2025-12-04T17:13:04.000-05:00
diff --git a/src/client/Microsoft.Identity.Client/ManagedIdentity/ImdsManagedIdentitySource.cs b/src/client/Microsoft.Identity.Client/ManagedIdentity/ImdsManagedIdentitySource.cs
@@ -243,7 +243,8 @@ public static string ImdsQueryParamsHelper(
 
         public static async Task<bool> ProbeImdsEndpointAsync(
             RequestContext requestContext,
-            ImdsVersion imdsVersion)
+            ImdsVersion imdsVersion,
+            CancellationToken cancellationToken)
         {
             string apiVersionQueryParam;
             string imdsApiVersion;
@@ -289,21 +290,18 @@ public static async Task<bool> ProbeImdsEndpointAsync(
 
             try
             {
-                using (var timeoutCts = new CancellationTokenSource(TimeSpan.FromSeconds(1)))
-                {
-                    response = await requestContext.ServiceBundle.HttpManager.SendRequestAsync(
-                        GetValidatedEndpoint(requestContext.Logger, imdsEndpoint, queryParams),
-                        headers,
-                        body: null,
-                        method: HttpMethod.Get,
-                        logger: requestContext.Logger,
-                        doNotThrow: false,
-                        mtlsCertificate: null,
-                        validateServerCertificate: null,
-                        cancellationToken: timeoutCts.Token,
-                        retryPolicy: retryPolicy)
-                    .ConfigureAwait(false);
-                }
+                response = await requestContext.ServiceBundle.HttpManager.SendRequestAsync(
+                    GetValidatedEndpoint(requestContext.Logger, imdsEndpoint, queryParams),
+                    headers,
+                    body: null,
+                    method: HttpMethod.Get,
+                    logger: requestContext.Logger,
+                    doNotThrow: false,
+                    mtlsCertificate: null,
+                    validateServerCertificate: null,
+                    cancellationToken: cancellationToken,
+                    retryPolicy: retryPolicy)
+                .ConfigureAwait(false);
             }
             catch (Exception ex)
             {
diff --git a/src/client/Microsoft.Identity.Client/ManagedIdentity/ManagedIdentityClient.cs b/src/client/Microsoft.Identity.Client/ManagedIdentity/ManagedIdentityClient.cs
@@ -40,12 +40,15 @@ internal async Task<ManagedIdentityResponse> SendTokenRequestForManagedIdentityA
             AcquireTokenForManagedIdentityParameters parameters,
             CancellationToken cancellationToken)
         {
-            AbstractManagedIdentity msi = await GetOrSelectManagedIdentitySourceAsync(requestContext, parameters.IsMtlsPopRequested).ConfigureAwait(false);
+            AbstractManagedIdentity msi = await GetOrSelectManagedIdentitySourceAsync(requestContext, parameters.IsMtlsPopRequested, cancellationToken).ConfigureAwait(false);
             return await msi.AuthenticateAsync(parameters, cancellationToken).ConfigureAwait(false);
         }
 
         // This method tries to create managed identity source for different sources, if none is created then defaults to IMDS.
-        private async Task<AbstractManagedIdentity> GetOrSelectManagedIdentitySourceAsync(RequestContext requestContext, bool isMtlsPopRequested)
+        private async Task<AbstractManagedIdentity> GetOrSelectManagedIdentitySourceAsync(
+            RequestContext requestContext,
+            bool isMtlsPopRequested,
+            CancellationToken cancellationToken)
         {
             using (requestContext.Logger.LogMethodDuration())
             {
@@ -57,7 +60,7 @@ private async Task<AbstractManagedIdentity> GetOrSelectManagedIdentitySourceAsyn
                 if (s_sourceName == ManagedIdentitySource.None)
                 {
                     // First invocation: detect and cache
-                    source = await GetManagedIdentitySourceAsync(requestContext, isMtlsPopRequested).ConfigureAwait(false);
+                    source = await GetManagedIdentitySourceAsync(requestContext, isMtlsPopRequested, cancellationToken).ConfigureAwait(false);
                 }
                 else
                 {
@@ -102,7 +105,8 @@ private async Task<AbstractManagedIdentity> GetOrSelectManagedIdentitySourceAsyn
         // This method is perf sensitive any changes should be benchmarked.
         internal async Task<ManagedIdentitySource> GetManagedIdentitySourceAsync(
             RequestContext requestContext,
-            bool isMtlsPopRequested)
+            bool isMtlsPopRequested,
+            CancellationToken cancellationToken)
         {
             // First check env vars to avoid the probe if possible
             ManagedIdentitySource source = GetManagedIdentitySourceNoImds(requestContext.Logger);
@@ -115,7 +119,7 @@ internal async Task<ManagedIdentitySource> GetManagedIdentitySourceAsync(
             // skip the ImdsV2 probe if MtlsPop was NOT requested
             if (isMtlsPopRequested)
             {
-                var imdsV2Response = await ImdsManagedIdentitySource.ProbeImdsEndpointAsync(requestContext, ImdsVersion.V2).ConfigureAwait(false);
+                var imdsV2Response = await ImdsManagedIdentitySource.ProbeImdsEndpointAsync(requestContext, ImdsVersion.V2, cancellationToken).ConfigureAwait(false);
                 if (imdsV2Response)
                 {
                     requestContext.Logger.Info("[Managed Identity] ImdsV2 detected.");
@@ -128,7 +132,7 @@ internal async Task<ManagedIdentitySource> GetManagedIdentitySourceAsync(
                 requestContext.Logger.Info("[Managed Identity] Mtls Pop was not requested; skipping ImdsV2 probe.");
             }
 
-            var imdsV1Response = await ImdsManagedIdentitySource.ProbeImdsEndpointAsync(requestContext, ImdsVersion.V1).ConfigureAwait(false);
+            var imdsV1Response = await ImdsManagedIdentitySource.ProbeImdsEndpointAsync(requestContext, ImdsVersion.V1, cancellationToken).ConfigureAwait(false);
             if (imdsV1Response)
             {
                 requestContext.Logger.Info("[Managed Identity] ImdsV1 detected.");
diff --git a/src/client/Microsoft.Identity.Client/ManagedIdentityApplication.cs b/src/client/Microsoft.Identity.Client/ManagedIdentityApplication.cs
@@ -56,18 +56,18 @@ public AcquireTokenForManagedIdentityParameterBuilder AcquireTokenForManagedIden
         }
 
         /// <inheritdoc/>
-        public async Task<ManagedIdentitySource> GetManagedIdentitySourceAsync()
+        public async Task<ManagedIdentitySource> GetManagedIdentitySourceAsync(CancellationToken cancellationToken)
         {
             if (ManagedIdentityClient.s_sourceName != ManagedIdentitySource.None)
             {
                 return ManagedIdentityClient.s_sourceName;
             }
 
             // Create a temporary RequestContext for the logger and the IMDS probe request.
-            var requestContext = new RequestContext(this.ServiceBundle, Guid.NewGuid(), null, CancellationToken.None);
+            var requestContext = new RequestContext(this.ServiceBundle, Guid.NewGuid(), null, cancellationToken);
 
             // GetManagedIdentitySourceAsync might return ImdsV2 = true, but it still requires .WithMtlsProofOfPossesion on the Managed Identity Application object to hit the ImdsV2 flow
-            return await ManagedIdentityClient.GetManagedIdentitySourceAsync(requestContext, isMtlsPopRequested: true).ConfigureAwait(false);
+            return await ManagedIdentityClient.GetManagedIdentitySourceAsync(requestContext, isMtlsPopRequested: true, cancellationToken).ConfigureAwait(false);
         }
 
         /// <summary>
diff --git a/src/client/Microsoft.Identity.Client/PublicApi/net462/PublicAPI.Shipped.txt b/src/client/Microsoft.Identity.Client/PublicApi/net462/PublicAPI.Shipped.txt
@@ -1086,7 +1086,6 @@ const Microsoft.Identity.Client.MsalError.MtlsNotSupportedForManagedIdentity = "
 const Microsoft.Identity.Client.MsalError.MtlsPopTokenNotSupportedinImdsV1 = "mtls_pop_token_not_supported_in_imds_v1" -> string
 Microsoft.Identity.Client.IMsalMtlsHttpClientFactory
 Microsoft.Identity.Client.IMsalMtlsHttpClientFactory.GetHttpClient(System.Security.Cryptography.X509Certificates.X509Certificate2 x509Certificate2) -> System.Net.Http.HttpClient
-Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync() -> System.Threading.Tasks.Task<Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource>
 Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource.ImdsV2 = 8 -> Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource
 Microsoft.Identity.Client.ManagedIdentityApplicationBuilder.WithExtraQueryParameters(System.Collections.Generic.IDictionary<string, string> extraQueryParameters) -> Microsoft.Identity.Client.ManagedIdentityApplicationBuilder
 static Microsoft.Identity.Client.ApplicationBase.ResetStateForTest() -> void
diff --git a/src/client/Microsoft.Identity.Client/PublicApi/net462/PublicAPI.Unshipped.txt b/src/client/Microsoft.Identity.Client/PublicApi/net462/PublicAPI.Unshipped.txt
@@ -1 +1,2 @@
 ﻿const Microsoft.Identity.Client.MsalError.ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable" -> string
+Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task<Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource>
diff --git a/src/client/Microsoft.Identity.Client/PublicApi/net472/PublicAPI.Shipped.txt b/src/client/Microsoft.Identity.Client/PublicApi/net472/PublicAPI.Shipped.txt
@@ -1086,7 +1086,6 @@ const Microsoft.Identity.Client.MsalError.MtlsNotSupportedForManagedIdentity = "
 const Microsoft.Identity.Client.MsalError.MtlsPopTokenNotSupportedinImdsV1 = "mtls_pop_token_not_supported_in_imds_v1" -> string
 Microsoft.Identity.Client.IMsalMtlsHttpClientFactory
 Microsoft.Identity.Client.IMsalMtlsHttpClientFactory.GetHttpClient(System.Security.Cryptography.X509Certificates.X509Certificate2 x509Certificate2) -> System.Net.Http.HttpClient
-Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync() -> System.Threading.Tasks.Task<Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource>
 Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource.ImdsV2 = 8 -> Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource
 Microsoft.Identity.Client.ManagedIdentityApplicationBuilder.WithExtraQueryParameters(System.Collections.Generic.IDictionary<string, string> extraQueryParameters) -> Microsoft.Identity.Client.ManagedIdentityApplicationBuilder
 static Microsoft.Identity.Client.ApplicationBase.ResetStateForTest() -> void
diff --git a/src/client/Microsoft.Identity.Client/PublicApi/net472/PublicAPI.Unshipped.txt b/src/client/Microsoft.Identity.Client/PublicApi/net472/PublicAPI.Unshipped.txt
@@ -1 +1,2 @@
 ﻿const Microsoft.Identity.Client.MsalError.ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable" -> string
+Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task<Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource>
diff --git a/src/client/Microsoft.Identity.Client/PublicApi/net8.0-android/PublicAPI.Shipped.txt b/src/client/Microsoft.Identity.Client/PublicApi/net8.0-android/PublicAPI.Shipped.txt
@@ -1052,7 +1052,6 @@ const Microsoft.Identity.Client.MsalError.MtlsNotSupportedForManagedIdentity = "
 const Microsoft.Identity.Client.MsalError.MtlsPopTokenNotSupportedinImdsV1 = "mtls_pop_token_not_supported_in_imds_v1" -> string
 Microsoft.Identity.Client.IMsalMtlsHttpClientFactory
 Microsoft.Identity.Client.IMsalMtlsHttpClientFactory.GetHttpClient(System.Security.Cryptography.X509Certificates.X509Certificate2 x509Certificate2) -> System.Net.Http.HttpClient
-Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync() -> System.Threading.Tasks.Task<Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource>
 Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource.ImdsV2 = 8 -> Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource
 Microsoft.Identity.Client.ManagedIdentityApplicationBuilder.WithExtraQueryParameters(System.Collections.Generic.IDictionary<string, string> extraQueryParameters) -> Microsoft.Identity.Client.ManagedIdentityApplicationBuilder
 static Microsoft.Identity.Client.ApplicationBase.ResetStateForTest() -> void
diff --git a/src/client/Microsoft.Identity.Client/PublicApi/net8.0-android/PublicAPI.Unshipped.txt b/src/client/Microsoft.Identity.Client/PublicApi/net8.0-android/PublicAPI.Unshipped.txt
@@ -1 +1,2 @@
 ﻿const Microsoft.Identity.Client.MsalError.ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable" -> string
+Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task<Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource>
diff --git a/src/client/Microsoft.Identity.Client/PublicApi/net8.0-ios/PublicAPI.Shipped.txt b/src/client/Microsoft.Identity.Client/PublicApi/net8.0-ios/PublicAPI.Shipped.txt
@@ -1054,7 +1054,6 @@ const Microsoft.Identity.Client.MsalError.MtlsNotSupportedForManagedIdentity = "
 const Microsoft.Identity.Client.MsalError.MtlsPopTokenNotSupportedinImdsV1 = "mtls_pop_token_not_supported_in_imds_v1" -> string
 Microsoft.Identity.Client.IMsalMtlsHttpClientFactory
 Microsoft.Identity.Client.IMsalMtlsHttpClientFactory.GetHttpClient(System.Security.Cryptography.X509Certificates.X509Certificate2 x509Certificate2) -> System.Net.Http.HttpClient
-Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync() -> System.Threading.Tasks.Task<Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource>
 Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource.ImdsV2 = 8 -> Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource
 Microsoft.Identity.Client.ManagedIdentityApplicationBuilder.WithExtraQueryParameters(System.Collections.Generic.IDictionary<string, string> extraQueryParameters) -> Microsoft.Identity.Client.ManagedIdentityApplicationBuilder
 static Microsoft.Identity.Client.ApplicationBase.ResetStateForTest() -> void
diff --git a/src/client/Microsoft.Identity.Client/PublicApi/net8.0-ios/PublicAPI.Unshipped.txt b/src/client/Microsoft.Identity.Client/PublicApi/net8.0-ios/PublicAPI.Unshipped.txt
@@ -1 +1,2 @@
 ﻿const Microsoft.Identity.Client.MsalError.ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable" -> string
+Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task<Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource>
diff --git a/src/client/Microsoft.Identity.Client/PublicApi/net8.0/PublicAPI.Shipped.txt b/src/client/Microsoft.Identity.Client/PublicApi/net8.0/PublicAPI.Shipped.txt
@@ -1048,7 +1048,6 @@ const Microsoft.Identity.Client.MsalError.MtlsNotSupportedForManagedIdentity = "
 const Microsoft.Identity.Client.MsalError.MtlsPopTokenNotSupportedinImdsV1 = "mtls_pop_token_not_supported_in_imds_v1" -> string
 Microsoft.Identity.Client.IMsalMtlsHttpClientFactory
 Microsoft.Identity.Client.IMsalMtlsHttpClientFactory.GetHttpClient(System.Security.Cryptography.X509Certificates.X509Certificate2 x509Certificate2) -> System.Net.Http.HttpClient
-Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync() -> System.Threading.Tasks.Task<Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource>
 Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource.ImdsV2 = 8 -> Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource
 Microsoft.Identity.Client.ManagedIdentityApplicationBuilder.WithExtraQueryParameters(System.Collections.Generic.IDictionary<string, string> extraQueryParameters) -> Microsoft.Identity.Client.ManagedIdentityApplicationBuilder
 static Microsoft.Identity.Client.ApplicationBase.ResetStateForTest() -> void
diff --git a/src/client/Microsoft.Identity.Client/PublicApi/net8.0/PublicAPI.Unshipped.txt b/src/client/Microsoft.Identity.Client/PublicApi/net8.0/PublicAPI.Unshipped.txt
@@ -1 +1,2 @@
 ﻿const Microsoft.Identity.Client.MsalError.ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable" -> string
+Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task<Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource>
diff --git a/src/client/Microsoft.Identity.Client/PublicApi/netstandard2.0/PublicAPI.Shipped.txt b/src/client/Microsoft.Identity.Client/PublicApi/netstandard2.0/PublicAPI.Shipped.txt
@@ -1048,7 +1048,6 @@ const Microsoft.Identity.Client.MsalError.MtlsNotSupportedForManagedIdentity = "
 const Microsoft.Identity.Client.MsalError.MtlsPopTokenNotSupportedinImdsV1 = "mtls_pop_token_not_supported_in_imds_v1" -> string
 Microsoft.Identity.Client.IMsalMtlsHttpClientFactory
 Microsoft.Identity.Client.IMsalMtlsHttpClientFactory.GetHttpClient(System.Security.Cryptography.X509Certificates.X509Certificate2 x509Certificate2) -> System.Net.Http.HttpClient
-Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync() -> System.Threading.Tasks.Task<Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource>
 Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource.ImdsV2 = 8 -> Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource
 Microsoft.Identity.Client.ManagedIdentityApplicationBuilder.WithExtraQueryParameters(System.Collections.Generic.IDictionary<string, string> extraQueryParameters) -> Microsoft.Identity.Client.ManagedIdentityApplicationBuilder
 static Microsoft.Identity.Client.ApplicationBase.ResetStateForTest() -> void
diff --git a/src/client/Microsoft.Identity.Client/PublicApi/netstandard2.0/PublicAPI.Unshipped.txt b/src/client/Microsoft.Identity.Client/PublicApi/netstandard2.0/PublicAPI.Unshipped.txt
@@ -1 +1,2 @@
 ﻿const Microsoft.Identity.Client.MsalError.ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable" -> string
+Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task<Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource>
diff --git a/tests/Microsoft.Identity.Test.Common/Core/Mocks/MockHttpMessageHandler.cs b/tests/Microsoft.Identity.Test.Common/Core/Mocks/MockHttpMessageHandler.cs
@@ -41,15 +41,8 @@ internal class MockHttpMessageHandler : HttpClientHandler
 
         public X509Certificate2 ExpectedMtlsBindingCertificate { get; set; }
 
-        public Func<HttpRequestMessage, CancellationToken, Task<HttpResponseMessage>> HandlerFunc { get; set; }
-
         protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
         {
-            if (HandlerFunc != null)
-            {
-                return await HandlerFunc(request, cancellationToken).ConfigureAwait(false);
-            }
-
             ActualRequestMessage = request;
 
             if (ExceptionToThrow != null)
diff --git a/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/AppServiceTests.cs b/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/AppServiceTests.cs
@@ -70,7 +70,7 @@ public async Task TestAppServiceUpgradeScenario(
 
                 ManagedIdentityApplication mi = miBuilder.Build() as ManagedIdentityApplication;
 
-                Assert.AreEqual(expectedManagedIdentitySource, await mi.GetManagedIdentitySourceAsync().ConfigureAwait(false));
+                Assert.AreEqual(expectedManagedIdentitySource, await mi.GetManagedIdentitySourceAsync(ManagedIdentityTests.ImdsProbesCancellationToken).ConfigureAwait(false));
             }
         }
     }
diff --git a/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsTests.cs b/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsTests.cs
@@ -4,6 +4,7 @@
 using System;
 using System.Net;
 using System.Net.Http;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client;
 using Microsoft.Identity.Client.AppConfig;
@@ -449,19 +450,11 @@ public async Task ProbeImdsEndpointAsync_TimesOutAfterOneSecond()
                 var managedIdentityApp = miBuilder.Build();
 
                 httpManager.AddMockHandler(MockHelpers.MockImdsProbeFailure(ImdsVersion.V2));
+                httpManager.AddMockHandler(MockHelpers.MockImdsProbe(ImdsVersion.V1));
 
-                var imdsV1Handler = new MockHttpMessageHandler
-                {
-                    HandlerFunc = async (request, cancellationToken) =>
-                    {
-                        // IMDS (V1 and V2) probe has a timeout after 1 second
-                        await Task.Delay(TimeSpan.FromSeconds(2), cancellationToken).ConfigureAwait(false);
-                        return new HttpResponseMessage(HttpStatusCode.OK);
-                    }
-                };
-                httpManager.AddMockHandler(imdsV1Handler);
-
-                var miSource = await (managedIdentityApp as ManagedIdentityApplication).GetManagedIdentitySourceAsync().ConfigureAwait(false);
+                var imdsProbesCancellationToken = new CancellationTokenSource(TimeSpan.FromSeconds(0)).Token; // timeout immediately
+                
+                var miSource = await (managedIdentityApp as ManagedIdentityApplication).GetManagedIdentitySourceAsync(imdsProbesCancellationToken).ConfigureAwait(false);
                 Assert.AreEqual(ManagedIdentitySource.None, miSource); // Probe timed out, no source available
 
                 var ex = await Assert.ThrowsExceptionAsync<MsalClientException>(async () =>
diff --git a/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsV2Tests.cs b/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsV2Tests.cs
diff --git a/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ManagedIdentityTests.cs b/tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ManagedIdentityTests.cs

Original file line number	Diff line number	Diff line change
`@@ -40,12 +40,15 @@ internal async Task<ManagedIdentityResponse> SendTokenRequestForManagedIdentityA`
`40`	`40`	`AcquireTokenForManagedIdentityParameters parameters,`
`41`	`41`	`CancellationToken cancellationToken)`
`42`	`42`	`{`
`43`		`- AbstractManagedIdentity msi = await GetOrSelectManagedIdentitySourceAsync(requestContext, parameters.IsMtlsPopRequested).ConfigureAwait(false);`
	`43`	`+ AbstractManagedIdentity msi = await GetOrSelectManagedIdentitySourceAsync(requestContext, parameters.IsMtlsPopRequested, cancellationToken).ConfigureAwait(false);`
`44`	`44`	`return await msi.AuthenticateAsync(parameters, cancellationToken).ConfigureAwait(false);`
`45`	`45`	`}`
`46`	`46`
`47`	`47`	`// This method tries to create managed identity source for different sources, if none is created then defaults to IMDS.`
`48`		`- private async Task<AbstractManagedIdentity> GetOrSelectManagedIdentitySourceAsync(RequestContext requestContext, bool isMtlsPopRequested)`
	`48`	`+ private async Task<AbstractManagedIdentity> GetOrSelectManagedIdentitySourceAsync(`
	`49`	`+ RequestContext requestContext,`
	`50`	`+ bool isMtlsPopRequested,`
	`51`	`+ CancellationToken cancellationToken)`
`49`	`52`	`{`
`50`	`53`	`using (requestContext.Logger.LogMethodDuration())`
`51`	`54`	`{`
`@@ -57,7 +60,7 @@ private async Task<AbstractManagedIdentity> GetOrSelectManagedIdentitySourceAsyn`
`57`	`60`	`if (s_sourceName == ManagedIdentitySource.None)`
`58`	`61`	`{`
`59`	`62`	`// First invocation: detect and cache`
`60`		`- source = await GetManagedIdentitySourceAsync(requestContext, isMtlsPopRequested).ConfigureAwait(false);`
	`63`	`+ source = await GetManagedIdentitySourceAsync(requestContext, isMtlsPopRequested, cancellationToken).ConfigureAwait(false);`
`61`	`64`	`}`
`62`	`65`	`else`
`63`	`66`	`{`
`@@ -102,7 +105,8 @@ private async Task<AbstractManagedIdentity> GetOrSelectManagedIdentitySourceAsyn`
`102`	`105`	`// This method is perf sensitive any changes should be benchmarked.`
`103`	`106`	`internal async Task<ManagedIdentitySource> GetManagedIdentitySourceAsync(`
`104`	`107`	`RequestContext requestContext,`
`105`		`- bool isMtlsPopRequested)`
	`108`	`+ bool isMtlsPopRequested,`
	`109`	`+ CancellationToken cancellationToken)`
`106`	`110`	`{`
`107`	`111`	`// First check env vars to avoid the probe if possible`
`108`	`112`	`ManagedIdentitySource source = GetManagedIdentitySourceNoImds(requestContext.Logger);`
`@@ -115,7 +119,7 @@ internal async Task<ManagedIdentitySource> GetManagedIdentitySourceAsync(`
`115`	`119`	`// skip the ImdsV2 probe if MtlsPop was NOT requested`
`116`	`120`	`if (isMtlsPopRequested)`
`117`	`121`	`{`
`118`		`- var imdsV2Response = await ImdsManagedIdentitySource.ProbeImdsEndpointAsync(requestContext, ImdsVersion.V2).ConfigureAwait(false);`
	`122`	`+ var imdsV2Response = await ImdsManagedIdentitySource.ProbeImdsEndpointAsync(requestContext, ImdsVersion.V2, cancellationToken).ConfigureAwait(false);`
`119`	`123`	`if (imdsV2Response)`
`120`	`124`	`{`
`121`	`125`	`requestContext.Logger.Info("[Managed Identity] ImdsV2 detected.");`
`@@ -128,7 +132,7 @@ internal async Task<ManagedIdentitySource> GetManagedIdentitySourceAsync(`
`128`	`132`	`requestContext.Logger.Info("[Managed Identity] Mtls Pop was not requested; skipping ImdsV2 probe.");`
`129`	`133`	`}`
`130`	`134`
`131`		`- var imdsV1Response = await ImdsManagedIdentitySource.ProbeImdsEndpointAsync(requestContext, ImdsVersion.V1).ConfigureAwait(false);`
	`135`	`+ var imdsV1Response = await ImdsManagedIdentitySource.ProbeImdsEndpointAsync(requestContext, ImdsVersion.V1, cancellationToken).ConfigureAwait(false);`
`132`	`136`	`if (imdsV1Response)`
`133`	`137`	`{`
`134`	`138`	`requestContext.Logger.Info("[Managed Identity] ImdsV1 detected.");`
Original file line number	Diff line number	Diff line change
`@@ -56,18 +56,18 @@ public AcquireTokenForManagedIdentityParameterBuilder AcquireTokenForManagedIden`
`56`	`56`	`}`
`57`	`57`
`58`	`58`	`/// <inheritdoc/>`
`59`		`- public async Task<ManagedIdentitySource> GetManagedIdentitySourceAsync()`
	`59`	`+ public async Task<ManagedIdentitySource> GetManagedIdentitySourceAsync(CancellationToken cancellationToken)`
`60`	`60`	`{`
`61`	`61`	`if (ManagedIdentityClient.s_sourceName != ManagedIdentitySource.None)`
`62`	`62`	`{`
`63`	`63`	`return ManagedIdentityClient.s_sourceName;`
`64`	`64`	`}`
`65`	`65`
`66`	`66`	`// Create a temporary RequestContext for the logger and the IMDS probe request.`
`67`		`- var requestContext = new RequestContext(this.ServiceBundle, Guid.NewGuid(), null, CancellationToken.None);`
	`67`	`+ var requestContext = new RequestContext(this.ServiceBundle, Guid.NewGuid(), null, cancellationToken);`
`68`	`68`
`69`	`69`	`// GetManagedIdentitySourceAsync might return ImdsV2 = true, but it still requires .WithMtlsProofOfPossesion on the Managed Identity Application object to hit the ImdsV2 flow`
`70`		`- return await ManagedIdentityClient.GetManagedIdentitySourceAsync(requestContext, isMtlsPopRequested: true).ConfigureAwait(false);`
	`70`	`+ return await ManagedIdentityClient.GetManagedIdentitySourceAsync(requestContext, isMtlsPopRequested: true, cancellationToken).ConfigureAwait(false);`
`71`	`71`	`}`
`72`	`72`
`73`	`73`	`/// <summary>`
Original file line number	Diff line number	Diff line change
`@@ -1 +1,2 @@`
`1`	`1`	`const Microsoft.Identity.Client.MsalError.ManagedIdentityAllSourcesUnavailable = "managed_identity_all_sources_unavailable" -> string`
	`2`	`+Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task<Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource>`