From 60b84aed84a172e2bc7f82fbda4a6ecd4a74d4a9 Mon Sep 17 00:00:00 2001 From: Jeremy Date: Tue, 27 Oct 2026 20:49:44 +0000 Subject: [PATCH] fix(ai-client): keep linked model folders out of downloads (#109) --- crates/ai-client/src/folders.rs | 54 +++++++------------ crates/ai-client/tests/folders.rs | 21 +++++--- crates/ai-client/tests/link_default_safety.rs | 54 +++++++++++++++++++ .../tests/native_setup_support/mod.rs | 3 +- crates/ai-core/src/model_folders.rs | 28 +++++----- docs/adopting/use-in-your-app.md | 4 +- tools/mutation/mutation_check.py | 2 +- tools/mutation/mutation_first_run.py | 9 ++-- 8 files changed, 112 insertions(+), 63 deletions(-) create mode 100644 crates/ai-client/tests/link_default_safety.rs diff --git a/crates/ai-client/src/folders.rs b/crates/ai-client/src/folders.rs index da8e3a1..ec74c70 100644 --- a/crates/ai-client/src/folders.rs +++ b/crates/ai-client/src/folders.rs @@ -33,12 +33,23 @@ impl FolderDesk { Self { shared } } + /// Reads the shared folder map without changing the registered folders. pub fn list(&self) -> Result, AiError> { Ok(self.shared.load()?.model_folders) } - /// Registers a folder the person chose and reads what is in it. The files are not touched. + /// Links a folder for discovery without choosing it for downloads. pub fn add(&self, path: &Path, label: Option<&str>) -> Result { + self.add_inner(path, label, false) + } + + /// Registers a scanned folder, optionally choosing it as the first download destination. + fn add_inner( + &self, + path: &Path, + label: Option<&str>, + choose_if_missing: bool, + ) -> Result { // The path is stored as text, so one that is not valid text cannot be kept faithfully: // refuse it rather than store a different folder from the one chosen. let text = path @@ -69,15 +80,11 @@ impl FolderDesk { )); } let id = unique_id(&name, &store.model_folders); - let default = !store - .model_folders - .iter() - .any(|f| f.default && f.origin == FolderOrigin::Chosen); let folder = ModelFolder { id, label: name.clone(), path: text.clone(), - default, + default: choose_if_missing && !store.model_folders.iter().any(|f| f.default), origin: FolderOrigin::Chosen, kinds: scan.kinds.clone(), }; @@ -101,7 +108,7 @@ impl FolderDesk { } std::fs::create_dir_all(path) .map_err(|_| AiError::new(ErrorKind::Internal, "that folder could not be created"))?; - self.add(path, label) + self.add_inner(path, label, true) } /// The folder new downloads go to, made if there is none yet: the folder the suite suggests @@ -128,9 +135,8 @@ impl FolderDesk { } /// [`ensure_default`](Self::ensure_default) with the suggested folder given. Returns the - /// folder a person already made the default as it is. With folders chosen but none marked - /// default, the first becomes the default. With none, `suggested` is made and registered. A - /// folder read from another program is never made the default. + /// folder a person already made the default as it is. Otherwise `suggested` is made and + /// registered. Linked folders are never promoted automatically. pub fn ensure_default_at(&self, suggested: &Path) -> Result { if let Some(folder) = self.chosen_default()? { return Ok(folder); @@ -138,23 +144,12 @@ impl FolderDesk { Ok(self.add_new(suggested, Some("ArtCraft models"))?.folder) } - /// The default among the folders the person chose: the one marked default, else the first - /// chosen folder, which is then marked. `None` when they chose none. Needs no suggestion, so - /// a computer that suggests no place still finds a folder the person already has. + /// Returns only an explicitly selected default, never a merely linked folder. fn chosen_default(&self) -> Result, AiError> { - let chosen: Vec = self + Ok(self .list()? .into_iter() - .filter(|f| f.origin == FolderOrigin::Chosen) - .collect(); - if let Some(default) = chosen.iter().find(|f| f.default) { - return Ok(Some(default.clone())); - } - if let Some(first) = chosen.first() { - self.set_default(&first.id)?; - return self.find(&first.id).map(Some); - } - Ok(None) + .find(|f| f.origin == FolderOrigin::Chosen && f.default)) } /// Reads the folder again, for files added or moved since. Replaces its map. @@ -185,20 +180,11 @@ impl FolderDesk { Ok(()) } - /// Forgets the folder. Its files stay where they are. The default moves to the next folder - /// the person chose, so downloads still have somewhere to go. + /// Forgets the folder without touching its files or choosing a replacement download target. pub fn remove(&self, id: &str) -> Result<(), AiError> { self.find(id)?; self.shared.update(|store| { store.model_folders.retain(|f| f.id != id); - let has_default = store.model_folders.iter().any(|f| f.default); - let next = store - .model_folders - .iter_mut() - .find(|f| f.origin == FolderOrigin::Chosen); - if let (false, Some(next)) = (has_default, next) { - next.default = true; - } })?; Ok(()) } diff --git a/crates/ai-client/tests/folders.rs b/crates/ai-client/tests/folders.rs index 55da128..10698c9 100644 --- a/crates/ai-client/tests/folders.rs +++ b/crates/ai-client/tests/folders.rs @@ -57,12 +57,13 @@ fn listing(root: &Path) -> Vec { out } +/// Discovery registers the folder without selecting it for downloads. #[test] -fn a_folder_is_read_stored_and_the_first_becomes_the_default() { +fn a_linked_folder_is_read_and_stored_without_becoming_default() { let w = world(); let added = desk(&w).add(w.models.path(), Some("My models")).unwrap(); assert_eq!(added.folder.id, "my-models"); - assert!(added.folder.default && added.folder.origin == FolderOrigin::Chosen); + assert!(!added.folder.default && added.folder.origin == FolderOrigin::Chosen); assert_eq!(added.folder.kinds["vae"], ["vae"]); assert_eq!(added.folder.kinds["loras"], ["Lora"]); assert_eq!(added.scan.files.len(), 2); @@ -119,6 +120,7 @@ fn a_rescan_sees_what_was_added_since() { assert_eq!(d.list().unwrap()[0].kinds["diffusion_models"], ["unet"]); } +/// Only an explicitly chosen folder may become the default. #[test] fn the_default_moves_and_a_linked_folder_can_never_be_it() { let w = world(); @@ -126,7 +128,8 @@ fn the_default_moves_and_a_linked_folder_can_never_be_it() { let second = tempfile::tempdir().unwrap(); let first = d.add(w.models.path(), Some("one")).unwrap().folder; let two = d.add(second.path(), Some("two")).unwrap().folder; - assert!(first.default && !two.default); + d.set_default(&first.id).unwrap(); + assert!(!first.default && !two.default); d.set_default(&two.id).unwrap(); let now = d.list().unwrap(); assert_eq!( @@ -155,17 +158,19 @@ fn the_default_moves_and_a_linked_folder_can_never_be_it() { assert!(d.set_default("nope").is_err()); } +/// Removing a default never turns a discovery-only folder into a destination. #[test] -fn removing_the_default_hands_it_to_the_next_chosen_folder() { +fn removing_the_default_does_not_promote_a_linked_folder() { let w = world(); let d = desk(&w); let other = tempfile::tempdir().unwrap(); let first = d.add(w.models.path(), Some("one")).unwrap().folder; d.add(other.path(), Some("two")).unwrap(); + d.set_default(&first.id).unwrap(); d.remove(&first.id).unwrap(); let left = d.list().unwrap(); assert_eq!(left.len(), 1); - assert!(left[0].default, "downloads still have somewhere to go"); + assert!(!left[0].default, "a linked folder is not silently promoted"); assert!(d.remove(&first.id).is_err()); } @@ -409,7 +414,8 @@ fn native_downloads_go_only_to_the_chosen_default_not_a_linked_folder() { .update(|store| store.model_folders[0].default = false) .unwrap(); let chosen = tempfile::tempdir().unwrap(); - d.add(chosen.path(), None).unwrap(); + let chosen_folder = d.add(chosen.path(), None).unwrap(); + d.set_default(&chosen_folder.folder.id).unwrap(); let dir = d.native_download_dir(&feature, "candle").unwrap(); assert_eq!(dir, chosen.path().join("sam3_transformers")); assert!( @@ -423,7 +429,8 @@ fn native_downloads_go_only_to_the_chosen_default_not_a_linked_folder() { fn native_install_requires_licence_and_rehashes_existing_files() { let w = world(); let d = desk(&w); - d.add(w.models.path(), None).unwrap(); + let selected = d.add(w.models.path(), None).unwrap(); + d.set_default(&selected.folder.id).unwrap(); let mut feature = small_native_feature(); let hashes = [ "2d711642b726b04401627ca9fbac32f5c8530fb1903cc4db02258717921a4881", diff --git a/crates/ai-client/tests/link_default_safety.rs b/crates/ai-client/tests/link_default_safety.rs new file mode 100644 index 0000000..e8ee74f --- /dev/null +++ b/crates/ai-client/tests/link_default_safety.rs @@ -0,0 +1,54 @@ +use craft_ai_client::FolderDesk; +use craft_ai_core::model_folders::download_target; +use craft_ai_native::shared::SharedStore; + +/// Linking an existing folder is discovery, not approval to write downloads into it. +#[test] +fn linked_folder_never_becomes_a_download_target_without_an_explicit_choice() { + let temp = tempfile::tempdir().unwrap(); + let linked = temp.path().join("existing"); + std::fs::create_dir(&linked).unwrap(); + let desk = FolderDesk::new(SharedStore::new(temp.path().join("connections.json"))); + std::fs::write(linked.join("keep.bin"), b"model").unwrap(); + let item = desk.add(&linked, None).unwrap().folder; + assert!(!item.default); + assert!(download_target(&desk.list().unwrap(), "checkpoints").is_none()); + let suggested = temp.path().join("new-models"); + let default = desk.ensure_default_at(&suggested).unwrap(); + assert_eq!(default.path, suggested.to_str().unwrap()); + assert_ne!(default.id, item.id); + assert_eq!(std::fs::read(linked.join("keep.bin")).unwrap(), b"model"); + desk.remove(&default.id).unwrap(); + assert!(download_target(&desk.list().unwrap(), "checkpoints").is_none()); + desk.set_default(&item.id).unwrap(); + assert_eq!( + download_target(&desk.list().unwrap(), "checkpoints") + .unwrap() + .0 + .id, + item.id + ); +} + +/// Linking another folder leaves an existing default unchanged across readers. +#[test] +fn linking_a_second_folder_preserves_the_existing_default() { + let temp = tempfile::tempdir().unwrap(); + let desk = FolderDesk::new(SharedStore::new(temp.path().join("connections.json"))); + let existing = desk + .add_new(&temp.path().join("install"), None) + .unwrap() + .folder; + let linked = temp.path().join("discovery"); + std::fs::create_dir(&linked).unwrap(); + let item = desk.add(&linked, None).unwrap().folder; + assert!(!item.default); + let other = FolderDesk::new(SharedStore::new(temp.path().join("connections.json"))); + assert_eq!( + download_target(&other.list().unwrap(), "checkpoints") + .unwrap() + .0 + .id, + existing.id + ); +} diff --git a/crates/ai-client/tests/native_setup_support/mod.rs b/crates/ai-client/tests/native_setup_support/mod.rs index c03af04..94ba548 100644 --- a/crates/ai-client/tests/native_setup_support/mod.rs +++ b/crates/ai-client/tests/native_setup_support/mod.rs @@ -32,7 +32,8 @@ pub fn world() -> (tempfile::TempDir, FolderDesk, PathBuf) { let models = dir.path().join("models"); std::fs::create_dir(&models).unwrap(); let desk = FolderDesk::new(SharedStore::new(dir.path().join("connections.json"))); - desk.add(&models, None).unwrap(); + let chosen = desk.add(&models, None).unwrap(); + desk.set_default(&chosen.folder.id).unwrap(); (dir, desk, models) } diff --git a/crates/ai-core/src/model_folders.rs b/crates/ai-core/src/model_folders.rs index 47465c4..1220212 100644 --- a/crates/ai-core/src/model_folders.rs +++ b/crates/ai-core/src/model_folders.rs @@ -197,19 +197,16 @@ pub fn validate_folders(folders: &[ModelFolder]) -> Result<(), FolderError> { Ok(()) } -/// Where a new file of `kind` should be saved: the default folder (else the first folder the -/// person chose), and inside it the first place already holding that kind, else a sub-folder -/// named for the kind. Linked folders never receive downloads. +/// Where a new file of `kind` should be saved: the explicitly chosen default folder, +/// inside its first place already holding that kind, else a sub-folder named for the kind. +/// Linked folders and non-default folders never receive downloads. pub fn download_target<'a>( folders: &'a [ModelFolder], kind: &str, ) -> Option<(&'a ModelFolder, String)> { - let writable = |f: &&ModelFolder| f.origin == FolderOrigin::Chosen; let folder = folders .iter() - .filter(writable) - .find(|f| f.default) - .or_else(|| folders.iter().find(writable))?; + .find(|f| f.origin == FolderOrigin::Chosen && f.default)?; let sub = match folder.kinds.get(kind).and_then(|dirs| dirs.first()) { Some(mapped) => mapped.clone(), // No place yet: the kind names a new sub-folder, so it must be a safe one. A caller passing @@ -327,6 +324,7 @@ mod tests { ); } + /// Imported or non-default folders cannot become download targets by fallback. #[test] fn downloads_go_to_the_default_folder_and_never_to_a_linked_one() { let mut linked = folder("comfy", "/c/models"); @@ -335,14 +333,9 @@ mod tests { let mut mine = folder("mine", "/m"); mine.kinds.insert("vae".into(), vec!["weights/vae".into()]); let list = [linked, folder("other", "/o"), mine]; - let (f, sub) = download_target(&list, "vae").unwrap(); - assert_eq!( - f.id, "other", - "the default is a linked folder, so the first chosen one is used" - ); - assert_eq!( - sub, "vae", - "a kind with no place yet gets a sub-folder named for it" + assert!( + download_target(&list, "vae").is_none(), + "a linked default must not select another folder" ); let mut list = list; list[2].default = true; @@ -359,9 +352,12 @@ mod tests { ); } + /// Reject unsafe kind names even when an explicit download destination exists. #[test] fn a_kind_that_would_leave_the_folder_is_never_a_place_to_save() { - let list = [folder("mine", "/m")]; + let mut mine = folder("mine", "/m"); + mine.default = true; + let list = [mine]; for bad in ["..", "../x", "/etc", "a/../b", "C:/x", ""] { assert!(download_target(&list, bad).is_none(), "{bad:?}"); } diff --git a/docs/adopting/use-in-your-app.md b/docs/adopting/use-in-your-app.md index 1b6d8ec..f1dbcaf 100644 --- a/docs/adopting/use-in-your-app.md +++ b/docs/adopting/use-in-your-app.md @@ -287,7 +287,9 @@ so you draw them and run the work on a thread: 2. Only when the person takes the suggested folder, call `FolderDesk::ensure_default`, then give `SetupSession::review` a fresh plan. It takes the suggested folder only when the person has none, and never replaces one they chose. Do not call it while drawing the first screen: it - makes the folder. + makes the folder. `FolderDesk::add` links an existing folder for discovery; it does not select + that folder for downloads. Call `set_default` only when the person explicitly chooses it as + the destination. Removing a default never promotes another linked folder automatically. 3. One button calls `SetupSession::start`. It checks the files already on disk when there are complete ones, and downloads when there are not. When it does nothing, `SetupSession::blocker` says why (busy, no folder, agreement missing, this app cannot run it). If a check of the diff --git a/tools/mutation/mutation_check.py b/tools/mutation/mutation_check.py index 8cb7705..97fdcba 100644 --- a/tools/mutation/mutation_check.py +++ b/tools/mutation/mutation_check.py @@ -321,7 +321,7 @@ V1 = [ ('status-reads-one-fixed-app', CLIENT, STATUS_RESOLVE, 'let resolved = match store.resolve(&AppId::new("launcher").unwrap_or_else(|_| self.app.clone()), role) {', CL, 'one_binding', ''), ('unbound-reads-as-broken', ROLE, "Self::Unbound(r) => write!(f, \"no model is chosen for '{r}'\"),", "Self::Unbound(r) => write!(f, \"'{r}' is broken\"),", CL, 'one_binding', ''), ('unbound-hides-the-tool', ROLE, "Self::Unbound(r) => write!(f, \"no model is chosen for '{r}'\"),", "Self::Unbound(_) => write!(f, \"no model is chosen\"),", CL, 'one_binding', ''), - ('choose-writes-nothing', TOOLS, ' set_default(store, binding);\n Ok::<(), AiError>(())', ' let _ = binding;\n Ok::<(), AiError>(())', CL, 'one_binding', ''), + ('choose-writes-nothing', TOOLS, ' set_default(store, binding);\n Ok(true)', ' let _ = binding;\n Ok(true)', CL, 'one_binding', ''), ] PICK_G = C + 'ai-client/src/picker_groups.rs' diff --git a/tools/mutation/mutation_first_run.py b/tools/mutation/mutation_first_run.py index 974c412..1a2ef40 100644 --- a/tools/mutation/mutation_first_run.py +++ b/tools/mutation/mutation_first_run.py @@ -13,9 +13,12 @@ FR = [ ('agreement-always-asked', F + 'recipe.rs', 'self.models.iter().any(|m| m.licence.accept_required)', 'true', CL, 'first_run', 'asks_for_agreement_only_when'), ('download-ignores-a-missing-agreement', C + 'setup_flow.rs', 'if (self.plan.needs_agreement() && !self.accepted)', 'if (false && !self.accepted)', CL, 'first_run', 'asks_for_agreement_only_when'), ('start-always-downloads', C + 'setup_flow.rs', 'if existing.is_some() && existing != self.failed_check.as_deref() {\n SetupTask::CheckExisting', 'if false {\n SetupTask::CheckExisting', CL, 'first_run', 'start_checks_files'), - ('suggested-folder-replaces-a-chosen-default', C + 'folders.rs', 'if let Some(default) = chosen.iter().find(|f| f.default) {', 'if let Some(default) = chosen.iter().find(|f| f.default && false) {', CL, 'first_run', 'suggested_folder_is_made_only_when'), + ('suggested-folder-replaces-a-chosen-default', C + 'folders.rs', '.find(|f| f.origin == FolderOrigin::Chosen && f.default))', '.find(|f| f.origin == FolderOrigin::Chosen && f.default && false))', CL, 'first_run', 'suggested_folder_is_made_only_when'), + ('link-silently-becomes-default', C + 'folders.rs', 'default: choose_if_missing && !store.model_folders.iter().any(|f| f.default),', 'default: !store.model_folders.iter().any(|f| f.default),', CL, 'link_default_safety', 'linked_folder_never_becomes'), + ('download-falls-back-to-linked', 'crates/ai-core/src/model_folders.rs', '.find(|f| f.origin == FolderOrigin::Chosen && f.default)?;', '.find(|f| f.origin == FolderOrigin::Chosen)?;', CL, 'link_default_safety', 'linked_folder_never_becomes'), + ('removal-promotes-a-link', C + 'folders.rs', 'store.model_folders.retain(|f| f.id != id);', 'store.model_folders.retain(|f| f.id != id);\n if let Some(f) = store.model_folders.first_mut() { f.default = true; }', CL, 'link_default_safety', 'linked_folder_never_becomes'), ('adopt-forgets-the-kept-choice', C + 'local.rs', ' self.view(&role)?.bound\n', ' None\n', CL, 'first_run', 'adopting_checked_files'), - ('sample-ignores-the-route', C + 'tools.rs', 'feature.engines().contains(&route.provider.as_str())', 'false && feature.engines().contains(&route.provider.as_str())', CL, 'sample_route', 'built_in_route_is_tried'), + ('sample-ignores-the-route', C + 'tools.rs', '.engine_recipe(&route.provider)', '.engine_recipe("not-a-runner")', CL, 'sample_route', 'built_in_route_is_tried'), ('nothing-selected-still-passes', C + 'tools.rs', 'is_some_and(|found| !found.matched)', 'is_some_and(|found| false && !found.matched)', CL, 'sample_route', 'chose_nothing'), ('status-ignores-a-missing-file', C + 'client.rs', '&& !std::path::Path::new(file).is_file()', '&& false', CL, 'first_run_state', 'file_is_gone_is_not_ready'), ('picker-lists-a-gone-model-twice', C + 'picker_groups.rs', 'if files_gone(connection, role, options) {', 'if false {', CL, 'first_run_state', 'file_is_gone_is_not_ready'), @@ -33,7 +36,7 @@ FR = [ ('missing-result-still-passes', C + 'tools.rs', 'if found.is_none() && role.capability.as_str() == segment::CAPABILITY {', 'if false && role.capability.as_str() == segment::CAPABILITY {', CL, 'sample_missing_result', 'no_result_beside_it'), ('small-card-calls-a-quick-engine-slow', F + 'engine.rs', 'Some(g) if too_small(g, min_vram_mb) && offer.fast_on_processor => (', 'Some(g) if false && too_small(g, min_vram_mb) && offer.fast_on_processor => (', FE, 'engine_small_card', 'too_small_does_not_make'), ('candle-always-reports-a-match', 'crates/ai-engine-candle/src/runner.rs', 'matched: score.is_some(),', 'matched: true,', 'craft-ai-engine-candle', 'lib', 'result_says_whether_anything'), - ('custom-connection-gets-the-click', C + 'tools.rs', '&& route.connection.as_str()\n == crate::local::native_connection_id(&feature.id, &route.provider)', '', CL, 'sample_route', 'custom_connection_with_a_built_in'), + ('custom-connection-gets-the-click', C + 'tools.rs', 'route.connection.as_str()\n == crate::local::native_connection_id(&feature.id, &route.provider)', 'true', CL, 'sample_route', 'custom_connection_with_a_built_in'), ('suggested-folder-path-hidden', 'crates/ai-egui/src/native_setup.rs', 'ui.label(format!("Suggested folder: {}", suggested.to_string_lossy()));', '', EG, 'native_setup_first_run', 'suggested_folder_is_shown'), ('suggested-folder-active-while-busy', 'crates/ai-egui/src/native_setup.rs', 'plan.available() && !session.busy()', 'plan.available()', EG, 'native_setup_first_run', 'busy_session_cannot_take'), ('recovery-uses-old-folder', C + 'setup_flow.rs', 'self.failed_check.as_ref().is_some_and(|failed| {', 'false && self.failed_check.as_ref().is_some_and(|failed| {', CL, 'setup_recovery_folder', 'bad_same_size_file'),