Emit x-ms-cosmos-hub-region-processing-only header (#4389)

Implements
[`HUB_REGION_PROCESSING_HEADER_SPEC.md`](https://github.com/Azure/azure-sdk-for-rust/blob/release/azure_data_cosmos-previews/sdk/cosmos/azure_data_cosmos/docs/HUB_REGION_PROCESSING_HEADER_SPEC.md)
(merged in #4320) for single-master data-plane Cosmos operations.

On the retry triggered by the first `404 / 1002
(READ_SESSION_NOT_AVAILABLE)` response, the driver now latches a
per-operation flag and emits the
`x-ms-cosmos-hub-region-processing-only: True` request header on every
subsequent transport attempt for that operation. This asks the backend
to route only to a region that has caught up to the requested LSN,
reducing the chance of a follow-up retry hitting another region whose
session is also behind.

The latch is sticky for the remainder of the operation and gated on
three conditions, which together scope the change exactly to the case
the spec calls out:

- data-plane scope (metadata-pipeline ops are out of scope per spec
§1.5)
- single-master accounts (multi-master accounts already have a different
recovery path)
- first 1002 within the operation (subsequent 1002s preserve the latch
but don't re-arm it)

Mirrors .NET parity
([Azure/azure-cosmos-dotnet-v3#5447](Azure/azure-cosmos-dotnet-v3#5447)).
Forward-compatible by design: backends that ignore the header preserve
the existing single-master 1002 retry behavior, so this lands safely
ahead of the server-side rollout.

## Changes by file

- **`cosmos_headers.rs`** — adds `HUB_REGION_PROCESSING_ONLY` const
inside `request_header_names`.
- **`components.rs`** — adds `is_dataplane: bool` and
`hub_region_processing_only: bool` fields to `OperationRetryState` with
LOAD-BEARING doc comments; `initial()` defaults both to `false`;
invariant note added to `session_token_retry_count` because the latch
trigger reads it pre-increment. `advance_*` propagates the flags
automatically via `..self` — no new positional parameters.
- **`operation_pipeline.rs`** — sets `retry_state.is_dataplane =
pipeline_type.is_data_plane()` once at the production call site (uses
the accessor because `PipelineType` is `#[non_exhaustive]` — a future
variant must not silently bypass an equality gate); extracts
`apply_hub_region_header(..)` as a small free function so the emission
rule can be tested without driving the full pipeline; updates 5 existing
exhaustive `OperationRetryState` struct-literal tests for the two new
fields.
- **`retry_evaluation.rs`** — extracts `build_session_retry_state(..)`
that latches `hub_region_processing_only` when all four trigger
conditions hold; gate placed after the existing `>= 2` abort check so it
only runs on the SessionRetry path; updates the one exhaustive
struct-literal test for the two new fields.
- **`azure_data_cosmos_driver/CHANGELOG.md`** and
**`azure_data_cosmos/CHANGELOG.md`** — entries under the unreleased
version with a back-link to #4320 and the spec doc.

## Tests added (10)

Trigger-side coverage in `retry_evaluation::tests` (8):
single-master/data-plane first 1002 latches (AC-1); multi-master never
latches (AC-4); latch stays set across subsequent 1002s (AC-2);
200/410/503 never latch (AC-5); count boundary at exhaustion (AC-3);
metadata-pipeline never latches (AC-8); independent operations don't
share state (AC-6); latch survives failover via `..self` propagation.

Emission-side coverage in `operation_pipeline::tests` (2): header set to
`True` when latched; header omitted when not latched.

All 876 lib tests + 57 in-memory-emulator tests + 40 emulator tests
pass; `cargo fmt` and `cargo clippy --all-features --tests` are clean.

## Acceptance criteria

AC-1 through AC-6 and AC-8 from the spec are covered by the tests above.
AC-7 (long-tail recovery / per-region telemetry) is deferred per the
spec's ALT-5 — no client-side scope.

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: NaluTripician

sdk/cosmos/azure_data_cosmos/CHANGELOG.md

@@ -5,12 +5,10 @@
 ### Features Added
 
 ### Breaking Changes
-- Removed the `request_url()` accessor (gated on the `fault_injection` feature) from `ItemResponse`/`ResourceResponse`/`BatchResponse`. Driver-routed operations never populated it, so it always returned `None` in current usage.
 
+- Removed the `request_url()` accessor (gated on the `fault_injection` feature) from `ItemResponse`/`ResourceResponse`/`BatchResponse`. Driver-routed operations never populated it, so it always returned `None` in current usage.
 - `CosmosClientBuilder::with_user_agent_suffix` (and `CosmosClientOptions::with_user_agent_suffix`) now take `UserAgentSuffix` instead of `impl Into<String>`. Callers passing a `&str` or `String` must construct the value explicitly via `UserAgentSuffix::new` (panics on invalid input) or `UserAgentSuffix::try_new` (returns `Option`). Validation rules (max 25 characters, HTTP-header-safe) are now enforced at the construction site instead of being applied silently inside the builder. ([#4368](https://github.com/Azure/azure-sdk-for-rust/pull/4368))
-
 - Replaced `CosmosDiagnostics` with `CosmosDiagnosticsContext` (a re-export of `azure_data_cosmos_driver::diagnostics::DiagnosticsContext`). All response types now return `Arc<CosmosDiagnosticsContext>` from `diagnostics()` (the returned `Arc` derefs transparently to `CosmosDiagnosticsContext` for read-only inspection, and can be retained alongside a consumed response body). The previous `activity_id() -> Option<&str>` and `server_duration_ms() -> Option<f64>` accessors on `CosmosDiagnostics` are replaced by `CosmosDiagnosticsContext::activity_id() -> &ActivityId` and per-request server timing via `CosmosDiagnosticsContext::requests()[i].server_duration_ms()`.
-
 - Removed `azure_data_cosmos::constants::SubStatusCode` and its `new`/`value`/`from_header_value`/`From`/`Display`/`Debug` API. The SDK no longer maintains a parallel sub-status-code type — fault-injection (the only remaining consumer) now uses `azure_data_cosmos_driver::models::SubStatusCode` directly. Callers that referenced the SDK type should switch to the driver re-export.
 
 ### Bugs Fixed

sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Features Added
 
+- Added support for the `x-ms-cosmos-hub-region-processing-only` request header on retries after a `404 / 1002 (READ_SESSION_NOT_AVAILABLE)` response on single-master data-plane Cosmos operations. The header asks the backend to route only to a region that has caught up to the requested LSN, reducing the chance of a follow-up retry hitting a region whose session is also behind. The header is scoped to single-master accounts (multi-master accounts already have a different recovery path) and to data-plane operations (metadata-pipeline operations are out of scope per the design spec). Once latched on the first 1002 within an operation, the header is emitted on every subsequent retry for that operation. ([#4389](https://github.com/Azure/azure-sdk-for-rust/pull/4389))
 - Added local query-plan generator scaffolding under `crate::query` (lexer, parser, AST, planner, and in-memory evaluator). The scaffolding is **not wired into the production query path** yet — production callers still issue Gateway query-plan requests via `CosmosOperation::query_plan`. The `__internal_testing` cargo feature exposes `query::__test_only_generate_query_plan_for_pk_paths`, `query::__TEST_ONLY_SUPPORTED_QUERY_FEATURES`, and `CosmosOperation::query_plan` for cross-crate gateway-comparison tests; this feature is intentionally unstable and **not covered by SemVer**.
 - Added per-partition automatic failover (PPAF) for writes on single-master accounts. On 403/3 WriteForbidden, 503 ServiceUnavailable, 429/3092 SystemResourceUnavailable, 410/1022 Gone, or 408 RequestTimeout from a region, the affected partition is failed over to the next preferred region; subsequent writes for that partition skip the failed region. ([#4156](https://github.com/Azure/azure-sdk-for-rust/pull/4156))
 - Added per-partition circuit breaker (PPCB) for reads (any account) and writes (multi-master accounts). Tracks failure counts per `(partition_key_range_id, region)` and routes to an alternate region once the threshold (default 10 reads, 5 writes) is exceeded. A background failback loop probes the original region for recovery. ([#4156](https://github.com/Azure/azure-sdk-for-rust/pull/4156))

sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/components.rs

@@ -75,13 +75,67 @@ pub(crate) struct OperationRetryState {
     /// Number of operation-level failover retries attempted.
     pub failover_retry_count: u32,
     /// Number of session-token retries attempted.
+    ///
+    /// **Invariant:** This counter is incremented only by
+    /// [`Self::advance_session_retry`] on the 404/1002 path. The
+    /// hub-region-processing-only latch trigger in
+    /// `retry_evaluation::try_handle_read_session_not_available` reads
+    /// `== 0` to detect the first 1002 within an operation; if a future
+    /// refactor adds another increment site, that trigger gate must be
+    /// re-validated.
     pub session_token_retry_count: u32,
     /// Maximum failover retries.
     pub max_failover_retries: u32,
     /// Maximum session retries.
     pub max_session_retries: u32,
     /// Whether multiple write locations can be used.
     pub can_use_multiple_write_locations: bool,
+    /// Whether this operation is on the data-plane pipeline (vs metadata).
+    ///
+    /// Set once at the production call site in `execute_operation_pipeline`
+    /// from `pipeline_type.is_data_plane()`. Used to gate the
+    /// hub-region-processing-only latch so metadata-pipeline operations
+    /// (which ride the same `execute_operation_pipeline` but are scoped out
+    /// of the spec per HUB_REGION_PROCESSING_HEADER_SPEC.md §1.5) never
+    /// emit the header.
+    ///
+    /// LOAD-BEARING for the metadata-pipeline scope gate (AC-8).
+    /// The production call site MUST use the
+    /// `PipelineType::is_data_plane()` accessor — NOT `==` matching —
+    /// because `PipelineType` is `#[non_exhaustive]` and a future variant
+    /// would silently bypass an equality gate.
+    pub is_dataplane: bool,
+    /// Hub-region-processing-only latch.
+    ///
+    /// Sticky within a single operation. Set on the retry triggered by the
+    /// FIRST `404 / 1002 (READ_SESSION_NOT_AVAILABLE)` on a single-master
+    /// data-plane account; once set, every subsequent transport attempt
+    /// for this operation emits the
+    /// `x-ms-cosmos-hub-region-processing-only: True` header.
+    ///
+    /// Bounded by operation lifetime — `OperationRetryState` is
+    /// constructed fresh per call to `execute_operation_pipeline` (single
+    /// production call site), so the latch never leaks across operations.
+    ///
+    /// LOAD-BEARING for SE-003 mitigation — see
+    /// HUB_REGION_PROCESSING_HEADER_SPEC.md AG-1..AG-4. If a future
+    /// refactor adds a second production construction site for
+    /// `OperationRetryState`, the SE-003 mitigation argument needs to be
+    /// re-validated.
+    ///
+    /// **Cross-region hedging coordination.** The orchestrator added in
+    /// `azure_data_cosmos_driver/docs/HEDGING_SPEC.md` §9.5 constructs an
+    /// `OperationRetryState` *per hedge*, so this per-state latch is
+    /// per-hedge by default. The hedging spec requires augmenting this
+    /// state with a `shared_hub_region_latch: Option<Arc<AtomicBool>>`
+    /// that is `Some` only when running inside `execute_with_hedging()`,
+    /// is CAS-set alongside this field in `build_session_retry_state`,
+    /// and is OR'd into the emission decision in `apply_hub_region_header`.
+    /// This mirrors .NET v3's `CrossRegionAvailabilityContext` shared
+    /// object introduced by azure-cosmos-dotnet-v3#5815. Any change to
+    /// the latch trigger or emission rule here MUST update both call
+    /// sites and §9.5 of the hedging spec.
+    pub hub_region_processing_only: bool,
     /// Regions excluded for this operation.
     pub excluded_regions: Vec<Region>,
     /// Session-retry routing override for read operations.
@@ -164,6 +218,8 @@ impl OperationRetryState {
             max_failover_retries,
             max_session_retries,
             can_use_multiple_write_locations,
+            is_dataplane: false,
+            hub_region_processing_only: false,
             excluded_regions,
             session_retry_routing: SessionRetryRouting::PreferredEndpoints,
             partition_key_range_id: None,

sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/operation_pipeline.rs

@@ -140,6 +140,18 @@ pub(crate) async fn execute_operation_pipeline(
         .per_partition_circuit_breaker_enabled
         && location_snapshot.account.preferred_read_endpoints.len() > 1;
 
+    // HUB_REGION_PROCESSING_HEADER_SPEC.md §1.5: gate the
+    // `x-ms-cosmos-hub-region-processing-only` latch on data-plane scope
+    // so metadata-pipeline operations (which ride the same
+    // `execute_operation_pipeline`) never emit the header.
+    //
+    // Use the `PipelineType::is_data_plane()` accessor — NOT `==` matching
+    // — because `PipelineType` is `#[non_exhaustive]` and a future variant
+    // would silently bypass an equality gate. Equivalently
+    // `!pipeline_type.is_metadata()` (the metadata pipeline is the only
+    // current variant that is out of spec scope).
+    retry_state.is_dataplane = pipeline_type.is_data_plane();
+
     let deadline = options
         .end_to_end_latency_policy()
         .map(|p| Instant::now() + p.timeout());
@@ -153,7 +165,7 @@ pub(crate) async fn execute_operation_pipeline(
             operation,
             &retry_state,
             &location,
-            pipeline_type == PipelineType::DataPlane,
+            pipeline_type.is_data_plane(),
             location_state_store.endpoint_unavailability_ttl(),
         );
 
@@ -185,6 +197,14 @@ pub(crate) async fn execute_operation_pipeline(
         };
         let mut transport_request = build_transport_request(operation, custom_headers, &ctx)?;
 
+        // HUB_REGION_PROCESSING_HEADER_SPEC.md §3 / public-spec §3.4:
+        // Emit the `x-ms-cosmos-hub-region-processing-only: True` header
+        // when the latch is set. The latch is flipped in
+        // `try_handle_read_session_not_available` on the first 1002 of a
+        // single-master data-plane operation, and is sticky for the
+        // remainder of the operation's transport attempts (AC-1, AC-2).
+        apply_hub_region_header(&mut transport_request, &retry_state);
+
         apply_optional_request_headers(&mut transport_request, operation, options);
 
         tracing::trace!(
@@ -923,6 +943,40 @@ fn compute_execution_context(retry_state: &OperationRetryState) -> ExecutionCont
     }
 }
 
+/// Conditionally emits the `x-ms-cosmos-hub-region-processing-only: True`
+/// header on the outbound transport request when the latch on
+/// `retry_state` is set.
+///
+/// Extracted as a free function (rather than left inline at the call site
+/// in `execute_operation_pipeline`) so that the emission rule can be
+/// exercised by unit tests without spinning up the full pipeline. The
+/// production call site is the loop iteration after `build_transport_request`
+/// and before `apply_optional_request_headers`.
+///
+/// HUB_REGION_PROCESSING_HEADER_SPEC.md §3 / public-spec §3.4. See
+/// `try_handle_read_session_not_available` for the latch trigger.
+///
+/// **Hedging coordination (future).** Per HEDGING_SPEC.md §9.5, the
+/// emission decision MUST OR in a `shared_hub_region_latch:
+/// Option<Arc<AtomicBool>>` field added to `OperationRetryState`, read
+/// with `Acquire` ordering. That field is set from `build_session_retry_state`
+/// the first time any hedge in the same `execute_with_hedging()`
+/// fan-out observes 1002 and is what makes the other (still-latch-clean)
+/// hedges immediately emit the header — the Rust counterpart of .NET v3's
+/// `CrossRegionAvailabilityContext.ShouldAddHubRegionProcessingOnlyHeader`
+/// from azure-cosmos-dotnet-v3#5815.
+fn apply_hub_region_header(
+    transport_request: &mut TransportRequest,
+    retry_state: &OperationRetryState,
+) {
+    if retry_state.hub_region_processing_only {
+        transport_request.headers.insert(
+            HeaderName::from_static(request_header_names::HUB_REGION_PROCESSING_ONLY),
+            HeaderValue::from_static("True"),
+        );
+    }
+}
+
 /// Applies operation-options-driven request headers that are only known
 /// after the request has been built by `build_transport_request`.
 ///
@@ -1333,6 +1387,8 @@ mod tests {
             max_failover_retries: 3,
             max_session_retries: 2,
             can_use_multiple_write_locations: false,
+            is_dataplane: false,
+            hub_region_processing_only: false,
             excluded_regions: Vec::new(),
             session_retry_routing:
                 crate::driver::pipeline::components::SessionRetryRouting::PreferredWriteEndpoints,
@@ -1387,6 +1443,8 @@ mod tests {
             max_failover_retries: 3,
             max_session_retries: 2,
             can_use_multiple_write_locations: false,
+            is_dataplane: false,
+            hub_region_processing_only: false,
             excluded_regions: Vec::new(),
             session_retry_routing:
                 crate::driver::pipeline::components::SessionRetryRouting::PreferredEndpoints,
@@ -1441,6 +1499,8 @@ mod tests {
             max_failover_retries: 3,
             max_session_retries: 2,
             can_use_multiple_write_locations: false,
+            is_dataplane: false,
+            hub_region_processing_only: false,
             excluded_regions: Vec::new(),
             session_retry_routing:
                 crate::driver::pipeline::components::SessionRetryRouting::PreferredEndpoints,
@@ -1502,6 +1562,8 @@ mod tests {
             max_failover_retries: 3,
             max_session_retries: 3,
             can_use_multiple_write_locations: true,
+            is_dataplane: false,
+            hub_region_processing_only: false,
             excluded_regions: Vec::new(),
             session_retry_routing:
                 crate::driver::pipeline::components::SessionRetryRouting::PreferredEndpoints,
@@ -1761,6 +1823,8 @@ mod tests {
             max_failover_retries: 3,
             max_session_retries: 2,
             can_use_multiple_write_locations: false,
+            is_dataplane: false,
+            hub_region_processing_only: false,
             excluded_regions: vec!["westus2".into()],
             session_retry_routing:
                 crate::driver::pipeline::components::SessionRetryRouting::PreferredEndpoints,
@@ -2736,6 +2800,65 @@ mod tests {
         ));
     }
 
+    // ── apply_hub_region_header ──────────────────────────────────────
+    //
+    // See HUB_REGION_PROCESSING_HEADER_SPEC.md §3.4 / public-spec §4.2.
+    // The emission logic itself is a 4-line conditional; these tests
+    // exercise both branches so AC-1/AC-5 don't drift on a refactor.
+
+    fn build_minimal_transport_request() -> super::TransportRequest {
+        let operation = CosmosOperation::read_all_databases(test_account());
+        let routing = test_routing();
+        let activity_id = ActivityId::from_string("hub-region-test".to_string());
+        let ctx = TransportRequestContext {
+            routing: &routing,
+            activity_id: &activity_id,
+            execution_context: ExecutionContext::Initial,
+            deadline: None,
+            resolved_session_token: None,
+            throughput_control: None,
+        };
+        build_transport_request(&operation, None, &ctx).expect("request should build")
+    }
+
+    /// T-6 — When the latch is set on `retry_state`, the helper emits
+    /// `x-ms-cosmos-hub-region-processing-only: True` on the transport
+    /// request (AC-1).
+    #[test]
+    fn transport_request_emits_hub_region_header_when_latched() {
+        let mut request = build_minimal_transport_request();
+        let mut state = super::OperationRetryState::initial(0, false, Vec::new(), 3, 1);
+        state.is_dataplane = true;
+        state.hub_region_processing_only = true;
+
+        super::apply_hub_region_header(&mut request, &state);
+
+        let value = request.headers.get_optional_str(&HeaderName::from_static(
+            request_header_names::HUB_REGION_PROCESSING_ONLY,
+        ));
+        assert_eq!(value, Some("True"));
+    }
+
+    /// T-7 — When the latch is NOT set, the helper does not emit the
+    /// header. Covers AC-5 / cross-operation isolation guarantee at the
+    /// emission layer.
+    #[test]
+    fn transport_request_omits_hub_region_header_when_not_latched() {
+        let mut request = build_minimal_transport_request();
+        let state = super::OperationRetryState::initial(0, false, Vec::new(), 3, 1);
+        assert!(!state.hub_region_processing_only);
+
+        super::apply_hub_region_header(&mut request, &state);
+
+        let value = request.headers.get_optional_str(&HeaderName::from_static(
+            request_header_names::HUB_REGION_PROCESSING_ONLY,
+        ));
+        assert!(
+            value.is_none(),
+            "hub-region header must not be present when latch is unset, got {value:?}",
+        );
+    }
+
     // ── apply_failover_delay ──────────────────────────────────────────
 
     #[tokio::test]