diff --git a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/user_agent.rs b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/user_agent.rs index b9a3e53bb1..a2c553ac85 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/user_agent.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/user_agent.rs @@ -447,3 +447,136 @@ async fn wrapping_sdk_identifier_appears_on_all_requests() { ); } } + +/// Expected cross-SDK feature-flag token on the wire for the emulator's +/// default client configuration: per-partition circuit breaker (PPCB, bit +/// `0x2`, enabled by default) OR HTTP/2 (bit `0x10`, the connection-pool +/// default) == `0x12`, encoded as `|F12`. +/// +/// This is the Rust side of the cross-SDK `User-Agent` feature-flag contract +/// (see `UserAgentFeatureFlags` in `azure_data_cosmos_driver`). +const EXPECTED_FEATURE_TOKEN: &str = "|F12"; + +/// Verifies that the cross-SDK feature-flag token (`|F`) is advertised in +/// the `User-Agent` header on data-plane requests, so backend telemetry can +/// bucket Rust traffic by enabled client feature. +#[tokio::test] +async fn user_agent_advertises_feature_flags_on_the_wire() { + let observer = RecordingObserver::new(); + let emulator = build_emulator(observer.clone()); + + perform_create_and_read(emulator, None).await; + + let snapshots = observer.snapshots(); + let data_plane: Vec<&RequestSnapshot> = snapshots + .iter() + .filter(|s| is_item_data_plane_request(s)) + .collect(); + + // Sanity-check that the data plane was actually exercised; otherwise the + // assertion below would be vacuously satisfied. + assert!( + data_plane.len() >= 2, + "expected at least one create_item POST and one read_item GET to reach the emulator; \ + captured requests: {:?}", + snapshots + .iter() + .map(|s| (s.method, s.url.as_str(), s.user_agent.as_deref())) + .collect::>(), + ); + + let missing: Vec<_> = data_plane + .iter() + .filter(|s| { + !s.user_agent + .as_deref() + .is_some_and(|ua| ua.ends_with(EXPECTED_FEATURE_TOKEN)) + }) + .map(|s| (s.method, s.url.as_str(), s.user_agent.as_deref())) + .collect(); + assert!( + missing.is_empty(), + "expected every data-plane request to end with feature-flag token {EXPECTED_FEATURE_TOKEN:?}; \ + requests missing the token: {missing:?}", + ); +} + +/// Verifies that when a [`UserAgentSuffix`] is configured, the feature-flag +/// token is appended *after* the suffix with no separating space — matching +/// the .NET/Java `userAgent + "|F" + hex` encoding — so both the operator +/// suffix and the feature bitmask survive on the wire. +#[tokio::test] +async fn user_agent_appends_feature_token_after_suffix_on_the_wire() { + const SUFFIX: &str = "myapp-westus2"; + + let observer = RecordingObserver::new(); + let emulator = build_emulator(observer.clone()); + + perform_create_and_read(emulator, Some(UserAgentSuffix::new(SUFFIX))).await; + + let snapshots = observer.snapshots(); + let data_plane: Vec<&RequestSnapshot> = snapshots + .iter() + .filter(|s| is_item_data_plane_request(s)) + .collect(); + + assert!( + data_plane.len() >= 2, + "expected at least one create_item POST and one read_item GET to reach the emulator; \ + captured requests: {:?}", + snapshots + .iter() + .map(|s| (s.method, s.url.as_str(), s.user_agent.as_deref())) + .collect::>(), + ); + + let expected_tail = format!("{SUFFIX}{EXPECTED_FEATURE_TOKEN}"); + let missing: Vec<_> = data_plane + .iter() + .filter(|s| { + !s.user_agent + .as_deref() + .is_some_and(|ua| ua.ends_with(&expected_tail)) + }) + .map(|s| (s.method, s.url.as_str(), s.user_agent.as_deref())) + .collect(); + assert!( + missing.is_empty(), + "expected every data-plane request to end with {expected_tail:?}; \ + requests missing it: {missing:?}", + ); +} + +/// Negative control: the feature-flag token must be a genuine, separately +/// computed artifact — assert that the default `User-Agent` (no suffix) ends +/// with exactly one `|F` token and nothing trails it. +#[tokio::test] +async fn feature_flag_token_is_the_trailing_user_agent_segment() { + let observer = RecordingObserver::new(); + let emulator = build_emulator(observer.clone()); + + perform_create_and_read(emulator, None).await; + + let snapshots = observer.snapshots(); + let data_plane: Vec<&RequestSnapshot> = snapshots + .iter() + .filter(|s| is_item_data_plane_request(s)) + .collect(); + + assert!(data_plane.len() >= 2, "data plane not exercised"); + + for snap in &data_plane { + let ua = snap + .user_agent + .as_deref() + .expect("data-plane request carried a User-Agent"); + // Exactly one feature token, and it is the final segment. + assert_eq!( + ua.matches("|F").count(), + 1, + "expected exactly one feature-flag token in {ua:?}", + ); + let token = &ua[ua.rfind("|F").unwrap()..]; + assert_eq!(token, EXPECTED_FEATURE_TOKEN, "unexpected token in {ua:?}"); + } +} diff --git a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index dd1185ea8c..7c342ffd19 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md @@ -18,6 +18,7 @@ ### Features Added +- The `User-Agent` header now advertises enabled client features (PPCB, HTTP/2) via the cross-SDK `|F` feature-flag token, consistent with the .NET and Java Cosmos SDKs. ([#4635](https://github.com/Azure/azure-sdk-for-rust/pull/4635)) - Added support for using a native query planning library to generate query plans locally, avoiding a Gateway round-trip on cross-partition queries. Gated behind the `__internal_native_query_plan` feature flag. ([#4554](https://github.com/Azure/azure-sdk-for-rust/pull/4554)) - Restructured the client / runtime options layering on the driver. Two new nested option groups, a per-client overrides surface on `DriverOptionsBuilder`, and a single canonical `AZURE_COSMOS_PPCB_*` namespace for partition-failover environment variables. The driver now consumes partition-failover configuration once at construction (`CosmosDriver::new` no longer fabricates an `OperationOptionsView` outside any operation context) ([#4588](https://github.com/Azure/azure-sdk-for-rust/pull/4588)): - Added new nested `OperationOptions::throughput_control` group (`ThroughputControlOptions` / `…Builder` / `…View`, mirroring the `ThrottlingRetryOptions` pattern). Exposes three layered fields ([#4588](https://github.com/Azure/azure-sdk-for-rust/pull/4588)): diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs index e6c4ec53cc..06da07d6be 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs @@ -26,6 +26,7 @@ use crate::{ effective_partition_key::EffectivePartitionKey, AccountEndpoint, AccountReference, ContainerProperties, ContainerReference, ContinuationToken, CosmosOperation, DatabaseReference, PartitionKey, ResolvedToken, ResourceType, UserAgent, + UserAgentFeatureFlags, }, options::{ ConnectionPoolOptions, DriverOptions, OperationOptions, OperationOptionsView, @@ -1105,16 +1106,29 @@ impl CosmosDriver { let account_endpoint = AccountEndpoint::from(&account); let default_endpoint = CosmosEndpoint::global(account.endpoint().clone()); - // Per-driver User-Agent: when the driver-level suffix override is unset, - // share the runtime's `Arc` (cheap atomic refcount bump); - // when set, compute a fresh `UserAgent` from the runtime's wrapping-SDK - // identifier and the driver's suffix, owned by this driver alone. + // Per-driver User-Agent: compute the cross-SDK feature flags advertised + // in the header from this driver's effective client configuration — + // HTTP/2 (runtime connection pool) and PPCB (this driver's partition + // failover options). When the driver-level suffix override is unset and + // the flags match the runtime's base flags (the common case), share the + // runtime's `Arc` (cheap atomic refcount bump). Otherwise + // compute a fresh `UserAgent` owned by this driver alone. + let feature_flags = UserAgentFeatureFlags::from_client_config( + runtime.connection_pool().is_http2_allowed(), + options + .partition_failover_options() + .circuit_breaker_enabled(), + ); let user_agent = match options.user_agent_suffix() { Some(suffix) => Arc::new(UserAgent::from_suffix( runtime.wrapping_sdk_identifier(), suffix, + feature_flags, )), - None => Arc::clone(runtime.user_agent()), + None if feature_flags == runtime.user_agent_feature_flags() => { + Arc::clone(runtime.user_agent()) + } + None => Arc::new(runtime.user_agent_with_feature_flags(feature_flags)), }; // Per-driver HTTP client factory: wrap with fault injection if rules @@ -4247,6 +4261,66 @@ mod tests { assert!(!driver.user_agent().as_str().contains("runtime-default")); } + #[tokio::test] + async fn driver_disabling_ppcb_recomputes_user_agent_without_ppcb_bit() { + // A driver that disables PPCB (with no per-driver suffix override) must + // NOT share the runtime's base `Arc`: its feature flags + // differ from the runtime's, so it recomputes its own `UserAgent` whose + // cross-SDK feature token drops the PPCB bit (0x2) while retaining + // HTTP/2 (0x10) -> `|F10`. This exercises the `None => recompute` branch + // in `CosmosDriver::new` and proves the emitted token tracks per-driver + // client configuration rather than a hardcoded value. + let factory = Arc::new(ScriptedFactory::new(std::iter::repeat_n( + ResponsePlan::Success, + 10, + ))); + let runtime = Arc::new( + CosmosDriverRuntimeBuilder::new() + .with_http_client_factory(factory) + .build() + .await + .unwrap(), + ); + + // The runtime's base header advertises HTTP/2 + PPCB by default (|F12). + assert_eq!( + runtime.user_agent_feature_flags(), + UserAgentFeatureFlags::HTTP2 | UserAgentFeatureFlags::PER_PARTITION_CIRCUIT_BREAKER, + ); + assert!( + runtime.user_agent().as_str().ends_with("|F12"), + "unexpected runtime User-Agent: {}", + runtime.user_agent().as_str() + ); + + let driver = CosmosDriver::new( + Arc::clone(&runtime), + DriverOptionsBuilder::new(signed_test_account( + "https://account.documents.azure.com:443/", + )) + .with_partition_failover_options( + crate::options::PartitionFailoverOptions::builder() + .with_circuit_breaker_enabled(false) + .build() + .unwrap(), + ) + .build(), + ) + .expect("CosmosDriver::new should succeed in tests"); + + // Distinct allocation (the recompute branch), not the shared runtime Arc. + assert!( + !Arc::ptr_eq(driver.user_agent(), runtime.user_agent()), + "driver disabling PPCB must own a distinct User-Agent Arc" + ); + // PPCB bit (0x2) dropped, HTTP/2 (0x10) retained -> |F10. + assert!( + driver.user_agent().as_str().ends_with("|F10"), + "expected driver User-Agent to drop the PPCB bit (|F10): {}", + driver.user_agent().as_str() + ); + } + // ========================================================================= // pre_resolve_partition_key_range_id — EPK-range seeding (#4611 fix) // diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/runtime.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/runtime.rs index a2286ab540..ac16daf8e7 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/runtime.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/runtime.rs @@ -14,10 +14,10 @@ use std::{ use crate::{ diagnostics::ProxyConfiguration, - models::{normalize_wrapping_sdk_identifier, UserAgent}, + models::{normalize_wrapping_sdk_identifier, UserAgent, UserAgentFeatureFlags}, options::{ parse_duration_millis_from_env, ConnectionPoolOptions, CorrelationId, DriverOptions, - OperationOptions, UserAgentSuffix, WorkloadId, + OperationOptions, PartitionFailoverOptions, UserAgentSuffix, WorkloadId, }, system::{CpuMemoryMonitor, VmMetadataService}, }; @@ -153,6 +153,15 @@ pub struct CosmosDriverRuntime { /// driver's own identifier. wrapping_sdk_identifier: Option, + /// Cross-SDK feature flags advertised in this runtime's base `User-Agent`. + /// + /// Computed from runtime-scoped client configuration (HTTP/2 transport and + /// the default per-partition circuit breaker setting). Drivers built from + /// this runtime compare their own computed flags against this value: when + /// they match (the common case), they share the runtime's `Arc`; + /// otherwise they recompute their own `User-Agent`. + user_agent_feature_flags: UserAgentFeatureFlags, + /// Shared container metadata cache used by drivers in this runtime. container_cache: ContainerCache, @@ -305,6 +314,32 @@ impl CosmosDriverRuntime { self.wrapping_sdk_identifier.as_deref() } + /// Returns the cross-SDK feature flags advertised in this runtime's base + /// `User-Agent` header. + pub(crate) fn user_agent_feature_flags(&self) -> UserAgentFeatureFlags { + self.user_agent_feature_flags + } + + /// Recomputes a `User-Agent` from this runtime's suffix source (suffix, + /// workload id, or correlation id, in priority order) plus the supplied + /// feature flags. + /// + /// Used by a driver that overrode a feature-affecting option (e.g. disabled + /// PPCB) without supplying its own suffix, so it cannot share the runtime's + /// shared `Arc`. + pub(crate) fn user_agent_with_feature_flags( + &self, + feature_flags: UserAgentFeatureFlags, + ) -> UserAgent { + compute_user_agent( + self.wrapping_sdk_identifier.as_deref(), + self.user_agent_suffix.as_ref(), + self.workload_id, + self.correlation_id.as_ref(), + feature_flags, + ) + } + /// Returns the effective correlation dimension. /// /// Returns `correlation_id` if set, otherwise falls back to `user_agent_suffix`. @@ -532,20 +567,28 @@ impl CosmosDriverRuntimeBuilder { /// configuration failure). /// pub async fn build(self) -> crate::error::Result> { + let connection_pool = self.connection_pool.unwrap_or_default(); + + // Compute the base feature flags advertised in the User-Agent from + // runtime-scoped client configuration. HTTP/2 comes from the connection + // pool; PPCB uses the driver default (its per-driver value is folded in + // later by `CosmosDriver::new`). PPAF is server-driven per-partition and + // therefore unknown here, so it is not advertised in the shared header. + let user_agent_feature_flags = UserAgentFeatureFlags::from_client_config( + connection_pool.is_http2_allowed(), + PartitionFailoverOptions::default().circuit_breaker_enabled(), + ); + // Compute user agent from suffix/workloadId/correlationId (in priority order), // optionally prepending a wrapping-SDK identifier. - let wrapping = self.wrapping_sdk_identifier.as_deref(); - let user_agent = Arc::new(if let Some(ref suffix) = self.user_agent_suffix { - UserAgent::from_suffix(wrapping, suffix) - } else if let Some(workload_id) = self.workload_id { - UserAgent::from_workload_id(wrapping, workload_id) - } else if let Some(ref correlation_id) = self.correlation_id { - UserAgent::from_correlation_id(wrapping, correlation_id) - } else { - UserAgent::from_wrapping_sdk_identifier(wrapping) - }); + let user_agent = Arc::new(compute_user_agent( + self.wrapping_sdk_identifier.as_deref(), + self.user_agent_suffix.as_ref(), + self.workload_id, + self.correlation_id.as_ref(), + user_agent_feature_flags, + )); - let connection_pool = self.connection_pool.unwrap_or_default(); let proxy_configuration = ProxyConfiguration::from_env(connection_pool.proxy_allowed()); let http_client_factory: Arc = { #[cfg(any( @@ -634,6 +677,7 @@ impl CosmosDriverRuntimeBuilder { correlation_id: self.correlation_id, user_agent_suffix: self.user_agent_suffix, wrapping_sdk_identifier: self.wrapping_sdk_identifier, + user_agent_feature_flags, container_cache: ContainerCache::new(), account_metadata_cache: Arc::new(AccountMetadataCache::new()), cpu_monitor, @@ -645,6 +689,30 @@ impl CosmosDriverRuntimeBuilder { static NEXT_RUNTIME_ID: AtomicUsize = AtomicUsize::new(0); +/// Builds a [`UserAgent`] from a suffix source (suffix, workload id, or +/// correlation id, in priority order) plus the supplied feature flags. +/// +/// Shared by [`CosmosDriverRuntimeBuilder::build`] and +/// [`CosmosDriverRuntime::user_agent_with_feature_flags`] so the priority +/// ordering is defined in exactly one place. +fn compute_user_agent( + wrapping_sdk_identifier: Option<&str>, + user_agent_suffix: Option<&UserAgentSuffix>, + workload_id: Option, + correlation_id: Option<&CorrelationId>, + feature_flags: UserAgentFeatureFlags, +) -> UserAgent { + if let Some(suffix) = user_agent_suffix { + UserAgent::from_suffix(wrapping_sdk_identifier, suffix, feature_flags) + } else if let Some(workload_id) = workload_id { + UserAgent::from_workload_id(wrapping_sdk_identifier, workload_id, feature_flags) + } else if let Some(correlation_id) = correlation_id { + UserAgent::from_correlation_id(wrapping_sdk_identifier, correlation_id, feature_flags) + } else { + UserAgent::from_wrapping_sdk_identifier(wrapping_sdk_identifier, feature_flags) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/models/mod.rs b/sdk/cosmos/azure_data_cosmos_driver/src/models/mod.rs index 4f6cea6b1d..3174e7d43d 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/models/mod.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/models/mod.rs @@ -74,8 +74,8 @@ pub use resource_reference::{ }; pub use response_body::ResponseBody; pub use session_token_segment::SessionTokenSegment; -pub(crate) use user_agent::normalize_wrapping_sdk_identifier; pub use user_agent::UserAgent; +pub(crate) use user_agent::{normalize_wrapping_sdk_identifier, UserAgentFeatureFlags}; pub(crate) use account_reference::AccountEndpoint; diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/models/user_agent.rs b/sdk/cosmos/azure_data_cosmos_driver/src/models/user_agent.rs index bc854b8469..5f6cc741a1 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/models/user_agent.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/models/user_agent.rs @@ -3,13 +3,143 @@ //! User agent string for HTTP requests to Cosmos DB. -use std::fmt; +use std::{ + fmt, + ops::{BitAnd, BitAndAssign, BitOr, BitOrAssign}, +}; use crate::options::{CorrelationId, UserAgentSuffix, WorkloadId}; /// Maximum length for the full user agent string (HTTP header limit). const MAX_USER_AGENT_LENGTH: usize = 255; +/// Bitmask of client-side features advertised in the `User-Agent` header. +/// +/// The Cosmos SDKs share a cross-language contract: enabled client features are +/// encoded as a `|F` token appended to the `User-Agent` string, where +/// `` is the uppercase hexadecimal representation of the OR-ed bit values. +/// This lets backend telemetry bucket traffic by feature regardless of which +/// language SDK produced the request. +/// +/// **The bit values below MUST stay consistent with the other Cosmos SDKs** +/// (.NET `UserAgentFeatureFlags`, Java `UserAgentFeatureFlags`). Do not +/// renumber existing bits — only append new ones. +/// +/// # Example +/// +/// ```ignore +/// let flags = UserAgentFeatureFlags::PER_PARTITION_CIRCUIT_BREAKER +/// | UserAgentFeatureFlags::HTTP2; +/// assert_eq!(flags.to_string(), "|F12"); // 0x2 | 0x10 == 0x12 +/// ``` +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)] +pub(crate) struct UserAgentFeatureFlags(u32); + +impl UserAgentFeatureFlags { + /// No features advertised. Renders to an empty token. + pub(crate) const NONE: Self = Self(0); + + /// Per-partition automatic failover (PPAF). Cross-SDK bit value `0x1`. + /// + /// Reserved to keep Rust's bit assignments aligned with the .NET and Java + /// Cosmos SDKs; this driver does not advertise it yet (PPAF is server-driven + /// and resolved per-partition at request time, so it is unknown when the + /// shared header value is computed). + #[allow(dead_code)] // Reserved cross-SDK bit; not advertised by this driver yet. + pub(crate) const PER_PARTITION_AUTOMATIC_FAILOVER: Self = Self(1); + + /// Per-partition circuit breaker (PPCB). Cross-SDK bit value `0x2`. + pub(crate) const PER_PARTITION_CIRCUIT_BREAKER: Self = Self(2); + + /// Thin client mode. Cross-SDK bit value `0x4`. + /// + /// Reserved for cross-SDK parity; not advertised by this driver yet. + #[allow(dead_code)] // Reserved cross-SDK bit; not advertised by this driver yet. + pub(crate) const THIN_CLIENT: Self = Self(4); + + /// Cosmos binary encoding. Cross-SDK bit value `0x8`. + /// + /// Reserved for cross-SDK parity; not advertised by this driver yet. + #[allow(dead_code)] // Reserved cross-SDK bit; not advertised by this driver yet. + pub(crate) const BINARY_ENCODING: Self = Self(8); + + /// HTTP/2 transport. Cross-SDK bit value `0x10`. + pub(crate) const HTTP2: Self = Self(16); + + /// Returns the raw bitmask value. + pub(crate) const fn bits(self) -> u32 { + self.0 + } + + /// Returns `true` when no feature bits are set. + pub(crate) const fn is_empty(self) -> bool { + self.0 == 0 + } + + /// Returns the union of two flag sets (every bit set in either operand). + pub(crate) const fn union(self, other: Self) -> Self { + Self(self.0 | other.0) + } + + /// Maps the statically-known client configuration to feature flags. + /// + /// Only features whose enablement is known at `User-Agent` construction + /// time are advertised. PPAF is intentionally excluded: it is server-driven + /// and resolved per-partition at request time, so it is not known when the + /// shared header value is computed. + pub(crate) fn from_client_config(is_http2_allowed: bool, ppcb_enabled: bool) -> Self { + let mut flags = Self::NONE; + if ppcb_enabled { + flags |= Self::PER_PARTITION_CIRCUIT_BREAKER; + } + if is_http2_allowed { + flags |= Self::HTTP2; + } + flags + } +} + +impl BitOr for UserAgentFeatureFlags { + type Output = Self; + + fn bitor(self, rhs: Self) -> Self { + self.union(rhs) + } +} + +impl BitOrAssign for UserAgentFeatureFlags { + fn bitor_assign(&mut self, rhs: Self) { + *self = self.union(rhs); + } +} + +impl BitAnd for UserAgentFeatureFlags { + type Output = Self; + + fn bitand(self, rhs: Self) -> Self { + Self(self.0 & rhs.0) + } +} + +impl BitAndAssign for UserAgentFeatureFlags { + fn bitand_assign(&mut self, rhs: Self) { + self.0 &= rhs.0; + } +} + +impl fmt::Display for UserAgentFeatureFlags { + /// Renders the cross-SDK `|F` token, or an empty string when no + /// features are set. The hex digits are uppercase with no leading zeros, + /// matching the .NET and Java encodings. + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + if self.is_empty() { + Ok(()) + } else { + write!(f, "|F{:X}", self.bits()) + } + } +} + /// Azure SDK user agent prefix. const AZSDK_USER_AGENT_PREFIX: &str = "azsdk-rust-"; @@ -30,15 +160,16 @@ const SDK_VERSION: &str = env!("CARGO_PKG_VERSION"); /// - Rust version (compile time) /// /// An optional suffix can be appended (typically from [`UserAgentSuffix`], -/// [`WorkloadId`], or [`CorrelationId`]). +/// [`WorkloadId`], or [`CorrelationId`]), followed by an optional cross-SDK +/// feature-flag token (`|F`). /// /// # Example /// /// Driver used directly, no suffix: /// `azsdk-rust-cosmos-driver/0.1.0 windows/x86_64 rustc/1.85.0` /// -/// Driver used directly, with suffix: -/// `azsdk-rust-cosmos-driver/0.1.0 windows/x86_64 rustc/1.85.0 myapp-westus2` +/// Driver used directly, with suffix and feature flags: +/// `azsdk-rust-cosmos-driver/0.1.0 windows/x86_64 rustc/1.85.0 myapp-westus2|F12` /// /// Wrapped by a higher-level SDK: /// `azsdk-rust-cosmos/0.34.0 azsdk-rust-cosmos-driver/0.1.0 windows/x86_64 rustc/1.85.0 myapp-westus2` @@ -53,7 +184,7 @@ pub struct UserAgent { impl Default for UserAgent { fn default() -> Self { - Self::new(None::<&str>, None::<&str>) + Self::new(None::<&str>, None::<&str>, UserAgentFeatureFlags::NONE) } } @@ -149,14 +280,19 @@ impl UserAgent { } } - /// Creates a new user agent with optional wrapping-SDK identifier and suffix. + /// Creates a new user agent with optional wrapping-SDK identifier, suffix, + /// and feature flags. /// /// The wrapping identifier is prepended to the driver's base prefix; the - /// suffix is appended after the base, separated by a space. If the - /// resulting string exceeds 255 characters, the suffix is truncated. + /// suffix is appended after the base, separated by a space; the feature + /// flag token (`|F`) is appended last with no separator, matching the + /// cross-SDK encoding. If the resulting string would exceed 255 characters, + /// the suffix is truncated first — the feature-flag token is preserved so + /// telemetry never loses it. fn new( wrapping_sdk_identifier: Option>, suffix: Option>, + feature_flags: UserAgentFeatureFlags, ) -> Self { // Normalize to ASCII once; this makes byte-length checks safe and avoids // reprocessing after we build the final string. @@ -165,7 +301,13 @@ impl UserAgent { )); let normalized_suffix = suffix.map(Into::into).map(|s| strip_non_ascii(&s)); - let max_suffix_len = MAX_USER_AGENT_LENGTH.saturating_sub(base.len() + 1); + // The feature-flag token is short, ASCII, and higher-priority telemetry + // than an arbitrarily long operator suffix, so reserve its length up + // front and let the suffix absorb any remaining truncation. + let feature_token = feature_flags.to_string(); + + let max_suffix_len = + MAX_USER_AGENT_LENGTH.saturating_sub(base.len() + 1 + feature_token.len()); let effective_suffix = normalized_suffix.and_then(|s| { if s.is_empty() || max_suffix_len == 0 { None @@ -175,13 +317,14 @@ impl UserAgent { }); let mut full_user_agent = String::with_capacity( - base.len() + effective_suffix.as_ref().map_or(0, |s| 1 + s.len()), + base.len() + effective_suffix.as_ref().map_or(0, |s| 1 + s.len()) + feature_token.len(), ); full_user_agent.push_str(&base); if let Some(s) = &effective_suffix { full_user_agent.push(' '); full_user_agent.push_str(s); } + full_user_agent.push_str(&feature_token); Self { full_user_agent, @@ -190,26 +333,36 @@ impl UserAgent { } /// Creates a user agent with only a wrapping-SDK identifier (no suffix). - pub(crate) fn from_wrapping_sdk_identifier(wrapping_sdk_identifier: Option<&str>) -> Self { - Self::new(wrapping_sdk_identifier, None::<&str>) + pub(crate) fn from_wrapping_sdk_identifier( + wrapping_sdk_identifier: Option<&str>, + feature_flags: UserAgentFeatureFlags, + ) -> Self { + Self::new(wrapping_sdk_identifier, None::<&str>, feature_flags) } /// Creates a user agent from a [`UserAgentSuffix`]. pub(crate) fn from_suffix( wrapping_sdk_identifier: Option<&str>, suffix: &UserAgentSuffix, + feature_flags: UserAgentFeatureFlags, ) -> Self { - Self::new(wrapping_sdk_identifier, Some(suffix.as_str())) + Self::new( + wrapping_sdk_identifier, + Some(suffix.as_str()), + feature_flags, + ) } /// Creates a user agent from a [`WorkloadId`]. pub(crate) fn from_workload_id( wrapping_sdk_identifier: Option<&str>, workload_id: WorkloadId, + feature_flags: UserAgentFeatureFlags, ) -> Self { Self::new( wrapping_sdk_identifier, Some(format!("w{}", workload_id.value())), + feature_flags, ) } @@ -217,8 +370,13 @@ impl UserAgent { pub(crate) fn from_correlation_id( wrapping_sdk_identifier: Option<&str>, correlation_id: &CorrelationId, + feature_flags: UserAgentFeatureFlags, ) -> Self { - Self::new(wrapping_sdk_identifier, Some(correlation_id.as_str())) + Self::new( + wrapping_sdk_identifier, + Some(correlation_id.as_str()), + feature_flags, + ) } /// Returns the full user agent string. @@ -288,7 +446,7 @@ mod tests { #[test] fn user_agent_with_suffix() { - let ua = UserAgent::new(None::<&str>, Some("my-app")); + let ua = UserAgent::new(None::<&str>, Some("my-app"), UserAgentFeatureFlags::NONE); assert!(ua.as_str().contains("my-app")); assert_eq!(ua.suffix(), Some("my-app")); } @@ -296,21 +454,21 @@ mod tests { #[test] fn user_agent_from_user_agent_suffix() { let suffix = UserAgentSuffix::new("myapp-westus2"); - let ua = UserAgent::from_suffix(None, &suffix); + let ua = UserAgent::from_suffix(None, &suffix, UserAgentFeatureFlags::NONE); assert!(ua.as_str().contains("myapp-westus2")); } #[test] fn user_agent_from_workload_id() { let workload_id = WorkloadId::new(25); - let ua = UserAgent::from_workload_id(None, workload_id); + let ua = UserAgent::from_workload_id(None, workload_id, UserAgentFeatureFlags::NONE); assert!(ua.as_str().contains("w25")); } #[test] fn user_agent_from_correlation_id() { let correlation_id = CorrelationId::new("aks-prod-eastus"); - let ua = UserAgent::from_correlation_id(None, &correlation_id); + let ua = UserAgent::from_correlation_id(None, &correlation_id, UserAgentFeatureFlags::NONE); assert!(ua.as_str().contains("aks-prod-eastus")); } @@ -324,7 +482,10 @@ mod tests { #[test] fn user_agent_with_wrapping_sdk_identifier_prepends() { - let ua = UserAgent::from_wrapping_sdk_identifier(Some("azsdk-rust-cosmos/0.34.0")); + let ua = UserAgent::from_wrapping_sdk_identifier( + Some("azsdk-rust-cosmos/0.34.0"), + UserAgentFeatureFlags::NONE, + ); assert!( ua.as_str() .starts_with("azsdk-rust-cosmos/0.34.0 azsdk-rust-cosmos-driver/"), @@ -337,7 +498,11 @@ mod tests { #[test] fn user_agent_wrapping_plus_suffix() { let suffix = UserAgentSuffix::new("myapp-westus2"); - let ua = UserAgent::from_suffix(Some("azsdk-rust-cosmos/0.34.0"), &suffix); + let ua = UserAgent::from_suffix( + Some("azsdk-rust-cosmos/0.34.0"), + &suffix, + UserAgentFeatureFlags::NONE, + ); let s = ua.as_str(); assert!( s.starts_with("azsdk-rust-cosmos/0.34.0 azsdk-rust-cosmos-driver/"), @@ -348,15 +513,20 @@ mod tests { #[test] fn user_agent_wrapping_identifier_strips_non_ascii() { - let ua = UserAgent::from_wrapping_sdk_identifier(Some("azsdk-rust-café/0.1.0")); + let ua = UserAgent::from_wrapping_sdk_identifier( + Some("azsdk-rust-café/0.1.0"), + UserAgentFeatureFlags::NONE, + ); assert!(ua.as_str().is_ascii()); assert!(ua.as_str().starts_with("azsdk-rust-caf_/0.1.0 ")); } #[test] fn user_agent_empty_wrapping_identifier_treated_as_absent() { - let ua_empty = UserAgent::from_wrapping_sdk_identifier(Some("")); - let ua_ws = UserAgent::from_wrapping_sdk_identifier(Some(" ")); + let ua_empty = + UserAgent::from_wrapping_sdk_identifier(Some(""), UserAgentFeatureFlags::NONE); + let ua_ws = + UserAgent::from_wrapping_sdk_identifier(Some(" "), UserAgentFeatureFlags::NONE); let ua_default = UserAgent::default(); assert_eq!(ua_empty.as_str(), ua_default.as_str()); assert_eq!(ua_ws.as_str(), ua_default.as_str()); @@ -367,7 +537,7 @@ mod tests { // Force a long wrapping identifier and a long suffix; total must still be capped. let long_wrap = format!("azsdk-rust-{}", "x".repeat(200)); let suffix = UserAgentSuffix::new("a".repeat(25)); - let ua = UserAgent::from_suffix(Some(&long_wrap), &suffix); + let ua = UserAgent::from_suffix(Some(&long_wrap), &suffix, UserAgentFeatureFlags::HTTP2); assert!( ua.as_str().len() <= MAX_USER_AGENT_LENGTH, "len={} value={}", @@ -384,7 +554,7 @@ mod tests { // truncated instead. let long_wrap = format!("azsdk-rust-{}", "x".repeat(500)); let suffix = UserAgentSuffix::new("myapp-westus2"); - let ua = UserAgent::from_suffix(Some(&long_wrap), &suffix); + let ua = UserAgent::from_suffix(Some(&long_wrap), &suffix, UserAgentFeatureFlags::NONE); assert!( ua.as_str().len() <= MAX_USER_AGENT_LENGTH, "exceeded cap: {}", @@ -397,4 +567,139 @@ mod tests { ua.as_str() ); } + + #[test] + fn feature_flags_token_matches_cross_sdk_encoding() { + // Bit values and `|F` encoding must stay consistent with .NET/Java. + assert_eq!(UserAgentFeatureFlags::NONE.to_string(), ""); + assert_eq!( + UserAgentFeatureFlags::PER_PARTITION_AUTOMATIC_FAILOVER.to_string(), + "|F1" + ); + assert_eq!( + UserAgentFeatureFlags::PER_PARTITION_CIRCUIT_BREAKER.to_string(), + "|F2" + ); + assert_eq!(UserAgentFeatureFlags::THIN_CLIENT.to_string(), "|F4"); + assert_eq!(UserAgentFeatureFlags::BINARY_ENCODING.to_string(), "|F8"); + assert_eq!(UserAgentFeatureFlags::HTTP2.to_string(), "|F10"); + // PPAF (1) + PPCB (2) -> 0x3 == "|F3" (matches the Java example). + let ppaf_ppcb = UserAgentFeatureFlags::PER_PARTITION_AUTOMATIC_FAILOVER + | UserAgentFeatureFlags::PER_PARTITION_CIRCUIT_BREAKER; + assert_eq!(ppaf_ppcb.to_string(), "|F3"); + // PPCB (2) + Http2 (16) -> 0x12 == "|F12". + let ppcb_http2 = + UserAgentFeatureFlags::PER_PARTITION_CIRCUIT_BREAKER | UserAgentFeatureFlags::HTTP2; + assert_eq!(ppcb_http2.to_string(), "|F12"); + } + + #[test] + fn feature_flags_bitwise_helpers() { + let mut flags = UserAgentFeatureFlags::NONE; + assert!(flags.is_empty()); + flags |= UserAgentFeatureFlags::HTTP2; + flags |= UserAgentFeatureFlags::PER_PARTITION_CIRCUIT_BREAKER; + assert!(!flags.is_empty()); + assert_eq!(flags.bits(), 0x12); + + // `union` matches the `|` operator. + assert_eq!( + UserAgentFeatureFlags::HTTP2 + .union(UserAgentFeatureFlags::PER_PARTITION_CIRCUIT_BREAKER), + flags + ); + + // `&` / `&=` mask down to the intersecting bits. + assert_eq!( + flags & UserAgentFeatureFlags::HTTP2, + UserAgentFeatureFlags::HTTP2 + ); + assert_eq!( + flags & UserAgentFeatureFlags::PER_PARTITION_AUTOMATIC_FAILOVER, + UserAgentFeatureFlags::NONE + ); + let mut masked = flags; + masked &= UserAgentFeatureFlags::HTTP2; + assert_eq!(masked, UserAgentFeatureFlags::HTTP2); + } + + #[test] + fn feature_flags_from_client_config() { + assert_eq!( + UserAgentFeatureFlags::from_client_config(false, false), + UserAgentFeatureFlags::NONE + ); + assert_eq!( + UserAgentFeatureFlags::from_client_config(true, false), + UserAgentFeatureFlags::HTTP2 + ); + assert_eq!( + UserAgentFeatureFlags::from_client_config(false, true), + UserAgentFeatureFlags::PER_PARTITION_CIRCUIT_BREAKER + ); + assert_eq!( + UserAgentFeatureFlags::from_client_config(true, true), + UserAgentFeatureFlags::HTTP2 | UserAgentFeatureFlags::PER_PARTITION_CIRCUIT_BREAKER + ); + } + + #[test] + fn user_agent_appends_feature_token_after_suffix() { + let suffix = UserAgentSuffix::new("myapp-westus2"); + let flags = + UserAgentFeatureFlags::PER_PARTITION_CIRCUIT_BREAKER | UserAgentFeatureFlags::HTTP2; + let ua = UserAgent::from_suffix(None, &suffix, flags); + // The token is appended directly to the suffix with no separating space, + // matching the .NET/Java `userAgent + "|F" + hex` encoding. + assert!( + ua.as_str().ends_with("myapp-westus2|F12"), + "unexpected user agent: {}", + ua.as_str() + ); + assert_eq!(ua.suffix(), Some("myapp-westus2")); + } + + #[test] + fn user_agent_appends_feature_token_without_suffix() { + let ua = UserAgent::from_wrapping_sdk_identifier(None, UserAgentFeatureFlags::HTTP2); + assert!( + ua.as_str().ends_with("|F10"), + "unexpected user agent: {}", + ua.as_str() + ); + assert!(ua.suffix().is_none()); + } + + #[test] + fn user_agent_no_feature_token_when_flags_empty() { + let ua = UserAgent::default(); + assert!( + !ua.as_str().contains("|F"), + "unexpected feature token in: {}", + ua.as_str() + ); + } + + #[test] + fn user_agent_keeps_feature_token_over_suffix_when_truncating() { + // The feature token is higher-priority telemetry than the operator + // suffix: when a pathologically long wrapping identifier leaves no room + // for the full suffix, the suffix is truncated but the token survives + // and the total stays within the cap. + let long_wrap = format!("azsdk-rust-{}", "x".repeat(500)); + let suffix = UserAgentSuffix::new("a".repeat(UserAgentSuffix::MAX_LENGTH)); + let flags = + UserAgentFeatureFlags::PER_PARTITION_CIRCUIT_BREAKER | UserAgentFeatureFlags::HTTP2; + let ua = UserAgent::from_suffix(Some(&long_wrap), &suffix, flags); + assert!( + ua.as_str().len() <= MAX_USER_AGENT_LENGTH, + "exceeded cap: {}", + ua.as_str() + ); + assert!( + ua.as_str().ends_with("|F12"), + "feature token lost: {}", + ua.as_str() + ); + } }