forge_workspace/
hub.rs

1//! The Hugging Face Hub calls the speaking model's install makes: the
2//! repository's current revision, its blobs with their LFS digests, and the
3//! download of one pinned file.
4//!
5//! The token is carried for a gated repository; the speaking model's own
6//! repository is open, so an absent token changes nothing.
7
8use std::io::Write;
9use std::path::Path;
10use std::time::Duration;
11
12use serde::Deserialize;
13use sha2::{Digest, Sha256};
14
15/// The Hub's connect and read bounds, a whole-request backstop sized so the
16/// largest variant (325 MB at 0.4 Mbps) still finishes, and the page cap on
17/// one listing before the install gives up.
18const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
19const READ_TIMEOUT: Duration = Duration::from_secs(60);
20const TOTAL_TIMEOUT: Duration = Duration::from_secs(2 * 60 * 60);
21const MAX_TREE_PAGES: usize = 20;
22
23/// One file the repository publishes.
24#[derive(Debug, Clone, PartialEq, Eq)]
25pub struct HubBlob {
26    pub path: String,
27    pub size: u64,
28    /// The LFS digest the repository publishes, when the file is stored by
29    /// LFS.
30    pub sha256: Option<String>,
31}
32
33/// The files one speaking variant is made of.
34#[derive(Debug, Clone, PartialEq, Eq)]
35pub struct SpeakingFiles {
36    pub model: HubBlob,
37    pub voices: Vec<HubBlob>,
38}
39
40/// A Hub client with the optional token every request carries.
41pub struct Hub {
42    client: reqwest::Client,
43    token: Option<String>,
44}
45
46impl Hub {
47    pub fn new(token: Option<String>) -> Result<Self, String> {
48        // A stalled Hub must fail the install rather than hold it in
49        // `Downloading` forever; the total is a backstop for one big file.
50        // The trust roots are the machine's, so a corporate MITM CA in
51        // `NODE_EXTRA_CA_CERTS` or `SSL_CERT_FILE` reaches the Hub too.
52        let client = forge_agent::http_trust::with_extra_roots(
53            reqwest::Client::builder()
54                .connect_timeout(CONNECT_TIMEOUT)
55                .read_timeout(READ_TIMEOUT)
56                .timeout(TOTAL_TIMEOUT),
57        )
58        .build()
59        .map_err(|error| format!("the Hub client did not build: {error}"))?;
60        Ok(Self { client, token })
61    }
62
63    /// One authorised GET, its status checked.
64    async fn get(&self, url: &str) -> Result<reqwest::Response, String> {
65        let request = self.client.get(url);
66        let request = match self.token.as_deref() {
67            Some(token) => request.bearer_auth(token),
68            None => request,
69        };
70        let response = request
71            .send()
72            .await
73            .map_err(|error| format!("the Hub did not answer for {url}: {error}"))?;
74        if !response.status().is_success() {
75            return Err(format!("the Hub answered {} for {url}", response.status()));
76        }
77        Ok(response)
78    }
79
80    /// The repository's current revision sha. Every file URL is pinned to it,
81    /// so one install fetches one coherent set.
82    pub async fn revision(&self, repo: &str) -> Result<String, String> {
83        let url = format!("https://huggingface.co/api/models/{repo}");
84        let response = self.get(&url).await?;
85        let info: ModelInfo = response
86            .json()
87            .await
88            .map_err(|error| format!("the Hub's answer for {repo} did not parse: {error}"))?;
89        Ok(info.sha)
90    }
91
92    /// Every file the revision publishes, with its LFS digest.
93    ///
94    /// **Recursive, because the model and its voices are not at the root.**
95    /// The plain tree call answers the root's own entries - `onnx` and
96    /// `voices` as directories - so a reader that dropped directories would
97    /// find no model at all. A page that caps follows its `Link` header.
98    pub async fn blobs(&self, repo: &str, revision: &str) -> Result<Vec<HubBlob>, String> {
99        let mut url = tree_url(repo, revision);
100        let mut blobs = Vec::new();
101        for _ in 0..MAX_TREE_PAGES {
102            let response = self.get(&url).await?;
103            let next = response
104                .headers()
105                .get(reqwest::header::LINK)
106                .and_then(|link| link.to_str().ok())
107                .and_then(next_page);
108            let raw = response
109                .bytes()
110                .await
111                .map_err(|error| format!("the Hub's listing for {repo} did not read: {error}"))?;
112            blobs.extend(parse_blobs(&raw)?);
113            match next {
114                Some(next) => url = next,
115                None => return Ok(blobs),
116            }
117        }
118        Err(format!("the Hub's listing for {repo} did not finish within {MAX_TREE_PAGES} pages"))
119    }
120
121    /// Fetch one pinned file, reporting the bytes written so far.
122    ///
123    /// The file lands at `dest` only once its byte length and, where the
124    /// repository publishes one, its digest match; a mismatch leaves nothing
125    /// behind.
126    pub async fn download(
127        &self,
128        repo: &str,
129        revision: &str,
130        blob: &HubBlob,
131        dest: &Path,
132        mut on_bytes: impl FnMut(u64),
133    ) -> Result<(), String> {
134        let url = format!("https://huggingface.co/{repo}/resolve/{revision}/{}", blob.path);
135        let request = self.client.get(&url);
136        let request = match self.token.as_deref() {
137            Some(token) => request.bearer_auth(token),
138            None => request,
139        };
140        let mut response = request
141            .send()
142            .await
143            .map_err(|error| format!("{} did not fetch: {error}", blob.path))?;
144        if !response.status().is_success() {
145            return Err(format!("{} answered {}", blob.path, response.status()));
146        }
147        if let Some(parent) = dest.parent() {
148            std::fs::create_dir_all(parent)
149                .map_err(|error| format!("{}: {error}", parent.display()))?;
150        }
151        let part = part_path(dest);
152        let mut file =
153            std::fs::File::create(&part).map_err(|error| format!("{}: {error}", part.display()))?;
154        let mut hasher = Sha256::new();
155        let mut written = 0_u64;
156        loop {
157            let chunk = match response.chunk().await {
158                Ok(Some(chunk)) => chunk,
159                Ok(None) => break,
160                Err(error) => {
161                    let _ = std::fs::remove_file(&part);
162                    return Err(format!("{} did not finish: {error}", blob.path));
163                }
164            };
165            if let Err(error) = file.write_all(&chunk) {
166                let _ = std::fs::remove_file(&part);
167                return Err(format!("{}: {error}", part.display()));
168            }
169            hasher.update(&chunk);
170            written += chunk.len() as u64;
171            on_bytes(written);
172        }
173        drop(file);
174        finish_file(&part, dest, blob, written, &hex::encode(hasher.finalize()))
175    }
176}
177
178/// Where a download is written while it is being checked.
179fn part_path(dest: &Path) -> std::path::PathBuf {
180    let name =
181        dest.file_name().map_or_else(String::new, |name| name.to_string_lossy().into_owned());
182    dest.with_file_name(format!("{name}.part"))
183}
184
185/// Check a finished download and move it into place, or remove it.
186pub fn finish_file(
187    part: &Path,
188    dest: &Path,
189    blob: &HubBlob,
190    written: u64,
191    digest: &str,
192) -> Result<(), String> {
193    if written != blob.size {
194        let _ = std::fs::remove_file(part);
195        return Err(format!(
196            "{} is {written} bytes and the entry declares {}",
197            blob.path, blob.size
198        ));
199    }
200    if let Some(expected) = blob.sha256.as_deref()
201        && !digest.eq_ignore_ascii_case(expected)
202    {
203        let _ = std::fs::remove_file(part);
204        return Err(format!("{} declares sha256 {expected}, and the file is {digest}", blob.path));
205    }
206    std::fs::rename(part, dest).map_err(|error| format!("{}: {error}", dest.display()))
207}
208
209/// The listing URL for one revision: recursive, because the model and its
210/// voices sit under `onnx/` and `voices/`.
211pub fn tree_url(repo: &str, revision: &str) -> String {
212    format!("https://huggingface.co/api/models/{repo}/tree/{revision}?recursive=true&limit=1000")
213}
214
215/// The next page a `Link` header names, when it names one.
216pub fn next_page(link: &str) -> Option<String> {
217    for part in link.split(',') {
218        if !part.contains("rel=\"next\"") {
219            continue;
220        }
221        let start = part.find('<')?;
222        let end = start + 1 + part[start + 1..].find('>')?;
223        return Some(part[start + 1..end].to_owned());
224    }
225    None
226}
227
228/// Refuse a listing whose model file disagrees with the compiled variant.
229///
230/// The digest the install checks a download against comes from this same
231/// listing, so without this the compile-time pin would be display only: a
232/// replaced upstream file would install under a new revision while the page
233/// still drew the old one.
234pub fn verify_release(
235    variant: &crate::tts::SpeakingVariant,
236    model: &HubBlob,
237) -> Result<(), String> {
238    if model.size != variant.size {
239        return Err(format!(
240            "{} is {} bytes at this revision and the compiled entry declares {}",
241            variant.path, model.size, variant.size
242        ));
243    }
244    match model.sha256.as_deref() {
245        Some(digest) if digest.eq_ignore_ascii_case(variant.sha256) => Ok(()),
246        Some(digest) => Err(format!(
247            "{} declares sha256 {digest} at this revision and the compiled entry declares {}",
248            variant.path, variant.sha256
249        )),
250        None => Err(format!(
251            "{} publishes no digest at this revision, so the compiled one cannot be checked",
252            variant.path
253        )),
254    }
255}
256
257#[derive(Deserialize)]
258struct ModelInfo {
259    sha: String,
260}
261
262#[derive(Deserialize)]
263struct TreeEntry {
264    #[serde(rename = "type")]
265    kind: String,
266    path: String,
267    size: Option<u64>,
268    lfs: Option<Lfs>,
269}
270
271#[derive(Deserialize)]
272struct Lfs {
273    oid: String,
274}
275
276/// The files a tree listing carries, directories dropped.
277pub fn parse_blobs(raw: &[u8]) -> Result<Vec<HubBlob>, String> {
278    let entries: Vec<TreeEntry> = serde_json::from_slice(raw)
279        .map_err(|error| format!("the Hub's listing did not parse: {error}"))?;
280    Ok(entries
281        .into_iter()
282        .filter(|entry| entry.kind == "file")
283        .map(|entry| HubBlob {
284            size: entry.size.unwrap_or(0),
285            sha256: entry.lfs.map(|lfs| lfs.oid),
286            path: entry.path,
287        })
288        .collect())
289}
290
291/// The voice id a style file's path names. Every named voice carries a
292/// `<language><gender>_<name>` stem; the repository's `voices/af.bin` blend
293/// does not, and is not one of the model's named voices.
294pub fn voice_id(path: &str) -> Option<&str> {
295    let file = std::path::Path::new(path.strip_prefix("voices/")?);
296    if !file.extension().is_some_and(|extension| extension.eq_ignore_ascii_case("bin")) {
297        return None;
298    }
299    let name = file.file_stem()?.to_str()?;
300    name.contains('_').then_some(name)
301}
302
303/// The model file and the style files one variant is made of.
304pub fn speaking_files(blobs: &[HubBlob], variant_path: &str) -> Result<SpeakingFiles, String> {
305    let model = blobs
306        .iter()
307        .find(|blob| blob.path == variant_path)
308        .cloned()
309        .ok_or_else(|| format!("the repository publishes no {variant_path}"))?;
310    let voices: Vec<HubBlob> =
311        blobs.iter().filter(|blob| voice_id(&blob.path).is_some()).cloned().collect();
312    if voices.is_empty() {
313        return Err("the repository publishes no voice files".to_owned());
314    }
315    Ok(SpeakingFiles { model, voices })
316}
317
318#[cfg(test)]
319mod tests {
320    use std::path::Path;
321
322    use super::{HubBlob, finish_file, parse_blobs, speaking_files, voice_id};
323
324    fn listing() -> Vec<HubBlob> {
325        parse_blobs(
326            br#"[
327                {"type":"file","path":"onnx/model_quantized.onnx","size":92,"lfs":{"oid":"aa"}},
328                {"type":"directory","path":"voices","size":0},
329                {"type":"file","path":"voices/af_heart.bin","size":5,"lfs":{"oid":"bb"}},
330                {"type":"file","path":"voices/af.bin","size":9,"lfs":{"oid":"cc"}},
331                {"type":"file","path":"README.md","size":10}
332            ]"#,
333        )
334        .expect("the listing parses")
335    }
336
337    /// The plain tree call answers only the root's directories; a listing
338    /// that is not recursive is an install that finds no model at all.
339    #[test]
340    fn the_tree_url_is_recursive() {
341        let url = super::tree_url("onnx-community/Kokoro-82M-v1.0-ONNX", "abc123");
342        assert!(url.contains("/tree/abc123?recursive=true"), "got: {url}");
343    }
344
345    #[test]
346    fn a_link_header_names_the_next_page_or_none() {
347        let next = super::next_page(
348            "<https://huggingface.co/api/models/o/r/tree/abc?recursive=true&cursor=xyz>; rel=\"next\"",
349        );
350        assert_eq!(
351            next.as_deref(),
352            Some("https://huggingface.co/api/models/o/r/tree/abc?recursive=true&cursor=xyz")
353        );
354        assert_eq!(super::next_page("<https://x>; rel=\"prev\""), None);
355    }
356
357    /// A recursive listing is flat: the model and its voices arrive by path
358    /// even though the root carries them as directories.
359    #[test]
360    fn a_recursive_listing_finds_the_nested_model_and_voices() {
361        let blobs = parse_blobs(
362            br#"[
363                {"type":"directory","path":"onnx","size":0},
364                {"type":"directory","path":"voices","size":0},
365                {"type":"file","path":"onnx/model_quantized.onnx","size":92,"lfs":{"oid":"aa"}},
366                {"type":"file","path":"voices/af_heart.bin","size":5,"lfs":{"oid":"bb"}},
367                {"type":"file","path":"voices/bf_emma.bin","size":5,"lfs":{"oid":"cc"}}
368            ]"#,
369        )
370        .expect("the listing parses");
371        let files = speaking_files(&blobs, "onnx/model_quantized.onnx").expect("the set");
372        assert_eq!(files.model.path, "onnx/model_quantized.onnx");
373        assert_eq!(files.voices.len(), 2);
374    }
375
376    /// The compiled variant's size and digest are the pin: a release that
377    /// disagrees is refused by the file it names.
378    #[test]
379    fn a_release_that_disagrees_with_the_compiled_pin_is_refused() {
380        let variant = crate::tts::SpeakingVariant::named(crate::tts::PINNED_VARIANT)
381            .expect("the pinned variant exists");
382        let pinned = HubBlob {
383            path: variant.path.to_owned(),
384            size: variant.size,
385            sha256: Some(variant.sha256.to_owned()),
386        };
387        super::verify_release(&variant, &pinned).expect("the compiled release is accepted");
388
389        let resized = HubBlob { size: variant.size + 1, ..pinned.clone() };
390        let error = super::verify_release(&variant, &resized).expect_err("a different size");
391        assert!(error.contains(variant.path), "got: {error}");
392
393        let rehashed = HubBlob { sha256: Some("00".repeat(32)), ..pinned.clone() };
394        let error = super::verify_release(&variant, &rehashed).expect_err("a different digest");
395        assert!(error.contains("sha256"), "got: {error}");
396
397        let undigested = HubBlob { sha256: None, ..pinned };
398        let error = super::verify_release(&variant, &undigested).expect_err("no digest");
399        assert!(error.contains("no digest"), "got: {error}");
400    }
401
402    #[test]
403    fn a_tree_listing_drops_directories_and_keeps_lfs_digests() {
404        let blobs = listing();
405        assert_eq!(blobs.len(), 4, "the directory is not a file");
406        let model = blobs
407            .iter()
408            .find(|blob| {
409                std::path::Path::new(&blob.path)
410                    .extension()
411                    .is_some_and(|extension| extension.eq_ignore_ascii_case("onnx"))
412            })
413            .expect("the model");
414        assert_eq!(model.size, 92);
415        assert_eq!(model.sha256.as_deref(), Some("aa"));
416        let readme = blobs.iter().find(|blob| blob.path == "README.md").expect("the readme");
417        assert_eq!(readme.sha256, None, "a file outside LFS has no digest");
418    }
419
420    #[test]
421    fn a_named_voice_is_a_stem_with_a_language_gender_prefix() {
422        assert_eq!(voice_id("voices/af_heart.bin"), Some("af_heart"));
423        assert_eq!(voice_id("voices/af.bin"), None, "the blend is not a named voice");
424        assert_eq!(voice_id("onnx/model.onnx"), None);
425    }
426
427    #[test]
428    fn a_variant_is_its_model_file_and_every_named_voice() {
429        let files = speaking_files(&listing(), "onnx/model_quantized.onnx").expect("the set");
430        assert_eq!(files.model.path, "onnx/model_quantized.onnx");
431        assert_eq!(files.voices.len(), 1, "only the named voice, not the blend");
432        assert_eq!(files.voices[0].path, "voices/af_heart.bin");
433    }
434
435    #[test]
436    fn a_variant_the_repository_does_not_publish_is_refused_by_name() {
437        let error = speaking_files(&listing(), "onnx/model_q4.onnx").expect_err("refused");
438        assert!(error.contains("onnx/model_q4.onnx"), "got: {error}");
439    }
440
441    #[test]
442    fn a_download_whose_digest_disagrees_is_taken_back() {
443        let dir = tempfile::tempdir().expect("a temp dir");
444        let part = dir.path().join("model.onnx.part");
445        let dest = dir.path().join("model.onnx");
446        std::fs::write(&part, b"weights").expect("the part is written");
447        let blob =
448            HubBlob { path: "onnx/model.onnx".to_owned(), size: 7, sha256: Some("00".repeat(32)) };
449        let error =
450            finish_file(&part, &dest, &blob, 7, "11".repeat(32).as_str()).expect_err("refused");
451        assert!(error.contains("sha256"), "got: {error}");
452        assert!(!part.exists(), "the partial file goes");
453        assert!(!dest.exists(), "nothing lands under the model's name");
454    }
455
456    #[test]
457    fn a_download_of_the_wrong_length_is_taken_back() {
458        let dir = tempfile::tempdir().expect("a temp dir");
459        let part = dir.path().join("model.onnx.part");
460        let dest = dir.path().join("model.onnx");
461        std::fs::write(&part, b"weights").expect("the part is written");
462        let blob = HubBlob { path: "onnx/model.onnx".to_owned(), size: 99, sha256: None };
463        let error =
464            finish_file(&part, &dest, &blob, 7, "11".repeat(32).as_str()).expect_err("refused");
465        assert!(error.contains("99"), "got: {error}");
466        assert!(!dest.exists());
467    }
468
469    #[test]
470    fn a_verified_download_lands_under_its_name() {
471        let dir = tempfile::tempdir().expect("a temp dir");
472        let part = dir.path().join("model.onnx.part");
473        let dest = dir.path().join("model.onnx");
474        std::fs::write(&part, b"weights").expect("the part is written");
475        // The digest the Hub would publish for these bytes, so the matching
476        // path is as pinned as the mismatch above.
477        let digest = {
478            use sha2::{Digest as _, Sha256};
479            let mut hasher = Sha256::new();
480            hasher.update(b"weights");
481            hex::encode(hasher.finalize())
482        };
483        let blob =
484            HubBlob { path: "onnx/model.onnx".to_owned(), size: 7, sha256: Some(digest.clone()) };
485        finish_file(&part, &dest, &blob, 7, digest.as_str()).expect("it lands");
486        assert_eq!(std::fs::read(&dest).expect("reads"), b"weights");
487        assert!(!Path::new(&part).exists());
488    }
489}
forge_workspace/hub.rs

forge_workspace/
hub.rs
