Forward x-ms-client-request-id and use AzPS exception types (#452)

- Forward x-ms-client-request-id header to acquirePolicyToken API call for log
  correlation (Celina's feedback)
- Replace InvalidOperationException with AzPSInvalidOperationException and
  AzPSCloudException for proper Azure PowerShell error handling patterns
- Add ErrorKind categorization (ServiceError, UserError) and desensitized
  messages for telemetry
- Improve error messages with request URI context
- Add E2E test for x-ms-client-request-id forwarding
- Update exception assertions in existing tests

Author: notyashhh

src/Common.Test/AcquirePolicyTokenHandlerTests.cs

@@ -12,6 +12,8 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
+using Microsoft.Azure.Commands.Common;
+using Microsoft.Azure.Commands.Common.Exceptions;
 using Microsoft.WindowsAzure.Commands.Common;
 using Newtonsoft.Json;
 using Newtonsoft.Json.Linq;
@@ -841,6 +843,28 @@ public async Task E2E_TokenRequest_ForwardsAuthHeaders()
             Assert.True(capturedTokenRequest.Headers.Contains("x-ms-authorization-auxiliary"));
         }
 
+        // ----- Token request forwards x-ms-client-request-id for log correlation -----
+
+        [Fact]
+        public async Task E2E_TokenRequest_ForwardsClientRequestId()
+        {
+            HttpRequestMessage capturedTokenRequest = null;
+            var tokenClient = CreateMockTokenServer(req => capturedTokenRequest = req);
+            var (handler, client) = CreateTestPipeline(tokenClient);
+
+            var request = new HttpRequestMessage(HttpMethod.Delete,
+                $"{BaseUrl}/subscriptions/{TestSubId}/resourceGroups/rg/providers/Microsoft.Storage/storageAccounts/sa1?api-version=2024-01-01");
+            request.Headers.Authorization = new System.Net.Http.Headers.AuthenticationHeaderValue("Bearer", "test-token");
+            request.Headers.Add("x-ms-client-request-id", "correlation-id-abc-123");
+
+            await client.SendAsync(request);
+
+            Assert.True(capturedTokenRequest.Headers.Contains("x-ms-client-request-id"),
+                "Token request must forward x-ms-client-request-id for log correlation.");
+            Assert.Equal("correlation-id-abc-123",
+                capturedTokenRequest.Headers.GetValues("x-ms-client-request-id").First());
+        }
+
         // ----- Token API failure returns clear error -----
 
         [Fact]
@@ -849,12 +873,13 @@ public async Task E2E_TokenApiFails_ThrowsWithMessage()
             var tokenClient = CreateFailingTokenServer(HttpStatusCode.Forbidden, "Access denied");
             var (handler, client) = CreateTestPipeline(tokenClient);
 
-            var ex = await Assert.ThrowsAsync<InvalidOperationException>(() =>
+            var ex = await Assert.ThrowsAsync<AzPSCloudException>(() =>
                 client.SendAsync(new HttpRequestMessage(HttpMethod.Delete,
                     $"{BaseUrl}/subscriptions/{TestSubId}/resourceGroups/rg/providers/Microsoft.Storage/storageAccounts/sa1?api-version=2024-01-01")));
 
-            Assert.Contains("Failed to acquire policy token", ex.Message);
+            Assert.Contains("Policy token acquisition failed", ex.Message);
             Assert.Contains("403", ex.Message);
+            Assert.Equal(ErrorKind.ServiceError, ex.ErrorKind);
         }
 
         // ----- Token API returns 200 but no token field -----
@@ -873,12 +898,13 @@ public async Task E2E_TokenApiReturns200ButNoToken_ThrowsWithResponse()
             var tokenClient = new HttpClient(handler);
             var (pipelineHandler, client) = CreateTestPipeline(tokenClient);
 
-            var ex = await Assert.ThrowsAsync<InvalidOperationException>(() =>
+            var ex = await Assert.ThrowsAsync<AzPSCloudException>(() =>
                 client.SendAsync(new HttpRequestMessage(HttpMethod.Delete,
                     $"{BaseUrl}/subscriptions/{TestSubId}/resourceGroups/rg/providers/Microsoft.Storage/storageAccounts/sa1?api-version=2024-01-01")));
 
-            Assert.Contains("No token returned", ex.Message);
+            Assert.Contains("no token was returned", ex.Message);
             Assert.Contains("validator error", ex.Message);
+            Assert.Equal(ErrorKind.ServiceError, ex.ErrorKind);
         }
 
         // ----- GET request is NOT intercepted even with flag -----

src/Common/AcquirePolicyTokenHandler.cs

@@ -21,6 +21,8 @@
 using System.Text.RegularExpressions;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Azure.Commands.Common;
+using Microsoft.Azure.Commands.Common.Exceptions;
 using Newtonsoft.Json;
 using Newtonsoft.Json.Linq;
 
@@ -112,10 +114,22 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
                     EnqueueDebug("No token returned (null/empty).");
                 }
             }
