Adds User-Agent header in cas_client requests to CAS (#546)

Adds User-Agent when making requests to CAS.

* sets to (project) / (version)
* version is picked from Cargo.toml
* project is hf-xet crates, git-xet crates, (also hard-coded xtool)

The reason for this change is to add better observability on the server
- so we can segment reqeusts by client and understand client versions in
the wild.

Author: rajatarya

cas_client/src/download_utils.rs

@@ -635,7 +635,7 @@ mod tests {
             MerkleHash::default(),
             file_range,
             server.base_url(),
-            Arc::new(build_http_client(RetryConfig::default(), "")?),
+            Arc::new(build_http_client(RetryConfig::default(), "", "")?),
         );
 
         fetch_info.query().await?;
@@ -682,7 +682,7 @@ mod tests {
             MerkleHash::default(),
             file_range_to_refresh,
             server.base_url(),
-            Arc::new(build_http_client(RetryConfig::default(), "")?),
+            Arc::new(build_http_client(RetryConfig::default(), "", "")?),
         ));
 
         // Spawn multiple tasks each calling into refresh with a different delay in
@@ -751,7 +751,7 @@ mod tests {
             MerkleHash::default(),
             file_range,
             server.base_url(),
-            Arc::new(build_http_client(RetryConfig::default(), "")?),
+            Arc::new(build_http_client(RetryConfig::default(), "", "")?),
         );
 
         let (offset_info_first_range, terms) = fetch_info.query().await?.unwrap();
@@ -762,7 +762,7 @@ mod tests {
                 range: x1range[0].range,
                 fetch_info: Arc::new(fetch_info),
                 chunk_cache: None,
-                client: Arc::new(build_http_client(RetryConfig::default(), "")?),
+                client: Arc::new(build_http_client(RetryConfig::default(), "", "")?),
                 range_download_single_flight: Arc::new(Group::new()),
             },
             term: terms[0].clone(),

cas_client/src/http_client.rs

@@ -88,29 +88,40 @@ impl RetryConfig<No429RetryStrategy> {
     }
 }
 
