Improve Cosmos perf diagnostics

sdk/cosmos/.cspell.json

@@ -4,6 +4,7 @@
   ],
   "ignoreWords": [
     "accountname",
+    "acked",
     "activityid",
     "ALPN",
     "apacsoutheast",
@@ -262,6 +263,7 @@
     "RUPM",
     "rwcache",
     "sbyte",
+    "schemaversion",
     "serviceunavailable",
     "serviceversion",
     "sess",

sdk/cosmos/azure_data_cosmos/CHANGELOG.md

@@ -10,6 +10,8 @@
 
 ### Other Changes
 
+- Existing `ResponseHeaders` accessors now return Gateway 2.0 backend duration, quota, item-count, and local-LSN response metadata. ([#4797](https://github.com/Azure/azure-sdk-for-rust/pull/4797))
+
 ## 0.37.0 (2026-07-20)
 
 ### Features Added

sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/end_to_end.rs

@@ -546,53 +546,59 @@ async fn sdk_query_metadata_databases_and_containers() {
     let db_query = Query::from("SELECT * FROM c WHERE c.id = @id")
         .with_parameter("@id", db_name.as_str())
         .unwrap();
-    let emu_databases: Vec<DatabaseProperties> = backend
-        .emulator_client
-        .query_databases(db_query.clone(), None)
-        .await
-        .unwrap()
-        .try_collect()
-        .await
-        .unwrap();
+    let emu_databases: Vec<DatabaseProperties> = Box::pin(
+        backend
+            .emulator_client
+            .query_databases(db_query.clone(), None),
+    )
+    .await
+    .unwrap()
+    .try_collect()
+    .await
+    .unwrap();
     assert_eq!(emu_databases.len(), 1);
     assert_eq!(emu_databases[0].id.as_deref(), Some(db_name.as_str()));
 
     if let Some(ref real_client) = backend.real_client {
-        let real_databases: Vec<DatabaseProperties> = real_client
-            .query_databases(db_query, None)
-            .await
-            .unwrap()
-            .try_collect()
-            .await
-            .unwrap();
+        let real_databases: Vec<DatabaseProperties> =
+            Box::pin(real_client.query_databases(db_query, None))
+                .await
+                .unwrap()
+                .try_collect()
+                .await
+                .unwrap();
         assert_eq!(real_databases.len(), emu_databases.len());
         assert_eq!(real_databases[0].id, emu_databases[0].id);
     }
 
     let container_query = Query::from("SELECT * FROM c WHERE c.id = @id")
         .with_parameter("@id", "testcoll")
         .unwrap();
-    let emu_containers: Vec<ContainerProperties> = backend
-        .emulator_client
-        .database_client(&db_name)
-        .query_containers(container_query.clone(), None)
+    let emu_containers: Vec<ContainerProperties> = Box::pin(
+        backend
+            .emulator_client
+            .database_client(&db_name)
+            .query_containers(container_query.clone(), None),
+    )
+    .await
+    .unwrap()
+    .try_collect()
+    .await
+    .unwrap();
+    assert_eq!(emu_containers.len(), 1);
+    assert_eq!(emu_containers[0].id, "testcoll");
+
+    if let Some(ref real_client) = backend.real_client {
+        let real_containers: Vec<ContainerProperties> = Box::pin(
+            real_client
+                .database_client(&db_name)
+                .query_containers(container_query, None),
+        )
         .await
         .unwrap()
         .try_collect()
         .await
         .unwrap();
-    assert_eq!(emu_containers.len(), 1);
-    assert_eq!(emu_containers[0].id, "testcoll");
-
-    if let Some(ref real_client) = backend.real_client {
-        let real_containers: Vec<ContainerProperties> = real_client
-            .database_client(&db_name)
-            .query_containers(container_query, None)
-            .await
-            .unwrap()
-            .try_collect()
-            .await
-            .unwrap();
         assert_eq!(real_containers.len(), emu_containers.len());
         assert_eq!(real_containers[0].id, emu_containers[0].id);
     }
@@ -1085,25 +1091,25 @@ async fn sdk_query_items_with_filter_and_projection() {
             .unwrap()
     }
 
-    let emu_items: Vec<QueryTestItem> = emu_container
-        .query_items(query(), FeedScope::partition("pk1"), None)
-        .await
-        .unwrap()
-        .try_collect()
-        .await
-        .unwrap();
-    assert_eq!(emu_items.len(), 2);
-    assert_eq!(emu_items[0].id, "query-1");
-    assert_eq!(emu_items[1].id, "query-2");
-
-    if let Some(ref real) = real_container {
-        let real_items: Vec<QueryTestItem> = real
-            .query_items(query(), FeedScope::partition("pk1"), None)
+    let emu_items: Vec<QueryTestItem> =
+        Box::pin(emu_container.query_items(query(), FeedScope::partition("pk1"), None))
             .await
             .unwrap()
             .try_collect()
             .await
             .unwrap();
+    assert_eq!(emu_items.len(), 2);
+    assert_eq!(emu_items[0].id, "query-1");
+    assert_eq!(emu_items[1].id, "query-2");
+
+    if let Some(ref real) = real_container {
+        let real_items: Vec<QueryTestItem> =
+            Box::pin(real.query_items(query(), FeedScope::partition("pk1"), None))
+                .await
+                .unwrap()
+                .try_collect()
+                .await
+                .unwrap();
         assert_eq!(real_items.len(), emu_items.len());
         assert_eq!(
             real_items.iter().map(|i| &i.id).collect::<Vec<_>>(),

sdk/cosmos/azure_data_cosmos/tests/multi_write_tests/cosmos_multi_write_fault_injection.rs

@@ -9,7 +9,10 @@ use azure_data_cosmos::fault_injection::{
     FaultInjectionRuleBuilder, FaultOperationType,
 };
 use azure_data_cosmos::models::{ContainerProperties, ThroughputProperties};
-use azure_data_cosmos::options::{ExcludedRegions, ItemReadOptions, OperationOptions};
+use azure_data_cosmos::options::{
+    ExcludedRegions, ItemReadOptions, OperationOptions, OperationOptionsBuilder,
+    ThrottlingRetryOptionsBuilder,
+};
 use framework::{
     assert_local_retry_attempted_on_region, assert_region_contacted_with_retry,
     assert_region_not_contacted, TestClient, TestOptions, HUB_REGION, SATELLITE_REGION,
@@ -88,9 +91,20 @@ async fn verify_read_fails_with_injected_error(
                 .expect("fault client should be available");
             let fault_db_client = fault_client.database_client(db_client.id());
             let fault_container_client = fault_db_client.container_client(&container_id).await?;
+            let read_options = (expected_status == StatusCode::TooManyRequests).then(|| {
+                ItemReadOptions::default().with_operation_options(
+                    OperationOptionsBuilder::new()
+                        .with_throttling_retry_options(
+                            ThrottlingRetryOptionsBuilder::new()
+                                .with_max_retry_count(0)
+                                .build(),
+                        )
+                        .build(),
+                )
+            });
 
             let result = run_context
-                .read_item(&fault_container_client, &pk, &item_id, None)
+                .read_item(&fault_container_client, &pk, &item_id, read_options)
                 .await;
 
             let err = result.expect_err(&format!(

sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md

@@ -10,6 +10,8 @@
 
 ### Other Changes
 
+- Gateway 2.0 responses now preserve backend duration, quota, item-count, quorum, replica, query, and physical-partition RNTBD metadata when converting to standard Cosmos response headers. ([#4797](https://github.com/Azure/azure-sdk-for-rust/pull/4797))
+
 ## 0.6.0 (2026-07-20)
 
 ### Features Added

sdk/cosmos/azure_data_cosmos_driver/src/driver/transport/gateway_v2_dispatch.rs

@@ -410,6 +410,84 @@ pub(crate) fn unwrap_response_for_gateway_v2(
     if let Some(charge) = response.request_charge {
         headers.insert(response_header_names::REQUEST_CHARGE, charge.to_string());
     }
+    if let Some(duration_ms) = response.backend_request_duration_ms {
+        headers.insert(
+            response_header_names::SERVER_DURATION_MS,
+            duration_ms.to_string(),
+        );
+    }
+    if let Some(value) = response.last_state_change_date_time {
+        headers.insert(response_header_names::LAST_STATE_CHANGE_UTC, value);
+    }
+    if let Some(value) = response.storage_max_resource_quota {
+        headers.insert(response_header_names::RESOURCE_QUOTA, value);
+    }
+    if let Some(value) = response.storage_resource_quota_usage {
+        headers.insert(response_header_names::RESOURCE_USAGE, value);
+    }
+    if let Some(value) = response.schema_version {
+        headers.insert(response_header_names::SCHEMA_VERSION, value);
+    }
+    if let Some(value) = response.item_count {
+        headers.insert(response_header_names::ITEM_COUNT, value.to_string());
+    }
+    if let Some(value) = response.owner_id {
+        headers.insert(response_header_names::OWNER_ID, value);
+    }
+    if let Some(value) = response.quorum_acked_lsn {
+        headers.insert(response_header_names::QUORUM_ACKED_LSN, value.to_string());
+    }
+    if let Some(value) = response.current_write_quorum {
+        headers.insert(
+            response_header_names::CURRENT_WRITE_QUORUM,
+            value.to_string(),
+        );
+    }
+    if let Some(value) = response.current_replica_set_size {
+        headers.insert(
+            response_header_names::CURRENT_REPLICA_SET_SIZE,
+            value.to_string(),
+        );
+    }
+    if let Some(value) = response.xp_role {
+        headers.insert(response_header_names::XP_ROLE, value.to_string());
+    }
+    if let Some(value) = response.number_of_read_regions {
+        headers.insert(
+            response_header_names::NUMBER_OF_READ_REGIONS,
+            value.to_string(),
+        );
+    }
+    if let Some(value) = response.local_lsn.filter(|value| *value >= 0) {
+        headers.insert(response_header_names::LOCAL_LSN, value.to_string());
+    }
+    if let Some(value) = response.quorum_acked_local_lsn {
+        headers.insert(
+            response_header_names::QUORUM_ACKED_LOCAL_LSN,
+            value.to_string(),
+        );
+    }
+    if let Some(value) = response.item_local_lsn.filter(|value| *value >= 0) {
+        headers.insert(response_header_names::ITEM_LOCAL_LSN, value.to_string());
+    }
+    if let Some(value) = response.query_execution_info {
+        headers.insert(response_header_names::QUERY_EXECUTION_INFO, value);
+    }
+    if let Some(value) = response.pending_pk_delete {
+        headers.insert(
+            response_header_names::PENDING_PK_DELETE,
+            if value { "True" } else { "False" },
+        );
+    }
+    if let Some(value) = response.physical_partition_id {
+        headers.insert(response_header_names::PHYSICAL_PARTITION_ID, value);
+    }
+    if let Some(value) = response.conflict_resolved_timestamp {
+        headers.insert(
+            response_header_names::CONFLICT_RESOLVED_TIMESTAMP,
+            value.to_string(),
+        );
+    }
     if let Some(token) = response.session_token {
         // GW2 surfaces only the vector portion of the session token, with the
         // partition key range id carried separately. Classic gateway emits
@@ -1933,8 +2011,27 @@ mod tests {
                 404,
                 activity_id,
                 |tokens| {
+                    write_small_string_token(tokens, 0x0002, "2026-07-21T00:00:00Z");
+                    write_string_token(tokens, 0x000E, "documentSize=10240;");
+                    write_string_token(tokens, 0x000F, "documentSize=1;");
+                    write_small_string_token(tokens, 0x0010, "1.0");
                     write_u32_token(tokens, 0x001C, 1002);
+                    write_u32_token(tokens, 0x0014, 1);
                     write_double_token(tokens, 0x0015, 3.5);
+                    write_string_token(tokens, 0x0018, "owner-rid");
+                    write_i64_token(tokens, 0x001A, 40);
+                    write_u32_token(tokens, 0x001E, 3);
+                    write_u32_token(tokens, 0x001F, 4);
+                    write_u32_token(tokens, 0x0026, 1);
+                    write_u32_token(tokens, 0x0030, 2);
+                    write_i64_token(tokens, 0x003A, 41);
+                    write_i64_token(tokens, 0x003B, 40);
+                    write_i64_token(tokens, 0x003C, 39);
+                    write_string_token(tokens, 0x0045, "{\"reverseRidEnabled\":false}");
+                    write_double_token(tokens, 0x0051, 12.75);
+                    write_byte_token(tokens, 0x0055, 0);
+                    write_string_token(tokens, 0x0063, "physical-0");
+                    write_u64_token(tokens, 0x0087, 1_234_567);
                     write_string_token(tokens, 0x003E, "1:2#3");
                     write_string_token(tokens, 0x0004, "\"etag\"");
                     write_string_token(tokens, 0x0003, "continuation");
@@ -1967,6 +2064,47 @@ mod tests {
                 .get_optional_str(&HeaderName::from_static("x-ms-request-charge")),
             Some("3.5")
         );
+        assert_eq!(
+            unwrapped.headers.get_optional_str(&HeaderName::from_static(
+                response_header_names::SERVER_DURATION_MS
+            )),
+            Some("12.75")
+        );
+        let parsed_headers = crate::models::CosmosResponseHeaders::from_headers(&unwrapped.headers);
+        assert_eq!(parsed_headers.server_duration_ms, Some(12.75));
+        assert_eq!(
+            parsed_headers.last_state_change_utc.as_deref(),
+            Some("2026-07-21T00:00:00Z")
+        );
+        assert_eq!(
+            parsed_headers.resource_quota.as_deref(),
+            Some("documentSize=10240;")
+        );
+        assert_eq!(
+            parsed_headers.resource_usage.as_deref(),
+            Some("documentSize=1;")
+        );
+        assert_eq!(parsed_headers.schema_version.as_deref(), Some("1.0"));
+        assert_eq!(parsed_headers.item_count, Some(1));
+        assert_eq!(parsed_headers.owner_id.as_deref(), Some("owner-rid"));
+        assert_eq!(parsed_headers.quorum_acked_lsn, Some(40));
+        assert_eq!(parsed_headers.current_write_quorum, Some(3));
+        assert_eq!(parsed_headers.current_replica_set_size, Some(4));
+        assert_eq!(parsed_headers.xp_role, Some(1));
+        assert_eq!(parsed_headers.number_of_read_regions, Some(2));
+        assert_eq!(parsed_headers.local_lsn, Some(41));
+        assert_eq!(parsed_headers.quorum_acked_local_lsn, Some(40));
+        assert_eq!(parsed_headers.item_local_lsn, Some(39));
+        assert_eq!(
+            parsed_headers.query_execution_info.as_deref(),
+            Some("{\"reverseRidEnabled\":false}")
+        );
+        assert_eq!(parsed_headers.pending_pk_delete, Some(false));
+        assert_eq!(
+            parsed_headers.physical_partition_id.as_deref(),
+            Some("physical-0")
+        );
+        assert_eq!(parsed_headers.conflict_resolved_timestamp, Some(1_234_567));
         assert_eq!(
             unwrapped
                 .headers
@@ -2142,6 +2280,19 @@ mod tests {
         bytes.extend_from_slice(value.as_bytes());
     }
 
+    fn write_small_string_token(bytes: &mut Vec<u8>, id: u16, value: &str) {
+        bytes.extend_from_slice(&id.to_le_bytes());
+        bytes.push(0x07);
+        bytes.push(u8::try_from(value.len()).unwrap());
+        bytes.extend_from_slice(value.as_bytes());
+    }
+
+    fn write_byte_token(bytes: &mut Vec<u8>, id: u16, value: u8) {
+        bytes.extend_from_slice(&id.to_le_bytes());
+        bytes.push(0x00);
+        bytes.push(value);
+    }
+
     fn write_u32_token(bytes: &mut Vec<u8>, id: u16, value: u32) {
         bytes.extend_from_slice(&id.to_le_bytes());
         bytes.push(0x02);
@@ -2154,6 +2305,12 @@ mod tests {
         bytes.extend_from_slice(&value.to_le_bytes());
     }
 
+    fn write_u64_token(bytes: &mut Vec<u8>, id: u16, value: u64) {
+        bytes.extend_from_slice(&id.to_le_bytes());
+        bytes.push(0x04);
+        bytes.extend_from_slice(&value.to_le_bytes());
+    }
+
     fn write_double_token(bytes: &mut Vec<u8>, id: u16, value: f64) {
         bytes.extend_from_slice(&id.to_le_bytes());
         bytes.push(0x0E);