Skip to main content

xet_client/cas_client/
remote_client.rs

1use std::sync::Arc;
2use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
3
4use anyhow::anyhow;
5use bytes::Bytes;
6use futures::TryStreamExt;
7use http::HeaderValue;
8use http::header::{CONTENT_LENGTH, HeaderMap, RANGE};
9use reqwest::{Body, Response, StatusCode, Url};
10use reqwest_middleware::ClientWithMiddleware;
11use tracing::{event, info, instrument};
12use xet_core_structures::merklehash::MerkleHash;
13use xet_core_structures::metadata_shard::file_structs::{FileDataSequenceEntry, FileDataSequenceHeader, MDBFileInfo};
14use xet_core_structures::xorb_object::SerializedXorbObject;
15use xet_runtime::core::XetContext;
16
17use super::adaptive_concurrency::{
18    AdaptiveConcurrencyController, ConnectionPermit, download_controller, upload_controller,
19};
20use super::auth::AuthConfig;
21use super::interface::URLProvider;
22use super::progress_tracked_streams::{
23    DownloadProgressStream, ProgressCallback, StreamProgressReporter, UploadProgressStream,
24};
25use super::retry_wrapper::{RetryWrapper, RetryableReqwestError};
26use super::{Client, INFORMATION_LOG_LEVEL};
27use crate::cas_types::{
28    BatchQueryReconstructionResponse, FileChunkHashesResponse, FileRange, HttpRange, Key, QueryReconstructionResponse,
29    QueryReconstructionResponseV2, UploadShardResponse, UploadShardResponseType, UploadXorbResponse,
30    X_RANGE_DIRTY_HEADER,
31};
32use crate::common::http_client::{self, Api};
33use crate::error::{ClientError, Result};
34
35pub const CAS_ENDPOINT: &str = "http://localhost:8080";
36pub const PREFIX_DEFAULT: &str = "default";
37
38use lazy_static::lazy_static;
39
40lazy_static! {
41    static ref FN_CALL_ID: AtomicU64 = AtomicU64::new(1);
42}
43
44pub struct RemoteClient {
45    pub(crate) ctx: XetContext,
46    endpoint: String,
47    dry_run: bool,
48    http_client: Arc<ClientWithMiddleware>,
49    authenticated_http_client: Arc<ClientWithMiddleware>,
50    /// Authenticated client with no read_timeout, used for shard uploads where server-side
51    /// processing time scales with file entry count and can exceed the global read_timeout.
52    #[cfg(not(target_family = "wasm"))]
53    shard_upload_http_client: Arc<ClientWithMiddleware>,
54    upload_concurrency_controller: Arc<AdaptiveConcurrencyController>,
55    download_concurrency_controller: Arc<AdaptiveConcurrencyController>,
56    /// Caches the discovered reconstruction API version (0 = not yet probed, 1 = V1, 2 = V2).
57    detected_reconstruction_api_version: AtomicU32,
58}
59
60impl RemoteClient {
61    /// Creates a new RemoteClient with an explicit Unix socket path.
62    ///
63    /// # Arguments
64    /// * `endpoint` - The CAS endpoint URL
65    /// * `auth` - Optional authentication configuration
66    /// * `session_id` - Session identifier
67    /// * `dry_run` - Whether to run in dry-run mode
68    /// * `unix_socket_path` - Optional Unix socket path for proxying connections (ignored on non-Unix platforms)
69    /// * `custom_headers` - Optional custom headers to include in HTTP requests (should include User-Agent)
70    pub fn new_with_socket(
71        ctx: XetContext,
72        endpoint: &str,
73        auth: &Option<AuthConfig>,
74        session_id: &str,
75        dry_run: bool,
76        unix_socket_path: Option<&str>,
77        custom_headers: Option<Arc<HeaderMap>>,
78    ) -> Arc<Self> {
79        Arc::new(Self {
80            ctx: ctx.clone(),
81            endpoint: endpoint.to_string(),
82            dry_run,
83            authenticated_http_client: Arc::new(
84                http_client::build_auth_http_client(&ctx, auth, session_id, unix_socket_path, custom_headers.clone())
85                    .unwrap(),
86            ),
87            http_client: Arc::new(
88                http_client::build_http_client(&ctx, session_id, unix_socket_path, custom_headers.clone()).unwrap(),
89            ),
90            #[cfg(not(target_family = "wasm"))]
91            shard_upload_http_client: Arc::new(
92                http_client::build_auth_http_client_no_read_timeout(
93                    &ctx,
94                    auth,
95                    session_id,
96                    unix_socket_path,
97                    custom_headers,
98                )
99                .unwrap(),
100            ),
101            upload_concurrency_controller: upload_controller(&ctx, endpoint),
102            download_concurrency_controller: download_controller(&ctx, endpoint),
103            detected_reconstruction_api_version: AtomicU32::new(0),
104        })
105    }
106
107    /// Creates a new RemoteClient.
108    ///
109    /// If `HF_XET_CLIENT_UNIX_SOCKET_PATH` is set in the configuration, this will
110    /// automatically use the Unix socket for connections (checked by build_http_client).
111    ///
112    /// # Arguments
113    /// * `endpoint` - The CAS endpoint URL
114    /// * `auth` - Optional authentication configuration
115    /// * `session_id` - Session identifier
116    /// * `dry_run` - Whether to run in dry-run mode
117    /// * `custom_headers` - Optional custom headers to include in HTTP requests (should include User-Agent)
118    pub fn new(
119        ctx: XetContext,
120        endpoint: &str,
121        auth: &Option<AuthConfig>,
122        session_id: &str,
123        dry_run: bool,
124        custom_headers: Option<Arc<HeaderMap>>,
125    ) -> Arc<Self> {
126        Self::new_with_socket(ctx, endpoint, auth, session_id, dry_run, None, custom_headers)
127    }
128
129    /// Get the endpoint URL.
130    pub fn endpoint(&self) -> &str {
131        &self.endpoint
132    }
133
134    #[cfg(feature = "simulation")]
135    pub(crate) fn http_client(&self) -> Arc<ClientWithMiddleware> {
136        self.http_client.clone()
137    }
138
139    async fn query_dedup_api(&self, prefix: &str, chunk_hash: &MerkleHash) -> Result<Option<Response>> {
140        // The API endpoint now only supports non-batched dedup request and
141        let key = Key {
142            prefix: prefix.into(),
143            hash: *chunk_hash,
144        };
145
146        let call_id = FN_CALL_ID.fetch_add(1, Ordering::Relaxed);
147        let url = Url::parse(&format!("{}/v1/chunks/{key}", self.endpoint))?;
148        event!(
149            INFORMATION_LOG_LEVEL,
150            call_id,
151            prefix,
152            %chunk_hash,
153            "Starting query_dedup API call",
154        );
155
156        let client = self.authenticated_http_client.clone();
157        let api_tag = "cas::query_dedup";
158
159        let result = RetryWrapper::new(self.ctx.clone(), api_tag)
160            .with_429_no_retry()
161            .with_expected_404()
162            .log_errors_as_info()
163            .run(move || client.get(url.clone()).with_extension(Api(api_tag)).send())
164            .await;
165
166        if result.as_ref().is_err_and(|e| e.status().is_some()) {
167            event!(
168                INFORMATION_LOG_LEVEL,
169                call_id,
170                prefix,
171                %chunk_hash,
172                result="not_found",
173                "Completed query_dedup API call",
174            );
175            return Ok(None);
176        }
177
178        event!(
179            INFORMATION_LOG_LEVEL,
180            call_id,
181            prefix,
182            %chunk_hash,
183            result="found",
184            "Completed query_dedup API call",
185        );
186        Ok(Some(result?))
187    }
188}
189
190impl RemoteClient {
191    async fn get_reconstruction_impl<T>(
192        &self,
193        file_id: &MerkleHash,
194        bytes_range: Option<FileRange>,
195        api_version: &str,
196    ) -> Result<Option<T>>
197    where
198        T: serde::de::DeserializeOwned + 'static,
199    {
200        let call_id = FN_CALL_ID.fetch_add(1, Ordering::Relaxed);
201        let url = Url::parse(&format!("{}/{api_version}/reconstructions/{}", self.endpoint, file_id.hex()))?;
202        let api_tag = match api_version {
203            "v1" => "cas::get_reconstruction_v1",
204            "v2" => "cas::get_reconstruction_v2",
205            _ => {
206                return Err(ClientError::InternalError(anyhow!(
207                    "unsupported reconstruction API version: {api_version}"
208                )));
209            },
210        };
211
212        event!(
213            INFORMATION_LOG_LEVEL,
214            call_id,
215            %file_id,
216            ?bytes_range,
217            api_version,
218            "Starting get_reconstruction API call",
219        );
220
221        let client = self.authenticated_http_client.clone();
222
223        let result: Result<T> = RetryWrapper::new(self.ctx.clone(), api_tag)
224            .with_expected_416()
225            .run_and_extract_json(move || {
226                let mut request = client.get(url.clone()).with_extension(Api(api_tag));
227                if let Some(range) = bytes_range {
228                    request = request.header(RANGE, HttpRange::from(range).range_header())
229                }
230                request.send()
231            })
232            .await;
233
234        match result {
235            Ok(response) => {
236                event!(
237                    INFORMATION_LOG_LEVEL,
238                    call_id,
239                    %file_id,
240                    ?bytes_range,
241                    api_version,
242                    "Completed get_reconstruction API call"
243                );
244                Ok(Some(response))
245            },
246            Err(ClientError::ReqwestError(ref e, _)) if e.status() == Some(StatusCode::RANGE_NOT_SATISFIABLE) => {
247                Ok(None)
248            },
249            Err(e) => Err(e),
250        }
251    }
252
253    /// V1 reconstruction: returns per-range presigned URLs.
254    pub async fn get_reconstruction_v1(
255        &self,
256        file_id: &MerkleHash,
257        bytes_range: Option<FileRange>,
258    ) -> Result<Option<QueryReconstructionResponse>> {
259        self.get_reconstruction_impl(file_id, bytes_range, "v1").await
260    }
261
262    /// V2 reconstruction: returns per-xorb multi-range fetch descriptors.
263    pub async fn get_reconstruction_v2(
264        &self,
265        file_id: &MerkleHash,
266        bytes_range: Option<FileRange>,
267    ) -> Result<Option<QueryReconstructionResponseV2>> {
268        self.get_reconstruction_impl(file_id, bytes_range, "v2").await
269    }
270
271    pub(crate) async fn get_reconstruction_with_version_override(
272        &self,
273        file_id: &MerkleHash,
274        bytes_range: Option<FileRange>,
275        forced_version: Option<u32>,
276    ) -> Result<Option<QueryReconstructionResponseV2>> {
277        // Prefer V2; fall back to V1 on 404/501; persist detected version to
278        // avoid repeated fallback attempts.
279        let version = match forced_version {
280            Some(v) => v,
281            None => {
282                let detected = self.detected_reconstruction_api_version.load(Ordering::Relaxed);
283                if detected != 0 { detected } else { 2 }
284            },
285        };
286
287        match version {
288            2 => match self.get_reconstruction_v2(file_id, bytes_range).await {
289                Ok(result) => {
290                    if forced_version.is_none() {
291                        self.detected_reconstruction_api_version.store(2, Ordering::Relaxed);
292                    }
293                    Ok(result)
294                },
295                Err(e)
296                    if forced_version.is_none()
297                        && matches!(e.status(), Some(StatusCode::NOT_FOUND) | Some(StatusCode::NOT_IMPLEMENTED)) =>
298                {
299                    info!(status = ?e.status(), "V2 reconstruction not available, falling back to V1");
300                    let result = self.get_reconstruction_v1(file_id, bytes_range).await?.map(Into::into);
301                    // Store after success to make sure we don't mess up on e.g. network failure.
302                    self.detected_reconstruction_api_version.store(1, Ordering::Relaxed);
303                    Ok(result)
304                },
305                Err(e) => Err(e),
306            },
307            1 => Ok(self.get_reconstruction_v1(file_id, bytes_range).await?.map(Into::into)),
308            other => Err(ClientError::InternalError(anyhow!("unsupported reconstruction API version: {other}"))),
309        }
310    }
311}
312
313#[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
314#[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
315impl Client for RemoteClient {
316    async fn get_reconstruction(
317        &self,
318        file_id: &MerkleHash,
319        bytes_range: Option<FileRange>,
320    ) -> Result<Option<QueryReconstructionResponseV2>> {
321        let forced_version = self.ctx.config.client.reconstruction_api_version;
322        self.get_reconstruction_with_version_override(file_id, bytes_range, forced_version)
323            .await
324    }
325
326    async fn batch_get_reconstruction(&self, file_ids: &[MerkleHash]) -> Result<BatchQueryReconstructionResponse> {
327        let mut url_str = format!("{}/reconstructions?", self.endpoint);
328        let mut is_first = true;
329        let mut file_id_list = Vec::new();
330        for hash in file_ids {
331            file_id_list.push(hash.hex());
332            if is_first {
333                is_first = false;
334            } else {
335                url_str.push('&');
336            }
337            url_str.push_str("file_id=");
338            url_str.push_str(hash.hex().as_str());
339        }
340        let url: Url = url_str.parse()?;
341
342        let call_id = FN_CALL_ID.fetch_add(1, Ordering::Relaxed);
343        info!(call_id, file_ids=?file_id_list, "Starting batch_get_reconstruction API call");
344
345        let api_tag = "cas::batch_get_reconstruction";
346        let client = self.authenticated_http_client.clone();
347
348        let response: BatchQueryReconstructionResponse = RetryWrapper::new(self.ctx.clone(), api_tag)
349            .run_and_extract_json(move || client.get(url.clone()).with_extension(Api(api_tag)).send())
350            .await?;
351
352        info!(call_id,
353            file_ids=?file_id_list,
354            response_count=response.files.len(),
355            "Completed batch_get_reconstruction API call",
356        );
357
358        Ok(response)
359    }
360
361    async fn acquire_download_permit(&self) -> Result<ConnectionPermit> {
362        self.download_concurrency_controller.acquire_connection_permit().await
363    }
364
365    async fn get_file_term_data(
366        &self,
367        url_info: Box<dyn URLProvider>,
368        download_permit: ConnectionPermit,
369        progress_callback: Option<ProgressCallback>,
370        uncompressed_size_if_known: Option<usize>,
371    ) -> Result<(Bytes, Vec<u32>)> {
372        let api_tag = "s3::get_range";
373        let http_client = self.http_client.clone();
374        let url_info = Arc::new(url_info);
375
376        let (_, url_ranges) = url_info.retrieve_url().await?;
377        let total_download_bytes: u64 = url_ranges.iter().map(|r| r.length()).sum();
378
379        let mut transfer_reporter = StreamProgressReporter::new(total_download_bytes)
380            .with_adaptive_concurrency_reporter(download_permit.get_partial_completion_reporting_function());
381        if let Some(cb) = progress_callback {
382            transfer_reporter = transfer_reporter.with_progress_callback(cb);
383        }
384
385        let result = RetryWrapper::new(self.ctx.clone(), api_tag)
386            .with_retry_on_403()
387            .with_connection_permit(download_permit, None)
388            .run_and_extract_custom(
389                move || {
390                    let http_client = http_client.clone();
391                    let url_info = url_info.clone();
392
393                    async move {
394                        let (url_string, url_ranges) = url_info
395                            .retrieve_url()
396                            .await
397                            .map_err(|e| reqwest_middleware::Error::Middleware(e.into()))?;
398                        let url =
399                            Url::parse(&url_string).map_err(|e| reqwest_middleware::Error::Middleware(e.into()))?;
400
401                        // RFC 7233 ยง2.1: single-range uses "bytes=S-E", multi-range uses "bytes=S1-E1,S2-E2,..."
402                        let range_header_value = if url_ranges.len() == 1 {
403                            url_ranges[0].range_header()
404                        } else {
405                            let joined = url_ranges
406                                .iter()
407                                .map(|r| format!("{}-{}", r.start, r.end))
408                                .collect::<Vec<_>>()
409                                .join(",");
410                            format!("bytes={joined}")
411                        };
412
413                        let response = http_client
414                            .get(url)
415                            .header(RANGE, range_header_value)
416                            .with_extension(Api(api_tag))
417                            .send()
418                            .await?;
419
420                        if response.status() == reqwest::StatusCode::FORBIDDEN {
421                            url_info
422                                .refresh_url()
423                                .await
424                                .map_err(|e| reqwest_middleware::Error::Middleware(e.into()))?;
425                        }
426
427                        Ok(response)
428                    }
429                },
430                move |resp: Response| {
431                    let transfer_reporter = transfer_reporter.clone();
432                    async move {
433                        let content_type = resp
434                            .headers()
435                            .get("content-type")
436                            .and_then(|v| v.to_str().ok())
437                            .unwrap_or("")
438                            .to_string();
439
440                        let is_multipart = content_type.contains("multipart/byteranges");
441
442                        if is_multipart {
443                            let body = resp
444                                .bytes()
445                                .await
446                                .map_err(|e| RetryableReqwestError::RetryableError(ClientError::from(e)))?;
447
448                            let multipart_parts = crate::cas_client::multipart::parse_multipart_byteranges(&content_type, body)
449                                .map_err(RetryableReqwestError::FatalError)?;
450
451                            let mut all_decompressed = Vec::with_capacity(uncompressed_size_if_known.unwrap_or(0));
452                            let mut all_chunk_indices = Vec::<u32>::new();
453                            let mut total_compressed_bytes = 0u64;
454
455                            for part in multipart_parts {
456                                total_compressed_bytes += part.data.len() as u64;
457
458                                let (data, chunk_indices) =
459                                    xet_core_structures::xorb_object::deserialize_chunks(&mut std::io::Cursor::new(part.data.as_ref()))
460                                        .map_err(|e| {
461                                            RetryableReqwestError::RetryableError(ClientError::FormatError(e))
462                                        })?;
463
464                                xet_core_structures::xorb_object::append_chunk_segment(
465                                    &mut all_decompressed,
466                                    &mut all_chunk_indices,
467                                    &data,
468                                    &chunk_indices,
469                                );
470
471                                transfer_reporter.report_progress(total_compressed_bytes as usize);
472                            }
473
474                            if let Some(expected) = uncompressed_size_if_known
475                                && expected != all_decompressed.len()
476                            {
477                                return Err(RetryableReqwestError::RetryableError(ClientError::Other(format!(
478                                    "get_file_term_data: expected {expected} uncompressed bytes, got {}",
479                                    all_decompressed.len()
480                                ))));
481                            }
482                            Ok((Bytes::from(all_decompressed), all_chunk_indices))
483                        } else {
484                            let incoming_stream = DownloadProgressStream::wrap_stream(
485                                resp.bytes_stream().map_err(std::io::Error::other),
486                                transfer_reporter,
487                            );
488
489                            let capacity = uncompressed_size_if_known.unwrap_or(0);
490                            let mut buffer = Vec::with_capacity(capacity);
491                            let mut writer = std::io::Cursor::new(&mut buffer);
492
493                            let result = xet_core_structures::xorb_object::deserialize_async::deserialize_chunks_to_writer_from_stream(
494                                incoming_stream,
495                                &mut writer,
496                            )
497                            .await;
498
499                            match result {
500                                Ok((_compressed_len, chunk_byte_indices)) => {
501                                    if let Some(expected) = uncompressed_size_if_known
502                                        && expected != buffer.len()
503                                    {
504                                        return Err(RetryableReqwestError::RetryableError(ClientError::Other(format!(
505                                            "get_file_term_data: expected {expected} uncompressed bytes, got {}",
506                                            buffer.len()
507                                        ))));
508                                    }
509                                    Ok((Bytes::from(buffer), chunk_byte_indices))
510                                },
511                                Err(e) => Err(RetryableReqwestError::RetryableError(ClientError::FormatError(e))),
512                            }
513                        }
514                    }
515                },
516            )
517            .await?;
518
519        Ok(result)
520    }
521
522    #[instrument(skip_all, name = "RemoteClient::get_file_reconstruction", fields(file.hash = file_hash.hex()
523    ))]
524    async fn get_file_reconstruction_info(
525        &self,
526        file_hash: &MerkleHash,
527    ) -> Result<Option<(MDBFileInfo, Option<MerkleHash>)>> {
528        let call_id = FN_CALL_ID.fetch_add(1, Ordering::Relaxed);
529        let url = Url::parse(&format!("{}/v1/reconstructions/{}", self.endpoint, file_hash.hex()))?;
530        event!(INFORMATION_LOG_LEVEL, call_id, %file_hash, "Starting get_file_reconstruction_info API call");
531
532        let api_tag = "cas::get_reconstruction_info";
533        let client = self.authenticated_http_client.clone();
534
535        let response: QueryReconstructionResponse = RetryWrapper::new(self.ctx.clone(), api_tag)
536            .run_and_extract_json(move || client.get(url.clone()).with_extension(Api(api_tag)).send())
537            .await?;
538
539        let terms_count = response.terms.len();
540        let result = Some((
541            MDBFileInfo {
542                metadata: FileDataSequenceHeader::new(*file_hash, terms_count, false, false),
543                segments: response
544                    .terms
545                    .into_iter()
546                    .map(|ce| {
547                        FileDataSequenceEntry::new(ce.hash.into(), ce.unpacked_length, ce.range.start, ce.range.end)
548                    })
549                    .collect(),
550                verification: vec![],
551                metadata_ext: None,
552            },
553            None,
554        ));
555
556        event!(INFORMATION_LOG_LEVEL, call_id, %file_hash, terms_count, "Completed get_file_reconstruction_info API call");
557
558        Ok(result)
559    }
560
561    async fn query_for_global_dedup_shard(&self, prefix: &str, chunk_hash: &MerkleHash) -> Result<Option<Bytes>> {
562        let Some(response) = self.query_dedup_api(prefix, chunk_hash).await? else {
563            return Ok(None);
564        };
565
566        Ok(Some(response.bytes().await?))
567    }
568
569    async fn acquire_upload_permit(&self) -> Result<ConnectionPermit> {
570        self.upload_concurrency_controller.acquire_connection_permit().await
571    }
572
573    #[instrument(skip_all, name = "RemoteClient::upload_shard", fields(shard.len = shard_data.len()))]
574    async fn upload_shard(&self, shard_data: Bytes, upload_permit: ConnectionPermit) -> Result<bool> {
575        if self.dry_run {
576            return Ok(true);
577        }
578
579        let call_id = FN_CALL_ID.fetch_add(1, Ordering::Relaxed);
580        let n_upload_bytes = shard_data.len();
581        event!(INFORMATION_LOG_LEVEL, call_id, size = n_upload_bytes, "Starting upload_shard API call",);
582
583        let api_tag = "cas::upload_shard";
584        let url = Url::parse(&format!("{}/v1/shards", self.endpoint))?;
585
586        // Use the no-read-timeout client for shard uploads. reqwest's per-request timeout()
587        // does NOT override the client-level read_timeout(), so we use a separate client
588        // with no read_timeout. Server-side shard processing scales linearly with file entry
589        // count and can exceed the global read_timeout (300s) for large shards.
590        #[cfg(not(target_family = "wasm"))]
591        let client = self.shard_upload_http_client.clone();
592
593        #[cfg(target_family = "wasm")]
594        let client = self.authenticated_http_client.clone();
595
596        let response: UploadShardResponse = RetryWrapper::new(self.ctx.clone(), api_tag)
597            .with_connection_permit(upload_permit, Some(shard_data.len() as u64))
598            .run_and_extract_json(move || {
599                client
600                    .post(url.clone())
601                    .with_extension(Api(api_tag))
602                    .body(shard_data.clone())
603                    .send()
604            })
605            .await?;
606
607        match response.result {
608            UploadShardResponseType::Exists => {
609                event!(
610                    INFORMATION_LOG_LEVEL,
611                    call_id,
612                    size = n_upload_bytes,
613                    result = "exists",
614                    "Completed upload_shard API call",
615                );
616                Ok(false)
617            },
618            UploadShardResponseType::SyncPerformed => {
619                event!(
620                    INFORMATION_LOG_LEVEL,
621                    call_id,
622                    size = n_upload_bytes,
623                    result = "sync performed",
624                    "Completed upload_shard API call",
625                );
626                Ok(true)
627            },
628        }
629    }
630
631    #[instrument(skip_all, name = "RemoteClient::upload_xorb", fields(key = Key{prefix : prefix.to_string(), hash : serialized_xorb_object.hash}.to_string(),
632                 xorb.len = serialized_xorb_object.serialized_data.len(), xorb.num_chunks = serialized_xorb_object.num_chunks
633    ))]
634    async fn upload_xorb(
635        &self,
636        prefix: &str,
637        serialized_xorb_object: SerializedXorbObject,
638        progress_callback: Option<ProgressCallback>,
639        upload_permit: ConnectionPermit,
640    ) -> Result<u64> {
641        let key = Key {
642            prefix: prefix.to_string(),
643            hash: serialized_xorb_object.hash,
644        };
645
646        let call_id = FN_CALL_ID.fetch_add(1, Ordering::Relaxed);
647        let url = Url::parse(&format!("{}/v1/xorbs/{key}", self.endpoint))?;
648
649        let n_upload_bytes = serialized_xorb_object.serialized_data.len() as u64;
650        event!(
651            INFORMATION_LOG_LEVEL,
652            call_id,
653            prefix,
654            hash=%serialized_xorb_object.hash,
655            size=n_upload_bytes,
656            num_chunks=serialized_xorb_object.num_chunks,
657            "Starting upload_xorb API call",
658        );
659
660        let n_transfer_bytes = serialized_xorb_object.serialized_data.len() as u64;
661
662        let serialized_data = serialized_xorb_object.serialized_data.clone();
663
664        #[cfg(not(target_family = "wasm"))]
665        let block_size = self.ctx.config.client.upload_reporting_block_size;
666
667        let mut upload_reporter = StreamProgressReporter::new(n_transfer_bytes)
668            .with_adaptive_concurrency_reporter(upload_permit.get_partial_completion_reporting_function());
669        if let Some(cb) = progress_callback {
670            upload_reporter = upload_reporter.with_progress_callback(cb);
671        }
672
673        let xorb_uploaded = {
674            if !self.dry_run {
675                let client = self.authenticated_http_client.clone();
676
677                let api_tag = "cas::upload_xorb";
678
679                let response: UploadXorbResponse = RetryWrapper::new(self.ctx.clone(), api_tag)
680                    .with_connection_permit(upload_permit, Some(n_transfer_bytes))
681                    .run_and_extract_json(move || {
682                        let url = url.clone();
683                        let serialized_data = serialized_data.clone();
684
685                        let request = {
686                            #[cfg(not(target_family = "wasm"))]
687                            {
688                                let upload_stream = UploadProgressStream::wrap_bytes_as_stream(
689                                    serialized_data,
690                                    block_size,
691                                    upload_reporter.clone(),
692                                );
693                                client
694                                    .post(url)
695                                    .with_extension(Api(api_tag))
696                                    .header(CONTENT_LENGTH, HeaderValue::from(n_upload_bytes)) // must be set because of streaming
697                                    .body(Body::wrap_stream(upload_stream))
698                            }
699
700                            // reqwest's wasm backend does not support streaming request bodies;
701                            // pass the raw Bytes directly (CONTENT_LENGTH is set by reqwest from the body length).
702                            #[cfg(target_family = "wasm")]
703                            {
704                                client.post(url).with_extension(Api(api_tag)).body(serialized_data)
705                            }
706                        };
707
708                        request.send()
709                    })
710                    .await?;
711
712                // Wasm has no per-chunk progress hook (no streaming body); emit one bulk
713                // event after success so the user callback and adaptive-concurrency
714                // reporter both observe the full byte count.
715                #[cfg(target_family = "wasm")]
716                upload_reporter.report_progress(n_transfer_bytes as usize);
717
718                response.was_inserted
719            } else {
720                true
721            }
722        };
723
724        if !xorb_uploaded {
725            event!(
726                INFORMATION_LOG_LEVEL,
727                call_id,
728                prefix,
729                hash=%serialized_xorb_object.hash,
730                result="not_inserted",
731                "Completed upload_xorb API call",
732            );
733        } else {
734            event!(
735                INFORMATION_LOG_LEVEL,
736                call_id,
737                prefix,
738                hash=%serialized_xorb_object.hash,
739                size=n_upload_bytes,
740                result="inserted",
741                "Completed upload_xorb API call",
742            );
743        }
744
745        Ok(n_upload_bytes)
746    }
747
748    #[instrument(skip_all, name = "RemoteClient::get_file_chunk_hashes", fields(file.hash = file_id.hex(), n_ranges = dirty_ranges.len()))]
749    async fn get_file_chunk_hashes(
750        &self,
751        file_id: &MerkleHash,
752        dirty_ranges: Vec<FileRange>,
753    ) -> Result<FileChunkHashesResponse> {
754        if dirty_ranges.is_empty() {
755            return Err(ClientError::Other("get_file_chunk_hashes requires at least one dirty range".into()));
756        }
757
758        let url = Url::parse(&format!("{}/v2/file-chunk-hashes/{}", self.endpoint, file_id.hex()))?;
759
760        // Multi-range `bytes=A-B,C-D` value. `HttpRange` is inclusive-end and `Display`s as
761        // `start-end`; conversion from `FileRange` does the +1/-1 for us.
762        let header_value = HeaderValue::from_str(&format!(
763            "bytes={}",
764            dirty_ranges
765                .iter()
766                .copied()
767                .map(HttpRange::from)
768                .map(|r| r.to_string())
769                .collect::<Vec<_>>()
770                .join(",")
771        ))
772        .map_err(|err| ClientError::Other(format!("invalid X-Range-Dirty header value: {err}")))?;
773
774        let api_tag = "cas::get_file_chunk_hashes";
775        let client = self.authenticated_http_client.clone();
776
777        let response: FileChunkHashesResponse = RetryWrapper::new(self.ctx.clone(), api_tag)
778            .run_and_extract_json(move || {
779                client
780                    .get(url.clone())
781                    .header(X_RANGE_DIRTY_HEADER, header_value.clone())
782                    .with_extension(Api(api_tag))
783                    .send()
784            })
785            .await?;
786
787        Ok(response)
788    }
789}
790
791#[cfg(test)]
792#[cfg(not(target_family = "wasm"))]
793mod tests {
794    use tracing_test::traced_test;
795    use xet_core_structures::xorb_object::CompressionScheme;
796    use xet_core_structures::xorb_object::xorb_format_test_utils::{
797        ChunkSize, build_and_verify_xorb_object, build_raw_xorb,
798    };
799
800    use super::*;
801
802    #[test]
803    fn test_clients_share_controllers_per_ctx_and_endpoint() {
804        let ctx = XetContext::default().unwrap();
805        let c1 = RemoteClient::new(ctx.clone(), "https://cas-a.example.com", &None, "", false, None);
806        let c2 = RemoteClient::new(ctx.clone(), "https://cas-a.example.com", &None, "", false, None);
807
808        // Same ctx + same endpoint: shared upload and download controllers.
809        assert!(Arc::ptr_eq(&c1.upload_concurrency_controller, &c2.upload_concurrency_controller));
810        assert!(Arc::ptr_eq(&c1.download_concurrency_controller, &c2.download_concurrency_controller));
811
812        // Same ctx, different endpoint: independent controllers.
813        let c3 = RemoteClient::new(ctx.clone(), "https://cas-b.example.com", &None, "", false, None);
814        assert!(!Arc::ptr_eq(&c1.upload_concurrency_controller, &c3.upload_concurrency_controller));
815
816        // Creating a second endpoint must not evict the first: re-fetching cas-a still shares with c1.
817        let c5 = RemoteClient::new(ctx.clone(), "https://cas-a.example.com", &None, "", false, None);
818        assert!(Arc::ptr_eq(&c1.upload_concurrency_controller, &c5.upload_concurrency_controller));
819
820        // Different ctx (different session), same endpoint: independent controllers.
821        let ctx2 = XetContext::default().unwrap();
822        let c4 = RemoteClient::new(ctx2, "https://cas-a.example.com", &None, "", false, None);
823        assert!(!Arc::ptr_eq(&c1.upload_concurrency_controller, &c4.upload_concurrency_controller));
824    }
825
826    #[ignore = "requires a running CAS server"]
827    #[traced_test]
828    #[test]
829    fn test_basic_put() {
830        // Arrange
831        let prefix = PREFIX_DEFAULT;
832        let raw_xorb = build_raw_xorb(3, ChunkSize::Random(512, 10248));
833
834        let ctx = XetContext::default().unwrap();
835        let client = RemoteClient::new(ctx.clone(), CAS_ENDPOINT, &None, "", false, None);
836
837        let xorb_obj = build_and_verify_xorb_object(raw_xorb, CompressionScheme::LZ4);
838
839        // Act
840        let result = ctx
841            .runtime
842            .bridge_sync(async move {
843                let permit = client.acquire_upload_permit().await.unwrap();
844                client.upload_xorb(prefix, xorb_obj, None, permit).await
845            })
846            .unwrap();
847
848        // Assert
849        assert!(result.is_ok());
850    }
851}