Files
ai-app/server/src/models.rs
T
irisandClaude Opus 5 79682f03a7 Condense the documentation and thin the server's comments
The markdown had accumulated a lot that was stale rather than wrong.
PLAN.md still described pi as the llama.cpp harness, a refcounted
LlamaServerManager, and a providers-by-hosts cross-product, all of which
were superseded or never built; it also carried a second copy of the HTTP
table that routes.rs owns. EXPLORER.md and TRANSCRIPT_CACHE.md held
implementation checklists for work that has since landed. AGENTS.md
restated most of PLAN.md's design instead of being the working-notes
layer it says it is. 3225 lines of markdown to 2180, with the stale
sections gone rather than reworded.

On the server, comments explaining what the code already says are out and
the ones recording a constraint, a measurement or an incident are kept but
cut to a few lines each: 5504 comment lines to 4586.

Four doc comments in session/mod.rs, and one each in process.rs and
usage.rs, had drifted onto the item above the one they describe --
functions were reordered without them, so `stop_session`'s doc sat on
`set_session_cwd`, `stat_of`'s on `struct Stat`, and `UsageMonitor`'s on
`type Cached`. Each is back on its own item.

routes.rs's module table also claimed later phases would add `/hosts`,
which setups replaced.

cargo test (127 passed), clippy --all-targets and fmt are clean.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-04 15:45:43 -04:00

644 lines
23 KiB
Rust

