Files
Lumotia/crates/transcription/src/model_manager.rs
Jake 42f4d07e48
Some checks failed
check / cargo check (macos-latest) (push) Has been cancelled
check / cargo check (ubuntu-22.04) (push) Has been cancelled
check / cargo check (windows-latest) (push) Has been cancelled
check / svelte build + lint (push) Has been cancelled
agent: lumotia-rebrand — drop MagnotiaError prefix, now lumotia_core::Error
Phase 3 of the rebrand cascade per locked decision D4. MagnotiaError ->
Error in crates/core/src/error.rs (the crate name already qualifies it).
92 usages across 14 .rs files renamed via word-boundary sed.

One collision required disambiguation: lumotia_storage already had its
own local Error type (introduced by the slop-pass Area A residuals work).
crates/storage/src/error.rs aliases the imported core error as CoreError
on import; the From<Error> for CoreError boundary impl and the
CoreError::Storage construction site use the alias.

cargo build --workspace passes. cargo test --workspace: 330 pass, 0 fail.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-13 08:58:05 +01:00

620 lines
22 KiB
Rust

use std::collections::HashSet;
use std::path::{Path, PathBuf};
use std::sync::{LazyLock, Mutex};
use lumotia_core::error::{Error, Result};
use lumotia_core::model_registry::{find_model, ModelFile};
use lumotia_core::types::{DownloadProgress, ModelId};
static ACTIVE_DOWNLOADS: LazyLock<Mutex<HashSet<String>>> =
LazyLock::new(|| Mutex::new(HashSet::new()));
struct DownloadReservation {
id: String,
}
impl DownloadReservation {
fn acquire(id: &ModelId) -> Result<Self> {
let id = id.as_str().to_string();
let mut active = ACTIVE_DOWNLOADS
.lock()
.map_err(|_| Error::DownloadFailed("download lock poisoned".into()))?;
if !active.insert(id.clone()) {
return Err(Error::DownloadFailed(format!(
"download already in progress for {id}"
)));
}
Ok(Self { id })
}
}
impl Drop for DownloadReservation {
fn drop(&mut self) {
if let Ok(mut active) = ACTIVE_DOWNLOADS.lock() {
active.remove(&self.id);
}
}
}
/// Resolve the models storage directory.
/// Windows: %LOCALAPPDATA%/magnotia/models
/// Unix: ~/.magnotia/models
pub fn models_dir() -> PathBuf {
lumotia_core::paths::app_paths().models_dir()
}
/// Get the directory path where a specific model's files are stored.
pub fn model_dir(id: &ModelId) -> PathBuf {
lumotia_core::paths::app_paths().speech_model_dir(id)
}
/// Check whether all files for a model have been downloaded.
pub fn is_downloaded(id: &ModelId) -> bool {
let entry = match find_model(id) {
Some(e) => e,
None => return false,
};
let dir = model_dir(id);
entry.files.iter().all(|f| dir.join(f.filename).exists())
&& verified_manifest_matches(entry, &dir)
}
/// List all downloaded model IDs.
pub fn list_downloaded() -> Vec<ModelId> {
lumotia_core::model_registry::all_models()
.iter()
.filter(|m| is_downloaded(&m.id))
.map(|m| m.id.clone())
.collect()
}
/// Download all files for a model, calling the progress callback per chunk.
/// Files are downloaded to a .part suffix and atomically renamed on completion.
///
/// For files that declare a `sha256` checksum we validate an existing
/// complete file before skipping the download — a truncated or
/// tampered file gets redownloaded automatically (pattern ported from
/// `lumotia-llm`'s model_manager, item #8 in the Whisper ecosystem brief).
pub async fn download(
id: &ModelId,
progress: impl Fn(DownloadProgress) + Send + 'static,
) -> Result<()> {
let _reservation = DownloadReservation::acquire(id)?;
let entry = find_model(id).ok_or_else(|| Error::ModelNotFound(id.clone()))?;
let dir = model_dir(id);
std::fs::create_dir_all(&dir)?;
for file in &entry.files {
let dest = dir.join(file.filename);
if dest.exists() {
// Validate the existing file. If the hash doesn't match,
// the file is corrupt (partial download, tampering, bit
// rot) and we must re-fetch it to avoid crashing on
// model load later.
match sha256_of_file(&dest) {
Ok(actual) if actual.eq_ignore_ascii_case(file.sha256) => continue,
Ok(_actual) => {
let _ = std::fs::remove_file(&dest);
}
Err(e) => {
return Err(Error::DownloadFailed(format!(
"failed to verify existing {}: {e}",
file.filename
)));
}
}
}
download_file(file, &dest, id, &progress).await?;
}
write_verified_manifest(entry, &dir)?;
Ok(())
}
fn verified_manifest_path(dir: &Path) -> PathBuf {
dir.join(".lumotia-verified")
}
fn verified_manifest_matches(
entry: &lumotia_core::model_registry::ModelEntry,
dir: &Path,
) -> bool {
let manifest = match std::fs::read_to_string(verified_manifest_path(dir)) {
Ok(contents) => contents,
Err(_) => return false,
};
for file in &entry.files {
let path = dir.join(file.filename);
let size = match std::fs::metadata(&path) {
Ok(metadata) => metadata.len(),
Err(_) => return false,
};
let expected_line = format!("{}\t{}\t{}", file.filename, file.sha256, size);
if !manifest.lines().any(|line| line == expected_line) {
return false;
}
}
true
}
fn write_verified_manifest(
entry: &lumotia_core::model_registry::ModelEntry,
dir: &Path,
) -> std::io::Result<()> {
let mut lines = Vec::with_capacity(entry.files.len() + 1);
lines.push("version\t1".to_string());
for file in &entry.files {
let size = std::fs::metadata(dir.join(file.filename))?.len();
lines.push(format!("{}\t{}\t{}", file.filename, file.sha256, size));
}
std::fs::write(
verified_manifest_path(dir),
format!("{}\n", lines.join("\n")),
)
}
/// Non-streaming SHA256 of a file on disk. Used by `download()` to
/// validate an existing complete file before trusting it.
fn sha256_of_file(path: &Path) -> std::io::Result<String> {
use sha2::{Digest, Sha256};
let mut hasher = Sha256::new();
let mut file = std::fs::File::open(path)?;
let mut buffer = [0u8; 8192];
loop {
let n = std::io::Read::read(&mut file, &mut buffer)?;
if n == 0 {
break;
}
hasher.update(&buffer[..n]);
}
Ok(format!("{:x}", hasher.finalize()))
}
/// Download a single file with HTTP Range resume and optional SHA256 verification.
///
/// Resume pattern from Buzz (chidiwilliams/buzz): if a .part file exists,
/// send a Range header to resume from where we left off. SHA256 is checked
/// incrementally during download — no second pass over the file.
async fn download_file(
file: &ModelFile,
dest: &Path,
model_id: &ModelId,
progress: &(impl Fn(DownloadProgress) + Send),
) -> Result<()> {
use futures_util::StreamExt;
use sha2::{Digest, Sha256};
let part_path = dest.with_extension(
dest.extension()
.map(|e| format!("{}.part", e.to_string_lossy()))
.unwrap_or_else(|| "part".to_string()),
);
let client = reqwest::Client::builder()
.connect_timeout(std::time::Duration::from_secs(30))
.build()
.map_err(|e| Error::DownloadFailed(e.to_string()))?;
// Check for existing partial download (resume support)
let existing_bytes = if part_path.exists() {
std::fs::metadata(&part_path).map(|m| m.len()).unwrap_or(0)
} else {
0
};
let mut request = client.get(file.url);
let resuming = existing_bytes > 0;
if resuming {
request = request.header("Range", format!("bytes={existing_bytes}-"));
}
let response = request
.send()
.await
.map_err(|e| Error::DownloadFailed(e.to_string()))?;
// If we requested Range but the server returned 200 (full file), the
// server does not support resume. Rather than blindly appending a
// full file on top of our partial bytes (which would produce a
// corrupt result), restart cleanly. This mirrors the lumotia-llm
// ResumeUnsupported branch — item #8 of the brief.
//
// For the non-resume path, we still have to validate the status:
// reqwest does not error on 4xx/5xx by default, so without this
// check a 404 or 500 would be streamed into `.part` and renamed
// over the destination as if the download succeeded
// (2026-04-22 review MAJOR).
let actually_resuming = if resuming {
match response.status().as_u16() {
206 => true,
200 => {
// Server ignored our Range header — treat as fresh start.
// The old .part bytes are discarded below.
false
}
other => {
return Err(Error::DownloadFailed(format!(
"resume request returned unexpected status {other}"
)));
}
}
} else {
if !response.status().is_success() {
return Err(Error::DownloadFailed(format!(
"download returned HTTP {} for {}",
response.status(),
file.filename
)));
}
false
};
let total_bytes = if actually_resuming {
// Content-Range: bytes START-END/TOTAL — extract TOTAL
response
.headers()
.get("content-range")
.and_then(|v| v.to_str().ok())
.and_then(|s| s.rsplit('/').next())
.and_then(|s| s.parse::<u64>().ok())
.unwrap_or(0)
} else {
response.content_length().unwrap_or(0)
};
let mut stream = response.bytes_stream();
let mut downloaded: u64 = if actually_resuming { existing_bytes } else { 0 };
let mut last_percent: u8 = 0;
// Open file for append (resume) or create (fresh start)
let mut out = if actually_resuming {
std::fs::OpenOptions::new().append(true).open(&part_path)?
} else {
std::fs::File::create(&part_path)?
};
let mut hasher = Sha256::new();
if actually_resuming {
let mut partial = std::fs::File::open(&part_path)?;
let mut buffer = [0u8; 8192];
loop {
let n = std::io::Read::read(&mut partial, &mut buffer)?;
if n == 0 {
break;
}
hasher.update(&buffer[..n]);
}
}
while let Some(chunk) = stream.next().await {
let chunk = chunk.map_err(|e| Error::DownloadFailed(e.to_string()))?;
std::io::Write::write_all(&mut out, &chunk)?;
hasher.update(&chunk);
downloaded += chunk.len() as u64;
let percent = if total_bytes > 0 {
((downloaded as f64 / total_bytes as f64) * 100.0) as u8
} else {
0
};
if percent != last_percent {
last_percent = percent;
progress(DownloadProgress {
model_id: model_id.clone(),
file_name: file.filename.to_string(),
bytes_downloaded: downloaded,
total_bytes,
percent,
});
}
}
drop(out);
let actual = format!("{:x}", hasher.finalize());
if actual != file.sha256 {
let _ = std::fs::remove_file(&part_path);
return Err(Error::DownloadFailed(format!(
"SHA256 mismatch for {}: expected {}, got {}",
file.filename, file.sha256, actual
)));
}
// Atomic rename — file is complete and verified
std::fs::rename(&part_path, dest)?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use sha2::Digest;
use tempfile::tempdir;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
#[test]
fn model_dir_returns_correct_path() {
let id = ModelId::new("whisper-tiny-en");
let path = model_dir(&id);
assert!(path.to_string_lossy().contains("whisper-tiny-en"));
}
#[test]
fn is_downloaded_returns_false_for_missing_model() {
let id = ModelId::new("nonexistent-model");
assert!(!is_downloaded(&id));
}
#[test]
fn list_downloaded_returns_empty_when_no_models() {
let list = list_downloaded();
// In test environment, no models are downloaded
// This just verifies the function doesn't panic
assert!(list.len() <= lumotia_core::model_registry::all_models().len());
}
#[test]
fn sha256_of_file_matches_sha2() {
let dir = tempdir().unwrap();
let path = dir.path().join("f.bin");
std::fs::write(&path, b"hello world").unwrap();
let expected = format!("{:x}", sha2::Sha256::digest(b"hello world"));
assert_eq!(sha256_of_file(&path).unwrap(), expected);
}
/// A minimal HTTP server that sends a Range response when a Range
/// header is present and otherwise sends the full body. Ported from
/// crates/llm/src/model_manager.rs to give the transcription
/// download stack the same fixture-backed coverage.
async fn spawn_range_server(content: Vec<u8>) -> std::net::SocketAddr {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.unwrap();
let mut buf = vec![0u8; 2048];
let size = socket.read(&mut buf).await.unwrap();
let request = String::from_utf8_lossy(&buf[..size]).to_lowercase();
let range_start = request
.lines()
.find_map(|line| line.strip_prefix("range: bytes="))
.and_then(|line| line.strip_suffix('-'))
.and_then(|line| line.trim().parse::<usize>().ok());
if let Some(start) = range_start {
let body = &content[start..];
let response = format!(
"HTTP/1.1 206 Partial Content\r\n\
Content-Length: {}\r\n\
Content-Range: bytes {}-{}/{}\r\n\
Accept-Ranges: bytes\r\n\r\n",
body.len(),
start,
content.len() - 1,
content.len(),
);
socket.write_all(response.as_bytes()).await.unwrap();
socket.write_all(body).await.unwrap();
} else {
let response = format!(
"HTTP/1.1 200 OK\r\n\
Content-Length: {}\r\n\
Accept-Ranges: bytes\r\n\r\n",
content.len(),
);
socket.write_all(response.as_bytes()).await.unwrap();
socket.write_all(&content).await.unwrap();
}
});
addr
}
/// A minimal HTTP server that responds with 200 + full body **iff**
/// the request actually carries a `Range` header, and 400 otherwise.
/// This models a mirror / proxy that accepts Range requests but
/// refuses to honour them (returning a fresh full body), which is
/// exactly the ResumeUnsupported branch `download_file` needs to
/// handle. The 400-on-missing-Range behaviour is load-bearing for
/// the test: it turns "client never sent Range" into a download
/// failure, so deleting the resume-detection logic causes the test
/// to fail rather than pass coincidentally through File::create's
/// truncation semantics.
async fn spawn_no_range_server(content: Vec<u8>) -> std::net::SocketAddr {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.unwrap();
let mut buf = vec![0u8; 2048];
let size = socket.read(&mut buf).await.unwrap();
let request = String::from_utf8_lossy(&buf[..size]).to_lowercase();
let saw_range_header = request
.lines()
.any(|line| line.trim_start().starts_with("range:"));
if !saw_range_header {
let response = "HTTP/1.1 400 Bad Request\r\n\
Content-Length: 0\r\n\r\n";
socket.write_all(response.as_bytes()).await.unwrap();
return;
}
let response = format!(
"HTTP/1.1 200 OK\r\n\
Content-Length: {}\r\n\r\n",
content.len(),
);
socket.write_all(response.as_bytes()).await.unwrap();
socket.write_all(&content).await.unwrap();
});
addr
}
/// ModelFile stores `&'static str` fields, so we leak the strings
/// once per test — tests are one-shot, so the cost is noise.
fn leak(s: String) -> &'static str {
Box::leak(s.into_boxed_str())
}
#[tokio::test]
async fn download_file_resumes_from_partial_and_verifies_sha() {
let body = b"resumable transcription payload".to_vec();
let expected_sha = format!("{:x}", sha2::Sha256::digest(&body));
let addr = spawn_range_server(body.clone()).await;
let dir = tempdir().unwrap();
let dest = dir.path().join("fixture.bin");
let part = dest.with_extension("bin.part");
// Pretend we already downloaded the first 7 bytes.
std::fs::write(&part, &body[..7]).unwrap();
let file = ModelFile {
filename: leak(dest.file_name().unwrap().to_string_lossy().into_owned()),
url: leak(format!("http://{addr}/fixture.bin")),
size: lumotia_core::types::Megabytes(0),
sha256: leak(expected_sha.clone()),
};
let id = ModelId::new("test-fixture");
download_file(&file, &dest, &id, &|_| ()).await.unwrap();
let bytes = std::fs::read(&dest).unwrap();
assert_eq!(bytes, body);
assert!(!part.exists());
assert_eq!(sha256_of_file(&dest).unwrap(), expected_sha);
}
#[tokio::test]
async fn download_file_restarts_when_server_ignores_range() {
// Covers the ResumeUnsupported branch documented in `download_file`:
// when a partial `.part` file exists and the server returns 200
// (full body) to our Range request, we must discard the stale
// partial bytes and write the fresh body from offset zero rather
// than appending on top.
let body = b"fresh transcription payload that replaces any stale partial".to_vec();
let expected_sha = format!("{:x}", sha2::Sha256::digest(&body));
let addr = spawn_no_range_server(body.clone()).await;
let dir = tempdir().unwrap();
let dest = dir.path().join("fixture.bin");
let part = dest.with_extension("bin.part");
// Pretend a previous attempt downloaded 12 bytes of something
// entirely unrelated. If the client naively appended the 200
// body, the final file would start with these bytes.
std::fs::write(&part, b"STALE_BYTES1").unwrap();
let file = ModelFile {
filename: leak(dest.file_name().unwrap().to_string_lossy().into_owned()),
url: leak(format!("http://{addr}/fixture.bin")),
size: lumotia_core::types::Megabytes(0),
sha256: leak(expected_sha),
};
let id = ModelId::new("test-fixture");
download_file(&file, &dest, &id, &|_| ()).await.unwrap();
let bytes = std::fs::read(&dest).unwrap();
assert_eq!(
bytes, body,
"server returned 200 to Range — downloader must discard stale .part and rewrite from scratch"
);
assert!(!part.exists(), ".part → dest rename must run after restart");
}
/// Always returns HTTP 500 with a short error body. Used to verify
/// the non-resume download path validates status codes rather than
/// writing error bodies into `.part` and renaming them over the
/// destination.
async fn spawn_500_server() -> std::net::SocketAddr {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.unwrap();
let mut buf = vec![0u8; 2048];
let _ = socket.read(&mut buf).await.unwrap();
let body = b"internal error";
let response = format!(
"HTTP/1.1 500 Internal Server Error\r\n\
Content-Length: {}\r\n\r\n",
body.len()
);
socket.write_all(response.as_bytes()).await.unwrap();
socket.write_all(body).await.unwrap();
});
addr
}
#[tokio::test]
async fn download_file_rejects_5xx_on_non_resume_path() {
// Regression for the 2026-04-22 review: reqwest does not
// auto-error on 4xx/5xx, and the non-resume branch previously
// streamed any status' body into `.part` and renamed it over
// the destination.
let addr = spawn_500_server().await;
let dir = tempdir().unwrap();
let dest = dir.path().join("fixture.bin");
let part = dest.with_extension("bin.part");
let file = ModelFile {
filename: leak(dest.file_name().unwrap().to_string_lossy().into_owned()),
url: leak(format!("http://{addr}/fixture.bin")),
size: lumotia_core::types::Megabytes(0),
sha256: leak("0".repeat(64)),
};
let id = ModelId::new("test-fixture");
let err = download_file(&file, &dest, &id, &|_| ())
.await
.expect_err("5xx must fail");
let msg = err.to_string();
assert!(
msg.contains("HTTP 500"),
"error should name the HTTP status, got: {msg}"
);
assert!(!dest.exists(), "5xx must not leave a destination file");
assert!(!part.exists(), "5xx must not leave a .part file");
}
#[tokio::test]
async fn download_file_fails_on_sha_mismatch_and_cleans_part_file() {
let body = b"speech-to-text fixture body".to_vec();
let addr = spawn_range_server(body.clone()).await;
let dir = tempdir().unwrap();
let dest = dir.path().join("fixture.bin");
let file = ModelFile {
filename: leak(dest.file_name().unwrap().to_string_lossy().into_owned()),
url: leak(format!("http://{addr}/fixture.bin")),
size: lumotia_core::types::Megabytes(0),
sha256: leak("deadbeef".repeat(8)),
};
let id = ModelId::new("test-fixture");
let err = download_file(&file, &dest, &id, &|_| ())
.await
.expect_err("mismatched sha must fail");
let msg = err.to_string();
assert!(msg.contains("SHA256 mismatch"), "unexpected error: {msg}");
assert!(
!dest.exists(),
".part → dest rename must not run on mismatch"
);
let part = dest.with_extension("bin.part");
assert!(!part.exists(), "failed hash must clean up the .part file");
}
}