iris/ai-app
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'main' of git.arirex.me:iris/ai-app

Browse Source 
# Conflicts:
#	AGENTS.md
#	PLAN.md
#	app/androidApp/src/main/kotlin/com/example/aiapp/SessionUsageBar.kt
#	app/androidApp/src/main/kotlin/com/example/aiapp/SpawnScreen.kt
#	server/src/config.rs
#	server/src/main.rs
#	server/src/routes.rs
#	server/src/session/echo.rs
#	server/src/session/llama.rs
#	server/src/session/transport.rs
#	server/src/ssh.rs
#	server/src/usage.rs
This commit is contained in:
iris committed 2026-09-04 17:56:50 -04:00
2 parents 127b25e60a 1fcaa2d72d
commit 3c0214ece8
94 files changed
+8299 -9454
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/models.rs
+84 -118
View File
@@ -1,28 +1,22 @@
//! 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:
//! 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, 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.
//! 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 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.
//! **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, rather
//! than a bar drawn from how long the last download took.
//! `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};
@@ -42,35 +36,33 @@ use crate::session::transport::{Launch, Transport};
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.
/// 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.
/// `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.
/// 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.
/// 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,
@@ -82,20 +74,17 @@ pub enum DownloadState {
#[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.
/// 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,
/// 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.
/// Absent means "unknown", never "zero".
#[serde(skip_serializing_if = "Option::is_none")]
pub total: Option<u64>,
/// Present only when [`DownloadState::Failed`], and it is the reason.
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
pub started: f64,
@@ -154,9 +143,8 @@ impl Run {
/// 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.
/// 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,
}
@@ -173,11 +161,10 @@ impl ModelStore {
/// 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.
/// 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('/')) {
@@ -193,11 +180,9 @@ impl ModelStore {
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.
/// 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);
@@ -213,12 +198,10 @@ impl ModelStore {
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.
/// 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)?;
@@ -253,8 +236,8 @@ impl ModelStore {
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.
// 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);
@@ -277,8 +260,8 @@ impl ModelStore {
Ok(status)
}
/// Asks a running download to stop. The partial file stays, so
/// starting again resumes rather than refetching.
/// 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 {
@@ -305,7 +288,6 @@ impl ModelStore {
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);
let identity = identity_of(target);
@@ -313,9 +295,8 @@ impl ModelStore {
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 -- so it is refetched rather than guessed at.
// 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());
@@ -333,13 +314,11 @@ impl ModelStore {
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 rather than 200 with the whole file. 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. Throw the partial away and
// ask again from zero.
// 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",
@@ -351,11 +330,9 @@ impl ModelStore {
etag = etag_of(&response);
}
// On a 206, Content-Length is the length of the *range*, not of
// the file -- it answers a different question than the one a
// progress bar asks, and 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 (`bytes 162000000-234074815/234074816`), which has
// 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
@@ -380,12 +357,10 @@ impl ModelStore {
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.
// `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)
@@ -399,9 +374,8 @@ impl ModelStore {
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. That is
// what makes it safe to keep one across a restart of this server.
// 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();
}
@@ -427,12 +401,10 @@ impl ModelStore {
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 download of
// this size has too many ways to go subtly wrong to take on
// trust, and because a wrong model is the kind of failure that
// surfaces as bad output rather than as an error.
// 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)?;
@@ -448,8 +420,8 @@ impl ModelStore {
}
}
// Renamed only once complete, so a file at its real name is always
// a whole model -- `list` needs no other way to tell.
// 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();
@@ -457,9 +429,8 @@ impl ModelStore {
}
}
/// 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 ever makes.
/// 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 =
@@ -473,8 +444,8 @@ fn sha256_of(path: &Path) -> Result<String> {
}
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.
// 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()
@@ -489,9 +460,8 @@ fn request(url: &str, from: u64) -> Result<(ureq::http::Response<ureq::Body>, bo
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.
// 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))
}
@@ -508,8 +478,8 @@ fn etag_of(response: &ureq::http::Response<ureq::Body>) -> Option<String> {
)
}
/// `x.gguf` -> `x.gguf.part.etag`, holding which revision the partial
/// beside it is a piece of.
/// `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");
@@ -646,18 +616,17 @@ pub struct RemoteRepo {
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.
/// 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.
/// 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",
@@ -685,11 +654,9 @@ pub fn search(query: &str) -> Result<Vec<RemoteRepo>> {
.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 on a download rather
/// than something we would have to compute a second source of truth for.
/// 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()?;
@@ -738,10 +705,9 @@ fn get_json(url: &str) -> Result<serde_json::Value> {
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.
/// 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()