diff --git a/Cargo.lock b/Cargo.lock index befb0adfcc85..78770e95650b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5129,6 +5129,7 @@ dependencies = [ "base64 0.22.1", "chrono", "encoding_rs", + "env-lock", "etcetera 0.11.0", "fs2", "futures", diff --git a/crates/goose-local-inference/Cargo.toml b/crates/goose-local-inference/Cargo.toml index c3ae34274720..402d6176dcff 100644 --- a/crates/goose-local-inference/Cargo.toml +++ b/crates/goose-local-inference/Cargo.toml @@ -52,3 +52,6 @@ safemlx-lm-utils = { optional = true, version = "0.1.2" } [target.'cfg(target_os = "macos")'.dependencies] llama-cpp-2 = { workspace = true, features = ["sampler", "metal", "mtmd"] } + +[dev-dependencies] +env-lock.workspace = true diff --git a/crates/goose-local-inference/src/lib.rs b/crates/goose-local-inference/src/lib.rs index e36ee9bed0fa..2bca144deed8 100644 --- a/crates/goose-local-inference/src/lib.rs +++ b/crates/goose-local-inference/src/lib.rs @@ -314,13 +314,13 @@ pub fn recommend_local_model(runtime: &InferenceRuntime) -> String { let mut models: Vec<_> = registry .list_models() .iter() - .filter(|m| is_featured_model(&m.id) && m.size_bytes > 0) + .filter(|m| is_featured_model(&m.id) && m.file_size() > 0) .collect(); - models.sort_by_key(|model| std::cmp::Reverse(model.size_bytes)); + models.sort_by_key(|model| std::cmp::Reverse(model.file_size())); // Return largest that fits in available memory for model in &models { - if available_memory >= model.size_bytes { + if available_memory >= model.file_size() { return model.id.clone(); } } diff --git a/crates/goose-local-inference/src/local_model_registry.rs b/crates/goose-local-inference/src/local_model_registry.rs index 893d4e369b48..e4491fd23395 100644 --- a/crates/goose-local-inference/src/local_model_registry.rs +++ b/crates/goose-local-inference/src/local_model_registry.rs @@ -517,7 +517,12 @@ impl LocalModelRegistry { let mut changed = false; for mut entry in featured_entries { - if !self.models.iter().any(|m| m.id == entry.id) { + if let Some(existing) = self.models.iter_mut().find(|m| m.id == entry.id) { + if existing.size_bytes == 0 && entry.size_bytes > 0 { + existing.size_bytes = entry.size_bytes; + changed = true; + } + } else { entry.enrich_with_featured_mmproj(); self.models.push(entry); changed = true; @@ -560,7 +565,6 @@ impl LocalModelRegistry { if let Some(entry) = self.models.iter_mut().find(|m| m.id == id) { entry.local_path = Paths::in_data_dir("models").join(&entry.filename); entry.storage = LocalModelStorage::GooseManaged; - entry.size_bytes = 0; entry.shard_files.clear(); } self.save() @@ -727,6 +731,67 @@ mod tests { get_download_manager().clear_completed(&download_id); } + #[test] + fn delete_featured_model_preserves_size_bytes() { + let root = tempfile::tempdir().unwrap(); + let _guard = env_lock::lock_env([("GOOSE_PATH_ROOT", Some(root.path().to_str().unwrap()))]); + + let mut entry = test_entry("unsloth/gemma-4-E4B-it-GGUF:Q4_K_M"); + entry.local_path = root.path().join("model.gguf"); + entry.size_bytes = 12345; + + let mut registry = LocalModelRegistry { + models: vec![entry], + }; + registry + .delete_model("unsloth/gemma-4-E4B-it-GGUF:Q4_K_M") + .unwrap(); + + let kept = registry + .get_model("unsloth/gemma-4-E4B-it-GGUF:Q4_K_M") + .expect("featured model entry should remain after deletion"); + assert_eq!(kept.size_bytes, 12345); + assert!(!kept.is_downloaded()); + } + + #[test] + fn sync_with_featured_backfills_size_without_touching_settings() { + let root = tempfile::tempdir().unwrap(); + let _guard = env_lock::lock_env([("GOOSE_PATH_ROOT", Some(root.path().to_str().unwrap()))]); + + let mut existing = test_entry("unsloth/gemma-4-E4B-it-GGUF:Q4_K_M"); + existing.settings.enable_thinking = true; + existing.storage = LocalModelStorage::HuggingFaceCache; + + let mut resolved = test_entry("unsloth/gemma-4-E4B-it-GGUF:Q4_K_M"); + resolved.size_bytes = 999; + resolved.settings.enable_thinking = false; + + let mut registry = LocalModelRegistry { + models: vec![existing], + }; + registry.sync_with_featured(vec![resolved]); + + let updated = registry + .get_model("unsloth/gemma-4-E4B-it-GGUF:Q4_K_M") + .unwrap(); + assert_eq!(updated.size_bytes, 999); + assert!(updated.settings.enable_thinking); + assert_eq!(updated.storage, LocalModelStorage::HuggingFaceCache); + + let mut refreshed = test_entry("unsloth/gemma-4-E4B-it-GGUF:Q4_K_M"); + refreshed.size_bytes = 777; + registry.sync_with_featured(vec![refreshed]); + assert_eq!( + registry + .get_model("unsloth/gemma-4-E4B-it-GGUF:Q4_K_M") + .unwrap() + .size_bytes, + 999, + "non-zero size should not be overwritten" + ); + } + #[test] fn enrich_with_featured_mmproj_preserves_existing_downloaded_path() { let existing_path = std::env::temp_dir().join(format!( diff --git a/crates/goose-local-inference/src/management.rs b/crates/goose-local-inference/src/management.rs index 31d90d53ba1b..4b8014d7d967 100644 --- a/crates/goose-local-inference/src/management.rs +++ b/crates/goose-local-inference/src/management.rs @@ -294,6 +294,7 @@ pub async fn ensure_featured_models_current() -> Result<()> { .map_err(|_| anyhow!("Failed to acquire registry lock"))?; if let Some(existing) = registry.get_model(&model_id) { let needs_backfill = existing.mmproj_path.is_none() && featured.mmproj.is_some(); + let needs_size = existing.size_bytes == 0 && !existing.is_downloaded(); let needs_download = existing.is_downloaded() && featured.mmproj.is_some() && !existing.mmproj_path.as_ref().is_some_and(|p| p.exists()); @@ -309,7 +310,7 @@ pub async fn ensure_featured_models_current() -> Result<()> { } } - if !needs_backfill { + if !needs_backfill && !needs_size { continue; } }