CosmosException: Fixes the headers not matching CosmosException property values

Microsoft.Azure.Cosmos/src/Authorization/TokenCredentialCache.cs

@@ -221,7 +221,10 @@ private async ValueTask RefreshCachedTokenWithRetryHelperAsync(ITrace trace)
 
                             throw CosmosExceptionFactory.CreateRequestTimeoutException(
                                 message: ClientResources.FailedToGetAadToken,
-                                subStatusCode: (int)SubStatusCodes.FailedToGetAadToken,
+                                headers: new Headers()
+                                {
+                                    SubStatusCode = SubStatusCodes.FailedToGetAadToken,
+                                },
                                 innerException: lastException,
                                 trace: getTokenTrace);
                         }
@@ -253,7 +256,10 @@ private async ValueTask RefreshCachedTokenWithRetryHelperAsync(ITrace trace)
 
                 throw CosmosExceptionFactory.CreateUnauthorizedException(
                     message: ClientResources.FailedToGetAadToken,
-                    subStatusCode: (int)SubStatusCodes.FailedToGetAadToken,
+                    headers: new Headers()
+                    {
+                        SubStatusCode = SubStatusCodes.FailedToGetAadToken,
+                    },
                     innerException: lastException,
                     trace: trace);
             }

Microsoft.Azure.Cosmos/src/FeedRange/FeedRanges/FeedRangePartitionKeyRange.cs

@@ -49,13 +49,12 @@ public FeedRangePartitionKeyRange(string partitionKeyRangeId)
             {
                 throw CosmosExceptionFactory.Create(
                     statusCode: HttpStatusCode.Gone,
-                    subStatusCode: (int)SubStatusCodes.PartitionKeyRangeGone,
                     message: $"The PartitionKeyRangeId: \"{this.PartitionKeyRangeId}\" is not valid for the current container {containerRid} .",
                     stackTrace: string.Empty,
-                    activityId: string.Empty,
-                    requestCharge: 0,
-                    retryAfter: null,
-                    headers: null,
+                    headers: new Headers()
+                    {
+                        SubStatusCode = SubStatusCodes.PartitionKeyRangeGone,
+                    },
                     error: null,
                     innerException: null,
                     trace: NoOpTrace.Singleton);

Microsoft.Azure.Cosmos/src/HttpClient/CosmosHttpClientCore.cs

@@ -292,6 +292,10 @@ private async Task<HttpResponseMessage> SendHttpHelperAsync(
                                             $"GatewayStoreClient Request Timeout. Start Time UTC:{startDateTimeUtc}; Total Duration:{(DateTime.UtcNow - startDateTimeUtc).TotalMilliseconds} Ms; Request Timeout {requestTimeout.TotalMilliseconds} Ms; Http Client Timeout:{this.httpClient.Timeout.TotalMilliseconds} Ms; Activity id: {System.Diagnostics.Trace.CorrelationManager.ActivityId};";
                                         throw CosmosExceptionFactory.CreateRequestTimeoutException(
                                             message,
+                                            headers: new Headers()
+                                            {
+                                                ActivityId = System.Diagnostics.Trace.CorrelationManager.ActivityId.ToString()
+                                            },
                                             innerException: operationCanceledException,
                                             trace: helperTrace);
                                     }

Microsoft.Azure.Cosmos/src/Query/Core/ExceptionToCosmosException.cs

@@ -49,13 +49,12 @@ public static CosmosException CreateFromException(Exception exception)
             }
 
             return CosmosExceptionFactory.CreateInternalServerErrorException(
-                subStatusCode: default,
                 message: exception.Message,
                 stackTrace: exception.StackTrace,
-                activityId: EmptyGuidString,
-                requestCharge: 0,
-                retryAfter: null,
-                headers: null,
+                headers: new Headers()
+                {
+                    ActivityId = EmptyGuidString,
+                },
                 trace: NoOpTrace.Singleton,
                 innerException: exception);
         }
@@ -81,13 +80,9 @@ private static CosmosException CreateFromExceptionWithStackTrace(ExceptionWithSt
             CosmosException cosmosException = ExceptionToCosmosException.CreateFromException(exceptionWithStackTrace.InnerException);
             return CosmosExceptionFactory.Create(
                 cosmosException.StatusCode,
-                cosmosException.SubStatusCode,
                 cosmosException.Message,
                 exceptionWithStackTrace.StackTrace,
-                cosmosException.ActivityId,
-                cosmosException.RequestCharge,
-                cosmosException.RetryAfter,
-                cosmosException.Headers,
+                headers: cosmosException.Headers,
                 cosmosException.Trace,
                 cosmosException.Error,
                 cosmosException.InnerException);
@@ -103,15 +98,18 @@ private QueryExceptionConverter()
 
             public override CosmosException Visit(MalformedContinuationTokenException malformedContinuationTokenException) => CosmosExceptionFactory.CreateBadRequestException(
                     message: malformedContinuationTokenException.Message,
+                    headers: new Headers(),
                     stackTrace: malformedContinuationTokenException.StackTrace,
                     innerException: malformedContinuationTokenException);
 
             public override CosmosException Visit(UnexpectedQueryPartitionProviderException unexpectedQueryPartitionProviderException) => CosmosExceptionFactory.CreateInternalServerErrorException(
                     message: $"{nameof(CosmosException)} due to {nameof(UnexpectedQueryPartitionProviderException)}",
+                    headers: new Headers(),
                     innerException: unexpectedQueryPartitionProviderException);
 
             public override CosmosException Visit(ExpectedQueryPartitionProviderException expectedQueryPartitionProviderException) => CosmosExceptionFactory.CreateBadRequestException(
                     message: expectedQueryPartitionProviderException.Message,
+                    headers: new Headers(),
                     stackTrace: expectedQueryPartitionProviderException.StackTrace,
                     innerException: expectedQueryPartitionProviderException);
         }
