feat: cache and reuse XetSession in HFClient

.github/workflows/ci.yml

@@ -8,29 +8,46 @@ on:
 
 env:
   CARGO_TERM_COLOR: always
-  HF_ENDPOINT: https://hub-ci.huggingface.co
-  RUST_BACKTRACE: 1
 
 jobs:
   fmt:
     name: Format
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
-      - uses: dtolnay/rust-toolchain@nightly
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6.0.2
+      - uses: dtolnay/rust-toolchain@3c5f7ea28cd621ae0bf5283f0e981fb97b8a7af9  # master
         with:
+          toolchain: nightly
           components: rustfmt
+      - uses: actions/cache@668228422ae6a00e4ad889ee87cd7109ec5666a7
+        with:
+          path: |
+            ~/.cargo/registry
+            ~/.cargo/git
+            target
+          key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
+          restore-keys: |
+            ${{ runner.os }}-cargo-
       - run: cargo +nightly fmt --all --check
 
   clippy:
     name: Lint
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
-      - uses: dtolnay/rust-toolchain@stable
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6.0.2
+      - uses: dtolnay/rust-toolchain@3c5f7ea28cd621ae0bf5283f0e981fb97b8a7af9  # master
         with:
+          toolchain: stable
           components: clippy
-      - uses: Swatinem/rust-cache@v2
+      - uses: actions/cache@668228422ae6a00e4ad889ee87cd7109ec5666a7
+        with:
+          path: |
+            ~/.cargo/registry
+            ~/.cargo/git
+            target
+          key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
+          restore-keys: |
+            ${{ runner.os }}-cargo-
       # lint with all features, and with only default features
       - run: cargo clippy -p huggingface-hub --all-features -- -D warnings
       - run: cargo clippy -p huggingface-hub -- -D warnings
@@ -39,9 +56,19 @@ jobs:
     name: Build (release)
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
-      - uses: dtolnay/rust-toolchain@stable
-      - uses: Swatinem/rust-cache@v2
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6.0.2
+      - uses: dtolnay/rust-toolchain@3c5f7ea28cd621ae0bf5283f0e981fb97b8a7af9  # master
+        with:
+          toolchain: stable
+      - uses: actions/cache@668228422ae6a00e4ad889ee87cd7109ec5666a7
+        with:
+          path: |
+            ~/.cargo/registry
+            ~/.cargo/git
+            target
+          key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
+          restore-keys: |
+            ${{ runner.os }}-cargo-
       - run: cargo build -p huggingface-hub --all-features --release
       - run: cargo build -p huggingface-hub --release
       - run: cargo build -p huggingface-hub --features xet --release
@@ -51,9 +78,19 @@ jobs:
     name: Test
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
-      - uses: dtolnay/rust-toolchain@stable
-      - uses: Swatinem/rust-cache@v2
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6.0.2
+      - uses: dtolnay/rust-toolchain@3c5f7ea28cd621ae0bf5283f0e981fb97b8a7af9  # master
+        with:
+          toolchain: stable
+      - uses: actions/cache@668228422ae6a00e4ad889ee87cd7109ec5666a7
+        with:
+          path: |
+            ~/.cargo/registry
+            ~/.cargo/git
+            target
+          key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
+          restore-keys: |
+            ${{ runner.os }}-cargo-
       - name: Set up Python for cache interop tests
         uses: actions/setup-python@v5
         with:
@@ -65,7 +102,8 @@ jobs:
           pip install huggingface_hub[cli]
       - name: All tests
         env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+          HF_CI_TOKEN: ${{ secrets.HF_CI_TOKEN }}
+          HF_PROD_TOKEN: ${{ secrets.HF_PROD_TOKEN }}
           HF_TEST_WRITE: "1"
         run: |
           source .venv/bin/activate

huggingface_hub/Cargo.toml

@@ -75,6 +75,7 @@ sha2 = "0.10"
 serial_test = "3"
 assert_cmd = "2"
 anyhow = "1"
+libc = "0.2"
 
 [[example]]
 name = "repo"

huggingface_hub/src/api/commits.rs

