Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 3 additions & 0 deletions crates/goose-local-inference/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment on lines +55 to +57

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this new dependency necesarry?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, the two new tests set GOOSE_PATH_ROOT to a tempdir so delete_model/sync_with_featured don't write to the real goose data dir, and env_lock::lock_env serializes the env-var mutation across parallel tests (plain std::env::set_var would race). It's dev-dependencies only, and env-lock is already a workspace dependency used with this same pattern in other crates (e.g. goose/src/config/declarative_providers.rs tests), so this just reuses the existing convention rather than adding anything new to the tree.

6 changes: 3 additions & 3 deletions crates/goose-local-inference/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
Expand Down
69 changes: 67 additions & 2 deletions crates/goose-local-inference/src/local_model_registry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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!(
Expand Down
3 changes: 2 additions & 1 deletion crates/goose-local-inference/src/management.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand All @@ -309,7 +310,7 @@ pub async fn ensure_featured_models_current() -> Result<()> {
}
}

if !needs_backfill {
if !needs_backfill && !needs_size {
continue;
}
}
Expand Down
Loading