1use crate::error::{EnvironmentError, ProviderError, Result, UserError};
4use futures_util::StreamExt;
5use indicatif::{ProgressBar, ProgressStyle};
6use sha2::{Digest, Sha256};
7use std::fs::{self, File, OpenOptions};
8use std::io::{Read, Write};
9use std::path::{Path, PathBuf};
10use std::sync::Arc;
11
12const HF_BASE: &str = "https://huggingface.co/ggerganov/whisper.cpp/resolve/main";
16pub const ARTIFACT_MANIFEST_VERSION: &str = "1";
18pub const ARTIFACT_MANIFEST_SOURCE: &str = "aurum-builtin-review";
20
21#[derive(Debug, Clone, Copy)]
23pub struct ModelInfo {
24 pub name: &'static str,
25 pub filename: &'static str,
26 pub approx_bytes: u64,
28 pub notes: &'static str,
30}
31
32pub const MODELS: &[ModelInfo] = &[
34 ModelInfo {
36 name: "tiny",
37 filename: "ggml-tiny.bin",
38 approx_bytes: 75_000_000,
39 notes: "fastest full-precision",
40 },
41 ModelInfo {
42 name: "tiny-q5_1",
43 filename: "ggml-tiny-q5_1.bin",
44 approx_bytes: 32_000_000,
45 notes: "quantized ~32MB — best first-run trial",
46 },
47 ModelInfo {
48 name: "tiny-q8_0",
49 filename: "ggml-tiny-q8_0.bin",
50 approx_bytes: 44_000_000,
51 notes: "quantized",
52 },
53 ModelInfo {
54 name: "tiny.en",
55 filename: "ggml-tiny.en.bin",
56 approx_bytes: 75_000_000,
57 notes: "english-only",
58 },
59 ModelInfo {
60 name: "tiny.en-q5_1",
61 filename: "ggml-tiny.en-q5_1.bin",
62 approx_bytes: 32_000_000,
63 notes: "english-only quantized",
64 },
65 ModelInfo {
67 name: "base",
68 filename: "ggml-base.bin",
69 approx_bytes: 142_000_000,
70 notes: "default full-precision",
71 },
72 ModelInfo {
73 name: "base-q5_1",
74 filename: "ggml-base-q5_1.bin",
75 approx_bytes: 60_000_000,
76 notes: "quantized ~60MB",
77 },
78 ModelInfo {
79 name: "base-q8_0",
80 filename: "ggml-base-q8_0.bin",
81 approx_bytes: 82_000_000,
82 notes: "quantized",
83 },
84 ModelInfo {
85 name: "base.en",
86 filename: "ggml-base.en.bin",
87 approx_bytes: 142_000_000,
88 notes: "english-only",
89 },
90 ModelInfo {
91 name: "base.en-q5_1",
92 filename: "ggml-base.en-q5_1.bin",
93 approx_bytes: 60_000_000,
94 notes: "english-only quantized",
95 },
96 ModelInfo {
98 name: "small",
99 filename: "ggml-small.bin",
100 approx_bytes: 466_000_000,
101 notes: "higher accuracy",
102 },
103 ModelInfo {
104 name: "small-q5_1",
105 filename: "ggml-small-q5_1.bin",
106 approx_bytes: 190_000_000,
107 notes: "quantized",
108 },
109 ModelInfo {
110 name: "small-q8_0",
111 filename: "ggml-small-q8_0.bin",
112 approx_bytes: 264_000_000,
113 notes: "quantized",
114 },
115 ModelInfo {
116 name: "small.en",
117 filename: "ggml-small.en.bin",
118 approx_bytes: 466_000_000,
119 notes: "english-only",
120 },
121 ModelInfo {
122 name: "small.en-q5_1",
123 filename: "ggml-small.en-q5_1.bin",
124 approx_bytes: 190_000_000,
125 notes: "english-only quantized",
126 },
127 ModelInfo {
129 name: "medium",
130 filename: "ggml-medium.bin",
131 approx_bytes: 1_500_000_000,
132 notes: "large download",
133 },
134 ModelInfo {
135 name: "medium.en",
136 filename: "ggml-medium.en.bin",
137 approx_bytes: 1_500_000_000,
138 notes: "english-only",
139 },
140 ModelInfo {
142 name: "large-v3",
143 filename: "ggml-large-v3.bin",
144 approx_bytes: 3_100_000_000,
145 notes: "highest quality",
146 },
147 ModelInfo {
148 name: "large-v3-q5_0",
149 filename: "ggml-large-v3-q5_0.bin",
150 approx_bytes: 1_080_000_000,
151 notes: "experimental — not recommended (degeneration risk; prefer large-v3-turbo)",
152 },
153 ModelInfo {
154 name: "large",
155 filename: "ggml-large-v3.bin",
156 approx_bytes: 3_100_000_000,
157 notes: "alias of large-v3",
158 },
159 ModelInfo {
160 name: "large-v3-turbo",
161 filename: "ggml-large-v3-turbo.bin",
162 approx_bytes: 1_600_000_000,
163 notes: "fast large",
164 },
165 ModelInfo {
166 name: "large-v3-turbo-q5_0",
167 filename: "ggml-large-v3-turbo-q5_0.bin",
168 approx_bytes: 574_000_000,
169 notes: "quantized turbo",
170 },
171 ModelInfo {
172 name: "turbo",
173 filename: "ggml-large-v3-turbo.bin",
174 approx_bytes: 1_600_000_000,
175 notes: "alias of large-v3-turbo",
176 },
177 ModelInfo {
178 name: "turbo-q5_0",
179 filename: "ggml-large-v3-turbo-q5_0.bin",
180 approx_bytes: 574_000_000,
181 notes: "alias of large-v3-turbo-q5_0",
182 },
183];
184
185pub fn available_model_names() -> String {
187 list_canonical_models()
188 .iter()
189 .map(|m| m.name)
190 .collect::<Vec<_>>()
191 .join(", ")
192}
193
194fn list_canonical_models() -> Vec<&'static ModelInfo> {
195 MODELS
196 .iter()
197 .filter(|m| !matches!(m.name, "large" | "turbo" | "turbo-q5_0"))
198 .collect()
199}
200
201pub fn lookup_model(name: &str) -> Result<&'static ModelInfo> {
202 let key = name.trim().to_ascii_lowercase();
203 MODELS.iter().find(|m| m.name == key).ok_or_else(|| {
204 UserError::InvalidModel {
205 model: name.to_string(),
206 available: available_model_names(),
207 }
208 .into()
209 })
210}
211
212pub fn models_dir(cache_dir: &Path) -> PathBuf {
214 cache_dir.join("models")
215}
216
217pub fn model_path(cache_dir: &Path, info: &ModelInfo) -> PathBuf {
219 models_dir(cache_dir).join(info.filename)
220}
221
222#[derive(Debug, Clone)]
224pub struct ModelStatus {
225 pub info: &'static ModelInfo,
226 pub cached: bool,
227 pub path: PathBuf,
228 pub size_bytes: Option<u64>,
229}
230
231pub fn list_models(cache_dir: &Path) -> Vec<ModelStatus> {
233 list_canonical_models()
234 .into_iter()
235 .map(|info| {
236 let path = model_path(cache_dir, info);
237 let (cached, size_bytes) = match fs::metadata(&path) {
238 Ok(m) if m.len() > 1_000_000 => (true, Some(m.len())),
239 _ => (false, None),
240 };
241 ModelStatus {
242 info,
243 cached,
244 path,
245 size_bytes,
246 }
247 })
248 .collect()
249}
250
251pub fn format_model_list(cache_dir: &Path) -> String {
253 let rows = list_models(cache_dir);
254 let mut out = String::from("Local whisper.cpp models (cache: ");
255 out.push_str(&models_dir(cache_dir).display().to_string());
256 out.push_str(")\n\n");
257 out.push_str(&format!(
258 "{:<22} {:>10} {:<8} {:<12} {}\n",
259 "NAME", "SIZE", "STATUS", "TIER", "NOTES"
260 ));
261 out.push_str(&format!(
262 "{:<22} {:>10} {:<8} {:<12} {}\n",
263 "----", "----", "------", "----", "-----"
264 ));
265 for row in rows {
266 let size = format_bytes(row.info.approx_bytes);
267 let status = if row.cached { "cached" } else { "—" };
268 let tier = match model_support_tier(row.info.name) {
269 ModelSupportTier::Supported => "supported",
270 ModelSupportTier::Experimental => "experimental",
271 };
272 out.push_str(&format!(
273 "{:<22} {:>10} {:<8} {:<12} {}\n",
274 row.info.name, size, status, tier, row.info.notes
275 ));
276 }
277 out.push_str(
278 "\nTip: first run downloads the selected model. Try `tiny-q5_1` (~32 MB) for a quick trial.\n",
279 );
280 out.push_str("Default model: `base` (~142 MB). Use --model <name> to choose.\n");
281 out.push_str(
282 "Guidance (single-clip dogfood, not formal WER): English lecture quality often \
283 favors `small.en` or `large-v3-turbo`; prefer `.en` variants for English-only audio. \
284 Avoid `large-v3-q5_0` for production (experimental).\n",
285 );
286 out
287}
288
289fn format_bytes(n: u64) -> String {
290 const KB: f64 = 1024.0;
291 const MB: f64 = KB * 1024.0;
292 const GB: f64 = MB * 1024.0;
293 let n = n as f64;
294 if n >= GB {
295 format!("{:.1} GB", n / GB)
296 } else if n >= MB {
297 format!("{:.0} MB", n / MB)
298 } else {
299 format!("{:.0} KB", n / KB)
300 }
301}
302
303#[derive(Debug, Clone)]
305pub struct DownloadProgress {
306 pub model: String,
307 pub downloaded_bytes: u64,
308 pub total_bytes: u64,
309}
310
311impl DownloadProgress {
312 pub fn fraction(&self) -> Option<f64> {
313 if self.total_bytes == 0 {
314 None
315 } else {
316 Some((self.downloaded_bytes as f64 / self.total_bytes as f64).clamp(0.0, 1.0))
317 }
318 }
319}
320
321pub type DownloadProgressCallback = Arc<dyn Fn(DownloadProgress) + Send + Sync>;
323
324#[derive(Clone, Default)]
326pub struct EnsureModelOptions {
327 pub local_only: bool,
329 pub show_progress: bool,
331 pub on_progress: Option<DownloadProgressCallback>,
333}
334
335impl EnsureModelOptions {
336 pub fn new() -> Self {
337 Self::default()
338 }
339
340 pub fn local_only(mut self, v: bool) -> Self {
341 self.local_only = v;
342 self
343 }
344
345 pub fn show_progress(mut self, v: bool) -> Self {
346 self.show_progress = v;
347 self
348 }
349
350 pub fn on_progress(mut self, cb: DownloadProgressCallback) -> Self {
351 self.on_progress = Some(cb);
352 self
353 }
354}
355
356pub fn is_model_cached(cache_dir: &Path, model_name: &str) -> bool {
358 let Ok(info) = lookup_model(model_name) else {
359 return false;
360 };
361 let path = model_path(cache_dir, info);
362 if !(path.exists()
363 && path
364 .metadata()
365 .map(|m| m.len() > 1_000_000)
366 .unwrap_or(false))
367 {
368 return false;
369 }
370 pinned_sha256(info.filename).is_some() && verify_model_basic(&path, info).is_ok()
372}
373
374pub async fn ensure_model(
376 cache_dir: &Path,
377 model_name: &str,
378 show_progress: bool,
379) -> Result<PathBuf> {
380 ensure_model_with_options(
381 cache_dir,
382 model_name,
383 EnsureModelOptions {
384 show_progress,
385 ..EnsureModelOptions::default()
386 },
387 )
388 .await
389}
390
391pub async fn ensure_model_with_options(
393 cache_dir: &Path,
394 model_name: &str,
395 opts: EnsureModelOptions,
396) -> Result<PathBuf> {
397 let info = lookup_model(model_name)?;
398 let path = model_path(cache_dir, info);
399
400 if path.exists()
401 && path
402 .metadata()
403 .map(|m| m.len() > 1_000_000)
404 .unwrap_or(false)
405 {
406 if pinned_sha256(info.filename).is_none() {
408 return Err(ProviderError::ModelDownload {
409 model: model_name.to_string(),
410 reason: format!(
411 "model `{}` has no reviewed SHA-256 pin — not available on the trusted path",
412 info.name
413 ),
414 }
415 .into());
416 }
417 if verify_model_basic(&path, info).is_ok() {
418 tracing::info!(model = info.name, path = %path.display(), "using cached model");
419 return Ok(path);
420 }
421 tracing::warn!(
422 model = info.name,
423 path = %path.display(),
424 "cached model failed integrity check; re-downloading"
425 );
426 let _ = fs::remove_file(&path);
427 }
428
429 if opts.local_only {
430 return Err(UserError::ModelNotCached {
431 model: model_name.to_string(),
432 }
433 .into());
434 }
435
436 fs::create_dir_all(models_dir(cache_dir)).map_err(|e| EnvironmentError::DirectoryAccess {
437 path: models_dir(cache_dir).display().to_string(),
438 reason: e.to_string(),
439 })?;
440
441 let lock_path = path.with_extension("bin.lock");
443 let lock_file = OpenOptions::new()
444 .create(true)
445 .read(true)
446 .write(true)
447 .truncate(false)
448 .open(&lock_path)
449 .map_err(|e| EnvironmentError::DirectoryAccess {
450 path: lock_path.display().to_string(),
451 reason: e.to_string(),
452 })?;
453
454 if opts.show_progress {
455 eprintln!("aurum: waiting for model download lock ({}) …", info.name);
456 }
457 lock_file.lock().map_err(|e| EnvironmentError::Other {
458 message: format!("failed to acquire model lock: {e}"),
459 })?;
460
461 if path.exists()
463 && path
464 .metadata()
465 .map(|m| m.len() > 1_000_000)
466 .unwrap_or(false)
467 && pinned_sha256(info.filename).is_some()
468 && verify_model_basic(&path, info).is_ok()
469 {
470 let _ = lock_file.unlock();
471 tracing::info!(model = info.name, "model appeared while waiting on lock");
472 return Ok(path);
473 }
474
475 if opts.local_only {
476 let _ = lock_file.unlock();
477 return Err(UserError::ModelNotCached {
478 model: model_name.to_string(),
479 }
480 .into());
481 }
482
483 if opts.show_progress {
484 eprintln!(
485 "aurum: downloading model `{}` ({}) — first run only …",
486 info.name,
487 format_bytes(info.approx_bytes)
488 );
489 }
490
491 let result = download_model(info, &path, opts.show_progress, opts.on_progress.as_ref()).await;
492 let _ = lock_file.unlock();
493 result?;
494 verify_model_basic(&path, info)?;
495 Ok(path)
496}
497
498#[allow(dead_code)]
501fn verify_cached_checksum(path: &Path) -> bool {
502 let checksum_path = path.with_extension("bin.sha256");
503 let Ok(contents) = fs::read_to_string(&checksum_path) else {
504 return false;
505 };
506 let Some(expected) = contents.split_whitespace().next() else {
507 return false;
508 };
509 verify_against_expected(path, expected)
510}
511
512async fn download_model(
513 info: &ModelInfo,
514 dest: &Path,
515 show_progress: bool,
516 on_progress: Option<&DownloadProgressCallback>,
517) -> Result<()> {
518 let url = format!("{HF_BASE}/{}?download=true", info.filename);
519 tracing::info!(model = info.name, %url, "downloading model");
520
521 let partial = dest.with_extension(format!(
522 "bin.partial.{}-{}",
523 std::process::id(),
524 std::time::SystemTime::now()
525 .duration_since(std::time::UNIX_EPOCH)
526 .map(|d| d.as_millis())
527 .unwrap_or(0)
528 ));
529 if partial.exists() {
530 let _ = fs::remove_file(&partial);
531 }
532
533 let client = reqwest::Client::builder()
537 .user_agent(concat!("aurum-core/", env!("CARGO_PKG_VERSION")))
538 .connect_timeout(std::time::Duration::from_secs(30))
539 .timeout(std::time::Duration::from_secs(30 * 60))
540 .redirect(reqwest::redirect::Policy::custom(hf_redirect_policy))
541 .build()
542 .map_err(|e| ProviderError::ModelDownload {
543 model: info.name.to_string(),
544 reason: format!("http client error: {e}"),
545 })?;
546
547 let response = client
548 .get(&url)
549 .send()
550 .await
551 .map_err(|e| ProviderError::ModelDownload {
552 model: info.name.to_string(),
553 reason: format!("request failed: {e}"),
554 })?;
555
556 if !response.status().is_success() {
557 return Err(ProviderError::ModelDownload {
558 model: info.name.to_string(),
559 reason: format!("HTTP {}", response.status()),
560 }
561 .into());
562 }
563
564 let total = response.content_length().unwrap_or(info.approx_bytes);
565
566 let pb = if show_progress {
567 let pb = ProgressBar::new(total);
568 pb.set_style(
569 ProgressStyle::with_template(
570 "{msg} [{bar:40.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})",
571 )
572 .unwrap_or_else(|_| ProgressStyle::default_bar())
573 .progress_chars("=>-"),
574 );
575 pb.set_message(format!("Downloading {}", info.name));
576 Some(pb)
577 } else {
578 None
579 };
580
581 let mut file = File::create(&partial).map_err(|e| EnvironmentError::DirectoryAccess {
582 path: partial.display().to_string(),
583 reason: e.to_string(),
584 })?;
585
586 let mut stream = response.bytes_stream();
587 let mut hasher = Sha256::new();
588 let mut downloaded: u64 = 0;
589
590 while let Some(chunk) = stream.next().await {
591 let chunk = chunk.map_err(|e| ProviderError::ModelDownload {
592 model: info.name.to_string(),
593 reason: format!("stream error: {e}"),
594 })?;
595 file.write_all(&chunk)
596 .map_err(|e| EnvironmentError::DiskSpace {
597 path: partial.display().to_string(),
598 reason: e.to_string(),
599 })?;
600 hasher.update(&chunk);
601 downloaded += chunk.len() as u64;
602 let max_allowed = crate::download::download_byte_cap(
604 info.approx_bytes,
605 pinned_exact_bytes(info.filename),
606 3,
607 );
608 if downloaded > max_allowed {
609 let _ = fs::remove_file(&partial);
610 return Err(ProviderError::ModelDownload {
611 model: info.name.to_string(),
612 reason: format!(
613 "download exceeded size cap ({downloaded} > {max_allowed} bytes) — aborting"
614 ),
615 }
616 .into());
617 }
618 if let Some(pb) = &pb {
619 pb.set_position(downloaded.min(total));
620 }
621 if let Some(cb) = on_progress {
622 cb(DownloadProgress {
623 model: info.name.to_string(),
624 downloaded_bytes: downloaded,
625 total_bytes: total,
626 });
627 }
628 }
629
630 file.flush().ok();
631 let _ = file.sync_all();
632 drop(file);
633
634 let digest = hex::encode(hasher.finalize());
635 tracing::info!(
636 model = info.name,
637 sha256 = %digest,
638 bytes = downloaded,
639 "model download complete (pre-publish)"
640 );
641
642 let Some(expected) = pinned_sha256(info.filename) else {
644 let _ = fs::remove_file(&partial);
645 return Err(ProviderError::ModelDownload {
646 model: info.name.to_string(),
647 reason: format!(
648 "no reviewed SHA-256 pin for {} — refusing to publish unauthenticated artifact",
649 info.filename
650 ),
651 }
652 .into());
653 };
654 if digest != expected {
655 let _ = fs::remove_file(&partial);
656 return Err(ProviderError::ModelDownload {
657 model: info.name.to_string(),
658 reason: format!(
659 "sha256 mismatch (got {digest}, expected {expected}) — refusing to publish"
660 ),
661 }
662 .into());
663 }
664 if let Some(exact) = pinned_exact_bytes(info.filename) {
665 if downloaded != exact {
666 let _ = fs::remove_file(&partial);
667 return Err(ProviderError::ModelDownload {
668 model: info.name.to_string(),
669 reason: format!(
670 "size mismatch (got {downloaded}, expected {exact}) — refusing to publish"
671 ),
672 }
673 .into());
674 }
675 }
676
677 if dest.exists() {
679 let _ = fs::remove_file(dest);
680 }
681 fs::rename(&partial, dest).map_err(|e| {
682 let _ = fs::remove_file(&partial);
683 EnvironmentError::DirectoryAccess {
684 path: dest.display().to_string(),
685 reason: e.to_string(),
686 }
687 })?;
688
689 if let Some(pb) = pb {
690 pb.finish_with_message(format!("Downloaded {} ({downloaded} bytes)", info.name));
691 }
692
693 let checksum_path = dest.with_extension("bin.sha256");
694 let _ = fs::write(&checksum_path, format!("{digest} {}\n", info.filename));
695
696 sweep_stale_partials(dest.parent().unwrap_or_else(|| Path::new(".")));
698
699 Ok(())
700}
701
702fn hf_redirect_policy(attempt: reqwest::redirect::Attempt<'_>) -> reqwest::redirect::Action {
704 let url = attempt.url();
705 let host = url.host_str().unwrap_or("").to_ascii_lowercase();
706 let allowed = host == "huggingface.co"
707 || host.ends_with(".huggingface.co")
708 || host.ends_with(".hf.co")
709 || host == "hf.co"
710 || host.ends_with(".cdn.hf.co");
711 if allowed && attempt.previous().len() < 8 {
712 attempt.follow()
713 } else {
714 attempt.stop()
715 }
716}
717
718pub fn pinned_sha256(filename: &str) -> Option<&'static str> {
724 match filename {
725 "ggml-tiny.bin" => Some("be07e048e1e599ad46341c8d2a135645097a538221678b7acdd1b1919c6e1b21"),
726 "ggml-tiny-q5_1.bin" => {
727 Some("818710568da3ca15689e31a743197b520007872ff9576237bda97bd1b469c3d7")
728 }
729 "ggml-tiny-q8_0.bin" => {
730 Some("c2085835d3f50733e2ff6e4b41ae8a2b8d8110461e18821b09a15c40c42d1cca")
731 }
732 "ggml-tiny.en.bin" => {
733 Some("921e4cf8686fdd993dcd081a5da5b6c365bfde1162e72b08d75ac75289920b1f")
734 }
735 "ggml-tiny.en-q5_1.bin" => {
736 Some("c77c5766f1cef09b6b7d47f21b546cbddd4157886b3b5d6d4f709e91e66c7c2b")
737 }
738 "ggml-base.bin" => Some("60ed5bc3dd14eea856493d334349b405782ddcaf0028d4b5df4088345fba2efe"),
739 "ggml-base-q5_1.bin" => {
740 Some("422f1ae452ade6f30a004d7e5c6a43195e4433bc370bf23fac9cc591f01a8898")
741 }
742 "ggml-base-q8_0.bin" => {
743 Some("c577b9a86e7e048a0b7eada054f4dd79a56bbfa911fbdacf900ac5b567cbb7d9")
744 }
745 "ggml-base.en.bin" => {
746 Some("a03779c86df3323075f5e796cb2ce5029f00ec8869eee3fdfb897afe36c6d002")
747 }
748 "ggml-base.en-q5_1.bin" => {
749 Some("4baf70dd0d7c4247ba2b81fafd9c01005ac77c2f9ef064e00dcf195d0e2fdd2f")
750 }
751 "ggml-small.bin" => {
752 Some("1be3a9b2063867b937e64e2ec7483364a79917e157fa98c5d94b5c1fffea987b")
753 }
754 "ggml-small-q5_1.bin" => {
755 Some("ae85e4a935d7a567bd102fe55afc16bb595bdb618e11b2fc7591bc08120411bb")
756 }
757 "ggml-small-q8_0.bin" => {
758 Some("49c8fb02b65e6049d5fa6c04f81f53b867b5ec9540406812c643f177317f779f")
759 }
760 "ggml-small.en.bin" => {
761 Some("c6138d6d58ecc8322097e0f987c32f1be8bb0a18532a3f88f734d1bbf9c41e5d")
762 }
763 "ggml-small.en-q5_1.bin" => {
764 Some("bfdff4894dcb76bbf647d56263ea2a96645423f1669176f4844a1bf8e478ad30")
765 }
766 "ggml-medium.bin" => {
767 Some("6c14d5adee5f86394037b4e4e8b59f1673b6cee10e3cf0b11bbdbee79c156208")
768 }
769 "ggml-medium.en.bin" => {
770 Some("cc37e93478338ec7700281a7ac30a10128929eb8f427dda2e865faa8f6da4356")
771 }
772 "ggml-large-v3.bin" => {
773 Some("64d182b440b98d5203c4f9bd541544d84c605196c4f7b845dfa11fb23594d1e2")
774 }
775 "ggml-large-v3-q5_0.bin" => {
776 Some("d75795ecff3f83b5faa89d1900604ad8c780abd5739fae406de19f23ecd98ad1")
777 }
778 "ggml-large-v3-turbo.bin" => {
779 Some("1fc70f774d38eb169993ac391eea357ef47c88757ef72ee5943879b7e8e2bc69")
780 }
781 "ggml-large-v3-turbo-q5_0.bin" => {
782 Some("394221709cd5ad1f40c46e6031ca61bce88931e6e088c188294c6d5a55ffa7e2")
783 }
784 _ => None,
785 }
786}
787
788#[derive(Debug, Clone, Copy, PartialEq, Eq)]
790pub enum ModelSupportTier {
791 Supported,
793 Experimental,
795}
796
797pub fn model_support_tier(name: &str) -> ModelSupportTier {
799 match name {
800 "large-v3-q5_0" => ModelSupportTier::Experimental,
801 _ => ModelSupportTier::Supported,
802 }
803}
804
805pub fn pinned_exact_bytes(filename: &str) -> Option<u64> {
807 match filename {
808 "ggml-tiny.bin" => Some(77_691_713),
809 "ggml-tiny-q5_1.bin" => Some(32_152_673),
810 "ggml-tiny-q8_0.bin" => Some(43_537_433),
811 "ggml-tiny.en.bin" => Some(77_704_715),
812 "ggml-tiny.en-q5_1.bin" => Some(32_166_155),
813 "ggml-base.bin" => Some(147_951_465),
814 "ggml-base-q5_1.bin" => Some(59_707_625),
815 "ggml-base-q8_0.bin" => Some(81_768_585),
816 "ggml-base.en.bin" => Some(147_964_211),
817 "ggml-base.en-q5_1.bin" => Some(59_721_011),
818 "ggml-small.bin" => Some(487_601_967),
819 "ggml-small-q5_1.bin" => Some(190_085_487),
820 "ggml-small-q8_0.bin" => Some(264_464_607),
821 "ggml-small.en.bin" => Some(487_614_201),
822 "ggml-small.en-q5_1.bin" => Some(190_098_681),
823 "ggml-medium.bin" => Some(1_533_763_059),
824 "ggml-medium.en.bin" => Some(1_533_774_781),
825 "ggml-large-v3.bin" => Some(3_095_033_483),
826 "ggml-large-v3-q5_0.bin" => Some(1_081_140_203),
827 "ggml-large-v3-turbo.bin" => Some(1_624_555_275),
828 "ggml-large-v3-turbo-q5_0.bin" => Some(574_041_195),
829 _ => None,
830 }
831}
832
833pub fn artifact_manifest_json(info: &ModelInfo) -> serde_json::Value {
835 serde_json::json!({
836 "manifest_version": ARTIFACT_MANIFEST_VERSION,
837 "source": ARTIFACT_MANIFEST_SOURCE,
838 "id": info.name,
839 "filename": info.filename,
840 "approx_bytes": info.approx_bytes,
841 "exact_bytes": pinned_exact_bytes(info.filename),
842 "sha256": pinned_sha256(info.filename),
843 "license": "MIT (whisper.cpp weights via OpenAI Whisper terms)",
844 "family": "whisper",
845 "download_url_template": format!("{HF_BASE}/{}", info.filename),
846 })
847}
848
849fn sweep_stale_partials(dir: &Path) {
850 let Ok(entries) = fs::read_dir(dir) else {
851 return;
852 };
853 let stale_after = std::time::Duration::from_secs(6 * 3600);
854 let now = std::time::SystemTime::now();
855 for ent in entries.flatten() {
856 let name = ent.file_name();
857 let name = name.to_string_lossy();
858 if !name.contains(".bin.partial.") {
860 continue;
861 }
862 let Ok(meta) = ent.metadata() else {
863 continue;
864 };
865 let Ok(modified) = meta.modified() else {
866 continue;
867 };
868 if now.duration_since(modified).unwrap_or_default() > stale_after {
869 let _ = fs::remove_file(ent.path());
870 }
871 }
872}
873
874pub fn ensure_model_verified_local(path: &Path, info: &ModelInfo) -> Result<()> {
876 verify_model_basic(path, info)
877}
878
879fn verify_model_basic(path: &Path, info: &ModelInfo) -> Result<()> {
881 let meta = fs::metadata(path).map_err(|e| ProviderError::ModelDownload {
882 model: info.name.to_string(),
883 reason: format!("missing after download: {e}"),
884 })?;
885
886 if meta.len() < 1_000_000 {
887 return Err(ProviderError::ModelDownload {
889 model: info.name.to_string(),
890 reason: format!(
891 "model file is only {} bytes — likely truncated or HTML error page",
892 meta.len()
893 ),
894 }
895 .into());
896 }
897
898 let mut hdr = [0u8; 4];
899 let mut f = File::open(path).map_err(|e| ProviderError::ModelDownload {
900 model: info.name.to_string(),
901 reason: e.to_string(),
902 })?;
903 f.read_exact(&mut hdr)
904 .map_err(|e| ProviderError::ModelDownload {
905 model: info.name.to_string(),
906 reason: format!("cannot read header: {e}"),
907 })?;
908
909 let magic_ok = matches!(
910 &hdr,
911 b"ggml"
912 | b"lmgg"
913 | b"ggmf"
914 | b"fmgg"
915 | b"ggjt"
916 | b"tjgg"
917 | b"ggjf"
918 | b"fjgg"
919 | b"gguf"
920 | b"fugg"
921 | b"GGUF"
922 );
923
924 if !magic_ok {
925 return Err(ProviderError::ModelDownload {
927 model: info.name.to_string(),
928 reason: format!(
929 "model header {:?} is not a recognized ggml/gguf magic — refusing to use file",
930 hdr
931 ),
932 }
933 .into());
934 }
935
936 let Some(expected) = pinned_sha256(info.filename) else {
938 return Err(ProviderError::ModelDownload {
939 model: info.name.to_string(),
940 reason: format!(
941 "no reviewed SHA-256 pin for {} — not trusted",
942 info.filename
943 ),
944 }
945 .into());
946 };
947 if !verify_against_expected(path, expected) {
948 return Err(ProviderError::ModelDownload {
949 model: info.name.to_string(),
950 reason: format!(
951 "cached model failed pinned sha256 check ({expected}); \
952 run `aurum cache verify` / quarantine repair"
953 ),
954 }
955 .into());
956 }
957 if let Some(exact) = pinned_exact_bytes(info.filename) {
958 if meta.len() != exact {
959 return Err(ProviderError::ModelDownload {
960 model: info.name.to_string(),
961 reason: format!(
962 "cached model size mismatch (got {}, expected {exact})",
963 meta.len()
964 ),
965 }
966 .into());
967 }
968 }
969
970 Ok(())
971}
972
973fn verify_against_expected(path: &Path, expected: &str) -> bool {
974 let Ok(mut file) = File::open(path) else {
975 return false;
976 };
977 let mut hasher = Sha256::new();
978 let mut buf = [0u8; 64 * 1024];
979 loop {
980 match file.read(&mut buf) {
981 Ok(0) => break,
982 Ok(n) => hasher.update(&buf[..n]),
983 Err(_) => return false,
984 }
985 }
986 hex::encode(hasher.finalize()) == expected
987}
988
989#[cfg(test)]
990mod tests {
991 use super::*;
992
993 #[test]
994 fn lookup_known_models() {
995 assert_eq!(lookup_model("base").unwrap().filename, "ggml-base.bin");
996 assert_eq!(
997 lookup_model("tiny-q5_1").unwrap().filename,
998 "ggml-tiny-q5_1.bin"
999 );
1000 assert_eq!(
1001 lookup_model("large-v3-turbo").unwrap().filename,
1002 "ggml-large-v3-turbo.bin"
1003 );
1004 assert_eq!(lookup_model("turbo").unwrap().name, "turbo");
1005 assert!(lookup_model("nope").is_err());
1006 }
1007
1008 #[test]
1009 fn model_path_joins() {
1010 let p = model_path(Path::new("/tmp/cache"), lookup_model("tiny").unwrap());
1011 assert_eq!(p, PathBuf::from("/tmp/cache/models/ggml-tiny.bin"));
1012 }
1013
1014 #[test]
1015 fn list_includes_quantized() {
1016 let list = format_model_list(Path::new("/tmp/aurum-cache-test"));
1017 assert!(list.contains("tiny-q5_1"));
1018 assert!(list.contains("base-q5_1"));
1019 assert!(list.contains("first run"));
1020 }
1021
1022 #[test]
1023 fn every_catalogue_file_has_exact_size_metadata() {
1024 let mut missing = Vec::new();
1026 for m in MODELS {
1027 if pinned_exact_bytes(m.filename).is_none() {
1028 missing.push(m.filename);
1029 }
1030 }
1031 assert!(
1032 missing.is_empty(),
1033 "missing exact_bytes pins for: {missing:?}"
1034 );
1035 }
1036
1037 #[test]
1038 fn reviewed_sha256_pins_cover_every_catalogue_filename() {
1039 let mut missing = Vec::new();
1041 let mut seen = std::collections::HashSet::new();
1042 for m in MODELS {
1043 if !seen.insert(m.filename) {
1044 continue; }
1046 if pinned_sha256(m.filename).is_none() {
1047 missing.push(m.filename);
1048 } else {
1049 let pin = pinned_sha256(m.filename).unwrap();
1050 assert_eq!(pin.len(), 64, "pin length for {}", m.filename);
1051 }
1052 }
1053 assert!(missing.is_empty(), "missing sha256 pins for: {missing:?}");
1054 }
1055
1056 #[test]
1057 fn aliases_share_canonical_artifact_pins() {
1058 let large = lookup_model("large").unwrap();
1059 let large_v3 = lookup_model("large-v3").unwrap();
1060 assert_eq!(large.filename, large_v3.filename);
1061 assert_eq!(
1062 pinned_sha256(large.filename),
1063 pinned_sha256(large_v3.filename)
1064 );
1065 }
1066
1067 #[test]
1068 fn large_v3_q5_0_is_experimental_tier() {
1069 assert_eq!(
1070 model_support_tier("large-v3-q5_0"),
1071 ModelSupportTier::Experimental
1072 );
1073 assert_eq!(model_support_tier("base"), ModelSupportTier::Supported);
1074 }
1075}