1use crate::download::{self, DownloadOptions, DownloadRequest};
4use crate::error::{EnvironmentError, ProviderError, Result, UserError};
5use sha2::{Digest, Sha256};
6use std::fs::{self, File, OpenOptions};
7use std::io::Read;
8use std::path::{Path, PathBuf};
9use std::sync::Arc;
10
11const HF_BASE: &str = "https://huggingface.co/ggerganov/whisper.cpp/resolve/main";
15pub const ARTIFACT_MANIFEST_VERSION: &str = "1";
17pub const ARTIFACT_MANIFEST_SOURCE: &str = "aurum-builtin-review";
19
20#[derive(Debug, Clone, Copy)]
22pub struct ModelInfo {
23 pub name: &'static str,
24 pub filename: &'static str,
25 pub approx_bytes: u64,
27 pub notes: &'static str,
29}
30
31pub const MODELS: &[ModelInfo] = &[
33 ModelInfo {
35 name: "tiny",
36 filename: "ggml-tiny.bin",
37 approx_bytes: 75_000_000,
38 notes: "fastest full-precision",
39 },
40 ModelInfo {
41 name: "tiny-q5_1",
42 filename: "ggml-tiny-q5_1.bin",
43 approx_bytes: 32_000_000,
44 notes: "quantized ~32MB — best first-run trial",
45 },
46 ModelInfo {
47 name: "tiny-q8_0",
48 filename: "ggml-tiny-q8_0.bin",
49 approx_bytes: 44_000_000,
50 notes: "quantized",
51 },
52 ModelInfo {
53 name: "tiny.en",
54 filename: "ggml-tiny.en.bin",
55 approx_bytes: 75_000_000,
56 notes: "english-only",
57 },
58 ModelInfo {
59 name: "tiny.en-q5_1",
60 filename: "ggml-tiny.en-q5_1.bin",
61 approx_bytes: 32_000_000,
62 notes: "english-only quantized",
63 },
64 ModelInfo {
66 name: "base",
67 filename: "ggml-base.bin",
68 approx_bytes: 142_000_000,
69 notes: "default full-precision",
70 },
71 ModelInfo {
72 name: "base-q5_1",
73 filename: "ggml-base-q5_1.bin",
74 approx_bytes: 60_000_000,
75 notes: "quantized ~60MB",
76 },
77 ModelInfo {
78 name: "base-q8_0",
79 filename: "ggml-base-q8_0.bin",
80 approx_bytes: 82_000_000,
81 notes: "quantized",
82 },
83 ModelInfo {
84 name: "base.en",
85 filename: "ggml-base.en.bin",
86 approx_bytes: 142_000_000,
87 notes: "english-only",
88 },
89 ModelInfo {
90 name: "base.en-q5_1",
91 filename: "ggml-base.en-q5_1.bin",
92 approx_bytes: 60_000_000,
93 notes: "english-only quantized",
94 },
95 ModelInfo {
97 name: "small",
98 filename: "ggml-small.bin",
99 approx_bytes: 466_000_000,
100 notes: "higher accuracy",
101 },
102 ModelInfo {
103 name: "small-q5_1",
104 filename: "ggml-small-q5_1.bin",
105 approx_bytes: 190_000_000,
106 notes: "quantized",
107 },
108 ModelInfo {
109 name: "small-q8_0",
110 filename: "ggml-small-q8_0.bin",
111 approx_bytes: 264_000_000,
112 notes: "quantized",
113 },
114 ModelInfo {
115 name: "small.en",
116 filename: "ggml-small.en.bin",
117 approx_bytes: 466_000_000,
118 notes: "english-only",
119 },
120 ModelInfo {
121 name: "small.en-q5_1",
122 filename: "ggml-small.en-q5_1.bin",
123 approx_bytes: 190_000_000,
124 notes: "english-only quantized",
125 },
126 ModelInfo {
128 name: "medium",
129 filename: "ggml-medium.bin",
130 approx_bytes: 1_500_000_000,
131 notes: "large download",
132 },
133 ModelInfo {
134 name: "medium.en",
135 filename: "ggml-medium.en.bin",
136 approx_bytes: 1_500_000_000,
137 notes: "english-only",
138 },
139 ModelInfo {
141 name: "large-v3",
142 filename: "ggml-large-v3.bin",
143 approx_bytes: 3_100_000_000,
144 notes: "highest quality",
145 },
146 ModelInfo {
147 name: "large-v3-q5_0",
148 filename: "ggml-large-v3-q5_0.bin",
149 approx_bytes: 1_080_000_000,
150 notes: "experimental — not recommended (degeneration risk; prefer large-v3-turbo)",
151 },
152 ModelInfo {
153 name: "large",
154 filename: "ggml-large-v3.bin",
155 approx_bytes: 3_100_000_000,
156 notes: "alias of large-v3",
157 },
158 ModelInfo {
159 name: "large-v3-turbo",
160 filename: "ggml-large-v3-turbo.bin",
161 approx_bytes: 1_600_000_000,
162 notes: "fast large",
163 },
164 ModelInfo {
165 name: "large-v3-turbo-q5_0",
166 filename: "ggml-large-v3-turbo-q5_0.bin",
167 approx_bytes: 574_000_000,
168 notes: "quantized turbo",
169 },
170 ModelInfo {
171 name: "turbo",
172 filename: "ggml-large-v3-turbo.bin",
173 approx_bytes: 1_600_000_000,
174 notes: "alias of large-v3-turbo",
175 },
176 ModelInfo {
177 name: "turbo-q5_0",
178 filename: "ggml-large-v3-turbo-q5_0.bin",
179 approx_bytes: 574_000_000,
180 notes: "alias of large-v3-turbo-q5_0",
181 },
182];
183
184pub fn available_model_names() -> String {
186 list_canonical_models()
187 .iter()
188 .map(|m| m.name)
189 .collect::<Vec<_>>()
190 .join(", ")
191}
192
193fn list_canonical_models() -> Vec<&'static ModelInfo> {
194 MODELS
195 .iter()
196 .filter(|m| !matches!(m.name, "large" | "turbo" | "turbo-q5_0"))
197 .collect()
198}
199
200pub fn lookup_model(name: &str) -> Result<&'static ModelInfo> {
201 let key = name.trim().to_ascii_lowercase();
202 MODELS.iter().find(|m| m.name == key).ok_or_else(|| {
203 UserError::InvalidModel {
204 model: name.to_string(),
205 available: available_model_names(),
206 }
207 .into()
208 })
209}
210
211pub fn models_dir(cache_dir: &Path) -> PathBuf {
213 cache_dir.join("models")
214}
215
216pub fn model_path(cache_dir: &Path, info: &ModelInfo) -> PathBuf {
218 models_dir(cache_dir).join(info.filename)
219}
220
221#[derive(Debug, Clone)]
223pub struct ModelStatus {
224 pub info: &'static ModelInfo,
225 pub cached: bool,
226 pub path: PathBuf,
227 pub size_bytes: Option<u64>,
228}
229
230pub fn list_models(cache_dir: &Path) -> Vec<ModelStatus> {
232 list_canonical_models()
233 .into_iter()
234 .map(|info| {
235 let path = model_path(cache_dir, info);
236 let (cached, size_bytes) = match fs::metadata(&path) {
237 Ok(m) if m.len() > 1_000_000 => (true, Some(m.len())),
238 _ => (false, None),
239 };
240 ModelStatus {
241 info,
242 cached,
243 path,
244 size_bytes,
245 }
246 })
247 .collect()
248}
249
250pub fn format_model_list(cache_dir: &Path) -> String {
252 let rows = list_models(cache_dir);
253 let mut out = String::from("Local whisper.cpp models (cache: ");
254 out.push_str(&models_dir(cache_dir).display().to_string());
255 out.push_str(")\n\n");
256 out.push_str(&format!(
257 "{:<22} {:>10} {:<8} {:<12} {}\n",
258 "NAME", "SIZE", "STATUS", "TIER", "NOTES"
259 ));
260 out.push_str(&format!(
261 "{:<22} {:>10} {:<8} {:<12} {}\n",
262 "----", "----", "------", "----", "-----"
263 ));
264 for row in rows {
265 let size = format_bytes(row.info.approx_bytes);
266 let status = if row.cached { "cached" } else { "—" };
267 let tier = match model_support_tier(row.info.name) {
268 ModelSupportTier::Supported => "supported",
269 ModelSupportTier::Experimental => "experimental",
270 };
271 out.push_str(&format!(
272 "{:<22} {:>10} {:<8} {:<12} {}\n",
273 row.info.name, size, status, tier, row.info.notes
274 ));
275 }
276 out.push_str(
277 "\nTip: first run downloads the selected model. Try `tiny-q5_1` (~32 MB) for a quick trial.\n",
278 );
279 out.push_str("Default model: `base` (~142 MB). Use --model <name> to choose.\n");
280 out.push_str(
281 "Guidance (single-clip dogfood, not formal WER): English lecture quality often \
282 favors `small.en` or `large-v3-turbo`; prefer `.en` variants for English-only audio. \
283 Avoid `large-v3-q5_0` for production (experimental).\n",
284 );
285 out
286}
287
288fn format_bytes(n: u64) -> String {
289 const KB: f64 = 1024.0;
290 const MB: f64 = KB * 1024.0;
291 const GB: f64 = MB * 1024.0;
292 let n = n as f64;
293 if n >= GB {
294 format!("{:.1} GB", n / GB)
295 } else if n >= MB {
296 format!("{:.0} MB", n / MB)
297 } else {
298 format!("{:.0} KB", n / KB)
299 }
300}
301
302#[derive(Debug, Clone)]
304pub struct DownloadProgress {
305 pub model: String,
306 pub downloaded_bytes: u64,
307 pub total_bytes: u64,
308}
309
310impl DownloadProgress {
311 pub fn fraction(&self) -> Option<f64> {
312 if self.total_bytes == 0 {
313 None
314 } else {
315 Some((self.downloaded_bytes as f64 / self.total_bytes as f64).clamp(0.0, 1.0))
316 }
317 }
318}
319
320pub type DownloadProgressCallback = Arc<dyn Fn(DownloadProgress) + Send + Sync>;
322
323#[derive(Clone, Default)]
325pub struct EnsureModelOptions {
326 pub local_only: bool,
328 pub show_progress: bool,
330 pub on_progress: Option<DownloadProgressCallback>,
332}
333
334impl EnsureModelOptions {
335 pub fn new() -> Self {
336 Self::default()
337 }
338
339 pub fn local_only(mut self, v: bool) -> Self {
340 self.local_only = v;
341 self
342 }
343
344 pub fn show_progress(mut self, v: bool) -> Self {
345 self.show_progress = v;
346 self
347 }
348
349 pub fn on_progress(mut self, cb: DownloadProgressCallback) -> Self {
350 self.on_progress = Some(cb);
351 self
352 }
353}
354
355pub fn is_model_cached(cache_dir: &Path, model_name: &str) -> bool {
357 let Ok(info) = lookup_model(model_name) else {
358 return false;
359 };
360 let path = model_path(cache_dir, info);
361 if !(path.exists()
362 && path
363 .metadata()
364 .map(|m| m.len() > 1_000_000)
365 .unwrap_or(false))
366 {
367 return false;
368 }
369 pinned_sha256(info.filename).is_some() && verify_model_basic(&path, info).is_ok()
371}
372
373pub async fn ensure_model(
375 cache_dir: &Path,
376 model_name: &str,
377 show_progress: bool,
378) -> Result<PathBuf> {
379 ensure_model_with_options(
380 cache_dir,
381 model_name,
382 EnsureModelOptions {
383 show_progress,
384 ..EnsureModelOptions::default()
385 },
386 )
387 .await
388}
389
390pub async fn ensure_model_with_options(
392 cache_dir: &Path,
393 model_name: &str,
394 opts: EnsureModelOptions,
395) -> Result<PathBuf> {
396 let info = lookup_model(model_name)?;
397 let path = model_path(cache_dir, info);
398
399 if path.exists()
400 && path
401 .metadata()
402 .map(|m| m.len() > 1_000_000)
403 .unwrap_or(false)
404 {
405 if pinned_sha256(info.filename).is_none() {
407 return Err(ProviderError::ModelDownload {
408 model: model_name.to_string(),
409 reason: format!(
410 "model `{}` has no reviewed SHA-256 pin — not available on the trusted path",
411 info.name
412 ),
413 }
414 .into());
415 }
416 if verify_model_basic(&path, info).is_ok() {
417 tracing::info!(model = info.name, path = %path.display(), "using cached model");
418 return Ok(path);
419 }
420 tracing::warn!(
421 model = info.name,
422 path = %path.display(),
423 "cached model failed integrity check; re-downloading"
424 );
425 let _ = fs::remove_file(&path);
426 }
427
428 if opts.local_only {
429 return Err(UserError::ModelNotCached {
430 model: model_name.to_string(),
431 }
432 .into());
433 }
434
435 fs::create_dir_all(models_dir(cache_dir)).map_err(|e| EnvironmentError::DirectoryAccess {
436 path: models_dir(cache_dir).display().to_string(),
437 reason: e.to_string(),
438 })?;
439
440 let lock_path = path.with_extension("bin.lock");
442 let lock_file = OpenOptions::new()
443 .create(true)
444 .read(true)
445 .write(true)
446 .truncate(false)
447 .open(&lock_path)
448 .map_err(|e| EnvironmentError::DirectoryAccess {
449 path: lock_path.display().to_string(),
450 reason: e.to_string(),
451 })?;
452
453 if opts.show_progress {
454 eprintln!("aurum: waiting for model download lock ({}) …", info.name);
455 }
456 lock_file.lock().map_err(|e| EnvironmentError::Other {
457 message: format!("failed to acquire model lock: {e}"),
458 })?;
459
460 if path.exists()
462 && path
463 .metadata()
464 .map(|m| m.len() > 1_000_000)
465 .unwrap_or(false)
466 && pinned_sha256(info.filename).is_some()
467 && verify_model_basic(&path, info).is_ok()
468 {
469 let _ = lock_file.unlock();
470 tracing::info!(model = info.name, "model appeared while waiting on lock");
471 return Ok(path);
472 }
473
474 if opts.local_only {
475 let _ = lock_file.unlock();
476 return Err(UserError::ModelNotCached {
477 model: model_name.to_string(),
478 }
479 .into());
480 }
481
482 if opts.show_progress {
483 eprintln!(
484 "aurum: downloading model `{}` ({}) — first run only …",
485 info.name,
486 format_bytes(info.approx_bytes)
487 );
488 }
489
490 let result = download_model(info, &path, opts.show_progress, opts.on_progress.as_ref()).await;
491 let _ = lock_file.unlock();
492 result?;
493 verify_model_basic(&path, info)?;
494 Ok(path)
495}
496
497async fn download_model(
498 info: &ModelInfo,
499 dest: &Path,
500 show_progress: bool,
501 on_progress: Option<&DownloadProgressCallback>,
502) -> Result<()> {
503 let Some(expected) = pinned_sha256(info.filename) else {
505 return Err(ProviderError::ModelDownload {
506 model: info.name.to_string(),
507 reason: format!(
508 "no reviewed SHA-256 pin for {} — refusing to publish unauthenticated artifact",
509 info.filename
510 ),
511 }
512 .into());
513 };
514
515 let url = format!("{HF_BASE}/{}?download=true", info.filename);
516 let req = DownloadRequest {
517 id: info.name,
518 filename: info.filename,
519 sha256: expected,
520 exact_bytes: pinned_exact_bytes(info.filename),
521 approx_bytes: info.approx_bytes,
522 url: &url,
523 };
524
525 let mut opts = DownloadOptions {
526 show_progress,
527 ..DownloadOptions::default()
528 };
529 if let Some(cb) = on_progress {
530 let model = info.name.to_string();
531 let cb = Arc::clone(cb);
532 opts.on_progress = Some(Arc::new(move |downloaded, total| {
533 cb(DownloadProgress {
534 model: model.clone(),
535 downloaded_bytes: downloaded,
536 total_bytes: total,
537 });
538 }));
539 }
540
541 download::download_verified_request(&req, dest, &opts).await?;
542
543 sweep_stale_partials(dest.parent().unwrap_or_else(|| Path::new(".")));
545 Ok(())
546}
547
548pub fn pinned_sha256(filename: &str) -> Option<&'static str> {
554 match filename {
555 "ggml-tiny.bin" => Some("be07e048e1e599ad46341c8d2a135645097a538221678b7acdd1b1919c6e1b21"),
556 "ggml-tiny-q5_1.bin" => {
557 Some("818710568da3ca15689e31a743197b520007872ff9576237bda97bd1b469c3d7")
558 }
559 "ggml-tiny-q8_0.bin" => {
560 Some("c2085835d3f50733e2ff6e4b41ae8a2b8d8110461e18821b09a15c40c42d1cca")
561 }
562 "ggml-tiny.en.bin" => {
563 Some("921e4cf8686fdd993dcd081a5da5b6c365bfde1162e72b08d75ac75289920b1f")
564 }
565 "ggml-tiny.en-q5_1.bin" => {
566 Some("c77c5766f1cef09b6b7d47f21b546cbddd4157886b3b5d6d4f709e91e66c7c2b")
567 }
568 "ggml-base.bin" => Some("60ed5bc3dd14eea856493d334349b405782ddcaf0028d4b5df4088345fba2efe"),
569 "ggml-base-q5_1.bin" => {
570 Some("422f1ae452ade6f30a004d7e5c6a43195e4433bc370bf23fac9cc591f01a8898")
571 }
572 "ggml-base-q8_0.bin" => {
573 Some("c577b9a86e7e048a0b7eada054f4dd79a56bbfa911fbdacf900ac5b567cbb7d9")
574 }
575 "ggml-base.en.bin" => {
576 Some("a03779c86df3323075f5e796cb2ce5029f00ec8869eee3fdfb897afe36c6d002")
577 }
578 "ggml-base.en-q5_1.bin" => {
579 Some("4baf70dd0d7c4247ba2b81fafd9c01005ac77c2f9ef064e00dcf195d0e2fdd2f")
580 }
581 "ggml-small.bin" => {
582 Some("1be3a9b2063867b937e64e2ec7483364a79917e157fa98c5d94b5c1fffea987b")
583 }
584 "ggml-small-q5_1.bin" => {
585 Some("ae85e4a935d7a567bd102fe55afc16bb595bdb618e11b2fc7591bc08120411bb")
586 }
587 "ggml-small-q8_0.bin" => {
588 Some("49c8fb02b65e6049d5fa6c04f81f53b867b5ec9540406812c643f177317f779f")
589 }
590 "ggml-small.en.bin" => {
591 Some("c6138d6d58ecc8322097e0f987c32f1be8bb0a18532a3f88f734d1bbf9c41e5d")
592 }
593 "ggml-small.en-q5_1.bin" => {
594 Some("bfdff4894dcb76bbf647d56263ea2a96645423f1669176f4844a1bf8e478ad30")
595 }
596 "ggml-medium.bin" => {
597 Some("6c14d5adee5f86394037b4e4e8b59f1673b6cee10e3cf0b11bbdbee79c156208")
598 }
599 "ggml-medium.en.bin" => {
600 Some("cc37e93478338ec7700281a7ac30a10128929eb8f427dda2e865faa8f6da4356")
601 }
602 "ggml-large-v3.bin" => {
603 Some("64d182b440b98d5203c4f9bd541544d84c605196c4f7b845dfa11fb23594d1e2")
604 }
605 "ggml-large-v3-q5_0.bin" => {
606 Some("d75795ecff3f83b5faa89d1900604ad8c780abd5739fae406de19f23ecd98ad1")
607 }
608 "ggml-large-v3-turbo.bin" => {
609 Some("1fc70f774d38eb169993ac391eea357ef47c88757ef72ee5943879b7e8e2bc69")
610 }
611 "ggml-large-v3-turbo-q5_0.bin" => {
612 Some("394221709cd5ad1f40c46e6031ca61bce88931e6e088c188294c6d5a55ffa7e2")
613 }
614 _ => None,
615 }
616}
617
618#[derive(Debug, Clone, Copy, PartialEq, Eq)]
620pub enum ModelSupportTier {
621 Supported,
623 Experimental,
625}
626
627pub fn model_support_tier(name: &str) -> ModelSupportTier {
629 match name {
630 "large-v3-q5_0" => ModelSupportTier::Experimental,
631 _ => ModelSupportTier::Supported,
632 }
633}
634
635pub fn pinned_exact_bytes(filename: &str) -> Option<u64> {
637 match filename {
638 "ggml-tiny.bin" => Some(77_691_713),
639 "ggml-tiny-q5_1.bin" => Some(32_152_673),
640 "ggml-tiny-q8_0.bin" => Some(43_537_433),
641 "ggml-tiny.en.bin" => Some(77_704_715),
642 "ggml-tiny.en-q5_1.bin" => Some(32_166_155),
643 "ggml-base.bin" => Some(147_951_465),
644 "ggml-base-q5_1.bin" => Some(59_707_625),
645 "ggml-base-q8_0.bin" => Some(81_768_585),
646 "ggml-base.en.bin" => Some(147_964_211),
647 "ggml-base.en-q5_1.bin" => Some(59_721_011),
648 "ggml-small.bin" => Some(487_601_967),
649 "ggml-small-q5_1.bin" => Some(190_085_487),
650 "ggml-small-q8_0.bin" => Some(264_464_607),
651 "ggml-small.en.bin" => Some(487_614_201),
652 "ggml-small.en-q5_1.bin" => Some(190_098_681),
653 "ggml-medium.bin" => Some(1_533_763_059),
654 "ggml-medium.en.bin" => Some(1_533_774_781),
655 "ggml-large-v3.bin" => Some(3_095_033_483),
656 "ggml-large-v3-q5_0.bin" => Some(1_081_140_203),
657 "ggml-large-v3-turbo.bin" => Some(1_624_555_275),
658 "ggml-large-v3-turbo-q5_0.bin" => Some(574_041_195),
659 _ => None,
660 }
661}
662
663pub fn artifact_manifest_json(info: &ModelInfo) -> serde_json::Value {
665 serde_json::json!({
666 "manifest_version": ARTIFACT_MANIFEST_VERSION,
667 "source": ARTIFACT_MANIFEST_SOURCE,
668 "id": info.name,
669 "filename": info.filename,
670 "approx_bytes": info.approx_bytes,
671 "exact_bytes": pinned_exact_bytes(info.filename),
672 "sha256": pinned_sha256(info.filename),
673 "license": "MIT (whisper.cpp weights via OpenAI Whisper terms)",
674 "family": "whisper",
675 "download_url_template": format!("{HF_BASE}/{}", info.filename),
676 })
677}
678
679fn sweep_stale_partials(dir: &Path) {
680 let Ok(entries) = fs::read_dir(dir) else {
681 return;
682 };
683 let stale_after = std::time::Duration::from_secs(6 * 3600);
684 let now = std::time::SystemTime::now();
685 for ent in entries.flatten() {
686 let name = ent.file_name();
687 let name = name.to_string_lossy();
688 if !name.ends_with(".aurum.partial") {
690 continue;
691 }
692 let Ok(meta) = ent.metadata() else {
693 continue;
694 };
695 let Ok(modified) = meta.modified() else {
696 continue;
697 };
698 if now.duration_since(modified).unwrap_or_default() > stale_after {
699 let _ = fs::remove_file(ent.path());
700 }
701 }
702}
703
704pub fn ensure_model_verified_local(path: &Path, info: &ModelInfo) -> Result<()> {
706 verify_model_basic(path, info)
707}
708
709fn verify_model_basic(path: &Path, info: &ModelInfo) -> Result<()> {
711 let meta = fs::metadata(path).map_err(|e| ProviderError::ModelDownload {
712 model: info.name.to_string(),
713 reason: format!("missing after download: {e}"),
714 })?;
715
716 if meta.len() < 1_000_000 {
717 return Err(ProviderError::ModelDownload {
719 model: info.name.to_string(),
720 reason: format!(
721 "model file is only {} bytes — likely truncated or HTML error page",
722 meta.len()
723 ),
724 }
725 .into());
726 }
727
728 let mut hdr = [0u8; 4];
729 let mut f = File::open(path).map_err(|e| ProviderError::ModelDownload {
730 model: info.name.to_string(),
731 reason: e.to_string(),
732 })?;
733 f.read_exact(&mut hdr)
734 .map_err(|e| ProviderError::ModelDownload {
735 model: info.name.to_string(),
736 reason: format!("cannot read header: {e}"),
737 })?;
738
739 let magic_ok = matches!(
740 &hdr,
741 b"ggml"
742 | b"lmgg"
743 | b"ggmf"
744 | b"fmgg"
745 | b"ggjt"
746 | b"tjgg"
747 | b"ggjf"
748 | b"fjgg"
749 | b"gguf"
750 | b"fugg"
751 | b"GGUF"
752 );
753
754 if !magic_ok {
755 return Err(ProviderError::ModelDownload {
757 model: info.name.to_string(),
758 reason: format!(
759 "model header {:?} is not a recognized ggml/gguf magic — refusing to use file",
760 hdr
761 ),
762 }
763 .into());
764 }
765
766 let Some(expected) = pinned_sha256(info.filename) else {
768 return Err(ProviderError::ModelDownload {
769 model: info.name.to_string(),
770 reason: format!(
771 "no reviewed SHA-256 pin for {} — not trusted",
772 info.filename
773 ),
774 }
775 .into());
776 };
777 if !verify_against_expected(path, expected) {
778 return Err(ProviderError::ModelDownload {
779 model: info.name.to_string(),
780 reason: format!(
781 "cached model failed pinned sha256 check ({expected}); \
782 run `aurum cache verify` / quarantine repair"
783 ),
784 }
785 .into());
786 }
787 if let Some(exact) = pinned_exact_bytes(info.filename) {
788 if meta.len() != exact {
789 return Err(ProviderError::ModelDownload {
790 model: info.name.to_string(),
791 reason: format!(
792 "cached model size mismatch (got {}, expected {exact})",
793 meta.len()
794 ),
795 }
796 .into());
797 }
798 }
799
800 Ok(())
801}
802
803fn verify_against_expected(path: &Path, expected: &str) -> bool {
804 let Ok(mut file) = File::open(path) else {
805 return false;
806 };
807 let mut hasher = Sha256::new();
808 let mut buf = [0u8; 64 * 1024];
809 loop {
810 match file.read(&mut buf) {
811 Ok(0) => break,
812 Ok(n) => hasher.update(&buf[..n]),
813 Err(_) => return false,
814 }
815 }
816 hex::encode(hasher.finalize()) == expected
817}
818
819#[cfg(test)]
820mod tests {
821 use super::*;
822
823 #[test]
824 fn lookup_known_models() {
825 assert_eq!(lookup_model("base").unwrap().filename, "ggml-base.bin");
826 assert_eq!(
827 lookup_model("tiny-q5_1").unwrap().filename,
828 "ggml-tiny-q5_1.bin"
829 );
830 assert_eq!(
831 lookup_model("large-v3-turbo").unwrap().filename,
832 "ggml-large-v3-turbo.bin"
833 );
834 assert_eq!(lookup_model("turbo").unwrap().name, "turbo");
835 assert!(lookup_model("nope").is_err());
836 }
837
838 #[test]
839 fn model_path_joins() {
840 let p = model_path(Path::new("/tmp/cache"), lookup_model("tiny").unwrap());
841 assert_eq!(p, PathBuf::from("/tmp/cache/models/ggml-tiny.bin"));
842 }
843
844 #[test]
845 fn list_includes_quantized() {
846 let list = format_model_list(Path::new("/tmp/aurum-cache-test"));
847 assert!(list.contains("tiny-q5_1"));
848 assert!(list.contains("base-q5_1"));
849 assert!(list.contains("first run"));
850 }
851
852 #[test]
853 fn every_catalogue_file_has_exact_size_metadata() {
854 let mut missing = Vec::new();
856 for m in MODELS {
857 if pinned_exact_bytes(m.filename).is_none() {
858 missing.push(m.filename);
859 }
860 }
861 assert!(
862 missing.is_empty(),
863 "missing exact_bytes pins for: {missing:?}"
864 );
865 }
866
867 #[test]
868 fn reviewed_sha256_pins_cover_every_catalogue_filename() {
869 let mut missing = Vec::new();
871 let mut seen = std::collections::HashSet::new();
872 for m in MODELS {
873 if !seen.insert(m.filename) {
874 continue; }
876 if pinned_sha256(m.filename).is_none() {
877 missing.push(m.filename);
878 } else {
879 let pin = pinned_sha256(m.filename).unwrap();
880 assert_eq!(pin.len(), 64, "pin length for {}", m.filename);
881 }
882 }
883 assert!(missing.is_empty(), "missing sha256 pins for: {missing:?}");
884 }
885
886 #[test]
887 fn aliases_share_canonical_artifact_pins() {
888 let large = lookup_model("large").unwrap();
889 let large_v3 = lookup_model("large-v3").unwrap();
890 assert_eq!(large.filename, large_v3.filename);
891 assert_eq!(
892 pinned_sha256(large.filename),
893 pinned_sha256(large_v3.filename)
894 );
895 }
896
897 #[test]
898 fn large_v3_q5_0_is_experimental_tier() {
899 assert_eq!(
900 model_support_tier("large-v3-q5_0"),
901 ModelSupportTier::Experimental
902 );
903 assert_eq!(model_support_tier("base"), ModelSupportTier::Supported);
904 }
905}