Add Gateway 2.0 implementation to Cosmos driver

sdk/cosmos/azure_data_cosmos/CHANGELOG.md

@@ -4,8 +4,12 @@
 
 ### Features Added
 
+- Added `CosmosClientBuilder::with_gateway20_disabled(bool)` to opt out of the new Gateway 2.0 transport, which is now enabled by default. Gateway 2.0 routes data-plane requests through a regional proxy that forwards RNTBD-over-HTTP/2 to the backend. Set this to `true` to fall back to the direct gateway transport — useful for workloads that depend on the published gateway latency SLAs (Gateway 2.0 is not currently covered by them) or that need the direct-gateway behavior for diagnostics. ([#4319](https://github.com/Azure/azure-sdk-for-rust/pull/4319))
+
 ### Breaking Changes
 
+- Consolidated SDK fault-injection types as re-exports from `azure_data_cosmos_driver::fault_injection`. `FaultInjectionRule`, `FaultInjectionCondition`, `FaultInjectionResult`, `CustomResponse`, `FaultInjectionErrorType`, `FaultOperationType`, and the matching builders are now provided by the driver crate. Field access is via accessor methods (e.g., `rule.id()`, `condition.region()`, `response.body()`) rather than direct field reads. The SDK retains only `FaultInjectionClientBuilder` (gateway-side transport wrapper). ([#4319](https://github.com/Azure/azure-sdk-for-rust/pull/4319))
+
 ### Bugs Fixed
 
 ### Other Changes

sdk/cosmos/azure_data_cosmos/build.rs

@@ -6,5 +6,7 @@
 // unknown cfg names are warned/denied unless explicitly declared via check-cfg.
 fn main() {
     // Allow `#[cfg_attr(not(test_category = "..."), ignore)]` in `tests/*.rs`.
-    println!("cargo:rustc-check-cfg=cfg(test_category, values(\"emulator\", \"multi_write\", \"split\"))");
+    println!(
+        "cargo:rustc-check-cfg=cfg(test_category, values(\"emulator\", \"multi_write\", \"split\", \"gateway20\"))"
+    );
 }

sdk/cosmos/azure_data_cosmos/docs/sdk-to-driver-cutover.md

@@ -315,42 +315,24 @@ The gateway pipeline tracked this via `CosmosRequest` (which held the final URL)
 
 ## Fault Injection Wiring
 
-When cutting `read_item` over to the driver, the SDK's fault injection tests initially failed because the two execution paths (gateway and driver) have **independent fault injection systems**. This section documents how they were connected.
+The SDK no longer ships a parallel fault-injection type system. All fault-injection types — [`FaultInjectionRule`], [`FaultInjectionCondition`], [`FaultInjectionResult`], [`CustomResponse`], [`FaultInjectionErrorType`], [`FaultOperationType`], and the matching builders — are re-exported directly from the driver crate (`azure_data_cosmos_driver::fault_injection`) by `azure_data_cosmos::fault_injection`. The SDK only owns:
 
-### Problem
+- [`FaultInjectionClientBuilder`] — produces the `azure_core::http::Transport` that the SDK pipeline plugs in (i.e., a `FaultClient` HTTP client wrapper that evaluates driver rules against in-flight gateway requests).
+- A small private `fault_operation_for_sdk(SdkOperationType, SdkResourceType) → Option<FaultOperationType>` adapter so `CosmosRequest::add_fault_injection_headers` can stamp the right operation tag on the outbound headers.
 
-The SDK and driver each have their own fault injection module (`azure_data_cosmos::fault_injection` and `azure_data_cosmos_driver::fault_injection`). They define parallel but separate types (`FaultInjectionRule`, `FaultInjectionCondition`, `FaultInjectionResult`, etc.) with identical variants but different Rust types. Prior to this work, only the gateway pipeline received fault injection rules — the driver was built without them.
-
-### Solution: Rule Translation with Shared State
-
-The bridge module (`driver_bridge.rs`) includes `sdk_fi_rules_to_driver_fi_rules()`, which translates SDK fault injection rules into driver fault injection rules. The translation covers:
-
-- `FaultOperationType` — variant-by-variant match (identical variant names)
-- `FaultInjectionErrorType` — variant-by-variant match
-- `FaultInjectionCondition` — `RegionName` → `Region`, operation type and container ID mapped directly
-- `FaultInjectionResult` — `Duration` → `Option<Duration>`, probability copied
-- Timing fields — `start_time: Instant` → `Option<Instant>`, `end_time` and `hit_limit` copied
-
-### Shared Mutable State
-
-SDK `FaultInjectionRule` has `enabled: Arc<AtomicBool>` and `hit_count: Arc<AtomicU32>` that tests mutate at runtime (`.disable()`, `.enable()`, `.hit_count()`). The driver's `FaultInjectionRuleBuilder` accepts external `Arc`s via `with_shared_state()`, so both the SDK gateway path and the driver path reference the **same atomic state**. This means:
-
-- Calling `.disable()` on the SDK rule also disables it in the driver
-- Hit counts are shared — both paths increment the same counter
-- Tests that toggle rules or assert hit counts work correctly across both paths
+Because both transports (gateway and driver) consume the **same** `Arc<FaultInjectionRule>` instances now, there is no translation step and no shared-state plumbing — toggling `enable()`/`disable()`, hit-count increments, and `hit_limit` enforcement all happen against one canonical rule object.
 
 ### Wiring in `CosmosClientBuilder`
 
 In `CosmosClientBuilder::build()`:
 
-1. Before the `FaultInjectionClientBuilder` is consumed for the gateway transport, `rules()` extracts a reference to the SDK rules
-2. `sdk_fi_rules_to_driver_fi_rules()` translates them to driver rules with shared state
-3. The translated rules are passed to `CosmosDriverRuntimeBuilder::with_fault_injection_rules()`
-4. The SDK's `fault_injection` Cargo feature now forwards to the driver's `fault_injection` feature
+1. The `FaultInjectionClientBuilder::rules()` accessor returns `&[Arc<FaultInjectionRule>]` — already the driver type, so the SDK simply clones the slice (`fault_builder.rules().to_vec()`).
+2. The cloned rules are passed to `CosmosDriverRuntimeBuilder::with_fault_injection_rules()` so the driver's own fault-injection HTTP client can evaluate them.
+3. The `FaultInjectionClientBuilder` is then consumed to build the gateway transport, which wraps the inner `HttpClient` with a `FaultClient` that evaluates the same rules.
 
 ### Test Patterns for Future Cutover
 
-When cutting over additional operations, **no additional fault injection wiring is needed** — it's handled once at the `CosmosClientBuilder` level. However, tests need to account for two behavioral differences:
+When cutting over additional operations, **no fault-injection wiring changes are needed** — it's all wired once at `CosmosClientBuilder::build()`. However, tests need to account for two behavioral differences between gateway-routed and driver-routed operations:
 
 **`request_url()` returns `None` for driver-routed operations:**
 
@@ -378,24 +360,7 @@ let rule = FaultInjectionRuleBuilder::new("test", error)
 
 This asymmetry will disappear once all operations are driver-routed, since there will be only one hit-counting path.
 
-### `custom_response` Translation
-
-Translation of `CustomResponse` (synthetic HTTP responses) is not yet implemented. None of the current tests use custom responses for `ReadItem` operations. When needed, the bridge function should be extended to translate `CustomResponse` fields (`status_code`, `headers`, `body`).
-
-### Consolidating to Driver Fault Injection After Cutover
-
-The current dual-system architecture (SDK fault injection + driver fault injection + translation bridge) exists only because the cutover is incremental — some operations still go through the gateway while others go through the driver. Once **all** operations are routed through the driver:
-
-1. **Drop `azure_data_cosmos::fault_injection`** — the SDK's HTTP-client-level fault interception module becomes unreachable. Delete the entire `src/fault_injection/` directory.
-2. **Re-export driver types** — the SDK re-exports the driver's fault injection types directly:
-
-   ```rust
-   #[cfg(feature = "fault_injection")]
-   pub use azure_data_cosmos_driver::fault_injection;
-   ```
+### Final State After Cutover
 
-3. **Remove the translation layer** — `sdk_fi_rules_to_driver_fi_rules()` in `driver_bridge.rs` and the `shared_enabled()`/`shared_hit_count()` accessors on the SDK rule are no longer needed.
-4. **Simplify `CosmosClientBuilder`** — `with_fault_injection()` accepts `Vec<Arc<driver::FaultInjectionRule>>` directly and passes them to `CosmosDriverRuntimeBuilder::with_fault_injection_rules()`. No translation, no cloning, no intermediary builder.
-5. **Update tests** — tests construct driver `FaultInjectionRule` directly (same builders, same API) instead of SDK rules.
+Once **all** operations are routed through the driver, the SDK-side `FaultInjectionClientBuilder` and `FaultClient` HTTP wrapper become unreachable too — the driver-runtime fault-injection HTTP client is the single source of truth. At that point `azure_data_cosmos::fault_injection` collapses into a pure `pub use azure_data_cosmos_driver::fault_injection;` re-export (or is dropped entirely).
 
-At that point the SDK has **no fault injection logic of its own** — it's a pass-through to the driver, matching the overall "SDK as thin wrapper" goal. The driver is the single source of truth for all transport-related concerns including fault injection.

sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client_builder.rs

@@ -93,6 +93,15 @@ pub struct CosmosClientBuilder {
     fault_injection_builder: Option<crate::fault_injection::FaultInjectionClientBuilder>,
     /// Fallback endpoints tried when the primary endpoint is unavailable.
     backup_endpoints: Vec<azure_core::http::Url>,
+    /// Operator override for the Gateway 2.0 transport.
+    ///
+    /// `None` (the default) leaves the underlying driver in charge of
+    /// routing — Gateway 2.0 is selected automatically whenever the
+    /// account advertises a Gateway 2.0 endpoint and HTTP/2 is allowed.
+    /// `Some(true)` forces every request through the standard gateway
+    /// transport via [`with_gateway20_disabled`](Self::with_gateway20_disabled);
+    /// `Some(false)` explicitly opts in (matching the default behaviour).
+    gateway20_disabled: Option<bool>,
 }
 
 impl CosmosClientBuilder {
@@ -168,6 +177,41 @@ impl CosmosClientBuilder {
         self
     }
 
+    /// Disables the Gateway 2.0 transport for this client.
+    ///
+    /// Gateway 2.0 is the next-generation Cosmos DB dataplane transport:
+    /// SDK connections terminate at a regional Gateway 2.0 proxy that
+    /// forwards RNTBD-over-HTTP/2 to the backend. **Gateway 2.0 is enabled
+    /// by default** — whenever the account advertises a Gateway 2.0 endpoint
+    /// the SDK routes eligible dataplane operations through it and falls
+    /// back to the standard gateway only for operations Gateway 2.0 cannot
+    /// serve (e.g. metadata requests or accounts that do not advertise a
+    /// Gateway 2.0 endpoint).
+    ///
+    /// Pass `true` to opt out and force every request through the standard
+    /// gateway transport. The standard gateway path remains supported and
+    /// stable — disabling Gateway 2.0 is the recommended workaround if you
+    /// hit a regression on the new transport.
+    ///
+    /// # Latency caveat
+    ///
+    /// Gateway 2.0 traffic flows through a proxy that is
+    /// **not currently covered by the regional Cosmos DB latency SLA**.
+    /// Workloads with strict P99 latency requirements should opt out via
+    /// `with_gateway20_disabled(true)` until the proxy reaches general
+    /// availability. The extra hop also means Gateway 2.0 may add measurable
+    /// latency relative to the standard gateway in some regions.
+    ///
+    /// # Arguments
+    ///
+    /// * `disabled` - `true` to suppress Gateway 2.0 and force the standard
+    ///   gateway transport; `false` (or leaving the builder untouched) keeps
+    ///   the default Gateway 2.0 behaviour.
+    pub fn with_gateway20_disabled(mut self, disabled: bool) -> Self {
+        self.gateway20_disabled = Some(disabled);
+        self
+    }
+
     /// Registers a throughput control group on the driver runtime.
     ///
     /// Groups define throughput policies (priority level, throughput bucket) that
@@ -287,9 +331,10 @@ impl CosmosClientBuilder {
             Option<azure_core::http::Transport>,
             Vec<std::sync::Arc<azure_data_cosmos_driver::fault_injection::FaultInjectionRule>>,
         ) = if let Some(fault_builder) = self.fault_injection_builder {
-            // Translate rules for the driver before the builder is consumed.
-            let driver_rules =
-                crate::driver_bridge::sdk_fi_rules_to_driver_fi_rules(fault_builder.rules());
+            // SDK fault-injection rules are now driver `FaultInjectionRule`s
+            // (re-exported through `crate::fault_injection`), so the driver
+            // can consume them directly without a translation step.
+            let driver_rules = fault_builder.rules().to_vec();
             let fault_builder = match base_client {
                 Some(client) => fault_builder.with_inner_client(client),
                 None => fault_builder,
@@ -425,6 +470,9 @@ impl CosmosClientBuilder {
                 EmulatorServerCertValidation::DangerousDisabled,
             );
         }
+        if let Some(disabled) = self.gateway20_disabled {
+            pool_builder = pool_builder.with_gateway20_disabled(disabled);
+        }
         driver_runtime_builder = driver_runtime_builder.with_connection_pool(pool_builder.build()?);
 
         #[cfg(feature = "fault_injection")]

sdk/cosmos/azure_data_cosmos/src/constants.rs

@@ -18,6 +18,8 @@ macro_rules! cosmos_headers {
         /// A list of all Cosmos DB specific headers that should be allowed in logging.
         pub const COSMOS_ALLOWED_HEADERS: &[&HeaderName] = &[
             $(&$name,)*
+            &azure_data_cosmos_driver::constants::GATEWAY20_OPERATION_TYPE,
+            &azure_data_cosmos_driver::constants::GATEWAY20_RESOURCE_TYPE,
         ];
     };
 }
@@ -185,9 +187,6 @@ cosmos_headers! {
     COSMOS_QUORUM_ACKED_LLSN => "x-ms-cosmos-quorum-acked-llsn",
     REQUEST_DURATION_MS => "x-ms-request-duration-ms",
     COSMOS_INTERNAL_PARTITION_ID => "x-ms-cosmos-internal-partition-id",
-    // Thin Client
-    THINCLIENT_PROXY_OPERATION_TYPE => "x-ms-thinclient-proxy-operation-type",
-    THINCLIENT_PROXY_RESOURCE_TYPE => "x-ms-thinclient-proxy-resource-type",
     // Client ID
     CLIENT_ID => "x-ms-client-id",
     // these are not actually sent but are used internally for fault injection

sdk/cosmos/azure_data_cosmos/src/cosmos_request.rs

@@ -2,7 +2,7 @@
 // Licensed under the MIT License.
 
 #[cfg(feature = "fault_injection")]
-use crate::fault_injection::FaultOperationType;
+use crate::fault_injection::fault_operation_for_sdk;
 use crate::operation_context::OperationType;
 use crate::options::ExcludedRegions;
 use crate::request_context::RequestContext;
@@ -153,10 +153,7 @@ impl CosmosRequest {
 
     #[cfg(feature = "fault_injection")]
     pub fn add_fault_injection_headers(&mut self) {
-        let fault_op = FaultOperationType::from_operation_and_resource(
-            &self.operation_type,
-            &self.resource_type,
-        );
+        let fault_op = fault_operation_for_sdk(&self.operation_type, &self.resource_type);
 
         if let Some(op) = fault_op {
             self.headers.insert(