From de247ad54f2b3b0fcecbfab02cdeda0bde9b0279 Mon Sep 17 00:00:00 2001 From: Andrey Vyrodov Date: Sat, 10 Oct 2026 01:20:51 +0200 Subject: [PATCH] Issue #345: model downloads in the activity stack Downloading the SAM 3, face and AI Denoise models shows a row with the model's name and the megabytes so far; its cross stops the download like the Settings button. The denoise download tests get their own scratch folders (they shared the face tests' names and could delete each other's files). Co-Authored-By: Claude Opus 5.5 --- crates/engine/src/cmd/denoise.rs | 2 +- crates/engine/src/cmd/face_models.rs | 8 ++- crates/engine/src/cmd/masks.rs | 2 +- crates/engine/src/face_download.rs | 80 ++++++++++++++++++++++++-- crates/engine/src/model_download.rs | 82 +++++++++++++++++++++++++-- crates/engine/src/segment/download.rs | 23 ++++++-- crates/engine/src/segment/mod.rs | 6 +- crates/engine/src/tests_segment.rs | 48 ++++++++++++++++ 8 files changed, 232 insertions(+), 19 deletions(-) diff --git a/crates/engine/src/cmd/denoise.rs b/crates/engine/src/cmd/denoise.rs index e1cc0900c..6962f9407 100644 --- a/crates/engine/src/cmd/denoise.rs +++ b/crates/engine/src/cmd/denoise.rs @@ -368,7 +368,7 @@ fn download(s: &mut Session, p: &Value) -> Result { return Err(bad(C, "that model is already installed")); } let host = lightcraft_denoise::known::host(&spec.url).to_string(); - s.denoise.downloads.start(spec, &dir).map_err(|e| bad(C, e))?; + s.denoise.downloads.start(spec, &dir, &s.activity, &k.manifest.name).map_err(|e| bad(C, e))?; Ok(json!({"started": id, "from": host})) } diff --git a/crates/engine/src/cmd/face_models.rs b/crates/engine/src/cmd/face_models.rs index 44b359e4a..3a8e401dd 100644 --- a/crates/engine/src/cmd/face_models.rs +++ b/crates/engine/src/cmd/face_models.rs @@ -356,7 +356,13 @@ fn download(s: &mut Session, p: &Value) -> Result { return Err(bad(C, "that model is already installed")); } let host = spec.host().to_string(); - s.face_downloads.start(spec, &dir).map_err(|e| bad(C, e))?; + let name = known::all() + .into_iter() + .find(|m| m.id == spec.id) + .map(|m| m.name) + .or_else(|| s.face_catalog.entries.iter().find(|e| e.manifest.id == spec.id).map(|e| e.manifest.name.clone())) + .unwrap_or_else(|| spec.id.clone()); + s.face_downloads.start(spec, &dir, &s.activity, &name).map_err(|e| bad(C, e))?; Ok(json!({"started": id, "from": host})) } diff --git a/crates/engine/src/cmd/masks.rs b/crates/engine/src/cmd/masks.rs index d72996523..24a94cbc8 100644 --- a/crates/engine/src/cmd/masks.rs +++ b/crates/engine/src/cmd/masks.rs @@ -464,7 +464,7 @@ pub fn specs() -> Vec { ), )); } - let started = s.segmenter.start_download().map_err(ai_err)?; + let started = s.segmenter.start_download(&s.activity).map_err(ai_err)?; Ok(json!({"started": started, "installed": s.segmenter.installed(), "downloading": s.segmenter.download_status().running})) } ), diff --git a/crates/engine/src/face_download.rs b/crates/engine/src/face_download.rs index 0c6ed4ee1..df79d5630 100644 --- a/crates/engine/src/face_download.rs +++ b/crates/engine/src/face_download.rs @@ -15,6 +15,8 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex, PoisonError}; use std::time::Duration; +use crate::activity::{Activity, Cancel, TaskHandle, Unit}; + use lightcraft_faces::known::Download; /// The folder inside the models folder where a downloaded file waits to be installed. Its name starts with a dot, so it @@ -42,15 +44,20 @@ pub enum State { struct Shared { state: Mutex, - cancel: AtomicBool, + cancel: Arc, + /// The download's row in the activity stack. + task: Option, } impl Shared { fn new(total: u64) -> Self { - Shared { state: Mutex::new(State::Running { bytes: 0, total }), cancel: AtomicBool::new(false) } + Shared { state: Mutex::new(State::Running { bytes: 0, total }), cancel: Arc::new(AtomicBool::new(false)), task: None } } fn set(&self, s: State) { + if let (State::Running { bytes, total }, Some(task)) = (&s, &self.task) { + task.progress(*bytes, *total); + } *self.state.lock().unwrap_or_else(PoisonError::into_inner) = s; } @@ -68,18 +75,27 @@ pub struct Downloads { impl Downloads { /// Start fetching `spec` into `/.downloads/`. Different models can download side by side; the same /// one is started once. - pub fn start(&mut self, spec: Download, models_dir: &Path) -> Result<(), String> { + /// The activity stack shows it (as `name`) until the thread ends; its ✕ stops it like [`Downloads::discard`]. + pub fn start(&mut self, spec: Download, models_dir: &Path, activity: &Activity, name: &str) -> Result<(), String> { if self.jobs.get(spec.id.as_str()).is_some_and(|j| matches!(j.get(), State::Running { .. })) { return Err("that model is already downloading".into()); } let staging = models_dir.join(STAGING); std::fs::create_dir_all(&staging).map_err(|e| format!("could not create the download folder: {e}"))?; - let shared = Arc::new(Shared::new(spec.size_bytes)); + let cancel = Arc::new(AtomicBool::new(false)); + let task = activity.start("download", "Downloading model", Cancel::Flag(cancel.clone())); + task.set_unit(Unit::Bytes); + task.detail(name); + task.progress(0, spec.size_bytes); + let total = spec.size_bytes; + let shared = Arc::new(Shared { state: Mutex::new(State::Running { bytes: 0, total }), cancel, task: Some(task.handle()) }); let worker = shared.clone(); let id = spec.id.clone(); std::thread::Builder::new() .name("face-model-download".into()) .spawn(move || { + // the row goes when the thread ends, however it ends + let _task = task; let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| run(&spec, spec.base(), &staging, &worker))); if outcome.is_err() { worker.set(State::Failed("the download stopped unexpectedly".into())); @@ -267,6 +283,62 @@ mod tests { base } + /// A local server that sends `body` a few kilobytes at a time: a download that keeps running for a while. + fn serve_slowly(body: Vec) -> String { + let l = TcpListener::bind("127.0.0.1:0").unwrap(); + let base = format!("http://{}/m", l.local_addr().unwrap()); + std::thread::spawn(move || { + for mut s in l.incoming().flatten() { + let mut r = BufReader::new(s.try_clone().unwrap()); + loop { + let mut line = String::new(); + if r.read_line(&mut line).unwrap_or(0) == 0 || line.trim().is_empty() { + break; + } + } + let _ = s.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len()).as_bytes()); + for chunk in body.chunks(4096) { + if s.write_all(chunk).is_err() { + break; + } + std::thread::sleep(Duration::from_millis(20)); + } + } + }); + base + } + + /// Poll until `done` holds (or fail after 10 s). + fn wait(what: &str, done: impl Fn() -> bool) { + let t = std::time::Instant::now(); + while !done() { + assert!(t.elapsed() < Duration::from_secs(10), "{what}"); + std::thread::sleep(Duration::from_millis(10)); + } + } + + #[test] + fn a_running_download_shows_in_activity_and_cancels() { + use crate::activity::{Activity, Unit}; + let dir = scratch("activity"); + let bytes: Vec = (0..400_000u32).map(|i| (i % 251) as u8).collect(); + let mut spec = spec_for(&bytes); + spec.url = format!("{}/{}", serve_slowly(bytes), spec.file_name); + let activity = Activity::default(); + let mut d = Downloads::default(); + d.start(spec, &dir, &activity, "Test Model").unwrap(); + let rows = activity.list(); + assert_eq!(rows.len(), 1); + assert_eq!((rows[0].kind, rows[0].label.as_str(), rows[0].detail.as_str()), ("download", "Downloading model", "Test Model")); + assert_eq!((rows[0].unit, rows[0].total), (Unit::Bytes, 400_000)); + assert!(rows[0].cancellable); + activity.cancel(rows[0].id).unwrap(); + wait("the download never stopped", || !d.running()); + assert_eq!(d.jobs["test-model"].get(), State::Cancelled); + wait("the row stayed", || activity.list().is_empty()); + let _ = std::fs::remove_dir_all(&dir); + } + #[test] fn fetches_checks_and_stages_a_file() { let dir = scratch("ok"); diff --git a/crates/engine/src/model_download.rs b/crates/engine/src/model_download.rs index 98707f78f..96f5eb1c7 100644 --- a/crates/engine/src/model_download.rs +++ b/crates/engine/src/model_download.rs @@ -15,6 +15,8 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex, PoisonError}; use std::time::Duration; +use crate::activity::{Activity, Cancel, TaskHandle, Unit}; + use lightcraft_denoise::known::Download; /// The folder inside the models folder where a downloaded file waits to be installed. Its name starts with a dot, so it @@ -43,15 +45,20 @@ pub enum State { struct Shared { state: Mutex, - cancel: AtomicBool, + cancel: Arc, + /// The download's row in the activity stack. + task: Option, } impl Shared { fn new(total: u64) -> Self { - Shared { state: Mutex::new(State::Running { bytes: 0, total }), cancel: AtomicBool::new(false) } + Shared { state: Mutex::new(State::Running { bytes: 0, total }), cancel: Arc::new(AtomicBool::new(false)), task: None } } fn set(&self, s: State) { + if let (State::Running { bytes, total }, Some(task)) = (&s, &self.task) { + task.progress(*bytes, *total); + } *self.state.lock().unwrap_or_else(PoisonError::into_inner) = s; } @@ -69,18 +76,27 @@ pub struct Downloads { impl Downloads { /// Start fetching `spec` into `/.downloads/`. Different models can download side by side; the same /// one is started once. - pub fn start(&mut self, spec: Download, models_dir: &Path) -> Result<(), String> { + /// The activity stack shows it (as `name`) until the thread ends; its ✕ stops it like [`Downloads::discard`]. + pub fn start(&mut self, spec: Download, models_dir: &Path, activity: &Activity, name: &str) -> Result<(), String> { if self.jobs.get(&spec.id).is_some_and(|j| matches!(j.get(), State::Running { .. })) { return Err("that model is already downloading".into()); } let staging = models_dir.join(STAGING); std::fs::create_dir_all(&staging).map_err(|e| format!("could not create the download folder: {e}"))?; - let shared = Arc::new(Shared::new(spec.size_bytes)); + let cancel = Arc::new(AtomicBool::new(false)); + let task = activity.start("download", "Downloading model", Cancel::Flag(cancel.clone())); + task.set_unit(Unit::Bytes); + task.detail(name); + task.progress(0, spec.size_bytes); + let total = spec.size_bytes; + let shared = Arc::new(Shared { state: Mutex::new(State::Running { bytes: 0, total }), cancel, task: Some(task.handle()) }); let worker = shared.clone(); let id = spec.id.clone(); std::thread::Builder::new() .name("denoise-model-download".into()) .spawn(move || { + // the row goes when the thread ends, however it ends + let _task = task; let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| run(&spec, &spec.url, &staging, &worker))); if outcome.is_err() { worker.set(State::Failed("the download stopped unexpectedly".into())); @@ -218,7 +234,7 @@ mod tests { use std::net::TcpListener; fn scratch(name: &str) -> PathBuf { - let d = std::env::temp_dir().join(format!("lc-face-dl-{name}-{}", std::process::id())); + let d = std::env::temp_dir().join(format!("lc-denoise-dl-{name}-{}", std::process::id())); let _ = std::fs::remove_dir_all(&d); std::fs::create_dir_all(&d).unwrap(); d @@ -254,6 +270,62 @@ mod tests { base } + /// A local server that sends `body` a few kilobytes at a time: a download that keeps running for a while. + fn serve_slowly(body: Vec) -> String { + let l = TcpListener::bind("127.0.0.1:0").unwrap(); + let base = format!("http://{}/m", l.local_addr().unwrap()); + std::thread::spawn(move || { + for mut s in l.incoming().flatten() { + let mut r = BufReader::new(s.try_clone().unwrap()); + loop { + let mut line = String::new(); + if r.read_line(&mut line).unwrap_or(0) == 0 || line.trim().is_empty() { + break; + } + } + let _ = s.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len()).as_bytes()); + for chunk in body.chunks(4096) { + if s.write_all(chunk).is_err() { + break; + } + std::thread::sleep(Duration::from_millis(20)); + } + } + }); + base + } + + /// Poll until `done` holds (or fail after 10 s). + fn wait(what: &str, done: impl Fn() -> bool) { + let t = std::time::Instant::now(); + while !done() { + assert!(t.elapsed() < Duration::from_secs(10), "{what}"); + std::thread::sleep(Duration::from_millis(10)); + } + } + + #[test] + fn a_running_download_shows_in_activity_and_cancels() { + use crate::activity::{Activity, Unit}; + let dir = scratch("activity"); + let bytes: Vec = (0..400_000u32).map(|i| (i % 251) as u8).collect(); + let mut spec = spec_for(&bytes); + spec.url = format!("{}/{}", serve_slowly(bytes), spec.file_name); + let activity = Activity::default(); + let mut d = Downloads::default(); + d.start(spec, &dir, &activity, "Test Model").unwrap(); + let rows = activity.list(); + assert_eq!(rows.len(), 1); + assert_eq!((rows[0].kind, rows[0].label.as_str(), rows[0].detail.as_str()), ("download", "Downloading model", "Test Model")); + assert_eq!((rows[0].unit, rows[0].total), (Unit::Bytes, 400_000)); + assert!(rows[0].cancellable); + activity.cancel(rows[0].id).unwrap(); + wait("the download never stopped", || !d.running()); + assert_eq!(d.jobs["test-model"].get(), State::Cancelled); + wait("the row stayed", || activity.list().is_empty()); + let _ = std::fs::remove_dir_all(&dir); + } + #[test] fn fetches_checks_and_stages_a_file() { let dir = scratch("ok"); diff --git a/crates/engine/src/segment/download.rs b/crates/engine/src/segment/download.rs index 875d652f8..7147db09d 100644 --- a/crates/engine/src/segment/download.rs +++ b/crates/engine/src/segment/download.rs @@ -8,6 +8,8 @@ use std::sync::{Arc, Mutex}; use lightcraft_segment::fetch::{self, Options, Progress}; +use crate::activity::{Activity, Cancel, Unit}; + /// What the session sees of a download. #[derive(Clone, Debug, Default, PartialEq, Eq, serde::Serialize)] pub struct Status { @@ -31,7 +33,7 @@ struct Text { #[derive(Default)] struct Shared { running: AtomicBool, - cancel: AtomicBool, + cancel: Arc, finished: AtomicBool, done: AtomicU64, total: AtomicU64, @@ -92,9 +94,17 @@ impl Downloader { running } - /// Start downloading `files` from `mirrors` into `dir` on a background thread. False when - /// one is already running. - pub fn start(&self, files: &'static [fetch::FileSpec], mirrors: Vec, dir: PathBuf, opts: Options) -> Result { + /// Start downloading `files` from `mirrors` into `dir` on a background thread, shown in the activity stack as + /// `name` (its ✕ cancels like [`Downloader::cancel`]). False when one is already running. + pub fn start( + &self, + files: &'static [fetch::FileSpec], + mirrors: Vec, + dir: PathBuf, + opts: Options, + activity: &Activity, + name: &str, + ) -> Result { let s = self.shared.clone(); if s.running.swap(true, Ordering::SeqCst) { return Ok(false); @@ -106,6 +116,10 @@ impl Downloader { // nothing else runs while `running` was false *lock(&s.text) = Text::default(); *lock(&self.last) = Text::default(); + let task = activity.start("download", "Downloading model", Cancel::Flag(s.cancel.clone())); + task.set_unit(Unit::Bytes); + task.detail(name); + task.progress(0, s.total.load(Ordering::SeqCst)); let guard = Running(s.clone()); let spawned = std::thread::Builder::new().name("sam3-download".into()).spawn(move || { let s = guard.0.clone(); @@ -113,6 +127,7 @@ impl Downloader { let mut on_progress = |p: &Progress| { s.done.store(p.done, Ordering::SeqCst); s.total.store(p.total, Ordering::SeqCst); + task.progress(p.done, p.total); if p.file != last_file { last_file.clone_from(&p.file); lock(&s.text).file.clone_from(&p.file); diff --git a/crates/engine/src/segment/mod.rs b/crates/engine/src/segment/mod.rs index c21873f55..ec2c4ba85 100644 --- a/crates/engine/src/segment/mod.rs +++ b/crates/engine/src/segment/mod.rs @@ -259,7 +259,7 @@ impl Segmenter { /// Start downloading the model (on a background thread; see `segment.model.download`). /// `Ok(false)` when it is installed or downloading already. - pub fn start_download(&self) -> Result { + pub fn start_download(&self, activity: &crate::activity::Activity) -> Result { if !Self::AVAILABLE { return Err("AI masks are not available in this build".into()); } @@ -273,11 +273,11 @@ impl Segmenter { if mirrors.is_empty() { return Err(lightcraft_segment::fetch::no_mirrors_message()); } - self.download.start(lightcraft_segment::fetch::SAM3_FILES, mirrors, dir, lightcraft_segment::fetch::options()) + self.download.start(lightcraft_segment::fetch::SAM3_FILES, mirrors, dir, lightcraft_segment::fetch::options(), activity, "SAM 3") } #[cfg(not(feature = "sam"))] { - let _ = dir; + let _ = (dir, activity); Ok(false) } } diff --git a/crates/engine/src/tests_segment.rs b/crates/engine/src/tests_segment.rs index 59aee18c3..ccca1cb23 100644 --- a/crates/engine/src/tests_segment.rs +++ b/crates/engine/src/tests_segment.rs @@ -230,3 +230,51 @@ fn object_clicks_are_capped() { assert!(!s.segmenter.busy()); let _ = std::fs::remove_dir_all(&dir); } + +#[cfg(feature = "sam")] +#[test] +fn a_running_model_download_shows_in_activity_and_cancels() { + use std::io::{BufRead, BufReader, Write}; + // a local mirror that sends its files slowly, so the download is still running when it is cancelled + let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + let base = format!("http://{}/sam3", l.local_addr().unwrap()); + std::thread::spawn(move || { + for mut c in l.incoming().flatten() { + let mut r = BufReader::new(c.try_clone().unwrap()); + let mut line = String::new(); + while r.read_line(&mut line).unwrap_or(0) > 2 { + line.clear(); + } + let _ = c.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 1000000\r\nConnection: close\r\n\r\n"); + for _ in 0..250 { + if c.write_all(&[b' '; 4000]).is_err() { + break; + } + std::thread::sleep(Duration::from_millis(20)); + } + } + }); + let dir = tmp("download-activity"); + let mirrors = dir.with_extension("activity-mirrors.txt"); + std::fs::write(&mirrors, format!("# test mirror\n{base}\n")).unwrap(); + let mut s = Session::with_demo(); + s.segmenter.dir = Some(dir.clone()); + s.segmenter.mirrors_file = Some(mirrors.clone()); + assert_eq!(s.execute("segment.model.download", &json!({"acknowledged": true})).unwrap()["started"], true); + let rows = s.activity.list(); + assert_eq!(rows.len(), 1); + assert_eq!((rows[0].kind, rows[0].label.as_str(), rows[0].detail.as_str()), ("download", "Downloading model", "SAM 3")); + assert_eq!(rows[0].unit, crate::activity::Unit::Bytes); + assert!(rows[0].cancellable); + s.execute("activity.cancel", &json!({"id": rows[0].id})).unwrap(); + let t = Instant::now(); + while s.segmenter.download_status().running || !s.activity.list().is_empty() { + assert!(t.elapsed() < Duration::from_secs(10), "the download never stopped"); + std::thread::sleep(Duration::from_millis(10)); + } + let e = s.segmenter.download_status().error.unwrap_or_default(); + assert!(e.contains("cancelled"), "{e}"); + assert!(!s.segmenter.installed()); + let _ = std::fs::remove_dir_all(&dir); + let _ = std::fs::remove_file(&mirrors); +}