vrr/goose
2
0
Fork 0
mirror of https://github.com/block/goose.git synced 2026-07-09 16:09:22 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 22

Switch the local inference provider MLX backend to use the safemlx crate (#10304)

Browse source
This commit is contained in:
Jasper 2026-07-07 14:01:15 -07:00 • committed by GitHub
parent f96f62d985
commit 1bced6616b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
58 changed files with 1835 additions and 646 deletions
Download patch file Download diff file Expand all files Collapse all files

318
Cargo.lock generated
View file

@ -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",

54
crates/goose-cli/src/session/mod.rs
View file

@ -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: {}",

6
crates/goose-cli/src/session/output.rs
View file

@ -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(&notification.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(&notification.msg);
}

8
crates/goose-local-inference/Cargo.toml
View file

@ -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"] }

1
crates/goose-local-inference/src/backend.rs
View file

@ -22,6 +22,7 @@ pub(super) struct LocalGenerationRequest<'a> {
pub temperature: Option<f32>,
pub max_tokens: Option<i32>,
pub context_limit: usize,
pub model_load_ms: Option<u64>,
pub resolved_model: &'a ResolvedModelPaths,
pub draft_model_path: Option<std::path::PathBuf>,
pub message_id: &'a str,

333
crates/goose-local-inference/src/hf_models.rs
View file

@ -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::<Vec<_>>();
let filenames = mlx_download_filenames(&siblings)
.into_iter()
.collect::<std::collections::HashSet<_>>();
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<Vec<H
if results.len() >= 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<Option<HfModelIn
async fn get_local_model_info_for_repo_with_client(
client: &HFClient,
repo_id: &str,
) -> Result<Option<HfModelInfo>> {
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<u64>,
) -> Result<Option<HfModelInfo>> {
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<Option<HfModelInfo>> {
@ -1508,7 +1594,11 @@ async fn get_exact_name_local_model_info(model_name: &str) -> Result<Option<HfMo
Ok(None)
}
async fn model_info_to_local_model_info(info: ModelInfo) -> Result<Option<HfModelInfo>> {
async fn model_info_to_local_model_info(
repo: &HFRepository<RepoTypeModel>,
info: ModelInfo,
downloads_hint: Option<u64>,
) -> Result<Option<HfModelInfo>> {
let repo_id = info.id.clone();
let mut variants: Vec<HfModelVariant> = get_repo_gguf_variants(&repo_id)
.await
@ -1516,7 +1606,13 @@ async fn model_info_to_local_model_info(info: ModelInfo) -> Result<Option<HfMode
.iter()
.map(|variant| variant.to_model_variant(&repo_id))
.collect();
variants.extend(mlx_variants_from_model_info(&repo_id, &info));
if is_mlx_compatible_model_info(&info) {
let mlx_config = load_repo_config_json(repo).await.unwrap_or_else(|error| {
tracing::debug!(repo_id, %error, "Failed to load MLX config.json; falling back to API config");
info.config.clone()
});
variants.extend(mlx_variants_from_model_info(&repo_id, &info, &mlx_config));
}
if variants.is_empty() {
return Ok(None);
@ -1530,17 +1626,49 @@ async fn model_info_to_local_model_info(info: ModelInfo) -> Result<Option<HfMode
.next_back()
.unwrap_or(&repo_id)
.to_string();
let downloads = match best_download_count(info.downloads, downloads_hint) {
Some(downloads) => 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<u64>, hint: Option<u64>) -> Option<u64> {
primary
.filter(|downloads| *downloads > 0)
.or_else(|| hint.filter(|downloads| *downloads > 0))
}
async fn get_repo_downloads(repo_id: &str) -> Result<Option<u64>> {
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<Vec<HfModelVariant>> {
let mut variants: Vec<HfModelVariant> = get_repo_gguf_variants(repo_id)
.await
@ -1570,25 +1698,39 @@ pub async fn get_repo_mlx_variants(repo_id: &str) -> Result<Vec<HfModelVariant>>
])
.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<HfModelVariant> {
async fn load_repo_config_json(
repo: &HFRepository<RepoTypeModel>,
) -> Result<Option<serde_json::Value>> {
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<serde_json::Value>,
) -> Vec<HfModelVariant> {
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<HfModelV
size_bytes,
filename: None,
download_url: None,
description: mlx_variant_description(&info.config),
description: mlx_variant_description(mlx_config),
quality_rank: 91,
sharded: siblings
.iter()
.filter(|s| s.rfilename.ends_with(".safetensors"))
.count()
> 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<serde_json::Value>, 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<serde_json::Value>, 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<serde_json::Value>) -> Option<&str> {
@ -1644,22 +1824,8 @@ fn mlx_model_type(config: &Option<serde_json::Value>) -> 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<serde_json::Value>) -> 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<serde_json::Value>) -> 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<serde_json::Value>) -> Option<String> {
@ -1670,16 +1836,23 @@ fn mlx_unsupported_reason(config: &Option<serde_json::Value>) -> Option<String>
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<serde_json::Value>) -> Option<String> {
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<String> {
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<String> {
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<serde_json::Value>) -> 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<ResolvedLo
let repo = client.model(owner.to_string(), name.to_string());
let info = repo
.info()
.expand(vec!["siblings".to_string()])
.expand(vec!["siblings".to_string(), "safetensors".to_string()])
.send()
.await?;
let siblings = info.siblings.as_deref().unwrap_or(&[]);
let filenames = mlx_download_filenames(siblings);
let total_size = filenames
.iter()
.filter_map(|filename| {
siblings
.iter()
.find(|s| s.rfilename == *filename)
.and_then(|s| s.size)
})
.sum();
let total_size = mlx_download_size_bytes(&info, siblings);
let progress = HfDownloadProgress::new(repo_id.to_string(), total_size);
progress.init();
let mut snapshot_path = None;
@ -2103,8 +2272,13 @@ impl HfDownloadProgress {
}
fn finish_file(&self, size_bytes: u64) {
let observed_size = self
.state
.lock()
.map(|state| state.current_file_total_bytes.max(state.bytes_downloaded))
.unwrap_or(0);
if let Ok(mut completed_bytes) = self.completed_bytes.lock() {
*completed_bytes = completed_bytes.saturating_add(size_bytes);
*completed_bytes = completed_bytes.saturating_add(size_bytes.max(observed_size));
}
if let Ok(mut state) = self.state.lock() {
state.bytes_downloaded = 0;
@ -2117,10 +2291,15 @@ impl HfDownloadProgress {
let completed_bytes = self.completed_bytes.lock().map(|value| *value).unwrap_or(0);
if let Ok(state) = self.state.lock() {
let bytes_downloaded = completed_bytes.saturating_add(state.bytes_downloaded);
let total_bytes =
self.total_bytes
.max(completed_bytes.saturating_add(
state.current_file_total_bytes.max(state.bytes_downloaded),
));
update_download_manager_progress(
&self.model_id,
bytes_downloaded.min(self.total_bytes),
self.total_bytes.max(state.current_file_total_bytes),
bytes_downloaded.min(total_bytes),
total_bytes,
state.speed_bps,
);
}
@ -2152,10 +2331,12 @@ impl ProgressHandler for HfDownloadProgress {
file.bytes_completed
}
})
.max()
.unwrap_or(0);
.sum();
let total_bytes = files.iter().map(|file| file.total_bytes).sum();
if let Ok(mut state) = self.state.lock() {
state.bytes_downloaded = state.bytes_downloaded.max(bytes_downloaded);
state.current_file_total_bytes =
state.current_file_total_bytes.max(total_bytes);
}
self.update_progress_from_state();
}
@ -2168,7 +2349,7 @@ impl ProgressHandler for HfDownloadProgress {
return;
}
if let Ok(mut state) = self.state.lock() {
state.bytes_downloaded = *bytes_completed;
state.bytes_downloaded = state.bytes_downloaded.max(*bytes_completed);
state.current_file_total_bytes =
(*total_bytes).max(state.current_file_total_bytes);
state.speed_bps = bytes_per_sec.map(|speed| speed as u64);

391
crates/goose-local-inference/src/lib.rs
View file

@ -14,6 +14,7 @@ mod mlx;
pub(crate) mod multimodal;
#[cfg(feature = "mlx")]
mod native_tool_parsing;
pub(crate) mod thinking_output;
#[cfg(feature = "mlx")]
mod tool_emulation;
mod tool_parsing;
@ -23,7 +24,9 @@ use async_stream::try_stream;
use async_trait::async_trait;
use backend::{BackendLoadedModel, LocalInferenceBackend};
use goose_provider_types::base::{MessageStream, Provider, ProviderDescriptor, ProviderMetadata};
use goose_provider_types::conversation::message::{Message, MessageContent};
use goose_provider_types::conversation::message::{
Message, MessageContent, SystemNotificationType,
};
use goose_provider_types::conversation::token_usage::{ProviderUsage, Usage};
use goose_provider_types::errors::ProviderError;
use goose_provider_types::images::ImageFormat;
@ -34,13 +37,33 @@ use local_model_registry::ChatTemplate;
use mlx::{MlxBackend, MLX_BACKEND_ID};
use rmcp::model::Tool;
use serde_json::{json, Value};
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use std::sync::{Arc, Mutex as StdMutex, Weak};
use tokio::sync::Mutex;
use std::sync::{Arc, Mutex as StdMutex};
use tokio::sync::{Mutex, Notify};
use uuid::Uuid;
type ModelSlot = Arc<Mutex<Option<Box<dyn BackendLoadedModel>>>>;
type ModelSlotHandle = Arc<ModelSlot>;
struct ModelSlot {
state: Mutex<ModelSlotState>,
notify: Notify,
}
enum ModelSlotState {
Empty,
Loading,
Loaded(Box<dyn BackendLoadedModel>),
}
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<HashMap<ModelCacheKey, ModelSlot>>,
models: StdMutex<HashMap<ModelCacheKey, ModelSlotHandle>>,
cold_load_lock: Mutex<()>,
backends: HashMap<&'static str, Arc<dyn LocalInferenceBackend>>,
}
@ -72,19 +96,17 @@ pub fn builtin_chat_template_names() -> Vec<String> {
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<Weak<InferenceRuntime>> = StdMutex::new(Weak::new());
static RUNTIME: StdMutex<Option<Arc<InferenceRuntime>>> = StdMutex::new(None);
fn current_runtime() -> Option<Arc<InferenceRuntime>> {
RUNTIME.lock().expect("runtime lock poisoned").clone()
}
impl InferenceRuntime {
pub fn get_or_init() -> Result<Arc<Self>> {
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<dyn LocalInferenceBackend> = Arc::new(LlamaCppBackend::new()?);
let mlx_backend: Arc<dyn LocalInferenceBackend> = 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<ModelSlot> {
fn model_slot(&self, key: &ModelCacheKey) -> Option<ModelSlotHandle> {
let map = self.models.lock().expect("model cache lock poisoned");
map.get(key).cloned()
}
fn other_model_slots(&self, keep_key: &ModelCacheKey) -> Vec<ModelSlotHandle> {
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<bool, ProviderError> {
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<HashSet<String>, 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::<Vec<_>>()
};
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<bool, ProviderError> {
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::<Vec<_>>()
};
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::<bool>("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<Message>, Option<ProviderUsage>), 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! {

5
crates/goose-local-inference/src/llamacpp/inference_emulated_tools.rs
View file

@ -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";

81
crates/goose-local-inference/src/llamacpp/inference_engine.rs
View file

@ -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<String> {
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<String>,

4
crates/goose-local-inference/src/llamacpp/inference_native_tools.rs
View file

@ -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(

27
crates/goose-local-inference/src/management.rs
View file

@ -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<LocalInferenceModelsListResponse> {
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<LocalInferenceModelDto> = 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<bool> {
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<LocalInferenceModelSettingsReadResponse> {
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<String>,
) -> 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

629
crates/goose-local-inference/src/mlx.rs
View file

@ -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<u32>,
generated_text: String,
draft_stats: Option<DraftStats>,
time_to_first_token_ms: Option<u64>,
streamed_response: bool,
}
struct MlxLoadedModel {
model: LoadedModel,
tokenizer: Tokenizer,
model_dir: PathBuf,
stop_token_ids: Vec<u32>,
}
impl BackendLoadedModel for MlxLoadedModel {
@ -242,6 +306,44 @@ mod imp {
}
}
fn mlx_stop_token_ids(model: &LoadedModel, model_dir: &Path) -> Vec<u32> {
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<u32> {
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::<serde_json::Value>(&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<u32> {
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<u32>, Option<DraftStats>), ProviderError> {
let mut cache: Vec<Option<ConcatKeyValueCache>> = Vec::new();
prng_key: Option<Array>,
stream: &Stream,
generation_started: std::time::Instant,
mut emitter: MlxStreamEmitter<'_>,
) -> Result<MlxGeneration, ProviderError> {
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::<u32>();
let token_id = token.item::<u32>(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<Option<&'a str>, 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<i32>,
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<Message>, Option<ProviderUsage>), ProviderError>,
>,
output_filter: ThinkingOutputFilter,
emulator_parser: Option<StreamingEmulatorParser>,
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<Message>, Option<ProviderUsage>), 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<bool, ProviderError> {
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<Message>, Option<ProviderUsage>), 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<u32>) {
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<u32>) -> Result<Option<Array>, 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 reasoning</think>visible answer",
"<|im_start|>assistant\n<think>\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 reasoning</think>visible answer";
let (content, thinking) =
split_generated_thinking(generated, "<|im_start|>assistant\n<think>\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"))]

81
crates/goose-local-inference/src/thinking_output.rs Normal file
View file

@ -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<String> {
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
}
}

1
crates/goose-provider-types/src/conversation/message.rs
View file

@ -253,6 +253,7 @@ pub struct FrontendToolRequest {
#[serde(rename_all = "camelCase")]
pub enum SystemNotificationType {
ThinkingMessage,
ProgressMessage,
InlineMessage,
CreditsExhausted,
}

1
crates/goose-provider-types/src/conversation/token_usage.rs
View file

@ -25,6 +25,7 @@ pub enum CostSource {
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ProviderStats {
pub time_to_first_token_ms: Option<u64>,
pub model_load_ms: Option<u64>,
pub elapsed_ms: Option<u64>,
pub output_tokens: Option<usize>,
pub draft: Option<DraftStats>,

11
crates/goose-sdk-types/src/custom_requests.rs
View file

@ -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",

5
crates/goose/acp-meta.json
View file

@ -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",

26
crates/goose/acp-schema.json
View file

@ -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": [
{

72
crates/goose/src/acp/server.rs
View file

@ -2072,6 +2072,23 @@ fn send_status_message_update(
Ok(())
}
fn send_progress_message_update(
cx: &ConnectionTo<Client>,
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<StatusMessage> {
@ -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<Client>,
acp_session_id: &SessionId,
session_id: &str,
agent: &Arc<Agent>,
) -> 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<Client>,
@ -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();

17
crates/goose/src/acp/server/config.rs
View file

@ -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<bool, agent_client_protocol::Error> {
#[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,

8
crates/goose/src/acp/server/custom_dispatch.rs
View file

@ -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<EmptyResponse, agent_client_protocol::Error> {
self.on_local_inference_model_evict(req).await
}
#[custom_method(LocalInferenceModelSettingsReadRequest)]
async fn dispatch_local_inference_model_settings_read(
&self,

20
crates/goose/src/acp/server/local_inference.rs
View file

@ -98,6 +98,26 @@ impl GooseAcpAgent {
}
}
pub(super) async fn on_local_inference_model_evict(
&self,
req: LocalInferenceModelEvictRequest,
) -> Result<EmptyResponse, agent_client_protocol::Error> {
#[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,

21
crates/goose/src/agents/agent.rs
View file

@ -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,
)
);

36
crates/goose/src/agents/reply_parts.rs
View file

@ -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);

2
ui/desktop/src/acp/__tests__/chatNotifications.test.ts
View file

@ -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,

1
ui/desktop/src/acp/__tests__/chatSessionController.test.ts
View file

@ -108,6 +108,7 @@ function snapshotWithActivePrompt(activePromptAttemptId: string | null): AcpChat
accumulatedTotalTokens: 0,
},
notifications: [],
progressMessage: undefined,
chatState: activePromptAttemptId ? ChatState.Streaming : ChatState.Idle,
sessionLoadError: undefined,
activePromptAttemptId,

10
ui/desktop/src/acp/__tests__/sessionNotificationAdapter.test.ts
View file

@ -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' },
]);
});
});

6
ui/desktop/src/acp/adapter/gooseSessionNotifications.ts
View file

@ -38,7 +38,9 @@ function applyStatusMessage(
sessionId: string,
update: Extract<GooseSessionNotification_unstable['update'], { sessionUpdate: 'status_message' }>
): 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,
},
],

1
ui/desktop/src/acp/adapter/shared.ts
View file

@ -5,6 +5,7 @@ import type { Message, NotificationEvent } from '../../types/message';
export type AcpChatStateChange =
| { type: 'messages'; messages: Message[] }
| { type: 'tokenState'; tokenState: Partial<TokenState> }
| { type: 'progressMessage'; message: string | undefined }
| {
type: 'sessionInfo';
name?: string;

30
ui/desktop/src/acp/chatSessionStore.ts
View file

@ -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,

5
ui/desktop/src/acp/local-inference.ts
View file

@ -56,6 +56,11 @@ export async function deleteLocalModel(modelId: string): Promise<void> {
await client.goose.localInferenceModelsDelete_unstable({ modelId });
}
export async function evictLocalModel(modelId: string): Promise<void> {
const client = await getAcpClient();
await client.goose.localInferenceModelsEvict_unstable({ modelId });
}
export async function getModelSettings(modelId: string): Promise<ModelSettings> {
const client = await getAcpClient();
const response = await client.goose.localInferenceModelsSettingsRead_unstable({ modelId });

17
ui/desktop/src/components/BaseChat.tsx
View file

@ -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 && (
<div className="absolute bottom-1 left-4 z-20 pointer-events-none">
<LoadingGoose
chatState={chatState}
message={
messages.length > 0
? getThinkingMessage(messages[messages.length - 1])
: undefined
}
/>
<LoadingGoose chatState={chatState} message={progressMessage} />
</div>
)}
</div>

60
ui/desktop/src/components/settings/localInference/LocalInferenceSettings.tsx
View file

@ -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<Map<string, DownloadModelRequest>>(
new Map()
);
const [evictingModelId, setEvictingModelId] = useState<string | null>(null);
const [showAllFeatured, setShowAllFeatured] = useState(false);
const [settingsOpenFor, setSettingsOpenFor] = useState<string | null>(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 (
<div className="space-y-6">
<div>
@ -526,18 +559,26 @@ export const LocalInferenceSettings = () => {
: 'border-border-subtle bg-background-default hover:border-border-default'
}`}
>
<div className="flex items-center justify-between">
<div className="flex items-center gap-2">
<div className="flex items-center justify-between gap-3">
<div className="flex min-w-0 items-center gap-2 flex-wrap">
<input
type="radio"
checked={isSelected}
onChange={() => selectModel(model.id)}
className="cursor-pointer"
/>
<span className="text-sm font-medium text-text-default">{model.id}</span>
<span className="text-sm font-medium text-text-default break-all">
{model.id}
</span>
<span className="text-xs text-text-muted">
{formatBytes(model.sizeBytes)}
</span>
{model.isLoaded && (
<span className="inline-flex items-center gap-1 text-xs text-green-400 bg-green-500/10 px-2 py-0.5 rounded">
<Cpu className="w-3 h-3" />
{intl.formatMessage(i18n.loadedInMemory)}
</span>
)}
{model.recommended && (
<span className="text-xs bg-blue-500 text-white px-2 py-0.5 rounded">
{intl.formatMessage(i18n.recommended)}
@ -545,7 +586,7 @@ export const LocalInferenceSettings = () => {
)}
<VisionBadge model={model} intl={intl} />
</div>
<div className="flex items-center gap-1">
<div className="flex shrink-0 items-center gap-1">
<Button
variant="ghost"
size="sm"
@ -554,6 +595,17 @@ export const LocalInferenceSettings = () => {
>
<Settings2 className="w-4 h-4" />
</Button>
{model.isLoaded && (
<Button
variant="ghost"
size="sm"
onClick={() => handleEvictModel(model.id)}
disabled={evictingModelId === model.id}
title={intl.formatMessage(i18n.evictFromMemory)}
>
<PowerOff className="w-4 h-4" />
</Button>
)}
<Button
variant="ghost"
size="sm"

2
ui/desktop/src/hooks/useChatSession.ts
View file

@ -61,6 +61,7 @@ export function useChatSession({
const messages = acpSnapshot?.messages ?? [];
const session = acpSnapshot?.session;
const chatState = acpSnapshot?.chatState ?? ChatState.LoadingConversation;
const progressMessage = acpSnapshot?.progressMessage;
const sessionLoadError = acpSnapshot?.sessionLoadError;
const tokenState = acpSnapshot?.tokenState ?? initialTokenState;
const queueProcessingBlocked = acpSnapshot?.pendingCancelPromptAttemptId != null;
@ -319,6 +320,7 @@ export function useChatSession({
messages,
session,
chatState,
progressMessage,
updateSession,
handleSubmit,
onSteerQueuedMessage,

1
ui/desktop/src/hooks/useChatSessionTypes.ts
View file

@ -13,6 +13,7 @@ export interface UseChatSessionResult {
session?: Session;
messages: Message[];
chatState: ChatState;
progressMessage?: string;
updateSession: (updater: (session: Session) => Session) => void;
handleSubmit: (input: UserInput) => Promise<void>;
onSteerQueuedMessage?: (input: UserInput) => Promise<boolean>;

6
ui/desktop/src/i18n/messages/de.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloading": {
"defaultMessage": "Wird heruntergeladen"
},
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": {
"defaultMessage": "Empfohlene Modelle"
},
"localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Melden Sie sich an, um die Ratenbegrenzungen beim Suchen und Herunterladen von Modellen zu erhöhen und auf private oder zugriffsbeschränkte Hugging-Face-Repositories zuzugreifen."
},
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": {
"defaultMessage": "Modelleinstellungen"
},

6
ui/desktop/src/i18n/messages/en.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloading": {
"defaultMessage": "Downloading"
},
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": {
"defaultMessage": "Featured Models"
},
"localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Sign in to increase rate limits when searching and downloading models, and to access private or gated Hugging Face repositories."
},
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": {
"defaultMessage": "Model Settings"
},

6
ui/desktop/src/i18n/messages/es.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloadProgress": {
"defaultMessage": "{downloaded} / {total} ({percent}%)"
},
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": {
"defaultMessage": "Modelos destacados"
},
"localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Inicia sesión para aumentar los límites de uso al buscar y descargar modelos, y para acceder a repositorios de Hugging Face privados o restringidos."
},
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": {
"defaultMessage": "Ajustes del modelo"
},

6
ui/desktop/src/i18n/messages/fr.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloading": {
"defaultMessage": "Téléchargement en cours"
},
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": {
"defaultMessage": "Modèles en vedette"
},
"localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Connectez-vous pour augmenter les limites de débit lors de la recherche et du téléchargement de modèles, et pour accéder aux dépôts Hugging Face privés ou à accès restreint."
},
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": {
"defaultMessage": "Paramètres du modèle"
},

6
ui/desktop/src/i18n/messages/hi.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloadProgress": {
"defaultMessage": "{downloaded} / {total} ({percent}%)"
},
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": {
"defaultMessage": "विशेष रुप से प्रदर्शित मॉडल"
},
"localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "मॉडल खोजते और डाउनलोड करते समय दर सीमा बढ़ाने के लिए और निजी या गेटेड हगिंग फेस रिपॉजिटरी तक पहुंचने के लिए साइन इन करें।"
},
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": {
"defaultMessage": "मॉडल Settings"
},

6
ui/desktop/src/i18n/messages/id.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloading": {
"defaultMessage": "Mengunduh"
},
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": {
"defaultMessage": "Model Unggulan"
},
"localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Masuk untuk meningkatkan batas laju saat mencari dan mengunduh model, serta untuk mengakses repositori Hugging Face yang privat atau terbatas."
},
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": {
"defaultMessage": "Pengaturan Model"
},

6
ui/desktop/src/i18n/messages/it.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloading": {
"defaultMessage": "Download in corso"
},
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": {
"defaultMessage": "Modelli in evidenza"
},
"localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Accedi per aumentare i limiti di frequenza durante la ricerca e il download dei modelli e per accedere ai repository Hugging Face privati o ad accesso limitato."
},
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": {
"defaultMessage": "Impostazioni modello"
},

6
ui/desktop/src/i18n/messages/ja.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloadProgress": {
"defaultMessage": "{downloaded} / {total} ({percent}%)"
},
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": {
"defaultMessage": "注目モデル"
},
"localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "モデルの検索・ダウンロード時のレート制限を緩和し、非公開またはゲート付きHugging Faceリポジトリにアクセスするにはサインインしてください。"
},
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": {
"defaultMessage": "モデル設定"
},

6
ui/desktop/src/i18n/messages/ko.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloadProgress": {
"defaultMessage": "{downloaded} / {total} ({percent}%)"
},
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": {
"defaultMessage": "주요 모델"
},
"localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "모델 검색 및 다운로드 시 속도 제한 한도를 높이고 비공개 또는 접근 제한된 Hugging Face 리포지토리에 액세스하려면 로그인하세요."
},
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": {
"defaultMessage": "모델 설정"
},

6
ui/desktop/src/i18n/messages/ms.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloading": {
"defaultMessage": "Memuat turun"
},
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": {
"defaultMessage": "Model Pilihan"
},
"localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Log masuk untuk meningkatkan had kadar semasa mencari dan memuat turun model, dan untuk mengakses repositori Hugging Face peribadi atau berpagar."
},
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": {
"defaultMessage": "Tetapan Model"
},

6
ui/desktop/src/i18n/messages/pt.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloading": {
"defaultMessage": "Baixando"
},
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": {
"defaultMessage": "Modelos em Destaque"
},
"localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Faça login para aumentar os limites de taxa ao pesquisar e baixar modelos, e para acessar repositórios privados ou restritos do Hugging Face."
},
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": {
"defaultMessage": "Configurações do Modelo"
},

6
ui/desktop/src/i18n/messages/ru.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloadProgress": {
"defaultMessage": "{downloaded} / {total} ({percent}%)"
},
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": {
"defaultMessage": "Рекомендуемые модели"
},
"localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Войдите, чтобы увеличить лимиты запросов при поиске и загрузке моделей, а также получить доступ к приватным или ограниченным репозиториям Hugging Face."
},
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": {
"defaultMessage": "Настройки модели"
},

6
ui/desktop/src/i18n/messages/tr.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloadProgress": {
"defaultMessage": "{downloaded} / {total} (%{percent})"
},
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": {
"defaultMessage": "Öne Çıkan Modeller"
},
"localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Model ararken ve indirirken hız sınırlarını artırmak ve özel ya da erişimi kısıtlı Hugging Face depolarına erişmek için oturum açın."
},
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": {
"defaultMessage": "Modeli Ayarları"
},

6
ui/desktop/src/i18n/messages/vi.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloading": {
"defaultMessage": "Đang tải xuống"
},
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": {
"defaultMessage": "Mô hình nổi bật"
},
"localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "Đăng nhập để tăng giới hạn tốc độ khi tìm kiếm và tải xuống mô hình, và để truy cập các kho lưu trữ Hugging Face riêng tư hoặc bị giới hạn."
},
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": {
"defaultMessage": "Cài đặt mô hình"
},

6
ui/desktop/src/i18n/messages/zh-CN.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloadProgress": {
"defaultMessage": "{downloaded} / {total}（{percent}%）"
},
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": {
"defaultMessage": "推荐模型"
},
"localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "登录后可在搜索和下载模型时提高速率限制，并访问私有或受限的 Hugging Face 仓库。"
},
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": {
"defaultMessage": "模型设置"
},

6
ui/desktop/src/i18n/messages/zh-TW.json
View file

@ -1736,12 +1736,18 @@
"localInferenceSettings.downloading": {
"defaultMessage": "下載中"
},
"localInferenceSettings.evictFromMemory": {
"defaultMessage": "Evict from memory"
},
"localInferenceSettings.featuredModels": {
"defaultMessage": "精選模型"
},
"localInferenceSettings.huggingFaceSignInNote": {
"defaultMessage": "登入以在搜尋與下載模型時提高速率上限，並存取私人或受限的 Hugging Face 儲存庫。"
},
"localInferenceSettings.loadedInMemory": {
"defaultMessage": "Loaded in memory"
},
"localInferenceSettings.modelSettings": {
"defaultMessage": "模型設定"
},

20
ui/desktop/src/types/message.ts
View file

@ -93,7 +93,11 @@ type ContentIcon = {
theme?: 'light' | 'dark' | JsonObject;
};
export type SystemNotificationType = 'thinkingMessage' | 'inlineMessage' | 'creditsExhausted';
export type SystemNotificationType =
| 'thinkingMessage'
| 'progressMessage'
| 'inlineMessage'
| 'creditsExhausted';
export type SystemNotificationContent = {
data?: unknown;
@ -443,17 +447,3 @@ export function hasCompletedToolCalls(message: Message): boolean {
const toolRequests = getToolRequests(message);
return toolRequests.length > 0;
}
export function getThinkingMessage(message: Message | undefined): string | undefined {
if (!message || message.role !== 'assistant') {
return undefined;
}
for (const content of message.content) {
if (content.type === 'systemNotification' && content.notificationType === 'thinkingMessage') {
return content.msg;
}
}
return undefined;
}

26
ui/desktop/src/utils/conversionUtils.test.ts Normal file
View file

@ -0,0 +1,26 @@
import { describe, expect, it } from 'vitest';
import { errorMessage } from './conversionUtils';
describe('errorMessage', () => {
it('prefers ACP JSON-RPC error data over generic messages', () => {
expect(
errorMessage({
error: {
message: 'Invalid params',
data: 'MLX backend error: failed to load model',
},
})
).toBe('MLX backend error: failed to load model');
});
it('prefers ACP JSON-RPC error data from Error instances', () => {
const error = Object.assign(new Error('Invalid params'), {
error: {
message: 'Invalid params',
data: 'MLX backend error: failed to load model',
},
});
expect(errorMessage(error)).toBe('MLX backend error: failed to load model');
});
});

18
ui/desktop/src/utils/conversionUtils.ts
View file

@ -13,6 +13,11 @@ export async function safeJsonParse<T>(
}
export function errorMessage(err: Error | unknown, default_value?: string) {
const acpData = acpErrorData(err);
if (typeof acpData === 'string') {
return acpData;
}
if (err instanceof Error) {
return err.message;
} else if (typeof err === 'object' && err !== null && 'message' in err) {
@ -22,6 +27,19 @@ export function errorMessage(err: Error | unknown, default_value?: string) {
}
}
function acpErrorData(err: unknown): unknown {
if (typeof err !== 'object' || err === null) {
return undefined;
}
const candidate = 'error' in err && isRecord(err.error) ? err.error : err;
return isRecord(candidate) ? candidate.data : undefined;
}
function isRecord(value: unknown): value is Record<string, unknown> {
return typeof value === 'object' && value !== null;
}
export function formatErrorForLogging(error: unknown): string {
if (error instanceof Error) {
return `${error.name}: ${error.message}${error.stack ? `\n${error.stack}` : ''}`;

10
ui/sdk/src/generated/client.gen.ts
View file

@ -121,6 +121,7 @@ import type {
LocalInferenceModelDownloadProgressResponse_unstable,
LocalInferenceModelDownloadRequest_unstable,
LocalInferenceModelDownloadResponse_unstable,
LocalInferenceModelEvictRequest_unstable,
LocalInferenceModelSettingsReadRequest_unstable,
LocalInferenceModelSettingsReadResponse_unstable,
LocalInferenceModelSettingsUpdateRequest_unstable,
@ -1377,6 +1378,15 @@ export class GooseExtClient {
);
}
async localInferenceModelsEvict_unstable(
params: LocalInferenceModelEvictRequest_unstable,
): Promise<void> {
await this.conn.extMethod(
"_goose/unstable/local-inference/models/evict",
params,
);
}
async localInferenceModelsSettingsRead_unstable(
params: LocalInferenceModelSettingsReadRequest_unstable,
): Promise<LocalInferenceModelSettingsReadResponse_unstable> {

7
ui/sdk/src/generated/index.ts
View file

File diff suppressed because one or more lines are too long

7
ui/sdk/src/generated/types.gen.ts
View file

@ -2329,6 +2329,7 @@ export type LocalInferenceModelDto = {
sizeBytes: number;
status: LocalInferenceModelDownloadStatusDto;
recommended: boolean;
isLoaded: boolean;
settings: LocalInferenceModelSettingsDto;
visionCapable: boolean;
mmprojStatus?: LocalInferenceModelDownloadStatusDto | null;
@ -2433,6 +2434,10 @@ export type LocalInferenceModelDeleteRequest_unstable = {
modelId: string;
};
export type LocalInferenceModelEvictRequest_unstable = {
modelId: string;
};
export type LocalInferenceModelSettingsReadRequest_unstable = {
modelId: string;
};
@ -2618,7 +2623,7 @@ export type RecipeParamsAction = 'submit' | 'cancel';
export type ExtRequest = {
id: string;
method: string;
params?: AddSessionExtensionRequest_unstable | RemoveSessionExtensionRequest_unstable | GetToolsRequest_unstable | SetToolPermissionsRequest_unstable | GooseToolCallRequest_unstable | ReadResourceRequest_unstable | AppsListRequest_unstable | AppsExportRequest_unstable | AppsImportRequest_unstable | UpdateWorkingDirRequest_unstable | SetSessionSystemPromptRequest_unstable | SteerSessionRequest_unstable | DiagnosticsGetRequest_unstable | ListPromptsRequest_unstable | GetPromptRequest_unstable | SavePromptRequest_unstable | ResetPromptRequest_unstable | DeleteSessionRequest | GetConfigExtensionsRequest_unstable | GetAvailableExtensionsRequest_unstable | AddConfigExtensionRequest_unstable | RemoveConfigExtensionRequest_unstable | SetConfigExtensionEnabledRequest_unstable | GetSessionExtensionsRequest_unstable | ListProvidersRequest_unstable | ProviderSupportedModelsListRequest_unstable | ProviderCatalogListRequest_unstable | ProviderSetupCatalogListRequest_unstable | ProviderCatalogTemplateRequest_unstable | CustomProviderCreateRequest_unstable | CustomProviderReadRequest_unstable | CustomProviderUpdateRequest_unstable | CustomProviderDeleteRequest_unstable | RefreshProviderInventoryRequest_unstable | ProviderConfigReadRequest_unstable | ProviderConfigStatusRequest_unstable | ProviderConfigSaveRequest_unstable | ProviderConfigDeleteRequest_unstable | ProviderConfigAuthenticateRequest_unstable | ProviderSecretsListRequest_unstable | ProviderSecretDeleteRequest_unstable | CanonicalModelInfoRequest_unstable | PreferencesReadRequest_unstable | PreferencesSaveRequest_unstable | PreferencesRemoveRequest_unstable | ConfigReadRequest_unstable | ConfigUpsertRequest_unstable | ConfigRemoveRequest_unstable | ConfigReadAllRequest_unstable | DefaultsReadRequest_unstable | DefaultsSaveRequest_unstable | DefaultsClearRequest_unstable | OnboardingImportScanRequest_unstable | OnboardingImportApplyRequest_unstable | ExportSessionRequest_unstable | ImportSessionRequest_unstable | ShareSessionNostrRequest_unstable | EncodeRecipeRequest_unstable | DecodeRecipeRequest_unstable | ScanRecipeRequest_unstable | ListRecipesRequest_unstable | DeleteRecipeRequest_unstable | ScheduleRecipeRequest_unstable | SetRecipeSlashCommandRequest_unstable | SaveRecipeRequest_unstable | ParseRecipeRequest_unstable | RecipeToYamlRequest_unstable | ListSchedulesRequest_unstable | ListScheduleSessionsRequest_unstable | CreateScheduleRequest_unstable | DeleteScheduleRequest_unstable | PauseScheduleRequest_unstable | UnpauseScheduleRequest_unstable | UpdateScheduleRequest_unstable | RunScheduleNowRequest_unstable | KillRunningJobRequest_unstable | InspectRunningJobRequest_unstable | GetSessionInfoRequest_unstable | TruncateSessionConversationRequest_unstable | UpdateSessionProjectRequest_unstable | RenameSessionRequest_unstable | ArchiveSessionRequest_unstable | UnarchiveSessionRequest_unstable | CreateSourceRequest_unstable | ListSourcesRequest_unstable | ListAgentMentionsRequest_unstable | ListSlashCommandsRequest_unstable | UpdateSourceRequest_unstable | DeleteSourceRequest_unstable | ExportSourceRequest_unstable | ImportSourcesRequest_unstable | DictationTranscribeRequest_unstable | DictationConfigRequest_unstable | DictationSecretSaveRequest_unstable | DictationSecretDeleteRequest_unstable | DictationModelsListRequest_unstable | DictationModelDownloadRequest_unstable | DictationModelDownloadProgressRequest_unstable | DictationModelCancelRequest_unstable | DictationModelDeleteRequest_unstable | DictationModelSelectRequest_unstable | LocalInferenceModelsListRequest_unstable | LocalInferenceModelDownloadRequest_unstable | LocalInferenceModelDownloadProgressRequest_unstable | LocalInferenceModelDownloadCancelRequest_unstable | LocalInferenceModelDeleteRequest_unstable | LocalInferenceModelSettingsReadRequest_unstable | LocalInferenceModelSettingsUpdateRequest_unstable | LocalInferenceHuggingFaceSearchRequest_unstable | LocalInferenceHuggingFaceRepoVariantsRequest_unstable | LocalInferenceBuiltinChatTemplatesListRequest_unstable | {
params?: AddSessionExtensionRequest_unstable | RemoveSessionExtensionRequest_unstable | GetToolsRequest_unstable | SetToolPermissionsRequest_unstable | GooseToolCallRequest_unstable | ReadResourceRequest_unstable | AppsListRequest_unstable | AppsExportRequest_unstable | AppsImportRequest_unstable | UpdateWorkingDirRequest_unstable | SetSessionSystemPromptRequest_unstable | SteerSessionRequest_unstable | DiagnosticsGetRequest_unstable | ListPromptsRequest_unstable | GetPromptRequest_unstable | SavePromptRequest_unstable | ResetPromptRequest_unstable | DeleteSessionRequest | GetConfigExtensionsRequest_unstable | GetAvailableExtensionsRequest_unstable | AddConfigExtensionRequest_unstable | RemoveConfigExtensionRequest_unstable | SetConfigExtensionEnabledRequest_unstable | GetSessionExtensionsRequest_unstable | ListProvidersRequest_unstable | ProviderSupportedModelsListRequest_unstable | ProviderCatalogListRequest_unstable | ProviderSetupCatalogListRequest_unstable | ProviderCatalogTemplateRequest_unstable | CustomProviderCreateRequest_unstable | CustomProviderReadRequest_unstable | CustomProviderUpdateRequest_unstable | CustomProviderDeleteRequest_unstable | RefreshProviderInventoryRequest_unstable | ProviderConfigReadRequest_unstable | ProviderConfigStatusRequest_unstable | ProviderConfigSaveRequest_unstable | ProviderConfigDeleteRequest_unstable | ProviderConfigAuthenticateRequest_unstable | ProviderSecretsListRequest_unstable | ProviderSecretDeleteRequest_unstable | CanonicalModelInfoRequest_unstable | PreferencesReadRequest_unstable | PreferencesSaveRequest_unstable | PreferencesRemoveRequest_unstable | ConfigReadRequest_unstable | ConfigUpsertRequest_unstable | ConfigRemoveRequest_unstable | ConfigReadAllRequest_unstable | DefaultsReadRequest_unstable | DefaultsSaveRequest_unstable | DefaultsClearRequest_unstable | OnboardingImportScanRequest_unstable | OnboardingImportApplyRequest_unstable | ExportSessionRequest_unstable | ImportSessionRequest_unstable | ShareSessionNostrRequest_unstable | EncodeRecipeRequest_unstable | DecodeRecipeRequest_unstable | ScanRecipeRequest_unstable | ListRecipesRequest_unstable | DeleteRecipeRequest_unstable | ScheduleRecipeRequest_unstable | SetRecipeSlashCommandRequest_unstable | SaveRecipeRequest_unstable | ParseRecipeRequest_unstable | RecipeToYamlRequest_unstable | ListSchedulesRequest_unstable | ListScheduleSessionsRequest_unstable | CreateScheduleRequest_unstable | DeleteScheduleRequest_unstable | PauseScheduleRequest_unstable | UnpauseScheduleRequest_unstable | UpdateScheduleRequest_unstable | RunScheduleNowRequest_unstable | KillRunningJobRequest_unstable | InspectRunningJobRequest_unstable | GetSessionInfoRequest_unstable | TruncateSessionConversationRequest_unstable | UpdateSessionProjectRequest_unstable | RenameSessionRequest_unstable | ArchiveSessionRequest_unstable | UnarchiveSessionRequest_unstable | CreateSourceRequest_unstable | ListSourcesRequest_unstable | ListAgentMentionsRequest_unstable | ListSlashCommandsRequest_unstable | UpdateSourceRequest_unstable | DeleteSourceRequest_unstable | ExportSourceRequest_unstable | ImportSourcesRequest_unstable | DictationTranscribeRequest_unstable | DictationConfigRequest_unstable | DictationSecretSaveRequest_unstable | DictationSecretDeleteRequest_unstable | DictationModelsListRequest_unstable | DictationModelDownloadRequest_unstable | DictationModelDownloadProgressRequest_unstable | DictationModelCancelRequest_unstable | DictationModelDeleteRequest_unstable | DictationModelSelectRequest_unstable | LocalInferenceModelsListRequest_unstable | LocalInferenceModelDownloadRequest_unstable | LocalInferenceModelDownloadProgressRequest_unstable | LocalInferenceModelDownloadCancelRequest_unstable | LocalInferenceModelDeleteRequest_unstable | LocalInferenceModelEvictRequest_unstable | LocalInferenceModelSettingsReadRequest_unstable | LocalInferenceModelSettingsUpdateRequest_unstable | LocalInferenceHuggingFaceSearchRequest_unstable | LocalInferenceHuggingFaceRepoVariantsRequest_unstable | LocalInferenceBuiltinChatTemplatesListRequest_unstable | {
[key: string]: unknown;
} | null;
};

6
ui/sdk/src/generated/zod.gen.ts
View file

@ -2505,6 +2505,7 @@ export const zLocalInferenceModelDto = z.object({
sizeBytes: z.number().int().gte(0),
status: zLocalInferenceModelDownloadStatusDto,
recommended: z.boolean(),
isLoaded: z.boolean(),
settings: zLocalInferenceModelSettingsDto,
visionCapable: z.boolean(),
mmprojStatus: z.union([
@ -2573,6 +2574,10 @@ export const zLocalInferenceModelDeleteRequest_unstable = z.object({
modelId: z.string()
});
export const zLocalInferenceModelEvictRequest_unstable = z.object({
modelId: z.string()
});
export const zLocalInferenceModelSettingsReadRequest_unstable = z.object({
modelId: z.string()
});
@ -2914,6 +2919,7 @@ export const zExtRequest = z.object({
zLocalInferenceModelDownloadProgressRequest_unstable,
zLocalInferenceModelDownloadCancelRequest_unstable,
zLocalInferenceModelDeleteRequest_unstable,
zLocalInferenceModelEvictRequest_unstable,
zLocalInferenceModelSettingsReadRequest_unstable,
zLocalInferenceModelSettingsUpdateRequest_unstable,
zLocalInferenceHuggingFaceSearchRequest_unstable,