@@ -127,6 +125,7 @@ private ChangeFeedExceptionConverter()
             internal override CosmosException Visit(
                 MalformedChangeFeedContinuationTokenException malformedChangeFeedContinuationTokenException) => CosmosExceptionFactory.CreateBadRequestException(
                     message: malformedChangeFeedContinuationTokenException.Message,
+                    headers: new Headers(),
                     stackTrace: malformedChangeFeedContinuationTokenException.StackTrace,
                     innerException: malformedChangeFeedContinuationTokenException);
         }

Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs

@@ -75,7 +75,9 @@ public static async Task<PartitionedQueryExecutionInfo> GetQueryPlanWithServiceI
 
                     throw CosmosExceptionFactory.CreateBadRequestException(
                         message: tryGetQueryPlan.Exception.ToString(),
-                        stackTrace: tryGetQueryPlan.Exception.StackTrace);
+                        headers: new Headers(),
+                        stackTrace: tryGetQueryPlan.Exception.StackTrace,
+                        trace: trace);
                 }
 
                 return tryGetQueryPlan.Result;

Microsoft.Azure.Cosmos/src/Query/Core/QueryResponseFactory.cs

@@ -46,13 +46,13 @@ public static QueryResponseCore CreateFromException(Exception exception)
                 else
                 {
                     CosmosException unkownCosmosException = CosmosExceptionFactory.CreateInternalServerErrorException(
-                        subStatusCode: default,
+                        headers: new Headers()
+                        {
+                            SubStatusCode = SubStatusCodes.PartitionKeyRangeGone,
+                            ActivityId = QueryResponseCore.EmptyGuidString
+                        },
                         message: exception.Message,
                         stackTrace: exception.StackTrace,
-                        activityId: QueryResponseCore.EmptyGuidString,
-                        requestCharge: 0,
-                        retryAfter: null,
-                        headers: null,
                         trace: NoOpTrace.Singleton,
                         innerException: exception);
 
@@ -113,17 +113,13 @@ private static QueryResponseCore CreateFromExceptionWithStackTrace(ExceptionWith
                 statusCode: queryResponseCore.StatusCode,
                 subStatusCodes: queryResponseCore.SubStatusCode,
                 cosmosException: CosmosExceptionFactory.Create(
-                    cosmosException.StatusCode,
-                    cosmosException.SubStatusCode,
-                    cosmosException.Message,
-                    exceptionWithStackTrace.StackTrace,
-                    cosmosException.ActivityId,
-                    cosmosException.RequestCharge,
-                    cosmosException.RetryAfter,
-                    cosmosException.Headers,
-                    cosmosException.Trace,
-                    cosmosException.Error,
-                    cosmosException.InnerException),
+                    statusCode: cosmosException.StatusCode,
+                    message: cosmosException.Message,
+                    stackTrace: exceptionWithStackTrace.StackTrace,
+                    headers: cosmosException.Headers,
+                    trace: cosmosException.Trace,
+                    error: cosmosException.Error,
+                    innerException: cosmosException.InnerException),
                 requestCharge: queryResponseCore.RequestCharge,
                 activityId: queryResponseCore.ActivityId);
 
