hf_hub/
client.rs

1use std::borrow::Cow;
2use std::sync::Arc;
3use std::time::Duration;
4
5use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, utf8_percent_encode};
6use reqwest::header::{AUTHORIZATION, HeaderMap, HeaderValue, USER_AGENT};
7use tracing::debug;
8use url::Url;
9
10use crate::constants;
11use crate::error::{HFError, HFResult, HttpErrorContext, NotFoundContext};
12use crate::retry::{self, RetryConfig};
13
14/// Append `path` to `url` as percent-encoded URL path segments.
15///
16/// Splits `path` on `/` and pushes each non-empty piece via
17/// [`Url::path_segments_mut`], which percent-encodes characters that aren't
18/// valid in a path segment (e.g. `#`, `?`, ` `, `%`).
19///
20/// Returns an error only if `url` is a cannot-be-a-base URL (e.g. `mailto:`,
21/// `data:`); all URLs hf-hub constructs are http(s) and so this never fails
22/// in practice.
23pub(crate) fn append_path_segments(url: &mut Url, path: &str) -> HFResult<()> {
24    let url_str = url.as_str().to_string();
25    let mut segments = url
26        .path_segments_mut()
27        .map_err(|()| HFError::InvalidParameter(format!("unexpected non-base URL: {url_str}")))?;
28    for seg in path.split('/').filter(|s| !s.is_empty()) {
29        segments.push(seg);
30    }
31    Ok(())
32}
33
34/// Everything except RFC 3986 unreserved characters (`A-Za-z0-9-._~`), mirroring
35/// Python's `urllib.parse.quote(s, safe="")` as used by `huggingface_hub`.
36const REF_ENCODE_SET: &AsciiSet = &NON_ALPHANUMERIC.remove(b'-').remove(b'.').remove(b'_').remove(b'~');
37
38/// Percent-encode a git revision (branch, tag, or commit SHA) for use as a single
39/// URL path segment.
40///
41/// Unlike [`append_path_segments`], `/` is encoded too (`feature/x` → `feature%2Fx`)
42/// so the revision occupies exactly one segment — the Hub API expects revisions
43/// fully encoded. Also safe for `<base>..<head>` compare specs: `.` is unreserved,
44/// so the `..` separator stays literal while each side is encoded.
45pub(crate) fn encode_ref(revision: &str) -> Cow<'_, str> {
46    utf8_percent_encode(revision, REF_ENCODE_SET).into()
47}
48
49/// Split a repo or bucket id into its `(owner, name)` parts.
50///
51/// A fully-qualified id (`"owner/name"`) splits on the first `/`. A short-form
52/// id with no owner (`"gpt2"`) yields an empty owner, matching the `owner, name`
53/// argument pair taken by [`HFClient::model`], [`HFClient::dataset`],
54/// [`HFClient::space`], and [`HFClient::repository`]. Only the first `/` is a
55/// separator, so nested names (`"owner/sub/name"`) keep their tail intact
56/// (`("owner", "sub/name")`).
57///
58/// ```
59/// use hf_hub::split_id;
60///
61/// assert_eq!(split_id("openai-community/gpt2"), ("openai-community", "gpt2"));
62/// assert_eq!(split_id("gpt2"), ("", "gpt2"));
63/// ```
64pub fn split_id(id: &str) -> (&str, &str) {
65    match id.split_once('/') {
66        Some((owner, name)) => (owner, name),
67        None => ("", id),
68    }
69}
70
71/// Whether a `Location` header value is a relative reference (no scheme, no
72/// host), like `/Snowflake/repo/resolve/main/config.json`. Protocol-relative
73/// URLs (`//host/path`) carry a host and are treated as absolute.
74fn is_relative_location(location: &str) -> bool {
75    Url::parse(location) == Err(url::ParseError::RelativeUrlWithoutBase) && !location.starts_with("//")
76}
77
78/// Async client for the Hugging Face Hub API.
79///
80/// `HFClient` wraps an `Arc<HFClientInner>` so it is cheap to clone — all clones
81/// share the same underlying HTTP connection pool, token, and configuration.
82///
83/// # Creating a client
84///
85/// ```rust,no_run
86/// use hf_hub::HFClient;
87///
88/// // Reads token and settings from the environment (HF_TOKEN, HF_ENDPOINT, …).
89/// let client = HFClient::new()?;
90///
91/// // Or configure explicitly:
92/// let client = HFClient::builder().token("hf_…").endpoint("https://huggingface.co").build()?;
93/// # Ok::<(), hf_hub::HFError>(())
94/// ```
95pub struct HFClient {
96    pub(crate) inner: Arc<HFClientInner>,
97}
98
99impl Clone for HFClient {
100    fn clone(&self) -> Self {
101        Self {
102            inner: Arc::clone(&self.inner),
103        }
104    }
105}
106
107impl std::fmt::Debug for HFClient {
108    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109        let mut s = f.debug_struct("HFClient");
110        s.field("endpoint", &self.inner.endpoint);
111        s.field("token_set", &self.inner.token.is_some());
112        s.field("cache_enabled", &self.inner.cache_enabled);
113        if self.inner.cache_enabled {
114            s.field("cache_dir", &self.inner.cache_dir);
115        }
116        s.finish()
117    }
118}
119
120pub(crate) struct HFClientInner {
121    pub(crate) client: reqwest::Client,
122    // The wasm32 reqwest backend ignores `redirect()` configuration — the browser
123    // (or fetch in Node) handles redirects opaquely — so we only carry a separate
124    // no-redirect client on native targets.
125    #[cfg(not(target_family = "wasm"))]
126    pub(crate) no_redirect_client: reqwest::Client,
127    pub(crate) retry_config: RetryConfig,
128    pub(crate) endpoint: String,
129    pub(crate) token: Option<String>,
130    pub(crate) cache_dir: std::path::PathBuf,
131    pub(crate) cache_enabled: bool,
132    pub(crate) xet_state: std::sync::Mutex<crate::xet::XetState>,
133}
134
135/// Builder for [`HFClient`].
136///
137/// Get one via [`HFClient::builder()`] or [`HFClientBuilder::new()`].
138/// Call [`build()`](HFClientBuilder::build) when all options are set.
139pub struct HFClientBuilder {
140    endpoint: Option<String>,
141    token: Option<String>,
142    user_agent: Option<String>,
143    headers: Option<HeaderMap>,
144    client: Option<reqwest::Client>,
145    cache_dir: Option<std::path::PathBuf>,
146    cache_enabled: Option<bool>,
147    retry_max_attempts: Option<usize>,
148    retry_base_delay: Option<Duration>,
149}
150
151impl HFClientBuilder {
152    /// Creates a builder with all options unset; defaults are applied at [`build`](Self::build) time.
153    pub fn new() -> Self {
154        Self {
155            endpoint: None,
156            token: None,
157            user_agent: None,
158            headers: None,
159            client: None,
160            cache_dir: None,
161            cache_enabled: None,
162            retry_max_attempts: None,
163            retry_base_delay: None,
164        }
165    }
166
167    /// Overrides the Hub base URL (default: `https://huggingface.co`, or `HF_ENDPOINT` env var).
168    pub fn endpoint(mut self, endpoint: impl Into<String>) -> Self {
169        self.endpoint = Some(endpoint.into());
170        self
171    }
172
173    /// Sets the authentication token. Without this, the client falls back to the `HF_TOKEN` env
174    /// var and the cached token file written by `huggingface-cli login`.
175    pub fn token(mut self, token: impl Into<String>) -> Self {
176        self.token = Some(token.into());
177        self
178    }
179
180    /// Overrides the `User-Agent` header sent with every request.
181    pub fn user_agent(mut self, user_agent: impl Into<String>) -> Self {
182        self.user_agent = Some(user_agent.into());
183        self
184    }
185
186    /// Merges additional default headers into every request. Ignored when a custom
187    /// [`client`](Self::client) is supplied (configure headers on that client directly).
188    pub fn headers(mut self, headers: HeaderMap) -> Self {
189        self.headers = Some(headers);
190        self
191    }
192
193    /// Supplies a pre-configured `reqwest::Client`. Retry middleware is still applied on top.
194    /// The caller is responsible for any default headers (including `User-Agent`) on this client;
195    /// the [`headers`](Self::headers) and [`user_agent`](Self::user_agent) options are ignored.
196    pub fn client(mut self, client: reqwest::Client) -> Self {
197        self.client = Some(client);
198        self
199    }
200
201    /// Sets the local cache directory. Defaults to `HF_HUB_CACHE` → `$HF_HOME/hub` →
202    /// `~/.cache/huggingface/hub`.
203    pub fn cache_dir(mut self, path: impl Into<std::path::PathBuf>) -> Self {
204        self.cache_dir = Some(path.into());
205        self
206    }
207
208    /// Enables or disables the local file cache. Caching is on by default.
209    pub fn cache_enabled(mut self, enabled: bool) -> Self {
210        self.cache_enabled = Some(enabled);
211        self
212    }
213
214    /// Overrides the maximum number of retry attempts after an initial failure (default: 3).
215    pub fn retry_max_attempts(mut self, n: usize) -> Self {
216        self.retry_max_attempts = Some(n);
217        self
218    }
219
220    /// Overrides the base delay for exponential backoff between retries (default: 100ms).
221    pub fn retry_base_delay(mut self, delay: Duration) -> Self {
222        self.retry_base_delay = Some(delay);
223        self
224    }
225
226    /// Builds the [`HFClient`].
227    ///
228    /// # Errors
229    ///
230    /// Returns an error if the endpoint URL is not a valid URL or if the `reqwest` client
231    /// cannot be constructed (e.g., an invalid `User-Agent` string was provided).
232    pub fn build(self) -> HFResult<HFClient> {
233        let endpoint = self
234            .endpoint
235            .or_else(|| std::env::var(constants::HF_ENDPOINT).ok())
236            .unwrap_or_else(|| constants::DEFAULT_HF_ENDPOINT.to_string());
237
238        let _ = url::Url::parse(&endpoint)?;
239
240        let token = self.token.or_else(resolve_token);
241
242        let cache_dir = self.cache_dir.unwrap_or_else(constants::resolve_cache_dir);
243
244        let mut default_headers = self.headers.unwrap_or_default();
245
246        let user_agent = self.user_agent.unwrap_or_else(|| {
247            let ua_origin = std::env::var(constants::HF_HUB_USER_AGENT_ORIGIN).ok();
248            match ua_origin {
249                Some(origin) => format!("hf-hub/{}; {origin}", env!("CARGO_PKG_VERSION")),
250                None => format!("hf-hub/{}", env!("CARGO_PKG_VERSION")),
251            }
252        });
253        default_headers.insert(
254            USER_AGENT,
255            HeaderValue::from_str(&user_agent)
256                .map_err(|e| HFError::InvalidParameter(format!("invalid user agent {user_agent:?}: {e}")))?,
257        );
258
259        let client = match self.client {
260            Some(c) => c,
261            None => reqwest::Client::builder().default_headers(default_headers.clone()).build()?,
262        };
263
264        #[cfg(not(target_family = "wasm"))]
265        let no_redirect_client = reqwest::Client::builder()
266            .default_headers(default_headers)
267            .redirect(reqwest::redirect::Policy::none())
268            .build()?;
269
270        let retry_config = RetryConfig {
271            max_attempts: self.retry_max_attempts.unwrap_or(retry::DEFAULT_MAX_ATTEMPTS),
272            base_delay: self.retry_base_delay.unwrap_or(retry::DEFAULT_BASE_DELAY),
273        };
274
275        Ok(HFClient {
276            inner: Arc::new(HFClientInner {
277                client,
278                #[cfg(not(target_family = "wasm"))]
279                no_redirect_client,
280                retry_config,
281                endpoint: endpoint.trim_end_matches('/').to_string(),
282                token,
283                cache_dir,
284                cache_enabled: self.cache_enabled.unwrap_or(true),
285                xet_state: std::sync::Mutex::new(crate::xet::XetState::default()),
286            }),
287        })
288    }
289
290    /// Builds the [`crate::blocking::HFClientSync`].
291    ///
292    /// Requires the `blocking` feature and is equivalent to calling
293    /// [`build`](Self::build) and then [`HFClientSync::from_inner`](crate::HFClientSync::from_inner).
294    ///
295    /// # Errors
296    ///
297    /// Returns an error if the endpoint URL is not a valid URL, if the `reqwest` client
298    /// cannot be constructed (e.g., an invalid `User-Agent` string was provided), or if the
299    /// tokio runtime handle could not be correctly created for the blocking client.
300    #[cfg(feature = "blocking")]
301    #[cfg_attr(docsrs, doc(cfg(feature = "blocking")))]
302    pub fn build_sync(self) -> HFResult<crate::blocking::HFClientSync> {
303        let async_client = self.build()?;
304        let client = crate::blocking::HFClientSync::from_inner(async_client)?;
305        Ok(client)
306    }
307}
308
309impl Default for HFClientBuilder {
310    fn default() -> Self {
311        Self::new()
312    }
313}
314
315impl HFClient {
316    /// Creates a client with default settings, reading the token and endpoint from the
317    /// environment. Equivalent to `HFClient::builder().build()`.
318    ///
319    /// # Errors
320    ///
321    /// Fails if the resolved endpoint URL is invalid or the HTTP client cannot be built.
322    pub fn new() -> HFResult<Self> {
323        HFClientBuilder::new().build()
324    }
325
326    /// Returns an [`HFClientBuilder`] for fine-grained configuration.
327    pub fn builder() -> HFClientBuilder {
328        HFClientBuilder::new()
329    }
330
331    pub(crate) fn http_client(&self) -> &reqwest::Client {
332        &self.inner.client
333    }
334
335    #[cfg(not(target_family = "wasm"))]
336    pub(crate) fn no_redirect_client(&self) -> &reqwest::Client {
337        &self.inner.no_redirect_client
338    }
339
340    pub(crate) fn retry_config(&self) -> &RetryConfig {
341        &self.inner.retry_config
342    }
343
344    /// Hub base URL this client targets, with any trailing slash trimmed.
345    ///
346    /// Resolved at [`build`](HFClientBuilder::build) time from
347    /// [`HFClientBuilder::endpoint`] → `HF_ENDPOINT` → the default
348    /// (`https://huggingface.co`).
349    pub fn endpoint(&self) -> &str {
350        &self.inner.endpoint
351    }
352
353    /// Local cache directory used for downloaded files.
354    ///
355    /// Resolved at [`build`](HFClientBuilder::build) time from
356    /// [`HFClientBuilder::cache_dir`] → `HF_HUB_CACHE` → `$HF_HOME/hub` →
357    /// `~/.cache/huggingface/hub`. Returned even when caching is disabled —
358    /// see [`cache_enabled`](Self::cache_enabled) to check that.
359    pub fn cache_dir(&self) -> &std::path::Path {
360        &self.inner.cache_dir
361    }
362
363    /// Whether the local file cache is enabled. Set via
364    /// [`HFClientBuilder::cache_enabled`]; defaults to `true`.
365    pub fn cache_enabled(&self) -> bool {
366        self.inner.cache_enabled
367    }
368
369    /// Build authorization headers for requests
370    pub(crate) fn auth_headers(&self) -> HeaderMap {
371        let mut headers = HeaderMap::new();
372        if let Some(ref token) = self.inner.token
373            && let Ok(val) = HeaderValue::from_str(&format!("Bearer {token}"))
374        {
375            headers.insert(AUTHORIZATION, val);
376        }
377        headers
378    }
379
380    /// Build a URL for the API: {endpoint}/api/{api_segment}/{repo_id}.
381    ///
382    /// `api_segment` is the plural form of the repo type — `"models"`, `"datasets"`,
383    /// `"spaces"`, or `"kernels"`. Pass [`crate::repository::RepoType::plural`] for the
384    /// concrete repo type at the call site.
385    pub(crate) fn api_url(&self, api_segment: &str, repo_id: &str) -> String {
386        format!("{}/api/{}/{}", self.endpoint(), api_segment, repo_id)
387    }
388
389    /// Build a download URL: `{endpoint}/{url_prefix}{repo_id}/resolve/{revision}/{filename}`.
390    ///
391    /// `url_prefix` is `""` for models or `"<plural>/"` for the other repo types. Pass
392    /// [`crate::repository::RepoType::url_prefix`] for the concrete repo type at the call
393    /// site. `revision` is fully percent-encoded (including `/`) so it occupies exactly one
394    /// path segment; `filename` is percent-encoded path-segment by path-segment, so values
395    /// like `subdir/file name #1.bin` produce a correctly escaped URL.
396    pub(crate) fn download_url(
397        &self,
398        url_prefix: &str,
399        repo_id: &str,
400        revision: &str,
401        filename: &str,
402    ) -> HFResult<String> {
403        let base = format!("{}/{}{}/resolve/{}", self.endpoint(), url_prefix, repo_id, encode_ref(revision));
404        let mut url = Url::parse(&base)?;
405        append_path_segments(&mut url, filename)?;
406        Ok(url.into())
407    }
408
409    /// Build a bucket API URL: `{endpoint}/api/buckets/{bucket_id}`
410    pub(crate) fn bucket_api_url(&self, bucket_id: &str) -> String {
411        format!("{}/api/buckets/{}", self.endpoint(), bucket_id)
412    }
413
414    /// Issue a HEAD request via the no-redirect client, manually following
415    /// *relative* redirects (Location without a host).
416    ///
417    /// The Hub 307s to the canonical repo path when a repo id differs in
418    /// casing (e.g. `snowflake/…` → `Snowflake/…`); that response carries no
419    /// ETag, so metadata extraction must happen after following it. Absolute
420    /// redirects (e.g. to the CDN) are returned as-is — their headers already
421    /// carry the linked metadata. Mirrors `huggingface_hub`'s
422    /// `follow_relative_redirects` behavior.
423    #[cfg(not(target_family = "wasm"))]
424    pub(crate) async fn head_with_relative_redirects(
425        &self,
426        url: &str,
427        headers: &HeaderMap,
428    ) -> HFResult<reqwest::Response> {
429        const MAX_RELATIVE_REDIRECTS: usize = 10;
430        let mut url = Url::parse(url)?;
431        for _ in 0..=MAX_RELATIVE_REDIRECTS {
432            let response = retry::retry(self.retry_config(), || {
433                self.no_redirect_client().head(url.clone()).headers(headers.clone()).send()
434            })
435            .await?;
436            if response.status().is_redirection()
437                && let Some(location) = response.headers().get(reqwest::header::LOCATION)
438                && let Ok(location) = location.to_str()
439                && is_relative_location(location)
440            {
441                url = url.join(location)?;
442                continue;
443            }
444            return Ok(response);
445        }
446        Err(HFError::malformed_response_at("too many relative redirects", url.to_string()))
447    }
448
449    /// Check an HTTP response and map error status codes to HFError variants.
450    /// Returns the response on success (2xx).
451    ///
452    /// `repo_id` and `not_found_ctx` control how 404s are mapped:
453    /// - `NotFoundContext::Repo` → `HFError::RepoNotFound`
454    /// - `NotFoundContext::Entry { path }` → `HFError::EntryNotFound`
455    /// - `NotFoundContext::Revision { revision }` → `HFError::RevisionNotFound`
456    /// - `NotFoundContext::Generic` → `HFError::Http`
457    pub(crate) async fn check_response(
458        &self,
459        response: reqwest::Response,
460        repo_id: Option<&str>,
461        not_found_ctx: NotFoundContext,
462    ) -> HFResult<reqwest::Response> {
463        let status = response.status();
464        if status.is_success() {
465            return Ok(response);
466        }
467
468        let retry_after = if status == reqwest::StatusCode::TOO_MANY_REQUESTS {
469            crate::error::parse_retry_after(response.headers())
470        } else {
471            None
472        };
473        let context = Box::new(HttpErrorContext::from_response(response).await);
474        let repo_id_str = repo_id.unwrap_or("").to_string();
475
476        match status.as_u16() {
477            401 => Err(HFError::AuthRequired { context }),
478            403 => Err(HFError::Forbidden { context }),
479            404 => match not_found_ctx {
480                NotFoundContext::Repo => Err(HFError::RepoNotFound {
481                    repo_id: repo_id_str,
482                    context: Some(context),
483                }),
484                NotFoundContext::Bucket => Err(HFError::BucketNotFound {
485                    bucket_id: repo_id_str,
486                    context: Some(context),
487                }),
488                NotFoundContext::Entry { path } => Err(HFError::EntryNotFound {
489                    path,
490                    repo_id: repo_id_str,
491                    context: Some(context),
492                }),
493                NotFoundContext::Revision { revision } => Err(HFError::RevisionNotFound {
494                    revision,
495                    repo_id: repo_id_str,
496                    context: Some(context),
497                }),
498                NotFoundContext::Generic => Err(HFError::Http { context }),
499            },
500            409 => Err(HFError::Conflict { context }),
501            429 => Err(HFError::RateLimited { retry_after, context }),
502            _ => Err(HFError::Http { context }),
503        }
504    }
505
506    /// Get or lazily create the cached XetSession.
507    ///
508    /// Returns `(session, generation)`. The generation is an opaque counter
509    /// that identifies which session instance this is. Pass it to
510    /// [`replace_xet_session`](Self::replace_xet_session) so that only the
511    /// caller that observed the error triggers a replacement — concurrent
512    /// callers that already got a session won't clobber it.
513    pub(crate) fn xet_session(&self) -> HFResult<(xet::xet_session::XetSession, u64)> {
514        let mut guard = self
515            .inner
516            .xet_state
517            .lock()
518            .map_err(|e| HFError::Other(format!("xet session mutex poisoned: {e}")))?;
519
520        if let Some(ref session) = guard.session {
521            return Ok((session.clone(), guard.generation));
522        }
523
524        #[cfg(not(target_family = "wasm"))]
525        let builder = xet::xet_session::XetSessionBuilder::new();
526        // Cap transfer concurrency on wasm to avoid exhausting linear memory.
527        #[cfg(target_family = "wasm")]
528        let builder = {
529            let mut config = xet::xet_session::XetConfig::new();
530            config.client.ac_max_upload_concurrency = 8;
531            config.client.ac_max_download_concurrency = 8;
532            xet::xet_session::XetSessionBuilder::new_with_config(config)
533        };
534        let session = builder
535            .build()
536            .map_err(|e| HFError::xet(crate::error::XetOperation::Session, e))?;
537        guard.session = Some(session.clone());
538        guard.generation += 1;
539        Ok((session, guard.generation))
540    }
541
542    /// Replace the cached XetSession only if the generation matches.
543    ///
544    /// Called by xet call sites when a factory method returns an error.
545    /// The generation check ensures that if another thread already replaced
546    /// the session, this call is a no-op rather than discarding the fresh one.
547    pub(crate) fn replace_xet_session(&self, generation: u64, err: &xet::error::XetError) {
548        tracing::warn!(error = %err, generation, "replacing cached XetSession");
549        let Ok(mut guard) = self.inner.xet_state.lock() else {
550            return;
551        };
552        if guard.generation == generation {
553            guard.session = None;
554        }
555    }
556}
557
558/// Resolve token from environment or a token file.
559/// Priority: HF_TOKEN env → HF_TOKEN_PATH file → $HF_HOME/token file.
560fn resolve_token() -> Option<String> {
561    if let Ok(val) = std::env::var(constants::HF_HUB_DISABLE_IMPLICIT_TOKEN)
562        && !val.is_empty()
563    {
564        debug!("implicit token disabled via HF_HUB_DISABLE_IMPLICIT_TOKEN");
565        return None;
566    }
567
568    if let Ok(token) = std::env::var(constants::HF_TOKEN)
569        && !token.is_empty()
570    {
571        debug!("resolved token from HF_TOKEN env var");
572        return Some(token);
573    }
574
575    if let Ok(path) = std::env::var(constants::HF_TOKEN_PATH)
576        && let Ok(token) = std::fs::read_to_string(&path)
577    {
578        let token = token.trim().to_string();
579        if !token.is_empty() {
580            debug!("resolved token from HF_TOKEN_PATH file");
581            return Some(token);
582        }
583    }
584
585    let hf_home = constants::hf_home();
586    let token_path = hf_home.join(constants::TOKEN_FILENAME);
587    if let Ok(token) = std::fs::read_to_string(&token_path) {
588        let token = token.trim().to_string();
589        if !token.is_empty() {
590            debug!("resolved token from stored token file");
591            return Some(token);
592        }
593    }
594
595    debug!("no token found");
596    None
597}
598
599#[cfg(test)]
600mod tests {
601    use serial_test::serial;
602    use url::Url;
603
604    use super::{HFClientBuilder, append_path_segments, encode_ref, is_relative_location, split_id};
605
606    #[test]
607    fn append_path_segments_encodes_special_chars() {
608        let mut url = Url::parse("https://example.com/base").unwrap();
609        append_path_segments(&mut url, "with space/and#hash?and+plus").unwrap();
610        assert_eq!(url.as_str(), "https://example.com/base/with%20space/and%23hash%3Fand+plus");
611    }
612
613    #[test]
614    fn append_path_segments_handles_leading_trailing_and_repeated_slashes() {
615        let mut url = Url::parse("https://example.com/base").unwrap();
616        append_path_segments(&mut url, "/a//b/c/").unwrap();
617        assert_eq!(url.as_str(), "https://example.com/base/a/b/c");
618    }
619
620    #[test]
621    fn append_path_segments_encodes_percent_literals() {
622        let mut url = Url::parse("https://example.com/base").unwrap();
623        append_path_segments(&mut url, "a%2Fb/c").unwrap();
624        assert_eq!(url.as_str(), "https://example.com/base/a%252Fb/c");
625    }
626
627    #[test]
628    fn append_path_segments_handles_unicode() {
629        let mut url = Url::parse("https://example.com/base").unwrap();
630        append_path_segments(&mut url, "ümlaut/файл.txt").unwrap();
631        assert_eq!(url.as_str(), "https://example.com/base/%C3%BCmlaut/%D1%84%D0%B0%D0%B9%D0%BB.txt");
632    }
633
634    #[test]
635    fn append_path_segments_empty_path_is_noop() {
636        let mut url = Url::parse("https://example.com/base").unwrap();
637        append_path_segments(&mut url, "").unwrap();
638        assert_eq!(url.as_str(), "https://example.com/base");
639    }
640
641    #[test]
642    fn encode_ref_encodes_slash_and_reserved_chars() {
643        assert_eq!(encode_ref("feature/x"), "feature%2Fx");
644        assert_eq!(encode_ref("fix#123"), "fix%23123");
645        assert_eq!(encode_ref("wip?draft"), "wip%3Fdraft");
646        assert_eq!(encode_ref("v1.0+beta"), "v1.0%2Bbeta");
647        assert_eq!(encode_ref("with space"), "with%20space");
648        assert_eq!(encode_ref("100%done"), "100%25done");
649    }
650
651    #[test]
652    fn encode_ref_preserves_unreserved_chars() {
653        assert_eq!(encode_ref("main"), "main");
654        assert_eq!(encode_ref("v1.0-rc.2_final~ok"), "v1.0-rc.2_final~ok");
655        assert_eq!(encode_ref("abc123def456"), "abc123def456");
656    }
657
658    #[test]
659    fn encode_ref_keeps_compare_dots_literal() {
660        assert_eq!(encode_ref("main..feature/x"), "main..feature%2Fx");
661        assert_eq!(encode_ref("v1.0...v2.0"), "v1.0...v2.0");
662    }
663
664    #[test]
665    fn encode_ref_encodes_unicode() {
666        assert_eq!(encode_ref("ветка"), "%D0%B2%D0%B5%D1%82%D0%BA%D0%B0");
667    }
668
669    #[test]
670    fn is_relative_location_classification() {
671        assert!(is_relative_location("/Snowflake/repo/resolve/main/config.json"));
672        assert!(is_relative_location("relative/path"));
673        assert!(!is_relative_location("https://cdn-lfs.huggingface.co/repos/x"));
674        assert!(!is_relative_location("//cdn-lfs.huggingface.co/repos/x"));
675    }
676
677    #[test]
678    fn download_url_encodes_revision() {
679        let client = HFClientBuilder::new().endpoint("https://huggingface.co").build().unwrap();
680        let url = client.download_url("", "owner/repo", "feature/x", "file.bin").unwrap();
681        assert_eq!(url, "https://huggingface.co/owner/repo/resolve/feature%2Fx/file.bin");
682
683        let url = client.download_url("", "owner/repo", "fix#123", "file.bin").unwrap();
684        assert_eq!(url, "https://huggingface.co/owner/repo/resolve/fix%23123/file.bin");
685    }
686
687    #[test]
688    fn download_url_encodes_filename_segments() {
689        let client = HFClientBuilder::new().endpoint("https://huggingface.co").build().unwrap();
690        let url = client
691            .download_url("", "owner/repo", "main", "subdir/file name #1.bin")
692            .unwrap();
693        assert_eq!(url, "https://huggingface.co/owner/repo/resolve/main/subdir/file%20name%20%231.bin");
694    }
695
696    #[test]
697    fn download_url_dataset_prefix() {
698        let client = HFClientBuilder::new().endpoint("https://huggingface.co").build().unwrap();
699        let url = client.download_url("datasets/", "owner/ds", "v1.0", "data/train.csv").unwrap();
700        assert_eq!(url, "https://huggingface.co/datasets/owner/ds/resolve/v1.0/data/train.csv");
701    }
702
703    #[test]
704    fn split_id_splits_owner_and_name() {
705        assert_eq!(split_id("openai-community/gpt2"), ("openai-community", "gpt2"));
706        assert_eq!(split_id("HuggingFaceFW/fineweb"), ("HuggingFaceFW", "fineweb"));
707    }
708
709    #[test]
710    fn split_id_ownerless_short_form() {
711        assert_eq!(split_id("gpt2"), ("", "gpt2"));
712        assert_eq!(split_id(""), ("", ""));
713    }
714
715    #[test]
716    fn split_id_only_splits_first_slash() {
717        assert_eq!(split_id("owner/sub/name"), ("owner", "sub/name"));
718    }
719
720    #[test]
721    fn split_id_empty_owner_or_name() {
722        assert_eq!(split_id("/name"), ("", "name"));
723        assert_eq!(split_id("owner/"), ("owner", ""));
724    }
725
726    #[test]
727    fn test_builder_cache_dir_explicit() {
728        let client = HFClientBuilder::new().cache_dir("/tmp/my-cache").build().unwrap();
729        assert_eq!(client.cache_dir(), std::path::Path::new("/tmp/my-cache"));
730    }
731
732    // `#[serial]` because the precedence tests below mutate `HF_HOME`, and the default
733    // cache dir is derived from it.
734    #[test]
735    #[serial]
736    fn test_builder_cache_dir_default() {
737        let client = HFClientBuilder::new().build().unwrap();
738        let path_str = client.cache_dir().to_string_lossy();
739        assert!(path_str.contains("huggingface") && path_str.ends_with("hub"));
740    }
741
742    #[test]
743    fn test_xet_session_lazy_creation() {
744        let client = HFClientBuilder::new().build().unwrap();
745        assert!(client.inner.xet_state.lock().unwrap().session.is_none());
746        let (_s1, _gen) = client.xet_session().unwrap();
747        assert!(client.inner.xet_state.lock().unwrap().session.is_some());
748    }
749
750    #[test]
751    fn test_xet_session_shared_across_clones() {
752        let client = HFClientBuilder::new().build().unwrap();
753        let clone = client.clone();
754        let (_s1, _gen) = client.xet_session().unwrap();
755        assert!(clone.inner.xet_state.lock().unwrap().session.is_some());
756    }
757
758    #[test]
759    fn test_xet_session_recovers_after_abort() {
760        let client = HFClientBuilder::new().build().unwrap();
761
762        let (session, generation) = client.xet_session().unwrap();
763        session.abort().unwrap();
764
765        match session.new_file_download_group() {
766            Ok(_) => panic!("expected error after abort"),
767            Err(e) => client.replace_xet_session(generation, &e),
768        }
769
770        let (recovered, _) = client.xet_session().unwrap();
771        assert!(recovered.new_file_download_group().is_ok());
772    }
773
774    #[test]
775    fn test_xet_session_recovers_after_sigint_abort() {
776        let client = HFClientBuilder::new().build().unwrap();
777
778        let (session, generation) = client.xet_session().unwrap();
779        session.sigint_abort().unwrap();
780
781        client.replace_xet_session(generation, &xet::error::XetError::KeyboardInterrupt);
782
783        let (recovered, _) = client.xet_session().unwrap();
784        assert!(recovered.new_file_download_group().is_ok());
785    }
786
787    /// Simulates the call-site retry pattern used in xet.rs:
788    /// 1. Get session + generation, factory call fails
789    /// 2. Call replace_xet_session(generation) to drop the bad session
790    /// 3. Get a fresh session, factory call succeeds
791    #[test]
792    fn test_replace_and_retry_after_abort() {
793        let client = HFClientBuilder::new().build().unwrap();
794
795        let (session, generation) = client.xet_session().unwrap();
796        assert!(session.new_file_download_group().is_ok());
797
798        session.abort().unwrap();
799
800        let group = match session.new_file_download_group() {
801            Ok(b) => b,
802            Err(e) => {
803                client.replace_xet_session(generation, &e);
804                client
805                    .xet_session()
806                    .unwrap()
807                    .0
808                    .new_file_download_group()
809                    .expect("fresh session factory call should succeed")
810            },
811        };
812        drop(group);
813    }
814
815    /// Verifies that replace_xet_session with a stale generation is a no-op.
816    #[test]
817    fn test_replace_with_stale_generation_is_noop() {
818        let client = HFClientBuilder::new().build().unwrap();
819
820        let (session, gen1) = client.xet_session().unwrap();
821        session.abort().unwrap();
822
823        // First replace succeeds
824        client.replace_xet_session(gen1, &xet::error::XetError::KeyboardInterrupt);
825
826        // Get the fresh session with a new generation
827        let (_fresh, gen2) = client.xet_session().unwrap();
828        assert_ne!(gen1, gen2);
829
830        // Attempting to replace with the old generation is a no-op
831        client.replace_xet_session(gen1, &xet::error::XetError::KeyboardInterrupt);
832
833        // The fresh session is still cached
834        let (still_fresh, gen3) = client.xet_session().unwrap();
835        assert_eq!(gen2, gen3);
836        assert!(still_fresh.new_file_download_group().is_ok());
837    }
838
839    #[test]
840    fn test_xet_session_reuse_without_replacement() {
841        let client = HFClientBuilder::new().build().unwrap();
842
843        let (s1, g1) = client.xet_session().unwrap();
844        let (s2, g2) = client.xet_session().unwrap();
845
846        assert_eq!(g1, g2);
847        assert!(s1.new_file_download_group().is_ok());
848        assert!(s2.new_file_download_group().is_ok());
849    }
850
851    /// Token resolution precedence tests.
852    ///
853    /// These tests mutate process-wide environment variables, so they must run
854    /// serially. Each test isolates `HF_HOME` to a tempdir so the developer's
855    /// real `~/.cache/huggingface/token` cannot leak into the result.
856    mod token_precedence {
857        use std::io::Write;
858
859        use serial_test::serial;
860        use tempfile::TempDir;
861
862        use super::HFClientBuilder;
863        use crate::constants::{HF_HOME, HF_HUB_DISABLE_IMPLICIT_TOKEN, HF_TOKEN, HF_TOKEN_PATH, TOKEN_FILENAME};
864
865        struct EnvGuard {
866            saved: Vec<(&'static str, Option<String>)>,
867            _hf_home: TempDir,
868        }
869
870        impl EnvGuard {
871            fn new() -> Self {
872                let hf_home = tempfile::tempdir().expect("tempdir for HF_HOME");
873                let keys = [HF_TOKEN, HF_TOKEN_PATH, HF_HOME, HF_HUB_DISABLE_IMPLICIT_TOKEN];
874                let saved = keys.iter().map(|k| (*k, std::env::var(*k).ok())).collect();
875                for k in keys {
876                    unsafe { std::env::remove_var(k) };
877                }
878                unsafe { std::env::set_var(HF_HOME, hf_home.path()) };
879                Self {
880                    saved,
881                    _hf_home: hf_home,
882                }
883            }
884        }
885
886        impl Drop for EnvGuard {
887            fn drop(&mut self) {
888                for (k, v) in &self.saved {
889                    match v {
890                        Some(val) => unsafe { std::env::set_var(k, val) },
891                        None => unsafe { std::env::remove_var(k) },
892                    }
893                }
894            }
895        }
896
897        fn write_token_file(dir: &std::path::Path, name: &str, contents: &str) -> std::path::PathBuf {
898            let path = dir.join(name);
899            let mut f = std::fs::File::create(&path).unwrap();
900            f.write_all(contents.as_bytes()).unwrap();
901            path
902        }
903
904        #[test]
905        #[serial]
906        fn test_explicit_token_overrides_env() {
907            let _g = EnvGuard::new();
908            unsafe { std::env::set_var(HF_TOKEN, "env-token") };
909
910            let client = HFClientBuilder::new().token("explicit-token").build().unwrap();
911            assert_eq!(client.inner.token.as_deref(), Some("explicit-token"));
912        }
913
914        #[test]
915        #[serial]
916        fn test_env_token_used_when_no_explicit() {
917            let _g = EnvGuard::new();
918            unsafe { std::env::set_var(HF_TOKEN, "env-token") };
919
920            let client = HFClientBuilder::new().build().unwrap();
921            assert_eq!(client.inner.token.as_deref(), Some("env-token"));
922        }
923
924        #[test]
925        #[serial]
926        fn test_env_token_overrides_token_path_file() {
927            let g = EnvGuard::new();
928            let dir = tempfile::tempdir().unwrap();
929            let token_file = write_token_file(dir.path(), "tok", "file-token");
930            unsafe { std::env::set_var(HF_TOKEN, "env-token") };
931            unsafe { std::env::set_var(HF_TOKEN_PATH, &token_file) };
932
933            let client = HFClientBuilder::new().build().unwrap();
934            assert_eq!(client.inner.token.as_deref(), Some("env-token"));
935            drop(g);
936        }
937
938        #[test]
939        #[serial]
940        fn test_token_path_file_used_when_no_env() {
941            let _g = EnvGuard::new();
942            let dir = tempfile::tempdir().unwrap();
943            let token_file = write_token_file(dir.path(), "tok", "file-token\n");
944            unsafe { std::env::set_var(HF_TOKEN_PATH, &token_file) };
945
946            let client = HFClientBuilder::new().build().unwrap();
947            assert_eq!(client.inner.token.as_deref(), Some("file-token"));
948        }
949
950        #[test]
951        #[serial]
952        fn test_token_from_hf_home_file() {
953            let g = EnvGuard::new();
954            // HF_HOME was set to a tempdir by EnvGuard. Place a token file inside it.
955            let hf_home = std::env::var(HF_HOME).unwrap();
956            write_token_file(std::path::Path::new(&hf_home), TOKEN_FILENAME, "home-token");
957
958            let client = HFClientBuilder::new().build().unwrap();
959            assert_eq!(client.inner.token.as_deref(), Some("home-token"));
960            drop(g);
961        }
962
963        #[test]
964        #[serial]
965        fn test_disable_implicit_token_returns_none() {
966            let _g = EnvGuard::new();
967            unsafe { std::env::set_var(HF_TOKEN, "env-token") };
968            unsafe { std::env::set_var(HF_HUB_DISABLE_IMPLICIT_TOKEN, "1") };
969
970            let client = HFClientBuilder::new().build().unwrap();
971            assert!(client.inner.token.is_none());
972        }
973
974        #[test]
975        #[serial]
976        fn test_disable_implicit_token_does_not_block_explicit() {
977            let _g = EnvGuard::new();
978            unsafe { std::env::set_var(HF_HUB_DISABLE_IMPLICIT_TOKEN, "1") };
979
980            let client = HFClientBuilder::new().token("explicit-token").build().unwrap();
981            assert_eq!(client.inner.token.as_deref(), Some("explicit-token"));
982        }
983
984        #[test]
985        #[serial]
986        fn test_no_token_anywhere_is_none() {
987            let _g = EnvGuard::new();
988            // EnvGuard isolates HF_HOME to a fresh tempdir with no token file inside.
989            let client = HFClientBuilder::new().build().unwrap();
990            assert!(client.inner.token.is_none());
991        }
992    }
993}
hf_hub/client.rs

hf_hub/
client.rs