@@ -5,11 +5,11 @@ use url::Url;
 use crate::constants;
 use crate::diff::HFFileDiff;
 use crate::error::Result;
-use crate::repository::{
-    HFRepository, RepoCreateBranchParams, RepoCreateTagParams, RepoDeleteBranchParams, RepoDeleteTagParams,
+use crate::repository::HFRepository;
+use crate::types::{
+    GitCommitInfo, GitRefs, RepoCreateBranchParams, RepoCreateTagParams, RepoDeleteBranchParams, RepoDeleteTagParams,
     RepoGetCommitDiffParams, RepoGetRawDiffParams, RepoListCommitsParams, RepoListRefsParams,
 };
-use crate::types::{GitCommitInfo, GitRefs};
 
 impl HFRepository {
     /// Stream commit history for the repository at a given revision.

huggingface_hub/src/api/files.rs

@@ -11,12 +11,13 @@ use tokio::io::AsyncWriteExt;
 use url::Url;
 
 use crate::error::{HFError, Result};
-use crate::repository::{
-    HFRepository, RepoCreateCommitParams, RepoDeleteFileParams, RepoDeleteFolderParams, RepoDownloadFileParams,
-    RepoDownloadFileStreamParams, RepoDownloadFileToBytesParams, RepoGetPathsInfoParams, RepoListFilesParams,
-    RepoListTreeParams, RepoSnapshotDownloadParams, RepoUploadFileParams, RepoUploadFolderParams,
+use crate::repository::HFRepository;
+use crate::types::{
+    AddSource, CommitInfo, CommitOperation, RepoCreateCommitParams, RepoDeleteFileParams, RepoDeleteFolderParams,
+    RepoDownloadFileParams, RepoDownloadFileStreamParams, RepoDownloadFileToBytesParams, RepoGetPathsInfoParams,
+    RepoListFilesParams, RepoListTreeParams, RepoSnapshotDownloadParams, RepoTreeEntry, RepoType, RepoUploadFileParams,
+    RepoUploadFolderParams,
 };
-use crate::types::{AddSource, CommitInfo, CommitOperation, RepoTreeEntry, RepoType};
 use crate::{cache, constants};
 
 impl HFRepository {

huggingface_hub/src/api/repo.rs

@@ -4,10 +4,10 @@ use url::Url;
 use crate::client::HFClient;
 use crate::constants;
 use crate::error::{HFError, Result};
-use crate::repository::{HFRepository, RepoFileExistsParams, RepoRevisionExistsParams, RepoUpdateSettingsParams};
+use crate::repository::HFRepository;
 use crate::types::{
     CreateRepoParams, DatasetInfo, DeleteRepoParams, ListDatasetsParams, ListModelsParams, ListSpacesParams, ModelInfo,
-    MoveRepoParams, RepoUrl, SpaceInfo,
+    MoveRepoParams, RepoFileExistsParams, RepoRevisionExistsParams, RepoUpdateSettingsParams, RepoUrl, SpaceInfo,
 };
 
 impl HFRepository {
@@ -471,7 +471,7 @@ sync_api_stream! {
 
 sync_api! {
     impl HFRepository -> HFRepositorySync {
-        fn info(&self, params: &crate::repository::RepoInfoParams) -> Result<crate::types::RepoInfo>;
+        fn info(&self, params: &crate::types::RepoInfoParams) -> Result<crate::types::RepoInfo>;
         fn exists(&self) -> Result<bool>;
         fn revision_exists(&self, params: &RepoRevisionExistsParams) -> Result<bool>;
         fn file_exists(&self, params: &RepoFileExistsParams) -> Result<bool>;

huggingface_hub/src/api/spaces.rs

@@ -1,10 +1,9 @@
-use crate::SpaceVariableDeleteParams;
 use crate::error::Result;
-use crate::repository::{
-    HFSpace, SpaceHardwareRequestParams, SpaceSecretDeleteParams, SpaceSecretParams, SpaceSleepTimeParams,
-    SpaceVariableParams,
+use crate::repository::HFSpace;
+use crate::types::{
+    DuplicateSpaceParams, RepoUrl, SpaceHardwareRequestParams, SpaceRuntime, SpaceSecretDeleteParams,
+    SpaceSecretParams, SpaceSleepTimeParams, SpaceVariableDeleteParams, SpaceVariableParams,
 };
-use crate::types::{DuplicateSpaceParams, RepoUrl, SpaceRuntime};
 
 impl HFSpace {
     /// Fetch the current runtime state of the Space (hardware, stage, URL, etc.).

huggingface_hub/src/client.rs

@@ -45,6 +45,8 @@ pub(crate) struct HFClientInner {
     pub(crate) token: Option<String>,
     pub(crate) cache_dir: std::path::PathBuf,
     pub(crate) cache_enabled: bool,
+    #[cfg(feature = "xet")]
+    pub(crate) xet_state: std::sync::Mutex<crate::xet::XetState>,
 }
 
 /// Builder for [`HFClient`].
@@ -182,6 +184,8 @@ impl HFClientBuilder {
                 token,
                 cache_dir,
                 cache_enabled: self.cache_enabled.unwrap_or(true),
+                #[cfg(feature = "xet")]
+                xet_state: std::sync::Mutex::new(crate::xet::XetState::default()),
             }),
         })
     }
@@ -300,6 +304,50 @@ impl HFClient {
     }
 }
 
+#[cfg(feature = "xet")]
+impl HFClient {
+    /// Get or lazily create the cached XetSession.
+    ///
+    /// Returns `(session, generation)`. The generation is an opaque counter
+    /// that identifies which session instance this is. Pass it to
+    /// [`replace_xet_session`](Self::replace_xet_session) so that only the
+    /// caller that observed the error triggers a replacement — concurrent
+    /// callers that already obtained a fresh session won't clobber it.
+    pub(crate) fn xet_session(&self) -> Result<(xet::xet_session::XetSession, u64)> {
+        let mut guard = self
+            .inner
+            .xet_state
+            .lock()
+            .map_err(|e| HFError::Other(format!("xet session mutex poisoned: {e}")))?;
+
+        if let Some(ref session) = guard.session {
+            return Ok((session.clone(), guard.generation));
+        }
+
+        let session = xet::xet_session::XetSessionBuilder::new()
+            .build()
+            .map_err(|e| HFError::Other(format!("Failed to build xet session: {e}")))?;
+        guard.session = Some(session.clone());
+        guard.generation += 1;
+        Ok((session, guard.generation))
+    }
+
+    /// Replace the cached XetSession only if the generation matches.
+    ///
+    /// Called by xet call sites when a factory method returns an error.
+    /// The generation check ensures that if another thread already replaced
+    /// the session, this call is a no-op rather than discarding the fresh one.
+    pub(crate) fn replace_xet_session(&self, generation: u64, err: &xet::error::XetError) {
+        tracing::warn!(error = %err, generation, "replacing cached XetSession");
+        let Ok(mut guard) = self.inner.xet_state.lock() else {
+            return;
+        };
+        if guard.generation == generation {
+            guard.session = None;
+        }
+    }
+}
+
 /// Resolve token from environment or token file.
 /// Priority: HF_TOKEN env → HF_TOKEN_PATH file → $HF_HOME/token file.
 fn resolve_token() -> Option<String> {
@@ -357,4 +405,120 @@ mod tests {
         let path_str = api.cache_dir().to_string_lossy();
         assert!(path_str.contains("huggingface") && path_str.ends_with("hub"));
     }
+
+    #[cfg(feature = "xet")]
+    #[test]
+    fn test_xet_session_lazy_creation() {
+        let client = HFClientBuilder::new().build().unwrap();
+        assert!(client.inner.xet_state.lock().unwrap().session.is_none());
+        let (_s1, _gen) = client.xet_session().unwrap();
+        assert!(client.inner.xet_state.lock().unwrap().session.is_some());
+    }
+
+    #[cfg(feature = "xet")]
+    #[test]
+    fn test_xet_session_shared_across_clones() {
+        let client = HFClientBuilder::new().build().unwrap();
+        let clone = client.clone();
+        let (_s1, _gen) = client.xet_session().unwrap();
+        assert!(clone.inner.xet_state.lock().unwrap().session.is_some());
+    }
+
+    #[cfg(feature = "xet")]
+    #[test]
+    fn test_xet_session_recovers_after_abort() {
+        let client = HFClientBuilder::new().build().unwrap();
+
+        let (session, generation) = client.xet_session().unwrap();
+        session.abort().unwrap();
+
+        match session.new_file_download_group() {
+            Ok(_) => panic!("expected error after abort"),
+            Err(e) => client.replace_xet_session(generation, &e),
+        }
+
+        let (recovered, _) = client.xet_session().unwrap();
+        assert!(recovered.new_file_download_group().is_ok());
+    }
+
+    #[cfg(feature = "xet")]
+    #[test]
+    fn test_xet_session_recovers_after_sigint_abort() {
+        let client = HFClientBuilder::new().build().unwrap();
+
+        let (session, generation) = client.xet_session().unwrap();
+        session.sigint_abort().unwrap();
+
+        client.replace_xet_session(generation, &xet::error::XetError::KeyboardInterrupt);
+
+        let (recovered, _) = client.xet_session().unwrap();
+        assert!(recovered.new_file_download_group().is_ok());
+    }
+
+    /// Simulates the call-site retry pattern used in xet.rs:
+    /// 1. Get session + generation, factory call fails
+    /// 2. Call replace_xet_session(generation) to drop the bad session
+    /// 3. Get fresh session, factory call succeeds
+    #[cfg(feature = "xet")]
+    #[test]
+    fn test_replace_and_retry_after_abort() {
+        let client = HFClientBuilder::new().build().unwrap();
+
+        let (session, generation) = client.xet_session().unwrap();
+        assert!(session.new_file_download_group().is_ok());
+
+        session.abort().unwrap();
+
+        let group = match session.new_file_download_group() {
+            Ok(b) => b,
+            Err(e) => {
+                client.replace_xet_session(generation, &e);
+                client
+                    .xet_session()
+                    .unwrap()
+                    .0
+                    .new_file_download_group()
+                    .expect("fresh session factory call should succeed")
+            },
+        };
+        drop(group);
+    }
+
+    /// Verifies that replace_xet_session with a stale generation is a no-op.
+    #[cfg(feature = "xet")]
+    #[test]
+    fn test_replace_with_stale_generation_is_noop() {
+        let client = HFClientBuilder::new().build().unwrap();
+
+        let (session, gen1) = client.xet_session().unwrap();
+        session.abort().unwrap();
+
+        // First replace succeeds
+        client.replace_xet_session(gen1, &xet::error::XetError::KeyboardInterrupt);
+
+        // Get the fresh session with a new generation
+        let (_fresh, gen2) = client.xet_session().unwrap();
+        assert_ne!(gen1, gen2);
+
+        // Attempting to replace with the old generation is a no-op
+        client.replace_xet_session(gen1, &xet::error::XetError::KeyboardInterrupt);
+
+        // The fresh session is still cached
+        let (still_fresh, gen3) = client.xet_session().unwrap();
+        assert_eq!(gen2, gen3);
+        assert!(still_fresh.new_file_download_group().is_ok());
+    }
+
+    #[cfg(feature = "xet")]
+    #[test]
+    fn test_xet_session_reuse_without_replacement() {
+        let client = HFClientBuilder::new().build().unwrap();
+
+        let (s1, g1) = client.xet_session().unwrap();
+        let (s2, g2) = client.xet_session().unwrap();
+
+        assert_eq!(g1, g2);
+        assert!(s1.new_file_download_group().is_ok());
+        assert!(s2.new_file_download_group().is_ok());
+    }
 }

huggingface_hub/src/lib.rs

@@ -32,6 +32,8 @@ pub mod types;
 #[cfg(feature = "xet")]
 pub mod xet;
 
+pub mod test_utils;
+
 #[cfg(feature = "blocking")]
 pub use blocking::{HFClientSync, HFRepoSync, HFRepositorySync, HFSpaceSync};
 pub use client::{HFClient, HFClientBuilder};