+            catch (AzPSInvalidOperationException)
+            {
+                throw;
+            }
+            catch (AzPSCloudException)
+            {
+                throw;
+            }
             catch (Exception ex)
             {
                 EnqueueDebug($"Exception: {ex.GetType().Name}: {ex.Message}");
-                throw new InvalidOperationException($"Failed to acquire policy token: {ex.Message}", ex);
+                throw new AzPSInvalidOperationException(
+                    $"Failed to acquire policy token for {request.Method} {request.RequestUri}: {ex.Message}",
+                    ErrorKind.ServiceError,
+                    ex,
+                    desensitizedMessage: "Failed to acquire policy token.");
             }
 
             return await base.SendAsync(request, cancellationToken).ConfigureAwait(false);
@@ -126,8 +140,11 @@ private async Task<string> AcquirePolicyTokenAsync(HttpRequestMessage originalRe
             var subscriptionId = ExtractSubscriptionId(originalRequest.RequestUri);
             if (string.IsNullOrEmpty(subscriptionId))
             {
-                EnqueueDebug("Failed: subscription id not found in URI.");
-                throw new InvalidOperationException("Unable to determine subscription ID for policy token acquisition.");
+                EnqueueDebug($"Failed: subscription id not found in URI {originalRequest.RequestUri}.");
+                throw new AzPSInvalidOperationException(
+                    $"Unable to determine subscription ID for policy token acquisition from URI: {originalRequest.RequestUri}",
+                    ErrorKind.UserError,
+                    desensitizedMessage: "Unable to determine subscription ID for policy token acquisition.");
             }
 
             var authority = originalRequest.RequestUri.GetLeftPart(UriPartial.Authority);
@@ -180,6 +197,10 @@ private async Task<string> AcquirePolicyTokenAsync(HttpRequestMessage originalRe
             {
                 tokenRequest.Headers.TryAddWithoutValidation("x-ms-authorization-auxiliary", auxValues);
             }
+            if (originalRequest.Headers.TryGetValues("x-ms-client-request-id", out var clientRequestIdValues))
+            {
+                tokenRequest.Headers.TryAddWithoutValidation("x-ms-client-request-id", clientRequestIdValues);
+            }
 
             var isOwnedClient = _tokenHttpClient == null;
             var http = _tokenHttpClient ?? new HttpClient();
@@ -198,21 +219,33 @@ private async Task<string> AcquirePolicyTokenAsync(HttpRequestMessage originalRe
                         if (string.IsNullOrEmpty(token))
                         {
                             EnqueueDebug("Response OK but token missing.");
-                            throw new InvalidOperationException($"No token returned. Response:{responseContent}");
+                            throw new AzPSCloudException(
+                                $"Policy token acquisition succeeded but no token was returned. Response: {responseContent}",
+                                ErrorKind.ServiceError,
+                                desensitizedMessage: "Policy token acquisition succeeded but no token was returned.");
                         }
                         return token;
                     }
-                    throw new InvalidOperationException("Empty response body when acquiring policy token.");
+                    throw new AzPSCloudException(
+                        "Policy token acquisition returned an empty response body.",
+                        ErrorKind.ServiceError,
+                        desensitizedMessage: "Policy token acquisition returned an empty response body.");
                 }
                 else if (response.StatusCode == HttpStatusCode.Accepted)
                 {
                     EnqueueDebug("202 Accepted received (async not supported).");
-                    throw new InvalidOperationException("Asynchronous policy token acquisition (202 Accepted) is not supported.");
+                    throw new AzPSCloudException(
+                        "Asynchronous policy token acquisition (202 Accepted) is not supported.",
+                        ErrorKind.ServiceError,
+                        desensitizedMessage: "Asynchronous policy token acquisition (202 Accepted) is not supported.");
                 }
                 else
                 {
-                    EnqueueDebug("Non-success status; will throw.");
-                    throw new InvalidOperationException($"Policy token acquisition failed with {(int)response.StatusCode} {response.StatusCode}: {responseContent}");
+                    EnqueueDebug($"Non-success status {(int)response.StatusCode}; will throw.");
+                    throw new AzPSCloudException(
+                        $"Policy token acquisition failed with {(int)response.StatusCode} {response.StatusCode}: {responseContent}",
+                        ErrorKind.ServiceError,
+                        desensitizedMessage: $"Policy token acquisition failed with status {(int)response.StatusCode}.");
                 }
             }
             finally