//! GGUF models on this machine, and the downloads that produce them.
//!
//! The registry pattern again: one owner, one lock, so what is on disk and what
//! this server believes cannot come apart. Three things shape the design, all
//! consequences of a model file being gigabytes rather than kilobytes:
//!
//! **A download belongs to the model, not to whoever asked for it.** It is
//! keyed by the model it produces and lives here, so any device can watch it --
//! including one that did not start it. State in a per-connection channel would
//! not survive the phone locking its screen, which for an hour-long download is
//! the normal case.
//!
//! **Every run has an id, and its outcome outlives it.** Without those, "not
//! downloading" is three answers at once -- it finished, it never started, or a
//! different run finished while you were away.
//!
//! **Progress is measured, never estimated.** `total` is whatever
//! `Content-Length` said and nothing else; when the server does not send one it
//! stays `None` and the phone shows that it does not know.
use std::collections::HashMap;
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use anyhow::{Context, Result, bail};
use serde::Serialize;
use wg_app_link::private;
/// Identifies this client to HuggingFace. They ask for one, and a request
/// without it is more likely to be rate-limited.
const USER_AGENT: &str = concat!("ai-server/", env!("CARGO_PKG_VERSION"));
/// Read size per loop iteration. Big enough that the syscall overhead is
/// nothing against a multi-gigabyte file, small enough that a cancel is noticed
/// promptly -- the flag is only checked between chunks.
const CHUNK: usize = 256 * 1024;
/// A model file sitting on this machine, ready to run.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct LocalModel {
/// `owner/repo/file.gguf` -- the HuggingFace coordinates, which are already
/// unique, so nothing has to invent an id.
pub key: String,
pub repo: String,
pub file: String,
pub bytes: u64,
}
/// What a run is doing, or did. Flat rather than a tagged enum carrying its
/// message, because the phone switches on this and a string it can compare is
/// easier to render than a variant it has to destructure.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum DownloadState {
Running,
/// Reading the finished file back to check it against the hash HuggingFace
/// publishes. Its own state because it takes real time on a multi-gigabyte
/// file and "still working" is the honest thing to show, rather than a bar
/// sitting at 100% for half a minute.
Verifying,
Finished,
Failed,
Cancelled,
}
/// One download run, as the phone sees it.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct DownloadStatus {
pub key: String,
/// Distinguishes this run from any earlier one for the same model, so a
/// device that was watching run 3 can tell it is now looking at run 4.
pub run: u64,
pub repo: String,
pub file: String,
pub state: DownloadState,
pub done: u64,
/// What `Content-Length` said, or absent when the server did not say.
/// Absent means "unknown", never "zero".
#[serde(skip_serializing_if = "Option::is_none")]
pub total: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
pub started: f64,
#[serde(skip_serializing_if = "Option::is_none")]
pub finished: Option<f64>,
}
/// The mutable half of a run, behind one lock.
#[derive(Debug)]
struct Progress {
state: DownloadState,
done: u64,
total: Option<u64>,
error: Option<String>,
started: f64,
finished: Option<f64>,
}
/// A run, shared between the thread doing the work and everyone watching.
struct Run {
id: u64,
key: String,
repo: String,
file: String,
progress: Mutex<Progress>,
/// Set by [`ModelStore::cancel`]; the download loop checks it between
/// chunks and stops, leaving the partial file for a later resume.
cancel: AtomicBool,
}
impl Run {
fn status(&self) -> DownloadStatus {
let p = self.progress.lock().unwrap();
DownloadStatus {
key: self.key.clone(),
run: self.id,
repo: self.repo.clone(),
file: self.file.clone(),
state: p.state,
done: p.done,
total: p.total,
error: p.error.clone(),
started: p.started,
finished: p.finished,
}
}
fn finish(&self, state: DownloadState, error: Option<String>) {
let mut p = self.progress.lock().unwrap();
p.state = state;
p.error = error;
p.finished = Some(crate::session::now());
}
}
/// Every model this machine has, and every download in flight or finished.
pub struct ModelStore {
dir: PathBuf,
/// Keyed by model key: one run per model at a time, and the last run for a
/// model stays here after it ends so its outcome can still be read.
runs: Mutex<HashMap<String, Arc<Run>>>,
next_run: AtomicU64,
}
impl ModelStore {
pub fn new(dir: PathBuf) -> Self {
Self {
dir,
runs: Mutex::new(HashMap::new()),
next_run: AtomicU64::new(1),
}
}
/// Where a model's file lives, refusing anything that would escape the
/// models directory.
///
/// The repo and file come from a phone, so they are treated as hostile:
/// every component must be an ordinary name. Rejecting rather than
/// sanitising, since a silently rewritten path would download the right
/// bytes to the wrong place.
fn path_for(&self, repo: &str, file: &str) -> Result<PathBuf> {
let mut path = self.dir.clone();
for part in repo.split('/').chain(file.split('/')) {
if part.is_empty() || part == "." || part == ".." || part.contains('\\') {
bail!("\"{repo}/{file}\" is not a name this can store: \"{part}\"");
}
path.push(part);
}
Ok(path)
}
pub fn key_for(repo: &str, file: &str) -> String {
format!("{repo}/{file}")
}
/// Every `.gguf` found under the models directory, newest first. Read from
/// disk on each call rather than cached: a file deleted by hand should stop
/// being offered.
pub fn list(&self) -> Vec<LocalModel> {
let mut found = Vec::new();
collect(&self.dir, &self.dir, &mut found);
found.sort_by(|a, b| a.key.cmp(&b.key));
found
}
/// The status of every run this server remembers.
pub fn downloads(&self) -> Vec<DownloadStatus> {
let runs = self.runs.lock().unwrap();
let mut all: Vec<_> = runs.values().map(|run| run.status()).collect();
all.sort_by_key(|status| std::cmp::Reverse(status.run));
all
}
/// Starts fetching `file` from `repo`, or returns the run already doing so.
/// Idempotent on purpose: a phone that lost its connection will press the
/// button again, and that must join the existing run rather than start a
/// second one writing the same file.
pub fn start(self: &Arc<Self>, repo: &str, file: &str) -> Result<DownloadStatus> {
let key = Self::key_for(repo, file);
let target = self.path_for(repo, file)?;
if target.is_file() {
bail!("{key} is already downloaded");
}
let mut runs = self.runs.lock().unwrap();
if let Some(existing) = runs.get(&key)
&& existing.progress.lock().unwrap().state == DownloadState::Running
{
return Ok(existing.status());
}
let run = Arc::new(Run {
id: self.next_run.fetch_add(1, Ordering::Relaxed),
key: key.clone(),
repo: repo.to_string(),
file: file.to_string(),
progress: Mutex::new(Progress {
state: DownloadState::Running,
done: 0,
total: None,
error: None,
started: crate::session::now(),
finished: None,
}),
cancel: AtomicBool::new(false),
});
runs.insert(key, Arc::clone(&run));
let status = run.status();
drop(runs);
// A dedicated thread rather than the blocking pool: this holds its
// thread for as long as the download takes, which is minutes to hours,
// and the pool exists for short work.
let store = Arc::clone(self);
std::thread::spawn(move || {
let outcome = store.fetch(&run, &target);
match outcome {
Ok(()) if run.cancel.load(Ordering::Relaxed) => {
run.finish(DownloadState::Cancelled, None);
tracing::info!("download {} cancelled", run.key);
}
Ok(()) => {
run.finish(DownloadState::Finished, None);
tracing::info!("download {} finished", run.key);
}
Err(err) => {
let message = format!("{err:#}");
tracing::warn!("download {} failed: {message}", run.key);
run.finish(DownloadState::Failed, Some(message));
}
}
});
Ok(status)
}
/// Asks a running download to stop. The partial file stays, so starting
/// again resumes rather than refetching.
pub fn cancel(&self, key: &str) -> Result<DownloadStatus> {
let runs = self.runs.lock().unwrap();
let Some(run) = runs.get(key) else {
bail!("no download for {key}");
};
run.cancel.store(true, Ordering::Relaxed);
Ok(run.status())
}
/// Removes a downloaded model, and any partial file for it.
pub fn delete(&self, key: &str) -> Result<()> {
let (repo, file) = key.rsplit_once('/').context("a key is repo/file")?;
let target = self.path_for(repo, file)?;
let partial = partial_of(&target);
if !target.is_file() && !partial.is_file() {
bail!("{key} is not downloaded");
}
for path in [&target, &partial] {
if path.is_file() {
std::fs::remove_file(path).with_context(|| format!("remove {}", path.display()))?;
}
}
self.runs.lock().unwrap().remove(key);
Ok(())
}
fn fetch(&self, run: &Run, target: &Path) -> Result<()> {
let partial = partial_of(target);
let identity = identity_of(target);
if let Some(parent) = target.parent() {
private::create_dir(parent)?;
}
// What we have, and what it was part of. A partial with no recorded
// identity is not resumable -- it could be a fragment of any revision.
let known = std::fs::read_to_string(&identity)
.ok()
.map(|s| s.trim().to_string());
let have = match known {
Some(_) => partial.metadata().map(|m| m.len()).unwrap_or(0),
None => 0,
};
let url = format!(
"https://huggingface.co/{}/resolve/main/{}",
run.repo,
run.file.replace(' ', "%20")
);
let (mut response, mut resumed) = request(&url, have)?;
let mut etag = etag_of(&response);
// HuggingFace's CDN ignores `If-Range` -- probed 2026-08-28: a
// deliberately stale validator still answers 206 with the ranged bytes.
// So the header cannot be relied on to restart us, and the check is done
// here instead: if what arrived is not the revision our partial belongs
// to, resuming would splice two files into something of exactly the
// right length and the wrong contents.
if resumed && etag.is_some() && etag != known {
tracing::info!(
"{} changed upstream since the partial was written -- starting again",
run.key,
);
let (fresh, fresh_resumed) = request(&url, 0)?;
response = fresh;
resumed = fresh_resumed;
etag = etag_of(&response);
}
// On a 206, Content-Length is the length of the *range*, not of the file
// -- taken at face value it would fill the bar at 72 MB of a 234 MB
// model. The whole size is the last field of Content-Range, which has
// the further merit of not depending on where the range began.
let total: Option<u64> = if resumed {
response
.headers()
.get("content-range")
.and_then(|v| v.to_str().ok())
.and_then(|v| {
v.rsplit_once('/')
.map(|(_, whole)| whole.trim().to_string())
})
.and_then(|whole| whole.parse().ok())
} else {
response
.headers()
.get("content-length")
.and_then(|v| v.to_str().ok()?.parse().ok())
};
let mut done = if resumed { have } else { 0 };
{
let mut p = run.progress.lock().unwrap();
p.done = done;
p.total = total;
}
// `truncate(false)` is the whole resume story: the file is opened to be
// seeked into and appended to, and truncating would throw away exactly
// the bytes the Range request just asked the server not to send again.
// Stated rather than left to the default.
let mut file = std::fs::OpenOptions::new()
.create(true)
.write(true)
.truncate(false)
.open(&partial)
.with_context(|| format!("open {}", partial.display()))?;
if resumed {
file.seek(SeekFrom::Start(have))
.context("seek to resume point")?;
} else {
file.set_len(0)
.context("truncate a partial we cannot resume onto")?;
}
// Written before the body, so an interrupted download leaves a partial
// that can still say which revision it belongs to.
if let Some(etag) = &etag {
std::fs::write(&identity, etag).ok();
}
let mut reader = response.body_mut().as_reader();
let mut buffer = vec![0u8; CHUNK];
loop {
if run.cancel.load(Ordering::Relaxed) {
file.flush().ok();
return Ok(());
}
let read = reader
.read(&mut buffer)
.context("reading from HuggingFace")?;
if read == 0 {
break;
}
file.write_all(&buffer[..read])
.context("writing the model file")?;
done += read as u64;
run.progress.lock().unwrap().done = done;
}
file.flush().context("flushing the model file")?;
drop(file);
// Checked before the rename, so a file that fails never gets the real
// name and `list` never offers it. With the identity check above this
// should not fire; it is here because a wrong model is the kind of
// failure that surfaces as bad output rather than as an error.
if let Some(expected) = published_sha256(&run.repo, &run.file) {
run.progress.lock().unwrap().state = DownloadState::Verifying;
let actual = sha256_of(&partial)?;
if actual != expected {
std::fs::remove_file(&partial).ok();
std::fs::remove_file(&identity).ok();
bail!(
"{} arrived corrupted -- HuggingFace publishes sha256 {expected}, what \
arrived hashes to {actual}. It has been deleted; downloading again \
starts clean.",
run.key,
);
}
}
// Renamed only once complete, so a file at its real name is always a
// whole model -- `list` needs no other way to tell.
std::fs::rename(&partial, target)
.with_context(|| format!("finish {}", target.display()))?;
std::fs::remove_file(&identity).ok();
Ok(())
}
}
/// The sha256 of a file, read in chunks -- these are gigabytes, and reading one
/// into memory to hash it would be the largest allocation this server makes.
fn sha256_of(path: &Path) -> Result<String> {
use sha2::{Digest, Sha256};
let mut file =
std::fs::File::open(path).with_context(|| format!("reopen {}", path.display()))?;
let mut hasher = Sha256::new();
let mut buffer = vec![0u8; CHUNK];
loop {
let read = file.read(&mut buffer).context("reading back to verify")?;
if read == 0 {
break;
}
hasher.update(&buffer[..read]);
}
// Hex by hand, as `wg_app_link::enroll::token_hash_hex` also has to, since
// this sha2 version's output type does not implement LowerHex.
Ok(hasher
.finalize()
.iter()
.map(|byte| format!("{byte:02x}"))
.collect())
}
/// One GET, ranged when there is something to resume onto.
fn request(url: &str, from: u64) -> Result<(ureq::http::Response<ureq::Body>, bool)> {
let mut get = ureq::get(url).header("User-Agent", USER_AGENT);
if from > 0 {
get = get.header("Range", &format!("bytes={from}-"));
}
let response = get.call().with_context(|| format!("GET {url}"))?;
// Trust the status, not the request: a server that ignores Range answers
// 200 with the whole file, and appending to that would corrupt it.
let resumed = response.status() == 206;
Ok((response, resumed))
}
fn etag_of(response: &ureq::http::Response<ureq::Body>) -> Option<String> {
Some(
response
.headers()
.get("etag")?
.to_str()
.ok()?
.trim()
.to_string(),
)
}
/// `x.gguf` -> `x.gguf.part.etag`, holding which revision the partial beside it
/// is a piece of.
fn identity_of(target: &Path) -> PathBuf {
let mut name = target.as_os_str().to_os_string();
name.push(".part.etag");
PathBuf::from(name)
}
/// `x.gguf` -> `x.gguf.part`, the in-progress name.
fn partial_of(target: &Path) -> PathBuf {
let mut name = target.as_os_str().to_os_string();
name.push(".part");
PathBuf::from(name)
}
/// Walks `dir` collecting `.gguf` files, keyed by their path under `root`.
fn collect(root: &Path, dir: &Path, found: &mut Vec<LocalModel>) {
let Ok(entries) = std::fs::read_dir(dir) else {
return;
};
for entry in entries.flatten() {
let path = entry.path();
if path.is_dir() {
collect(root, &path, found);
continue;
}
if path.extension().is_none_or(|e| e != "gguf") {
continue;
}
let Ok(relative) = path.strip_prefix(root) else {
continue;
};
let key = relative.to_string_lossy().replace('\\', "/");
let Some((repo, file)) = key.rsplit_once('/') else {
continue;
};
found.push(LocalModel {
key: key.clone(),
repo: repo.to_string(),
file: file.to_string(),
bytes: entry.metadata().map(|m| m.len()).unwrap_or(0),
});
}
}
/// A model repository on HuggingFace, as the browse screen shows it.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct RemoteRepo {
/// `owner/name`, which is what everything else here is keyed by.
pub id: String,
pub downloads: u64,
pub likes: u64,
}
/// One downloadable file inside a repository.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct RemoteFile {
pub path: String,
pub bytes: u64,
/// Already on this machine, so the phone can say so rather than offering to
/// fetch it again.
pub have: bool,
}
/// Searches HuggingFace for GGUF repositories matching `query`.
///
/// Proxied through this server rather than called from the phone, for two
/// reasons that both matter: the app trusts exactly one certificate -- this
/// server's -- and has no general internet trust to spend on huggingface.co,
/// and the machine that has to do the downloading is this one.
pub fn search(query: &str) -> Result<Vec<RemoteRepo>> {
let url = format!(
"https://huggingface.co/api/models?search={}&filter=gguf&limit=25&sort=downloads&direction=-1",
urlencode(query)
);
let body = get_json(&url)?;
let list = body
.as_array()
.context("HuggingFace returned something that is not a list")?;
Ok(list
.iter()
.filter_map(|m| {
Some(RemoteRepo {
id: m.get("id")?.as_str()?.to_string(),
downloads: m
.get("downloads")
.and_then(serde_json::Value::as_u64)
.unwrap_or(0),
likes: m
.get("likes")
.and_then(serde_json::Value::as_u64)
.unwrap_or(0),
})
})
.collect())
}
/// The sha256 HuggingFace publishes for one file, if it publishes one. It is
/// the LFS object id, which for these repositories is the sha256 of the content
/// -- so it is a free integrity check rather than a second source of truth.
fn published_sha256(repo: &str, file: &str) -> Option<String> {
let url = format!("https://huggingface.co/api/models/{repo}/tree/main?expand=true");
let body = get_json(&url).ok()?;
body.as_array()?.iter().find_map(|f| {
(f.get("path")?.as_str()? == file)
.then(|| f.get("lfs")?.get("oid")?.as_str().map(str::to_string))?
})
}
/// The GGUF files in one repository, largest last, with the ones already
/// downloaded marked.
pub fn files(repo: &str, store: &ModelStore) -> Result<Vec<RemoteFile>> {
let url = format!("https://huggingface.co/api/models/{repo}/tree/main");
let body = get_json(&url)?;
let list = body
.as_array()
.context("HuggingFace returned something that is not a list")?;
let have: std::collections::HashSet<String> = store.list().into_iter().map(|m| m.key).collect();
let mut files: Vec<RemoteFile> = list
.iter()
.filter_map(|f| {
let path = f.get("path")?.as_str()?.to_string();
if !path.ends_with(".gguf") {
return None;
}
Some(RemoteFile {
bytes: f
.get("size")
.and_then(serde_json::Value::as_u64)
.unwrap_or(0),
have: have.contains(&ModelStore::key_for(repo, &path)),
path,
})
})
.collect();
files.sort_by_key(|f| f.bytes);
Ok(files)
}
fn get_json(url: &str) -> Result<serde_json::Value> {
let text = ureq::get(url)
.header("User-Agent", USER_AGENT)
.call()
.and_then(|mut r| r.body_mut().read_to_string())
.with_context(|| format!("GET {url}"))?;
serde_json::from_str(&text).with_context(|| format!("{url} did not return JSON"))
}
/// Percent-encodes a query string. Deliberately minimal -- this escapes what a
/// model search actually contains rather than implementing the whole rule set,
/// and anything unexpected becomes `%XX` rather than being passed through.
fn urlencode(value: &str) -> String {
value
.bytes()
.map(|b| match b {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
(b as char).to_string()
}
b' ' => "+".to_string(),
other => format!("%{other:02X}"),
})
.collect()
}