From 1bced6616bf57ea98ff9069f0d0e93a90fd02867 Mon Sep 17 00:00:00 2001 From: Jasper Date: Tue, 7 Jul 2026 14:01:15 -0700 Subject: [PATCH] Switch the local inference provider MLX backend to use the safemlx crate (#10304) --- Cargo.lock | 318 +++------ crates/goose-cli/src/session/mod.rs | 54 +- crates/goose-cli/src/session/output.rs | 6 +- crates/goose-local-inference/Cargo.toml | 8 +- crates/goose-local-inference/src/backend.rs | 1 + crates/goose-local-inference/src/hf_models.rs | 333 +++++++--- crates/goose-local-inference/src/lib.rs | 391 ++++++++--- .../src/llamacpp/inference_emulated_tools.rs | 5 +- .../src/llamacpp/inference_engine.rs | 81 --- .../src/llamacpp/inference_native_tools.rs | 4 +- .../goose-local-inference/src/management.rs | 27 +- crates/goose-local-inference/src/mlx.rs | 629 +++++++++++++++--- .../src/thinking_output.rs | 81 +++ .../src/conversation/message.rs | 1 + .../src/conversation/token_usage.rs | 1 + crates/goose-sdk-types/src/custom_requests.rs | 11 + crates/goose/acp-meta.json | 5 + crates/goose/acp-schema.json | 26 + crates/goose/src/acp/server.rs | 72 +- crates/goose/src/acp/server/config.rs | 17 +- .../goose/src/acp/server/custom_dispatch.rs | 8 + .../goose/src/acp/server/local_inference.rs | 20 + crates/goose/src/agents/agent.rs | 21 +- crates/goose/src/agents/reply_parts.rs | 36 +- .../acp/__tests__/chatNotifications.test.ts | 2 + .../__tests__/chatSessionController.test.ts | 1 + .../sessionNotificationAdapter.test.ts | 10 +- .../acp/adapter/gooseSessionNotifications.ts | 6 +- ui/desktop/src/acp/adapter/shared.ts | 1 + ui/desktop/src/acp/chatSessionStore.ts | 30 + ui/desktop/src/acp/local-inference.ts | 5 + ui/desktop/src/components/BaseChat.tsx | 17 +- .../localInference/LocalInferenceSettings.tsx | 60 +- ui/desktop/src/hooks/useChatSession.ts | 2 + ui/desktop/src/hooks/useChatSessionTypes.ts | 1 + ui/desktop/src/i18n/messages/de.json | 6 + ui/desktop/src/i18n/messages/en.json | 6 + ui/desktop/src/i18n/messages/es.json | 6 + ui/desktop/src/i18n/messages/fr.json | 6 + ui/desktop/src/i18n/messages/hi.json | 6 + ui/desktop/src/i18n/messages/id.json | 6 + ui/desktop/src/i18n/messages/it.json | 6 + ui/desktop/src/i18n/messages/ja.json | 6 + ui/desktop/src/i18n/messages/ko.json | 6 + ui/desktop/src/i18n/messages/ms.json | 6 + ui/desktop/src/i18n/messages/pt.json | 6 + ui/desktop/src/i18n/messages/ru.json | 6 + ui/desktop/src/i18n/messages/tr.json | 6 + ui/desktop/src/i18n/messages/vi.json | 6 + ui/desktop/src/i18n/messages/zh-CN.json | 6 + ui/desktop/src/i18n/messages/zh-TW.json | 6 + ui/desktop/src/types/message.ts | 20 +- ui/desktop/src/utils/conversionUtils.test.ts | 26 + ui/desktop/src/utils/conversionUtils.ts | 18 + ui/sdk/src/generated/client.gen.ts | 10 + ui/sdk/src/generated/index.ts | 7 +- ui/sdk/src/generated/types.gen.ts | 7 +- ui/sdk/src/generated/zod.gen.ts | 6 + 58 files changed, 1835 insertions(+), 646 deletions(-) create mode 100644 crates/goose-local-inference/src/thinking_output.rs create mode 100644 ui/desktop/src/utils/conversionUtils.test.ts diff --git a/Cargo.lock b/Cargo.lock index 543db1ae36..630febf1b0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1201,7 +1201,7 @@ dependencies = [ "bincode", "bytesize", "clircle", - "console 0.16.3", + "console", "content_inspector", "encoding_rs", "flate2", @@ -1255,26 +1255,6 @@ dependencies = [ "serde", ] -[[package]] -name = "bindgen" -version = "0.70.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f49d8fed880d473ea71efb9bf597651e77201bdd4893efe54c9e5d65ae04ce6f" -dependencies = [ - "bitflags 2.11.1", - "cexpr", - "clang-sys", - "itertools 0.13.0", - "log", - "prettyplease", - "proc-macro2", - "quote", - "regex", - "rustc-hash 1.1.0", - "shlex 1.3.0", - "syn 2.0.117", -] - [[package]] name = "bindgen" version = "0.72.1" @@ -2302,8 +2282,8 @@ version = "0.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4529f45438fc25ca048b242d5c48e2d3ce9a521e2a5a9123d9737d8520b030dd" dependencies = [ - "console 0.16.3", - "indicatif 0.18.4", + "console", + "indicatif", "once_cell", "strsim", "textwrap", @@ -2478,19 +2458,6 @@ dependencies = [ "winnow 1.0.3", ] -[[package]] -name = "console" -version = "0.15.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "054ccb5b10f9f2cbf51eb355ca1d05c2d279ce1804688d0db74b4733a5aeafd8" -dependencies = [ - "encode_unicode", - "libc", - "once_cell", - "unicode-width 0.2.2", - "windows-sys 0.59.0", -] - [[package]] name = "console" version = "0.16.3" @@ -5024,7 +4991,7 @@ dependencies = [ "clap_mangen", "cliclack", "comfy-table", - "console 0.16.3", + "console", "dotenvy", "env-lock", "etcetera 0.11.0", @@ -5032,7 +4999,7 @@ dependencies = [ "goose", "goose-mcp", "goose-providers", - "indicatif 0.18.4", + "indicatif", "open", "rand 0.10.1", "regex", @@ -5089,17 +5056,17 @@ dependencies = [ "goose-download-manager", "goose-provider-types", "goose-sdk-types", - "hf-hub 1.0.0-rc.1", + "hf-hub", "include_dir", "llama-cpp-2", "llama-cpp-sys-2", "minijinja", - "mlx-lm", - "mlx-lm-utils", - "mlx-rs", "regex", "reqwest 0.13.4", "rmcp", + "safemlx", + "safemlx-lm", + "safemlx-lm-utils", "serde", "serde_json", "tempfile", @@ -5447,25 +5414,6 @@ dependencies = [ "arrayvec", ] -[[package]] -name = "hf-hub" -version = "0.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "629d8f3bbeda9d148036d6b0de0a3ab947abd08ce90626327fc3547a49d59d97" -dependencies = [ - "dirs", - "http 1.4.2", - "indicatif 0.17.11", - "libc", - "log", - "rand 0.9.4", - "serde", - "serde_json", - "thiserror 2.0.18", - "ureq", - "windows-sys 0.60.2", -] - [[package]] name = "hf-hub" version = "1.0.0-rc.1" @@ -6070,26 +6018,13 @@ dependencies = [ "serde_core", ] -[[package]] -name = "indicatif" -version = "0.17.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235" -dependencies = [ - "console 0.15.11", - "number_prefix", - "portable-atomic", - "unicode-width 0.2.2", - "web-time", -] - [[package]] name = "indicatif" version = "0.18.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "25470f23803092da7d239834776d653104d551bc4d7eacaf31e6837854b8e9eb" dependencies = [ - "console 0.16.3", + "console", "portable-atomic", "unicode-width 0.2.2", "unit-prefix", @@ -6611,7 +6546,7 @@ version = "0.1.146" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b291e4bc2d10c43cd8dec16d49b6104cb3cb125f596ec380a753a5db1d965dd" dependencies = [ - "bindgen 0.72.1", + "bindgen", "cc", "cmake", "find_cuda_helper", @@ -6862,99 +6797,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "mlx-internal-macros" -version = "0.25.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a7c4444d624bf6b93db5cc22ebff4fdfa13593fd56154fe33b1f302a557c2c6" -dependencies = [ - "darling 0.21.3", - "itertools 0.14.0", - "proc-macro2", - "quote", - "syn 2.0.117", -] - -[[package]] -name = "mlx-lm" -version = "0.0.1" -source = "git+https://github.com/jh-block/mlx-lm#4cb8572045f55f170eb6e41005a3f8a730673c42" -dependencies = [ - "anyhow", - "clap", - "idna_adapter", - "minijinja", - "mlx-lm-utils", - "mlx-macros", - "mlx-rs", - "serde", - "serde_json", - "thiserror 2.0.18", - "tokenizers 0.22.2", -] - -[[package]] -name = "mlx-lm-utils" -version = "0.0.1" -source = "git+https://github.com/jh-block/mlx-lm#4cb8572045f55f170eb6e41005a3f8a730673c42" -dependencies = [ - "minijinja", - "minijinja-contrib", - "serde", - "serde_json", - "thiserror 2.0.18", - "tokenizers 0.22.2", -] - -[[package]] -name = "mlx-macros" -version = "0.25.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a819ee8b4434690572b6feb9c3ef0b6e90137e4190b340cf00150703b410aaf9" -dependencies = [ - "darling 0.21.3", - "proc-macro2", - "quote", - "syn 2.0.117", -] - -[[package]] -name = "mlx-rs" -version = "0.25.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3a0f592c5839b0237b3072530b1d3c503923579a2009ee0761df3edd1b1f27b" -dependencies = [ - "bytemuck", - "dyn-clone", - "half", - "itertools 0.14.0", - "libc", - "mach-sys", - "mlx-internal-macros", - "mlx-macros", - "mlx-sys", - "num-complex", - "num-traits", - "num_enum", - "parking_lot", - "paste", - "safetensors 0.6.2", - "smallvec", - "strum 0.27.2", - "thiserror 2.0.18", -] - -[[package]] -name = "mlx-sys" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e3bc3880111918b2d5018f845d48fd995f9901f16efc81d1fcfd2f4210b8219" -dependencies = [ - "bindgen 0.70.1", - "cc", - "cmake", -] - [[package]] name = "mockall" version = "0.15.0" @@ -7398,19 +7240,13 @@ dependencies = [ "syn 2.0.117", ] -[[package]] -name = "number_prefix" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" - [[package]] name = "oauth2" version = "5.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "51e219e79014df21a225b1860a479e2dcd7cbd9130f4defd4bd0e191ea31d67d" dependencies = [ - "base64 0.21.7", + "base64 0.22.1", "chrono", "getrandom 0.2.17", "http 1.4.2", @@ -8564,7 +8400,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" dependencies = [ "anyhow", - "itertools 0.13.0", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -9347,7 +9183,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ef86cd5876211988985292b91c96a8f2d298df24e75989a43a3c73f2d4d8168b" dependencies = [ "aws-lc-rs", - "log", "once_cell", "ring", "rustls-pki-types", @@ -9462,6 +9297,99 @@ version = "0.11.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3944826ff8fa8093089aba3acb4ef44b9446a99a16f3bf4e74af3f77d340ab7d" +[[package]] +name = "safemlx" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f501c4f9748dfe9ccc905e7e13ae88873d96a00735268c69258e158da07a777f" +dependencies = [ + "bytemuck", + "dyn-clone", + "half", + "itertools 0.14.0", + "libc", + "mach-sys", + "num-complex", + "num-traits", + "num_enum", + "parking_lot", + "paste", + "safemlx-internal-macros", + "safemlx-macros", + "safemlx-sys", + "safetensors 0.6.2", + "smallvec", + "strum 0.27.2", + "thiserror 2.0.18", +] + +[[package]] +name = "safemlx-internal-macros" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9708b1312176963a4ea6412859e6655e56a83fe6ab936142715aa0e9fd662fe5" +dependencies = [ + "darling 0.21.3", + "itertools 0.14.0", + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "safemlx-lm" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78f450654386009003b745adfdae7af2c7f44ae749d75621a2a0b07607945e9f" +dependencies = [ + "anyhow", + "clap", + "idna_adapter", + "minijinja", + "safemlx", + "safemlx-lm-utils", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokenizers 0.22.2", +] + +[[package]] +name = "safemlx-lm-utils" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "583d39124dae0b4a93ac7528d16d543f5e717c21365e6b034a8a74d8e1b4eb3a" +dependencies = [ + "minijinja", + "minijinja-contrib", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokenizers 0.22.2", +] + +[[package]] +name = "safemlx-macros" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d549c5fdbe69b904ef84c5cf352add203de1093b3c636e8133d94c24cfc33ce0" +dependencies = [ + "darling 0.21.3", + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "safemlx-sys" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "552285b007dab687890f5c49382e52e69a3bb1359bb2c8f2c2743743509af071" +dependencies = [ + "cc", + "cmake", +] + [[package]] name = "safetensors" version = "0.4.5" @@ -10317,17 +10245,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "socks" -version = "0.3.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f0c3dbbd9ae980613c6dd8e28a9407b50509d3803b57624d5dfe8315218cd58b" -dependencies = [ - "byteorder", - "libc", - "winapi", -] - [[package]] name = "sourcemap" version = "9.3.2" @@ -11777,8 +11694,6 @@ dependencies = [ "derive_builder", "esaxx-rs", "getrandom 0.3.4", - "hf-hub 0.4.3", - "indicatif 0.18.4", "itertools 0.14.0", "log", "macro_rules_attribute", @@ -12810,25 +12725,6 @@ version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" -[[package]] -name = "ureq" -version = "2.12.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02d1a66277ed75f640d608235660df48c8e3c19f3b4edb6a263315626cc3c01d" -dependencies = [ - "base64 0.22.1", - "flate2", - "log", - "once_cell", - "rustls", - "rustls-pki-types", - "serde", - "serde_json", - "socks", - "url", - "webpki-roots 0.26.11", -] - [[package]] name = "url" version = "2.5.8" @@ -12948,7 +12844,7 @@ version = "145.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6f84468f393984251db025e944544590a8fb429d5088b78102bd72f09232f0da" dependencies = [ - "bindgen 0.72.1", + "bindgen", "bitflags 2.11.1", "fslock", "gzip-header", diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index 783ca87651..1c02ececb9 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/crates/goose-cli/src/session/mod.rs @@ -53,7 +53,7 @@ use std::collections::{HashMap, HashSet}; use std::io::IsTerminal; use std::path::PathBuf; use std::sync::Arc; -use std::time::Instant; +use std::time::{Duration, Instant}; use tokio; use tokio_util::sync::CancellationToken; use tracing::warn; @@ -1769,25 +1769,54 @@ fn print_run_stats( usage: Option<&ProviderUsage>, ) { let elapsed = run_started.elapsed(); + let stats = usage.and_then(|usage| usage.stats.as_ref()); + let generation_elapsed = stats + .and_then(|stats| stats.elapsed_ms) + .map(Duration::from_millis); let output_tokens = usage .and_then(|usage| usage.usage.output_tokens) .and_then(|tokens| usize::try_from(tokens).ok()) - .or_else(|| usage.and_then(|usage| usage.stats.as_ref()?.output_tokens)); + .or_else(|| stats.and_then(|stats| stats.output_tokens)); let tokens_per_second = output_tokens.map(|tokens| { - if elapsed.as_secs_f64() > 0.0 { - tokens as f64 / elapsed.as_secs_f64() + let rate_elapsed = generation_elapsed.unwrap_or(elapsed); + if rate_elapsed.as_secs_f64() > 0.0 { + tokens as f64 / rate_elapsed.as_secs_f64() } else { 0.0 } }); + let model_load_ms = stats.and_then(|stats| stats.model_load_ms); + let generation_time_to_first_token_ms = stats.and_then(|stats| stats.time_to_first_token_ms); eprintln!("\nStats:"); - match first_token_at { - Some(first) => eprintln!( - " Time to first token: {:.2}s", - first.duration_since(run_started).as_secs_f64() - ), - None => eprintln!(" Time to first token: unavailable"), + if let Some(ms) = model_load_ms { + eprintln!(" Model load: {:.2}s", ms as f64 / 1000.0); + } + if model_load_ms.is_some() { + match generation_time_to_first_token_ms { + Some(ms) => eprintln!( + " Generation time to first token: {:.2}s", + ms as f64 / 1000.0 + ), + None => eprintln!(" Generation time to first token: unavailable"), + } + match first_token_at { + Some(first) => eprintln!( + " End-to-end time to first token: {:.2}s", + first.duration_since(run_started).as_secs_f64() + ), + None => eprintln!(" End-to-end time to first token: unavailable"), + } + } else if let Some(ms) = generation_time_to_first_token_ms { + eprintln!(" Time to first token: {:.2}s", ms as f64 / 1000.0); + } else { + match first_token_at { + Some(first) => eprintln!( + " Time to first token: {:.2}s", + first.duration_since(run_started).as_secs_f64() + ), + None => eprintln!(" Time to first token: unavailable"), + } } match tokens_per_second { Some(rate) => eprintln!(" Tokens/sec: {:.2}", rate), @@ -1797,10 +1826,7 @@ fn print_run_stats( eprintln!(" Output tokens: {tokens}"); } - if let Some(draft) = usage - .and_then(|usage| usage.stats.as_ref()) - .and_then(|stats| stats.draft.as_ref()) - { + if let Some(draft) = stats.and_then(|stats| stats.draft.as_ref()) { eprintln!(" Draft accept rate: {:.1}%", draft.accept_rate * 100.0); eprintln!( " Draft tokens: {} accepted: {} target verified: {} rounds: {}", diff --git a/crates/goose-cli/src/session/output.rs b/crates/goose-cli/src/session/output.rs index 7d1b56ea54..17648a0892 100644 --- a/crates/goose-cli/src/session/output.rs +++ b/crates/goose-cli/src/session/output.rs @@ -248,7 +248,8 @@ pub fn render_message(message: &Message, debug: bool) { } MessageContent::SystemNotification(notification) => { match notification.notification_type { - SystemNotificationType::ThinkingMessage => { + SystemNotificationType::ThinkingMessage + | SystemNotificationType::ProgressMessage => { show_thinking(); set_thinking_message(¬ification.msg); } @@ -330,7 +331,8 @@ pub fn render_message_streaming( } MessageContent::SystemNotification(notification) => { match notification.notification_type { - SystemNotificationType::ThinkingMessage => { + SystemNotificationType::ThinkingMessage + | SystemNotificationType::ProgressMessage => { show_thinking(); set_thinking_message(¬ification.msg); } diff --git a/crates/goose-local-inference/Cargo.toml b/crates/goose-local-inference/Cargo.toml index b2c6c42d53..c3ae342747 100644 --- a/crates/goose-local-inference/Cargo.toml +++ b/crates/goose-local-inference/Cargo.toml @@ -15,7 +15,7 @@ workspace = true default = [] cuda = ["llama-cpp-2/cuda"] vulkan = ["llama-cpp-2/vulkan"] -mlx = ["dep:mlx-rs", "dep:mlx-lm", "dep:mlx-lm-utils"] +mlx = ["dep:safemlx", "dep:safemlx-lm", "dep:safemlx-lm-utils"] [dependencies] anyhow = { workspace = true } @@ -46,9 +46,9 @@ utoipa = { workspace = true, features = ["chrono"] } uuid = { workspace = true, features = ["v4", "std"] } tempfile = { workspace = true } -mlx-rs = { version = "0.25.3", default-features = false, features = ["accelerate", "metal", "safetensors"], optional = true } -mlx-lm = { git = "https://github.com/jh-block/mlx-lm", optional = true } -mlx-lm-utils = { git = "https://github.com/jh-block/mlx-lm", optional = true } +safemlx = { default-features = false, features = ["accelerate", "metal", "safetensors"], optional = true, version = "0.1.2" } +safemlx-lm = { optional = true, version = "0.1.5" } +safemlx-lm-utils = { optional = true, version = "0.1.2" } [target.'cfg(target_os = "macos")'.dependencies] llama-cpp-2 = { workspace = true, features = ["sampler", "metal", "mtmd"] } diff --git a/crates/goose-local-inference/src/backend.rs b/crates/goose-local-inference/src/backend.rs index 2187945a77..c36420c366 100644 --- a/crates/goose-local-inference/src/backend.rs +++ b/crates/goose-local-inference/src/backend.rs @@ -22,6 +22,7 @@ pub(super) struct LocalGenerationRequest<'a> { pub temperature: Option, pub max_tokens: Option, pub context_limit: usize, + pub model_load_ms: Option, pub resolved_model: &'a ResolvedModelPaths, pub draft_model_path: Option, pub message_id: &'a str, diff --git a/crates/goose-local-inference/src/hf_models.rs b/crates/goose-local-inference/src/hf_models.rs index 8f458b3233..a63d34a122 100644 --- a/crates/goose-local-inference/src/hf_models.rs +++ b/crates/goose-local-inference/src/hf_models.rs @@ -1085,20 +1085,13 @@ mod tests { } #[test] - fn mlx_compatible_repo_accepts_supported_tokenizer_formats() { + fn mlx_compatible_repo_accepts_tokenizer_json() { let config = Some(serde_json::json!({ "model_type": "llama" })); - for tokenizer_files in [ - vec!["tokenizer.json"], - vec!["tokenizer.model"], - vec!["tokenizer.tiktoken"], - vec!["vocab.json", "merges.txt"], - ] { - assert!( - is_mlx_compatible_repo(&config, &mlx_siblings(&tokenizer_files)), - "{tokenizer_files:?}" - ); - } + assert!(is_mlx_compatible_repo( + &config, + &mlx_siblings(&["tokenizer.json"]) + )); } #[test] @@ -1107,8 +1100,11 @@ mod tests { for tokenizer_files in [ vec!["tokenizer_config.json"], + vec!["tokenizer.model"], + vec!["tokenizer.tiktoken"], vec!["vocab.json"], vec!["merges.txt"], + vec!["vocab.json", "merges.txt"], ] { assert!( !is_mlx_compatible_repo(&config, &mlx_siblings(&tokenizer_files)), @@ -1117,6 +1113,75 @@ mod tests { } } + #[test] + fn mlx_download_filenames_include_fp8_snapshot_metadata() { + let siblings = [ + "config.json", + "configuration.json", + "generation_config.json", + "model.safetensors.index.json", + "layers-0.safetensors", + "outside.safetensors", + "tokenizer.json", + "tokenizer_config.json", + "chat_template.jinja", + "preprocessor_config.json", + "video_preprocessor_config.json", + "README.md", + ] + .into_iter() + .map(sibling) + .collect::>(); + + let filenames = mlx_download_filenames(&siblings) + .into_iter() + .collect::>(); + + for filename in [ + "configuration.json", + "chat_template.jinja", + "preprocessor_config.json", + "video_preprocessor_config.json", + ] { + assert!(filenames.contains(filename), "{filename}"); + } + assert!(!filenames.contains("README.md")); + } + + #[test] + fn mlx_download_size_uses_safetensors_metadata_when_sibling_sizes_are_missing() { + let info: ModelInfo = serde_json::from_value(serde_json::json!({ + "id": "owner/repo", + "safetensors": { + "parameters": { + "BF16": 10, + "F8_E4M3": 20 + }, + "total": 30 + } + })) + .unwrap(); + let siblings = vec![ + RepoSibling { + rfilename: "config.json".to_string(), + size: None, + lfs: None, + }, + RepoSibling { + rfilename: "model.safetensors".to_string(), + size: None, + lfs: None, + }, + ]; + + assert_eq!(mlx_download_size_bytes(&info, &siblings), 40); + } + + #[test] + fn mlx_variant_id_detects_fp8_repo_name() { + assert_eq!(mlx_variant_id("Qwen/Qwen3.6-35B-A3B-FP8", &None), "fp8"); + } + fn test_model(repo_id: &str) -> HfModelInfo { HfModelInfo { repo_id: repo_id.to_string(), @@ -1157,6 +1222,13 @@ mod tests { assert_eq!(results[0].repo_id, "gguf/repo"); } + #[test] + fn best_download_count_ignores_zero_primary() { + assert_eq!(best_download_count(Some(0), Some(42)), Some(42)); + assert_eq!(best_download_count(Some(7), Some(42)), Some(7)); + assert_eq!(best_download_count(Some(0), Some(0)), None); + } + #[test] fn test_recommend_variant() { let variants = vec![ @@ -1468,7 +1540,13 @@ async fn search_mlx_models_with_query(query: &str, limit: usize) -> Result= limit { break; } - if let Some(model) = get_local_model_info_for_repo_with_client(&client, &info.id).await? { + if let Some(model) = get_local_model_info_for_repo_with_client_and_downloads( + &client, + &info.id, + info.downloads, + ) + .await? + { results.push(model); } } @@ -1484,6 +1562,14 @@ async fn get_local_model_info_for_repo(repo_id: &str) -> Result Result> { + get_local_model_info_for_repo_with_client_and_downloads(client, repo_id, None).await +} + +async fn get_local_model_info_for_repo_with_client_and_downloads( + client: &HFClient, + repo_id: &str, + downloads_hint: Option, ) -> Result> { let repo = model_repo(client, repo_id)?; let info = repo @@ -1495,7 +1581,7 @@ async fn get_local_model_info_for_repo_with_client( ]) .send() .await?; - model_info_to_local_model_info(info).await + model_info_to_local_model_info(&repo, info, downloads_hint).await } async fn get_exact_name_local_model_info(model_name: &str) -> Result> { @@ -1508,7 +1594,11 @@ async fn get_exact_name_local_model_info(model_name: &str) -> Result Result> { +async fn model_info_to_local_model_info( + repo: &HFRepository, + info: ModelInfo, + downloads_hint: Option, +) -> Result> { let repo_id = info.id.clone(); let mut variants: Vec = get_repo_gguf_variants(&repo_id) .await @@ -1516,7 +1606,13 @@ async fn model_info_to_local_model_info(info: ModelInfo) -> Result Result downloads, + None => get_repo_downloads(&repo_id).await?.unwrap_or(0), + }; Ok(Some(HfModelInfo { repo_id, author, model_name, - downloads: info.downloads.unwrap_or(0), + downloads, gguf_files: Vec::new(), variants, })) } +fn best_download_count(primary: Option, hint: Option) -> Option { + primary + .filter(|downloads| *downloads > 0) + .or_else(|| hint.filter(|downloads| *downloads > 0)) +} + +async fn get_repo_downloads(repo_id: &str) -> Result> { + let client = reqwest::Client::new(); + let token = optional_hf_token(huggingface_auth::resolve_token_async()).await; + let url = format!("{}/{}", HF_API_BASE, repo_id); + + let response = apply_hf_auth(client.get(&url), token.as_deref()) + .header("User-Agent", "goose-ai-agent") + .send() + .await?; + + if !response.status().is_success() { + bail!( + "HuggingFace API returned status {} for repo {}", + response.status(), + repo_id + ); + } + + let model: HfApiModel = response.json().await?; + Ok(model.downloads) +} + pub async fn get_repo_local_variants(repo_id: &str) -> Result> { let mut variants: Vec = get_repo_gguf_variants(repo_id) .await @@ -1570,25 +1698,39 @@ pub async fn get_repo_mlx_variants(repo_id: &str) -> Result> ]) .send() .await?; - Ok(mlx_variants_from_model_info(repo_id, &info)) + if !is_mlx_compatible_model_info(&info) { + return Ok(Vec::new()); + } + let mlx_config = load_repo_config_json(&repo) + .await + .unwrap_or_else(|_| info.config.clone()); + Ok(mlx_variants_from_model_info(repo_id, &info, &mlx_config)) } -fn mlx_variants_from_model_info(repo_id: &str, info: &ModelInfo) -> Vec { +async fn load_repo_config_json( + repo: &HFRepository, +) -> Result> { + let config_path = repo + .download_file() + .filename("config.json".to_string()) + .send() + .await?; + let config_json = tokio::fs::read_to_string(config_path).await?; + Ok(Some(serde_json::from_str(&config_json)?)) +} + +fn mlx_variants_from_model_info( + repo_id: &str, + info: &ModelInfo, + mlx_config: &Option, +) -> Vec { let siblings = info.siblings.as_deref().unwrap_or(&[]); if !is_mlx_compatible_repo(&info.config, siblings) { return Vec::new(); } - let size_bytes = mlx_download_filenames(siblings) - .into_iter() - .filter_map(|filename| { - siblings - .iter() - .find(|s| s.rfilename == filename) - .and_then(|s| s.size) - }) - .sum(); + let size_bytes = mlx_download_size_bytes(info, siblings); let variant_id = mlx_variant_id(repo_id, &info.config); vec![HfModelVariant { @@ -1601,20 +1743,24 @@ fn mlx_variants_from_model_info(repo_id: &str, info: &ModelInfo) -> Vec 1, - supported: is_mlx_runtime_supported(&info.config) + supported: is_mlx_runtime_supported(mlx_config) && cfg!(target_os = "macos") && cfg!(feature = "mlx"), - unsupported_reason: mlx_unsupported_reason(&info.config), + unsupported_reason: mlx_unsupported_reason(mlx_config), }] } +fn is_mlx_compatible_model_info(info: &ModelInfo) -> bool { + is_mlx_compatible_repo(&info.config, info.siblings.as_deref().unwrap_or_default()) +} + fn is_mlx_compatible_repo(config: &Option, siblings: &[RepoSibling]) -> bool { let has_config = siblings.iter().any(|s| s.rfilename == "config.json"); let has_tokenizer = has_mlx_tokenizer(siblings); @@ -1625,16 +1771,50 @@ fn is_mlx_compatible_repo(config: &Option, siblings: &[RepoSi has_config && has_tokenizer && has_safetensors && mlx_model_type(config).is_some() } +fn mlx_download_size_bytes(info: &ModelInfo, siblings: &[RepoSibling]) -> u64 { + let sibling_size: u64 = mlx_download_filenames(siblings) + .into_iter() + .filter_map(|filename| { + siblings + .iter() + .find(|s| s.rfilename == filename) + .and_then(|s| s.size) + }) + .sum(); + sibling_size.max(estimated_safetensors_size_bytes(info)) +} + +fn estimated_safetensors_size_bytes(info: &ModelInfo) -> u64 { + info.safetensors + .as_ref() + .map(|safetensors| { + safetensors + .parameters + .iter() + .map(|(dtype, count)| count.saturating_mul(dtype_size_bytes(dtype))) + .sum() + }) + .unwrap_or(0) +} + +fn dtype_size_bytes(dtype: &str) -> u64 { + match dtype.to_ascii_uppercase().as_str() { + "BOOL" | "I8" | "U8" | "F8_E4M3" | "F8_E4M3FN" | "F8_E5M2" | "F8_E5M2FNUZ" => 1, + "BF16" | "F16" | "I16" | "U16" => 2, + "F32" | "I32" | "U32" => 4, + "F64" | "I64" | "U64" => 8, + _ => 0, + } +} + fn has_mlx_tokenizer(siblings: &[RepoSibling]) -> bool { - let has_file = |filename: &str| siblings.iter().any(|s| s.rfilename == filename); siblings .iter() .any(|s| is_standalone_mlx_tokenizer_file(&s.rfilename)) - || (has_file("vocab.json") && has_file("merges.txt")) } fn is_standalone_mlx_tokenizer_file(filename: &str) -> bool { - filename == "tokenizer.json" || filename.ends_with(".model") || filename.ends_with(".tiktoken") + filename == "tokenizer.json" } fn mlx_model_type(config: &Option) -> Option<&str> { @@ -1644,22 +1824,8 @@ fn mlx_model_type(config: &Option) -> Option<&str> { .and_then(|value| value.as_str()) } -fn is_mlx_runtime_supported_model_type(model_type: &str) -> bool { - matches!(model_type, "gemma4" | "gemma4_text" | "llama" | "qwen3") -} - -fn is_mlx_moe_model(config: &Option) -> bool { - config - .as_ref() - .and_then(|config| config.get("text_config")) - .and_then(|text_config| text_config.get("enable_moe_block")) - .and_then(|value| value.as_bool()) - .unwrap_or(false) -} - fn is_mlx_runtime_supported(config: &Option) -> bool { - mlx_model_type(config).is_some_and(is_mlx_runtime_supported_model_type) - && !is_mlx_moe_model(config) + mlx_config_support(config).is_none() } fn mlx_unsupported_reason(config: &Option) -> Option { @@ -1670,16 +1836,23 @@ fn mlx_unsupported_reason(config: &Option) -> Option return Some("MLX support was not compiled in".to_string()); } - let model_type = mlx_model_type(config)?; - if !is_mlx_runtime_supported_model_type(model_type) { - return Some(format!( - "MLX backend does not support '{}' models yet", - model_type - )); - } - if is_mlx_moe_model(config) { - return Some("MLX backend does not support Gemma 4 MoE models yet".to_string()); - } + mlx_config_support(config) +} + +fn mlx_config_support(config: &Option) -> Option { + let config = config.as_ref()?; + mlx_config_support_for_value(config) +} + +#[cfg(feature = "mlx")] +fn mlx_config_support_for_value(config: &serde_json::Value) -> Option { + safemlx_lm::check_model_config(config) + .unsupported_reason() + .map(str::to_string) +} + +#[cfg(not(feature = "mlx"))] +fn mlx_config_support_for_value(_config: &serde_json::Value) -> Option { None } @@ -1771,6 +1944,10 @@ fn should_download_for_mlx(filename: &str) -> bool { || is_standalone_mlx_tokenizer_file(filename) || filename == "tokenizer_config.json" || filename == "generation_config.json" + || filename == "configuration.json" + || filename == "chat_template.jinja" + || filename == "preprocessor_config.json" + || filename == "video_preprocessor_config.json" || filename == "special_tokens_map.json" || filename == "model.safetensors.index.json" || filename == "vocab.json" @@ -1780,7 +1957,7 @@ fn should_download_for_mlx(filename: &str) -> bool { fn mlx_variant_id(repo_id: &str, config: &Option) -> String { let repo_lower = repo_id.to_lowercase(); - for marker in ["bf16", "f16", "fp16", "f32", "fp32", "4bit", "8bit"] { + for marker in ["bf16", "f16", "fp16", "f32", "fp32", "fp8", "4bit", "8bit"] { if repo_lower.contains(marker) { return marker.to_string(); } @@ -1954,20 +2131,12 @@ async fn resolve_mlx_model(repo_id: &str, variant_id: &str) -> Result>>>; +type ModelSlotHandle = Arc; + +struct ModelSlot { + state: Mutex, + notify: Notify, +} + +enum ModelSlotState { + Empty, + Loading, + Loaded(Box), +} + +impl ModelSlot { + fn new() -> Self { + Self { + state: Mutex::new(ModelSlotState::Empty), + notify: Notify::new(), + } + } +} #[derive(Clone, Debug, Eq, Hash, PartialEq)] struct ModelCacheKey { @@ -64,7 +87,8 @@ impl ModelCacheKey { } pub struct InferenceRuntime { - models: StdMutex>, + models: StdMutex>, + cold_load_lock: Mutex<()>, backends: HashMap<&'static str, Arc>, } @@ -72,19 +96,17 @@ pub fn builtin_chat_template_names() -> Vec { llamacpp::builtin_chat_template_names() } -/// Global weak reference used to share a single `InferenceRuntime` across -/// all providers and management APIs. Only a `Weak` is stored here — strong -/// `Arc`s live in providers and the local-inference management layer. When all -/// strong refs drop (normal shutdown), the runtime is deallocated and the -/// backend freed. The `Weak` left behind is inert during `__cxa_finalize`, so no -/// ggml statics race. -static RUNTIME: StdMutex> = StdMutex::new(Weak::new()); +static RUNTIME: StdMutex>> = StdMutex::new(None); + +fn current_runtime() -> Option> { + RUNTIME.lock().expect("runtime lock poisoned").clone() +} impl InferenceRuntime { pub fn get_or_init() -> Result> { let mut guard = RUNTIME.lock().expect("runtime lock poisoned"); - if let Some(runtime) = guard.upgrade() { - return Ok(runtime); + if let Some(runtime) = guard.as_ref() { + return Ok(runtime.clone()); } let llamacpp_backend: Arc = Arc::new(LlamaCppBackend::new()?); let mlx_backend: Arc = Arc::new(MlxBackend::new()); @@ -93,9 +115,10 @@ impl InferenceRuntime { backends.insert(MLX_BACKEND_ID, mlx_backend); let runtime = Arc::new(Self { models: StdMutex::new(HashMap::new()), + cold_load_lock: Mutex::new(()), backends, }); - *guard = Arc::downgrade(&runtime); + *guard = Some(runtime.clone()); Ok(runtime) } @@ -122,14 +145,19 @@ impl InferenceRuntime { }) } - fn get_or_create_model_slot(&self, key: ModelCacheKey) -> ModelSlot { + fn get_or_create_model_slot(&self, key: ModelCacheKey) -> ModelSlotHandle { let mut map = self.models.lock().expect("model cache lock poisoned"); map.entry(key) - .or_insert_with(|| Arc::new(Mutex::new(None))) + .or_insert_with(|| Arc::new(ModelSlot::new())) .clone() } - fn other_model_slots(&self, keep_key: &ModelCacheKey) -> Vec { + fn model_slot(&self, key: &ModelCacheKey) -> Option { + let map = self.models.lock().expect("model cache lock poisoned"); + map.get(key).cloned() + } + + fn other_model_slots(&self, keep_key: &ModelCacheKey) -> Vec { let map = self.models.lock().expect("model cache lock poisoned"); map.iter() .filter(|(key, _)| *key != keep_key) @@ -138,6 +166,76 @@ impl InferenceRuntime { } } +pub async fn is_model_loaded(model_name: &str) -> Result { + let resolved = match resolve_model_path(model_name) { + Some(resolved) => resolved, + None => return Ok(false), + }; + let runtime = InferenceRuntime::get_or_init().map_err(|error| { + ProviderError::ExecutionError(format!("Failed to initialize local inference: {error}")) + })?; + let backend = runtime.backend_for_model(&resolved)?; + let key = ModelCacheKey::new( + backend.id(), + model_name.to_string(), + resolved.settings.chat_template, + ); + let Some(slot) = runtime.model_slot(&key) else { + return Ok(false); + }; + + let state = slot.state.lock().await; + Ok(matches!(*state, ModelSlotState::Loaded(_))) +} + +pub async fn loaded_model_ids() -> Result, ProviderError> { + let Some(runtime) = current_runtime() else { + return Ok(HashSet::new()); + }; + let slots = { + let map = runtime.models.lock().expect("model cache lock poisoned"); + map.iter() + .map(|(key, slot)| (key.model_id.clone(), slot.clone())) + .collect::>() + }; + + let mut loaded = HashSet::new(); + for (model_id, slot) in slots { + if let Ok(state) = slot.state.try_lock() { + if matches!(*state, ModelSlotState::Loaded(_)) { + loaded.insert(model_id); + } + } else { + loaded.insert(model_id); + } + } + Ok(loaded) +} + +pub async fn evict_model(model_name: &str) -> Result { + let Some(runtime) = current_runtime() else { + return Ok(false); + }; + let slots = { + let map = runtime.models.lock().expect("model cache lock poisoned"); + map.iter() + .filter(|(key, _)| key.model_id == model_name) + .map(|(_, slot)| slot.clone()) + .collect::>() + }; + + let mut evicted = false; + for slot in slots { + let mut state = slot.state.lock().await; + if matches!(*state, ModelSlotState::Loaded(_)) { + *state = ModelSlotState::Empty; + evicted = true; + slot.notify.notify_waiters(); + } + } + Ok(evicted) +} + const PROVIDER_NAME: &str = "local"; const DEFAULT_MODEL: &str = "bartowski/Llama-3.2-1B-Instruct-GGUF:Q4_K_M"; @@ -565,41 +663,9 @@ impl Provider for LocalInferenceProvider { })?; let backend = self.runtime.backend_for_model(&resolved)?; let model_context_limit = resolved.context_limit; - let model_settings = resolved.settings.clone(); - let cache_key = ModelCacheKey::new( - backend.id(), - model_config.model_name.clone(), - model_settings.chat_template.clone(), - ); - let model_slot = self.runtime.get_or_create_model_slot(cache_key.clone()); - - // Ensure model is loaded — unload any other models first to free memory. - { - let mut model_lock = model_slot.lock().await; - if model_lock.is_none() { - for slot in self.runtime.other_model_slots(&cache_key) { - let mut other = slot.lock().await; - if other.is_some() { - tracing::info!("Unloading previous model to free memory"); - *other = None; - } - } - - let model_id = model_config.model_name.clone(); - let resolved_for_load = resolved.clone(); - let settings_for_load = model_settings.clone(); - let backend_for_load = backend.clone(); - let loaded = tokio::task::spawn_blocking(move || { - backend_for_load.load_model(&model_id, &resolved_for_load, &settings_for_load) - }) - .await - .map_err(|e| ProviderError::ExecutionError(e.to_string()))??; - *model_lock = Some(loaded); - } - } // Allow request_params to override thinking - let mut model_settings = model_settings; + let mut model_settings = resolved.settings.clone(); if let Some(false) = model_config .request_param::("enable_thinking") .or_else(|| { @@ -611,6 +677,15 @@ impl Provider for LocalInferenceProvider { model_settings.enable_thinking = false; } + let cache_key = ModelCacheKey::new( + backend.id(), + model_config.model_name.clone(), + model_settings.chat_template.clone(), + ); + let model_slot = self.runtime.get_or_create_model_slot(cache_key.clone()); + let runtime = self.runtime.clone(); + + let cache_key = cache_key.clone(); let model_arc = model_slot.clone(); let backend = backend.clone(); let model_name = model_config.model_name.clone(); @@ -639,17 +714,172 @@ impl Provider for LocalInferenceProvider { }, }); - let mut log = start_log(model_config, &log_payload)?; - let (tx, mut rx) = tokio::sync::mpsc::channel::< Result<(Option, Option), ProviderError>, >(32); + let mut log = start_log(model_config, &log_payload)?; - tokio::task::spawn_blocking(move || { - // Macro to log errors before sending them through the channel - macro_rules! send_err { - ($err:expr) => {{ - let err = $err; + tokio::spawn(async move { + let mut model_load_ms = None; + + // Ensure model is loaded — unload any other models first to free memory. + loop { + let state = model_slot.state.lock().await; + match &*state { + ModelSlotState::Loaded(_) => break, + ModelSlotState::Loading => { + let notified = model_slot.notify.notified(); + drop(state); + notified.await; + } + ModelSlotState::Empty => { + drop(state); + + let cold_load_guard = runtime.cold_load_lock.lock().await; + let mut state = model_slot.state.lock().await; + match &*state { + ModelSlotState::Loaded(_) => break, + ModelSlotState::Loading => { + let notified = model_slot.notify.notified(); + drop(state); + drop(cold_load_guard); + notified.await; + continue; + } + ModelSlotState::Empty => {} + } + *state = ModelSlotState::Loading; + drop(state); + + let loading_message = Message::assistant().with_system_notification( + SystemNotificationType::ProgressMessage, + format!("Loading local model {model_name}..."), + ); + if tx.send(Ok((Some(loading_message), None))).await.is_err() { + let mut state = model_slot.state.lock().await; + *state = ModelSlotState::Empty; + model_slot.notify.notify_waiters(); + return; + } + + let other_model_slots = runtime.other_model_slots(&cache_key); + for slot in other_model_slots { + let mut other = slot.state.lock().await; + if matches!(*other, ModelSlotState::Loaded(_)) { + tracing::info!("Unloading previous model to free memory"); + *other = ModelSlotState::Empty; + } + } + + let model_id = model_name.clone(); + let resolved_for_load = resolved_model.clone(); + let settings_for_load = settings.clone(); + let backend_for_load = backend.clone(); + let load_started = std::time::Instant::now(); + let loaded = match tokio::task::spawn_blocking(move || { + backend_for_load.load_model( + &model_id, + &resolved_for_load, + &settings_for_load, + ) + }) + .await + { + Ok(Ok(loaded)) => loaded, + Ok(Err(err)) => { + let mut state = model_slot.state.lock().await; + *state = ModelSlotState::Empty; + model_slot.notify.notify_waiters(); + let _ = log.error(&err); + let _ = tx.send(Err(err)).await; + return; + } + Err(err) => { + let mut state = model_slot.state.lock().await; + *state = ModelSlotState::Empty; + model_slot.notify.notify_waiters(); + let err = ProviderError::ExecutionError(err.to_string()); + let _ = log.error(&err); + let _ = tx.send(Err(err)).await; + return; + } + }; + let elapsed_ms = + u64::try_from(load_started.elapsed().as_millis()).unwrap_or(u64::MAX); + model_load_ms = Some(elapsed_ms); + tracing::info!( + backend = backend.id(), + model = %model_name, + model_load_ms = elapsed_ms, + "Loaded local inference model" + ); + let _ = log.write( + &json!({ + "path": "model_load", + "backend": backend.id(), + "model": &model_name, + "model_load_ms": elapsed_ms, + }), + None, + ); + + let mut state = model_slot.state.lock().await; + *state = ModelSlotState::Loaded(loaded); + model_slot.notify.notify_waiters(); + drop(cold_load_guard); + break; + } + } + } + + tokio::task::spawn_blocking(move || { + // Macro to log errors before sending them through the channel + macro_rules! send_err { + ($err:expr) => {{ + let err = $err; + let msg = match &err { + ProviderError::ExecutionError(s) => s.as_str(), + ProviderError::ContextLengthExceeded(s) => s.as_str(), + _ => "unknown error", + }; + let _ = log.error(msg); + let _ = tx.blocking_send(Err(err)); + return; + }}; + } + + let mut model_guard = model_arc.state.blocking_lock(); + let loaded = match &mut *model_guard { + ModelSlotState::Loaded(loaded) => loaded.as_mut(), + ModelSlotState::Empty | ModelSlotState::Loading => { + send_err!(ProviderError::ExecutionError( + "Model not loaded".to_string() + )); + } + }; + + let message_id = Uuid::new_v4().to_string(); + + let request = backend::LocalGenerationRequest { + model_name, + system: &system, + messages: &messages, + tools: &tools, + settings: &settings, + temperature, + max_tokens, + context_limit, + model_load_ms, + resolved_model: &resolved_model, + draft_model_path: resolved_model.draft_model_path.clone(), + message_id: &message_id, + tx: &tx, + log: &mut log, + }; + + let result = backend.generate(loaded, request); + + if let Err(err) = result { let msg = match &err { ProviderError::ExecutionError(s) => s.as_str(), ProviderError::ContextLengthExceeded(s) => s.as_str(), @@ -657,49 +887,8 @@ impl Provider for LocalInferenceProvider { }; let _ = log.error(msg); let _ = tx.blocking_send(Err(err)); - return; - }}; - } - - let mut model_guard = model_arc.blocking_lock(); - let loaded = match model_guard.as_mut() { - Some(l) => l, - None => { - send_err!(ProviderError::ExecutionError( - "Model not loaded".to_string() - )); } - }; - - let message_id = Uuid::new_v4().to_string(); - - let request = backend::LocalGenerationRequest { - model_name, - system: &system, - messages: &messages, - tools: &tools, - settings: &settings, - temperature, - max_tokens, - context_limit, - resolved_model: &resolved_model, - draft_model_path: resolved_model.draft_model_path.clone(), - message_id: &message_id, - tx: &tx, - log: &mut log, - }; - - let result = backend.generate(loaded.as_mut(), request); - - if let Err(err) = result { - let msg = match &err { - ProviderError::ExecutionError(s) => s.as_str(), - ProviderError::ContextLengthExceeded(s) => s.as_str(), - _ => "unknown error", - }; - let _ = log.error(msg); - let _ = tx.blocking_send(Err(err)); - } + }); }); Ok(Box::pin(try_stream! { diff --git a/crates/goose-local-inference/src/llamacpp/inference_emulated_tools.rs b/crates/goose-local-inference/src/llamacpp/inference_emulated_tools.rs index ba0b2e7d81..27f7d4d0cb 100644 --- a/crates/goose-local-inference/src/llamacpp/inference_emulated_tools.rs +++ b/crates/goose-local-inference/src/llamacpp/inference_emulated_tools.rs @@ -27,10 +27,9 @@ use serde_json::json; use std::borrow::Cow; use uuid::Uuid; -use super::super::{finalize_usage, StreamSender}; +use super::super::{finalize_usage, thinking_output::ThinkingOutputFilter, StreamSender}; use super::inference_engine::{ - generation_loop, prepare_generation, GenerationContext, StopSuffixTrimmer, - ThinkingOutputFilter, TokenAction, + generation_loop, prepare_generation, GenerationContext, StopSuffixTrimmer, TokenAction, }; const SHELL_TOOL: &str = "developer__shell"; diff --git a/crates/goose-local-inference/src/llamacpp/inference_engine.rs b/crates/goose-local-inference/src/llamacpp/inference_engine.rs index 95c66c3ac9..4b3e083b57 100644 --- a/crates/goose-local-inference/src/llamacpp/inference_engine.rs +++ b/crates/goose-local-inference/src/llamacpp/inference_engine.rs @@ -3,7 +3,6 @@ use crate::local_model_registry::ModelSettings; use crate::multimodal::ExtractedImage; use goose_provider_types::errors::ProviderError; use goose_provider_types::request_log::{LoggerHandleExt, RequestLogHandle}; -use goose_provider_types::thinking::{FilterOut, ThinkFilter}; use llama_cpp_2::context::params::LlamaContextParams; use llama_cpp_2::llama_batch::LlamaBatch; use llama_cpp_2::model::{AddBos, ChatTemplateResult, LlamaChatTemplate, LlamaModel}; @@ -48,86 +47,6 @@ pub(super) struct PreparedGeneration<'model> { pub effective_ctx: usize, } -pub(super) struct ThinkingOutputFilter { - enabled: bool, - saw_structured_reasoning: bool, - think_filter: ThinkFilter, - pending_inline_thinking: String, - accumulated_thinking: String, -} - -impl ThinkingOutputFilter { - pub(super) fn new(enable_thinking: bool, generation_prompt: &str) -> Self { - let mut think_filter = ThinkFilter::new(); - if enable_thinking && !generation_prompt.is_empty() { - let _ = think_filter.push(generation_prompt); - } - - Self { - enabled: enable_thinking, - saw_structured_reasoning: false, - think_filter, - pending_inline_thinking: String::new(), - accumulated_thinking: String::new(), - } - } - - pub(super) fn push_structured_reasoning(&mut self, reasoning: &str) -> Option { - if reasoning.is_empty() { - return None; - } - - self.saw_structured_reasoning = true; - self.pending_inline_thinking.clear(); - self.think_filter = ThinkFilter::new(); - self.accumulated_thinking.push_str(reasoning); - Some(reasoning.to_string()) - } - - pub(super) fn push_text(&mut self, text: &str) -> FilterOut { - if !self.enabled { - return FilterOut { - content: text.to_string(), - thinking: String::new(), - }; - } - - let mut filtered = self.think_filter.push(text); - if self.saw_structured_reasoning { - filtered.thinking.clear(); - } else if !filtered.thinking.is_empty() { - self.pending_inline_thinking.push_str(&filtered.thinking); - filtered.thinking.clear(); - } - filtered - } - - pub(super) fn finish(&mut self) -> FilterOut { - let mut filtered = if self.enabled && !self.saw_structured_reasoning { - std::mem::take(&mut self.think_filter).finish() - } else { - FilterOut::default() - }; - - if !self.saw_structured_reasoning { - let mut thinking = std::mem::take(&mut self.pending_inline_thinking); - thinking.push_str(&filtered.thinking); - if !thinking.is_empty() { - self.accumulated_thinking.push_str(&thinking); - } - filtered.thinking = thinking; - } else { - filtered.thinking.clear(); - } - - filtered - } - - pub(super) fn accumulated_thinking(&self) -> &str { - &self.accumulated_thinking - } -} - pub(super) struct StopSuffixTrimmer { pending: String, stops: Vec, diff --git a/crates/goose-local-inference/src/llamacpp/inference_native_tools.rs b/crates/goose-local-inference/src/llamacpp/inference_native_tools.rs index 5ac76d0598..ed44c78403 100644 --- a/crates/goose-local-inference/src/llamacpp/inference_native_tools.rs +++ b/crates/goose-local-inference/src/llamacpp/inference_native_tools.rs @@ -6,9 +6,9 @@ use std::borrow::Cow; use uuid::Uuid; use super::super::finalize_usage; +use super::super::thinking_output::ThinkingOutputFilter; use super::inference_engine::{ - generation_loop, prepare_generation, GenerationContext, StopSuffixTrimmer, - ThinkingOutputFilter, TokenAction, + generation_loop, prepare_generation, GenerationContext, StopSuffixTrimmer, TokenAction, }; pub(super) fn generate_with_native_tools( diff --git a/crates/goose-local-inference/src/management.rs b/crates/goose-local-inference/src/management.rs index 4517124d6d..31d90d53ba 100644 --- a/crates/goose-local-inference/src/management.rs +++ b/crates/goose-local-inference/src/management.rs @@ -26,6 +26,7 @@ use goose_sdk_types::custom_requests::{ LocalInferenceModelSettingsReadResponse, LocalInferenceModelSettingsUpdateResponse, LocalInferenceModelsListResponse, LocalInferenceSamplingConfig, LocalInferenceToolCallingMode, }; +use std::collections::HashSet; use std::path::PathBuf; use std::sync::{Arc, OnceLock}; @@ -44,13 +45,16 @@ pub async fn list_models() -> Result { let runtime = management_runtime()?; let recommended_id = recommend_local_model(&runtime); + let loaded_model_ids = crate::loaded_model_ids() + .await + .map_err(|error| anyhow!(error.to_string()))?; let registry = get_registry() .lock() .map_err(|_| anyhow!("Failed to acquire registry lock"))?; let mut models: Vec = registry .list_models() .iter() - .map(|entry| local_model_to_dto(entry, &recommended_id)) + .map(|entry| local_model_to_dto(entry, &recommended_id, &loaded_model_ids)) .collect(); models.sort_by(|a, b| { @@ -204,6 +208,20 @@ pub fn delete_model(model_id: &str) -> Result<()> { registry.delete_model(model_id) } +pub fn model_exists(model_id: &str) -> Result { + let registry = get_registry() + .lock() + .map_err(|_| anyhow!("Failed to acquire registry lock"))?; + Ok(registry.get_model(model_id).is_some()) +} + +pub async fn evict_model(model_id: &str) -> Result<()> { + crate::evict_model(model_id) + .await + .map(|_| ()) + .map_err(|error| anyhow!(error.to_string())) +} + pub fn get_model_settings(model_id: &str) -> Result { let registry = get_registry() .lock() @@ -412,7 +430,11 @@ pub async fn ensure_featured_models_current() -> Result<()> { Ok(()) } -fn local_model_to_dto(entry: &LocalModelEntry, recommended_id: &str) -> LocalInferenceModelDto { +fn local_model_to_dto( + entry: &LocalModelEntry, + recommended_id: &str, + loaded_model_ids: &HashSet, +) -> LocalInferenceModelDto { let vision_capable = entry.settings.vision_capable; LocalInferenceModelDto { id: entry.id.clone(), @@ -422,6 +444,7 @@ fn local_model_to_dto(entry: &LocalModelEntry, recommended_id: &str) -> LocalInf size_bytes: entry.file_size(), status: model_download_status_to_dto(entry.download_status()), recommended: recommended_id == entry.id, + is_loaded: loaded_model_ids.contains(&entry.id), settings: model_settings_to_dto(&entry.settings), vision_capable, mmproj_status: vision_capable diff --git a/crates/goose-local-inference/src/mlx.rs b/crates/goose-local-inference/src/mlx.rs index 2deb4b3f2f..6b6ddb0fd4 100644 --- a/crates/goose-local-inference/src/mlx.rs +++ b/crates/goose-local-inference/src/mlx.rs @@ -3,23 +3,24 @@ mod imp { use std::any::Any; use std::path::{Path, PathBuf}; - use mlx_lm::cache::ConcatKeyValueCache; - use mlx_lm::gemma4_mtp::generate_gemma4_mtp; - use mlx_lm::models::{gemma4_assistant::load_gemma4_assistant_model, LoadedModel, Model}; - use mlx_lm_utils::tokenizer::{Chat, Conversation, Role}; - use mlx_rs::transforms::eval; + use safemlx::transforms::eval; + use safemlx::{random, Array, Device, DeviceType, Stream}; + use safemlx_lm::gemma4_mtp::generate_gemma4_mtp; + use safemlx_lm::models::{gemma4_assistant::load_gemma4_assistant_model, LoadedModel, Model}; + use safemlx_lm_utils::tokenizer::{Chat, Conversation, Role, Tokenizer}; use serde_json::json; use crate::backend::{BackendLoadedModel, LocalGenerationRequest, LocalInferenceBackend}; use crate::local_model_registry::{ModelSettings, ToolCallingMode}; use crate::native_tool_parsing::message_from_native_tool_text; use crate::provider_utils::filter_extensions_from_system_prompt; + use crate::thinking_output::ThinkingOutputFilter; use crate::tool_emulation::{ build_emulator_tool_description, load_tiny_model_prompt, message_for_emulator_action, StreamingEmulatorParser, CODE_EXECUTION_TOOL, }; use crate::{extract_text_content, ResolvedModelPaths}; - use goose_provider_types::conversation::message::Message; + use goose_provider_types::conversation::message::{Message, MessageContent}; use goose_provider_types::conversation::token_usage::{ DraftStats, ProviderStats, ProviderUsage, Usage, }; @@ -57,14 +58,25 @@ mod imp { } let model_dir = model_dir_from_path(&resolved.model_path)?; - let model = LoadedModel::load(&model_dir).map_err(mlx_error)?; + let stream = Stream::new_with_device(&Device::new(DeviceType::Gpu, 0)); + let weights_stream = Stream::new_with_device(&Device::new(DeviceType::Cpu, 0)); + let model = + LoadedModel::load(&model_dir, &stream, &weights_stream).map_err(mlx_error)?; + let tokenizer = + Tokenizer::from_file(model_dir.join("tokenizer.json")).map_err(mlx_error)?; tracing::info!( backend = self.id(), model_id, model_type = model.model_type(), "MLX model loaded successfully" ); - Ok(Box::new(MlxLoadedModel { model, model_dir })) + let stop_token_ids = mlx_stop_token_ids(&model, &model_dir); + Ok(Box::new(MlxLoadedModel { + model, + tokenizer, + model_dir, + stop_token_ids, + })) } fn generate( @@ -79,6 +91,7 @@ mod imp { ProviderError::ExecutionError("Loaded model backend mismatch".to_string()) })?; + let stream = Stream::new_with_device(&Device::new(DeviceType::Gpu, 0)); let tool_mode = if request.tools.is_empty() { ToolMode::None } else { @@ -110,77 +123,114 @@ mod imp { let prompt_array = loaded .model - .encode_to_array(&prompt, false) + .encode_to_array(&prompt, false, &stream) .map_err(mlx_error)?; - let max_tokens = request - .settings - .max_output_tokens - .or_else(|| { - request - .max_tokens - .and_then(|tokens| usize::try_from(tokens).ok()) - }) - .unwrap_or(512); - let temp = request - .temperature - .unwrap_or_else(|| temperature(request.settings)); - let eos_token_ids = loaded.model.eos_token_ids().to_vec(); + let max_tokens = mlx_max_tokens( + request.settings, + request.max_tokens, + request.context_limit, + prompt_tokens.len(), + ); + let (settings_temp, seed) = sampling(request.settings); + let temp = request.temperature.unwrap_or(settings_temp); + let prng_key = prng_key(temp, seed)?; + let eos_token_ids = loaded.stop_token_ids.clone(); let generation_started = std::time::Instant::now(); - let (generated_ids, draft_stats) = - if let Some(draft_model_path) = &request.draft_model_path { - if matches!(loaded.model.model_mut(), Model::Gemma4(_)) { - let mut assistant = - load_gemma4_assistant_model(draft_model_path).map_err(|error| { + let MlxGeneration { + generated_ids, + generated_text, + draft_stats, + time_to_first_token_ms, + streamed_response, + } = if let Some(draft_model_path) = &request.draft_model_path { + if matches!(loaded.model.model_mut(), Model::Gemma4(_)) { + let weights_stream = Stream::new_with_device(&Device::new(DeviceType::Cpu, 0)); + let mut assistant = + load_gemma4_assistant_model(draft_model_path, &stream, &weights_stream) + .map_err(|error| { mlx_error(format!("failed to load MLX draft model: {error}")) })?; - let target = match loaded.model.model_mut() { - Model::Gemma4(target) => target, - _ => unreachable!(), - }; - let (ids, stats) = generate_gemma4_mtp( - target, - &mut assistant, - &prompt_array, - &eos_token_ids, - max_tokens, - temp, - ) - .map_err(mlx_error)?; - ( - ids, - Some(DraftStats { - model: Some(draft_model_path.display().to_string()), - draft_tokens: stats.draft_tokens, - accepted_tokens: stats.accepted_tokens, - target_tokens: stats.target_tokens, - rounds: stats.rounds, - accept_rate: stats.accept_rate(), - }), - ) - } else { - generate_single_model( - &mut loaded.model, - &prompt_array, - &eos_token_ids, - max_tokens, - temp, - )? - } - } else { - generate_single_model( - &mut loaded.model, + let target = match loaded.model.model_mut() { + Model::Gemma4(target) => target, + _ => unreachable!(), + }; + let (ids, stats) = generate_gemma4_mtp( + target, + &mut assistant, &prompt_array, &eos_token_ids, max_tokens, temp, + prng_key, + &stream, + ) + .map_err(mlx_error)?; + let generated_text = loaded.tokenizer.decode(&ids, true).map_err(mlx_error)?; + MlxGeneration { + generated_ids: ids, + generated_text, + draft_stats: Some(DraftStats { + model: Some(draft_model_path.display().to_string()), + draft_tokens: stats.draft_tokens, + accepted_tokens: stats.accepted_tokens, + target_tokens: stats.target_tokens, + rounds: stats.rounds, + accept_rate: stats.accept_rate(), + }), + time_to_first_token_ms: None, + streamed_response: false, + } + } else { + generate_single_model( + &mut loaded.model, + &loaded.tokenizer, + &prompt_array, + &eos_token_ids, + max_tokens, + temp, + prng_key, + &stream, + generation_started, + MlxStreamEmitter::new( + request.message_id, + tool_mode, + request.settings.enable_thinking, + &prompt, + request.tx, + ), )? - }; + } + } else { + generate_single_model( + &mut loaded.model, + &loaded.tokenizer, + &prompt_array, + &eos_token_ids, + max_tokens, + temp, + prng_key, + &stream, + generation_started, + MlxStreamEmitter::new( + request.message_id, + tool_mode, + request.settings.enable_thinking, + &prompt, + request.tx, + ), + )? + }; - let generated_text = loaded - .model - .decode(&generated_ids, true) - .map_err(mlx_error)?; - emit_generated_response(&generated_text, request.message_id, tool_mode, request.tx)?; + if !streamed_response { + emit_generated_response( + &generated_text, + &prompt, + request.settings.enable_thinking, + request.message_id, + tool_mode, + request.tx, + )?; + } let output_tokens = generated_ids.len() as i32; let input_tokens = prompt_tokens.len() as i32; @@ -194,12 +244,16 @@ mod imp { "model_dir": loaded.model_dir, "prompt_tokens": input_tokens, "output_tokens": output_tokens, + "model_load_ms": request.model_load_ms, + "time_to_first_token_ms": time_to_first_token_ms, + "elapsed_ms": generation_started.elapsed().as_millis() as u64, "generated_text": generated_text, "draft": draft_stats, }); let _ = request.log.write(&log_json, Some(&usage)); let stats = ProviderStats { - time_to_first_token_ms: None, + time_to_first_token_ms, + model_load_ms: request.model_load_ms, elapsed_ms: Some(generation_started.elapsed().as_millis() as u64), output_tokens: Some(generated_ids.len()), draft: draft_stats, @@ -221,9 +275,19 @@ mod imp { Emulated { code_mode_enabled: bool }, } + struct MlxGeneration { + generated_ids: Vec, + generated_text: String, + draft_stats: Option, + time_to_first_token_ms: Option, + streamed_response: bool, + } + struct MlxLoadedModel { model: LoadedModel, + tokenizer: Tokenizer, model_dir: PathBuf, + stop_token_ids: Vec, } impl BackendLoadedModel for MlxLoadedModel { @@ -242,6 +306,44 @@ mod imp { } } + fn mlx_stop_token_ids(model: &LoadedModel, model_dir: &Path) -> Vec { + let mut ids = model.eos_token_ids().to_vec(); + for id in generation_config_eos_token_ids(model_dir) { + if !ids.contains(&id) { + ids.push(id); + } + } + ids + } + + fn generation_config_eos_token_ids(model_dir: &Path) -> Vec { + let Ok(config_json) = std::fs::read_to_string(model_dir.join("generation_config.json")) + else { + return Vec::new(); + }; + let Ok(config) = serde_json::from_str::(&config_json) else { + return Vec::new(); + }; + match config.get("eos_token_id") { + Some(value) => token_id_or_ids(value), + None => Vec::new(), + } + } + + fn token_id_or_ids(value: &serde_json::Value) -> Vec { + if let Some(id) = value.as_u64().and_then(|id| u32::try_from(id).ok()) { + return vec![id]; + } + value + .as_array() + .map(|ids| { + ids.iter() + .filter_map(|id| id.as_u64().and_then(|id| u32::try_from(id).ok())) + .collect() + }) + .unwrap_or_default() + } + fn build_prompt( model: &mut LoadedModel, model_name: &str, @@ -316,28 +418,107 @@ mod imp { fn generate_single_model( model: &mut LoadedModel, - prompt_array: &mlx_rs::Array, + tokenizer: &Tokenizer, + prompt_array: &Array, eos_token_ids: &[u32], max_tokens: usize, temp: f32, - ) -> Result<(Vec, Option), ProviderError> { - let mut cache: Vec> = Vec::new(); + prng_key: Option, + stream: &Stream, + generation_started: std::time::Instant, + mut emitter: MlxStreamEmitter<'_>, + ) -> Result { + let mut cache = model.new_cache(); let mut generated_ids = Vec::new(); + let mut streamed_text = String::new(); + let mut time_to_first_token_ms = None; + let stream_generation = emitter.can_stream(); + let mut decode_stream = tokenizer.decode_stream(true); { let generator = model - .generate(&mut cache, temp, prompt_array) + .generate_with_cache(&mut cache, temp, prompt_array, prng_key, stream) .take(max_tokens); for token in generator { let token = token.map_err(mlx_error)?; eval([&token]).map_err(mlx_error)?; - let token_id = token.item::(); + let token_id = token.item::(stream); + time_to_first_token_ms.get_or_insert_with(|| { + u64::try_from(generation_started.elapsed().as_millis()).unwrap_or(u64::MAX) + }); if eos_token_ids.contains(&token_id) { break; } generated_ids.push(token_id); + if stream_generation { + if let Some(piece) = decode_stream.step(token_id).map_err(mlx_error)? { + if !piece.is_empty() { + let should_continue = emitter.push_text(&piece)?; + streamed_text.push_str(&piece); + if !should_continue { + break; + } + } + } + } } } - Ok((generated_ids, None)) + let generated_text = tokenizer.decode(&generated_ids, true).map_err(mlx_error)?; + let streamed_response = if stream_generation { + match final_stream_suffix(&generated_text, &streamed_text)? { + Some(suffix) => { + if !suffix.is_empty() { + emitter.push_text(suffix)?; + } + true + } + None => false, + } + } else { + false + }; + if streamed_response { + emitter.finish()?; + } + Ok(MlxGeneration { + generated_ids, + generated_text, + draft_stats: None, + time_to_first_token_ms, + streamed_response, + }) + } + + fn final_stream_suffix<'a>( + generated_text: &'a str, + streamed_text: &str, + ) -> Result, ProviderError> { + if streamed_text.is_empty() { + return Ok(None); + } + + generated_text + .strip_prefix(streamed_text) + .map(Some) + .ok_or_else(|| mlx_error("streamed MLX decode did not match final tokenizer decode")) + } + + fn mlx_max_tokens( + settings: &ModelSettings, + request_max_tokens: Option, + context_limit: usize, + prompt_tokens: usize, + ) -> usize { + let configured_max = settings + .max_output_tokens + .or_else(|| request_max_tokens.and_then(|tokens| usize::try_from(tokens).ok())); + if context_limit == 0 { + return configured_max.unwrap_or(4096); + } + + let context_headroom = context_limit.saturating_sub(prompt_tokens); + configured_max + .map(|max| max.min(context_headroom)) + .unwrap_or(context_headroom) } fn is_gemma4(model: &LoadedModel) -> bool { @@ -463,6 +644,8 @@ mod imp { fn emit_generated_response( generated_text: &str, + generation_prompt: &str, + enable_thinking: bool, message_id: &str, tool_mode: ToolMode, tx: &tokio::sync::mpsc::Sender< @@ -473,30 +656,27 @@ mod imp { return Ok(()); } + let (content, thinking) = + split_generated_thinking(generated_text, generation_prompt, enable_thinking); + match tool_mode { ToolMode::None => { - let mut msg = Message::assistant().with_text(generated_text); - msg.id = Some(message_id.to_string()); - tx.blocking_send(Ok((Some(msg), None))).map_err(|_| { - ProviderError::ExecutionError("Failed to stream MLX response".to_string()) - })?; + emit_assistant_message(message_id, &thinking, &content, tx)?; } ToolMode::Native => { - if let Some(message) = message_from_native_tool_text(generated_text, message_id)? { + if let Some(mut message) = message_from_native_tool_text(&content, message_id)? { + prepend_thinking(&mut message, &thinking); tx.blocking_send(Ok((Some(message), None))).map_err(|_| { ProviderError::ExecutionError("Failed to stream MLX response".to_string()) })?; } else { - let mut msg = Message::assistant().with_text(generated_text); - msg.id = Some(message_id.to_string()); - tx.blocking_send(Ok((Some(msg), None))).map_err(|_| { - ProviderError::ExecutionError("Failed to stream MLX response".to_string()) - })?; + emit_assistant_message(message_id, &thinking, &content, tx)?; } } ToolMode::Emulated { code_mode_enabled } => { + emit_assistant_message(message_id, &thinking, "", tx)?; let mut parser = StreamingEmulatorParser::new(code_mode_enabled); - let mut actions = parser.process_chunk(generated_text); + let mut actions = parser.process_chunk(&content); actions.extend(parser.flush()); for action in actions { @@ -510,14 +690,191 @@ mod imp { Ok(()) } - fn temperature(settings: &ModelSettings) -> f32 { - match &settings.sampling { - crate::local_model_registry::SamplingConfig::Greedy => 0.0, - crate::local_model_registry::SamplingConfig::Temperature { temperature, .. } => { - *temperature + struct MlxStreamEmitter<'a> { + message_id: &'a str, + tool_mode: ToolMode, + tx: &'a tokio::sync::mpsc::Sender< + Result<(Option, Option), ProviderError>, + >, + output_filter: ThinkingOutputFilter, + emulator_parser: Option, + stop_after_tool_call: bool, + } + + impl<'a> MlxStreamEmitter<'a> { + fn new( + message_id: &'a str, + tool_mode: ToolMode, + enable_thinking: bool, + generation_prompt: &str, + tx: &'a tokio::sync::mpsc::Sender< + Result<(Option, Option), ProviderError>, + >, + ) -> Self { + let emulator_parser = match tool_mode { + ToolMode::Emulated { code_mode_enabled } => { + Some(StreamingEmulatorParser::new(code_mode_enabled)) + } + ToolMode::None | ToolMode::Native => None, + }; + Self { + message_id, + tool_mode, + tx, + output_filter: ThinkingOutputFilter::new(enable_thinking, generation_prompt), + emulator_parser, + stop_after_tool_call: false, } - crate::local_model_registry::SamplingConfig::MirostatV2 { .. } => 0.0, } + + fn can_stream(&self) -> bool { + !matches!(self.tool_mode, ToolMode::Native) + } + + fn push_text(&mut self, text: &str) -> Result { + let filtered = self.output_filter.push_text(text); + if !filtered.content.is_empty() { + self.emit_content(&filtered.content)?; + } + Ok(!self.stop_after_tool_call) + } + + fn finish(&mut self) -> Result<(), ProviderError> { + self.flush_filtered_output()?; + let actions = self + .emulator_parser + .as_mut() + .map(StreamingEmulatorParser::flush) + .unwrap_or_default(); + for action in actions { + let (message, is_tool) = message_for_emulator_action(&action, self.message_id); + if is_tool { + self.flush_filtered_output()?; + } + self.send(message)?; + self.stop_after_tool_call |= is_tool; + if is_tool { + break; + } + } + Ok(()) + } + + fn flush_filtered_output(&mut self) -> Result<(), ProviderError> { + let filtered = self.output_filter.finish(); + if !filtered.thinking.is_empty() { + let mut message = Message::assistant().with_thinking(filtered.thinking, ""); + message.id = Some(self.message_id.to_string()); + self.send(message)?; + } + if !filtered.content.is_empty() { + self.emit_content(&filtered.content)?; + } + Ok(()) + } + + fn emit_content(&mut self, content: &str) -> Result<(), ProviderError> { + match self.tool_mode { + ToolMode::None => { + let mut message = Message::assistant().with_text(content); + message.id = Some(self.message_id.to_string()); + self.send(message) + } + ToolMode::Emulated { .. } => { + let actions = self + .emulator_parser + .as_mut() + .map(|parser| parser.process_chunk(content)) + .unwrap_or_default(); + for action in actions { + let (message, is_tool) = + message_for_emulator_action(&action, self.message_id); + if is_tool { + self.flush_filtered_output()?; + } + self.send(message)?; + self.stop_after_tool_call |= is_tool; + if is_tool { + break; + } + } + Ok(()) + } + ToolMode::Native => Ok(()), + } + } + + fn send(&self, message: Message) -> Result<(), ProviderError> { + self.tx + .blocking_send(Ok((Some(message), None))) + .map_err(|_| { + ProviderError::ExecutionError("Failed to stream MLX response".to_string()) + }) + } + } + + fn split_generated_thinking( + generated_text: &str, + generation_prompt: &str, + enable_thinking: bool, + ) -> (String, String) { + let mut filter = ThinkingOutputFilter::new(enable_thinking, generation_prompt); + let mut filtered = filter.push_text(generated_text); + let final_filtered = filter.finish(); + filtered.content.push_str(&final_filtered.content); + filtered.thinking.push_str(&final_filtered.thinking); + (filtered.content, filtered.thinking) + } + + fn emit_assistant_message( + message_id: &str, + thinking: &str, + content: &str, + tx: &tokio::sync::mpsc::Sender< + Result<(Option, Option), ProviderError>, + >, + ) -> Result<(), ProviderError> { + if thinking.is_empty() && content.is_empty() { + return Ok(()); + } + + let mut message = Message::assistant(); + if !thinking.is_empty() { + message = message.with_thinking(thinking, ""); + } + if !content.is_empty() { + message = message.with_text(content); + } + message.id = Some(message_id.to_string()); + tx.blocking_send(Ok((Some(message), None))) + .map_err(|_| ProviderError::ExecutionError("Failed to stream MLX response".to_string())) + } + + fn prepend_thinking(message: &mut Message, thinking: &str) { + if !thinking.is_empty() { + message + .content + .insert(0, MessageContent::thinking(thinking, "")); + } + } + + fn sampling(settings: &ModelSettings) -> (f32, Option) { + match &settings.sampling { + crate::local_model_registry::SamplingConfig::Greedy => (0.0, None), + crate::local_model_registry::SamplingConfig::Temperature { + temperature, seed, .. + } => (*temperature, *seed), + crate::local_model_registry::SamplingConfig::MirostatV2 { seed, .. } => (0.0, *seed), + } + } + + fn prng_key(temp: f32, seed: Option) -> Result, ProviderError> { + if temp == 0.0 { + return Ok(None); + } + random::key(seed.unwrap_or(0) as u64) + .map(Some) + .map_err(mlx_error) } fn render_prompt(system: &str, messages: &[Message]) -> String { @@ -547,6 +904,90 @@ mod imp { fn mlx_error(error: impl std::fmt::Display) -> ProviderError { ProviderError::ExecutionError(format!("MLX backend error: {}", error)) } + + #[cfg(test)] + mod tests { + use super::{ + final_stream_suffix, mlx_max_tokens, split_generated_thinking, token_id_or_ids, + }; + use crate::local_model_registry::ModelSettings; + use serde_json::json; + + #[test] + fn extracts_thinking_started_by_generation_prompt() { + let (content, thinking) = split_generated_thinking( + "hidden reasoningvisible answer", + "<|im_start|>assistant\n\n", + true, + ); + + assert_eq!(thinking.trim(), "hidden reasoning"); + assert_eq!(content, "visible answer"); + } + + #[test] + fn leaves_think_tags_as_content_when_thinking_disabled() { + let generated = "hidden reasoningvisible answer"; + let (content, thinking) = + split_generated_thinking(generated, "<|im_start|>assistant\n\n", false); + + assert!(thinking.is_empty()); + assert_eq!(content, generated); + } + + #[test] + fn parses_single_and_multiple_eos_token_ids() { + assert_eq!(token_id_or_ids(&json!(248044)), vec![248044]); + assert_eq!( + token_id_or_ids(&json!([248046, 248044])), + vec![248046, 248044] + ); + } + + #[test] + fn final_stream_suffix_flushes_append_only_suffix() { + assert_eq!( + final_stream_suffix("hello world", "hello").unwrap(), + Some(" world") + ); + } + + #[test] + fn final_stream_suffix_does_not_replay_fully_streamed_text() { + assert_eq!( + final_stream_suffix("run tool", "run tool").unwrap(), + Some("") + ); + } + + #[test] + fn final_stream_suffix_allows_unstreamed_fallback() { + assert_eq!(final_stream_suffix("hello world", "").unwrap(), None); + } + + #[test] + fn final_stream_suffix_rejects_rewritten_streamed_prefix() { + assert!(final_stream_suffix("corrected response", "stale prefix").is_err()); + } + + #[test] + fn max_tokens_defaults_to_context_headroom() { + let settings = ModelSettings::default(); + + assert_eq!(mlx_max_tokens(&settings, None, 128_000, 1_752), 126_248); + } + + #[test] + fn max_tokens_respects_configured_caps() { + let mut settings = ModelSettings::default(); + settings.max_output_tokens = Some(2048); + + assert_eq!(mlx_max_tokens(&settings, None, 128_000, 1_752), 2048); + + let settings = ModelSettings::default(); + assert_eq!(mlx_max_tokens(&settings, Some(1024), 128_000, 1_752), 1024); + } + } } #[cfg(not(feature = "mlx"))] diff --git a/crates/goose-local-inference/src/thinking_output.rs b/crates/goose-local-inference/src/thinking_output.rs new file mode 100644 index 0000000000..b47d23bc9f --- /dev/null +++ b/crates/goose-local-inference/src/thinking_output.rs @@ -0,0 +1,81 @@ +use goose_provider_types::thinking::{FilterOut, ThinkFilter}; + +pub(crate) struct ThinkingOutputFilter { + enabled: bool, + saw_structured_reasoning: bool, + think_filter: ThinkFilter, + pending_inline_thinking: String, + accumulated_thinking: String, +} + +impl ThinkingOutputFilter { + pub(crate) fn new(enable_thinking: bool, generation_prompt: &str) -> Self { + let mut think_filter = ThinkFilter::new(); + if enable_thinking && !generation_prompt.is_empty() { + let _ = think_filter.push(generation_prompt); + } + + Self { + enabled: enable_thinking, + saw_structured_reasoning: false, + think_filter, + pending_inline_thinking: String::new(), + accumulated_thinking: String::new(), + } + } + + pub(crate) fn push_structured_reasoning(&mut self, reasoning: &str) -> Option { + if reasoning.is_empty() { + return None; + } + + self.saw_structured_reasoning = true; + self.pending_inline_thinking.clear(); + self.think_filter = ThinkFilter::new(); + self.accumulated_thinking.push_str(reasoning); + Some(reasoning.to_string()) + } + + pub(crate) fn push_text(&mut self, text: &str) -> FilterOut { + if !self.enabled { + return FilterOut { + content: text.to_string(), + thinking: String::new(), + }; + } + + let mut filtered = self.think_filter.push(text); + if self.saw_structured_reasoning { + filtered.thinking.clear(); + } else if !filtered.thinking.is_empty() { + self.pending_inline_thinking.push_str(&filtered.thinking); + filtered.thinking.clear(); + } + filtered + } + + pub(crate) fn finish(&mut self) -> FilterOut { + let mut filtered = if self.enabled && !self.saw_structured_reasoning { + std::mem::take(&mut self.think_filter).finish() + } else { + FilterOut::default() + }; + + if !self.saw_structured_reasoning { + let mut thinking = std::mem::take(&mut self.pending_inline_thinking); + thinking.push_str(&filtered.thinking); + if !thinking.is_empty() { + self.accumulated_thinking.push_str(&thinking); + } + filtered.thinking = thinking; + } else { + filtered.thinking.clear(); + } + + filtered + } + + pub(crate) fn accumulated_thinking(&self) -> &str { + &self.accumulated_thinking + } +} diff --git a/crates/goose-provider-types/src/conversation/message.rs b/crates/goose-provider-types/src/conversation/message.rs index 092ce0dd28..1c095089c7 100644 --- a/crates/goose-provider-types/src/conversation/message.rs +++ b/crates/goose-provider-types/src/conversation/message.rs @@ -253,6 +253,7 @@ pub struct FrontendToolRequest { #[serde(rename_all = "camelCase")] pub enum SystemNotificationType { ThinkingMessage, + ProgressMessage, InlineMessage, CreditsExhausted, } diff --git a/crates/goose-provider-types/src/conversation/token_usage.rs b/crates/goose-provider-types/src/conversation/token_usage.rs index 925228a3d0..4fa2db5c52 100644 --- a/crates/goose-provider-types/src/conversation/token_usage.rs +++ b/crates/goose-provider-types/src/conversation/token_usage.rs @@ -25,6 +25,7 @@ pub enum CostSource { #[derive(Debug, Clone, Serialize, Deserialize, Default)] pub struct ProviderStats { pub time_to_first_token_ms: Option, + pub model_load_ms: Option, pub elapsed_ms: Option, pub output_tokens: Option, pub draft: Option, diff --git a/crates/goose-sdk-types/src/custom_requests.rs b/crates/goose-sdk-types/src/custom_requests.rs index b1c8626912..f48cc50b71 100644 --- a/crates/goose-sdk-types/src/custom_requests.rs +++ b/crates/goose-sdk-types/src/custom_requests.rs @@ -1910,6 +1910,7 @@ pub struct LocalInferenceModelDto { pub size_bytes: u64, pub status: LocalInferenceModelDownloadStatusDto, pub recommended: bool, + pub is_loaded: bool, pub settings: LocalInferenceModelSettingsDto, pub vision_capable: bool, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -2031,6 +2032,16 @@ pub struct LocalInferenceModelDeleteRequest { pub model_id: String, } +#[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcRequest)] +#[request( + method = "_goose/unstable/local-inference/models/evict", + response = EmptyResponse +)] +#[serde(rename_all = "camelCase")] +pub struct LocalInferenceModelEvictRequest { + pub model_id: String, +} + #[derive(Debug, Default, Clone, Serialize, Deserialize, JsonSchema, JsonRpcRequest)] #[request( method = "_goose/unstable/local-inference/models/settings/read", diff --git a/crates/goose/acp-meta.json b/crates/goose/acp-meta.json index 11260c7dcd..d0cce9de4e 100644 --- a/crates/goose/acp-meta.json +++ b/crates/goose/acp-meta.json @@ -530,6 +530,11 @@ "requestType": "LocalInferenceModelDeleteRequest_unstable", "responseType": "EmptyResponse" }, + { + "method": "_goose/unstable/local-inference/models/evict", + "requestType": "LocalInferenceModelEvictRequest_unstable", + "responseType": "EmptyResponse" + }, { "method": "_goose/unstable/local-inference/models/settings/read", "requestType": "LocalInferenceModelSettingsReadRequest_unstable", diff --git a/crates/goose/acp-schema.json b/crates/goose/acp-schema.json index 87e8509a0b..21dfb7afce 100644 --- a/crates/goose/acp-schema.json +++ b/crates/goose/acp-schema.json @@ -5734,6 +5734,9 @@ "recommended": { "type": "boolean" }, + "isLoaded": { + "type": "boolean" + }, "settings": { "$ref": "#/$defs/LocalInferenceModelSettingsDto" }, @@ -5759,6 +5762,7 @@ "sizeBytes", "status", "recommended", + "isLoaded", "settings", "visionCapable" ] @@ -6217,6 +6221,19 @@ "x-side": "agent", "x-method": "_goose/unstable/local-inference/models/delete" }, + "LocalInferenceModelEvictRequest_unstable": { + "type": "object", + "properties": { + "modelId": { + "type": "string" + } + }, + "required": [ + "modelId" + ], + "x-side": "agent", + "x-method": "_goose/unstable/local-inference/models/evict" + }, "LocalInferenceModelSettingsReadRequest_unstable": { "type": "object", "properties": { @@ -7789,6 +7806,15 @@ "description": "Params for _goose/unstable/local-inference/models/delete", "title": "LocalInferenceModelDeleteRequest_unstable" }, + { + "allOf": [ + { + "$ref": "#/$defs/LocalInferenceModelEvictRequest_unstable" + } + ], + "description": "Params for _goose/unstable/local-inference/models/evict", + "title": "LocalInferenceModelEvictRequest_unstable" + }, { "allOf": [ { diff --git a/crates/goose/src/acp/server.rs b/crates/goose/src/acp/server.rs index e6e5ac6266..ed8d55153d 100644 --- a/crates/goose/src/acp/server.rs +++ b/crates/goose/src/acp/server.rs @@ -2072,6 +2072,23 @@ fn send_status_message_update( Ok(()) } +fn send_progress_message_update( + cx: &ConnectionTo, + supports_goose_custom_notifications: bool, + session_id: &str, + message: String, +) -> Result<(), agent_client_protocol::Error> { + if supports_goose_custom_notifications { + cx.send_notification(GooseSessionNotification { + session_id: session_id.to_string(), + update: GooseSessionUpdate::StatusMessage(StatusMessageUpdate { + status: StatusMessage::Progress { message }, + }), + })?; + } + Ok(()) +} + fn status_message_from_system_notification( notification: &SystemNotificationContent, ) -> Option { @@ -2079,9 +2096,11 @@ fn status_message_from_system_notification( SystemNotificationType::InlineMessage => Some(StatusMessage::Notice { message: notification.msg.clone(), }), - SystemNotificationType::ThinkingMessage => Some(StatusMessage::Progress { - message: notification.msg.clone(), - }), + SystemNotificationType::ThinkingMessage | SystemNotificationType::ProgressMessage => { + Some(StatusMessage::Progress { + message: notification.msg.clone(), + }) + } SystemNotificationType::CreditsExhausted => None, } } @@ -2482,6 +2501,44 @@ impl GooseAcpAgent { )) } + async fn send_local_inference_progress_update( + &self, + cx: &ConnectionTo, + acp_session_id: &SessionId, + session_id: &str, + agent: &Arc, + ) -> Result<(), agent_client_protocol::Error> { + let Ok(provider) = agent.provider().await else { + return Ok(()); + }; + if provider.get_name() != "local" { + return Ok(()); + } + + let model_config = agent.model_config_for_session(session_id).await.ok(); + let model_name = model_config + .as_ref() + .map(|config| config.model_name.clone()) + .unwrap_or_else(|| "local model".to_string()); + + #[cfg(feature = "local-inference")] + if let Some(model_config) = model_config.as_ref() { + if crate::providers::local_inference::is_model_loaded(&model_config.model_name) + .await + .unwrap_or(false) + { + return Ok(()); + } + } + + send_progress_message_update( + cx, + self.supports_goose_custom_notifications(), + acp_session_id.0.as_ref(), + format!("Loading local model {model_name}..."), + ) + } + async fn on_load_session( &self, cx: &ConnectionTo, @@ -2524,6 +2581,15 @@ impl GooseAcpAgent { return Err(error); } + if let Err(error) = self + .send_local_inference_progress_update(cx, &args.session_id, &session_id, &agent) + .await + { + self.clear_active_run(&session_id, &run_id).await; + let _ = Self::send_active_run_update(cx, &args.session_id, None); + return Err(error); + } + let user_message = Self::convert_acp_prompt_to_message(&args.prompt); let message_text = user_message.as_concat_text(); diff --git a/crates/goose/src/acp/server/config.rs b/crates/goose/src/acp/server/config.rs index e70ddf98df..5ce4324ca6 100644 --- a/crates/goose/src/acp/server/config.rs +++ b/crates/goose/src/acp/server/config.rs @@ -216,7 +216,8 @@ impl GooseAcpAgent { if let Some(model_id) = model_id.as_deref() { let model_exists = entry.default_model == model_id - || entry.models.iter().any(|model| model.id == model_id); + || entry.models.iter().any(|model| model.id == model_id) + || (provider_id == "local" && local_inference_model_exists(model_id)?); if !model_exists { return Err(agent_client_protocol::Error::invalid_params().data(format!( "Model '{model_id}' is not available for provider '{provider_id}'" @@ -254,6 +255,20 @@ impl GooseAcpAgent { } } +fn local_inference_model_exists(model_id: &str) -> Result { + #[cfg(feature = "local-inference")] + { + crate::providers::local_inference::management::model_exists(model_id) + .internal_err_ctx("Failed to read local inference models") + } + + #[cfg(not(feature = "local-inference"))] + { + let _ = model_id; + Ok(false) + } +} + struct PreferenceDef { key: PreferenceKey, config_key: &'static str, diff --git a/crates/goose/src/acp/server/custom_dispatch.rs b/crates/goose/src/acp/server/custom_dispatch.rs index 0e78d2b5f4..0b499e7303 100644 --- a/crates/goose/src/acp/server/custom_dispatch.rs +++ b/crates/goose/src/acp/server/custom_dispatch.rs @@ -875,6 +875,14 @@ impl GooseAcpAgent { self.on_local_inference_model_delete(req).await } + #[custom_method(LocalInferenceModelEvictRequest)] + async fn dispatch_local_inference_model_evict( + &self, + req: LocalInferenceModelEvictRequest, + ) -> Result { + self.on_local_inference_model_evict(req).await + } + #[custom_method(LocalInferenceModelSettingsReadRequest)] async fn dispatch_local_inference_model_settings_read( &self, diff --git a/crates/goose/src/acp/server/local_inference.rs b/crates/goose/src/acp/server/local_inference.rs index 9796c200b4..6b733060e0 100644 --- a/crates/goose/src/acp/server/local_inference.rs +++ b/crates/goose/src/acp/server/local_inference.rs @@ -98,6 +98,26 @@ impl GooseAcpAgent { } } + pub(super) async fn on_local_inference_model_evict( + &self, + req: LocalInferenceModelEvictRequest, + ) -> Result { + #[cfg(feature = "local-inference")] + { + crate::providers::local_inference::configure_huggingface_auth(); + crate::providers::local_inference::management::evict_model(&req.model_id) + .await + .invalid_params_err()?; + Ok(EmptyResponse {}) + } + + #[cfg(not(feature = "local-inference"))] + { + let _ = req; + Err(local_inference_unavailable()) + } + } + pub(super) async fn on_local_inference_model_settings_read( &self, req: LocalInferenceModelSettingsReadRequest, diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 95220cf74e..e50e75ad5f 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -68,7 +68,7 @@ use tracing::{debug, error, info, instrument, warn}; const DEFAULT_MAX_TURNS: u32 = 1000; const DEFAULT_STOP_HOOK_BLOCK_CAP: u32 = 8; -const COMPACTION_THINKING_TEXT: &str = "goose is compacting the conversation..."; +const COMPACTION_PROGRESS_TEXT: &str = "goose is compacting the conversation..."; const MAX_TURNS_MESSAGE: &str = "I've reached the maximum number of actions I can do without user input. Would you like me to continue?"; const DEFAULT_FRONTEND_INSTRUCTIONS: &str = "The following tools are provided directly by the frontend and will be executed by the frontend when called."; @@ -1769,8 +1769,8 @@ impl Agent { yield AgentEvent::Message( Message::assistant().with_system_notification( - SystemNotificationType::ThinkingMessage, - COMPACTION_THINKING_TEXT, + SystemNotificationType::ProgressMessage, + COMPACTION_PROGRESS_TEXT, ) ); @@ -2060,6 +2060,17 @@ impl Agent { } if let Some(response) = response { + if !response.content.is_empty() + && response + .content + .iter() + .all(|content| matches!(content, MessageContent::SystemNotification(_))) + { + yield AgentEvent::Message(response); + tokio::task::yield_now().await; + continue; + } + let ToolCategorizeResult { frontend_requests, remaining_requests, @@ -2427,8 +2438,8 @@ impl Agent { ); yield AgentEvent::Message( Message::assistant().with_system_notification( - SystemNotificationType::ThinkingMessage, - COMPACTION_THINKING_TEXT, + SystemNotificationType::ProgressMessage, + COMPACTION_PROGRESS_TEXT, ) ); diff --git a/crates/goose/src/agents/reply_parts.rs b/crates/goose/src/agents/reply_parts.rs index ee5571b211..e504d48566 100644 --- a/crates/goose/src/agents/reply_parts.rs +++ b/crates/goose/src/agents/reply_parts.rs @@ -161,6 +161,13 @@ fn fill_stream_timing( } } +fn message_has_timing_content(message: &Message) -> bool { + message + .content + .iter() + .any(|content| !matches!(content, MessageContent::SystemNotification(_))) +} + impl Agent { pub async fn prepare_tools_and_prompt( &self, @@ -345,7 +352,7 @@ impl Agent { let (msg_opt, usage_opt) = result?; if let Some(msg) = msg_opt { - if first_content_at.is_none() { + if first_content_at.is_none() && message_has_timing_content(&msg) { first_content_at = Some(std::time::Instant::now()); } accumulated_message = Some(match accumulated_message { @@ -394,7 +401,9 @@ impl Agent { while let Some(result) = stream.next().await { let (message, mut usage) = result?; - if message.is_some() && first_content_at.is_none() { + if first_content_at.is_none() + && message.as_ref().is_some_and(message_has_timing_content) + { first_content_at = Some(std::time::Instant::now()); } if let Some(usage) = usage.as_mut() { @@ -656,7 +665,7 @@ pub fn is_tool_visible_to_model(tool: &Tool) -> bool { mod tests { use super::*; use crate::config::GooseMode; - use crate::conversation::message::Message; + use crate::conversation::message::{Message, SystemNotificationType}; use crate::providers::base::Provider; use crate::session::session_manager::SessionType; use async_trait::async_trait; @@ -1036,6 +1045,27 @@ mod tests { usage } + #[test] + fn message_has_timing_content_ignores_system_notification_only_messages() { + let message = Message::assistant().with_system_notification( + SystemNotificationType::ProgressMessage, + "Loading local model test-model...", + ); + + assert!(!message_has_timing_content(&message)); + } + + #[test] + fn message_has_timing_content_counts_user_visible_messages() { + let text_message = Message::assistant().with_text("hello"); + let mixed_message = Message::assistant() + .with_system_notification(SystemNotificationType::ProgressMessage, "Loading...") + .with_text("ready"); + + assert!(message_has_timing_content(&text_message)); + assert!(message_has_timing_content(&mixed_message)); + } + #[test] fn fill_stream_timing_fills_both_fields_when_stats_absent() { let request_started = Instant::now() - Duration::from_millis(100); diff --git a/ui/desktop/src/acp/__tests__/chatNotifications.test.ts b/ui/desktop/src/acp/__tests__/chatNotifications.test.ts index 71d83482a0..f1d9424052 100644 --- a/ui/desktop/src/acp/__tests__/chatNotifications.test.ts +++ b/ui/desktop/src/acp/__tests__/chatNotifications.test.ts @@ -81,6 +81,7 @@ function snapshotWithName(name: string): AcpChatSessionSnapshot { accumulatedTotalTokens: 0, }, notifications: [], + progressMessage: undefined, chatState: ChatState.Idle, sessionLoadError: undefined, activePromptAttemptId: null, @@ -102,6 +103,7 @@ function snapshotWithoutSession(): AcpChatSessionSnapshot { accumulatedTotalTokens: 0, }, notifications: [], + progressMessage: undefined, chatState: ChatState.Idle, sessionLoadError: undefined, activePromptAttemptId: null, diff --git a/ui/desktop/src/acp/__tests__/chatSessionController.test.ts b/ui/desktop/src/acp/__tests__/chatSessionController.test.ts index 2c2dfc1552..20afdb1b0d 100644 --- a/ui/desktop/src/acp/__tests__/chatSessionController.test.ts +++ b/ui/desktop/src/acp/__tests__/chatSessionController.test.ts @@ -108,6 +108,7 @@ function snapshotWithActivePrompt(activePromptAttemptId: string | null): AcpChat accumulatedTotalTokens: 0, }, notifications: [], + progressMessage: undefined, chatState: activePromptAttemptId ? ChatState.Streaming : ChatState.Idle, sessionLoadError: undefined, activePromptAttemptId, diff --git a/ui/desktop/src/acp/__tests__/sessionNotificationAdapter.test.ts b/ui/desktop/src/acp/__tests__/sessionNotificationAdapter.test.ts index 22e0546776..c9037f678c 100644 --- a/ui/desktop/src/acp/__tests__/sessionNotificationAdapter.test.ts +++ b/ui/desktop/src/acp/__tests__/sessionNotificationAdapter.test.ts @@ -583,13 +583,9 @@ describe('createAcpSessionNotificationAdapter', () => { status: { type: 'progress', message: 'Still working' }, }) ); - messages = expectOnlyMessagesChange(progressStateChanges); - - expect(firstContent(messages[2])).toMatchObject({ - type: 'systemNotification', - notificationType: 'thinkingMessage', - msg: 'Still working', - }); + expect(progressStateChanges).toEqual([ + { type: 'progressMessage', message: 'Still working' }, + ]); }); }); diff --git a/ui/desktop/src/acp/adapter/gooseSessionNotifications.ts b/ui/desktop/src/acp/adapter/gooseSessionNotifications.ts index c7153386ab..b2f11f471a 100644 --- a/ui/desktop/src/acp/adapter/gooseSessionNotifications.ts +++ b/ui/desktop/src/acp/adapter/gooseSessionNotifications.ts @@ -38,7 +38,9 @@ function applyStatusMessage( sessionId: string, update: Extract ): AcpChatStateChange[] { - const notificationType = update.status.type === 'notice' ? 'inlineMessage' : 'thinkingMessage'; + if (update.status.type === 'progress') { + return [{ type: 'progressMessage', message: update.status.message }]; + } state.messages.push({ id: `acp_status_${sessionId}_${Date.now()}_${Math.random().toString(36).slice(2, 10)}`, @@ -47,7 +49,7 @@ function applyStatusMessage( content: [ { type: 'systemNotification', - notificationType, + notificationType: 'inlineMessage', msg: update.status.message, }, ], diff --git a/ui/desktop/src/acp/adapter/shared.ts b/ui/desktop/src/acp/adapter/shared.ts index 073962701b..c52eb43d3c 100644 --- a/ui/desktop/src/acp/adapter/shared.ts +++ b/ui/desktop/src/acp/adapter/shared.ts @@ -5,6 +5,7 @@ import type { Message, NotificationEvent } from '../../types/message'; export type AcpChatStateChange = | { type: 'messages'; messages: Message[] } | { type: 'tokenState'; tokenState: Partial } + | { type: 'progressMessage'; message: string | undefined } | { type: 'sessionInfo'; name?: string; diff --git a/ui/desktop/src/acp/chatSessionStore.ts b/ui/desktop/src/acp/chatSessionStore.ts index c93f78c5dc..b773f260b7 100644 --- a/ui/desktop/src/acp/chatSessionStore.ts +++ b/ui/desktop/src/acp/chatSessionStore.ts @@ -19,6 +19,7 @@ export interface AcpChatSessionSnapshot { messages: Message[]; tokenState: TokenState; notifications: NotificationEvent[]; + progressMessage: string | undefined; chatState: ChatState; sessionLoadError: string | undefined; activePromptAttemptId: string | null; @@ -155,6 +156,7 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { messages: [], tokenState: { ...initialTokenState }, notifications: [], + progressMessage: undefined, chatState: ChatState.Idle, sessionLoadError: undefined, activePromptAttemptId: null, @@ -190,6 +192,7 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { const entry = getOrCreateEntry(sessionId); resetReplayState(entry); entry.sessionLoadError = undefined; + entry.progressMessage = undefined; entry.chatState = ChatState.LoadingConversation; return notify(sessionId, entry); }; @@ -198,6 +201,7 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { const entry = getOrCreateEntry(sessionId); entry.session = session; entry.sessionLoadError = undefined; + entry.progressMessage = undefined; entry.chatState = entry.activePromptAttemptId ? ChatState.Streaming : ChatState.Idle; return notify(sessionId, entry); }; @@ -208,6 +212,7 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { ) => { const entry = getOrCreateEntry(sessionId); entry.sessionLoadError = sessionLoadError; + entry.progressMessage = undefined; entry.chatState = ChatState.Idle; return notify(sessionId, entry); }; @@ -237,6 +242,9 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { const setChatState: AcpChatSessionActions['setChatState'] = (sessionId, chatState) => { const entry = getOrCreateEntry(sessionId); + if (chatState === ChatState.Idle) { + entry.progressMessage = undefined; + } entry.chatState = chatState; return notify(sessionId, entry); }; @@ -287,6 +295,7 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { entry.chatState = ChatState.Streaming; entry.sessionLoadError = undefined; entry.notifications = []; + entry.progressMessage = undefined; return notify(sessionId, entry); }; @@ -309,6 +318,7 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { entry.pendingCancelPromptAttemptId = promptAttemptId; entry.pendingUserInputRequestIds.clear(); discardPendingLocalSteerMessages(entry); + entry.progressMessage = undefined; entry.chatState = ChatState.Idle; return notify(sessionId, entry); }; @@ -385,6 +395,7 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { entry.promptCancellationRestoreState = null; entry.pendingUserInputRequestIds.clear(); discardPendingLocalSteerMessages(entry); + entry.progressMessage = undefined; entry.chatState = ChatState.Idle; entry.sessionLoadError = error; notify(sessionId, entry); @@ -416,6 +427,9 @@ function createAcpChatSessionStoreInternal(): AcpChatSessionStoreInternal { notification ) => { const entry = getOrCreateEntry(notification.sessionId); + if (shouldClearProgressMessage(notification)) { + entry.progressMessage = undefined; + } const changes = entry.adapter.apply(notification); applyChatStateChanges(entry, changes); return notify(notification.sessionId, entry); @@ -580,6 +594,9 @@ function applyChatStateChanges(entry: StoreEntry, changes: AcpChatStateChange[]) case 'tokenState': entry.tokenState = { ...entry.tokenState, ...change.tokenState }; break; + case 'progressMessage': + entry.progressMessage = change.message; + break; case 'sessionInfo': if (change.name && entry.session) { entry.session = { ...entry.session, name: change.name }; @@ -598,10 +615,22 @@ function applyChatStateChanges(entry: StoreEntry, changes: AcpChatStateChange[]) } } +function shouldClearProgressMessage(notification: SessionNotification): boolean { + switch (notification.update.sessionUpdate) { + case 'agent_message_chunk': + case 'agent_thought_chunk': + case 'tool_call': + return true; + default: + return false; + } +} + function resetReplayState(entry: StoreEntry): void { entry.messages = []; entry.tokenState = { ...initialTokenState }; entry.notifications = []; + entry.progressMessage = undefined; entry.activeRunId = null; entry.pendingCancelPromptAttemptId = null; entry.promptCancellationRestoreState = null; @@ -675,6 +704,7 @@ function snapshotFromEntry(entry: StoreEntry): AcpChatSessionSnapshot { messages: cloneMessages(entry.messages), tokenState: { ...entry.tokenState }, notifications: [...entry.notifications], + progressMessage: entry.progressMessage, chatState: entry.chatState, sessionLoadError: entry.sessionLoadError, activePromptAttemptId: entry.activePromptAttemptId, diff --git a/ui/desktop/src/acp/local-inference.ts b/ui/desktop/src/acp/local-inference.ts index 9003a4580d..5c0d589593 100644 --- a/ui/desktop/src/acp/local-inference.ts +++ b/ui/desktop/src/acp/local-inference.ts @@ -56,6 +56,11 @@ export async function deleteLocalModel(modelId: string): Promise { await client.goose.localInferenceModelsDelete_unstable({ modelId }); } +export async function evictLocalModel(modelId: string): Promise { + const client = await getAcpClient(); + await client.goose.localInferenceModelsEvict_unstable({ modelId }); +} + export async function getModelSettings(modelId: string): Promise { const client = await getAcpClient(); const response = await client.goose.localInferenceModelsSettingsRead_unstable({ modelId }); diff --git a/ui/desktop/src/components/BaseChat.tsx b/ui/desktop/src/components/BaseChat.tsx index e6e3bf8f38..d501f820c4 100644 --- a/ui/desktop/src/components/BaseChat.tsx +++ b/ui/desktop/src/components/BaseChat.tsx @@ -23,12 +23,7 @@ import { RecipeWarningModal } from './ui/RecipeWarningModal'; import { scanRecipe } from '../recipe'; import type { Recipe } from '../recipe'; import RecipeActivities from './recipes/RecipeActivities'; -import { - getThinkingMessage, - getTextAndImageContent, - type Message, - type UserInput, -} from '../types/message'; +import { getTextAndImageContent, type Message, type UserInput } from '../types/message'; import { substituteParameters } from '../utils/parameterSubstitution'; import { useAutoSubmit } from '../hooks/useAutoSubmit'; import { Goose } from './icons'; @@ -92,6 +87,7 @@ export default function BaseChat({ session, messages, chatState, + progressMessage, updateSession, handleSubmit, onSteerQueuedMessage, @@ -469,14 +465,7 @@ export default function BaseChat({ {chatState !== ChatState.Idle && (
- 0 - ? getThinkingMessage(messages[messages.length - 1]) - : undefined - } - /> +
)} diff --git a/ui/desktop/src/components/settings/localInference/LocalInferenceSettings.tsx b/ui/desktop/src/components/settings/localInference/LocalInferenceSettings.tsx index b9af91f84f..2c8befd800 100644 --- a/ui/desktop/src/components/settings/localInference/LocalInferenceSettings.tsx +++ b/ui/desktop/src/components/settings/localInference/LocalInferenceSettings.tsx @@ -8,6 +8,8 @@ import { Settings2, Eye, RefreshCw, + Cpu, + PowerOff, } from 'lucide-react'; import { Button } from '../../ui/button'; import { useModelAndProvider } from '../../ModelAndProviderContext'; @@ -18,6 +20,7 @@ import { getLocalModelDownloadProgress, cancelLocalModelDownload, deleteLocalModel, + evictLocalModel, type DownloadProgress, type DownloadModelRequest, type LocalModelResponse, @@ -106,6 +109,14 @@ const i18n = defineMessages({ id: 'localInferenceSettings.modelSettingsTitle', defaultMessage: 'Model settings', }, + loadedInMemory: { + id: 'localInferenceSettings.loadedInMemory', + defaultMessage: 'Loaded in memory', + }, + evictFromMemory: { + id: 'localInferenceSettings.evictFromMemory', + defaultMessage: 'Evict from memory', + }, vision: { id: 'localInferenceSettings.vision', defaultMessage: 'Vision', @@ -181,6 +192,7 @@ export const LocalInferenceSettings = () => { const [downloadRequests, setDownloadRequests] = useState>( new Map() ); + const [evictingModelId, setEvictingModelId] = useState(null); const [showAllFeatured, setShowAllFeatured] = useState(false); const [settingsOpenFor, setSettingsOpenFor] = useState(null); const { currentModel, currentProvider, refreshCurrentModelAndProvider } = useModelAndProvider(); @@ -371,6 +383,18 @@ export const LocalInferenceSettings = () => { } }; + const handleEvictModel = async (modelId: string) => { + setEvictingModelId(modelId); + try { + await evictLocalModel(modelId); + await loadModels(); + } catch (error) { + console.error('Failed to evict model:', error); + } finally { + setEvictingModelId(null); + } + }; + const handleHfDownloadStarted = (modelId: string, request: DownloadModelRequest) => { setDownloadRequests((prev) => new Map(prev).set(modelId, request)); pollDownloadProgress(modelId); @@ -393,6 +417,15 @@ export const LocalInferenceSettings = () => { .map(([modelId]) => modelId) ); + useEffect(() => { + if (downloadedModels.length === 0) return; + + const interval = setInterval(() => { + loadModels(); + }, 3000); + return () => clearInterval(interval); + }, [downloadedModels.length, loadModels]); + return (
@@ -526,18 +559,26 @@ export const LocalInferenceSettings = () => { : 'border-border-subtle bg-background-default hover:border-border-default' }`} > -
-
+
+
selectModel(model.id)} className="cursor-pointer" /> - {model.id} + + {model.id} + {formatBytes(model.sizeBytes)} + {model.isLoaded && ( + + + {intl.formatMessage(i18n.loadedInMemory)} + + )} {model.recommended && ( {intl.formatMessage(i18n.recommended)} @@ -545,7 +586,7 @@ export const LocalInferenceSettings = () => { )}
-
+
+ {model.isLoaded && ( + + )}