diff --git a/server/src/main.rs b/server/src/main.rs index 7cc59e9..05fe403 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -17,6 +17,7 @@ mod auth; mod certs; mod config; mod media; +mod models; mod private; mod routes; mod session; @@ -88,6 +89,11 @@ struct Args { #[arg(long)] data_dir: Option, + /// Directory for downloaded GGUF models. Defaults to + /// `$XDG_DATA_HOME/ai-app/models`. + #[arg(long)] + models_dir: Option, + /// Directory holding the TLS certificates, generated here on first /// start. Defaults to `$XDG_CONFIG_HOME/ai-app/certs`. #[arg(long)] @@ -179,11 +185,19 @@ async fn main() -> Result<()> { let data_dir = args .data_dir .unwrap_or_else(|| data_home().join("sessions")); + // Beside the session data rather than under it: models outlive every + // session and are shared by all of them, so deleting a session must + // never take a multi-gigabyte download with it. + let models_dir = args + .models_dir + .unwrap_or_else(|| data_home().join("models")); + let models = Arc::new(models::ModelStore::new(models_dir.clone())); let manager = Arc::new( SessionManager::new(config_path.clone(), data_dir) .with_context(|| format!("failed to load {}", config_path.display()))?, ); tracing::info!("config: {}", config_path.display()); + tracing::info!("models: {}", models_dir.display()); for provider in manager.providers() { tracing::info!(" provider {} ({:?})", provider.name, provider.kind); } @@ -261,6 +275,7 @@ async fn main() -> Result<()> { // auth. Zero unauthenticated endpoints. let app = routes::router(Arc::clone(&manager)) .merge(routes::usage_router(monitor)) + .merge(routes::models_router(Arc::clone(&models))) .layer(axum::middleware::from_fn_with_state( Arc::clone(&manager), auth::require_token, diff --git a/server/src/models.rs b/server/src/models.rs new file mode 100644 index 0000000..acd3c0d --- /dev/null +++ b/server/src/models.rs @@ -0,0 +1,534 @@ +//! GGUF models on this machine, and the downloads that produce them. +//! +//! The registry pattern again (see `session`): one owner, one lock, so what +//! is on disk and what this server believes cannot come apart. +//! +//! Three things shape the design, all of them 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, and one that opened the app +//! after it finished. State in a per-connection channel would not survive +//! the phone locking its screen, which for an hour-long download is the +//! normal case rather than an edge one. +//! +//! **Every run has an id, and its outcome outlives it.** Without those, +//! "not downloading" is three different answers at once -- it finished, +//! it never started, or a different run finished while you were away -- +//! and over an hour that ambiguity is certain to be hit. A device compares +//! the run it was watching against the run reported now. +//! +//! **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, rather +//! than a bar drawn from how long the last download took. + +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 crate::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, + 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. + /// A device that was watching run 3 can tell that what it is looking + /// at now is run 4 rather than assuming its own run ended. + pub run: u64, + pub repo: String, + pub file: String, + pub state: DownloadState, + /// Bytes on disk, including any carried over from a resumed attempt. + pub done: u64, + /// What `Content-Length` said, or absent when the server did not say. + /// Absent means "unknown", never "zero" -- see this module's doc. + #[serde(skip_serializing_if = "Option::is_none")] + pub total: Option, + /// Present only when [`DownloadState::Failed`], and it is the reason. + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, + pub started: f64, + #[serde(skip_serializing_if = "Option::is_none")] + pub finished: Option, +} + +/// The mutable half of a run, behind one lock. +#[derive(Debug)] +struct Progress { + state: DownloadState, + done: u64, + total: Option, + error: Option, + started: f64, + finished: Option, +} + +/// A run, shared between the thread doing the work and everyone watching. +struct Run { + id: u64, + key: String, + repo: String, + file: String, + progress: Mutex, + /// 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) { + 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. Bounded by how many distinct models have been asked for. + runs: Mutex>>, + 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, and this server runs as the + /// user who started it, so they are treated as hostile: every + /// component must be an ordinary name. Rejecting is deliberate 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 { + 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, and the directory is small enough + /// that walking it costs nothing next to loading a model. + pub fn list(&self) -> Vec { + 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 { + 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 and came + /// back 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, repo: &str, file: &str) -> Result { + 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 { + 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(()) + } + + /// The download loop: resume where a partial left off, write, report. + fn fetch(&self, run: &Run, target: &Path) -> Result<()> { + let partial = partial_of(target); + if let Some(parent) = target.parent() { + private::create_dir(parent)?; + } + let have = partial.metadata().map(|m| m.len()).unwrap_or(0); + + let url = format!( + "https://huggingface.co/{}/resolve/main/{}", + run.repo, + run.file.replace(' ', "%20") + ); + let mut request = ureq::get(&url).header("User-Agent", USER_AGENT); + if have > 0 { + request = request.header("Range", &format!("bytes={have}-")); + } + let mut response = request.call().with_context(|| format!("GET {url}"))?; + + // A server that ignores Range answers 200 with the whole file, and + // appending to what we have would corrupt it -- so trust the + // status, not the request. + let resumed = response.status() == 206; + let length: Option = response + .headers() + .get("content-length") + .and_then(|v| v.to_str().ok()?.parse().ok()); + let (mut done, total) = if resumed { + (have, length.map(|l| have + l)) + } else { + (0, length) + }; + { + 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 here would + // throw away exactly the bytes the Range request just asked the + // server not to send again. Stated rather than left to the + // default, because the default is what a reader would have to + // remember. + 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 the server would not resume")?; + } + + 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); + + // 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()))?; + Ok(()) + } +} + +/// `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) { + 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, so it is also the one whose view of what exists is relevant. +pub fn search(query: &str) -> Result> { + 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 GGUF files in one repository, largest last, with the ones already +/// downloaded marked. +pub fn files(repo: &str, store: &ModelStore) -> Result> { + 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 = store.list().into_iter().map(|m| m.key).collect(); + let mut files: Vec = 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 { + 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() +} diff --git a/server/src/routes.rs b/server/src/routes.rs index 7bad22e..0470277 100644 --- a/server/src/routes.rs +++ b/server/src/routes.rs @@ -456,3 +456,112 @@ async fn send_event( tx.send(SseEvent::default().id(entry.seq.to_string()).data(data)) .await } + +/// Separate router because its state is the model store, like `usage`'s. +/// +/// Keys are `owner/repo/file.gguf` and so contain slashes, which is why +/// nothing here puts one in the path: a key travels in the body or a query +/// string, and the routes stay addressable without escaping rules nobody +/// would get right from a phone. +pub fn models_router(store: Arc) -> Router { + Router::new() + .route("/models", get(list_models)) + .route("/models/search", get(search_models)) + .route("/models/files", get(repo_files)) + .route("/models/download", post(start_download)) + .route("/models/cancel", post(cancel_download)) + .route("/models/delete", post(delete_model)) + .with_state(store) +} + +/// What this machine has and what it is fetching, in one answer. +/// +/// Both together deliberately: a phone showing the model list needs both +/// to draw one screen, and two routes would let it render a model as +/// absent while its download sits at 99%. +#[derive(serde::Serialize)] +#[serde(rename_all = "camelCase")] +struct ModelsResponse { + local: Vec, + downloads: Vec, +} + +async fn list_models( + State(store): State>, +) -> Result, ApiError> { + let listing = tokio::task::spawn_blocking(move || ModelsResponse { + local: store.list(), + downloads: store.downloads(), + }) + .await + .context("listing models panicked")?; + Ok(axum::Json(listing)) +} + +#[derive(Deserialize)] +struct SearchQuery { + q: String, +} + +async fn search_models( + Query(query): Query, +) -> Result>, ApiError> { + // Blocking HTTP, like the usage fetch: off the request workers. + let found = tokio::task::spawn_blocking(move || crate::models::search(&query.q)) + .await + .context("model search panicked")? + .map_err(bad_request)?; + Ok(axum::Json(found)) +} + +#[derive(Deserialize)] +struct RepoQuery { + repo: String, +} + +async fn repo_files( + State(store): State>, + Query(query): Query, +) -> Result>, ApiError> { + let files = tokio::task::spawn_blocking(move || crate::models::files(&query.repo, &store)) + .await + .context("listing repository files panicked")? + .map_err(bad_request)?; + Ok(axum::Json(files)) +} + +#[derive(Deserialize)] +struct DownloadRequest { + repo: String, + file: String, +} + +/// Starts a download, or rejoins the one already running for that model. +async fn start_download( + State(store): State>, + axum::Json(body): axum::Json, +) -> Result, ApiError> { + let status = store.start(&body.repo, &body.file).map_err(bad_request)?; + Ok(axum::Json(status)) +} + +#[derive(Deserialize)] +struct KeyRequest { + key: String, +} + +async fn cancel_download( + State(store): State>, + axum::Json(body): axum::Json, +) -> Result, ApiError> { + let status = store.cancel(&body.key).map_err(bad_request)?; + Ok(axum::Json(status)) +} + +async fn delete_model( + State(store): State>, + axum::Json(body): axum::Json, +) -> Result { + store.delete(&body.key).map_err(bad_request)?; + Ok(StatusCode::NO_CONTENT) +}