Fix KeyNotFoundException when headers don't contain correlation ID

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

Author: Copilot

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

@@ -140,17 +140,24 @@ public async Task<HttpResponse> SendRequestAsync(
 
                 if (headers != null && headers.Count > 0)
                 {
-                    var correlationId = headers[OAuth2Header.CorrelationId];
-                    string correlationIdMsg = headers.ContainsKey(OAuth2Header.CorrelationId) ?
-                                                $" CorrelationId: {correlationId}" :
-                                                string.Empty;
+                    string correlationIdMsg = string.Empty;
+                    string correlationId = null;
+
+                    if (headers.ContainsKey(OAuth2Header.CorrelationId))
+                    {
+                        correlationId = headers[OAuth2Header.CorrelationId];
+                        correlationIdMsg = $" CorrelationId: {correlationId}";
+                    }
 
                     var ex = new MsalServiceException(
                                 MsalError.RequestTimeout,
                                 msg + correlationIdMsg,
                                 timeoutException);
 
-                    ex.CorrelationId = correlationId;
+                    if (correlationId != null)
+                    {
+                        ex.CorrelationId = correlationId;
+                    }
 
                     throw ex;
                 }

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

@@ -581,6 +581,47 @@ public async Task TestCorrelationIdWithRetryOnTimeoutFailureAsync(bool addCorrel
             }
         }
 
+        // Regression test for https://github.com/AzureAD/microsoft-authentication-library-for-dotnet/issues/XXXX
+        // Verifies that timeout with headers but without correlation ID does not throw KeyNotFoundException
+        [TestMethod]
+        public async Task TestRetryOnTimeoutWithHeadersButNoCorrelationIdAsync()
+        {
+            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);
+                }
+
+                // Create headers without correlation ID - this was causing KeyNotFoundException
+                var headers = new Dictionary<string, string>
+                {
+                    ["some-other-header"] = "some-value"
+                };
+
+                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);
+
+                // Should get timeout error without KeyNotFoundException
+                Assert.AreEqual(MsalError.RequestTimeout, exc.ErrorCode);
+                Assert.AreEqual("Request to the endpoint timed out.", exc.Message);
+                Assert.IsNull(exc.CorrelationId);
+            }
+        }
+
         [TestMethod]
         public async Task TestWithCorrelationId_RetryOnTimeoutFailureAsync()
         {