-fn reqwest_client() -> Result<reqwest::Client, CasClientError> {
+fn reqwest_client(user_agent: &str) -> Result<reqwest::Client, CasClientError> {
     // custom dns resolver not supported in WASM. no access to getaddrinfo/any other dns interface.
     #[cfg(target_family = "wasm")]
     {
-        static CLIENT: std::sync::LazyLock<reqwest::Client> = std::sync::LazyLock::new(|| reqwest::Client::new());
-        Ok((&*CLIENT).clone())
+        // For WASM, create a new client with the specified user_agent
+        // Note: we could cache this, but user_agent can vary, so we create per-call
+        let mut builder = reqwest::Client::builder();
+        if !user_agent.is_empty() {
+            builder = builder.user_agent(user_agent);
+        }
+        Ok(builder.build()?)
     }
 
     #[cfg(not(target_family = "wasm"))]
     {
         use xet_runtime::XetRuntime;
 
         let client = XetRuntime::get_or_create_reqwest_client(|| {
-            reqwest::Client::builder()
+            let mut builder = reqwest::Client::builder()
                 .pool_idle_timeout(*CLIENT_IDLE_CONNECTION_TIMEOUT)
                 .pool_max_idle_per_host(*CLIENT_MAX_IDLE_CONNECTIONS)
-                .http1_only() // high throughput parallel I/O has been shown to bottleneck with http2
-                .build()
+                .http1_only(); // high throughput parallel I/O has been shown to bottleneck with http2
+
+            if !user_agent.is_empty() {
+                builder = builder.user_agent(user_agent);
+            }
+
+            builder.build()
         })?;
 
         info!(
             idle_timeout=?*CLIENT_IDLE_CONNECTION_TIMEOUT,
             max_idle_connections=*CLIENT_MAX_IDLE_CONNECTIONS,
+            user_agent=?if user_agent.is_empty() { None } else { Some(user_agent) },
             "HTTP client configured"
         );
 
@@ -124,12 +135,13 @@ pub fn build_auth_http_client<R: RetryableStrategy + Send + Sync + 'static>(
     auth_config: &Option<AuthConfig>,
     retry_config: RetryConfig<R>,
     session_id: &str,
+    user_agent: &str,
 ) -> Result<ClientWithMiddleware, CasClientError> {
     let auth_middleware = auth_config.as_ref().map(AuthMiddleware::from);
     let logging_middleware = Some(LoggingMiddleware);
     let session_middleware = (!session_id.is_empty()).then(|| SessionMiddleware(session_id.to_owned()));
 
-    let client = ClientBuilder::new(reqwest_client()?)
+    let client = ClientBuilder::new(reqwest_client(user_agent)?)
         .maybe_with(auth_middleware)
         .with(get_retry_middleware(retry_config))
         .maybe_with(logging_middleware)
@@ -142,11 +154,12 @@ pub fn build_auth_http_client<R: RetryableStrategy + Send + Sync + 'static>(
 pub fn build_auth_http_client_no_retry(
     auth_config: &Option<AuthConfig>,
     session_id: &str,
+    user_agent: &str,
 ) -> Result<ClientWithMiddleware, CasClientError> {
     let auth_middleware = auth_config.as_ref().map(AuthMiddleware::from).info_none("CAS auth disabled");
     let logging_middleware = Some(LoggingMiddleware);
     let session_middleware = (!session_id.is_empty()).then(|| SessionMiddleware(session_id.to_owned()));
-    Ok(ClientBuilder::new(reqwest_client()?)
+    Ok(ClientBuilder::new(reqwest_client(user_agent)?)
         .maybe_with(auth_middleware)
         .maybe_with(logging_middleware)
         .maybe_with(session_middleware)
@@ -158,14 +171,15 @@ pub fn build_auth_http_client_no_retry(
 pub fn build_http_client<R: RetryableStrategy + Send + Sync + 'static>(
     retry_config: RetryConfig<R>,
     session_id: &str,
+    user_agent: &str,
 ) -> Result<ClientWithMiddleware, CasClientError> {
-    build_auth_http_client(&None, retry_config, session_id)
+    build_auth_http_client(&None, retry_config, session_id, user_agent)
 }
 
 /// Builds HTTP Client to talk to CAS.
 /// Includes retry middleware with exponential backoff.
-pub fn build_http_client_no_retry(session_id: &str) -> Result<ClientWithMiddleware, CasClientError> {
-    build_auth_http_client_no_retry(&None, session_id)
+pub fn build_http_client_no_retry(session_id: &str, user_agent: &str) -> Result<ClientWithMiddleware, CasClientError> {
+    build_auth_http_client_no_retry(&None, session_id, user_agent)
 }
 
 /// RetryStrategy
@@ -386,7 +400,7 @@ mod tests {
                 max_retry_interval_ms: 3000,
                 strategy: DefaultRetryableStrategy,
             };
-            let client = build_auth_http_client(&None, retry_config, "").unwrap();
+            let client = build_auth_http_client(&None, retry_config, "", "").unwrap();
 
             // Act & Assert - should retry and log
             let response = client.get(server.url("/data")).send().await.unwrap();
@@ -412,7 +426,7 @@ mod tests {
                 max_retry_interval_ms: 3000,
                 strategy: No429RetryStrategy,
             };
-            let client = build_auth_http_client(&None, retry_config, "").unwrap();
+            let client = build_auth_http_client(&None, retry_config, "", "").unwrap();
 
             // Act & Assert - should retry and log
             let response = client.get(server.url("/data")).send().await.unwrap();
@@ -442,7 +456,7 @@ mod tests {
                 max_retry_interval_ms: 3000,
                 strategy: DefaultRetryableStrategy,
             };
-            let client = build_auth_http_client(&None, retry_config, "").unwrap();
+            let client = build_auth_http_client(&None, retry_config, "", "").unwrap();
 
             // Act & Assert - should retry and log
             let response = client.get(server.url("/data")).send().await.unwrap();
@@ -468,7 +482,7 @@ mod tests {
                 max_retry_interval_ms: 3000,
                 strategy: No429RetryStrategy,
             };
-            let client = build_auth_http_client(&None, retry_config, "").unwrap();
+            let client = build_auth_http_client(&None, retry_config, "", "").unwrap();
 
             // Act & Assert - should retry and log
             let response = client.get(server.url("/data")).send().await.unwrap();
@@ -499,7 +513,7 @@ mod tests {
                 max_retry_interval_ms: 6000,
                 strategy: DefaultRetryableStrategy,
             };
-            let client = build_auth_http_client(&None, retry_config, "").unwrap();
+            let client = build_auth_http_client(&None, retry_config, "", "").unwrap();
 
             // Act & Assert - should retry and log
             let response = client.get(server.url("/data")).send().await.unwrap();
@@ -527,7 +541,7 @@ mod tests {
                 max_retry_interval_ms: 6000,
                 strategy: No429RetryStrategy,
             };
-            let client = build_auth_http_client(&None, retry_config, "").unwrap();
+            let client = build_auth_http_client(&None, retry_config, "", "").unwrap();
 
             // Act & Assert - should retry and log
             let response = client.get(server.url("/data")).send().await.unwrap();
@@ -557,7 +571,7 @@ mod tests {
             max_retry_interval_ms: 6000,
             strategy: No429RetryStrategy,
         };
-        let client = build_auth_http_client(&None, retry_config, "").unwrap();
+        let client = build_auth_http_client(&None, retry_config, "", "").unwrap();
 
         // Act & Assert - should retry and log
         let response = client.get(server.url("/data")).send().await.unwrap();

cas_client/src/remote_client.rs

@@ -209,6 +209,7 @@ impl RemoteClient {
         shard_cache_directory: Option<PathBuf>,
         session_id: &str,
         dry_run: bool,
+        user_agent: &str,
     ) -> Self {
         // use disk cache if cache_config provided.
         let chunk_cache = if let Some(cache_config) = cache_config {
@@ -229,13 +230,13 @@ impl RemoteClient {
             endpoint: endpoint.to_string(),
             dry_run,
             authenticated_http_client_with_retry: Arc::new(
-                http_client::build_auth_http_client(auth, RetryConfig::default(), session_id).unwrap(),
+                http_client::build_auth_http_client(auth, RetryConfig::default(), session_id, user_agent).unwrap(),
             ),
             authenticated_http_client: Arc::new(
-                http_client::build_auth_http_client_no_retry(auth, session_id).unwrap(),
+                http_client::build_auth_http_client_no_retry(auth, session_id, user_agent).unwrap(),
             ),
             http_client_with_retry: Arc::new(
-                http_client::build_http_client(RetryConfig::default(), session_id).unwrap(),
+                http_client::build_http_client(RetryConfig::default(), session_id, user_agent).unwrap(),
             ),
             chunk_cache,
             #[cfg(not(target_family = "wasm"))]
@@ -941,7 +942,7 @@ mod tests {
         let raw_xorb = build_raw_xorb(3, ChunkSize::Random(512, 10248));
 
         let threadpool = XetRuntime::new().unwrap();
-        let client = RemoteClient::new(CAS_ENDPOINT, &None, &None, None, "", false);
+        let client = RemoteClient::new(CAS_ENDPOINT, &None, &None, None, "", false, "");
 
         let cas_object = build_and_verify_cas_object(raw_xorb, Some(CompressionScheme::LZ4));
 
@@ -1315,7 +1316,7 @@ mod tests {
 
         // test reconstruct and sequential write
         let test = test_case.clone();
-        let client = RemoteClient::new(endpoint, &None, &None, None, "", false);
+        let client = RemoteClient::new(endpoint, &None, &None, None, "", false, "");
         let buf = ThreadSafeBuffer::default();
         let provider = SequentialOutput::from(buf.clone());
         let resp = threadpool.external_run_async_task(async move {
@@ -1337,7 +1338,7 @@ mod tests {
 
         // test reconstruct and parallel write
         let test = test_case;
-        let client = RemoteClient::new(endpoint, &None, &None, None, "", false);
+        let client = RemoteClient::new(endpoint, &None, &None, None, "", false, "");
         let buf = ThreadSafeBuffer::default();
         let provider = SeekingOutputProvider::from(buf.clone());
         let resp = threadpool.external_run_async_task(async move {

data/src/bin/xtool.rs

@@ -18,6 +18,7 @@ use walkdir::WalkDir;
 use xet_runtime::XetRuntime;
 
 const DEFAULT_HF_ENDPOINT: &str = "https://huggingface.co";
+const USER_AGENT: &str = concat!("xtool", "/", env!("CARGO_PKG_VERSION"));
 
 #[derive(Parser)]
 struct XCommand {
@@ -60,7 +61,7 @@ impl XCommand {
             &endpoint,
             RepoInfo::try_from(&self.overrides.repo_type, &self.overrides.repo_id)?,
             Some("main".to_owned()),
-            "xtool",
+            USER_AGENT,
             "",
             cred_helper,
         )?;
@@ -209,6 +210,7 @@ async fn query_reconstruction(
         None,
         Some((jwt_info.access_token, jwt_info.exp)),
         Some(token_refresher),
+        USER_AGENT.to_string(),
     )?;
     let cas_storage_config = &config.data_config;
     let remote_client = RemoteClient::new(
@@ -218,6 +220,7 @@ async fn query_reconstruction(
         Some(config.shard_config.cache_directory.clone()),
         "",
         true,
+        &cas_storage_config.user_agent,
     );
 
     remote_client

data/src/configurations.rs

@@ -23,6 +23,7 @@ pub struct DataConfig {
     pub prefix: String,
     pub cache_config: CacheConfig,
     pub staging_directory: Option<PathBuf>,
+    pub user_agent: String,
 }
 
 #[derive(Debug)]
@@ -98,6 +99,7 @@ impl TranslatorConfig {
                     cache_size: *CHUNK_CACHE_SIZE_BYTES,
                 },
                 staging_directory: None,
+                user_agent: String::new(),
             },
             shard_config: ShardConfig {
                 prefix: PREFIX_DEFAULT.into(),