Ensure correlationId is added to request_timeout exception when available (#5364)

* Ensure correlationId is added to retry exception when available

* Fixing test

* Add additional test

* Update httpManager

---------

Co-authored-by: trwalke <[email protected]>

Author: trwalke

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

@@ -14,6 +14,7 @@
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.Core;
 using Microsoft.Identity.Client.Http.Retry;
+using Microsoft.Identity.Client.OAuth2;
 
 namespace Microsoft.Identity.Client.Http
 {
@@ -134,10 +135,30 @@ public async Task<HttpResponse> SendRequestAsync(
             logger.Warning("Request retry failed.");
             if (timeoutException != null)
             {
+                //If the correlation id is available, include it in the exception message
+                string msg = MsalErrorMessage.RequestTimeOut;
+
+                if (headers != null && headers.Count > 0)
+                {
+                    var correlationId = headers[OAuth2Header.CorrelationId];
+                    string correlationIdMsg = headers.ContainsKey(OAuth2Header.CorrelationId) ?
+                                                $" CorrelationId: {correlationId}" :
+                                                string.Empty;
+
+                    var ex = new MsalServiceException(
+                                MsalError.RequestTimeout,
+                                msg + correlationIdMsg,
+                                timeoutException);
+
+                    ex.CorrelationId = correlationId;
+
+                    throw ex;
+                }
+
                 throw new MsalServiceException(
-                    MsalError.RequestTimeout,
-                    "Request to the endpoint timed out.",
-                    timeoutException);
+                                MsalError.RequestTimeout,
+                                msg,
+                                timeoutException);
             }
 
             if (doNotThrow)

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

@@ -438,5 +438,6 @@ public static string InvalidTokenProviderResponseValue(string invalidValueName)
         public const string MtlsNonTenantedAuthorityNotAllowedMessage = "mTLS authentication requires a tenanted authority. Using 'common', 'organizations', or similar non-tenanted authorities is not allowed. Please provide an authority with a specific tenant ID (e.g., 'https://login.microsoftonline.com/{tenantId}'). See https://aka.ms/msal-net-pop for details.";
         public const string RegionRequiredForMtlsPopMessage = "Regional auto-detect failed. mTLS Proof-of-Possession requires a region to be specified, as there is no global endpoint for mTLS. See https://aka.ms/msal-net-pop for details.";
         public const string ForceRefreshAndTokenHasNotCompatible = "Cannot specify ForceRefresh and AccessTokenSha256ToRefresh in the same request.";
+        public const string RequestTimeOut = "Request to the endpoint timed out.";
     }
 }

tests/Microsoft.Identity.Test.Unit/CoreTests/HttpTests/HttpManagerTests.cs

@@ -13,6 +13,7 @@
 using Microsoft.Identity.Client;
 using Microsoft.Identity.Client.Core;
 using Microsoft.Identity.Client.Http.Retry;
+using Microsoft.Identity.Client.OAuth2;
 using Microsoft.Identity.Test.Common;
 using Microsoft.Identity.Test.Common.Core.Helpers;
 using Microsoft.Identity.Test.Common.Core.Mocks;
@@ -536,6 +537,94 @@ public async Task TestSendPostWithRetryOnTimeoutFailureAsync()
             }
         }
 
+        [TestMethod]
+        [DataRow(true)]
+        [DataRow(false)]
+        public async Task TestCorrelationIdWithRetryOnTimeoutFailureAsync(bool addCorrelationId)
+        {
+            using (var httpManager = new MockHttpManager())
+            {
+                // Simulate permanent errors (to trigger the maximum number of retries)
+                const int Num500Errors = 1 + TestDefaultRetryPolicy.DefaultStsMaxRetries; // initial request + maximum number of retries
+                for (int i = 0; i < Num500Errors; i++)
+                {
+                    httpManager.AddRequestTimeoutResponseMessageMockHandler(HttpMethod.Post);
+                }
+
+                Guid correlationId = Guid.NewGuid();
+                var headers = new Dictionary<string, string>();
+
+                if (addCorrelationId)
+                {
+                    headers.Add(OAuth2Header.CorrelationId, correlationId.ToString());
+                }
+
+                var exc = await AssertException.TaskThrowsAsync<MsalServiceException>(() =>
+                    httpManager.SendRequestAsync(
+                        new Uri(TestConstants.AuthorityHomeTenant + "oauth2/token"),
+                        headers: headers,
+                        body: new FormUrlEncodedContent(new Dictionary<string, string>()),
+                        method: HttpMethod.Post,
+                        logger: Substitute.For<ILoggerAdapter>(),
+                        doNotThrow: false,
+                        mtlsCertificate: null,
+                        validateServerCert: null,
+                        cancellationToken: default,
+                        retryPolicy: _stsRetryPolicy))
+                    .ConfigureAwait(false);
+
+                Assert.AreEqual(MsalError.RequestTimeout, exc.ErrorCode);
+
+                if (addCorrelationId)
+                {
+                    Assert.AreEqual($"Request to the endpoint timed out. CorrelationId: {correlationId.ToString()}", exc.Message);
+                    Assert.AreEqual(correlationId.ToString(), exc.CorrelationId);
+                }
+                else
+                {
+                    Assert.AreEqual("Request to the endpoint timed out.", exc.Message);
+                }
+            }
+        }
+
+        [TestMethod]
+        public async Task TestWithCorrelationId_RetryOnTimeoutFailureAsync()
+        {
+            // Arrange
+            using (var httpManager = new MockHttpManager())
+            {
+                httpManager.AddInstanceDiscoveryMockHandler();
+
+                // Simulate permanent errors (to trigger the maximum number of retries)
+                const int Num500Errors = 1 + TestDefaultRetryPolicy.DefaultStsMaxRetries; // initial request + maximum number of retries
+                for (int i = 0; i < Num500Errors; i++)
+                {
+                    httpManager.AddRequestTimeoutResponseMessageMockHandler(HttpMethod.Post);
+                }
+                Guid correlationId = Guid.NewGuid();
+
+                var app = ConfidentialClientApplicationBuilder
+                            .Create(TestConstants.ClientId)
+                            .WithAuthority(TestConstants.AuthorityTestTenant)
+                            .WithHttpManager(httpManager)
+                            .WithClientSecret(TestConstants.ClientSecret)
+                            .Build();
+
+                var userAssertion = new UserAssertion(TestConstants.DefaultAccessToken);
+
+                // Act
+                var exc = await AssertException.TaskThrowsAsync<MsalServiceException>(() =>
+                                    app.AcquireTokenForClient(TestConstants.s_scope)
+                                    .WithCorrelationId(correlationId)
+                                    .ExecuteAsync())
+                                    .ConfigureAwait(false);
+
+                // Assert
+                Assert.AreEqual($"Request to the endpoint timed out. CorrelationId: {correlationId.ToString()}", exc.Message);
+                Assert.AreEqual(correlationId.ToString(), exc.CorrelationId);
+            }
+        }
+
         private class CapturingHandler : HttpMessageHandler
         {
             public HttpRequestMessage CapturedRequest { get; private set; }