@ondeinference / onde
Code Commits Issues Pulls Cloud Agent CI/CD
copilot/fix-onde-npm-sdk-v1-1-2-release
copilot/assess-inference-engine-complexity copilot/fix-onde-npm-sdk-v1-1-2-release development feature/diffusion-vision main
rs 1,333 lines 47.8 KB
Raw
1 //! HuggingFace hub cache inspection, repair, and model download.
2 //!
3 //! This module owns all on-disk state management for locally cached
4 //! HuggingFace models and exposes a progress-callback-based download API
5 //! that is completely decoupled from any application framework.
6 //!
7 //! 1. Call [`list_local_hf_models`] / [`list_supported_hf_models`] /
8 //! [`delete_local_hf_model`] directly from your command handlers.
9 //! 2. Call [`download_model`] with an `Arc<dyn Fn(ModelDownloadProgress)>`
10 //! that forwards progress to whatever event system your app uses.
11
12 use {
13 crate::inference::models::{SUPPORTED_MODELS, SUPPORTED_MODEL_INFO},
14 log::{debug, info, warn},
15 serde::{Deserialize, Serialize},
16 std::{
17 fs,
18 path::{Path, PathBuf},
19 },
20 tsync::tsync,
21 };
22
23 // ---------------------------------------------------------------------------
24 // Public types
25 // ---------------------------------------------------------------------------
26
27 /// Information about a locally cached HuggingFace model.
28 #[tsync]
29 #[derive(Debug, Clone, Serialize, Deserialize)]
30 pub struct LocalHfModel {
31 /// The full model identifier, e.g. "bartowski/Qwen2.5-1.5B-Instruct-GGUF".
32 pub model_id: String,
33 /// The organisation or user that published the model.
34 pub org: String,
35 /// The model name without the org prefix.
36 pub name: String,
37 /// Total size on disk in bytes.
38 pub size_bytes: u64,
39 /// Human-readable size string, e.g. "4.2 GB".
40 pub size_display: String,
41 /// Absolute path to the model cache directory.
42 pub path: String,
43 /// List of snapshot revisions that are locally available.
44 pub revisions: Vec<String>,
45 }
46
47 /// Response returned by [`list_local_hf_models`].
48 #[tsync]
49 #[derive(Debug, Clone, Serialize, Deserialize)]
50 pub struct LocalHfModelsResponse {
51 /// The models found in the local HuggingFace cache.
52 pub models: Vec<LocalHfModel>,
53 /// The path that was scanned.
54 pub cache_path: String,
55 /// Total size of all cached models in bytes.
56 pub total_size_bytes: u64,
57 /// Human-readable total size string.
58 pub total_size_display: String,
59 }
60
61 /// A supported model that may or may not be downloaded locally.
62 #[tsync]
63 #[derive(Debug, Clone, Serialize, Deserialize)]
64 pub struct SupportedHfModel {
65 /// The full model identifier, e.g. "bartowski/Qwen2.5-1.5B-Instruct-GGUF".
66 pub model_id: String,
67 /// Human-friendly display name.
68 pub name: String,
69 /// Organisation or publisher display name.
70 pub org: String,
71 /// Short description of the model.
72 pub description: String,
73 /// Whether this model is fully downloaded locally.
74 pub is_downloaded: bool,
75 /// Whether a partial (incomplete) download exists on disk.
76 /// True when a cache directory exists but its size is less than 99 %
77 /// of `expected_size_bytes`.
78 pub is_incomplete: bool,
79 /// Bytes currently on disk for this model (0 when not present at all).
80 pub local_size_bytes: u64,
81 /// Human-readable local size string, e.g. "4.2 GB".
82 pub local_size_display: String,
83 /// Approximate total size in bytes when fully downloaded.
84 pub expected_size_bytes: u64,
85 /// Human-readable expected size string, e.g. "23.80 GB".
86 pub expected_size_display: String,
87 }
88
89 /// Response returned by [`list_supported_hf_models`].
90 #[tsync]
91 #[derive(Debug, Clone, Serialize, Deserialize)]
92 pub struct SupportedHfModelsResponse {
93 /// All models the app supports, with download status.
94 pub models: Vec<SupportedHfModel>,
95 }
96
97 /// Progress payload passed to the [`download_model`] callback on every poll
98 /// tick and on completion.
99 ///
100 /// Forward this value to your app's event system so the frontend can render
101 /// a progress bar.
102 #[tsync]
103 #[derive(Debug, Clone, Serialize, Deserialize)]
104 pub struct ModelDownloadProgress {
105 /// The model being downloaded.
106 pub model_id: String,
107 /// Bytes downloaded so far (observed from cache directory size).
108 pub downloaded_bytes: u64,
109 /// Human-readable downloaded size, e.g. "1.2 GB".
110 pub downloaded_display: String,
111 /// Expected total size in bytes (approximate).
112 pub total_bytes: u64,
113 /// Human-readable total size, e.g. "23.80 GB".
114 pub total_display: String,
115 /// Download progress as a value between 0.0 and 1.0.
116 pub progress: f64,
117 /// Whether the download has completed.
118 pub done: bool,
119 }
120
121 // ---------------------------------------------------------------------------
122 // Internal constants
123 // ---------------------------------------------------------------------------
124
125 /// Fraction of `expected_size_bytes` that must be present on disk for a
126 /// model to be considered fully downloaded (99 %).
127 const DOWNLOAD_COMPLETE_THRESHOLD: f64 = 0.99;
128
129 // ---------------------------------------------------------------------------
130 // Private helpers
131 // ---------------------------------------------------------------------------
132
133 /// Resolve the HuggingFace hub cache directory.
134 ///
135 /// Respects the `HF_HOME` / `HUGGINGFACE_HUB_CACHE` environment variables.
136 /// Falls back to `~/.cache/huggingface/hub`.
137 ///
138 /// On iOS the `HF_HOME` variable is set during app setup (see
139 /// `setup_application_filesystem`) to a path inside the app's writable
140 /// `Library/Application Support/` container. We create the directory if it
141 /// doesn't exist yet so that first-run and freshly-installed apps work
142 /// correctly.
143 fn hf_cache_dir() -> Option<PathBuf> {
144 // 1. Explicit hub cache override
145 if let Ok(cache) = std::env::var("HUGGINGFACE_HUB_CACHE") {
146 let p = PathBuf::from(cache);
147 let _ = fs::create_dir_all(&p);
148 if p.is_dir() {
149 return Some(p);
150 }
151 }
152
153 // 2. HF_HOME override (hub is a subdirectory)
154 if let Ok(hf_home) = std::env::var("HF_HOME") {
155 let p = PathBuf::from(hf_home).join("hub");
156 let _ = fs::create_dir_all(&p);
157 if p.is_dir() {
158 return Some(p);
159 }
160 }
161
162 // 3. Default: ~/.cache/huggingface/hub
163 if let Some(home) = home::home_dir() {
164 let p = home.join(".cache").join("huggingface").join("hub");
165 let _ = fs::create_dir_all(&p);
166 if p.is_dir() {
167 return Some(p);
168 }
169 }
170
171 None
172 }
173
174 /// Recursively compute the total size of a directory in bytes.
175 ///
176 /// Follows symlinks to count the real file size and also counts in-progress
177 /// `.incomplete` staging files so that progress polling reflects partial
178 /// downloads before they are renamed to their final blob name.
179 fn dir_size(path: &PathBuf) -> u64 {
180 let mut total: u64 = 0;
181 if let Ok(entries) = fs::read_dir(path) {
182 for entry in entries.flatten() {
183 let entry_path = entry.path();
184 if entry_path.is_dir() {
185 total += dir_size(&entry_path);
186 } else {
187 // Count both completed blobs AND in-progress staging files
188 // (hf-hub downloads to `<sha256>.incomplete` and renames on
189 // completion — without this, progress stays at 0% for the
190 // entire download and jumps to 100% at the end).
191 if let Ok(meta) = entry_path.symlink_metadata() {
192 if meta.is_symlink() {
193 // Follow the symlink to get the real file size.
194 if let Ok(real_meta) = entry_path.metadata() {
195 total += real_meta.len();
196 }
197 } else {
198 total += meta.len();
199 }
200 }
201 }
202 }
203 }
204 total
205 }
206
207 /// Format a byte count into a human-readable string.
208 fn format_size(bytes: u64) -> String {
209 const KB: u64 = 1024;
210 const MB: u64 = KB * 1024;
211 const GB: u64 = MB * 1024;
212 const TB: u64 = GB * 1024;
213
214 if bytes >= TB {
215 format!("{:.2} TB", bytes as f64 / TB as f64)
216 } else if bytes >= GB {
217 format!("{:.2} GB", bytes as f64 / GB as f64)
218 } else if bytes >= MB {
219 format!("{:.1} MB", bytes as f64 / MB as f64)
220 } else if bytes >= KB {
221 format!("{:.0} KB", bytes as f64 / KB as f64)
222 } else {
223 format!("{} B", bytes)
224 }
225 }
226
227 /// List the snapshot revisions available for a cached model directory.
228 fn list_revisions(model_dir: &Path) -> Vec<String> {
229 let snapshots_dir = model_dir.join("snapshots");
230 if !snapshots_dir.is_dir() {
231 return Vec::new();
232 }
233
234 let mut revisions = Vec::new();
235 if let Ok(entries) = fs::read_dir(&snapshots_dir) {
236 for entry in entries.flatten() {
237 if entry.path().is_dir() {
238 if let Some(name) = entry.file_name().to_str() {
239 revisions.push(name.to_string());
240 }
241 }
242 }
243 }
244 revisions.sort();
245 revisions
246 }
247
248 /// Build a [`ModelDownloadProgress`] payload from raw byte counts.
249 #[cfg(any(
250 target_os = "macos",
251 target_os = "ios",
252 target_os = "tvos",
253 target_os = "visionos",
254 target_os = "watchos",
255 target_os = "android"
256 ))]
257 fn make_progress(model_id: &str, downloaded: u64, total: u64, done: bool) -> ModelDownloadProgress {
258 let progress = if total > 0 {
259 (downloaded as f64 / total as f64).min(1.0)
260 } else {
261 0.0
262 };
263 ModelDownloadProgress {
264 model_id: model_id.to_string(),
265 downloaded_bytes: downloaded,
266 downloaded_display: format_size(downloaded),
267 total_bytes: total,
268 total_display: format_size(total),
269 progress,
270 done,
271 }
272 }
273
274 // ---------------------------------------------------------------------------
275 // Public cache utilities
276 // ---------------------------------------------------------------------------
277
278 /// Return the cache sub-directory path for a given model ID.
279 ///
280 /// E.g. for `"bartowski/Qwen2.5-1.5B-Instruct-GGUF"` the cache directory is
281 /// `<hf_cache>/models--bartowski--Qwen2.5-1.5B-Instruct-GGUF`.
282 pub fn model_cache_path(model_id: &str) -> Option<PathBuf> {
283 let cache_dir = hf_cache_dir()?;
284 let dir_name = format!("models--{}", model_id.replace('/', "--"));
285 Some(cache_dir.join(dir_name))
286 }
287
288 /// Pre-create the HuggingFace cache directory tree for `model_id` so that the
289 /// progress-monitor thread can find a valid path immediately, even before
290 /// hf-hub creates it during the first download.
291 ///
292 /// hf-hub won't be confused by a pre-existing (empty) directory — it will
293 /// simply populate it as normal.
294 pub fn ensure_model_cache_dir(model_id: &str) {
295 if let Some(model_path) = model_cache_path(model_id) {
296 let blobs_dir = model_path.join("blobs");
297 if let Err(e) = fs::create_dir_all(&blobs_dir) {
298 warn!(
299 "Could not pre-create model cache dir {}: {}",
300 blobs_dir.display(),
301 e
302 );
303 } else {
304 info!("Pre-created model cache dir: {}", model_path.display());
305 }
306 }
307 }
308
309 /// Remove stale `.lock` and orphaned `.part` files from the HF cache
310 /// `blobs/` directory for a model.
311 ///
312 /// The HF hub client creates `.lock` files during download. If a download is
313 /// interrupted these locks persist and prevent any subsequent access (download
314 /// **or** model load). Calling this before loading or resuming a download
315 /// allows the hub client to re-acquire the lock and proceed normally.
316 pub fn clean_stale_lock_files(model_id: &str) {
317 if let Some(cache_path) = model_cache_path(model_id) {
318 let blobs_dir = cache_path.join("blobs");
319 if blobs_dir.is_dir() {
320 // Collect the names of complete blobs (no extension) so we can
321 // decide whether a `.part` file is stale.
322 let complete_blobs: Vec<String> = fs::read_dir(&blobs_dir)
323 .into_iter()
324 .flatten()
325 .flatten()
326 .filter_map(|e| {
327 let p = e.path();
328 if p.is_file() && p.extension().is_none() {
329 p.file_name().map(|n| n.to_string_lossy().to_string())
330 } else {
331 None
332 }
333 })
334 .collect();
335
336 if let Ok(entries) = fs::read_dir(&blobs_dir) {
337 for entry in entries.flatten() {
338 let path = entry.path();
339 let ext = path.extension().and_then(|e| e.to_str());
340
341 match ext {
342 Some("lock") => {
343 info!("Removing stale lock file: {}", path.display());
344 if let Err(e) = fs::remove_file(&path) {
345 warn!("Failed to remove lock file {}: {}", path.display(), e);
346 }
347 }
348 Some("part") => {
349 // A `.part` file is a partial download left by hf-hub.
350 // If the corresponding complete blob exists (same name
351 // without `.part`), the partial file is stale and
352 // causes hf-hub to issue a spurious Range-resume
353 // download instead of using the complete blob.
354 let stem = path.file_stem().and_then(|s| s.to_str()).unwrap_or("");
355 if complete_blobs.iter().any(|b| b == stem) {
356 let part_size = fs::metadata(&path).map(|m| m.len()).unwrap_or(0);
357 info!(
358 "Removing stale .part file ({} bytes, complete blob exists): {}",
359 part_size,
360 path.display()
361 );
362 if let Err(e) = fs::remove_file(&path) {
363 warn!("Failed to remove .part file {}: {}", path.display(), e);
364 }
365 } else {
366 let part_size = fs::metadata(&path).map(|m| m.len()).unwrap_or(0);
367 debug!(
368 "Keeping .part file ({} bytes, no complete blob yet): {}",
369 part_size,
370 path.display()
371 );
372 }
373 }
374 _ => {}
375 }
376 }
377 }
378 }
379 }
380 }
381
382 /// Log the complete HF cache structure for a model to aid debugging.
383 ///
384 /// This is called before model load so the developer can see exactly what's
385 /// present in the cache, whether `refs/main` resolves, whether symlinks are
386 /// intact, and whether any `.part` files are left over from interrupted
387 /// downloads.
388 pub fn diagnose_hf_cache(model_id: &str) {
389 let cache_path = match model_cache_path(model_id) {
390 Some(p) => p,
391 None => {
392 info!("diagnose_hf_cache({}): cache path not found", model_id);
393 return;
394 }
395 };
396
397 if !cache_path.is_dir() {
398 info!(
399 "diagnose_hf_cache({}): directory does not exist: {}",
400 model_id,
401 cache_path.display()
402 );
403 return;
404 }
405
406 info!(
407 "diagnose_hf_cache({}): scanning {}",
408 model_id,
409 cache_path.display()
410 );
411
412 // ── refs/ ────────────────────────────────────────────────────────────
413 let refs_dir = cache_path.join("refs");
414 if refs_dir.is_dir() {
415 if let Ok(entries) = fs::read_dir(&refs_dir) {
416 for entry in entries.flatten() {
417 let name = entry.file_name().to_string_lossy().to_string();
418 match fs::read_to_string(entry.path()) {
419 Ok(content) => {
420 let hash = content.trim();
421 info!(
422 " refs/{} → \"{}\" (len={})",
423 name,
424 if hash.len() > 12 {
425 format!("{}…", &hash[..12])
426 } else {
427 hash.to_string()
428 },
429 hash.len()
430 );
431 }
432 Err(e) => {
433 warn!(" refs/{} → <unreadable: {}>", name, e);
434 }
435 }
436 }
437 }
438 } else {
439 warn!(" refs/ directory missing!");
440 }
441
442 // ── blobs/ ───────────────────────────────────────────────────────────
443 let blobs_dir = cache_path.join("blobs");
444 if blobs_dir.is_dir() {
445 if let Ok(entries) = fs::read_dir(&blobs_dir) {
446 for entry in entries.flatten() {
447 let path = entry.path();
448 let name = entry.file_name().to_string_lossy().to_string();
449 let size = fs::metadata(&path).map(|m| m.len()).unwrap_or(0);
450 let kind = if name.ends_with(".lock") {
451 "LOCK"
452 } else if name.ends_with(".part") {
453 "PART (incomplete)"
454 } else {
455 "blob"
456 };
457 info!(
458 " blobs/{} — {} ({} bytes / {:.1} MB)",
459 name,
460 kind,
461 size,
462 size as f64 / 1_048_576.0
463 );
464 }
465 }
466 } else {
467 warn!(" blobs/ directory missing!");
468 }
469
470 // ── snapshots/ ───────────────────────────────────────────────────────
471 let snapshots_dir = cache_path.join("snapshots");
472 if snapshots_dir.is_dir() {
473 if let Ok(revs) = fs::read_dir(&snapshots_dir) {
474 for rev_entry in revs.flatten() {
475 let rev_name = rev_entry.file_name().to_string_lossy().to_string();
476 let rev_display = if rev_name.len() > 12 {
477 format!("{}…", &rev_name[..12])
478 } else {
479 rev_name.clone()
480 };
481 info!(" snapshots/{}/", rev_display);
482
483 if let Ok(files) = fs::read_dir(rev_entry.path()) {
484 for file_entry in files.flatten() {
485 let file_name = file_entry.file_name().to_string_lossy().to_string();
486 let file_path = file_entry.path();
487
488 if file_path.is_symlink() {
489 let target = fs::read_link(&file_path)
490 .map(|t| t.to_string_lossy().to_string())
491 .unwrap_or_else(|_| "<unreadable>".to_string());
492 let accessible = fs::metadata(&file_path).is_ok();
493 let status = if accessible { "OK" } else { "BROKEN" };
494 info!(" {} → {} [symlink: {}]", file_name, target, status);
495 } else if file_path.is_file() {
496 let size = fs::metadata(&file_path).map(|m| m.len()).unwrap_or(0);
497 info!(
498 " {} [{} bytes / {:.1} MB, real file]",
499 file_name,
500 size,
501 size as f64 / 1_048_576.0
502 );
503 } else {
504 warn!(" {} [not a file or symlink!]", file_name);
505 }
506 }
507 }
508 }
509 }
510 } else {
511 warn!(" snapshots/ directory missing!");
512 }
513
514 info!("diagnose_hf_cache({}): done", model_id);
515 }
516
517 /// Repair broken symlinks in the HF hub cache for a given model.
518 ///
519 /// The `hf-hub` crate stores downloaded files as blobs keyed by SHA-256 and
520 /// creates **symlinks** from `snapshots/<rev>/<filename>` → `../../blobs/<hash>`.
521 /// When the cache directory is copied into an iOS app container (e.g. via
522 /// Xcode "Replace Container" or `ideviceinstaller`), symlinks are often
523 /// silently dropped or broken. The `hf-hub` client then thinks the file is
524 /// missing and re-downloads it — even though the 940 MB blob is right there.
525 ///
526 /// This function walks every snapshot entry for `model_id`. If a path is a
527 /// broken symlink **or** a missing file, and a matching blob exists, it
528 /// replaces the symlink with a hard copy of the blob so that `hf-hub`'s
529 /// `Cache::get()` finds a real file at the expected path.
530 ///
531 /// It also ensures a `refs/main` file exists pointing at the snapshot
532 /// revision, since that file is needed for `hf-hub` to resolve the "main"
533 /// ref to a concrete snapshot directory.
534 ///
535 /// This is idempotent and safe to call on every app launch.
536 pub fn repair_hf_cache_symlinks(model_id: &str) {
537 let cache_path = match model_cache_path(model_id) {
538 Some(p) if p.is_dir() => p,
539 _ => return,
540 };
541
542 let blobs_dir = cache_path.join("blobs");
543 let snapshots_dir = cache_path.join("snapshots");
544 let refs_dir = cache_path.join("refs");
545
546 if !blobs_dir.is_dir() || !snapshots_dir.is_dir() {
547 return;
548 }
549
550 // Collect blob filenames (excluding .lock files) for lookup.
551 let blob_names: Vec<String> = match fs::read_dir(&blobs_dir) {
552 Ok(entries) => entries
553 .flatten()
554 .filter_map(|e| {
555 let name = e.file_name().to_string_lossy().to_string();
556 if name.ends_with(".lock") {
557 None
558 } else if e.path().is_file() {
559 Some(name)
560 } else {
561 None
562 }
563 })
564 .collect(),
565 Err(_) => return,
566 };
567
568 if blob_names.is_empty() {
569 return;
570 }
571
572 // Walk each snapshot revision directory.
573 let snapshot_revs: Vec<(String, PathBuf)> = match fs::read_dir(&snapshots_dir) {
574 Ok(entries) => entries
575 .flatten()
576 .filter(|e| e.path().is_dir())
577 .map(|e| {
578 let rev = e.file_name().to_string_lossy().to_string();
579 (rev, e.path())
580 })
581 .collect(),
582 Err(_) => return,
583 };
584
585 for (rev, snap_dir) in &snapshot_revs {
586 let entries = match fs::read_dir(snap_dir) {
587 Ok(e) => e,
588 Err(_) => continue,
589 };
590
591 for entry in entries.flatten() {
592 let snap_file = entry.path();
593 let file_name = entry.file_name().to_string_lossy().to_string();
594
595 // Check if the file is accessible (broken symlink → metadata fails).
596 let needs_repair = if snap_file.is_symlink() {
597 // Symlink exists but may be broken — check if target is readable.
598 fs::metadata(&snap_file).is_err()
599 } else if snap_file.exists() {
600 // Real file, no repair needed.
601 false
602 } else {
603 // Path doesn't exist at all (dangling entry in readdir).
604 true
605 };
606
607 if !needs_repair {
608 continue;
609 }
610
611 // Try to find a matching blob. For GGUF repos there's usually
612 // just one blob; for multi-file repos we match by the symlink
613 // target name if we can read it, otherwise fall back to the
614 // single-blob heuristic.
615 let blob_name = if let Ok(target) = fs::read_link(&snap_file) {
616 target
617 .file_name()
618 .map(|n| n.to_string_lossy().to_string())
619 .filter(|n| blob_names.contains(n))
620 } else {
621 None
622 };
623
624 let resolved_blob = if let Some(ref name) = blob_name {
625 Some(blobs_dir.join(name))
626 } else if blob_names.len() == 1 {
627 // Single-blob repo (typical for GGUF) — safe to assume it
628 // corresponds to this snapshot entry.
629 Some(blobs_dir.join(&blob_names[0]))
630 } else {
631 None
632 };
633
634 if let Some(blob_path) = resolved_blob {
635 if !blob_path.is_file() {
636 continue;
637 }
638
639 // Remove the broken symlink before copying.
640 if snap_file.is_symlink() || snap_file.exists() {
641 let _ = fs::remove_file(&snap_file);
642 }
643
644 info!(
645 "repair_hf_cache_symlinks: replacing broken symlink with copy: {} → {}",
646 blob_path.display(),
647 snap_file.display()
648 );
649
650 // Prefer hard link (instant, no extra disk space) over copy.
651 if fs::hard_link(&blob_path, &snap_file).is_err() {
652 if let Err(e) = fs::copy(&blob_path, &snap_file) {
653 warn!(
654 "repair_hf_cache_symlinks: failed to copy blob to snapshot: {}",
655 e
656 );
657 }
658 }
659 } else {
660 warn!(
661 "repair_hf_cache_symlinks: broken symlink for '{}' but no matching blob found",
662 file_name
663 );
664 }
665 }
666
667 // Ensure refs/main exists and points to this revision so that
668 // `hf-hub` can resolve the snapshot.
669 if !refs_dir.join("main").is_file() {
670 let _ = fs::create_dir_all(&refs_dir);
671 if let Err(e) = fs::write(refs_dir.join("main"), rev.as_bytes()) {
672 warn!("repair_hf_cache_symlinks: failed to write refs/main: {}", e);
673 } else {
674 info!("repair_hf_cache_symlinks: created refs/main → {}", rev);
675 }
676 }
677 }
678 }
679
680 /// Ensure snapshot entries exist for known GGUF files.
681 ///
682 /// When a download completes but the process is killed before hf-hub creates
683 /// the snapshot symlink, the blob exists but hf-hub considers it a cache miss
684 /// and re-downloads. This function creates hard links from the snapshot
685 /// directory to the blob directory for any missing entries.
686 pub fn ensure_snapshot_entries(model_id: &str, expected_files: &[&str]) {
687 let cache_path = match model_cache_path(model_id) {
688 Some(p) if p.is_dir() => p,
689 _ => return,
690 };
691
692 let blobs_dir = cache_path.join("blobs");
693 let snapshots_dir = cache_path.join("snapshots");
694 let refs_dir = cache_path.join("refs");
695
696 if !blobs_dir.is_dir() {
697 return;
698 }
699
700 // Collect non-lock, non-part blob files.
701 let blob_entries: Vec<(String, u64)> = match fs::read_dir(&blobs_dir) {
702 Ok(entries) => entries
703 .flatten()
704 .filter_map(|e| {
705 let name = e.file_name().to_string_lossy().to_string();
706 if name.ends_with(".lock") || name.ends_with(".part") {
707 return None;
708 }
709 let meta = e.metadata().ok()?;
710 if meta.is_file() && meta.len() > 0 {
711 Some((name, meta.len()))
712 } else {
713 None
714 }
715 })
716 .collect(),
717 Err(_) => return,
718 };
719
720 if blob_entries.is_empty() {
721 return;
722 }
723
724 // Find the revision from refs/main, or from an existing snapshot dir.
725 let revision = fs::read_to_string(refs_dir.join("main"))
726 .ok()
727 .map(|s| s.trim().to_string())
728 .or_else(|| {
729 fs::read_dir(&snapshots_dir)
730 .ok()?
731 .flatten()
732 .find(|e| e.path().is_dir())
733 .map(|e| e.file_name().to_string_lossy().to_string())
734 });
735
736 let revision = match revision {
737 Some(r) if !r.is_empty() => r,
738 _ => return,
739 };
740
741 let snap_dir = snapshots_dir.join(&revision);
742 let _ = fs::create_dir_all(&snap_dir);
743
744 for expected_file in expected_files {
745 let snap_file = snap_dir.join(expected_file);
746 if snap_file.exists() {
747 continue;
748 }
749
750 // For single-blob repos, link the only blob. For multi-blob repos,
751 // we can't determine the mapping without metadata.
752 let blob_to_link = if blob_entries.len() == 1 {
753 Some(&blob_entries[0].0)
754 } else {
755 None
756 };
757
758 if let Some(blob_name) = blob_to_link {
759 let blob_path = blobs_dir.join(blob_name);
760 if fs::hard_link(&blob_path, &snap_file).is_ok() {
761 info!(
762 "ensure_snapshot_entries: hard-linked {} → {}",
763 expected_file, blob_name
764 );
765 } else if fs::copy(&blob_path, &snap_file).is_ok() {
766 info!(
767 "ensure_snapshot_entries: copied {} → {}",
768 expected_file, blob_name
769 );
770 } else {
771 warn!(
772 "ensure_snapshot_entries: failed to link/copy blob for {}",
773 expected_file
774 );
775 }
776 }
777 }
778 }
779
780 // ---------------------------------------------------------------------------
781 // Query / mutation functions
782 // ---------------------------------------------------------------------------
783
784 /// Scan the local HuggingFace hub cache and return all downloaded models that
785 /// the inference engine supports.
786 ///
787 /// The cache directory is typically `~/.cache/huggingface/hub/` and contains
788 /// sub-directories named `models--{org}--{model}` for each cached model.
789 pub fn list_local_hf_models() -> LocalHfModelsResponse {
790 let cache_dir = match hf_cache_dir() {
791 Some(dir) => dir,
792 None => {
793 debug!("HuggingFace cache directory not found.");
794 return LocalHfModelsResponse {
795 models: Vec::new(),
796 cache_path: String::new(),
797 total_size_bytes: 0,
798 total_size_display: "0 B".to_string(),
799 };
800 }
801 };
802
803 let cache_path_str = cache_dir.to_string_lossy().to_string();
804 debug!("Scanning HuggingFace cache at: {}", cache_path_str);
805
806 let mut models: Vec<LocalHfModel> = Vec::new();
807
808 let entries = match fs::read_dir(&cache_dir) {
809 Ok(entries) => entries,
810 Err(e) => {
811 warn!("Failed to read HuggingFace cache directory: {:?}", e);
812 return LocalHfModelsResponse {
813 models: Vec::new(),
814 cache_path: cache_path_str,
815 total_size_bytes: 0,
816 total_size_display: "0 B".to_string(),
817 };
818 }
819 };
820
821 for entry in entries.flatten() {
822 let entry_path = entry.path();
823 if !entry_path.is_dir() {
824 continue;
825 }
826
827 let dir_name = match entry.file_name().to_str() {
828 Some(name) => name.to_string(),
829 None => continue,
830 };
831
832 // HuggingFace hub caches model repos as "models--{org}--{name}"
833 if !dir_name.starts_with("models--") {
834 continue;
835 }
836
837 // Strip the "models--" prefix and split by "--" to get org/name
838 let remainder = &dir_name["models--".len()..];
839 let parts: Vec<&str> = remainder.splitn(2, "--").collect();
840 if parts.len() != 2 {
841 // Some models have no org (single segment), handle gracefully
842 let model_id = remainder.replace("--", "/");
843 let size = dir_size(&entry_path);
844 let revisions = list_revisions(&entry_path);
845
846 models.push(LocalHfModel {
847 model_id: model_id.clone(),
848 org: String::new(),
849 name: model_id,
850 size_bytes: size,
851 size_display: format_size(size),
852 path: entry_path.to_string_lossy().to_string(),
853 revisions,
854 });
855 continue;
856 }
857
858 let org = parts[0].to_string();
859 let name = parts[1].to_string();
860 let model_id = format!("{}/{}", org, name);
861 let size = dir_size(&entry_path);
862 let revisions = list_revisions(&entry_path);
863
864 models.push(LocalHfModel {
865 model_id,
866 org,
867 name,
868 size_bytes: size,
869 size_display: format_size(size),
870 path: entry_path.to_string_lossy().to_string(),
871 revisions,
872 });
873 }
874
875 // Only keep models that the inference engine actually supports.
876 models.retain(|m| SUPPORTED_MODELS.iter().any(|&s| s == m.model_id));
877
878 // Sort by model_id for a stable, predictable order.
879 models.sort_by_key(|a| a.model_id.to_lowercase());
880
881 let total_size_bytes: u64 = models.iter().map(|m| m.size_bytes).sum();
882
883 debug!(
884 "Found {} local HuggingFace model(s), total size: {}",
885 models.len(),
886 format_size(total_size_bytes)
887 );
888
889 LocalHfModelsResponse {
890 models,
891 cache_path: cache_path_str,
892 total_size_bytes,
893 total_size_display: format_size(total_size_bytes),
894 }
895 }
896
897 /// Delete a locally cached HuggingFace model by removing its directory from
898 /// the hub cache.
899 ///
900 /// `model_id` should be the full identifier, e.g. "bartowski/Qwen2.5-1.5B-Instruct-GGUF".
901 /// Returns `Ok(())` on success or an error string describing what went wrong.
902 pub fn delete_local_hf_model(model_id: String) -> Result<(), String> {
903 let cache_dir = hf_cache_dir().ok_or("HuggingFace cache directory not found.".to_string())?;
904
905 // Convert model_id (e.g. "org/name") to the cache directory name ("models--org--name")
906 let dir_name = format!("models--{}", model_id.replace('/', "--"));
907 let model_path = cache_dir.join(&dir_name);
908
909 if !model_path.exists() {
910 return Err(format!(
911 "Model cache directory not found: {}",
912 model_path.display()
913 ));
914 }
915
916 if !model_path.is_dir() {
917 return Err(format!(
918 "Expected a directory but found a file: {}",
919 model_path.display()
920 ));
921 }
922
923 debug!(
924 "Deleting local HuggingFace model: {} at {:?}",
925 model_id, model_path
926 );
927
928 fs::remove_dir_all(&model_path).map_err(|e| {
929 let msg = format!("Failed to delete model {}: {}", model_id, e);
930 warn!("{}", msg);
931 msg
932 })?;
933
934 debug!("Successfully deleted local model: {}", model_id);
935 Ok(())
936 }
937
938 /// Return the full list of models the inference engine supports, together with
939 /// flags indicating whether each one is fully downloaded, partially downloaded,
940 /// or not present at all.
941 pub fn list_supported_hf_models() -> SupportedHfModelsResponse {
942 let models = SUPPORTED_MODEL_INFO
943 .iter()
944 .map(|info| {
945 let cache_path = model_cache_path(info.id);
946 let local_size = cache_path
947 .as_ref()
948 .filter(|p| p.exists())
949 .map(dir_size)
950 .unwrap_or(0);
951
952 let has_cache = local_size > 0;
953 let is_complete = info.expected_size_bytes > 0
954 && local_size as f64
955 >= info.expected_size_bytes as f64 * DOWNLOAD_COMPLETE_THRESHOLD;
956
957 SupportedHfModel {
958 model_id: info.id.to_string(),
959 name: info.name.to_string(),
960 org: info.org.to_string(),
961 description: info.description.to_string(),
962 is_downloaded: has_cache && is_complete,
963 is_incomplete: has_cache && !is_complete,
964 local_size_bytes: local_size,
965 local_size_display: format_size(local_size),
966 expected_size_bytes: info.expected_size_bytes,
967 expected_size_display: format_size(info.expected_size_bytes),
968 }
969 })
970 .collect();
971
972 SupportedHfModelsResponse { models }
973 }
974
975 // ---------------------------------------------------------------------------
976 // Model download — platform-specific implementations
977 // ---------------------------------------------------------------------------
978 //
979 // The download function accepts a generic progress callback so that callers
980 // can forward progress to any event system without coupling to this crate.
981 //
982 // Callback signature: `Fn(ModelDownloadProgress) + Send + Sync + Clone + 'static`
983 // - Called every ~2 s from a native monitor thread.
984 // - Called once with `done: true` after the build completes.
985
986 // ── Apple (macOS + iOS): Metal-accelerated, all model types ─────────────────
987
988 #[cfg(any(
989 target_os = "macos",
990 target_os = "ios",
991 target_os = "tvos",
992 target_os = "visionos",
993 target_os = "watchos"
994 ))]
995 pub async fn download_model<F>(
996 model_id: String,
997 on_progress: F,
998 app_data_dir: Option<PathBuf>,
999 ) -> Result<(), String>
1000 where
1001 F: Fn(ModelDownloadProgress) + Send + Sync + Clone + 'static,
1002 {
1003 use crate::inference::models::{
1004 BARTOWSKI_QWEN25_1_5B_INSTRUCT_GGUF, BARTOWSKI_QWEN25_3B_INSTRUCT_GGUF,
1005 QWEN25_1_5B_GGUF_FILE, QWEN25_3B_GGUF_FILE,
1006 };
1007 use crate::inference::token::hf_token_source;
1008 use mistralrs::GgufModelBuilder;
1009 use std::sync::{
1010 atomic::{AtomicBool, Ordering},
1011 Arc,
1012 };
1013 use std::time::Duration;
1014
1015 // When app_data_dir is provided, seed the HF env vars so that hf-hub,
1016 // mistral.rs, and hf_cache_dir() all write inside the app's writable
1017 // container. This is required on every sandboxed platform:
1018 // - iOS: ~/.cache is outside the container (os error 1).
1019 // - macOS App Store: the app is sandboxed, ~/.cache is inaccessible.
1020 // It is harmless on non-sandboxed desktop builds (the cache simply
1021 // lives under the app data dir instead of ~/.cache).
1022 if let Some(ref data_dir) = app_data_dir {
1023 let hf_home = data_dir.join("models");
1024 let hf_hub_cache = hf_home.join("hub");
1025 fs::create_dir_all(&hf_hub_cache)
1026 .map_err(|e| format!("Cannot create HF cache dir: {e}"))?;
1027 std::env::set_var("HF_HUB_CACHE", &hf_hub_cache);
1028 std::env::set_var("HF_HOME", &hf_home);
1029 info!(
1030 "HF hub cache resolved to app data path: {}",
1031 hf_hub_cache.display()
1032 );
1033 }
1034
1035 if !SUPPORTED_MODELS.contains(&model_id.as_str()) {
1036 return Err(format!("Model {} is not supported.", model_id));
1037 }
1038
1039 let expected_size = SUPPORTED_MODEL_INFO
1040 .iter()
1041 .find(|m| m.id == model_id)
1042 .map(|m| m.expected_size_bytes)
1043 .unwrap_or(0);
1044
1045 info!(
1046 "Starting download for model: {} (expected ~{})",
1047 model_id,
1048 format_size(expected_size)
1049 );
1050
1051 clean_stale_lock_files(&model_id);
1052 repair_hf_cache_symlinks(&model_id);
1053 ensure_model_cache_dir(&model_id);
1054
1055 // Recover from interrupted downloads where the blob exists but the
1056 // snapshot entry was never created (process killed after download,
1057 // before hf-hub finalised the cache structure).
1058 let expected_files: Vec<&str> = match model_id.as_str() {
1059 BARTOWSKI_QWEN25_1_5B_INSTRUCT_GGUF => vec![QWEN25_1_5B_GGUF_FILE],
1060 BARTOWSKI_QWEN25_3B_INSTRUCT_GGUF => vec![QWEN25_3B_GGUF_FILE],
1061 _ => vec![],
1062 };
1063 if !expected_files.is_empty() {
1064 ensure_snapshot_entries(&model_id, &expected_files);
1065 }
1066
1067 // Emit an initial progress event based on the actual cache size so that
1068 // resumed downloads don't visually restart from 0%.
1069 let initial_bytes = model_cache_path(&model_id)
1070 .filter(|p| p.exists())
1071 .map(|p| dir_size(&p))
1072 .unwrap_or(0);
1073 on_progress(make_progress(
1074 &model_id,
1075 initial_bytes,
1076 expected_size,
1077 false,
1078 ));
1079
1080 // Shared flag so the monitor task knows when to stop.
1081 let finished = Arc::new(AtomicBool::new(false));
1082 let finished_clone = Arc::clone(&finished);
1083 let monitor_model_id = model_id.clone();
1084 let monitor_cb = on_progress.clone();
1085
1086 // Spawn a native OS thread to poll the cache directory size every 2 seconds.
1087 // We use std::thread instead of tokio::spawn because the model builder's
1088 // .build().await internally calls the synchronous hf-hub API (api.get()),
1089 // which blocks the tokio worker thread. This starves any tokio-spawned
1090 // monitor tasks, preventing progress updates. A native thread runs
1091 // independently of the tokio runtime and is immune to this starvation.
1092 let monitor_handle = std::thread::spawn(move || loop {
1093 std::thread::sleep(Duration::from_secs(2));
1094
1095 if finished_clone.load(Ordering::Relaxed) {
1096 break;
1097 }
1098
1099 let current_bytes = model_cache_path(&monitor_model_id)
1100 .filter(|p| p.exists())
1101 .map(|p| dir_size(&p))
1102 .unwrap_or(0);
1103
1104 monitor_cb(make_progress(
1105 &monitor_model_id,
1106 current_bytes,
1107 expected_size,
1108 false,
1109 ));
1110 });
1111
1112 // Build (and immediately drop) the model using the GGUF chat builder.
1113 // Building triggers the HuggingFace hub download.
1114 let result: Result<mistralrs::Model, anyhow::Error> = match model_id.as_str() {
1115 BARTOWSKI_QWEN25_1_5B_INSTRUCT_GGUF => {
1116 GgufModelBuilder::new(&model_id, vec![QWEN25_1_5B_GGUF_FILE])
1117 .with_token_source(hf_token_source())
1118 .build()
1119 .await
1120 }
1121 BARTOWSKI_QWEN25_3B_INSTRUCT_GGUF => {
1122 GgufModelBuilder::new(&model_id, vec![QWEN25_3B_GGUF_FILE])
1123 .with_token_source(hf_token_source())
1124 .build()
1125 .await
1126 }
1127 _ => Err(anyhow::anyhow!(
1128 "Model '{}' is not part of the chat-only release.",
1129 model_id
1130 )),
1131 };
1132
1133 finished.store(true, Ordering::Relaxed);
1134 let _ = monitor_handle.join();
1135
1136 match result {
1137 Ok(_model) => {
1138 let final_bytes = model_cache_path(&model_id)
1139 .filter(|p| p.exists())
1140 .map(|p| dir_size(&p))
1141 .unwrap_or(expected_size);
1142 on_progress(make_progress(&model_id, final_bytes, final_bytes, true));
1143 info!("Successfully downloaded model: {}", model_id);
1144 Ok(())
1145 }
1146 Err(e) => {
1147 let msg = format!("Failed to download model {}: {}", model_id, e);
1148 warn!("{}", msg);
1149 Err(msg)
1150 }
1151 }
1152 }
1153
1154 // ── Android: GGUF-only download ───────────────────────────────────────────
1155 //
1156 // On Android there is no Metal GPU, so only pre-quantized GGUF models
1157 // (CPU inference via candle) are supported. We mirror the Apple
1158 // implementation but restrict to GGUF presets and skip all Metal / ISQ
1159 // model builders.
1160 //
1161 // `app_data_dir` is **required** on Android: `dirs::home_dir()` returns
1162 // `None` inside the Android sandbox, which causes `hf_hub::Cache::default()`
1163 // to panic. The caller resolves the app data directory and passes it here.
1164
1165 #[cfg(target_os = "android")]
1166 pub async fn download_model<F>(
1167 model_id: String,
1168 on_progress: F,
1169 app_data_dir: Option<PathBuf>,
1170 ) -> Result<(), String>
1171 where
1172 F: Fn(ModelDownloadProgress) + Send + Sync + Clone + 'static,
1173 {
1174 use crate::inference::models::{
1175 BARTOWSKI_QWEN25_1_5B_INSTRUCT_GGUF, BARTOWSKI_QWEN25_3B_INSTRUCT_GGUF,
1176 QWEN25_1_5B_GGUF_FILE, QWEN25_1_5B_TOK_MODEL_ID, QWEN25_3B_GGUF_FILE,
1177 QWEN25_3B_TOK_MODEL_ID,
1178 };
1179 use crate::inference::token::hf_token_source;
1180 use hf_hub::Cache;
1181 use mistralrs::GgufModelBuilder;
1182 use mistralrs_core::GLOBAL_HF_CACHE;
1183 use std::sync::{
1184 atomic::{AtomicBool, Ordering},
1185 Arc,
1186 };
1187 use std::time::Duration;
1188
1189 let (gguf_file, tok_model_id) = match model_id.as_str() {
1190 id if id == BARTOWSKI_QWEN25_1_5B_INSTRUCT_GGUF => {
1191 (QWEN25_1_5B_GGUF_FILE, QWEN25_1_5B_TOK_MODEL_ID)
1192 }
1193 id if id == BARTOWSKI_QWEN25_3B_INSTRUCT_GGUF => {
1194 (QWEN25_3B_GGUF_FILE, QWEN25_3B_TOK_MODEL_ID)
1195 }
1196 _ => {
1197 return Err(format!(
1198 "Model '{}' is not supported on Android. \
1199 Only Qwen 2.5 1.5B or 3B (GGUF) are available.",
1200 model_id
1201 ));
1202 }
1203 };
1204
1205 // ── Resolve the HF cache path ─────────────────────────────────────────
1206 //
1207 // `app_data_dir` is required on Android — `dirs::home_dir()` returns None
1208 // in the Android sandbox and causes Cache::default() to panic.
1209 let resolved_app_data = app_data_dir
1210 .ok_or_else(|| "app_data_dir is required on Android for HF cache resolution".to_string())?;
1211
1212 let hf_home = resolved_app_data.join("models");
1213 let hf_hub_cache = hf_home.join("hub");
1214
1215 fs::create_dir_all(&hf_hub_cache)
1216 .map_err(|e| format!("Failed to create HF hub cache dir: {e}"))?;
1217
1218 // Seed GLOBAL_HF_CACHE so get_paths_gguf! never falls back to Cache::default().
1219 GLOBAL_HF_CACHE.get_or_init(|| Cache::new(hf_hub_cache.clone()));
1220
1221 // Keep env vars in sync so hf_hub_cache_dir() and hf_cache_dir() agree.
1222 std::env::set_var("HF_HUB_CACHE", &hf_hub_cache);
1223 std::env::set_var("HF_HOME", &hf_home);
1224
1225 info!(
1226 "Android HF hub cache resolved to: {}",
1227 hf_hub_cache.display()
1228 );
1229
1230 let expected_size = SUPPORTED_MODEL_INFO
1231 .iter()
1232 .find(|m| m.id == model_id)
1233 .map(|m| m.expected_size_bytes)
1234 .unwrap_or(0);
1235
1236 info!(
1237 "Starting Android download for model: {} (expected ~{})",
1238 model_id,
1239 format_size(expected_size)
1240 );
1241
1242 clean_stale_lock_files(&model_id);
1243 ensure_model_cache_dir(&model_id);
1244
1245 let initial_bytes = model_cache_path(&model_id)
1246 .filter(|p| p.exists())
1247 .map(|p| dir_size(&p))
1248 .unwrap_or(0);
1249 on_progress(make_progress(
1250 &model_id,
1251 initial_bytes,
1252 expected_size,
1253 false,
1254 ));
1255
1256 // Native monitor thread — same reasoning as the Apple implementation.
1257 let finished = Arc::new(AtomicBool::new(false));
1258 let finished_clone = Arc::clone(&finished);
1259 let monitor_model_id = model_id.clone();
1260 let monitor_cb = on_progress.clone();
1261
1262 let monitor_handle = std::thread::spawn(move || loop {
1263 std::thread::sleep(Duration::from_secs(2));
1264 if finished_clone.load(Ordering::Relaxed) {
1265 break;
1266 }
1267 let current_bytes = model_cache_path(&monitor_model_id)
1268 .filter(|p| p.exists())
1269 .map(|p| dir_size(&p))
1270 .unwrap_or(0);
1271 monitor_cb(make_progress(
1272 &monitor_model_id,
1273 current_bytes,
1274 expected_size,
1275 false,
1276 ));
1277 });
1278
1279 let result = GgufModelBuilder::new(&model_id, vec![gguf_file])
1280 .with_tok_model_id(tok_model_id)
1281 .with_token_source(hf_token_source())
1282 .with_logging()
1283 .build()
1284 .await;
1285
1286 finished.store(true, Ordering::Relaxed);
1287 let _ = monitor_handle.join();
1288
1289 match result {
1290 Ok(_model) => {
1291 let final_bytes = model_cache_path(&model_id)
1292 .filter(|p| p.exists())
1293 .map(|p| dir_size(&p))
1294 .unwrap_or(expected_size);
1295 on_progress(make_progress(&model_id, final_bytes, final_bytes, true));
1296 info!("Successfully downloaded model on Android: {}", model_id);
1297 Ok(())
1298 }
1299 Err(e) => {
1300 let msg = format!("Failed to download model {}: {}", model_id, e);
1301 warn!("{}", msg);
1302 Err(msg)
1303 }
1304 }
1305 }
1306
1307 // ── Other platforms: no-op ────────────────────────────────────────────────
1308
1309 #[cfg(not(any(
1310 target_os = "macos",
1311 target_os = "ios",
1312 target_os = "tvos",
1313 target_os = "visionos",
1314 target_os = "watchos",
1315 target_os = "android"
1316 )))]
1317 pub async fn download_model<F>(
1318 model_id: String,
1319 _on_progress: F,
1320 _app_data_dir: Option<PathBuf>,
1321 ) -> Result<(), String>
1322 where
1323 F: Fn(ModelDownloadProgress) + Send + Sync + Clone + 'static,
1324 {
1325 debug!(
1326 "download_model: not supported on this platform, ignoring request for {}.",
1327 model_id
1328 );
1329 Err(format!(
1330 "Model downloads are not supported on this platform (requested: {}).",
1331 model_id
1332 ))
1333 }