@@ -142,6 +138,7 @@ public override CosmosException Visit(MalformedContinuationTokenException malfor
             {
                 return CosmosExceptionFactory.CreateBadRequestException(
                     message: malformedContinuationTokenException.Message,
+                    headers: new Headers(),
                     stackTrace: malformedContinuationTokenException.StackTrace,
                     innerException: malformedContinuationTokenException);
             }
@@ -150,13 +147,15 @@ public override CosmosException Visit(UnexpectedQueryPartitionProviderException
             {
                 return CosmosExceptionFactory.CreateInternalServerErrorException(
                     message: $"{nameof(CosmosException)} due to {nameof(UnexpectedQueryPartitionProviderException)}",
+                    headers: new Headers(),
                     innerException: unexpectedQueryPartitionProviderException);
             }
 
             public override CosmosException Visit(ExpectedQueryPartitionProviderException expectedQueryPartitionProviderException)
             {
                 return CosmosExceptionFactory.CreateBadRequestException(
                     message: expectedQueryPartitionProviderException.Message,
+                    headers: new Headers(),
                     stackTrace: expectedQueryPartitionProviderException.StackTrace,
                     innerException: expectedQueryPartitionProviderException);
             }

Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.cs

@@ -286,6 +286,7 @@ public async Task<IReadOnlyList<FeedRange>> GetFeedRangesAsync(
                 {
                     throw CosmosExceptionFactory.CreateInternalServerErrorException(
                         $"Container rid {containerRId} did not have a partition key range after refresh",
+                        headers: new Headers(),
                         trace: trace);
                 }
 
@@ -299,6 +300,7 @@ public async Task<IReadOnlyList<FeedRange>> GetFeedRangesAsync(
                 {
                     throw CosmosExceptionFactory.CreateInternalServerErrorException(
                         $"Container rid {containerRId} returned partitionKeyRanges null after Container RID refresh",
+                        headers: new Headers(),
                         trace: trace);
                 }
             }

Microsoft.Azure.Cosmos/src/Resource/CosmosExceptions/CosmosException.cs

@@ -21,29 +21,20 @@ public class CosmosException : Exception
         internal CosmosException(
             HttpStatusCode statusCode,
             string message,
-            int subStatusCode,
             string stackTrace,
-            string activityId,
-            double requestCharge,
-            TimeSpan? retryAfter,
             Headers headers,
             ITrace trace,
             Error error,
             Exception innerException)
             : base(CosmosException.GetMessageHelper(
                 statusCode,
-                subStatusCode,
-                message,
-                activityId), innerException)
+                headers,
+                message), innerException)
         {
             this.ResponseBody = message;
             this.stackTrace = stackTrace;
-            this.ActivityId = activityId;
             this.StatusCode = statusCode;
-            this.SubStatusCode = subStatusCode;
-            this.RetryAfter = retryAfter;
-            this.RequestCharge = requestCharge;
-            this.Headers = headers;
+            this.Headers = headers ?? new Headers();
             this.Error = error;
             this.Trace = trace;
         }
@@ -65,12 +56,19 @@ public CosmosException(
             : base(message)
         {
             this.stackTrace = null;
-            this.SubStatusCode = subStatusCode;
             this.StatusCode = statusCode;
-            this.RequestCharge = requestCharge;
-            this.ActivityId = activityId;
-            this.Headers = new Headers();
+            this.ResponseBody = message;
             this.Trace = NoOpTrace.Singleton;
+            this.Headers = new Headers()
+            {
+                SubStatusCode = (SubStatusCodes)subStatusCode,
+                RequestCharge = requestCharge,
+            };
+
+            if (!string.IsNullOrEmpty(activityId))
+            {
+                this.Headers.ActivityId = activityId;
+            }
         }
 
         /// <summary>
@@ -88,28 +86,28 @@ public CosmosException(
         /// Gets the request completion sub status code from the Azure Cosmos DB service.
         /// </summary>
         /// <value>The request completion status code</value>
-        public virtual int SubStatusCode { get; }
+        public virtual int SubStatusCode => Headers.GetIntValueOrDefault(this.Headers.SubStatusCodeLiteral);
 
         /// <summary>
         /// Gets the request charge for this request from the Azure Cosmos DB service.
         /// </summary>
         /// <value>
         /// The request charge measured in request units.
         /// </value>
-        public virtual double RequestCharge { get; }
+        public virtual double RequestCharge => this.Headers.RequestCharge;
 
         /// <summary>
         /// Gets the activity ID for the request from the Azure Cosmos DB service.
         /// </summary>
         /// <value>
         /// The activity ID for the request.
         /// </value>
-        public virtual string ActivityId { get; }
+        public virtual string ActivityId => this.Headers.ActivityId;
 
         /// <summary>
         /// Gets the retry after time. This tells how long a request should wait before doing a retry.
         /// </summary>
-        public virtual TimeSpan? RetryAfter { get; }
+        public virtual TimeSpan? RetryAfter => this.Headers.RetryAfter;
 
         /// <summary>
         /// Gets the response headers
@@ -184,28 +182,23 @@ internal ResponseMessage ToCosmosResponseMessage(RequestMessage request)
                  cosmosException: this,
                  statusCode: this.StatusCode,
                  trace: this.Trace);
-            if (this.SubStatusCode != 0)
-            {
-                responseMessage.Headers.SubStatusCode = (SubStatusCodes)this.SubStatusCode;
-            }
 
             return responseMessage;
         }
 
         private static string GetMessageHelper(
             HttpStatusCode statusCode,
-            int subStatusCode,
-            string responseBody,
-            string activityId)
+            Headers headers,
+            string responseBody)
         {
             StringBuilder stringBuilder = new StringBuilder();
 
             stringBuilder.Append($"Response status code does not indicate success: ");
             stringBuilder.Append($"{statusCode} ({(int)statusCode})");
             stringBuilder.Append("; Substatus: ");
-            stringBuilder.Append(subStatusCode);
+            stringBuilder.Append(headers?.SubStatusCodeLiteral ?? "0" );
             stringBuilder.Append("; ActivityId: ");
-            stringBuilder.Append(activityId ?? string.Empty);
+            stringBuilder.Append(headers?.ActivityId ?? string.Empty);
             stringBuilder.Append("; Reason: (");
             stringBuilder.Append(responseBody ?? string.Empty);
             stringBuilder.Append(");");