1use std::collections::HashMap;
2use std::future::Future;
3use std::path::{Path, PathBuf};
4use std::sync::{Arc, Mutex, OnceLock, Weak};
5use std::time::Instant;
6
7use console::Term;
8use hf_hub::api::tokio::{Api, ApiBuilder, ApiError, Progress};
9use hf_hub::{Cache, Repo, RepoType};
10use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle};
11use thiserror::Error;
12
13use crate::manifest::{paths_from_downloads, ModelComponent, ModelFile, ModelManifest};
14use crate::ModelPaths;
15
16fn hf_model_repo(repo_id: &str) -> Repo {
17 if let Some(revision) = crate::minimax_h3::repo_revision(repo_id) {
18 Repo::with_revision(repo_id.to_string(), RepoType::Model, revision.to_string())
19 } else {
20 Repo::new(repo_id.to_string(), RepoType::Model)
21 }
22}
23
24#[derive(Debug, Clone)]
26pub enum DownloadProgressEvent {
27 FileStart {
29 filename: String,
30 file_index: usize,
31 total_files: usize,
32 size_bytes: u64,
33 batch_bytes_downloaded: u64,
34 batch_bytes_total: u64,
35 batch_elapsed_ms: u64,
36 },
37 FileProgress {
39 filename: String,
40 file_index: usize,
41 bytes_downloaded: u64,
42 bytes_total: u64,
43 batch_bytes_downloaded: u64,
44 batch_bytes_total: u64,
45 batch_elapsed_ms: u64,
46 },
47 Status { message: String },
49 FileDone {
51 filename: String,
52 file_index: usize,
53 total_files: usize,
54 batch_bytes_downloaded: u64,
55 batch_bytes_total: u64,
56 batch_elapsed_ms: u64,
57 },
58}
59
60pub type DownloadProgressCallback = Arc<dyn Fn(DownloadProgressEvent) + Send + Sync>;
62
63#[derive(Debug, Clone, Default)]
65pub struct PullOptions {
66 pub skip_verify: bool,
68}
69
70#[derive(Debug, Error)]
71pub enum DownloadError {
72 #[error(transparent)]
73 ModelActivation(#[from] crate::ModelActivationError),
74
75 #[error(
76 "Model requires access approval on HuggingFace.\n\n 1. Visit: https://huggingface.co/{repo}\n 2. Accept the license agreement\n 3. Create a token at: https://huggingface.co/settings/tokens\n 4. Set: export HF_TOKEN=hf_...\n 5. Retry: mold pull {model}"
77 )]
78 GatedModel { repo: String, model: String },
79
80 #[error(
81 "Authentication required for repository {repo}.\n\n 1. Create a token at: https://huggingface.co/settings/tokens\n (select at least \"Read\" access)\n 2. Set: export HF_TOKEN=hf_...\n Or run: huggingface-cli login\n 3. Retry: mold pull {model}\n\n If HF_TOKEN is already set, it may be invalid or expired."
82 )]
83 Unauthorized { repo: String, model: String },
84
85 #[error("Download failed for {filename} from {repo}: {source}")]
86 DownloadFailed {
87 repo: String,
88 filename: String,
89 source: ApiError,
90 },
91
92 #[error("SHA-256 mismatch for {filename}\n Expected: {expected}\n Got: {actual}\n\nThe corrupted file has been removed. Re-run: mold pull {model}\nIf the file was intentionally updated on HuggingFace, use: mold pull {model} --skip-verify")]
93 Sha256Mismatch {
94 filename: String,
95 expected: String,
96 actual: String,
97 model: String,
98 },
99
100 #[error("Failed to build HuggingFace API client: {0}")]
101 ApiSetup(#[from] ApiError),
102
103 #[error("Failed to build sync HuggingFace API client: {0}")]
104 SyncApiSetup(String),
105
106 #[error("Sync download failed for {filename} from {repo}: {message}")]
107 SyncDownloadFailed {
108 repo: String,
109 filename: String,
110 message: String,
111 },
112
113 #[error("Missing component after download — this is a bug")]
114 MissingComponent,
115
116 #[error("{0}")]
117 Other(String),
118
119 #[error("IO error during file placement: {0}")]
120 FilePlacement(String),
121
122 #[error("Unknown model '{model}'. No manifest found.")]
123 UnknownModel { model: String },
124
125 #[error("Failed to save config: {0}")]
126 ConfigSave(String),
127
128 #[error("Recipe destination path '{dest}' escapes the per-recipe subdirectory")]
129 RecipePathTraversal { dest: String },
130
131 #[error("Civitai download requires CIVITAI_TOKEN.\n\n 1. Create a token at: https://civitai.com/user/account (Add API Key)\n 2. Set: export CIVITAI_TOKEN=...\n 3. Retry: mold pull {id}")]
132 MissingCivitaiToken { id: String },
133
134 #[error("Recipe HTTP fetch failed for {url}: status {status}{}", .body.as_ref().map(|b| format!(" — {b}")).unwrap_or_default())]
135 RecipeHttp {
136 url: String,
137 status: u16,
138 body: Option<String>,
139 },
140
141 #[error("Recipe transport error for {url}: {source}")]
142 RecipeTransport {
143 url: String,
144 #[source]
145 source: reqwest::Error,
146 },
147}
148
149fn gguf_header_contains_tensor(path: &std::path::Path, needle: &str) -> bool {
157 use std::io::Read;
158 let Ok(mut f) = std::fs::File::open(path) else {
159 return false;
160 };
161 let mut buf = vec![0u8; 4 * 1024 * 1024];
162 let Ok(n) = f.read(&mut buf) else {
163 return false;
164 };
165 buf.truncate(n);
166 if buf.len() < 4 || &buf[..4] != b"GGUF" {
167 return false;
168 }
169 buf.windows(needle.len()).any(|w| w == needle.as_bytes())
170}
171
172fn flux_reference_warning(manifest: &ModelManifest, models_dir: &Path) -> Option<String> {
178 if manifest.family != "flux" {
179 return None;
180 }
181 let xformer_file = manifest.files.iter().find(|f| {
182 f.component == ModelComponent::Transformer
183 && f.hf_filename.to_lowercase().ends_with(".gguf")
184 })?;
185 let xformer_path = models_dir.join(crate::manifest::storage_path(manifest, xformer_file));
186 if !xformer_path.exists() {
187 return None;
188 }
189 if gguf_header_contains_tensor(&xformer_path, "img_in.weight") {
191 return None;
192 }
193
194 let needs_guidance = !manifest.defaults.is_schnell;
195 let reference_candidates: &[&str] = if needs_guidance {
196 &["flux-dev:q8", "flux-dev:q6", "flux-dev:q4"]
197 } else {
198 &[
199 "flux-dev:q8",
200 "flux-dev:q6",
201 "flux-dev:q4",
202 "flux-schnell:q8",
203 "flux-schnell:q4",
204 ]
205 };
206 let have_reference = reference_candidates.iter().any(|name| {
207 let Some(m) = crate::manifest::find_manifest(name) else {
208 return false;
209 };
210 let Some(xf) = m
211 .files
212 .iter()
213 .find(|f| f.component == ModelComponent::Transformer)
214 else {
215 return false;
216 };
217 let path = models_dir.join(crate::manifest::storage_path(m, xf));
218 path.exists()
219 && gguf_header_contains_tensor(&path, "img_in.weight")
220 && (!needs_guidance
221 || gguf_header_contains_tensor(&path, "guidance_in.in_layer.weight"))
222 });
223 if have_reference {
224 return None;
225 }
226
227 let fix_cmd = if needs_guidance {
228 "mold pull flux-dev:q8"
229 } else {
230 "mold pull flux-dev:q8 (or flux-schnell:q8)"
231 };
232 Some(format!(
233 "Heads up: {} is a city96-format GGUF — it ships only the diffusion blocks. \
234 FLUX input embedding layers{} must be patched from a separate reference \
235 model at load time, and none is downloaded yet. Run `{fix_cmd}` before \
236 generating with {}.",
237 xformer_file.hf_filename,
238 if needs_guidance {
239 " (including dev-only guidance_in)"
240 } else {
241 ""
242 },
243 manifest.name,
244 ))
245}
246
247fn warn_if_flux_gguf_needs_reference(
256 manifest: &ModelManifest,
257 callback: Option<&DownloadProgressCallback>,
258) {
259 let Some(msg) = flux_reference_warning(manifest, &models_dir()) else {
260 return;
261 };
262 if let Some(cb) = callback {
263 cb(DownloadProgressEvent::Status {
264 message: format!("⚠ {msg}"),
265 });
266 } else {
267 let _ = console::Term::stderr().write_line(&format!("\n⚠ {msg}\n"));
268 }
269}
270
271fn resolve_hf_token() -> Option<String> {
274 if let Ok(token) = std::env::var("HF_TOKEN") {
275 let token = token.trim().to_string();
276 if !token.is_empty() {
277 return Some(token);
278 }
279 }
280 Cache::new(hf_cache_dir())
281 .token()
282 .or_else(|| Cache::from_env().token())
283}
284
285fn resolve_hf_token_for(explicit_token: Option<&str>) -> Option<String> {
286 explicit_token
287 .map(str::trim)
288 .filter(|token| !token.is_empty())
289 .map(str::to_string)
290 .or_else(resolve_hf_token)
291}
292
293fn models_dir() -> PathBuf {
304 static DIR: OnceLock<PathBuf> = OnceLock::new();
305 DIR.get_or_init(|| {
306 let dir = crate::Config::load_or_default().resolved_models_dir();
307 let _ = std::fs::create_dir_all(&dir);
308 dir
309 })
310 .clone()
311}
312
313fn hf_cache_dir() -> PathBuf {
316 static DIR: OnceLock<PathBuf> = OnceLock::new();
317 DIR.get_or_init(|| {
318 let dir = models_dir().join(".hf-cache");
319 let _ = std::fs::create_dir_all(&dir);
320 dir
321 })
322 .clone()
323}
324
325fn hardlink_or_copy(src: &std::path::Path, dst: &std::path::Path) -> Result<(), DownloadError> {
331 let real_src = src.canonicalize().map_err(|e| {
334 DownloadError::FilePlacement(format!(
335 "source file not found after download: {} ({e})",
336 src.display()
337 ))
338 })?;
339
340 if dst.exists() {
343 if let (Ok(src_meta), Ok(dst_meta)) = (real_src.metadata(), dst.metadata()) {
344 if src_meta.len() == dst_meta.len() {
345 return Ok(());
346 }
347 }
348 }
349
350 if dst.symlink_metadata().is_ok() {
356 let _ = std::fs::remove_file(dst);
357 }
358
359 if let Some(parent) = dst.parent() {
360 std::fs::create_dir_all(parent).map_err(|e| {
361 DownloadError::FilePlacement(format!(
362 "failed to create directory {}: {e}",
363 parent.display()
364 ))
365 })?;
366 }
367 match std::fs::hard_link(&real_src, dst) {
369 Ok(()) => return Ok(()),
370 Err(_e) => {
371 }
373 }
374 std::fs::copy(&real_src, dst).map_err(|e| {
376 DownloadError::FilePlacement(format!(
377 "failed to copy {} → {}: {e}",
378 real_src.display(),
379 dst.display()
380 ))
381 })?;
382 Ok(())
383}
384
385pub fn compute_sha256(path: &std::path::Path) -> anyhow::Result<String> {
387 use sha2::{Digest, Sha256};
388
389 let mut file = std::fs::File::open(path)?;
390 let mut hasher = Sha256::new();
391 std::io::copy(&mut file, &mut hasher)?;
392 Ok(format!("{:x}", hasher.finalize()))
393}
394
395pub fn verify_sha256(path: &std::path::Path, expected: &str) -> anyhow::Result<bool> {
403 Ok(compute_sha256(path)?.eq_ignore_ascii_case(expected))
404}
405
406pub fn pulling_marker_rel_path(model_name: &str) -> PathBuf {
410 let canonical = crate::manifest::resolve_model_name(model_name);
411 PathBuf::from(canonical.replace(':', "-")).join(".pulling")
412}
413
414pub fn pulling_marker_path_in(models_dir: &Path, model_name: &str) -> PathBuf {
416 models_dir.join(pulling_marker_rel_path(model_name))
417}
418
419fn pulling_marker_path(model_name: &str) -> PathBuf {
421 pulling_marker_path_in(&models_dir(), model_name)
422}
423
424fn write_pulling_marker(model_name: &str) -> Result<(), DownloadError> {
426 let path = pulling_marker_path(model_name);
427 if let Some(parent) = path.parent() {
428 std::fs::create_dir_all(parent).map_err(|e| {
429 DownloadError::FilePlacement(format!(
430 "failed to create directory for pull marker {}: {e}",
431 parent.display()
432 ))
433 })?;
434 }
435 std::fs::write(&path, model_name).map_err(|e| {
436 DownloadError::FilePlacement(format!(
437 "failed to write pull marker {}: {e}",
438 path.display()
439 ))
440 })
441}
442
443pub fn remove_pulling_marker(model_name: &str) {
445 let path = pulling_marker_path(model_name);
446 let _ = std::fs::remove_file(path);
447}
448
449pub fn has_pulling_marker(model_name: &str) -> bool {
451 let canonical = crate::manifest::resolve_model_name(model_name);
452 pulling_marker_path(&canonical).exists()
453}
454
455pub const SHA256_VERIFIED_SUFFIX: &str = ".sha256-verified";
468
469pub const RECIPE_PROGRESS_THROTTLE_MS: u64 = 250;
475
476pub fn sha256_marker_path(path: &Path) -> PathBuf {
479 let mut marker = path.as_os_str().to_os_string();
480 marker.push(SHA256_VERIFIED_SUFFIX);
481 PathBuf::from(marker)
482}
483
484pub fn has_sha256_marker(path: &Path) -> bool {
486 sha256_marker_path(path).exists()
487}
488
489pub fn write_sha256_marker(path: &Path, digest: &str) -> std::io::Result<()> {
494 let marker = sha256_marker_path(path);
495 let tmp = marker.with_extension(format!("sha256-verified.tmp.{}", std::process::id()));
496 std::fs::write(&tmp, format!("{digest}\n"))?;
497 std::fs::rename(&tmp, &marker)
498}
499
500fn verify_file_integrity(
514 clean_path: &std::path::Path,
515 file: &ModelFile,
516 model_name: &str,
517 skip_verify: bool,
518) -> Result<(), DownloadError> {
519 if skip_verify {
520 return Ok(());
521 }
522 let actual = match compute_sha256(clean_path) {
523 Ok(d) => d,
524 Err(e) => {
525 eprintln!(
529 "warning: failed to verify SHA-256 for {}: {e}",
530 file.hf_filename
531 );
532 return Ok(());
533 }
534 };
535 if let Some(expected) = file.sha256 {
536 if !actual.eq_ignore_ascii_case(expected) {
537 let _ = std::fs::remove_file(clean_path);
538 return Err(DownloadError::Sha256Mismatch {
539 filename: file.hf_filename.clone(),
540 expected: expected.to_string(),
541 actual,
542 model: model_name.to_string(),
543 });
544 }
545 }
546 if let Err(e) = write_sha256_marker(clean_path, &actual) {
547 eprintln!(
551 "warning: failed to write .sha256-verified marker for {}: {e}",
552 file.hf_filename
553 );
554 }
555 Ok(())
556}
557
558fn truncate_filename(name: &str, max_len: usize) -> String {
560 if name.len() <= max_len || max_len < 8 {
561 return name.to_string();
562 }
563 let suffix_len = max_len - 3; let start = name.len() - suffix_len;
566 format!("...{}", &name[start..])
567}
568
569fn filename_column_width() -> usize {
573 let term_width = Term::stderr().size().1 as usize;
574 term_width.saturating_sub(75).max(12)
575}
576
577#[derive(Clone)]
579struct DownloadProgress {
580 bar: ProgressBar,
581 max_msg_len: usize,
582 filename: String,
583}
584
585impl DownloadProgress {
586 fn new(bar: ProgressBar, max_msg_len: usize) -> Self {
587 Self {
588 bar,
589 max_msg_len,
590 filename: String::new(),
591 }
592 }
593}
594
595impl Progress for DownloadProgress {
596 async fn init(&mut self, size: usize, filename: &str) {
597 self.bar.set_length(size as u64);
598 self.filename = truncate_filename(filename, self.max_msg_len);
599 self.bar.set_message(self.filename.clone());
600 }
601
602 async fn update(&mut self, size: usize) {
603 self.bar.inc(size as u64);
604 }
605
606 async fn finish(&mut self) {
607 self.bar.finish_with_message(self.filename.clone());
608 }
609}
610
611#[derive(Clone)]
614struct CallbackProgress {
615 callback: DownloadProgressCallback,
616 file_index: usize,
617 total_files: usize,
618 batch_bytes_before_current: u64,
619 batch_bytes_total: u64,
620 batch_started_at: Instant,
621 shared: Arc<Mutex<CallbackProgressState>>,
622}
623
624struct CallbackProgressState {
625 accumulated: u64,
626 total: u64,
627 filename: String,
628 last_emit: Instant,
629}
630
631impl CallbackProgress {
632 fn new(
633 callback: DownloadProgressCallback,
634 file_index: usize,
635 total_files: usize,
636 batch_bytes_before_current: u64,
637 batch_bytes_total: u64,
638 batch_started_at: Instant,
639 ) -> Self {
640 Self {
641 callback,
642 file_index,
643 total_files,
644 batch_bytes_before_current,
645 batch_bytes_total,
646 batch_started_at,
647 shared: Arc::new(Mutex::new(CallbackProgressState {
648 accumulated: 0,
649 total: 0,
650 filename: String::new(),
651 last_emit: Instant::now(),
652 })),
653 }
654 }
655}
656
657impl Progress for CallbackProgress {
658 async fn init(&mut self, size: usize, filename: &str) {
659 let (fname, total) = {
660 let mut shared = self
661 .shared
662 .lock()
663 .expect("download progress mutex poisoned");
664 shared.total = size as u64;
665 shared.accumulated = 0;
666 shared.filename = filename.to_string();
667 shared.last_emit = Instant::now();
668 (shared.filename.clone(), shared.total)
669 };
670 (self.callback)(DownloadProgressEvent::FileStart {
671 filename: fname,
672 file_index: self.file_index,
673 total_files: self.total_files,
674 size_bytes: total,
675 batch_bytes_downloaded: self.batch_bytes_before_current,
676 batch_bytes_total: self.batch_bytes_total,
677 batch_elapsed_ms: self.batch_started_at.elapsed().as_millis() as u64,
678 });
679 }
680
681 async fn update(&mut self, size: usize) {
682 let mut shared = self
683 .shared
684 .lock()
685 .expect("download progress mutex poisoned");
686 shared.accumulated += size as u64;
687
688 let now = Instant::now();
689 let should_emit = now.duration_since(shared.last_emit).as_millis() >= 250
690 || shared.accumulated >= shared.total;
691 if !should_emit {
692 return;
693 }
694
695 shared.last_emit = now;
696 let filename = shared.filename.clone();
697 let accumulated = shared.accumulated;
698 let total = shared.total;
699 drop(shared);
700
701 (self.callback)(DownloadProgressEvent::FileProgress {
702 filename,
703 file_index: self.file_index,
704 bytes_downloaded: accumulated,
705 bytes_total: total,
706 batch_bytes_downloaded: self.batch_bytes_before_current + accumulated,
707 batch_bytes_total: self.batch_bytes_total,
708 batch_elapsed_ms: self.batch_started_at.elapsed().as_millis() as u64,
709 });
710 }
711
712 async fn finish(&mut self) {
713 let (fname, total) = {
714 let shared = self
715 .shared
716 .lock()
717 .expect("download progress mutex poisoned");
718 (shared.filename.clone(), shared.total)
719 };
720 (self.callback)(DownloadProgressEvent::FileDone {
721 filename: fname,
722 file_index: self.file_index,
723 total_files: self.total_files,
724 batch_bytes_downloaded: self.batch_bytes_before_current + total,
725 batch_bytes_total: self.batch_bytes_total,
726 batch_elapsed_ms: self.batch_started_at.elapsed().as_millis() as u64,
727 });
728 }
729}
730
731struct SyncDownloadProgress {
734 bar: ProgressBar,
735 max_msg_len: usize,
736 filename: String,
737}
738
739impl SyncDownloadProgress {
740 fn new(bar: ProgressBar, max_msg_len: usize) -> Self {
741 Self {
742 bar,
743 max_msg_len,
744 filename: String::new(),
745 }
746 }
747}
748
749impl hf_hub::api::Progress for SyncDownloadProgress {
750 fn init(&mut self, size: usize, filename: &str) {
751 self.bar.set_length(size as u64);
752 self.filename = truncate_filename(filename, self.max_msg_len);
753 self.bar.set_message(self.filename.clone());
754 }
755
756 fn update(&mut self, size: usize) {
757 self.bar.inc(size as u64);
758 }
759
760 fn finish(&mut self) {
761 self.bar.finish_with_message(self.filename.clone());
762 }
763}
764
765struct SyncCallbackProgress {
768 callback: DownloadProgressCallback,
769 started_at: Instant,
770 filename: String,
771 accumulated: u64,
772 total: u64,
773 last_emit: Instant,
774}
775
776impl SyncCallbackProgress {
777 fn new(callback: DownloadProgressCallback) -> Self {
778 Self {
779 callback,
780 started_at: Instant::now(),
781 filename: String::new(),
782 accumulated: 0,
783 total: 0,
784 last_emit: Instant::now(),
785 }
786 }
787}
788
789impl hf_hub::api::Progress for SyncCallbackProgress {
790 fn init(&mut self, size: usize, filename: &str) {
791 self.filename = filename.to_string();
792 self.accumulated = 0;
793 self.total = size as u64;
794 self.last_emit = Instant::now();
795 (self.callback)(DownloadProgressEvent::FileStart {
796 filename: self.filename.clone(),
797 file_index: 0,
798 total_files: 1,
799 size_bytes: self.total,
800 batch_bytes_downloaded: 0,
801 batch_bytes_total: self.total,
802 batch_elapsed_ms: self.started_at.elapsed().as_millis() as u64,
803 });
804 }
805
806 fn update(&mut self, size: usize) {
807 self.accumulated = self.accumulated.saturating_add(size as u64);
808 let now = Instant::now();
809 if now.duration_since(self.last_emit).as_millis() < 250 && self.accumulated < self.total {
810 return;
811 }
812 self.last_emit = now;
813 (self.callback)(DownloadProgressEvent::FileProgress {
814 filename: self.filename.clone(),
815 file_index: 0,
816 bytes_downloaded: self.accumulated,
817 bytes_total: self.total,
818 batch_bytes_downloaded: self.accumulated,
819 batch_bytes_total: self.total,
820 batch_elapsed_ms: self.started_at.elapsed().as_millis() as u64,
821 });
822 }
823
824 fn finish(&mut self) {
825 (self.callback)(DownloadProgressEvent::FileDone {
826 filename: self.filename.clone(),
827 file_index: 0,
828 total_files: 1,
829 batch_bytes_downloaded: self.total,
830 batch_bytes_total: self.total,
831 batch_elapsed_ms: self.started_at.elapsed().as_millis() as u64,
832 });
833 }
834}
835
836fn is_already_placed(
842 clean_path: &std::path::Path,
843 file: &ModelFile,
844 model_name: &str,
845 skip_verify: bool,
846) -> bool {
847 let size_ok = clean_path
848 .metadata()
849 .map(|m| m.len() == file.size_bytes)
850 .unwrap_or(false);
851 if !size_ok {
852 return false;
853 }
854 verify_file_integrity(clean_path, file, model_name, skip_verify).is_ok()
856}
857
858fn find_existing_placed_file(
861 models_dir: &std::path::Path,
862 manifest: &ModelManifest,
863 file: &ModelFile,
864 skip_verify: bool,
865) -> Result<Option<PathBuf>, DownloadError> {
866 let canonical_rel = crate::manifest::storage_path(manifest, file);
867 let canonical_path = models_dir.join(&canonical_rel);
868
869 for candidate_rel in crate::manifest::storage_path_candidates(manifest, file) {
870 let candidate_path = models_dir.join(candidate_rel);
871 if !is_already_placed(&candidate_path, file, &manifest.name, skip_verify) {
872 continue;
873 }
874 if candidate_path != canonical_path {
875 hardlink_or_copy(&candidate_path, &canonical_path)?;
876 verify_file_integrity(&canonical_path, file, &manifest.name, skip_verify)?;
877 }
878 return Ok(Some(canonical_path));
879 }
880
881 Ok(None)
882}
883
884fn require_manifest_acquisition(manifest: &ModelManifest) -> Result<(), DownloadError> {
894 let contains_gated_identity =
895 crate::require_model_activation(&manifest.name, Some(&manifest.family)).is_err()
896 || manifest.files.iter().any(|file| {
897 crate::require_model_activation(&file.hf_repo, Some(&manifest.family)).is_err()
898 || crate::require_model_activation(&file.hf_filename, Some(&manifest.family))
899 .is_err()
900 });
901 if contains_gated_identity {
902 crate::require_model_acquisition(&manifest.name, Some(&manifest.family))?;
903 let reviewed = crate::manifest::find_manifest(&manifest.name);
904 if !reviewed.is_some_and(|reviewed| std::ptr::eq(reviewed, manifest)) {
905 crate::require_model_activation("minimax-h3", Some("minimax-h3"))?;
906 }
907 } else {
908 crate::require_model_acquisition(&manifest.name, Some(&manifest.family))?;
909 }
910 Ok(())
911}
912
913pub async fn pull_model(
914 manifest: &ModelManifest,
915 opts: &PullOptions,
916) -> Result<ModelPaths, DownloadError> {
917 pull_model_with_hf_token(manifest, opts, None).await
918}
919
920async fn pull_model_with_hf_token(
921 manifest: &ModelManifest,
922 opts: &PullOptions,
923 hf_token: Option<&str>,
924) -> Result<ModelPaths, DownloadError> {
925 require_manifest_acquisition(manifest)?;
926 write_pulling_marker(&manifest.name)?;
927
928 let mut builder = ApiBuilder::from_env().with_cache_dir(hf_cache_dir());
929 if let Some(token) = resolve_hf_token_for(hf_token) {
930 builder = builder.with_token(Some(token));
931 }
932 let api = builder.build()?;
933
934 let multi = MultiProgress::with_draw_target(ProgressDrawTarget::stderr());
935 let msg_width = filename_column_width();
936 let bar_style = ProgressStyle::with_template(&format!(
937 " {{msg:<{msg_width}}} [{{bar:30.cyan/dim}}] {{bytes}}/{{total_bytes}} ({{bytes_per_sec}}, {{eta}})"
938 ))
939 .unwrap()
940 .progress_chars("━╸─");
941
942 let mdir = models_dir();
943 let mut downloads: Vec<(ModelComponent, PathBuf)> = Vec::new();
944
945 for file in &manifest.files {
946 if let Some(clean_path) =
947 find_existing_placed_file(&mdir, manifest, file, opts.skip_verify)?
948 {
949 downloads.push((file.component, clean_path));
950 continue;
951 }
952
953 let clean_path = mdir.join(crate::manifest::storage_path(manifest, file));
954
955 let bar = multi.add(ProgressBar::new(file.size_bytes));
956 bar.set_style(bar_style.clone());
957 bar.set_message(truncate_filename(&file.hf_filename, msg_width));
958
959 let clean_path = download_and_place_file(
960 &api,
961 file,
962 DownloadProgress::new(bar, msg_width),
963 &manifest.name,
964 &clean_path,
965 opts.skip_verify,
966 )
967 .await?;
968
969 downloads.push((file.component, clean_path));
970 }
971
972 warn_if_flux_gguf_needs_reference(manifest, None);
973
974 remove_pulling_marker(&manifest.name);
975 paths_from_downloads(&downloads, &manifest.family).ok_or(DownloadError::MissingComponent)
976}
977
978pub async fn pull_model_with_callback(
983 manifest: &ModelManifest,
984 callback: DownloadProgressCallback,
985 opts: &PullOptions,
986) -> Result<ModelPaths, DownloadError> {
987 pull_model_with_callback_and_hf_token(manifest, callback, opts, None).await
988}
989
990async fn pull_model_with_callback_and_hf_token(
991 manifest: &ModelManifest,
992 callback: DownloadProgressCallback,
993 opts: &PullOptions,
994 hf_token: Option<&str>,
995) -> Result<ModelPaths, DownloadError> {
996 require_manifest_acquisition(manifest)?;
997 write_pulling_marker(&manifest.name)?;
998
999 let mut builder = ApiBuilder::from_env().with_cache_dir(hf_cache_dir());
1000 if let Some(token) = resolve_hf_token_for(hf_token) {
1001 builder = builder.with_token(Some(token));
1002 }
1003 let api = builder.build()?;
1004
1005 let mdir = models_dir();
1006 let mut downloads: Vec<(ModelComponent, PathBuf)> = Vec::new();
1007
1008 let manifest_clone = manifest.clone();
1012 let skip_verify = opts.skip_verify;
1013 let mdir_clone = mdir.clone();
1014 let cb = callback.clone();
1015 let file_status: Vec<bool> = tokio::task::spawn_blocking(move || {
1016 let total = manifest_clone.files.len();
1017 manifest_clone
1018 .files
1019 .iter()
1020 .enumerate()
1021 .map(|(i, file)| {
1022 cb(DownloadProgressEvent::Status {
1023 message: format!(
1024 "Verifying file [{}/{}] {}...",
1025 i + 1,
1026 total,
1027 file.hf_filename
1028 ),
1029 });
1030 find_existing_placed_file(&mdir_clone, &manifest_clone, file, skip_verify)
1031 .map(|p| p.is_some())
1032 .unwrap_or(false)
1033 })
1034 .collect()
1035 })
1036 .await
1037 .map_err(|e| DownloadError::Other(format!("pre-scan task failed: {e}")))?;
1038
1039 let total_bytes_to_download: u64 = manifest
1040 .files
1041 .iter()
1042 .zip(file_status.iter())
1043 .filter(|(_, &placed)| !placed)
1044 .map(|(file, _)| file.size_bytes)
1045 .sum();
1046 let total_files_count = manifest.files.len();
1047 let mut completed_bytes = 0u64;
1048 let batch_started_at = Instant::now();
1049
1050 for (file_pos, (file, &already_placed)) in
1051 manifest.files.iter().zip(file_status.iter()).enumerate()
1052 {
1053 let clean_path = mdir.join(crate::manifest::storage_path(manifest, file));
1054
1055 if already_placed {
1056 let elapsed = batch_started_at.elapsed().as_millis() as u64;
1058 (callback)(DownloadProgressEvent::FileStart {
1059 filename: file.hf_filename.clone(),
1060 file_index: file_pos,
1061 total_files: total_files_count,
1062 size_bytes: file.size_bytes,
1063 batch_bytes_downloaded: completed_bytes,
1064 batch_bytes_total: total_bytes_to_download,
1065 batch_elapsed_ms: elapsed,
1066 });
1067 (callback)(DownloadProgressEvent::FileDone {
1068 filename: file.hf_filename.clone(),
1069 file_index: file_pos,
1070 total_files: total_files_count,
1071 batch_bytes_downloaded: completed_bytes,
1072 batch_bytes_total: total_bytes_to_download,
1073 batch_elapsed_ms: elapsed,
1074 });
1075 downloads.push((file.component, clean_path));
1076 continue;
1077 }
1078
1079 let progress = CallbackProgress::new(
1080 callback.clone(),
1081 file_pos,
1082 total_files_count,
1083 completed_bytes,
1084 total_bytes_to_download,
1085 batch_started_at,
1086 );
1087 let clean_path = download_and_place_file(
1088 &api,
1089 file,
1090 progress,
1091 &manifest.name,
1092 &clean_path,
1093 opts.skip_verify,
1094 )
1095 .await?;
1096
1097 downloads.push((file.component, clean_path));
1098 completed_bytes += file.size_bytes;
1099 }
1100
1101 warn_if_flux_gguf_needs_reference(manifest, Some(&callback));
1102
1103 remove_pulling_marker(&manifest.name);
1104 paths_from_downloads(&downloads, &manifest.family).ok_or(DownloadError::MissingComponent)
1105}
1106
1107async fn pull_model_files_only(
1112 manifest: &ModelManifest,
1113 opts: &PullOptions,
1114) -> Result<(), DownloadError> {
1115 pull_model_files_only_with_hf_token(manifest, opts, None).await
1116}
1117
1118async fn pull_model_files_only_with_hf_token(
1119 manifest: &ModelManifest,
1120 opts: &PullOptions,
1121 hf_token: Option<&str>,
1122) -> Result<(), DownloadError> {
1123 require_manifest_acquisition(manifest)?;
1124 write_pulling_marker(&manifest.name)?;
1125
1126 let mut builder = ApiBuilder::from_env().with_cache_dir(hf_cache_dir());
1127 if let Some(token) = resolve_hf_token_for(hf_token) {
1128 builder = builder.with_token(Some(token));
1129 }
1130 let api = builder.build()?;
1131
1132 let multi = MultiProgress::with_draw_target(ProgressDrawTarget::stderr());
1133 let msg_width = filename_column_width();
1134 let bar_style = ProgressStyle::with_template(&format!(
1135 " {{msg:<{msg_width}}} [{{bar:30.cyan/dim}}] {{bytes}}/{{total_bytes}} ({{bytes_per_sec}}, {{eta}})"
1136 ))
1137 .unwrap()
1138 .progress_chars("━╸─");
1139
1140 let mdir = models_dir();
1141
1142 for file in &manifest.files {
1143 if find_existing_placed_file(&mdir, manifest, file, opts.skip_verify)?.is_some() {
1144 continue;
1145 }
1146
1147 let clean_path = mdir.join(crate::manifest::storage_path(manifest, file));
1148
1149 let bar = multi.add(ProgressBar::new(file.size_bytes));
1150 bar.set_style(bar_style.clone());
1151 bar.set_message(truncate_filename(&file.hf_filename, msg_width));
1152
1153 download_and_place_file(
1154 &api,
1155 file,
1156 DownloadProgress::new(bar, msg_width),
1157 &manifest.name,
1158 &clean_path,
1159 opts.skip_verify,
1160 )
1161 .await?;
1162 }
1163
1164 remove_pulling_marker(&manifest.name);
1165 Ok(())
1166}
1167
1168async fn pull_model_files_only_with_callback_and_hf_token(
1169 manifest: &ModelManifest,
1170 callback: DownloadProgressCallback,
1171 opts: &PullOptions,
1172 hf_token: Option<&str>,
1173) -> Result<(), DownloadError> {
1174 require_manifest_acquisition(manifest)?;
1175 write_pulling_marker(&manifest.name)?;
1176
1177 let mut builder = ApiBuilder::from_env().with_cache_dir(hf_cache_dir());
1178 if let Some(token) = resolve_hf_token_for(hf_token) {
1179 builder = builder.with_token(Some(token));
1180 }
1181 let api = builder.build()?;
1182
1183 let mdir = models_dir();
1184
1185 let manifest_clone = manifest.clone();
1186 let skip_verify = opts.skip_verify;
1187 let mdir_clone = mdir.clone();
1188 let cb = callback.clone();
1189 let file_status: Vec<bool> = tokio::task::spawn_blocking(move || {
1190 let total = manifest_clone.files.len();
1191 manifest_clone
1192 .files
1193 .iter()
1194 .enumerate()
1195 .map(|(i, file)| {
1196 cb(DownloadProgressEvent::Status {
1197 message: format!(
1198 "Verifying file [{}/{}] {}...",
1199 i + 1,
1200 total,
1201 file.hf_filename
1202 ),
1203 });
1204 find_existing_placed_file(&mdir_clone, &manifest_clone, file, skip_verify)
1205 .map(|p| p.is_some())
1206 .unwrap_or(false)
1207 })
1208 .collect()
1209 })
1210 .await
1211 .map_err(|e| DownloadError::Other(format!("pre-scan task failed: {e}")))?;
1212 let total_bytes_to_download: u64 = manifest
1213 .files
1214 .iter()
1215 .zip(file_status.iter())
1216 .filter(|(_, &placed)| !placed)
1217 .map(|(file, _)| file.size_bytes)
1218 .sum();
1219 let total_files_count = manifest.files.len();
1220 let mut completed_bytes = 0u64;
1221 let batch_started_at = Instant::now();
1222
1223 for (file_pos, (file, &already_placed)) in
1224 manifest.files.iter().zip(file_status.iter()).enumerate()
1225 {
1226 let clean_path = mdir.join(crate::manifest::storage_path(manifest, file));
1227
1228 if already_placed {
1229 let elapsed = batch_started_at.elapsed().as_millis() as u64;
1230 (callback)(DownloadProgressEvent::FileStart {
1231 filename: file.hf_filename.clone(),
1232 file_index: file_pos,
1233 total_files: total_files_count,
1234 size_bytes: file.size_bytes,
1235 batch_bytes_downloaded: completed_bytes,
1236 batch_bytes_total: total_bytes_to_download,
1237 batch_elapsed_ms: elapsed,
1238 });
1239 (callback)(DownloadProgressEvent::FileDone {
1240 filename: file.hf_filename.clone(),
1241 file_index: file_pos,
1242 total_files: total_files_count,
1243 batch_bytes_downloaded: completed_bytes,
1244 batch_bytes_total: total_bytes_to_download,
1245 batch_elapsed_ms: elapsed,
1246 });
1247 continue;
1248 }
1249
1250 let progress = CallbackProgress::new(
1251 callback.clone(),
1252 file_pos,
1253 total_files_count,
1254 completed_bytes,
1255 total_bytes_to_download,
1256 batch_started_at,
1257 );
1258
1259 download_and_place_file(
1260 &api,
1261 file,
1262 progress,
1263 &manifest.name,
1264 &clean_path,
1265 opts.skip_verify,
1266 )
1267 .await?;
1268 completed_bytes += file.size_bytes;
1269 }
1270
1271 remove_pulling_marker(&manifest.name);
1272 Ok(())
1273}
1274
1275fn extract_http_status(err: &ApiError) -> Option<u16> {
1277 if let ApiError::RequestError(reqwest_err) = err {
1278 reqwest_err.status().map(|s| s.as_u16())
1279 } else {
1280 None
1281 }
1282}
1283
1284type HfFileDownloadFlight = tokio::sync::Mutex<()>;
1285type HfFileDownloadFlights =
1286 tokio::sync::Mutex<HashMap<(String, String), Weak<HfFileDownloadFlight>>>;
1287
1288async fn hf_file_download_flight(repo: &str, filename: &str) -> Arc<HfFileDownloadFlight> {
1295 static FLIGHTS: OnceLock<HfFileDownloadFlights> = OnceLock::new();
1296
1297 let flights = FLIGHTS.get_or_init(|| tokio::sync::Mutex::new(HashMap::new()));
1298 let mut flights = flights.lock().await;
1299 flights.retain(|_, flight| flight.strong_count() > 0);
1300
1301 let key = (repo.to_string(), filename.to_string());
1302 if let Some(flight) = flights.get(&key).and_then(Weak::upgrade) {
1303 return flight;
1304 }
1305
1306 let flight = Arc::new(tokio::sync::Mutex::new(()));
1307 flights.insert(key, Arc::downgrade(&flight));
1308 flight
1309}
1310
1311async fn with_hf_file_download_flight<T, Fut>(repo: &str, filename: &str, operation: Fut) -> T
1312where
1313 Fut: Future<Output = T>,
1314{
1315 let flight = hf_file_download_flight(repo, filename).await;
1316 let _flight_guard = flight.lock().await;
1317 operation.await
1318}
1319
1320async fn download_and_place_file<P: Progress + Clone + Send + Sync + 'static>(
1321 api: &Api,
1322 file: &ModelFile,
1323 progress: P,
1324 model_name: &str,
1325 clean_path: &Path,
1326 skip_verify: bool,
1327) -> Result<PathBuf, DownloadError> {
1328 with_hf_file_download_flight(&file.hf_repo, &file.hf_filename, async move {
1329 let repo = api.repo(hf_model_repo(&file.hf_repo));
1330 let hf_path = match repo
1331 .download_with_progress(&file.hf_filename, progress)
1332 .await
1333 {
1334 Ok(path) => path,
1335 Err(e) => {
1336 let status = extract_http_status(&e);
1337 let err_str = e.to_string();
1338 if status == Some(401)
1339 || err_str.contains("401")
1340 || err_str.contains("Unauthorized")
1341 {
1342 return Err(DownloadError::Unauthorized {
1343 repo: file.hf_repo.clone(),
1344 model: model_name.to_string(),
1345 });
1346 } else if status == Some(403)
1347 || err_str.contains("403")
1348 || err_str.contains("Forbidden")
1349 || err_str.contains("gated")
1350 || err_str.contains("Access denied")
1351 {
1352 return Err(DownloadError::GatedModel {
1353 repo: file.hf_repo.clone(),
1354 model: model_name.to_string(),
1355 });
1356 } else {
1357 return Err(DownloadError::DownloadFailed {
1358 repo: file.hf_repo.clone(),
1359 filename: file.hf_filename.clone(),
1360 source: e,
1361 });
1362 }
1363 }
1364 };
1365
1366 hardlink_or_copy(&hf_path, clean_path)?;
1367 verify_file_integrity(clean_path, file, model_name, skip_verify)?;
1368 Ok(clean_path.to_path_buf())
1369 })
1370 .await
1371}
1372
1373fn require_single_file_acquisition(
1376 hf_repo: &str,
1377 hf_filename: &str,
1378 target_subdir: Option<&str>,
1379) -> Result<(), DownloadError> {
1380 crate::require_model_activation(hf_repo, None)?;
1381 crate::require_model_activation(hf_filename, None)?;
1382 if let Some(target_subdir) = target_subdir {
1383 crate::require_model_activation(target_subdir, None)?;
1384 }
1385 Ok(())
1386}
1387
1388pub fn download_single_file_sync(
1396 hf_repo: &str,
1397 hf_filename: &str,
1398 target_subdir: Option<&str>,
1399) -> Result<PathBuf, DownloadError> {
1400 require_single_file_acquisition(hf_repo, hf_filename, target_subdir)?;
1401
1402 let msg_width = filename_column_width();
1403 let bar_style = ProgressStyle::with_template(&format!(
1404 " {{msg:<{msg_width}}} [{{bar:30.cyan/dim}}] {{bytes}}/{{total_bytes}} ({{bytes_per_sec}}, {{eta}})"
1405 ))
1406 .unwrap()
1407 .progress_chars("━╸─");
1408 let bar = ProgressBar::new(0);
1409 bar.set_style(bar_style);
1410 bar.set_message(truncate_filename(hf_filename, msg_width));
1411 let progress = SyncDownloadProgress::new(bar, msg_width);
1412 download_single_file_sync_with_adapter(
1413 &models_dir(),
1414 hf_repo,
1415 hf_filename,
1416 target_subdir,
1417 progress,
1418 )
1419}
1420
1421pub fn download_single_file_sync_with_progress(
1426 hf_repo: &str,
1427 hf_filename: &str,
1428 target_subdir: Option<&str>,
1429 callback: DownloadProgressCallback,
1430) -> Result<PathBuf, DownloadError> {
1431 require_single_file_acquisition(hf_repo, hf_filename, target_subdir)?;
1432
1433 download_single_file_sync_with_adapter(
1434 &models_dir(),
1435 hf_repo,
1436 hf_filename,
1437 target_subdir,
1438 SyncCallbackProgress::new(callback),
1439 )
1440}
1441
1442pub fn download_single_file_sync_with_progress_in(
1446 models_root: &Path,
1447 hf_repo: &str,
1448 hf_filename: &str,
1449 target_subdir: Option<&str>,
1450 callback: DownloadProgressCallback,
1451) -> Result<PathBuf, DownloadError> {
1452 require_single_file_acquisition(hf_repo, hf_filename, target_subdir)?;
1453
1454 download_single_file_sync_with_adapter(
1455 models_root,
1456 hf_repo,
1457 hf_filename,
1458 target_subdir,
1459 SyncCallbackProgress::new(callback),
1460 )
1461}
1462
1463pub fn planned_single_file_path(hf_filename: &str, target_subdir: &str) -> PathBuf {
1470 planned_single_file_path_in(&models_dir(), hf_filename, target_subdir)
1471}
1472
1473pub fn planned_single_file_path_in(
1476 models_root: &Path,
1477 hf_filename: &str,
1478 target_subdir: &str,
1479) -> PathBuf {
1480 let leaf = hf_filename.rsplit('/').next().unwrap_or(hf_filename);
1481 models_root.join(target_subdir).join(leaf)
1482}
1483
1484fn download_single_file_sync_with_adapter<P>(
1485 models_root: &Path,
1486 hf_repo: &str,
1487 hf_filename: &str,
1488 target_subdir: Option<&str>,
1489 progress: P,
1490) -> Result<PathBuf, DownloadError>
1491where
1492 P: hf_hub::api::Progress,
1493{
1494 require_single_file_acquisition(hf_repo, hf_filename, target_subdir)?;
1498
1499 use hf_hub::api::sync::ApiBuilder;
1500
1501 let mut builder = ApiBuilder::from_env()
1502 .with_cache_dir(models_root.join(".hf-cache"))
1503 .with_progress(false);
1504 if let Some(token) = resolve_hf_token() {
1505 builder = builder.with_token(Some(token));
1506 }
1507 let api = builder
1508 .build()
1509 .map_err(|e| DownloadError::SyncApiSetup(e.to_string()))?;
1510 let repo = api.repo(hf_model_repo(hf_repo));
1511 let hf_path = repo
1512 .download_with_progress(hf_filename, progress)
1513 .map_err(|e| {
1514 let err_str = e.to_string();
1515 if err_str.contains("401") || err_str.contains("Unauthorized") {
1516 DownloadError::Unauthorized {
1517 repo: hf_repo.to_string(),
1518 model: String::new(),
1519 }
1520 } else if err_str.contains("403")
1521 || err_str.contains("Forbidden")
1522 || err_str.contains("gated")
1523 || err_str.contains("Access denied")
1524 {
1525 DownloadError::GatedModel {
1526 repo: hf_repo.to_string(),
1527 model: String::new(),
1528 }
1529 } else {
1530 DownloadError::SyncDownloadFailed {
1531 repo: hf_repo.to_string(),
1532 filename: hf_filename.to_string(),
1533 message: err_str,
1534 }
1535 }
1536 })?;
1537
1538 if let Some(subdir) = target_subdir {
1540 let clean_path = planned_single_file_path_in(models_root, hf_filename, subdir);
1541 hardlink_or_copy(&hf_path, &clean_path)?;
1542 Ok(clean_path)
1543 } else {
1544 Ok(hf_path)
1545 }
1546}
1547
1548pub fn cached_file_path_in_mold_cache(hf_repo: &str, hf_filename: &str) -> Option<PathBuf> {
1560 let cache = Cache::new(hf_cache_dir());
1561 let repo = cache.repo(hf_model_repo(hf_repo));
1562 repo.get(hf_filename)
1563}
1564
1565pub fn cached_file_path(
1571 hf_repo: &str,
1572 hf_filename: &str,
1573 target_subdir: Option<&str>,
1574) -> Option<PathBuf> {
1575 cached_file_path_in(&models_dir(), hf_repo, hf_filename, target_subdir)
1576}
1577
1578pub fn cached_file_path_in(
1582 models_root: &Path,
1583 hf_repo: &str,
1584 hf_filename: &str,
1585 target_subdir: Option<&str>,
1586) -> Option<PathBuf> {
1587 if let Some(subdir) = target_subdir {
1589 let clean_path = planned_single_file_path_in(models_root, hf_filename, subdir);
1590 if clean_path.exists() {
1591 return Some(clean_path);
1592 }
1593 }
1594
1595 let new_cache = Cache::new(models_root.join(".hf-cache"));
1597 let new_repo = new_cache.repo(hf_model_repo(hf_repo));
1598 if let Some(path) = new_repo.get(hf_filename) {
1599 return Some(path);
1600 }
1601
1602 let old_cache = Cache::new(models_root.to_path_buf());
1604 let old_repo = old_cache.repo(hf_model_repo(hf_repo));
1605 if let Some(path) = old_repo.get(hf_filename) {
1606 return Some(path);
1607 }
1608
1609 let default_cache = Cache::from_env();
1611 let default_repo = default_cache.repo(hf_model_repo(hf_repo));
1612 default_repo.get(hf_filename)
1613}
1614
1615pub fn cached_file_path_existing_only(
1622 models_root: &Path,
1623 hf_repo: &str,
1624 hf_filename: &str,
1625 target_subdir: Option<&str>,
1626) -> Option<PathBuf> {
1627 if let Some(subdir) = target_subdir {
1628 let clean_path = planned_single_file_path_in(models_root, hf_filename, subdir);
1629 if clean_path.exists() {
1630 return Some(clean_path);
1631 }
1632 }
1633
1634 let lookup = |root: &Path| {
1635 root.is_dir().then(|| {
1636 Cache::new(root.to_path_buf())
1637 .repo(hf_model_repo(hf_repo))
1638 .get(hf_filename)
1639 })?
1640 };
1641 lookup(&models_root.join(".hf-cache"))
1642 .or_else(|| lookup(models_root))
1643 .or_else(|| {
1644 let cache = Cache::from_env();
1645 cache
1646 .path()
1647 .is_dir()
1648 .then(|| cache.repo(hf_model_repo(hf_repo)).get(hf_filename))?
1649 })
1650}
1651
1652fn manifest_uses_files_only_pull(manifest: &ModelManifest) -> bool {
1660 manifest.is_utility()
1661 || matches!(
1662 manifest.family.as_str(),
1663 "ltx2-control" | "ltx2-camera-control"
1664 )
1665}
1666
1667pub async fn pull_and_configure(
1671 model: &str,
1672 opts: &PullOptions,
1673) -> Result<(crate::Config, Option<ModelPaths>), DownloadError> {
1674 use crate::config::Config;
1675 use crate::manifest::{find_manifest, resolve_model_name};
1676
1677 let canonical = resolve_model_name(model);
1678
1679 let manifest = find_manifest(&canonical).ok_or_else(|| DownloadError::UnknownModel {
1680 model: model.to_string(),
1681 })?;
1682
1683 if manifest_uses_files_only_pull(manifest) {
1687 pull_model_files_only(manifest, opts).await?;
1688 let config = Config::load_or_default();
1689 return Ok((config, None));
1690 }
1691
1692 if manifest.is_upscaler() {
1695 pull_model_files_only(manifest, opts).await?;
1696
1697 let mdir = models_dir();
1699 let weights_file = manifest
1700 .files
1701 .iter()
1702 .find(|f| f.component == crate::manifest::ModelComponent::Upscaler)
1703 .ok_or(DownloadError::MissingComponent)?;
1704 let weights_path = mdir.join(crate::manifest::storage_path(manifest, weights_file));
1705
1706 let mut config = Config::load_or_default();
1707 let model_config = crate::config::ModelConfig {
1708 transformer: Some(weights_path.to_string_lossy().to_string()),
1709 family: Some("upscaler".to_string()),
1710 ..Default::default()
1711 };
1712 config.upsert_model(manifest.name.clone(), model_config);
1713 config
1714 .save()
1715 .map_err(|e| DownloadError::ConfigSave(e.to_string()))?;
1716
1717 return Ok((config, None));
1718 }
1719
1720 let paths = pull_model(manifest, opts).await?;
1721
1722 let mut config = Config::load_or_default();
1723 let model_config = manifest.to_model_config(&paths);
1724
1725 if !Config::exists_on_disk() {
1727 config.default_model = manifest.name.clone();
1728 }
1729
1730 config.upsert_model(manifest.name.clone(), model_config);
1731 config
1732 .save()
1733 .map_err(|e| DownloadError::ConfigSave(e.to_string()))?;
1734
1735 Ok((config, Some(paths)))
1736}
1737
1738pub async fn pull_and_configure_with_callback(
1741 model: &str,
1742 callback: DownloadProgressCallback,
1743 opts: &PullOptions,
1744) -> Result<(crate::Config, Option<ModelPaths>), DownloadError> {
1745 pull_and_configure_with_callback_and_hf_token(model, callback, opts, None).await
1746}
1747
1748pub async fn pull_and_configure_with_callback_and_hf_token(
1752 model: &str,
1753 callback: DownloadProgressCallback,
1754 opts: &PullOptions,
1755 hf_token: Option<&str>,
1756) -> Result<(crate::Config, Option<ModelPaths>), DownloadError> {
1757 use crate::config::Config;
1758 use crate::manifest::{find_manifest, resolve_model_name};
1759
1760 let canonical = resolve_model_name(model);
1761
1762 let manifest = find_manifest(&canonical).ok_or_else(|| DownloadError::UnknownModel {
1763 model: model.to_string(),
1764 })?;
1765
1766 if manifest_uses_files_only_pull(manifest) {
1770 pull_model_files_only_with_callback_and_hf_token(manifest, callback, opts, hf_token)
1771 .await?;
1772 let config = Config::load_or_default();
1773 return Ok((config, None));
1774 }
1775
1776 if manifest.is_upscaler() {
1778 pull_model_files_only_with_callback_and_hf_token(manifest, callback, opts, hf_token)
1779 .await?;
1780
1781 let mdir = models_dir();
1782 let weights_file = manifest
1783 .files
1784 .iter()
1785 .find(|f| f.component == crate::manifest::ModelComponent::Upscaler)
1786 .ok_or(DownloadError::MissingComponent)?;
1787 let weights_path = mdir.join(crate::manifest::storage_path(manifest, weights_file));
1788
1789 let mut config = Config::load_or_default();
1790 let model_config = crate::config::ModelConfig {
1791 transformer: Some(weights_path.to_string_lossy().to_string()),
1792 family: Some("upscaler".to_string()),
1793 ..Default::default()
1794 };
1795 config.upsert_model(manifest.name.clone(), model_config);
1796 config
1797 .save()
1798 .map_err(|e| DownloadError::ConfigSave(e.to_string()))?;
1799
1800 return Ok((config, None));
1801 }
1802
1803 let paths = pull_model_with_callback_and_hf_token(manifest, callback, opts, hf_token).await?;
1804
1805 let mut config = Config::load_or_default();
1806 let model_config = manifest.to_model_config(&paths);
1807
1808 if !Config::exists_on_disk() {
1809 config.default_model = manifest.name.clone();
1810 }
1811
1812 config.upsert_model(manifest.name.clone(), model_config);
1813 config
1814 .save()
1815 .map_err(|e| DownloadError::ConfigSave(e.to_string()))?;
1816
1817 Ok((config, Some(paths)))
1818}
1819
1820pub fn resolve_civitai_token() -> Option<String> {
1827 std::env::var("CIVITAI_TOKEN").ok().and_then(|t| {
1828 let trimmed = t.trim().to_string();
1829 if trimmed.is_empty() {
1830 None
1831 } else {
1832 Some(trimmed)
1833 }
1834 })
1835}
1836
1837pub fn civitai_auth_or_error(id: &str) -> Result<RecipeAuth, DownloadError> {
1841 match resolve_civitai_token() {
1842 Some(t) => Ok(RecipeAuth::Bearer(t)),
1843 None => Err(DownloadError::MissingCivitaiToken { id: id.to_string() }),
1844 }
1845}
1846
1847pub fn companion_present_on_disk(
1863 models_dir: &Path,
1864 manifest: &crate::manifest::ModelManifest,
1865) -> bool {
1866 if pulling_marker_path_in(models_dir, &manifest.name).exists() {
1867 return false;
1868 }
1869 manifest.files.iter().all(|f| {
1870 let storage = crate::manifest::storage_path(manifest, f);
1871 let path = models_dir.join(storage);
1872 if !path.exists() {
1873 return false;
1874 }
1875 if f.sha256.is_some() {
1876 return sha256_marker_path(&path).exists();
1877 }
1878 if f.size_bytes > 0 {
1879 return std::fs::metadata(&path)
1880 .map(|m| m.len() == f.size_bytes)
1881 .unwrap_or(false);
1882 }
1883 true
1884 })
1885}
1886
1887fn recipe_file_is_placed(dest: &Path, file: &RecipeFetchFile<'_>) -> bool {
1901 if !dest.exists() {
1902 return false;
1903 }
1904 if has_sha256_marker(dest) {
1905 return true;
1906 }
1907 match (file.sha256, file.size_bytes) {
1908 (Some(_), _) => sha256_marker_path(dest).exists(),
1909 (None, Some(expected)) => std::fs::metadata(dest)
1910 .map(|m| m.len() == expected)
1911 .unwrap_or(false),
1912 (None, None) => sha256_marker_path(dest).exists(),
1913 }
1914}
1915
1916pub fn catalog_entry_installed(models_dir: &Path, id: &str, files: &[RecipeFetchFile<'_>]) -> bool {
1932 if files.is_empty() {
1933 return false;
1934 }
1935 if pulling_marker_path_in(models_dir, id).exists() {
1936 return false;
1937 }
1938 let sanitized = sanitize_recipe_id(id);
1939 let subdir_root = models_dir.join(&sanitized);
1940 files.iter().all(|f| {
1941 let Ok(dest) = resolve_recipe_dest(&subdir_root, f.dest) else {
1942 return false;
1943 };
1944 recipe_file_is_placed(&dest, f)
1945 })
1946}
1947
1948pub fn missing_companions_from_json(
1958 companions_json: Option<&str>,
1959 models_dir: &Path,
1960) -> Vec<&'static crate::manifest::ModelManifest> {
1961 let Some(json) = companions_json else {
1962 return Vec::new();
1963 };
1964 let names: Vec<String> = match serde_json::from_str(json) {
1965 Ok(n) => n,
1966 Err(_) => return Vec::new(),
1967 };
1968 missing_companions(&names, models_dir)
1969}
1970
1971pub fn missing_companions(
1974 names: &[String],
1975 models_dir: &Path,
1976) -> Vec<&'static crate::manifest::ModelManifest> {
1977 let mut out = Vec::with_capacity(names.len());
1978 for name in names {
1979 let Some(manifest) = crate::manifest::find_manifest(name) else {
1980 tracing::warn!(
1981 companion = %name,
1982 "skipping companion with no synthetic manifest in this build",
1983 );
1984 continue;
1985 };
1986 if companion_present_on_disk(models_dir, manifest) {
1987 continue;
1988 }
1989 out.push(manifest);
1990 }
1991 out
1992}
1993
1994#[derive(Debug, Clone)]
2008pub struct RecipeFetchFile<'a> {
2009 pub url: &'a str,
2011 pub dest: &'a str,
2015 pub sha256: Option<&'a str>,
2017 pub size_bytes: Option<u64>,
2020}
2021
2022#[derive(Debug, Clone, PartialEq, Eq)]
2024pub enum RecipeAuth {
2025 None,
2027 Bearer(String),
2031}
2032
2033pub fn sanitize_recipe_id(id: &str) -> String {
2037 id.replace(':', "-")
2038}
2039
2040fn resolve_recipe_dest(subdir_root: &Path, dest: &str) -> Result<PathBuf, DownloadError> {
2044 let candidate = Path::new(dest);
2045 if candidate.is_absolute() {
2046 return Err(DownloadError::RecipePathTraversal {
2047 dest: dest.to_string(),
2048 });
2049 }
2050 for component in candidate.components() {
2051 match component {
2052 std::path::Component::Normal(_) => {}
2053 _ => {
2057 return Err(DownloadError::RecipePathTraversal {
2058 dest: dest.to_string(),
2059 });
2060 }
2061 }
2062 }
2063 Ok(subdir_root.join(candidate))
2064}
2065
2066pub async fn fetch_recipe(
2076 id: &str,
2077 files: &[RecipeFetchFile<'_>],
2078 auth: RecipeAuth,
2079 models_dir: &Path,
2080 progress: Option<DownloadProgressCallback>,
2081 opts: &PullOptions,
2082) -> Result<Vec<PathBuf>, DownloadError> {
2083 crate::require_model_activation(id, None)?;
2088 for file in files {
2089 crate::require_model_activation(file.url, None)?;
2090 crate::require_model_activation(file.dest, None)?;
2091 }
2092
2093 let sanitized = sanitize_recipe_id(id);
2094 let subdir_root = models_dir.join(&sanitized);
2095
2096 let resolved: Vec<PathBuf> = files
2099 .iter()
2100 .map(|f| resolve_recipe_dest(&subdir_root, f.dest))
2101 .collect::<Result<Vec<_>, _>>()?;
2102
2103 std::fs::create_dir_all(&subdir_root).map_err(|e| {
2104 DownloadError::FilePlacement(format!(
2105 "failed to create recipe subdir {}: {e}",
2106 subdir_root.display()
2107 ))
2108 })?;
2109
2110 let marker = pulling_marker_path_in(models_dir, id);
2111 if let Some(parent) = marker.parent() {
2112 let _ = std::fs::create_dir_all(parent);
2113 }
2114 std::fs::write(&marker, id).map_err(|e| {
2115 DownloadError::FilePlacement(format!(
2116 "failed to write recipe marker {}: {e}",
2117 marker.display()
2118 ))
2119 })?;
2120
2121 let result = fetch_recipe_inner(id, files, &resolved, auth, progress, opts).await;
2122 let _ = std::fs::remove_file(&marker);
2125 result
2126}
2127
2128async fn fetch_recipe_inner(
2129 id: &str,
2130 files: &[RecipeFetchFile<'_>],
2131 resolved: &[PathBuf],
2132 auth: RecipeAuth,
2133 progress: Option<DownloadProgressCallback>,
2134 opts: &PullOptions,
2135) -> Result<Vec<PathBuf>, DownloadError> {
2136 use std::io::Write;
2137
2138 let client = reqwest::Client::builder()
2139 .user_agent(concat!("mold/", env!("CARGO_PKG_VERSION")))
2140 .build()
2141 .map_err(|e| DownloadError::Other(format!("failed to build HTTP client: {e}")))?;
2142
2143 let total_files = files.len();
2144 let batch_bytes_total: u64 = files.iter().filter_map(|f| f.size_bytes).sum();
2145 let mut batch_bytes_downloaded: u64 = 0;
2146 let started = Instant::now();
2147
2148 for (file_index, (file, dest_path)) in files.iter().zip(resolved.iter()).enumerate() {
2149 if let Some(parent) = dest_path.parent() {
2150 std::fs::create_dir_all(parent).map_err(|e| {
2151 DownloadError::FilePlacement(format!(
2152 "failed to create directory {}: {e}",
2153 parent.display()
2154 ))
2155 })?;
2156 }
2157
2158 let already_placed = recipe_file_is_placed(dest_path, file);
2171 if already_placed {
2172 let size_bytes = file
2173 .size_bytes
2174 .unwrap_or_else(|| std::fs::metadata(dest_path).map(|m| m.len()).unwrap_or(0));
2175 if let Some(cb) = progress.as_deref() {
2176 cb(DownloadProgressEvent::FileStart {
2177 filename: file.dest.to_string(),
2178 file_index,
2179 total_files,
2180 size_bytes,
2181 batch_bytes_downloaded,
2182 batch_bytes_total,
2183 batch_elapsed_ms: started.elapsed().as_millis() as u64,
2184 });
2185 }
2186 batch_bytes_downloaded = batch_bytes_downloaded.saturating_add(size_bytes);
2187 if let Some(cb) = progress.as_deref() {
2188 cb(DownloadProgressEvent::FileDone {
2189 filename: file.dest.to_string(),
2190 file_index,
2191 total_files,
2192 batch_bytes_downloaded,
2193 batch_bytes_total,
2194 batch_elapsed_ms: started.elapsed().as_millis() as u64,
2195 });
2196 }
2197 continue;
2198 }
2199
2200 let mut req = client.get(file.url);
2201 if let RecipeAuth::Bearer(token) = &auth {
2202 req = req.bearer_auth(token);
2203 }
2204 let resp = req
2205 .send()
2206 .await
2207 .map_err(|e| DownloadError::RecipeTransport {
2208 url: file.url.to_string(),
2209 source: e,
2210 })?;
2211 if !resp.status().is_success() {
2212 let status = resp.status().as_u16();
2213 let body = resp.text().await.ok().map(|b| {
2214 let mut t = b.trim().to_string();
2215 if t.len() > 200 {
2216 t.truncate(200);
2217 }
2218 t
2219 });
2220 return Err(DownloadError::RecipeHttp {
2221 url: file.url.to_string(),
2222 status,
2223 body,
2224 });
2225 }
2226
2227 let content_length = resp.content_length();
2228 let size_bytes = file.size_bytes.or(content_length).unwrap_or(0);
2229
2230 if let Some(cb) = progress.as_deref() {
2231 cb(DownloadProgressEvent::FileStart {
2232 filename: file.dest.to_string(),
2233 file_index,
2234 total_files,
2235 size_bytes,
2236 batch_bytes_downloaded,
2237 batch_bytes_total,
2238 batch_elapsed_ms: started.elapsed().as_millis() as u64,
2239 });
2240 }
2241
2242 let mut bytes_downloaded: u64 = 0;
2243 let mut out = std::fs::File::create(dest_path).map_err(|e| {
2244 DownloadError::FilePlacement(format!("failed to create {}: {e}", dest_path.display()))
2245 })?;
2246 let mut resp = resp;
2247 let mut last_emit = Instant::now();
2252 let mut last_emit_bytes: u64 = 0;
2253 while let Some(chunk) = resp
2254 .chunk()
2255 .await
2256 .map_err(|e| DownloadError::RecipeTransport {
2257 url: file.url.to_string(),
2258 source: e,
2259 })?
2260 {
2261 out.write_all(&chunk).map_err(|e| {
2262 DownloadError::FilePlacement(format!(
2263 "failed to write to {}: {e}",
2264 dest_path.display()
2265 ))
2266 })?;
2267 bytes_downloaded += chunk.len() as u64;
2268 batch_bytes_downloaded += chunk.len() as u64;
2269 if let Some(cb) = progress.as_deref() {
2270 let now = Instant::now();
2271 let elapsed = now.duration_since(last_emit).as_millis();
2272 if elapsed >= RECIPE_PROGRESS_THROTTLE_MS as u128 {
2273 last_emit = now;
2274 last_emit_bytes = bytes_downloaded;
2275 cb(DownloadProgressEvent::FileProgress {
2276 filename: file.dest.to_string(),
2277 file_index,
2278 bytes_downloaded,
2279 bytes_total: size_bytes,
2280 batch_bytes_downloaded,
2281 batch_bytes_total,
2282 batch_elapsed_ms: started.elapsed().as_millis() as u64,
2283 });
2284 }
2285 }
2286 }
2287 if let Some(cb) = progress.as_deref() {
2292 if bytes_downloaded > last_emit_bytes {
2293 cb(DownloadProgressEvent::FileProgress {
2294 filename: file.dest.to_string(),
2295 file_index,
2296 bytes_downloaded,
2297 bytes_total: size_bytes,
2298 batch_bytes_downloaded,
2299 batch_bytes_total,
2300 batch_elapsed_ms: started.elapsed().as_millis() as u64,
2301 });
2302 }
2303 }
2304 drop(out);
2306
2307 if !opts.skip_verify {
2315 let actual = compute_sha256(dest_path).map_err(|e| {
2316 DownloadError::Other(format!(
2317 "failed to compute SHA-256 for {}: {e}",
2318 dest_path.display()
2319 ))
2320 })?;
2321 if let Some(expected) = file.sha256 {
2322 if !actual.eq_ignore_ascii_case(expected) {
2323 let _ = std::fs::remove_file(dest_path);
2324 return Err(DownloadError::Sha256Mismatch {
2325 filename: file.dest.to_string(),
2326 expected: expected.to_string(),
2327 actual,
2328 model: id.to_string(),
2329 });
2330 }
2331 }
2332 if let Err(e) = write_sha256_marker(dest_path, &actual) {
2333 eprintln!(
2334 "warning: failed to write .sha256-verified marker for {}: {e}",
2335 file.dest
2336 );
2337 }
2338 }
2339
2340 if let Some(cb) = progress.as_deref() {
2341 cb(DownloadProgressEvent::FileDone {
2342 filename: file.dest.to_string(),
2343 file_index,
2344 total_files,
2345 batch_bytes_downloaded,
2346 batch_bytes_total,
2347 batch_elapsed_ms: started.elapsed().as_millis() as u64,
2348 });
2349 }
2350 }
2351
2352 Ok(resolved.to_vec())
2353}
2354
2355#[cfg(test)]
2356mod tests {
2357 use super::*;
2358
2359 #[tokio::test]
2360 async fn identical_hf_files_serialize_acquisition_and_placement() {
2361 use std::sync::atomic::{AtomicUsize, Ordering};
2362
2363 let filename = format!("shared-{}.safetensors", uuid::Uuid::new_v4());
2364 let temp = tempfile::tempdir().unwrap();
2365 let source = temp.path().join("blob");
2366 let destination = temp.path().join("shared").join("encoder.safetensors");
2367 std::fs::write(&source, b"shared encoder").unwrap();
2368
2369 let active = Arc::new(AtomicUsize::new(0));
2370 let peak = Arc::new(AtomicUsize::new(0));
2371 let operation = || {
2372 let active = active.clone();
2373 let peak = peak.clone();
2374 let source = source.clone();
2375 let destination = destination.clone();
2376 async move {
2377 let now_active = active.fetch_add(1, Ordering::SeqCst) + 1;
2378 peak.fetch_max(now_active, Ordering::SeqCst);
2379 tokio::time::sleep(std::time::Duration::from_millis(25)).await;
2380 let result = hardlink_or_copy(&source, &destination);
2381 active.fetch_sub(1, Ordering::SeqCst);
2382 result
2383 }
2384 };
2385
2386 let first = with_hf_file_download_flight("Qwen/Qwen-Image-2512", &filename, operation());
2387 let second = with_hf_file_download_flight("Qwen/Qwen-Image-2512", &filename, operation());
2388 let (first_result, second_result) = tokio::join!(first, second);
2389
2390 first_result.unwrap();
2391 second_result.unwrap();
2392 assert_eq!(peak.load(Ordering::SeqCst), 1);
2393 assert_eq!(std::fs::read(destination).unwrap(), b"shared encoder");
2394 }
2395
2396 #[test]
2397 fn hidden_ltx2_adapters_use_files_only_pulls() {
2398 let adapters = crate::manifest::known_manifests()
2399 .iter()
2400 .filter(|manifest| {
2401 matches!(
2402 manifest.family.as_str(),
2403 "ltx2-control" | "ltx2-camera-control"
2404 )
2405 })
2406 .collect::<Vec<_>>();
2407 assert!(!adapters.is_empty());
2408 assert!(adapters
2409 .iter()
2410 .any(|manifest| manifest.name == "ltx2-camera-control-dolly-right-19b"));
2411 for manifest in adapters {
2412 assert!(manifest.is_auxiliary());
2413 assert!(manifest_uses_files_only_pull(manifest), "{}", manifest.name);
2414 }
2415
2416 assert!(!manifest_uses_files_only_pull(
2417 crate::manifest::find_manifest("ltx-2-19b-distilled:fp8").unwrap()
2418 ));
2419 assert!(!manifest_uses_files_only_pull(
2420 crate::manifest::find_manifest("controlnet-canny-sd15:fp16").unwrap()
2421 ));
2422 }
2423
2424 #[test]
2425 fn truncate_short_name_unchanged() {
2426 assert_eq!(truncate_filename("ae.safetensors", 45), "ae.safetensors");
2427 }
2428
2429 #[test]
2430 fn truncate_exact_fit_unchanged() {
2431 let name = "x".repeat(30);
2432 assert_eq!(truncate_filename(&name, 30), name);
2433 }
2434
2435 #[test]
2436 fn truncate_long_name_keeps_suffix() {
2437 let result = truncate_filename("unet/diffusion_pytorch_model.fp16.safetensors", 30);
2438 assert_eq!(result.len(), 30);
2439 assert!(result.starts_with("..."));
2440 assert!(result.ends_with(".fp16.safetensors"));
2441 }
2442
2443 #[test]
2444 fn truncate_very_small_max_returns_original() {
2445 let name = "something.safetensors";
2447 assert_eq!(truncate_filename(name, 5), name);
2448 }
2449
2450 #[test]
2451 fn sync_callback_progress_reports_real_accumulated_bytes() {
2452 let events = Arc::new(Mutex::new(Vec::new()));
2453 let events_for_callback = events.clone();
2454 let callback: DownloadProgressCallback = Arc::new(move |event| {
2455 events_for_callback.lock().unwrap().push(event);
2456 });
2457 let mut progress = SyncCallbackProgress::new(callback);
2458 hf_hub::api::Progress::init(&mut progress, 100, "encoder.gguf");
2459 hf_hub::api::Progress::update(&mut progress, 40);
2460 hf_hub::api::Progress::update(&mut progress, 60);
2461 hf_hub::api::Progress::finish(&mut progress);
2462
2463 let events = events.lock().unwrap();
2464 assert!(matches!(
2465 &events[0],
2466 DownloadProgressEvent::FileStart {
2467 filename,
2468 size_bytes: 100,
2469 ..
2470 } if filename == "encoder.gguf"
2471 ));
2472 assert!(events.iter().any(|event| matches!(
2473 event,
2474 DownloadProgressEvent::FileProgress {
2475 bytes_downloaded: 100,
2476 bytes_total: 100,
2477 ..
2478 }
2479 )));
2480 assert!(matches!(
2481 events.last(),
2482 Some(DownloadProgressEvent::FileDone {
2483 batch_bytes_downloaded: 100,
2484 batch_bytes_total: 100,
2485 ..
2486 })
2487 ));
2488 }
2489
2490 #[tokio::test]
2491 async fn callback_progress_clones_share_accumulated_bytes() {
2492 let events = Arc::new(Mutex::new(Vec::new()));
2493 let events_for_cb = events.clone();
2494 let callback: DownloadProgressCallback = Arc::new(move |event| {
2495 events_for_cb
2496 .lock()
2497 .expect("events mutex poisoned")
2498 .push(event);
2499 });
2500
2501 let mut progress = CallbackProgress::new(callback, 1, 3, 1_000, 10_000, Instant::now());
2502 progress.init(1_024, "weights.safetensors").await;
2503
2504 let mut chunk_a = progress.clone();
2505 let mut chunk_b = progress.clone();
2506 chunk_a.update(512).await;
2507 chunk_b.update(512).await;
2508 progress.finish().await;
2509
2510 let events = events.lock().expect("events mutex poisoned");
2511 assert!(events.iter().any(|event| matches!(
2512 event,
2513 DownloadProgressEvent::FileProgress {
2514 bytes_downloaded: 1_024,
2515 bytes_total: 1_024,
2516 batch_bytes_downloaded: 2_024,
2517 ..
2518 }
2519 )));
2520 }
2521
2522 #[test]
2523 fn download_error_gated_message() {
2524 let err = DownloadError::GatedModel {
2525 repo: "black-forest-labs/FLUX.1-dev".to_string(),
2526 model: "flux-dev:q8".to_string(),
2527 };
2528 let msg = err.to_string();
2529 assert!(msg.contains("huggingface.co/black-forest-labs/FLUX.1-dev"));
2530 assert!(msg.contains("HF_TOKEN"));
2531 assert!(msg.contains("mold pull flux-dev:q8"));
2532 }
2533
2534 #[test]
2535 fn download_error_unauthorized_message() {
2536 let err = DownloadError::Unauthorized {
2537 repo: "black-forest-labs/FLUX.1-schnell".to_string(),
2538 model: "flux-schnell:q8".to_string(),
2539 };
2540 let msg = err.to_string();
2541 assert!(msg.contains("Authentication required"));
2542 assert!(msg.contains("black-forest-labs/FLUX.1-schnell"));
2543 assert!(msg.contains("HF_TOKEN"));
2544 assert!(msg.contains("huggingface-cli login"));
2545 assert!(msg.contains("mold pull flux-schnell:q8"));
2546 }
2547
2548 fn compliance_gated_manifest(name: &str, family: &str, repo: &str) -> ModelManifest {
2549 use crate::manifest::{ManifestDefaults, ModelFile};
2550
2551 ModelManifest {
2552 name: name.to_string(),
2553 family: family.to_string(),
2554 description: "policy fixture".to_string(),
2555 files: vec![ModelFile {
2556 hf_repo: repo.to_string(),
2557 hf_filename: "weights.safetensors".to_string(),
2558 component: ModelComponent::Transformer,
2559 size_bytes: 1,
2560 gated: false,
2561 sha256: None,
2562 }],
2563 defaults: ManifestDefaults {
2564 steps: 1,
2565 guidance: 1.0,
2566 width: 32,
2567 height: 32,
2568 is_schnell: false,
2569 scheduler: None,
2570 negative_prompt: None,
2571 frames: None,
2572 fps: None,
2573 source_image: None,
2574 },
2575 hidden: true,
2576 }
2577 }
2578
2579 #[test]
2580 fn reviewed_h3_manifest_is_accepted_for_upstream_acquisition() {
2581 let manifest = crate::manifest::find_manifest(crate::minimax_h3::FL2VA_COMFY).unwrap();
2582 require_manifest_acquisition(manifest).unwrap();
2583 crate::require_model_activation(&manifest.name, Some(&manifest.family)).unwrap();
2584 }
2585
2586 #[test]
2587 fn h3_repo_identity_cannot_bypass_the_pinned_manifest() {
2588 let manifest = compliance_gated_manifest("renamed-model", "custom", "Comfy-Org/MiniMax-H3");
2589 assert!(require_manifest_acquisition(&manifest).is_err());
2590 }
2591
2592 #[test]
2593 fn arbitrary_h3_single_files_remain_outside_pinned_manifest_acquisition() {
2594 for (repo, filename, target) in [
2595 ("MiniMaxAI/MiniMax-H3", "weights.safetensors", None),
2596 (
2597 "example/renamed-model",
2598 "MiniMax-H3/weights.safetensors",
2599 None,
2600 ),
2601 (
2602 "example/renamed-model",
2603 "weights.safetensors",
2604 Some("shared/MiniMax-H3"),
2605 ),
2606 ] {
2607 assert!(require_single_file_acquisition(repo, filename, target).is_err());
2608 }
2609 }
2610
2611 #[test]
2612 fn sync_download_policy_keeps_h3_lookalikes_available() {
2613 for (repo, filename, target) in [
2614 ("example/minimax-h30", "weights.safetensors", None),
2615 ("example/model", "minimaxh30.safetensors", None),
2616 ("example/model", "weights.safetensors", Some("shared/h3")),
2617 ] {
2618 require_single_file_acquisition(repo, filename, target)
2619 .unwrap_or_else(|_| panic!("lookalike must remain available: {repo}/{filename}"));
2620 }
2621 }
2622
2623 static HF_TOKEN_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
2626
2627 fn write_fake_gguf(path: &std::path::Path, tensor_names: &[&str]) {
2637 if let Some(parent) = path.parent() {
2638 std::fs::create_dir_all(parent).unwrap();
2639 }
2640 let mut buf = Vec::with_capacity(4096);
2641 buf.extend_from_slice(b"GGUF");
2642 buf.extend(std::iter::repeat_n(0u8, 256));
2645 for name in tensor_names {
2646 buf.extend_from_slice(name.as_bytes());
2647 buf.push(0);
2648 }
2649 std::fs::write(path, &buf).unwrap();
2650 }
2651
2652 fn tmp_dir(tag: &str) -> std::path::PathBuf {
2653 let dir = std::env::temp_dir().join(format!(
2654 "mold-dl-{tag}-{}-{}",
2655 std::process::id(),
2656 std::time::SystemTime::now()
2657 .duration_since(std::time::UNIX_EPOCH)
2658 .unwrap()
2659 .as_nanos()
2660 ));
2661 std::fs::create_dir_all(&dir).unwrap();
2662 dir
2663 }
2664
2665 fn fake_flux_gguf_manifest(name: &str, filename: &str, is_schnell: bool) -> ModelManifest {
2666 use crate::manifest::{ManifestDefaults, ModelFile};
2667 ModelManifest {
2668 name: name.to_string(),
2669 family: "flux".to_string(),
2670 description: "test".to_string(),
2671 files: vec![ModelFile {
2672 hf_repo: "test/repo".to_string(),
2673 hf_filename: filename.to_string(),
2674 component: ModelComponent::Transformer,
2675 size_bytes: 0,
2676 gated: false,
2677 sha256: None,
2678 }],
2679 defaults: ManifestDefaults {
2680 steps: 20,
2681 guidance: 3.5,
2682 width: 1024,
2683 height: 1024,
2684 is_schnell,
2685 scheduler: None,
2686 negative_prompt: None,
2687 frames: None,
2688 fps: None,
2689 source_image: None,
2690 },
2691 hidden: false,
2692 }
2693 }
2694
2695 #[test]
2696 fn gguf_header_contains_tensor_false_for_missing_file() {
2697 let path = std::env::temp_dir().join(format!(
2698 "mold-dl-nofile-{}-{}.gguf",
2699 std::process::id(),
2700 std::time::SystemTime::now()
2701 .duration_since(std::time::UNIX_EPOCH)
2702 .unwrap()
2703 .as_nanos()
2704 ));
2705 assert!(!gguf_header_contains_tensor(&path, "img_in.weight"));
2706 }
2707
2708 #[test]
2709 fn gguf_header_contains_tensor_false_for_non_gguf_magic() {
2710 let dir = tmp_dir("nonmagic");
2711 let path = dir.join("not-a-gguf.gguf");
2712 std::fs::write(&path, b"SAFE\0\0\0\0img_in.weight\0").unwrap();
2713 assert!(!gguf_header_contains_tensor(&path, "img_in.weight"));
2714 std::fs::remove_dir_all(&dir).ok();
2715 }
2716
2717 #[test]
2718 fn gguf_header_contains_tensor_finds_needle_after_magic() {
2719 let dir = tmp_dir("finds");
2720 let path = dir.join("has.gguf");
2721 write_fake_gguf(&path, &["img_in.weight", "time_in.in_layer.weight"]);
2722 assert!(gguf_header_contains_tensor(&path, "img_in.weight"));
2723 assert!(gguf_header_contains_tensor(
2724 &path,
2725 "time_in.in_layer.weight"
2726 ));
2727 assert!(!gguf_header_contains_tensor(
2728 &path,
2729 "guidance_in.in_layer.weight"
2730 ));
2731 std::fs::remove_dir_all(&dir).ok();
2732 }
2733
2734 #[test]
2735 fn flux_reference_warning_noop_for_non_flux_family() {
2736 use crate::manifest::{ManifestDefaults, ModelFile};
2737 let dir = tmp_dir("non-flux");
2738 let manifest = ModelManifest {
2739 name: "sd15:fp16".to_string(),
2740 family: "sd15".to_string(),
2741 description: "test".to_string(),
2742 files: vec![ModelFile {
2743 hf_repo: "test/repo".to_string(),
2744 hf_filename: "model.gguf".to_string(),
2745 component: ModelComponent::Transformer,
2746 size_bytes: 0,
2747 gated: false,
2748 sha256: None,
2749 }],
2750 defaults: ManifestDefaults {
2751 steps: 25,
2752 guidance: 7.5,
2753 width: 512,
2754 height: 512,
2755 is_schnell: false,
2756 scheduler: None,
2757 negative_prompt: None,
2758 frames: None,
2759 fps: None,
2760 source_image: None,
2761 },
2762 hidden: false,
2763 };
2764 assert!(flux_reference_warning(&manifest, &dir).is_none());
2765 std::fs::remove_dir_all(&dir).ok();
2766 }
2767
2768 #[test]
2769 fn flux_reference_warning_noop_for_safetensors_transformer() {
2770 let dir = tmp_dir("safetensors");
2771 let manifest = fake_flux_gguf_manifest("ultra-test:bf16", "model.safetensors", false);
2773 assert!(flux_reference_warning(&manifest, &dir).is_none());
2774 std::fs::remove_dir_all(&dir).ok();
2775 }
2776
2777 #[test]
2778 fn flux_reference_warning_noop_when_file_absent() {
2779 let dir = tmp_dir("absent");
2780 let manifest = fake_flux_gguf_manifest("ultra-absent:q8", "ultra-absent-q8.gguf", false);
2781 assert!(flux_reference_warning(&manifest, &dir).is_none());
2783 std::fs::remove_dir_all(&dir).ok();
2784 }
2785
2786 #[test]
2787 fn flux_reference_warning_noop_when_transformer_is_complete() {
2788 let dir = tmp_dir("complete");
2789 let manifest =
2790 fake_flux_gguf_manifest("ultra-complete:q8", "ultra-complete-q8.gguf", false);
2791 let xformer = dir.join(crate::manifest::storage_path(&manifest, &manifest.files[0]));
2793 write_fake_gguf(&xformer, &["img_in.weight", "guidance_in.in_layer.weight"]);
2794 assert!(flux_reference_warning(&manifest, &dir).is_none());
2795 std::fs::remove_dir_all(&dir).ok();
2796 }
2797
2798 #[test]
2799 fn flux_reference_warning_fires_for_city96_dev_without_reference() {
2800 let dir = tmp_dir("city96-dev");
2801 let manifest = fake_flux_gguf_manifest("ultra-v4:q8", "ultra-v4-q8.gguf", false);
2802 let xformer = dir.join(crate::manifest::storage_path(&manifest, &manifest.files[0]));
2803 write_fake_gguf(&xformer, &["double_blocks.0.img_mod.lin.weight"]);
2805
2806 let msg = flux_reference_warning(&manifest, &dir)
2807 .expect("city96-format dev GGUF without reference must emit warning");
2808 assert!(msg.contains("ultra-v4-q8.gguf"));
2809 assert!(msg.contains("ultra-v4:q8"));
2810 assert!(msg.contains("mold pull flux-dev:q8"));
2811 assert!(
2812 msg.contains("guidance_in"),
2813 "dev target message must mention guidance_in: {msg}"
2814 );
2815 std::fs::remove_dir_all(&dir).ok();
2816 }
2817
2818 #[test]
2819 fn flux_reference_warning_fires_for_city96_schnell_without_reference() {
2820 let dir = tmp_dir("city96-schnell");
2821 let manifest = fake_flux_gguf_manifest("ultra-schnell:q8", "ultra-schnell-q8.gguf", true);
2822 let xformer = dir.join(crate::manifest::storage_path(&manifest, &manifest.files[0]));
2823 write_fake_gguf(&xformer, &["double_blocks.0.img_mod.lin.weight"]);
2824
2825 let msg = flux_reference_warning(&manifest, &dir)
2826 .expect("city96-format schnell GGUF without reference must emit warning");
2827 assert!(msg.contains("ultra-schnell-q8.gguf"));
2829 assert!(msg.contains("mold pull flux-dev:q8"));
2830 assert!(!msg.contains("guidance_in"));
2831 std::fs::remove_dir_all(&dir).ok();
2832 }
2833
2834 #[test]
2835 fn flux_reference_warning_silenced_when_dev_reference_exists() {
2836 let dir = tmp_dir("has-dev-ref");
2837 let manifest = fake_flux_gguf_manifest("ultra-v4:q8", "ultra-v4-q8.gguf", false);
2838 let xformer = dir.join(crate::manifest::storage_path(&manifest, &manifest.files[0]));
2839 write_fake_gguf(&xformer, &["double_blocks.0.img_mod.lin.weight"]);
2840
2841 let dev_manifest = crate::manifest::find_manifest("flux-dev:q8")
2843 .expect("flux-dev:q8 must exist in the static manifest catalog");
2844 let dev_xformer_file = dev_manifest
2845 .files
2846 .iter()
2847 .find(|f| f.component == ModelComponent::Transformer)
2848 .expect("flux-dev:q8 must declare a Transformer file");
2849 let dev_path = dir.join(crate::manifest::storage_path(
2850 dev_manifest,
2851 dev_xformer_file,
2852 ));
2853 write_fake_gguf(&dev_path, &["img_in.weight", "guidance_in.in_layer.weight"]);
2854
2855 assert!(
2856 flux_reference_warning(&manifest, &dir).is_none(),
2857 "warning must be silenced when a complete flux-dev reference is downloaded"
2858 );
2859 std::fs::remove_dir_all(&dir).ok();
2860 }
2861
2862 #[test]
2863 fn flux_reference_warning_rejects_schnell_as_reference_for_dev_target() {
2864 let dir = tmp_dir("schnell-only-for-dev");
2867 let manifest = fake_flux_gguf_manifest("ultra-v4:q8", "ultra-v4-q8.gguf", false);
2868 let xformer = dir.join(crate::manifest::storage_path(&manifest, &manifest.files[0]));
2869 write_fake_gguf(&xformer, &["double_blocks.0.img_mod.lin.weight"]);
2870
2871 let schnell_manifest = crate::manifest::find_manifest("flux-schnell:q8")
2873 .expect("flux-schnell:q8 must exist in the static manifest catalog");
2874 let schnell_xformer_file = schnell_manifest
2875 .files
2876 .iter()
2877 .find(|f| f.component == ModelComponent::Transformer)
2878 .expect("flux-schnell:q8 must declare a Transformer file");
2879 let schnell_path = dir.join(crate::manifest::storage_path(
2880 schnell_manifest,
2881 schnell_xformer_file,
2882 ));
2883 write_fake_gguf(&schnell_path, &["img_in.weight"]);
2884
2885 let msg = flux_reference_warning(&manifest, &dir)
2886 .expect("dev target must not accept schnell as reference; warning should fire");
2887 assert!(msg.contains("mold pull flux-dev:q8"));
2888 std::fs::remove_dir_all(&dir).ok();
2889 }
2890
2891 #[test]
2892 fn resolve_hf_token_reads_env_var() {
2893 let _guard = HF_TOKEN_LOCK.lock().unwrap();
2894 let original = std::env::var("HF_TOKEN").ok();
2895 std::env::set_var("HF_TOKEN", "hf_test_token_123");
2896 let token = resolve_hf_token();
2897 match &original {
2899 Some(v) => std::env::set_var("HF_TOKEN", v),
2900 None => std::env::remove_var("HF_TOKEN"),
2901 }
2902 assert_eq!(token, Some("hf_test_token_123".to_string()));
2903 }
2904
2905 #[test]
2906 fn resolve_hf_token_ignores_empty_env() {
2907 let _guard = HF_TOKEN_LOCK.lock().unwrap();
2908 let original = std::env::var("HF_TOKEN").ok();
2909 std::env::set_var("HF_TOKEN", " ");
2910 let token = resolve_hf_token();
2911 match &original {
2913 Some(v) => std::env::set_var("HF_TOKEN", v),
2914 None => std::env::remove_var("HF_TOKEN"),
2915 }
2916 assert_ne!(token, Some(" ".to_string()));
2918 }
2919
2920 #[test]
2921 fn compute_sha256_correct_digest() {
2922 let dir = std::env::temp_dir().join("mold_test_sha256_compute");
2923 let _ = std::fs::create_dir_all(&dir);
2924 let path = dir.join("test_file.bin");
2925 std::fs::write(&path, b"hello world").unwrap();
2926 let digest = compute_sha256(&path).unwrap();
2927 assert_eq!(
2928 digest,
2929 "b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9"
2930 );
2931 let _ = std::fs::remove_dir_all(&dir);
2932 }
2933
2934 #[test]
2935 fn verify_sha256_matches() {
2936 let dir = std::env::temp_dir().join("mold_test_sha256_match");
2937 let _ = std::fs::create_dir_all(&dir);
2938 let path = dir.join("test_file.bin");
2939 std::fs::write(&path, b"hello world").unwrap();
2940 let expected = "b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9";
2942 assert!(verify_sha256(&path, expected).unwrap());
2943 let _ = std::fs::remove_dir_all(&dir);
2944 }
2945
2946 #[test]
2947 fn verify_sha256_mismatch() {
2948 let dir = std::env::temp_dir().join("mold_test_sha256_mismatch");
2949 let _ = std::fs::create_dir_all(&dir);
2950 let path = dir.join("test_file.bin");
2951 std::fs::write(&path, b"hello world").unwrap();
2952 let wrong = "0000000000000000000000000000000000000000000000000000000000000000";
2953 assert!(!verify_sha256(&path, wrong).unwrap());
2954 let _ = std::fs::remove_dir_all(&dir);
2955 }
2956
2957 #[test]
2964 fn verify_sha256_is_hex_case_insensitive() {
2965 let dir = std::env::temp_dir().join("mold_test_sha256_case");
2966 let _ = std::fs::create_dir_all(&dir);
2967 let path = dir.join("test_file.bin");
2968 std::fs::write(&path, b"hello world").unwrap();
2969 let lower = "b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9";
2970 let upper = "B94D27B9934D3E08A52E52D7DA7DABFAC484EFE37A5380EE9088F7ACE2EFCDE9";
2971 let mixed = "B94d27b9934D3e08a52E52d7Da7dabfac484EFE37A5380ee9088f7Ace2efcDE9";
2972 assert!(
2973 verify_sha256(&path, lower).unwrap(),
2974 "lowercase digest must match"
2975 );
2976 assert!(
2977 verify_sha256(&path, upper).unwrap(),
2978 "uppercase digest must match (Civitai-style)",
2979 );
2980 assert!(
2981 verify_sha256(&path, mixed).unwrap(),
2982 "mixed-case must match"
2983 );
2984 let _ = std::fs::remove_dir_all(&dir);
2985 }
2986
2987 #[test]
2988 fn verify_file_integrity_deletes_on_mismatch() {
2989 use crate::manifest::{ModelComponent, ModelFile};
2990 let dir = std::env::temp_dir().join("mold_test_integrity_mismatch");
2991 let _ = std::fs::create_dir_all(&dir);
2992 let path = dir.join("corrupted.bin");
2993 std::fs::write(&path, b"corrupted data").unwrap();
2994
2995 let file = ModelFile {
2996 hf_repo: "test/repo".to_string(),
2997 hf_filename: "corrupted.bin".to_string(),
2998 component: ModelComponent::Transformer,
2999 size_bytes: 14,
3000 gated: false,
3001 sha256: Some("0000000000000000000000000000000000000000000000000000000000000000"),
3002 };
3003
3004 let result = verify_file_integrity(&path, &file, "test-model:q8", false);
3005 assert!(result.is_err());
3006 assert!(matches!(
3007 result.unwrap_err(),
3008 DownloadError::Sha256Mismatch { .. }
3009 ),);
3010 assert!(!path.exists());
3012 let _ = std::fs::remove_dir_all(&dir);
3013 }
3014
3015 #[test]
3016 fn verify_file_integrity_skip_verify_ignores_mismatch() {
3017 use crate::manifest::{ModelComponent, ModelFile};
3018 let dir = std::env::temp_dir().join("mold_test_integrity_skip");
3019 let _ = std::fs::create_dir_all(&dir);
3020 let path = dir.join("file.bin");
3021 std::fs::write(&path, b"some data").unwrap();
3022
3023 let file = ModelFile {
3024 hf_repo: "test/repo".to_string(),
3025 hf_filename: "file.bin".to_string(),
3026 component: ModelComponent::Transformer,
3027 size_bytes: 9,
3028 gated: false,
3029 sha256: Some("0000000000000000000000000000000000000000000000000000000000000000"),
3030 };
3031
3032 let result = verify_file_integrity(&path, &file, "test-model:q8", true);
3033 assert!(result.is_ok());
3034 assert!(path.exists());
3036 let _ = std::fs::remove_dir_all(&dir);
3037 }
3038
3039 #[test]
3040 fn verify_file_integrity_no_hash_is_ok() {
3041 use crate::manifest::{ModelComponent, ModelFile};
3042 let dir = std::env::temp_dir().join("mold_test_integrity_nohash");
3043 let _ = std::fs::create_dir_all(&dir);
3044 let path = dir.join("file.bin");
3045 std::fs::write(&path, b"data").unwrap();
3046
3047 let file = ModelFile {
3048 hf_repo: "test/repo".to_string(),
3049 hf_filename: "file.bin".to_string(),
3050 component: ModelComponent::Transformer,
3051 size_bytes: 4,
3052 gated: false,
3053 sha256: None,
3054 };
3055
3056 assert!(verify_file_integrity(&path, &file, "test:q8", false).is_ok());
3057 let _ = std::fs::remove_dir_all(&dir);
3058 }
3059
3060 #[test]
3063 fn sha256_marker_path_appends_suffix() {
3064 let p = std::path::Path::new("/tmp/foo/model.safetensors");
3065 let marker = sha256_marker_path(p);
3066 assert_eq!(
3067 marker,
3068 std::path::PathBuf::from("/tmp/foo/model.safetensors.sha256-verified")
3069 );
3070 }
3071
3072 #[test]
3073 fn sha256_marker_path_handles_dotted_filenames() {
3074 let p = std::path::Path::new("/tmp/.hidden.bin");
3075 let marker = sha256_marker_path(p);
3076 assert_eq!(
3077 marker,
3078 std::path::PathBuf::from("/tmp/.hidden.bin.sha256-verified")
3079 );
3080 }
3081
3082 #[test]
3083 fn write_sha256_marker_creates_file_with_digest() {
3084 let dir = std::env::temp_dir().join("mold_test_marker_write");
3085 let _ = std::fs::create_dir_all(&dir);
3086 let path = dir.join("file.bin");
3087 std::fs::write(&path, b"hello world").unwrap();
3088 let digest = "b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9";
3089 write_sha256_marker(&path, digest).unwrap();
3090
3091 let marker = sha256_marker_path(&path);
3092 assert!(marker.exists(), "marker should exist next to file");
3093 let content = std::fs::read_to_string(&marker).unwrap();
3094 assert!(
3095 content.contains(digest),
3096 "marker content should contain the digest, got: {content:?}"
3097 );
3098 let _ = std::fs::remove_dir_all(&dir);
3099 }
3100
3101 #[test]
3102 fn write_sha256_marker_is_idempotent() {
3103 let dir = std::env::temp_dir().join("mold_test_marker_idempotent");
3104 let _ = std::fs::create_dir_all(&dir);
3105 let path = dir.join("file.bin");
3106 std::fs::write(&path, b"x").unwrap();
3107 let digest = "2d711642b726b04401627ca9fbac32f5c8530fb1903cc4db02258717921a4881";
3108 write_sha256_marker(&path, digest).unwrap();
3109 write_sha256_marker(&path, digest).unwrap();
3111 let _ = std::fs::remove_dir_all(&dir);
3112 }
3113
3114 #[test]
3115 fn has_sha256_marker_reflects_existence() {
3116 let dir = std::env::temp_dir().join("mold_test_marker_has");
3117 let _ = std::fs::create_dir_all(&dir);
3118 let path = dir.join("file.bin");
3119 std::fs::write(&path, b"x").unwrap();
3120 assert!(!has_sha256_marker(&path), "no marker yet");
3121 write_sha256_marker(&path, "deadbeef").unwrap();
3122 assert!(has_sha256_marker(&path), "marker should exist");
3123 let _ = std::fs::remove_dir_all(&dir);
3124 }
3125
3126 #[test]
3129 fn verify_file_integrity_writes_marker_on_match() {
3130 use crate::manifest::{ModelComponent, ModelFile};
3131 let dir = std::env::temp_dir().join("mold_test_integrity_writes_marker");
3132 let _ = std::fs::create_dir_all(&dir);
3133 let path = dir.join("ok.bin");
3134 std::fs::write(&path, b"hello world").unwrap();
3135 let file = ModelFile {
3136 hf_repo: "test/repo".to_string(),
3137 hf_filename: "ok.bin".to_string(),
3138 component: ModelComponent::Transformer,
3139 size_bytes: 11,
3140 gated: false,
3141 sha256: Some("b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9"),
3142 };
3143 verify_file_integrity(&path, &file, "test:q8", false).unwrap();
3144 assert!(
3145 has_sha256_marker(&path),
3146 "marker should be written after a successful verify"
3147 );
3148 let _ = std::fs::remove_dir_all(&dir);
3149 }
3150
3151 #[test]
3152 fn verify_file_integrity_writes_marker_when_no_hash_declared() {
3153 use crate::manifest::{ModelComponent, ModelFile};
3154 let dir = std::env::temp_dir().join("mold_test_integrity_no_hash_marker");
3155 let _ = std::fs::create_dir_all(&dir);
3156 let path = dir.join("ok.bin");
3157 std::fs::write(&path, b"data").unwrap();
3158 let file = ModelFile {
3159 hf_repo: "test/repo".to_string(),
3160 hf_filename: "ok.bin".to_string(),
3161 component: ModelComponent::Transformer,
3162 size_bytes: 4,
3163 gated: false,
3164 sha256: None,
3165 };
3166 verify_file_integrity(&path, &file, "test:q8", false).unwrap();
3167 assert!(
3168 has_sha256_marker(&path),
3169 "marker must be written even when manifest declares no expected hash \
3170 (the marker still proves the file finished writing)"
3171 );
3172 let _ = std::fs::remove_dir_all(&dir);
3173 }
3174
3175 #[test]
3176 fn verify_file_integrity_no_marker_on_mismatch() {
3177 use crate::manifest::{ModelComponent, ModelFile};
3178 let dir = std::env::temp_dir().join("mold_test_integrity_no_marker_on_miss");
3179 let _ = std::fs::create_dir_all(&dir);
3180 let path = dir.join("bad.bin");
3181 std::fs::write(&path, b"corrupted").unwrap();
3182 let file = ModelFile {
3183 hf_repo: "test/repo".to_string(),
3184 hf_filename: "bad.bin".to_string(),
3185 component: ModelComponent::Transformer,
3186 size_bytes: 9,
3187 gated: false,
3188 sha256: Some("0000000000000000000000000000000000000000000000000000000000000000"),
3189 };
3190 let result = verify_file_integrity(&path, &file, "test:q8", false);
3191 assert!(result.is_err(), "mismatch should error");
3192 assert!(
3195 !has_sha256_marker(&path),
3196 "no marker may exist after a hash mismatch"
3197 );
3198 let _ = std::fs::remove_dir_all(&dir);
3199 }
3200
3201 #[test]
3202 fn verify_file_integrity_skip_verify_does_not_write_marker() {
3203 use crate::manifest::{ModelComponent, ModelFile};
3204 let dir = std::env::temp_dir().join("mold_test_integrity_skip_no_marker");
3205 let _ = std::fs::create_dir_all(&dir);
3206 let path = dir.join("file.bin");
3207 std::fs::write(&path, b"some data").unwrap();
3208 let file = ModelFile {
3209 hf_repo: "test/repo".to_string(),
3210 hf_filename: "file.bin".to_string(),
3211 component: ModelComponent::Transformer,
3212 size_bytes: 9,
3213 gated: false,
3214 sha256: Some("0000000000000000000000000000000000000000000000000000000000000000"),
3215 };
3216 verify_file_integrity(&path, &file, "test:q8", true).unwrap();
3218 assert!(
3219 !has_sha256_marker(&path),
3220 "skip_verify must not produce a marker — we have no integrity guarantee"
3221 );
3222 let _ = std::fs::remove_dir_all(&dir);
3223 }
3224
3225 #[test]
3226 fn pulling_marker_roundtrip() {
3227 let dir = std::env::temp_dir().join("mold_test_marker_roundtrip");
3228 let _ = std::fs::create_dir_all(&dir);
3229 let marker = dir.join(".pulling");
3230
3231 std::fs::write(&marker, "test-model:q8").unwrap();
3233 assert!(marker.exists());
3234
3235 let _ = std::fs::remove_file(&marker);
3237 assert!(!marker.exists());
3238
3239 let _ = std::fs::remove_dir_all(&dir);
3240 }
3241
3242 #[test]
3243 fn sha256_mismatch_error_message() {
3244 let err = DownloadError::Sha256Mismatch {
3245 filename: "transformer.gguf".to_string(),
3246 expected: "aaa".to_string(),
3247 actual: "bbb".to_string(),
3248 model: "flux-dev:q8".to_string(),
3249 };
3250 let msg = err.to_string();
3251 assert!(msg.contains("SHA-256 mismatch"));
3252 assert!(msg.contains("transformer.gguf"));
3253 assert!(msg.contains("mold pull flux-dev:q8"));
3254 assert!(msg.contains("--skip-verify"));
3255 }
3256
3257 static CIVITAI_TOKEN_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
3260
3261 #[test]
3262 fn resolve_civitai_token_reads_env_var() {
3263 let _guard = CIVITAI_TOKEN_LOCK.lock().unwrap();
3264 let original = std::env::var("CIVITAI_TOKEN").ok();
3265 std::env::set_var("CIVITAI_TOKEN", "cv_test_token_abc");
3266 let token = resolve_civitai_token();
3267 match &original {
3268 Some(v) => std::env::set_var("CIVITAI_TOKEN", v),
3269 None => std::env::remove_var("CIVITAI_TOKEN"),
3270 }
3271 assert_eq!(token, Some("cv_test_token_abc".to_string()));
3272 }
3273
3274 #[test]
3275 fn resolve_civitai_token_ignores_empty() {
3276 let _guard = CIVITAI_TOKEN_LOCK.lock().unwrap();
3277 let original = std::env::var("CIVITAI_TOKEN").ok();
3278 std::env::set_var("CIVITAI_TOKEN", " ");
3279 let token = resolve_civitai_token();
3280 match &original {
3281 Some(v) => std::env::set_var("CIVITAI_TOKEN", v),
3282 None => std::env::remove_var("CIVITAI_TOKEN"),
3283 }
3284 assert_eq!(token, None);
3285 }
3286
3287 #[test]
3288 fn civitai_auth_or_error_returns_bearer_when_set() {
3289 let _guard = CIVITAI_TOKEN_LOCK.lock().unwrap();
3290 let original = std::env::var("CIVITAI_TOKEN").ok();
3291 std::env::set_var("CIVITAI_TOKEN", "cv_secret_xyz");
3292 let auth = civitai_auth_or_error("cv:123");
3293 match &original {
3294 Some(v) => std::env::set_var("CIVITAI_TOKEN", v),
3295 None => std::env::remove_var("CIVITAI_TOKEN"),
3296 }
3297 match auth {
3298 Ok(RecipeAuth::Bearer(t)) => assert_eq!(t, "cv_secret_xyz"),
3299 other => panic!("expected Bearer, got {other:?}"),
3300 }
3301 }
3302
3303 #[test]
3304 fn civitai_auth_or_error_returns_missing_token_error_when_unset() {
3305 let _guard = CIVITAI_TOKEN_LOCK.lock().unwrap();
3306 let original = std::env::var("CIVITAI_TOKEN").ok();
3307 std::env::remove_var("CIVITAI_TOKEN");
3308 let err = civitai_auth_or_error("cv:618692").unwrap_err();
3309 if let Some(v) = &original {
3310 std::env::set_var("CIVITAI_TOKEN", v);
3311 }
3312 match err {
3313 DownloadError::MissingCivitaiToken { id } => {
3314 assert_eq!(id, "cv:618692");
3315 }
3316 other => panic!("expected MissingCivitaiToken, got {other:?}"),
3317 }
3318 }
3319
3320 #[test]
3321 fn missing_civitai_token_error_message_points_at_env_var() {
3322 let err = DownloadError::MissingCivitaiToken {
3323 id: "cv:618692".to_string(),
3324 };
3325 let msg = err.to_string();
3326 assert!(
3327 msg.contains("CIVITAI_TOKEN"),
3328 "msg should name the env var: {msg}"
3329 );
3330 assert!(
3331 msg.contains("mold pull cv:618692"),
3332 "msg should suggest the retry command verbatim: {msg}"
3333 );
3334 assert!(msg.contains("https://civitai.com"));
3335 }
3336
3337 fn stage_complete_companion(models_dir: &std::path::Path, name: &str) {
3340 let manifest = crate::manifest::find_manifest(name)
3341 .unwrap_or_else(|| panic!("companion manifest {name} must exist"));
3342 for f in &manifest.files {
3343 let dest = models_dir.join(crate::manifest::storage_path(manifest, f));
3344 if let Some(parent) = dest.parent() {
3345 std::fs::create_dir_all(parent).unwrap();
3346 }
3347 std::fs::File::create(&dest)
3348 .unwrap()
3349 .set_len(f.size_bytes)
3350 .unwrap();
3351 if f.sha256.is_some() {
3352 std::fs::write(sha256_marker_path(&dest), "verified").unwrap();
3353 }
3354 }
3355 }
3356
3357 #[test]
3358 fn companion_present_returns_false_when_files_missing() {
3359 let models_dir = recipe_tmp_dir("companion_missing");
3360 let manifest =
3361 crate::manifest::find_manifest("clip-l").expect("clip-l manifest must exist");
3362 assert!(!companion_present_on_disk(&models_dir, manifest));
3363 let _ = std::fs::remove_dir_all(&models_dir);
3364 }
3365
3366 #[test]
3367 fn companion_present_returns_true_when_files_present() {
3368 let models_dir = recipe_tmp_dir("companion_present");
3369 stage_complete_companion(&models_dir, "clip-l");
3370 let manifest = crate::manifest::find_manifest("clip-l").unwrap();
3371 assert!(companion_present_on_disk(&models_dir, manifest));
3372 let _ = std::fs::remove_dir_all(&models_dir);
3373 }
3374
3375 #[test]
3376 fn companion_present_returns_false_for_unverified_sha_file() {
3377 let models_dir = recipe_tmp_dir("companion_unverified_sha");
3378 let manifest = crate::manifest::find_manifest("sdxl-vae").unwrap();
3379 let file = &manifest.files[0];
3380 let dest = models_dir.join(crate::manifest::storage_path(manifest, file));
3381 std::fs::create_dir_all(dest.parent().unwrap()).unwrap();
3382 std::fs::File::create(&dest)
3383 .unwrap()
3384 .set_len(file.size_bytes)
3385 .unwrap();
3386 assert!(
3387 !companion_present_on_disk(&models_dir, manifest),
3388 "SHA-declared companion files need the verification marker before repair skips them"
3389 );
3390 std::fs::write(sha256_marker_path(&dest), "verified").unwrap();
3391 assert!(companion_present_on_disk(&models_dir, manifest));
3392 let _ = std::fs::remove_dir_all(&models_dir);
3393 }
3394
3395 #[test]
3396 fn companion_present_returns_false_when_pulling_marker_present() {
3397 let models_dir = recipe_tmp_dir("companion_marker");
3398 stage_complete_companion(&models_dir, "clip-l");
3399 let marker = pulling_marker_path_in(&models_dir, "clip-l");
3400 if let Some(parent) = marker.parent() {
3401 std::fs::create_dir_all(parent).unwrap();
3402 }
3403 std::fs::write(&marker, "in-progress").unwrap();
3404 let manifest = crate::manifest::find_manifest("clip-l").unwrap();
3405 assert!(
3406 !companion_present_on_disk(&models_dir, manifest),
3407 "marker must override on-disk completeness"
3408 );
3409 let _ = std::fs::remove_dir_all(&models_dir);
3410 }
3411
3412 #[test]
3413 fn missing_companions_skips_unknown_names() {
3414 let models_dir = recipe_tmp_dir("companion_unknown");
3415 let json = r#"["clip-l","future-encoder-9000"]"#;
3417 let missing = missing_companions_from_json(Some(json), &models_dir);
3418 assert_eq!(missing.len(), 1);
3419 assert_eq!(missing[0].name, "clip-l");
3420 let _ = std::fs::remove_dir_all(&models_dir);
3421 }
3422
3423 #[test]
3424 fn missing_companions_resolves_zimage_text_encoder() {
3425 let models_dir = recipe_tmp_dir("companion_zimage_te");
3426 let json = r#"["z-image-te"]"#;
3427 let missing = missing_companions_from_json(Some(json), &models_dir);
3428 assert_eq!(missing.len(), 1);
3429 assert_eq!(missing[0].name, "z-image-te");
3430 let _ = std::fs::remove_dir_all(&models_dir);
3431 }
3432
3433 #[test]
3434 fn missing_companions_skips_present_returns_only_missing() {
3435 let models_dir = recipe_tmp_dir("companion_skip_present");
3436 stage_complete_companion(&models_dir, "clip-l");
3437 let json = r#"["clip-l","sdxl-vae"]"#;
3439 let missing = missing_companions_from_json(Some(json), &models_dir);
3440 assert_eq!(missing.len(), 1);
3441 assert_eq!(missing[0].name, "sdxl-vae");
3442 let _ = std::fs::remove_dir_all(&models_dir);
3443 }
3444
3445 #[test]
3446 fn missing_companions_preserves_input_order() {
3447 let models_dir = recipe_tmp_dir("companion_order");
3448 let json = r#"["sdxl-vae","clip-l","clip-g"]"#;
3449 let missing = missing_companions_from_json(Some(json), &models_dir);
3450 let names: Vec<&str> = missing.iter().map(|m| m.name.as_str()).collect();
3451 assert_eq!(names, vec!["sdxl-vae", "clip-l", "clip-g"]);
3452 let _ = std::fs::remove_dir_all(&models_dir);
3453 }
3454
3455 #[test]
3456 fn missing_companions_returns_empty_for_none_or_invalid() {
3457 let models_dir = recipe_tmp_dir("companion_empty");
3458 assert!(missing_companions_from_json(None, &models_dir).is_empty());
3459 assert!(missing_companions_from_json(Some("not json"), &models_dir).is_empty());
3460 assert!(missing_companions_from_json(Some("[]"), &models_dir).is_empty());
3461 let _ = std::fs::remove_dir_all(&models_dir);
3462 }
3463
3464 fn recipe_tmp_dir(label: &str) -> std::path::PathBuf {
3467 let dir = std::env::temp_dir().join(format!(
3468 "mold_recipe_{label}_{}",
3469 uuid::Uuid::new_v4().simple()
3470 ));
3471 std::fs::create_dir_all(&dir).unwrap();
3472 dir
3473 }
3474
3475 #[tokio::test]
3476 async fn h3_recipe_fetch_rejects_id_url_and_dest_before_any_side_effect() {
3477 use std::sync::atomic::{AtomicUsize, Ordering};
3478 use wiremock::matchers::method;
3479 use wiremock::{Mock, MockServer, ResponseTemplate};
3480
3481 let server = MockServer::start().await;
3482 Mock::given(method("GET"))
3483 .respond_with(ResponseTemplate::new(200).set_body_bytes(b"must not be fetched"))
3484 .expect(0)
3485 .mount(&server)
3486 .await;
3487
3488 let ordinary_url = format!("{}/weights.safetensors", server.uri());
3489 let gated_url = format!("{}/MiniMax-H3/weights.safetensors", server.uri());
3490 let cases = [
3491 (
3492 "id",
3493 "hf:MiniMaxAI/MiniMax-H3".to_string(),
3494 ordinary_url.clone(),
3495 "weights.safetensors".to_string(),
3496 ),
3497 (
3498 "url",
3499 "cv:opaque".to_string(),
3500 gated_url,
3501 "weights.safetensors".to_string(),
3502 ),
3503 (
3504 "dest",
3505 "cv:opaque".to_string(),
3506 ordinary_url,
3507 "MiniMax-H3/weights.safetensors".to_string(),
3508 ),
3509 ];
3510
3511 for (field, id, url, dest) in cases {
3512 let models_dir = std::env::temp_dir().join(format!(
3513 "mold_recipe_h3_{field}_{}",
3514 uuid::Uuid::new_v4().simple()
3515 ));
3516 assert!(!models_dir.exists(), "test path must begin absent");
3517
3518 let files = [RecipeFetchFile {
3519 url: &url,
3520 dest: &dest,
3521 sha256: None,
3522 size_bytes: Some(1),
3523 }];
3524 let progress_count = Arc::new(AtomicUsize::new(0));
3525 let observed = progress_count.clone();
3526 let progress: DownloadProgressCallback = Arc::new(move |_| {
3527 observed.fetch_add(1, Ordering::SeqCst);
3528 });
3529
3530 let error = fetch_recipe(
3531 &id,
3532 &files,
3533 RecipeAuth::None,
3534 &models_dir,
3535 Some(progress),
3536 &PullOptions::default(),
3537 )
3538 .await
3539 .expect_err("H3 recipe input must be compliance-gated");
3540
3541 assert!(
3542 matches!(error, DownloadError::ModelActivation(_)),
3543 "{field}"
3544 );
3545 assert_eq!(progress_count.load(Ordering::SeqCst), 0, "{field}");
3546 assert!(
3547 !models_dir.exists(),
3548 "{field} rejection must not create the models directory"
3549 );
3550 }
3551
3552 server.verify().await;
3553 }
3554
3555 #[tokio::test]
3556 async fn recipe_fetcher_writes_files_under_models_dir() {
3557 use wiremock::matchers::{method, path};
3558 use wiremock::{Mock, MockServer, ResponseTemplate};
3559
3560 let server = MockServer::start().await;
3561 Mock::given(method("GET"))
3562 .and(path("/file1.safetensors"))
3563 .respond_with(ResponseTemplate::new(200).set_body_bytes(b"hello".as_ref()))
3564 .mount(&server)
3565 .await;
3566 Mock::given(method("GET"))
3567 .and(path("/sub/file2.safetensors"))
3568 .respond_with(ResponseTemplate::new(200).set_body_bytes(b"world".as_ref()))
3569 .mount(&server)
3570 .await;
3571
3572 let models_dir = recipe_tmp_dir("writes");
3573 let url1 = format!("{}/file1.safetensors", server.uri());
3574 let url2 = format!("{}/sub/file2.safetensors", server.uri());
3575 let files = vec![
3576 RecipeFetchFile {
3577 url: &url1,
3578 dest: "file1.safetensors",
3579 sha256: None,
3580 size_bytes: None,
3581 },
3582 RecipeFetchFile {
3583 url: &url2,
3584 dest: "sub/file2.safetensors",
3585 sha256: None,
3586 size_bytes: None,
3587 },
3588 ];
3589
3590 let written = fetch_recipe(
3591 "cv:42",
3592 &files,
3593 RecipeAuth::None,
3594 &models_dir,
3595 None,
3596 &PullOptions::default(),
3597 )
3598 .await
3599 .expect("fetch_recipe ok");
3600
3601 let f1 = models_dir.join("cv-42").join("file1.safetensors");
3602 let f2 = models_dir
3603 .join("cv-42")
3604 .join("sub")
3605 .join("file2.safetensors");
3606 assert_eq!(written, vec![f1.clone(), f2.clone()]);
3607 assert_eq!(std::fs::read(&f1).unwrap(), b"hello");
3608 assert_eq!(std::fs::read(&f2).unwrap(), b"world");
3609
3610 let _ = std::fs::remove_dir_all(&models_dir);
3611 }
3612
3613 #[tokio::test]
3614 async fn recipe_fetcher_verifies_sha256_when_present_match() {
3615 use wiremock::matchers::{method, path};
3616 use wiremock::{Mock, MockServer, ResponseTemplate};
3617
3618 let server = MockServer::start().await;
3619 let body = b"hello world";
3620 let expected = "b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9";
3622 Mock::given(method("GET"))
3623 .and(path("/m.safetensors"))
3624 .respond_with(ResponseTemplate::new(200).set_body_bytes(body.as_ref()))
3625 .mount(&server)
3626 .await;
3627
3628 let models_dir = recipe_tmp_dir("sha_match");
3629 let url = format!("{}/m.safetensors", server.uri());
3630 let files = vec![RecipeFetchFile {
3631 url: &url,
3632 dest: "m.safetensors",
3633 sha256: Some(expected),
3634 size_bytes: None,
3635 }];
3636 fetch_recipe(
3637 "cv:1",
3638 &files,
3639 RecipeAuth::None,
3640 &models_dir,
3641 None,
3642 &PullOptions::default(),
3643 )
3644 .await
3645 .expect("matching SHA must succeed");
3646
3647 let _ = std::fs::remove_dir_all(&models_dir);
3648 }
3649
3650 #[tokio::test]
3651 async fn recipe_fetcher_verifies_sha256_when_present_mismatch() {
3652 use wiremock::matchers::{method, path};
3653 use wiremock::{Mock, MockServer, ResponseTemplate};
3654
3655 let server = MockServer::start().await;
3656 Mock::given(method("GET"))
3657 .and(path("/bad.safetensors"))
3658 .respond_with(ResponseTemplate::new(200).set_body_bytes(b"hello".as_ref()))
3659 .mount(&server)
3660 .await;
3661
3662 let models_dir = recipe_tmp_dir("sha_mismatch");
3663 let url = format!("{}/bad.safetensors", server.uri());
3664 let files = vec![RecipeFetchFile {
3666 url: &url,
3667 dest: "bad.safetensors",
3668 sha256: Some("0000000000000000000000000000000000000000000000000000000000000000"),
3669 size_bytes: None,
3670 }];
3671
3672 let err = fetch_recipe(
3673 "cv:2",
3674 &files,
3675 RecipeAuth::None,
3676 &models_dir,
3677 None,
3678 &PullOptions::default(),
3679 )
3680 .await
3681 .expect_err("mismatched SHA must error");
3682
3683 match err {
3684 DownloadError::Sha256Mismatch { filename, .. } => {
3685 assert_eq!(filename, "bad.safetensors");
3686 }
3687 other => panic!("expected Sha256Mismatch, got {other:?}"),
3688 }
3689 let bad = models_dir.join("cv-2").join("bad.safetensors");
3691 assert!(
3692 !bad.exists(),
3693 "corrupted file should be removed on mismatch"
3694 );
3695
3696 let _ = std::fs::remove_dir_all(&models_dir);
3697 }
3698
3699 #[tokio::test]
3700 async fn recipe_fetcher_marker_lifecycle() {
3701 use wiremock::matchers::{method, path};
3702 use wiremock::{Mock, MockServer, ResponseTemplate};
3703
3704 let server = MockServer::start().await;
3705 Mock::given(method("GET"))
3706 .and(path("/x.safetensors"))
3707 .respond_with(ResponseTemplate::new(200).set_body_bytes(b"x".as_ref()))
3708 .mount(&server)
3709 .await;
3710
3711 let models_dir = recipe_tmp_dir("marker");
3712 let url = format!("{}/x.safetensors", server.uri());
3713 let files = vec![RecipeFetchFile {
3714 url: &url,
3715 dest: "x.safetensors",
3716 sha256: None,
3717 size_bytes: None,
3718 }];
3719 let marker = pulling_marker_path_in(&models_dir, "cv:7");
3720 assert!(!marker.exists(), "marker should not exist before fetch");
3721
3722 fetch_recipe(
3723 "cv:7",
3724 &files,
3725 RecipeAuth::None,
3726 &models_dir,
3727 None,
3728 &PullOptions::default(),
3729 )
3730 .await
3731 .expect("ok");
3732
3733 assert!(
3734 !marker.exists(),
3735 "marker should be removed after successful fetch"
3736 );
3737 let _ = std::fs::remove_dir_all(&models_dir);
3738 }
3739
3740 #[tokio::test]
3741 async fn recipe_fetcher_skips_files_with_matching_size() {
3742 use wiremock::matchers::{method, path};
3743 use wiremock::{Mock, MockServer, ResponseTemplate};
3744
3745 let server = MockServer::start().await;
3746 let body = b"hello world";
3747 Mock::given(method("GET"))
3748 .and(path("/m.safetensors"))
3749 .respond_with(ResponseTemplate::new(200).set_body_bytes(body.as_ref()))
3750 .expect(1)
3752 .mount(&server)
3753 .await;
3754
3755 let models_dir = recipe_tmp_dir("idempotent_size");
3756 let url = format!("{}/m.safetensors", server.uri());
3757 let files = vec![RecipeFetchFile {
3758 url: &url,
3759 dest: "m.safetensors",
3760 sha256: None,
3761 size_bytes: Some(body.len() as u64),
3762 }];
3763
3764 fetch_recipe(
3765 "cv:idemp",
3766 &files,
3767 RecipeAuth::None,
3768 &models_dir,
3769 None,
3770 &PullOptions::default(),
3771 )
3772 .await
3773 .expect("first fetch ok");
3774
3775 fetch_recipe(
3778 "cv:idemp",
3779 &files,
3780 RecipeAuth::None,
3781 &models_dir,
3782 None,
3783 &PullOptions::default(),
3784 )
3785 .await
3786 .expect("second fetch ok (skip path)");
3787
3788 server.verify().await;
3791
3792 let _ = std::fs::remove_dir_all(&models_dir);
3793 }
3794
3795 #[tokio::test]
3796 async fn recipe_fetcher_skips_files_with_sha256_marker_when_size_unknown() {
3797 use wiremock::matchers::{method, path};
3798 use wiremock::{Mock, MockServer, ResponseTemplate};
3799
3800 let server = MockServer::start().await;
3801 Mock::given(method("GET"))
3802 .and(path("/m.safetensors"))
3803 .respond_with(ResponseTemplate::new(200).set_body_bytes(b"x".as_ref()))
3804 .expect(1)
3805 .mount(&server)
3806 .await;
3807
3808 let models_dir = recipe_tmp_dir("idempotent_marker");
3809 let url = format!("{}/m.safetensors", server.uri());
3810 let files = vec![RecipeFetchFile {
3811 url: &url,
3812 dest: "m.safetensors",
3813 sha256: None,
3814 size_bytes: None,
3816 }];
3817
3818 fetch_recipe(
3820 "cv:idemp_marker",
3821 &files,
3822 RecipeAuth::None,
3823 &models_dir,
3824 None,
3825 &PullOptions::default(),
3826 )
3827 .await
3828 .expect("first fetch ok");
3829
3830 let dest = models_dir.join("cv-idemp_marker").join("m.safetensors");
3832 assert!(
3833 sha256_marker_path(&dest).exists(),
3834 "first fetch should have written the .sha256-verified marker"
3835 );
3836
3837 fetch_recipe(
3838 "cv:idemp_marker",
3839 &files,
3840 RecipeAuth::None,
3841 &models_dir,
3842 None,
3843 &PullOptions::default(),
3844 )
3845 .await
3846 .expect("second fetch ok");
3847
3848 server.verify().await;
3849 let _ = std::fs::remove_dir_all(&models_dir);
3850 }
3851
3852 #[tokio::test]
3853 async fn recipe_fetcher_refetches_when_sha256_declared_but_marker_missing() {
3854 use wiremock::matchers::{method, path};
3855 use wiremock::{Mock, MockServer, ResponseTemplate};
3856
3857 let server = MockServer::start().await;
3858 let body = b"correct";
3859 let expected = "15a596e3c98c407e043751ff3b21ff0358a1bdfdf3fe948b1523893a8e5de2e8";
3861 Mock::given(method("GET"))
3862 .and(path("/m.safetensors"))
3863 .respond_with(ResponseTemplate::new(200).set_body_bytes(body.as_ref()))
3864 .expect(1)
3867 .mount(&server)
3868 .await;
3869
3870 let models_dir = recipe_tmp_dir("idempotent_no_marker");
3871 let subdir = models_dir.join("cv-idemp_no_marker");
3872 std::fs::create_dir_all(&subdir).unwrap();
3873 let dest = subdir.join("m.safetensors");
3874 std::fs::write(&dest, b"BADBYTE").unwrap();
3878 assert!(!sha256_marker_path(&dest).exists());
3879
3880 let url = format!("{}/m.safetensors", server.uri());
3881 let files = vec![RecipeFetchFile {
3882 url: &url,
3883 dest: "m.safetensors",
3884 sha256: Some(expected),
3885 size_bytes: Some(body.len() as u64),
3886 }];
3887
3888 fetch_recipe(
3889 "cv:idemp_no_marker",
3890 &files,
3891 RecipeAuth::None,
3892 &models_dir,
3893 None,
3894 &PullOptions::default(),
3895 )
3896 .await
3897 .expect("fetch ok");
3898
3899 assert_eq!(std::fs::read(&dest).unwrap(), body);
3901 assert!(sha256_marker_path(&dest).exists());
3902 server.verify().await;
3903 let _ = std::fs::remove_dir_all(&models_dir);
3904 }
3905
3906 #[test]
3907 fn catalog_entry_installed_returns_true_for_complete_recipe() {
3908 let models_dir = recipe_tmp_dir("installed_complete");
3909 let subdir = models_dir.join("cv-installed_a");
3910 std::fs::create_dir_all(&subdir).unwrap();
3911 let dest = subdir.join("m.safetensors");
3912 std::fs::write(&dest, b"hello").unwrap();
3913
3914 let files = vec![RecipeFetchFile {
3915 url: "https://example.invalid/m.safetensors",
3916 dest: "m.safetensors",
3917 sha256: None,
3918 size_bytes: Some(5),
3919 }];
3920
3921 assert!(catalog_entry_installed(
3922 &models_dir,
3923 "cv:installed_a",
3924 &files
3925 ));
3926 let _ = std::fs::remove_dir_all(&models_dir);
3927 }
3928
3929 #[test]
3930 fn catalog_entry_installed_returns_false_when_any_file_missing() {
3931 let models_dir = recipe_tmp_dir("installed_partial");
3932 let subdir = models_dir.join("cv-installed_b");
3933 std::fs::create_dir_all(&subdir).unwrap();
3934 std::fs::write(subdir.join("a.safetensors"), b"present").unwrap();
3935 let files = vec![
3938 RecipeFetchFile {
3939 url: "https://example.invalid/a.safetensors",
3940 dest: "a.safetensors",
3941 sha256: None,
3942 size_bytes: Some(7),
3943 },
3944 RecipeFetchFile {
3945 url: "https://example.invalid/b.safetensors",
3946 dest: "b.safetensors",
3947 sha256: None,
3948 size_bytes: Some(7),
3949 },
3950 ];
3951
3952 assert!(!catalog_entry_installed(
3953 &models_dir,
3954 "cv:installed_b",
3955 &files
3956 ));
3957 let _ = std::fs::remove_dir_all(&models_dir);
3958 }
3959
3960 #[test]
3961 fn catalog_entry_installed_returns_false_on_size_mismatch() {
3962 let models_dir = recipe_tmp_dir("installed_mismatch");
3963 let subdir = models_dir.join("cv-installed_c");
3964 std::fs::create_dir_all(&subdir).unwrap();
3965 std::fs::write(subdir.join("m.safetensors"), b"WRONG").unwrap();
3966
3967 let files = vec![RecipeFetchFile {
3968 url: "https://example.invalid/m.safetensors",
3969 dest: "m.safetensors",
3970 sha256: None,
3971 size_bytes: Some(99),
3972 }];
3973
3974 assert!(!catalog_entry_installed(
3975 &models_dir,
3976 "cv:installed_c",
3977 &files
3978 ));
3979 let _ = std::fs::remove_dir_all(&models_dir);
3980 }
3981
3982 #[test]
3983 fn catalog_entry_installed_accepts_marker_when_declared_size_is_stale() {
3984 let models_dir = recipe_tmp_dir("installed_stale_size_marker");
3985 let subdir = models_dir.join("cv-installed_c2");
3986 std::fs::create_dir_all(&subdir).unwrap();
3987 let dest = subdir.join("m.safetensors");
3988 std::fs::write(&dest, b"new larger bytes").unwrap();
3989 write_sha256_marker(&dest, "deadbeef").unwrap();
3990
3991 let files = vec![RecipeFetchFile {
3992 url: "https://example.invalid/m.safetensors",
3993 dest: "m.safetensors",
3994 sha256: None,
3995 size_bytes: Some(5),
3996 }];
3997
3998 assert!(catalog_entry_installed(
3999 &models_dir,
4000 "cv:installed_c2",
4001 &files
4002 ));
4003 let _ = std::fs::remove_dir_all(&models_dir);
4004 }
4005
4006 #[test]
4007 fn catalog_entry_installed_uses_marker_when_size_unknown() {
4008 let models_dir = recipe_tmp_dir("installed_marker");
4009 let subdir = models_dir.join("cv-installed_d");
4010 std::fs::create_dir_all(&subdir).unwrap();
4011 let dest = subdir.join("m.safetensors");
4012 std::fs::write(&dest, b"hello").unwrap();
4013 write_sha256_marker(&dest, "deadbeef").unwrap();
4014
4015 let files = vec![RecipeFetchFile {
4016 url: "https://example.invalid/m.safetensors",
4017 dest: "m.safetensors",
4018 sha256: None,
4019 size_bytes: None,
4020 }];
4021
4022 assert!(catalog_entry_installed(
4023 &models_dir,
4024 "cv:installed_d",
4025 &files
4026 ));
4027 let _ = std::fs::remove_dir_all(&models_dir);
4028 }
4029
4030 #[test]
4031 fn catalog_entry_installed_returns_false_without_marker_and_without_size() {
4032 let models_dir = recipe_tmp_dir("installed_nomarker");
4033 let subdir = models_dir.join("cv-installed_e");
4034 std::fs::create_dir_all(&subdir).unwrap();
4035 std::fs::write(subdir.join("m.safetensors"), b"hello").unwrap();
4036 let files = vec![RecipeFetchFile {
4039 url: "https://example.invalid/m.safetensors",
4040 dest: "m.safetensors",
4041 sha256: None,
4042 size_bytes: None,
4043 }];
4044
4045 assert!(!catalog_entry_installed(
4046 &models_dir,
4047 "cv:installed_e",
4048 &files
4049 ));
4050 let _ = std::fs::remove_dir_all(&models_dir);
4051 }
4052
4053 #[test]
4054 fn catalog_entry_installed_returns_false_when_pulling_marker_present() {
4055 let models_dir = recipe_tmp_dir("installed_pulling");
4056 let subdir = models_dir.join("cv-installed_f");
4057 std::fs::create_dir_all(&subdir).unwrap();
4058 std::fs::write(subdir.join("m.safetensors"), b"hello").unwrap();
4059
4060 let marker = pulling_marker_path_in(&models_dir, "cv:installed_f");
4061 if let Some(parent) = marker.parent() {
4062 std::fs::create_dir_all(parent).unwrap();
4063 }
4064 std::fs::write(&marker, "in-progress").unwrap();
4065
4066 let files = vec![RecipeFetchFile {
4067 url: "https://example.invalid/m.safetensors",
4068 dest: "m.safetensors",
4069 sha256: None,
4070 size_bytes: Some(5),
4071 }];
4072
4073 assert!(
4074 !catalog_entry_installed(&models_dir, "cv:installed_f", &files),
4075 "active .pulling marker must override on-disk completeness"
4076 );
4077 let _ = std::fs::remove_dir_all(&models_dir);
4078 }
4079
4080 #[test]
4081 fn catalog_entry_installed_rejects_path_traversal() {
4082 let models_dir = recipe_tmp_dir("installed_traversal");
4083
4084 let files = vec![RecipeFetchFile {
4085 url: "https://example.invalid/m.safetensors",
4086 dest: "../escape.safetensors",
4087 sha256: None,
4088 size_bytes: Some(5),
4089 }];
4090
4091 assert!(
4092 !catalog_entry_installed(&models_dir, "cv:installed_g", &files),
4093 "path traversal must be treated as not-installed, not as a panic"
4094 );
4095 let _ = std::fs::remove_dir_all(&models_dir);
4096 }
4097
4098 #[test]
4099 fn catalog_entry_installed_returns_false_for_empty_files() {
4100 let models_dir = recipe_tmp_dir("installed_empty");
4101 assert!(
4102 !catalog_entry_installed(&models_dir, "cv:installed_h", &[]),
4103 "empty file slice means no recipe to verify; must refuse to claim install"
4104 );
4105 let _ = std::fs::remove_dir_all(&models_dir);
4106 }
4107
4108 #[test]
4109 fn catalog_entry_installed_returns_true_for_multi_file_complete_recipe() {
4110 let models_dir = recipe_tmp_dir("installed_multi");
4111 let subdir = models_dir.join("cv-installed_i");
4112 std::fs::create_dir_all(&subdir).unwrap();
4113 std::fs::write(subdir.join("a.safetensors"), b"present").unwrap();
4114 std::fs::write(subdir.join("b.safetensors"), b"present_too").unwrap();
4115 std::fs::write(subdir.join("c.safetensors"), b"third").unwrap();
4116
4117 let files = vec![
4118 RecipeFetchFile {
4119 url: "https://example.invalid/a.safetensors",
4120 dest: "a.safetensors",
4121 sha256: None,
4122 size_bytes: Some(7),
4123 },
4124 RecipeFetchFile {
4125 url: "https://example.invalid/b.safetensors",
4126 dest: "b.safetensors",
4127 sha256: None,
4128 size_bytes: Some(11),
4129 },
4130 RecipeFetchFile {
4131 url: "https://example.invalid/c.safetensors",
4132 dest: "c.safetensors",
4133 sha256: None,
4134 size_bytes: Some(5),
4135 },
4136 ];
4137
4138 assert!(
4139 catalog_entry_installed(&models_dir, "cv:installed_i", &files),
4140 "every file present at declared size — must report installed"
4141 );
4142 let _ = std::fs::remove_dir_all(&models_dir);
4143 }
4144
4145 #[test]
4146 fn catalog_entry_installed_returns_false_when_file_larger_than_declared() {
4147 let models_dir = recipe_tmp_dir("installed_too_big");
4152 let subdir = models_dir.join("cv-installed_j");
4153 std::fs::create_dir_all(&subdir).unwrap();
4154 std::fs::write(
4155 subdir.join("m.safetensors"),
4156 b"this is much longer than five bytes",
4157 )
4158 .unwrap();
4159
4160 let files = vec![RecipeFetchFile {
4161 url: "https://example.invalid/m.safetensors",
4162 dest: "m.safetensors",
4163 sha256: None,
4164 size_bytes: Some(5),
4165 }];
4166
4167 assert!(!catalog_entry_installed(
4168 &models_dir,
4169 "cv:installed_j",
4170 &files
4171 ));
4172 let _ = std::fs::remove_dir_all(&models_dir);
4173 }
4174
4175 #[test]
4176 fn catalog_entry_installed_requires_marker_when_sha256_declared() {
4177 let models_dir = recipe_tmp_dir("installed_no_marker");
4184 let subdir = models_dir.join("cv-installed_k");
4185 std::fs::create_dir_all(&subdir).unwrap();
4186 std::fs::write(subdir.join("m.safetensors"), b"hello").unwrap();
4187 let files = vec![RecipeFetchFile {
4190 url: "https://example.invalid/m.safetensors",
4191 dest: "m.safetensors",
4192 sha256: Some("deadbeef00000000000000000000000000000000000000000000000000000000"),
4193 size_bytes: Some(5),
4194 }];
4195
4196 assert!(
4197 !catalog_entry_installed(&models_dir, "cv:installed_k", &files),
4198 "size matches but no marker AND sha256 declared — must refuse to claim install",
4199 );
4200 let _ = std::fs::remove_dir_all(&models_dir);
4201 }
4202
4203 #[test]
4204 fn catalog_entry_installed_trusts_marker_over_stale_size_bytes() {
4205 let models_dir = recipe_tmp_dir("installed_stale_size");
4211 let subdir = models_dir.join("cv-installed_stale");
4212 std::fs::create_dir_all(&subdir).unwrap();
4213 let dest = subdir.join("m.safetensors");
4214 std::fs::write(&dest, b"hello").unwrap();
4216 write_sha256_marker(
4217 &dest,
4218 "deadbeef00000000000000000000000000000000000000000000000000000000",
4219 )
4220 .unwrap();
4221
4222 let files = vec![RecipeFetchFile {
4223 url: "https://example.invalid/m.safetensors",
4224 dest: "m.safetensors",
4225 sha256: Some("deadbeef00000000000000000000000000000000000000000000000000000000"),
4226 size_bytes: Some(99), }];
4228
4229 assert!(
4230 catalog_entry_installed(&models_dir, "cv:installed_stale", &files),
4231 "sha256 marker present → installed despite stale size_bytes",
4232 );
4233 let _ = std::fs::remove_dir_all(&models_dir);
4234 }
4235
4236 #[tokio::test]
4237 async fn recipe_fetcher_pulls_when_size_mismatch() {
4238 use wiremock::matchers::{method, path};
4239 use wiremock::{Mock, MockServer, ResponseTemplate};
4240
4241 let server = MockServer::start().await;
4242 Mock::given(method("GET"))
4243 .and(path("/m.safetensors"))
4244 .respond_with(ResponseTemplate::new(200).set_body_bytes(b"correct".as_ref()))
4245 .expect(1)
4247 .mount(&server)
4248 .await;
4249
4250 let models_dir = recipe_tmp_dir("idempotent_mismatch");
4251 let subdir = models_dir.join("cv-idemp_mismatch");
4252 std::fs::create_dir_all(&subdir).unwrap();
4253 let dest = subdir.join("m.safetensors");
4254 std::fs::write(&dest, b"WRNG").unwrap();
4256
4257 let url = format!("{}/m.safetensors", server.uri());
4258 let files = vec![RecipeFetchFile {
4259 url: &url,
4260 dest: "m.safetensors",
4261 sha256: None,
4262 size_bytes: Some(7),
4263 }];
4264
4265 fetch_recipe(
4266 "cv:idemp_mismatch",
4267 &files,
4268 RecipeAuth::None,
4269 &models_dir,
4270 None,
4271 &PullOptions::default(),
4272 )
4273 .await
4274 .expect("ok");
4275
4276 assert_eq!(std::fs::read(&dest).unwrap(), b"correct");
4278 server.verify().await;
4279 let _ = std::fs::remove_dir_all(&models_dir);
4280 }
4281
4282 #[tokio::test]
4283 async fn recipe_fetcher_rejects_path_traversal_in_dest() {
4284 let models_dir = recipe_tmp_dir("traversal");
4285 let files = vec![RecipeFetchFile {
4286 url: "http://example.invalid/should-not-be-fetched",
4287 dest: "../etc/passwd",
4288 sha256: None,
4289 size_bytes: None,
4290 }];
4291 let err = fetch_recipe(
4292 "cv:8",
4293 &files,
4294 RecipeAuth::None,
4295 &models_dir,
4296 None,
4297 &PullOptions::default(),
4298 )
4299 .await
4300 .expect_err("traversal must be rejected");
4301 match err {
4302 DownloadError::RecipePathTraversal { dest } => {
4303 assert_eq!(dest, "../etc/passwd");
4304 }
4305 other => panic!("expected RecipePathTraversal, got {other:?}"),
4306 }
4307 assert!(
4309 !models_dir.join("cv-8").exists()
4310 || std::fs::read_dir(models_dir.join("cv-8"))
4311 .map(|d| d.count())
4312 .unwrap_or(0)
4313 == 0
4314 );
4315 let _ = std::fs::remove_dir_all(&models_dir);
4316 }
4317
4318 #[tokio::test]
4319 async fn recipe_fetcher_rejects_absolute_dest() {
4320 let models_dir = recipe_tmp_dir("absolute");
4321 let files = vec![RecipeFetchFile {
4322 url: "http://example.invalid/should-not-be-fetched",
4323 dest: "/etc/passwd",
4324 sha256: None,
4325 size_bytes: None,
4326 }];
4327 let err = fetch_recipe(
4328 "cv:9",
4329 &files,
4330 RecipeAuth::None,
4331 &models_dir,
4332 None,
4333 &PullOptions::default(),
4334 )
4335 .await
4336 .expect_err("absolute dest must be rejected");
4337 assert!(matches!(err, DownloadError::RecipePathTraversal { .. }));
4338 let _ = std::fs::remove_dir_all(&models_dir);
4339 }
4340
4341 #[tokio::test]
4342 async fn recipe_fetcher_sends_bearer_token_when_auth_set() {
4343 use wiremock::matchers::{header, method, path};
4344 use wiremock::{Mock, MockServer, ResponseTemplate};
4345
4346 let server = MockServer::start().await;
4347 Mock::given(method("GET"))
4348 .and(path("/civitai.safetensors"))
4349 .and(header("authorization", "Bearer secret-cv-token"))
4350 .respond_with(ResponseTemplate::new(200).set_body_bytes(b"ok".as_ref()))
4351 .mount(&server)
4352 .await;
4353
4354 let models_dir = recipe_tmp_dir("bearer");
4355 let url = format!("{}/civitai.safetensors", server.uri());
4356 let files = vec![RecipeFetchFile {
4357 url: &url,
4358 dest: "civitai.safetensors",
4359 sha256: None,
4360 size_bytes: None,
4361 }];
4362 fetch_recipe(
4363 "cv:618692",
4364 &files,
4365 RecipeAuth::Bearer("secret-cv-token".to_string()),
4366 &models_dir,
4367 None,
4368 &PullOptions::default(),
4369 )
4370 .await
4371 .expect("authenticated request must succeed");
4372
4373 let _ = std::fs::remove_dir_all(&models_dir);
4374 }
4375}