Skip to main content

mold_core/
download.rs

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/// Callback-based download progress event.
25#[derive(Debug, Clone)]
26pub enum DownloadProgressEvent {
27    /// A file download has started.
28    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    /// Bytes downloaded for the current file.
38    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 (e.g. "Verifying cached files...").
48    Status { message: String },
49    /// A file download completed.
50    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
60/// Callback type for download progress reporting.
61pub type DownloadProgressCallback = Arc<dyn Fn(DownloadProgressEvent) + Send + Sync>;
62
63/// Options controlling model pull behavior.
64#[derive(Debug, Clone, Default)]
65pub struct PullOptions {
66    /// Skip SHA-256 verification after download (use when HF updated a file).
67    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
149/// Does a GGUF file's header contain the given tensor name?
150///
151/// Scans the first 4 MiB of the file — enough to cover tensor_infos for any
152/// real FLUX GGUF (~800 tensors × ~100 B per entry). Tensor names are stored
153/// as UTF-8 in the header, so a substring search is reliable: the needle is
154/// length-prefixed by a u64, so accidental coincidences in the scanned region
155/// would need to match a 31+ character needle exactly.
156fn 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
172/// Decide whether to emit the pull-time "city96-format, needs reference" warning.
173///
174/// Pure logic, no process-global state — `models_dir` is always passed in so
175/// tests can use a temp dir. Returns `Some(message)` when the warning should
176/// fire, `None` otherwise.
177fn 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    // img_in is present in schnell and in complete dev GGUFs; missing from city96-format
190    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
247/// Warn the operator if the downloaded transformer is a city96-format GGUF
248/// that will need an additional reference pull before inference will run.
249///
250/// Community FLUX fine-tune GGUFs ship only the diffusion blocks; their input
251/// embedding layers (img_in / time_in / vector_in / guidance_in) are inherited
252/// from base flux-dev and must be patched in from a locally-downloaded
253/// reference. This check surfaces the dependency at pull time so users don't
254/// discover it on the first generation attempt.
255fn 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
271/// Resolve HuggingFace token: `HF_TOKEN` env var takes precedence over
272/// the token file (`~/.cache/huggingface/token` from `huggingface-cli login`).
273fn 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
293/// Resolve the mold models directory. Computed once from config on first access.
294/// Resolution order: `MOLD_MODELS_DIR` env var → config `models_dir` → `~/.mold/models`.
295///
296/// This is the clean model storage root. Actual model files live at clean paths like
297/// `models/flux-schnell-q8/transformer.gguf` and `models/shared/flux/ae.safetensors`.
298///
299/// **OnceLock caching**: The directory is resolved once on the first call and cached
300/// for the entire process lifetime. Changing `MOLD_MODELS_DIR` or the config file
301/// after the first call has no effect. This is by design — model paths recorded in
302/// config must remain stable within a single process run.
303fn 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
313/// Internal hf-hub cache directory: `<models_dir>/.hf-cache/`.
314/// Hidden from users; files get hardlinked to clean paths after download.
315fn 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
325/// Hardlink `src` to `dst`, falling back to copy if hardlink fails (cross-filesystem).
326/// Idempotent: skips if `dst` already exists with the same size as `src`.
327///
328/// The source path is canonicalized to resolve hf-hub's symlink chain
329/// (`snapshots/<sha>/file → ../../blobs/<hash>`) before any filesystem ops.
330fn hardlink_or_copy(src: &std::path::Path, dst: &std::path::Path) -> Result<(), DownloadError> {
331    // Resolve symlinks — hf-hub cache returns symlink paths that can cause
332    // ENOENT on some filesystems when passed directly to hard_link or copy.
333    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    // Check if dst already has the correct content (idempotent skip).
341    // Use metadata() which follows symlinks — only skip if the real target matches.
342    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    // Remove stale destination before placement. A previous hard_link on an
351    // hf-hub symlink creates a relative symlink that dangles from the new
352    // location (e.g. shared/sd3/file → ../../blobs/hash, which doesn't exist
353    // relative to shared/sd3/). symlink_metadata() sees these even though
354    // exists() returns false for dangling symlinks.
355    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    // Try hardlink first (zero extra disk space, instant)
368    match std::fs::hard_link(&real_src, dst) {
369        Ok(()) => return Ok(()),
370        Err(_e) => {
371            // Expected on cross-filesystem setups; fall through to copy
372        }
373    }
374    // Fall back to copy (cross-filesystem or hard_link unsupported)
375    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
385/// Compute the SHA-256 hex digest of a file.
386pub 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
395/// Verify the SHA-256 digest of a file against an expected hex string.
396/// Comparison is hex-case-insensitive — Civitai's API publishes uppercase
397/// hashes and `compute_sha256` produces lowercase, so a literal `==`
398/// would false-mismatch on bit-identical files.
399///
400/// Returns `Ok(true)` when the digest matches, `Ok(false)` on mismatch.
401/// Errors only on I/O failures (e.g. file not found).
402pub fn verify_sha256(path: &std::path::Path, expected: &str) -> anyhow::Result<bool> {
403    Ok(compute_sha256(path)?.eq_ignore_ascii_case(expected))
404}
405
406// ── Pull marker file (.pulling) ──────────────────────────────────────────────
407
408/// Relative path to a model's `.pulling` marker: `<sanitized-name>/.pulling`.
409pub 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
414/// Path to the `.pulling` marker for a model under an explicit models dir.
415pub fn pulling_marker_path_in(models_dir: &Path, model_name: &str) -> PathBuf {
416    models_dir.join(pulling_marker_rel_path(model_name))
417}
418
419/// Path to the `.pulling` marker for a model: `<models_dir>/<sanitized-name>/.pulling`.
420fn pulling_marker_path(model_name: &str) -> PathBuf {
421    pulling_marker_path_in(&models_dir(), model_name)
422}
423
424/// Write a `.pulling` marker to signal an in-progress download.
425fn 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
443/// Remove the `.pulling` marker (best-effort, ignores errors).
444pub fn remove_pulling_marker(model_name: &str) {
445    let path = pulling_marker_path(model_name);
446    let _ = std::fs::remove_file(path);
447}
448
449/// Check whether a model has an active `.pulling` marker (incomplete download).
450pub 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
455/// Filename suffix for the "this file is fully written and integrity-checked"
456/// sidecar marker: `model.safetensors` → `model.safetensors.sha256-verified`.
457///
458/// The marker is written by [`verify_file_integrity`] on a successful pull
459/// (or by the post-startup backfill sweep for pre-marker installs). Two
460/// downstream consumers depend on it:
461///
462/// 1. `cleanup_partials_in_dir` (in `mold-server`) preserves any file that
463///    has a sibling marker — those are known-good and survive cancel/retry.
464/// 2. `Config::manifest_files_exist` requires the marker before reporting a
465///    model as "downloaded" — eliminates the existence-only race that let
466///    truncated files masquerade as complete installs.
467pub const SHA256_VERIFIED_SUFFIX: &str = ".sha256-verified";
468
469/// Minimum interval between `FileProgress` events emitted by the recipe-pull
470/// path (`fetch_recipe_inner`). The manifest-pull path's `CallbackProgress`
471/// throttles to the same cadence; this constant keeps them in sync. 250ms
472/// matches a comfortable UI refresh rate (~4 Hz) without flooding SSE
473/// subscribers when downloads run at multi-MB/s chunk rates.
474pub const RECIPE_PROGRESS_THROTTLE_MS: u64 = 250;
475
476/// Build the marker path for a downloaded file. `model.safetensors` →
477/// `model.safetensors.sha256-verified` in the same directory.
478pub 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
484/// True iff `<path>.sha256-verified` exists.
485pub fn has_sha256_marker(path: &Path) -> bool {
486    sha256_marker_path(path).exists()
487}
488
489/// Atomically write the `.sha256-verified` marker for `path` recording the
490/// computed digest. Atomic via tempfile-then-rename so a crash mid-write
491/// never leaves a half-populated marker (which would otherwise read as a
492/// successfully-installed file).
493pub 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
500/// Verify SHA-256 integrity of a downloaded file and write the
501/// `.sha256-verified` marker on success.
502///
503/// - Manifest declares `sha256`: compute, compare, on match write marker
504///   (containing the verified digest); on mismatch delete the corrupted
505///   file and return `Sha256Mismatch`.
506/// - Manifest does not declare a hash: still compute and write the marker
507///   so the file is positively attested as "fully written." This is the
508///   load-bearing change for the gallery race — `Config::manifest_files_exist`
509///   consults marker presence, so unmarked-but-present files no longer
510///   appear in the available-models list.
511/// - `skip_verify = true`: respected from the original contract — no read,
512///   no marker. The caller has explicitly asked us to trust the bytes.
513fn 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            // I/O failure during hashing — log and move on without a marker.
526            // The downstream `manifest_files_exist` check will report the
527            // file incomplete, prompting a retry rather than a silent pass.
528            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        // Marker-write failure isn't fatal to this attempt — the file is
548        // good. But it does mean the next `manifest_files_exist` check will
549        // report incomplete. Log loudly so users can see why.
550        eprintln!(
551            "warning: failed to write .sha256-verified marker for {}: {e}",
552            file.hf_filename
553        );
554    }
555    Ok(())
556}
557
558/// Truncate a string to fit within `max_len`, replacing the middle with "..." if needed.
559fn truncate_filename(name: &str, max_len: usize) -> String {
560    if name.len() <= max_len || max_len < 8 {
561        return name.to_string();
562    }
563    // Keep the end of the filename (the unique part) and trim the start
564    let suffix_len = max_len - 3; // "..." prefix
565    let start = name.len() - suffix_len;
566    format!("...{}", &name[start..])
567}
568
569/// Maximum characters for the filename column in progress bars.
570/// Derived from terminal width minus the fixed overhead of the bar template:
571/// 2 (indent) + 1 (space) + 1 ([) + 30 (bar) + 1 (]) + ~40 (bytes/speed/eta) = ~75 chars overhead.
572fn 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/// Progress adapter bridging hf-hub's `Progress` trait to an `indicatif::ProgressBar`.
578#[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/// Progress adapter that dispatches to a callback instead of indicatif.
612/// Throttles `FileProgress` events to ~4/sec per file to avoid flooding SSE.
613#[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
731/// Sync progress adapter bridging hf-hub's sync `Progress` trait to our
732/// local `indicatif::ProgressBar`.
733struct 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
765/// Synchronous hf-hub progress adapter used by pre-admission dependency
766/// preparation running on Tokio's blocking pool.
767struct 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
836/// Returns `true` if the file already exists at `clean_path` with the correct
837/// size and (if a SHA-256 is available) the correct digest.
838///
839/// **Side-effect**: if the file exists with matching size but failing integrity,
840/// `verify_file_integrity` will delete the corrupted file before returning `false`.
841fn 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 integrity — a same-size but corrupted file must not be accepted
855    verify_file_integrity(clean_path, file, model_name, skip_verify).is_ok()
856}
857
858/// Return an existing valid clean path for a manifest file, migrating from a
859/// legacy location when needed.
860fn 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
884/// Download all files for a model manifest, returning resolved paths.
885///
886/// Downloads go to a hidden hf-hub cache (`.hf-cache/`) for resume/dedup support,
887/// then files are hardlinked to clean paths:
888/// - Transformers → `<model-name>/<filename>`
889/// - Shared components → `shared/<family>/<filename>`
890///
891/// A `.pulling` marker file is written before downloads begin and removed on
892/// success. If the pull is interrupted, the marker signals an incomplete state.
893fn 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
978/// Download all files for a model manifest, reporting progress via callback.
979///
980/// Same as `pull_model` but uses a callback instead of indicatif progress bars.
981/// Suitable for server-side downloads where terminal bars are not appropriate.
982pub 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    // Pre-compute which files need downloading vs already cached.
1009    // Run in spawn_blocking because SHA-256 verification of multi-GB cached
1010    // files blocks the async runtime and prevents SSE event delivery.
1011    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            // Emit events for cached files so the TUI shows checkmarks.
1057            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
1107/// Download all files for a utility model (no ModelPaths, no config writing).
1108///
1109/// Used for models like qwen3-expand that are not diffusion models and don't
1110/// have a VAE. Files are downloaded and placed at their standard storage paths.
1111async 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
1275/// Extract HTTP status code from an async `ApiError`, if available.
1276fn 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
1288/// Return the process-wide flight for one Hugging Face repository file.
1289///
1290/// Different model variants often reuse the same large encoder or VAE blob.
1291/// The server deliberately runs unrelated pulls in parallel, so coordinate
1292/// only identical HF files here before entering hf-hub's short-lived file
1293/// lock. Weak entries keep completed identities from accumulating forever.
1294async 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
1373// ── Synchronous single-file download (for use from spawn_blocking) ───────────
1374
1375fn 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
1388/// Download a single file from HuggingFace, returning its path.
1389/// Uses the sync hf-hub API — safe to call from `spawn_blocking`.
1390/// Returns immediately if already cached.
1391///
1392/// If `target_subdir` is provided (e.g., `"shared/t5-gguf"`), the file is hardlinked
1393/// from the hf-cache to `<models_dir>/<target_subdir>/<leaf_filename>` and that clean
1394/// path is returned. If `None`, the raw hf-cache path is returned.
1395pub 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
1421/// Callback-reporting counterpart to [`download_single_file_sync`].
1422///
1423/// It remains a blocking function by design; callers must use
1424/// `tokio::task::spawn_blocking`.
1425pub 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
1442/// Explicit-root counterpart used when a caller owns an immutable config
1443/// snapshot. Both the managed Hugging Face cache and clean target stay under
1444/// `models_root`.
1445pub 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
1463/// Deterministic clean path that [`download_single_file_sync`] will populate
1464/// for a dependency with a target subdirectory.
1465///
1466/// This performs no I/O and does not imply that the file is present. It is
1467/// used by read-only placement previews to build the same engine input shape
1468/// that admission will materialize later.
1469pub 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
1473/// No-I/O counterpart used by read-only previews with their exact config
1474/// snapshot. The root is never created.
1475pub 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    // Keep the policy at the lowest download boundary as a defense against a
1495    // future internal caller bypassing the public wrappers. The wrappers also
1496    // check before constructing progress adapters or resolving managed paths.
1497    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    // Place at clean path if target_subdir specified
1539    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
1548/// Check whether a file is present in mold's managed hf-hub cache
1549/// (`<models_dir>/.hf-cache/`). Narrower than [`cached_file_path`] — does
1550/// not consult the system-wide `~/.cache/huggingface/`, the legacy mold
1551/// models cache, or any clean-path location. Used as a layout-agnostic
1552/// fallback by `Config::discovered_manifest_paths` so a single shard set
1553/// downloaded by a manifest install can also satisfy a catalog companion
1554/// that expects the same files under a different canonical layout (e.g.
1555/// the Gemma TE shared by `ltx-2.3-22b-distilled:fp8` and the catalog
1556/// `ltx2-te` companion). Tests that intentionally set up a "model not
1557/// downloaded" world are unaffected because they only override
1558/// `MOLD_MODELS_DIR`, not the user's home HF cache.
1559pub 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
1565/// Check if a file is already cached locally (no download).
1566///
1567/// If `target_subdir` is provided, checks the clean path first
1568/// (`<models_dir>/<target_subdir>/<leaf_filename>`). Then checks the hf-cache,
1569/// old mold models dir (backward compat), and default HF cache.
1570pub 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
1578/// Explicit-root cache lookup for admission. This preserves the historical
1579/// fallback search while keeping the managed cache and clean target bound to
1580/// the caller's config snapshot.
1581pub 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    // 1. Check clean path (if target_subdir specified)
1588    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    // 2. Check new hf-cache location (~/.mold/models/.hf-cache/)
1596    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    // 3. Check old mold models dir (backward compat — HF cached here before .hf-cache/)
1603    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    // 4. Check default HF cache (~/.cache/huggingface/hub/)
1610    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
1615/// Strictly read-only cache inspection for placement previews.
1616///
1617/// Every cache constructor is gated on an already-existing root, and the
1618/// clean destination is derived from the caller's immutable config snapshot.
1619/// This function never creates the models root, `.hf-cache`, or the default
1620/// Hugging Face cache.
1621pub 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
1652// ── Pull and configure (shared between CLI and server) ───────────────────────
1653
1654/// Hidden LTX-2 adapters are complete, runnable assets by themselves. They
1655/// must use the files-only pull path: treating their single LoRA tensor as a
1656/// standalone diffusion checkpoint makes `paths_from_downloads` require a
1657/// VAE after the verified file has already landed. Upstream publishes the
1658/// camera controls as individual `.safetensors` LoRAs (LTX-2 README:72-83).
1659fn 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
1667/// Download a model and save its paths to config. Returns the updated config
1668/// and resolved model paths. Used by both the CLI `pull` command and the
1669/// server's auto-pull logic.
1670pub 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    // Utility models and hidden LTX-2 control adapters have no standalone
1684    // runtime config entry. The selected control is frozen onto the request
1685    // as a concrete LoRA path after this download completes.
1686    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    // Upscaler models have a single weights file (no VAE, no encoders).
1693    // Download files and create a minimal config entry with the weights path.
1694    if manifest.is_upscaler() {
1695        pull_model_files_only(manifest, opts).await?;
1696
1697        // Resolve the weights path from the manifest storage path
1698        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    // Auto-set default_model if no config existed before
1726    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
1738/// Download a model and save its paths to config, reporting progress via callback.
1739/// Same as `pull_and_configure` but uses a callback instead of indicatif bars.
1740pub 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
1748/// Download a model with a request-scoped Hugging Face token. The explicit
1749/// token takes precedence over environment and token-file credentials for this
1750/// call only, without expanding the stable public [`PullOptions`] struct.
1751pub 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    // Utility models and hidden LTX-2 control adapters have no standalone
1767    // runtime config entry. The selected control is frozen onto the request
1768    // as a concrete LoRA path after this download completes.
1769    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    // Upscaler models: download files, create minimal config with weights path.
1777    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
1820// ── Civitai token resolution ────────────────────────────────────────────────
1821
1822/// Resolve `CIVITAI_TOKEN` from the environment. Mirrors `resolve_hf_token`'s
1823/// shape, but Civitai has no token-file convention — just the env var. An
1824/// empty / whitespace-only env var resolves to `None` so a stale shell can't
1825/// silently send blank `Authorization: Bearer ` headers.
1826pub 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
1837/// Build the [`RecipeAuth`] required for a Civitai-gated recipe. Returns
1838/// [`DownloadError::MissingCivitaiToken`] when no token is set; the error
1839/// message names the env var so the CLI/server can surface a clear remediation.
1840pub 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
1847// ── Companion presence helpers ──────────────────────────────────────────────
1848//
1849// Civitai single-file checkpoints ship without their text encoders / VAE,
1850// so the catalog scanner records `companions: ["clip-l", "sdxl-vae", ...]`
1851// on those entries. Both server (`POST /api/catalog/:id/download`) and CLI
1852// (`mold pull cv:<id>`) need to enqueue/pull missing companions before the
1853// primary entry. The on-disk presence check + name-resolution loop lives
1854// here so they share one implementation; the server's
1855// `enqueue_missing_companions` consumes this through `DownloadQueue`, the
1856// CLI through `pull_and_configure_with_callback`.
1857
1858/// True when every file the companion's synthetic manifest declares is
1859/// present under `models_dir` AND no `.pulling` marker for the manifest's
1860/// canonical name exists. A leftover marker means a previous pull was
1861/// interrupted and the on-disk content can't be trusted yet.
1862pub 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
1887/// True iff the recipe file at `dest` should be considered already
1888/// placed — used by both the catalog API's `installed: bool` predicate
1889/// AND the `fetch_recipe_inner` skip path so they cannot drift apart.
1890///
1891/// Acceptance rule:
1892/// - `sha256` declared → `.sha256-verified` marker is the sole criterion.
1893///   The marker is written only after cryptographic verification at download
1894///   time, so it is more authoritative than `size_bytes` (which can be stale
1895///   in the catalog DB when a model is re-uploaded under the same sha256 with
1896///   a different compressed size).  A file at the exact declared size but
1897///   without the marker is still rejected.
1898/// - `sha256` absent, `size_bytes` known → on-disk length must equal declared.
1899/// - Neither declared → marker is the only attestation; require it.
1900fn 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
1916/// True iff every file in the recipe is present at its declared size (or,
1917/// when the recipe omits the size, has a `.sha256-verified` marker from
1918/// a prior verified pull) AND no `.pulling` marker for the catalog id is
1919/// present. Used by the catalog API to set `installed: bool` on each
1920/// wire entry so the SPA can hide the Download button and show Repair
1921/// instead.
1922///
1923/// Empty file slice returns `false` — callers (see `catalog_row_to_wire`
1924/// in mold-server) use this for Civitai-style recipe rows. HF rows
1925/// without a recipe go through `Config::manifest_model_is_downloaded`
1926/// instead, so an empty input here means "no recipe to walk" and we
1927/// refuse to claim install.
1928///
1929/// `id` is the catalog id (`cv:1234` / `hf:author/name`) — same string
1930/// the recipe-pull path uses to derive its marker and subdir name.
1931pub 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
1948/// Parse a `Vec<String>` of companion names out of `companions_json` and
1949/// return the ones that (a) resolve to a known synthetic manifest and (b)
1950/// aren't already fully present under `models_dir`.
1951///
1952/// Order is preserved — callers depend on companion-first ordering. Unknown
1953/// companions (no synthetic manifest in this build) are silently skipped:
1954/// catalog scanners may ship new canonical names ahead of the binary, and
1955/// surfacing those as errors would break catalog rows older builds can
1956/// never satisfy. `None` / unparsable JSON returns an empty vec.
1957pub 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
1971/// `Vec<String>`-shaped variant for callers that already have a typed
1972/// companion list (e.g. live-fetched `CatalogEntry::companions`).
1973pub 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// ── Recipe-driven downloads (Civitai single-file checkpoints) ───────────────
1995//
1996// Catalog rows for `cv:<id>` entries carry a `download_recipe.files` list of
1997// `(url, dest, sha256, size_bytes)` tuples that the manifest path can't
1998// express (the manifest assumes HF repos). The recipe fetcher lives here so
1999// it can share `compute_sha256`, `verify_sha256`, and the `.pulling` marker
2000// lifecycle with the manifest path. mold-core takes a plain
2001// `&[RecipeFetchFile]` slice + `RecipeAuth`; CLI/server callers translate
2002// from `mold_catalog::DownloadRecipe` at the boundary so this crate stays
2003// catalog-free (`mold-catalog` already depends on `mold-core`).
2004
2005/// Plain recipe-input shape for [`fetch_recipe`]. Callers translate from
2006/// `mold_catalog::DownloadRecipe` to a slice of these.
2007#[derive(Debug, Clone)]
2008pub struct RecipeFetchFile<'a> {
2009    /// HTTP URL the file is fetched from.
2010    pub url: &'a str,
2011    /// Destination path relative to the per-recipe subdirectory
2012    /// (`<models_dir>/<sanitized-id>/`). May contain forward slashes for
2013    /// nested layouts; `..` and absolute paths are rejected.
2014    pub dest: &'a str,
2015    /// Optional SHA-256 hex digest. Verified after download when present.
2016    pub sha256: Option<&'a str>,
2017    /// Optional declared file size, used for progress reporting before the
2018    /// `Content-Length` header arrives.
2019    pub size_bytes: Option<u64>,
2020}
2021
2022/// Authentication required for the recipe's URLs.
2023#[derive(Debug, Clone, PartialEq, Eq)]
2024pub enum RecipeAuth {
2025    /// No bearer token required.
2026    None,
2027    /// Send the given bearer token as `Authorization: Bearer <token>`.
2028    /// Used for Civitai (`needs_token: Civitai`); callers resolve the
2029    /// token from `CIVITAI_TOKEN` / config before calling.
2030    Bearer(String),
2031}
2032
2033/// Sanitize a catalog id (e.g. `cv:618692`) into a filesystem-safe subdir
2034/// name (`cv-618692`). Mirrors the manifest path's `replace(':', "-")`
2035/// rule so both sides land under the same models-dir subtree convention.
2036pub fn sanitize_recipe_id(id: &str) -> String {
2037    id.replace(':', "-")
2038}
2039
2040/// Verify that a recipe `dest` stays under the per-recipe subdir. Rejects
2041/// absolute paths and any segment that traverses upward (`..`). Returns the
2042/// resolved per-file path under `subdir_root` on success.
2043fn 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            // ParentDir / Prefix / RootDir / CurDir all escape or are
2054            // pointless. CurDir (`./`) is harmless but signals a malformed
2055            // recipe — reject for consistency.
2056            _ => {
2057                return Err(DownloadError::RecipePathTraversal {
2058                    dest: dest.to_string(),
2059                });
2060            }
2061        }
2062    }
2063    Ok(subdir_root.join(candidate))
2064}
2065
2066/// Fetch a recipe-driven download. Writes each file under
2067/// `models_dir/<sanitized-id>/<dest>`, verifies SHA-256 when present, and
2068/// manages the `.pulling` marker lifecycle.
2069///
2070/// The marker is written before the first byte and removed only after every
2071/// file has been integrity-checked. On any error the marker is removed
2072/// best-effort so callers can retry; partial files are NOT cleaned up here
2073/// (callers wire that into their failure path the same way the manifest
2074/// path does, via `cleanup_partials_in_dir`).
2075pub 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    // The recipe path is a low-level public ingress used by both the CLI and
2084    // server. Apply the activation policy before deriving paths, creating the
2085    // recipe directory/marker, constructing an HTTP client, or reporting any
2086    // progress so callers cannot bypass a catalog-level gate with an opaque id.
2087    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    // Pre-flight: validate every dest before touching the network. A bad
2097    // `..` in any file aborts the whole recipe with no side-effects.
2098    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    // Marker removed on success and best-effort on error. Cleanup of
2123    // partial files is the caller's responsibility (matches manifest path).
2124    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        // Idempotency: skip the HTTP fetch when the file is already on disk
2159        // with the declared size, or (when no size is declared) when the
2160        // post-download .sha256-verified marker is present from a prior
2161        // run. Mirrors `is_already_placed` from the manifest path so a
2162        // recipe re-pull (Repair, double-clicked Download, retry-after-
2163        // partial-companion-failure) costs zero bytes when nothing's missing.
2164        //
2165        // The acceptance rule is centralized in `recipe_file_is_placed` so
2166        // that this skip path and `catalog_entry_installed` (the catalog
2167        // API's `installed: bool` predicate) cannot drift apart — otherwise
2168        // the SPA's Repair button would silently re-pull a model the
2169        // predicate just claimed was installed.
2170        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        // Throttle FileProgress to once per RECIPE_PROGRESS_THROTTLE_MS so SSE
2248        // subscribers and reactive UIs aren't drowned in chunk-rate events
2249        // (a multi-GB Civitai pull emits hundreds of thousands of chunks).
2250        // Mirrors the throttle in the manifest-pull `CallbackProgress::update`.
2251        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        // Final progress emit so the file's last few chunks aren't swallowed
2288        // by the throttle (FileDone fires below, but it doesn't carry the
2289        // intermediate bytes_downloaded value — drawers that key off
2290        // FileProgress for their byte counter would otherwise stall short).
2291        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 file handle so the SHA-256 read sees a flushed file.
2305        drop(out);
2306
2307        // Hash-and-mark on success. Mirror of the manifest-pull path: when
2308        // the recipe declares an expected hash we compare and bail on
2309        // mismatch; either way we end up writing the `.sha256-verified`
2310        // marker so `Config::manifest_files_exist` recognises this file
2311        // as a positively-attested install (not just "exists on disk").
2312        // Skipped under `skip_verify` — the user has explicitly asked us
2313        // not to read the file, so we have nothing to attest.
2314        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        // max_len < 8 returns unchanged to avoid degenerate "..." output
2446        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    /// Mutex to serialize tests that mutate `HF_TOKEN` — `set_var`/`remove_var`
2624    /// are process-global and not thread-safe, so parallel tests race.
2625    static HF_TOKEN_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
2626
2627    // ---------------------------------------------------------------------
2628    // FLUX city96-format GGUF pull-time warning tests.
2629    // `gguf_header_contains_tensor` just does a bounded substring scan after
2630    // validating the "GGUF" magic — no real GGUF parsing — so tests can write
2631    // synthetic files that only satisfy those two properties.
2632    // `flux_reference_warning` is pure over (manifest, models_dir), so we can
2633    // drive every branch without touching process-global state.
2634    // ---------------------------------------------------------------------
2635
2636    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        // Pad a couple hundred bytes of synthetic header bytes, then include
2643        // every tensor name as a plain UTF-8 substring so the scanner finds it.
2644        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        // Non-GGUF filename is ignored even when everything else matches.
2772        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        // Transformer file not written — function should silently return None.
2782        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        // A "complete" GGUF has img_in.weight, so no patching needed.
2792        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        // city96-format: diffusion blocks but no embedding layers.
2804        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        // Schnell target: message accepts flux-schnell OR flux-dev as reference.
2828        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        // Place a fake "downloaded" complete flux-dev:q8 alongside.
2842        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        // Regression: schnell has img_in but not guidance_in. Pre-fix, it was
2865        // accepted as a reference; then ensure_gguf_embeddings failed mid-patch.
2866        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        // Drop a schnell GGUF that looks "valid" (has img_in, lacks guidance_in).
2872        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        // Restore before asserting so we don't leak on panic
2898        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        // Restore before asserting
2912        match &original {
2913            Some(v) => std::env::set_var("HF_TOKEN", v),
2914            None => std::env::remove_var("HF_TOKEN"),
2915        }
2916        // Should fall through to file-based token (which may or may not exist)
2917        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        // SHA-256 of "hello world"
2941        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    /// Civitai's API returns SHA-256 hashes in uppercase hex
2958    /// (`DD08FA32...`), while `compute_sha256` formats with `{:x}` so it
2959    /// produces lowercase. A literal string comparison treats these as
2960    /// distinct, so every Civitai pull bailed out with a "mismatch" even
2961    /// when the file was bit-identical to what was advertised. The
2962    /// verifier must be hex-case-insensitive.
2963    #[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        // File should be deleted
3011        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        // File should still exist
3035        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    // ── .sha256-verified marker helpers (B1) ─────────────────────────────
3061
3062    #[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        // Second call must not fail.
3110        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    // ── verify_file_integrity now writes a marker on success (B2) ────────
3127
3128    #[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        // The corrupted file is removed by verify_file_integrity, but more
3193        // importantly: there must be no marker pointing at the bad bytes.
3194        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        // skip_verify = true: we don't know the file is good, so no marker.
3217        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        // Write
3232        std::fs::write(&marker, "test-model:q8").unwrap();
3233        assert!(marker.exists());
3234
3235        // Remove
3236        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    // ── Civitai token resolution (round 3) ──────────────────────────────
3258
3259    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    // ── Companion presence helpers (round 2) ────────────────────────────
3338
3339    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        // "clip-l" is real; "future-encoder-9000" doesn't exist.
3416        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        // clip-l is staged, sdxl-vae is not.
3438        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    // ── Recipe fetcher (round 1) ────────────────────────────────────────
3465
3466    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        // SHA-256 of "hello world"
3621        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        // Wrong digest — file content is "hello".
3665        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        // Corrupted file should be deleted.
3690        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            // First call serves the body; any second call is an unexpected re-fetch.
3751            .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        // Second call must skip the HTTP fetch entirely because the file is on
3776        // disk with the declared size.
3777        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        // wiremock's `.expect(1)` is verified on `MockServer::drop`; explicit
3789        // verify here gives a clearer failure message at the assertion site.
3790        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 intentionally None — fall through to marker check.
3815            size_bytes: None,
3816        }];
3817
3818        // First call writes the marker via the existing post-download codepath.
3819        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        // Confirm marker is in place (sanity check for the test setup).
3831        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        // SHA-256 of "correct"
3860        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            // The pre-staged file has the right size but no marker, so the
3865            // skip path must refuse it and re-fetch exactly once.
3866            .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        // Pre-stage a file at the declared size — but NO marker, and bytes
3875        // don't actually match the declared sha256. A size-only skip would
3876        // accept this; the tightened predicate must not.
3877        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        // After re-fetch the bytes match the server response and the marker exists.
3900        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        // b.safetensors is intentionally missing.
3936
3937        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        // No marker, no declared size — refuse to claim install.
4037
4038        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        // Mutation guard: pins == (not >=) for the size comparison. A
4148        // 99-byte file declared as 5 bytes is just as wrong as a 5-byte file
4149        // declared as 99 bytes — the existing `_size_mismatch` test only
4150        // exercises the file-too-small direction.
4151        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        // Cross-consistency: catalog_entry_installed and the inline skip
4178        // path inside fetch_recipe_inner must agree on what "placed" means.
4179        // A file at the right size with no marker — and a declared sha256
4180        // — would otherwise be reported `installed=true` by the catalog
4181        // API while the fetch path re-downloads it on Repair. Pins the
4182        // shared `recipe_file_is_placed` rule.
4183        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        // No marker.
4188
4189        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        // Regression guard: catalog DB can have stale size_bytes (e.g. model
4206        // re-uploaded with same sha256 but different compressed size).  When a
4207        // sha256 is declared and the marker exists, the file is verified —
4208        // reject it only on size would cause installed models to disappear from
4209        // the settings modal.
4210        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        // File is 5 bytes, but we'll declare size as 99 (stale) in the recipe.
4215        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), // stale — actual file is 5 bytes
4227        }];
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            // Size mismatch must trigger the fetch.
4246            .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        // Pre-stage a wrong-size file (4 bytes vs. the recipe's declared 7).
4255        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        // File should now match the server response.
4277        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        // Sanity: nothing should have been created outside the per-id subdir.
4308        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}