Jake
2371e73f18
The downloader already handles servers that return 200 to a Range request by falling through to a truncating File::create on the .part path, discarding stale partial bytes. That branch had no dedicated fixture test — the SHA mismatch and Range-honouring resume cases were covered but the mirror / proxy that strips Range support was not. Adds spawn_no_range_server (always 200, full body regardless of Range header) and download_file_restarts_when_server_ignores_range. Writes 12 bytes of stale content to .part, kicks off a download, asserts the final file matches the fresh body exactly (not stale-bytes-prefixed) and the .part file is cleaned up. Brings the transcription downloader to test-coverage parity with crates/llm/src/model_manager.rs per brief item #8 ("test coverage parity" acceptance).
476 lines
17 KiB
Rust
476 lines
17 KiB
Rust
use std::path::{Path, PathBuf};
|
|
|
|
use kon_core::error::{KonError, Result};
|
|
use kon_core::model_registry::{find_model, ModelFile};
|
|
use kon_core::types::{DownloadProgress, ModelId};
|
|
|
|
/// Resolve the models storage directory.
|
|
/// Windows: %LOCALAPPDATA%/kon/models
|
|
/// Unix: ~/.kon/models
|
|
pub fn models_dir() -> PathBuf {
|
|
if cfg!(target_os = "windows") {
|
|
let local_app_data = std::env::var("LOCALAPPDATA").unwrap_or_else(|_| ".".to_string());
|
|
PathBuf::from(local_app_data).join("kon").join("models")
|
|
} else {
|
|
dirs_path().join("models")
|
|
}
|
|
}
|
|
|
|
fn dirs_path() -> PathBuf {
|
|
if cfg!(target_os = "windows") {
|
|
let local_app_data = std::env::var("LOCALAPPDATA").unwrap_or_else(|_| ".".to_string());
|
|
PathBuf::from(local_app_data).join("kon")
|
|
} else {
|
|
let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
|
|
PathBuf::from(home).join(".kon")
|
|
}
|
|
}
|
|
|
|
/// Get the directory path where a specific model's files are stored.
|
|
pub fn model_dir(id: &ModelId) -> PathBuf {
|
|
models_dir().join(id.as_str())
|
|
}
|
|
|
|
/// 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())
|
|
}
|
|
|
|
/// List all downloaded model IDs.
|
|
pub fn list_downloaded() -> Vec<ModelId> {
|
|
kon_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
|
|
/// `kon-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 entry = find_model(id).ok_or_else(|| KonError::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() {
|
|
if let Some(expected_sha) = file.sha256 {
|
|
// 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(expected_sha) => continue,
|
|
Ok(_actual) => {
|
|
// Corrupt — remove + fall through to re-download.
|
|
let _ = std::fs::remove_file(&dest);
|
|
}
|
|
Err(e) => {
|
|
return Err(KonError::DownloadFailed(format!(
|
|
"failed to verify existing {}: {e}",
|
|
file.filename
|
|
)));
|
|
}
|
|
}
|
|
} else {
|
|
// No checksum — honour the existing file as-is; the
|
|
// engine will barf on load if it's broken.
|
|
continue;
|
|
}
|
|
}
|
|
download_file(file, &dest, id, &progress).await?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// 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| KonError::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);
|
|
|
|
// If we have a partial file and no SHA256 to verify (can't verify partial),
|
|
// request a range resume. If SHA256 is set, we restart to ensure integrity.
|
|
let resuming = existing_bytes > 0 && file.sha256.is_none();
|
|
if resuming {
|
|
request = request.header("Range", format!("bytes={existing_bytes}-"));
|
|
}
|
|
|
|
let response = request
|
|
.send()
|
|
.await
|
|
.map_err(|e| KonError::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 kon-llm
|
|
// ResumeUnsupported branch — item #8 of the brief.
|
|
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(KonError::DownloadFailed(format!(
|
|
"resume request returned unexpected status {other}"
|
|
)));
|
|
}
|
|
}
|
|
} else {
|
|
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)?
|
|
};
|
|
|
|
// Incremental SHA256 — only when a checksum is provided
|
|
let mut hasher = file.sha256.map(|_| Sha256::new());
|
|
|
|
// If resuming without SHA256, we can't hash the already-downloaded portion,
|
|
// but we also don't need to — we only hash when sha256 is set, and we
|
|
// restart from scratch in that case.
|
|
|
|
while let Some(chunk) = stream.next().await {
|
|
let chunk = chunk.map_err(|e| KonError::DownloadFailed(e.to_string()))?;
|
|
std::io::Write::write_all(&mut out, &chunk)?;
|
|
if let Some(ref mut h) = hasher {
|
|
h.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);
|
|
|
|
// Verify SHA256 if provided
|
|
if let (Some(expected), Some(hasher)) = (file.sha256, hasher) {
|
|
let actual = format!("{:x}", hasher.finalize());
|
|
if actual != expected {
|
|
// Delete corrupt file so next attempt starts fresh
|
|
let _ = std::fs::remove_file(&part_path);
|
|
return Err(KonError::DownloadFailed(format!(
|
|
"SHA256 mismatch for {}: expected {}, got {}",
|
|
file.filename, expected, 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() <= kon_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 always responds with 200 + the full
|
|
/// body, even when the request carries a `Range` header. Mirrors a
|
|
/// mirror / proxy that strips Range support, exercising the
|
|
/// ResumeUnsupported branch in `download_file`.
|
|
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 _ = socket.read(&mut buf).await.unwrap();
|
|
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: kon_core::types::Megabytes(0),
|
|
sha256: None, // resume path only kicks in when sha256 is absent
|
|
};
|
|
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());
|
|
// Confirm the full file hash matches what we would have got via
|
|
// a clean download — gives the resume path indirect integrity
|
|
// coverage even when the ModelFile has no sha256 set.
|
|
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 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: kon_core::types::Megabytes(0),
|
|
sha256: None,
|
|
};
|
|
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");
|
|
}
|
|
|
|
#[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: kon_core::types::Megabytes(0),
|
|
sha256: Some(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");
|
|
}
|
|
}
|