diff --git a/crates/goose-cli/src/commands/configure.rs b/crates/goose-cli/src/commands/configure.rs index ed0c20a29..092332639 100644 --- a/crates/goose-cli/src/commands/configure.rs +++ b/crates/goose-cli/src/commands/configure.rs @@ -2215,10 +2215,11 @@ fn add_provider() -> anyhow::Result<()> { }) .interact()?; - let models: Vec = models_input + let models: Vec = models_input .split(',') - .map(|s| s.trim().to_string()) - .filter(|s| !s.is_empty()) + .map(str::trim) + .filter(|name| !name.is_empty()) + .map(goose_providers::base::ModelInfo::new) .collect(); let supports_streaming = cliclack::confirm("Does this provider support streaming responses?") diff --git a/crates/goose-cli/src/commands/term.rs b/crates/goose-cli/src/commands/term.rs index 41a39ee22..9246042a3 100644 --- a/crates/goose-cli/src/commands/term.rs +++ b/crates/goose-cli/src/commands/term.rs @@ -342,9 +342,11 @@ pub async fn handle_term_info() -> Result<()> { .unwrap_or(0) as usize; let config = goose::config::Config::global(); - let model_name = config - .get_goose_model() - .ok() + let model_name = session + .as_ref() + .and_then(|session| session.model_config.as_ref()) + .map(|model| model.model_name.clone()) + .or_else(|| config.get_goose_model().ok()) .map(|name| { let short = name.rsplit('/').next().unwrap_or(&name); if let Some(stripped) = short.strip_prefix("goose-") { @@ -355,16 +357,14 @@ pub async fn handle_term_info() -> Result<()> { }) .unwrap_or_else(|| "?".to_string()); - let context_limit = config - .get_goose_model() - .ok() - .and_then(|model_name| { - config.get_goose_provider().ok().and_then(|provider_name| { - goose::model_config::model_config_from_user_config(&provider_name, &model_name).ok() - }) + let context_limit = session + .as_ref() + .and_then(|session| { + let provider_name = session.provider_name.as_deref()?; + let model = session.model_config.as_ref()?; + goose::context_limit::get_local_context_limit(provider_name, &model.model_name).ok() }) - .map(|mc| mc.context_limit()) - .unwrap_or(128_000); + .unwrap_or(goose_providers::model::DEFAULT_CONTEXT_LIMIT); let percentage = if context_limit > 0 { ((total_tokens as f64 / context_limit as f64) * 100.0).round() as usize diff --git a/crates/goose-cli/src/session/mod.rs b/crates/goose-cli/src/session/mod.rs index 2b8e73a85..bccffba0c 100644 --- a/crates/goose-cli/src/session/mod.rs +++ b/crates/goose-cli/src/session/mod.rs @@ -13,7 +13,6 @@ use crate::session::task_execution_display::{ format_task_execution_notification, TASK_EXECUTION_NOTIFICATION_TYPE, }; use goose::conversation::{fix_conversation, merge_consecutive_messages_for_request, Conversation}; -use std::env; use std::io::Write; use std::str::FromStr; use tokio::signal::ctrl_c; @@ -21,6 +20,7 @@ use tokio_util::task::AbortOnDropHandle; pub use builder::{build_session, ExtensionFailure, SessionBuilderConfig}; use console::Color; + use goose::agents::platform_extensions::developer::shell::{ parse_shell_output_notification, ShellOutputNotificationParams, ShellOutputStream, }; @@ -65,7 +65,6 @@ use tokio; use tokio_util::sync::CancellationToken; use tracing::warn; -const GOOSE_PLANNER_CONTEXT_LIMIT: &str = "GOOSE_PLANNER_CONTEXT_LIMIT"; const SHELL_STATUS_FALLBACK_WIDTH: usize = 120; const SHELL_STATUS_MAX_LINES: usize = 3; const SHELL_STATUS_RESERVED_WIDTH: usize = 2; @@ -1079,26 +1078,11 @@ impl CliSession { return Ok(()); } - if let Some(model_info) = target_entry - .metadata() - .known_models - .iter() - .find(|m| m.name == target_model_name) - { - if model_info.context_limit < current_model_config.context_limit.unwrap_or(0) { - eprintln!( - "{}", - console::style(format!( - "Warning: '{}' has a smaller context window ({} tokens) than the current session ({} tokens). \ - You may need to use /compact.", - target_model_name, - model_info.context_limit, - current_model_config.context_limit.unwrap_or(0) - )) - .yellow() - ); - } - } + let current_context_limit = goose::context_limit::get_context_limit( + provider.as_ref(), + ¤t_model_config.model_name, + ) + .await?; let extensions = self.agent.get_extension_configs().await; let new_provider = match goose::providers::create(target_provider_name, extensions).await { @@ -1122,6 +1106,22 @@ impl CliSession { return Ok(()); } + let new_context_limit = goose::context_limit::get_context_limit( + new_provider.as_ref(), + &new_model_config.model_name, + ) + .await?; + if new_context_limit < current_context_limit { + eprintln!( + "{}", + console::style(format!( + "Warning: '{}' has a smaller context window ({} tokens) than the current session ({} tokens). You may need to use /compact.", + target_model_name, new_context_limit, current_context_limit + )) + .yellow() + ); + } + self.agent .update_provider(new_provider, new_model_config, &self.session_id) .await?; @@ -2092,10 +2092,9 @@ impl CliSession { .agent .model_config_for_session(&self.session_id) .await?; - let context_limit = provider - .get_context_limit(&model_config) - .await - .unwrap_or_else(|_| model_config.context_limit()); + let context_limit = + goose::context_limit::get_context_limit(provider.as_ref(), &model_config.model_name) + .await?; let config = Config::global(); let show_cost = config @@ -2859,19 +2858,8 @@ async fn get_reasoner( .expect("No model configured. Run 'goose configure' first") }; - let planner_context_limit = match env::var(GOOSE_PLANNER_CONTEXT_LIMIT) - .ok() - .map(|v| v.parse::()) - { - Some(Ok(n)) if n >= 4096 => Some(n), - Some(Ok(_)) => anyhow::bail!("{} must be at least 4096", GOOSE_PLANNER_CONTEXT_LIMIT), - Some(Err(e)) => anyhow::bail!("{}: {}", GOOSE_PLANNER_CONTEXT_LIMIT, e), - None => None, - }; - let model_config = - goose::model_config::model_config_from_user_config(&provider, model.as_str())? - .with_context_limit(planner_context_limit); + goose::model_config::model_config_from_user_config(&provider, model.as_str())?; let extensions = goose::config::extensions::get_enabled_extensions_with_config(config); let reasoner = create(&provider, extensions).await?; diff --git a/crates/goose-context-management/src/provider.rs b/crates/goose-context-management/src/provider.rs index 3744a99fc..78335d96d 100644 --- a/crates/goose-context-management/src/provider.rs +++ b/crates/goose-context-management/src/provider.rs @@ -96,8 +96,8 @@ impl Provider for CompactingProvider { } } - async fn get_context_limit(&self, model_config: &ModelConfig) -> Result { - self.inner.get_context_limit(model_config).await + async fn get_context_limit(&self, model: &str, override_limit: Option) -> usize { + self.inner.get_context_limit(model, override_limit).await } fn manages_own_context(&self) -> bool { diff --git a/crates/goose-local-inference/src/lib.rs b/crates/goose-local-inference/src/lib.rs index 99aad9799..bb231575c 100644 --- a/crates/goose-local-inference/src/lib.rs +++ b/crates/goose-local-inference/src/lib.rs @@ -261,6 +261,12 @@ fn resolve_model_local_path(model_id: &str) -> Option { .map(|entry| entry.local_path.clone()) } +pub fn local_context_limit(model_id: &str) -> Option { + resolve_model_path(model_id) + .map(|resolved| resolved.context_limit) + .filter(|limit| *limit > 0) +} + /// Resolve model path, context limit, settings, and mmproj path for a model ID from the registry. fn resolve_model_path(model_id: &str) -> Option { use crate::local_model_registry::{default_settings_for_model, get_registry}; @@ -619,6 +625,14 @@ impl Provider for LocalInferenceProvider { &self.name } + async fn get_context_limit(&self, model: &str, override_limit: Option) -> usize { + goose_provider_types::context_limit::ContextLimitResolver::new(&self.name) + .resolve(model, override_limit, || async { + Ok(resolve_model_path(model).map(|resolved| resolved.context_limit)) + }) + .await + } + async fn fetch_supported_models(&self) -> Result, ProviderError> { use crate::local_model_registry::get_registry; diff --git a/crates/goose-provider-types/src/base.rs b/crates/goose-provider-types/src/base.rs index 55321c495..bf08b0c92 100644 --- a/crates/goose-provider-types/src/base.rs +++ b/crates/goose-provider-types/src/base.rs @@ -273,7 +273,7 @@ pub struct ModelInfo { #[serde(default, skip_serializing_if = "Option::is_none")] pub resolved_model: Option, /// The maximum context length this model supports - pub context_limit: usize, + pub context_limit: Option, /// Cost per token for input in USD (optional) pub input_token_cost: Option, /// Cost per token for output in USD (optional) @@ -293,12 +293,11 @@ pub struct ModelInfo { } impl ModelInfo { - /// Create a new ModelInfo with just name and context limit - pub fn new(name: impl Into, context_limit: usize) -> Self { + pub fn new(name: impl Into) -> Self { Self { name: name.into(), resolved_model: None, - context_limit, + context_limit: None, input_token_cost: None, output_token_cost: None, currency: None, @@ -309,6 +308,16 @@ impl ModelInfo { } } + pub fn with_context_limit(mut self, context_limit: usize) -> Self { + self.context_limit = Some(context_limit); + self + } + + pub fn with_optional_context_limit(mut self, context_limit: Option) -> Self { + self.context_limit = context_limit; + self + } + /// Create a new ModelInfo with cost information (per token) pub fn with_cost( name: impl Into, @@ -319,7 +328,7 @@ impl ModelInfo { Self { name: name.into(), resolved_model: None, - context_limit, + context_limit: Some(context_limit), input_token_cost: Some(input_cost), output_token_cost: Some(output_cost), currency: Some("$".to_string()), @@ -364,9 +373,7 @@ pub fn model_info_for_provider_model(provider_name: &str, model_name: &str) -> M ModelInfo { name: model_name.to_string(), resolved_model: None, - context_limit: ModelConfig::new(model_name) - .with_canonical_limits(provider_name) - .context_limit(), + context_limit: canonical.as_ref().map(|model| model.limit.context), input_token_cost: None, output_token_cost: None, currency: None, @@ -509,13 +516,15 @@ pub trait Provider: Send + Sync { collect_stream(stream).await } - /// Resolve the effective context limit for a model config. + /// Resolve the effective context limit for a model. /// - /// Providers may override this to enrich the limit with provider-specific - /// metadata (e.g. cached model info or a value captured from a remote - /// session). The default returns the limit derived from the model config. - async fn get_context_limit(&self, model_config: &ModelConfig) -> Result { - Ok(model_config.context_limit()) + /// `override_limit` is consumer policy and takes precedence over provider + /// configuration and discovery. The method is infallible because providers + /// fall through to canonical metadata and the global default. + async fn get_context_limit(&self, model: &str, override_limit: Option) -> usize { + crate::context_limit::ContextLimitResolver::new(self.get_name()) + .resolve(model, override_limit, || async { Ok(None) }) + .await } fn retry_config(&self) -> RetryConfig { @@ -1043,7 +1052,7 @@ mod tests { let info = ModelInfo { name: "test-model".to_string(), resolved_model: None, - context_limit: 1000, + context_limit: Some(1000), input_token_cost: None, output_token_cost: None, currency: None, @@ -1052,13 +1061,13 @@ mod tests { thinking_preservation_format: None, request_params: None, }; - assert_eq!(info.context_limit, 1000); + assert_eq!(info.context_limit, Some(1000)); // Test equality let info2 = ModelInfo { name: "test-model".to_string(), resolved_model: None, - context_limit: 1000, + context_limit: Some(1000), input_token_cost: None, output_token_cost: None, currency: None, @@ -1073,7 +1082,7 @@ mod tests { let info3 = ModelInfo { name: "test-model".to_string(), resolved_model: None, - context_limit: 2000, + context_limit: Some(2000), input_token_cost: None, output_token_cost: None, currency: None, @@ -1116,7 +1125,7 @@ mod tests { fn test_model_info_with_cost() { let info = ModelInfo::with_cost("gpt-4o", 128000, 0.0000025, 0.00001); assert_eq!(info.name, "gpt-4o"); - assert_eq!(info.context_limit, 128000); + assert_eq!(info.context_limit, Some(128000)); assert_eq!(info.input_token_cost, Some(0.0000025)); assert_eq!(info.output_token_cost, Some(0.00001)); assert_eq!(info.currency, Some("$".to_string())); diff --git a/crates/goose-provider-types/src/canonical/name_builder.rs b/crates/goose-provider-types/src/canonical/name_builder.rs index a4a4fed20..91381e420 100644 --- a/crates/goose-provider-types/src/canonical/name_builder.rs +++ b/crates/goose-provider-types/src/canonical/name_builder.rs @@ -44,7 +44,7 @@ pub(crate) fn is_meta_provider(provider: &str) -> bool { pub fn map_provider_name(provider: &str) -> &str { match provider { // Goose provider names that differ from models.dev names - "xai" => "x-ai", + "xai" | "xai_oauth" => "x-ai", "azure_openai" | "azure_foundry" => "azure", "aws_bedrock" => "amazon-bedrock", "gcp_vertex_ai" => "google-vertex", @@ -338,6 +338,10 @@ mod tests { map_to_canonical_model("openai", "gpt-4o-latest", r), Some("openai/gpt-4o".to_string()) ); + assert_eq!( + map_to_canonical_model("xai_oauth", "grok-4.5", r), + Some("x-ai/grok-4.5".to_string()) + ); assert_eq!( map_to_canonical_model("openai", "gpt-4-turbo-2024-04-09", r), Some("openai/gpt-4-turbo".to_string()) diff --git a/crates/goose-provider-types/src/context_limit.rs b/crates/goose-provider-types/src/context_limit.rs new file mode 100644 index 000000000..3eb1ee92c --- /dev/null +++ b/crates/goose-provider-types/src/context_limit.rs @@ -0,0 +1,195 @@ +use std::collections::HashMap; +use std::future::Future; + +use crate::canonical::maybe_get_canonical_model; +use crate::errors::ProviderError; +use crate::model::DEFAULT_CONTEXT_LIMIT; + +#[derive(Debug, Clone, Default)] +pub struct ContextLimitResolver { + provider_name: String, + configured_limits: HashMap, +} + +impl ContextLimitResolver { + pub fn new(provider_name: impl Into) -> Self { + Self { + provider_name: provider_name.into(), + configured_limits: HashMap::new(), + } + } + + pub fn with_configured_limits( + mut self, + configured_limits: impl IntoIterator, + ) -> Self { + self.configured_limits = configured_limits + .into_iter() + .filter(|(_, context_limit)| *context_limit > 0) + .collect(); + self + } + + fn configured_limit(&self, model: &str) -> Option { + self.configured_limits.get(model).copied().or_else(|| { + let mut matches = self + .configured_limits + .iter() + .filter(|(configured_model, _)| configured_model.eq_ignore_ascii_case(model)); + let (_, limit) = matches.next()?; + matches.next().is_none().then_some(*limit) + }) + } + + pub fn resolve_local(&self, model: &str, override_limit: Option) -> usize { + override_limit + .or_else(|| self.configured_limit(model)) + .or_else(|| { + maybe_get_canonical_model(&self.provider_name, model) + .map(|canonical| canonical.limit.context) + }) + .unwrap_or(DEFAULT_CONTEXT_LIMIT) + } + + pub async fn resolve( + &self, + model: &str, + override_limit: Option, + discover: F, + ) -> usize + where + F: FnOnce() -> Fut, + Fut: Future, ProviderError>>, + { + if let Some(limit) = override_limit { + return limit; + } + + if let Some(limit) = self.configured_limit(model) { + return limit; + } + + match discover().await { + Ok(Some(limit)) if limit > 0 => return limit, + Ok(Some(_) | None) => {} + Err(error) => tracing::warn!( + provider = self.provider_name, + model, + %error, + "Context-limit discovery failed; falling back" + ), + } + + maybe_get_canonical_model(&self.provider_name, model) + .map(|canonical| canonical.limit.context) + .unwrap_or(DEFAULT_CONTEXT_LIMIT) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn applies_precedence() { + let resolver = ContextLimitResolver::new("anthropic") + .with_configured_limits([("claude-sonnet-4-5".to_string(), 64_000)]); + + assert_eq!( + resolver + .resolve("claude-sonnet-4-5", Some(32_000), || async { + Ok(Some(16_000)) + }) + .await, + 32_000 + ); + assert_eq!( + resolver + .resolve("claude-sonnet-4-5", None, || async { Ok(Some(16_000)) }) + .await, + 64_000 + ); + } + + #[test] + fn configured_limits_match_case_insensitively() { + let resolver = ContextLimitResolver::new("unknown-provider") + .with_configured_limits([("MyModel".to_string(), 64_000)]); + + assert_eq!(resolver.resolve_local("mymodel", None), 64_000); + } + + #[test] + fn exact_match_wins_over_case_insensitive_matches() { + let resolver = ContextLimitResolver::new("unknown-provider").with_configured_limits([ + ("MyModel".to_string(), 64_000), + ("mymodel".to_string(), 32_000), + ]); + + assert_eq!(resolver.resolve_local("MyModel", None), 64_000); + assert_eq!(resolver.resolve_local("mymodel", None), 32_000); + assert_eq!( + resolver.resolve_local("MYMODEL", None), + DEFAULT_CONTEXT_LIMIT + ); + } + + #[test] + fn ignores_zero_configured_limits() { + let resolver = ContextLimitResolver::new("unknown-provider") + .with_configured_limits([("configured".to_string(), 0)]); + + assert_eq!( + resolver.resolve_local("configured", None), + DEFAULT_CONTEXT_LIMIT + ); + } + + #[test] + fn local_resolution_skips_discovery() { + let resolver = ContextLimitResolver::new("anthropic") + .with_configured_limits([("configured".to_string(), 64_000)]); + + assert_eq!(resolver.resolve_local("configured", None), 64_000); + assert_eq!(resolver.resolve_local("claude-sonnet-4-5", None), 1_000_000); + assert_eq!( + resolver.resolve_local("unknown", None), + DEFAULT_CONTEXT_LIMIT + ); + } + + #[tokio::test] + async fn ignores_zero_discovered_limits() { + let resolver = ContextLimitResolver::new("unknown-provider"); + + assert_eq!( + resolver + .resolve("unknown-model", None, || async { Ok(Some(0)) }) + .await, + DEFAULT_CONTEXT_LIMIT + ); + } + + #[tokio::test] + async fn uses_discovery_then_canonical_then_default() { + let resolver = ContextLimitResolver::new("anthropic"); + assert_eq!( + resolver + .resolve("runtime-model", None, || async { Ok(Some(24_000)) }) + .await, + 24_000 + ); + assert_eq!( + resolver + .resolve("claude-sonnet-4-5", None, || async { Ok(None) }) + .await, + 1_000_000 + ); + assert_eq!( + resolver + .resolve("unknown-model", None, || async { Ok(None) }) + .await, + DEFAULT_CONTEXT_LIMIT + ); + } +} diff --git a/crates/goose-provider-types/src/lib.rs b/crates/goose-provider-types/src/lib.rs index dd7a87dda..e9b73cb97 100644 --- a/crates/goose-provider-types/src/lib.rs +++ b/crates/goose-provider-types/src/lib.rs @@ -1,6 +1,7 @@ pub mod base; pub mod cache_semantics; pub mod canonical; +pub mod context_limit; pub mod conversation; pub mod errors; pub mod formats; diff --git a/crates/goose-provider-types/src/model.rs b/crates/goose-provider-types/src/model.rs index 453e6dc00..00730a289 100644 --- a/crates/goose-provider-types/src/model.rs +++ b/crates/goose-provider-types/src/model.rs @@ -39,6 +39,7 @@ pub fn is_goose_internal_request_param(key: &str) -> bool { #[derive(Debug, Clone, Serialize)] pub struct ModelConfig { pub model_name: String, + #[serde(skip)] pub context_limit: Option, pub temperature: Option, pub max_tokens: Option, @@ -65,7 +66,8 @@ impl<'de> Deserialize<'de> for ModelConfig { #[derive(Deserialize)] struct RawModelConfig { model_name: String, - context_limit: Option, + #[serde(rename = "context_limit")] + _context_limit: Option, temperature: Option, max_tokens: Option, toolshim: bool, @@ -81,7 +83,7 @@ impl<'de> Deserialize<'de> for ModelConfig { let raw = RawModelConfig::deserialize(deserializer)?; let mut config = Self { model_name: raw.model_name, - context_limit: raw.context_limit, + context_limit: None, temperature: raw.temperature, max_tokens: raw.max_tokens, toolshim: raw.toolshim, @@ -131,9 +133,6 @@ impl ModelConfig { ); if let Some(canonical) = canonical { - if self.context_limit.is_none() { - self.context_limit = Some(canonical.limit.context); - } if self.max_tokens.is_none() { self.max_tokens = canonical .limit @@ -696,25 +695,10 @@ mod tests { ("GOOSE_CONTEXT_LIMIT", None::<&str>), ]); let config = ModelConfig::new("gpt-4o").with_canonical_limits("openai"); - - assert_eq!(config.context_limit, Some(128_000)); assert_eq!(config.max_tokens, Some(16_384)); assert_eq!(config.reasoning, Some(false)); } - #[test] - fn does_not_override_existing_context_limit() { - let _guard = env_lock::lock_env([ - ("GOOSE_MAX_TOKENS", None::<&str>), - ("GOOSE_CONTEXT_LIMIT", None::<&str>), - ]); - let mut config = ModelConfig::new("gpt-4o"); - config.context_limit = Some(64_000); - let config = config.with_canonical_limits("openai"); - - assert_eq!(config.context_limit, Some(64_000)); - } - #[test] fn does_not_override_existing_max_tokens() { let _guard = env_lock::lock_env([ @@ -735,8 +719,6 @@ mod tests { ("GOOSE_CONTEXT_LIMIT", None::<&str>), ]); let config = ModelConfig::new("moonshotai/kimi-k2.6").with_canonical_limits("nvidia"); - - assert_eq!(config.context_limit, Some(262_144)); assert_eq!(config.max_tokens, None); assert_eq!(config.max_output_tokens(), 4_096); } @@ -749,8 +731,6 @@ mod tests { ]); let config = ModelConfig::new("global.anthropic.claude-sonnet-5") .with_canonical_limits("aws_bedrock"); - - assert_eq!(config.context_limit, Some(1_000_000)); assert_eq!(config.max_tokens, Some(128_000)); assert_eq!(config.reasoning, Some(true)); } @@ -775,18 +755,8 @@ mod tests { ("GOOSE_CONTEXT_LIMIT", None::<&str>), ]); - // "databricks-gpt-5.4-high" should resolve via "databricks-gpt-5.4" - let config = - ModelConfig::new("databricks-gpt-5.4-high").with_canonical_limits("databricks"); - assert_eq!(config.context_limit, Some(1_050_000)); - - // "gpt-5.4-xhigh" should resolve via "gpt-5.4" - let config = ModelConfig::new("gpt-5.4-xhigh").with_canonical_limits("openai"); - assert_eq!(config.context_limit, Some(1_050_000)); - // "gpt-5.6-sol-xhigh" should resolve via "gpt-5.6-sol" let config = ModelConfig::new("gpt-5.6-sol-xhigh").with_canonical_limits("openai"); - assert_eq!(config.context_limit, Some(1_050_000)); assert_eq!(config.max_tokens, Some(128_000)); assert_eq!(config.reasoning, Some(true)); let canonical = crate::canonical::maybe_get_canonical_model("openai", "gpt-5.6-sol") @@ -794,13 +764,8 @@ mod tests { assert_eq!(canonical.temperature, Some(false)); let config = ModelConfig::new("gpt-5.6-sol").with_canonical_limits("chatgpt_codex"); - assert_eq!(config.context_limit, Some(1_050_000)); assert_eq!(config.max_tokens, Some(128_000)); assert_eq!(config.reasoning, Some(true)); - - // "gpt-5.4-nano-low" should resolve via "gpt-5.4-nano" - let config = ModelConfig::new("gpt-5.4-nano-low").with_canonical_limits("openai"); - assert_eq!(config.context_limit, Some(400_000)); } #[test] diff --git a/crates/goose-providers/src/anthropic.rs b/crates/goose-providers/src/anthropic.rs index 9853d7c07..38dcdb05c 100644 --- a/crates/goose-providers/src/anthropic.rs +++ b/crates/goose-providers/src/anthropic.rs @@ -65,7 +65,7 @@ pub struct AnthropicProvider { api_client: ApiClient, supports_streaming: bool, name: String, - custom_models: Option>, + custom_models: Option>, dynamic_models: Option, skip_canonical_filtering: bool, #[serde(skip)] @@ -82,7 +82,7 @@ pub struct AnthropicProviderBuilder { api_client: ApiClient, supports_streaming: bool, name: String, - custom_models: Option>, + custom_models: Option>, dynamic_models: Option, skip_canonical_filtering: bool, format_options: AnthropicFormatOptions, @@ -129,7 +129,7 @@ impl AnthropicProviderBuilder { self } - pub fn custom_models(mut self, custom_models: Option>) -> Self { + pub fn custom_models(mut self, custom_models: Option>) -> Self { self.custom_models = custom_models; self } @@ -284,7 +284,7 @@ impl ProviderDescriptor for AnthropicProvider { fn metadata() -> ProviderMetadata { let models: Vec = ANTHROPIC_KNOWN_MODELS .iter() - .map(|&model_name| ModelInfo::new(model_name, 200_000)) + .map(|&model_name| ModelInfo::new(model_name).with_context_limit(200_000)) .collect(); ProviderMetadata::with_models( @@ -331,10 +331,25 @@ impl Provider for AnthropicProvider { self.skip_canonical_filtering } + async fn get_context_limit(&self, model: &str, override_limit: Option) -> usize { + let configured_limits = self + .custom_models + .iter() + .flatten() + .filter_map(|model| model.context_limit.map(|limit| (model.name.clone(), limit))); + crate::context_limit::ContextLimitResolver::new(&self.name) + .with_configured_limits(configured_limits) + .resolve(model, override_limit, || async { Ok(None) }) + .await + } + async fn fetch_supported_models(&self) -> Result, ProviderError> { if let Some(custom_models) = &self.custom_models { if self.dynamic_models == Some(false) { - return Ok(custom_models.clone()); + return Ok(custom_models + .iter() + .map(|model| model.name.clone()) + .collect()); } match self.fetch_models_from_api().await { Ok(models) => return Ok(models), @@ -344,7 +359,10 @@ impl Provider for AnthropicProvider { self.name, e ); - return Ok(custom_models.clone()); + return Ok(custom_models + .iter() + .map(|model| model.name.clone()) + .collect()); } Err(e) => return Err(e), } @@ -389,13 +407,7 @@ pub fn from_declarative_config( key_resolver: impl KeyResolver, ) -> Result { let custom_models = if !config.models.is_empty() { - Some( - config - .models - .iter() - .map(|m| m.name.clone()) - .collect::>(), - ) + Some(config.models.clone()) } else { None }; @@ -530,7 +542,7 @@ mod tests { .unwrap(), supports_streaming: true, name: "test-provider".to_string(), - custom_models: Some(custom_models), + custom_models: Some(custom_models.into_iter().map(ModelInfo::new).collect()), dynamic_models: Some(true), skip_canonical_filtering: false, format_options: AnthropicFormatOptions::default(), diff --git a/crates/goose-providers/src/azure_foundry.rs b/crates/goose-providers/src/azure_foundry.rs index 536d62915..6325adb5c 100644 --- a/crates/goose-providers/src/azure_foundry.rs +++ b/crates/goose-providers/src/azure_foundry.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; use std::sync::Mutex; +use std::time::{Duration, Instant}; use anyhow::Result; use async_trait::async_trait; @@ -12,7 +13,7 @@ use crate::base::{ }; use crate::conversation::message::Message; use crate::errors::ProviderError; -use crate::formats::openai::is_openai_responses_model; +use crate::formats::openai::{extract_reasoning_effort, is_openai_responses_model}; use crate::model::ModelConfig; use crate::openai::{OpenAiProvider, OpenAiProviderBuilder}; use crate::openai_compatible::{handle_response_openai_compat, OpenAiCompatibleProvider}; @@ -22,6 +23,9 @@ pub const AZURE_FOUNDRY_DEFAULT_MODEL: &str = "Phi-4"; pub const AZURE_FOUNDRY_DOC_URL: &str = "https://learn.microsoft.com/azure/ai-foundry/foundry-models/how-to/inference"; +const DEPLOYMENT_METADATA_TIMEOUT_SECS: u64 = 5; +const DEPLOYMENT_METADATA_TTL_SECS: u64 = 60; + pub const AZURE_FOUNDRY_KNOWN_MODELS: &[&str] = &[ "Phi-4", "Phi-4-mini", @@ -80,6 +84,27 @@ struct DeploymentMetadata { model_name: String, } +#[derive(Default)] +struct DeploymentCache { + deployments: HashMap, + fetched_at: Option, + fetch_failed: bool, +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum DeploymentMetadataLookup { + ContextDiscovery, + InferenceRouting, +} + +impl DeploymentCache { + fn applies_to(&self, lookup: DeploymentMetadataLookup) -> bool { + self.fetched_at.is_some_and(|fetched_at| { + fetched_at.elapsed() < Duration::from_secs(DEPLOYMENT_METADATA_TTL_SECS) + }) && (lookup == DeploymentMetadataLookup::ContextDiscovery || !self.fetch_failed) + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum InferenceRoute { MaasChatCompletions, @@ -134,7 +159,7 @@ pub struct AzureFoundryProvider { endpoint: String, api_version: Option, maas_model: Option, - deployments: Mutex>, + deployments: Mutex, } impl ProviderDescriptor for AzureFoundryProvider { @@ -249,7 +274,7 @@ impl AzureFoundryProvider { endpoint, api_version, maas_model, - deployments: Mutex::new(HashMap::new()), + deployments: Mutex::new(DeploymentCache::default()), }) } @@ -309,23 +334,41 @@ impl AzureFoundryProvider { Ok((models, deployments)) } - async fn deployment_for(&self, deployment_name: &str) -> Option { - if let Some(deployment) = self - .deployments - .lock() - .expect("Azure Foundry deployment cache poisoned") - .get(deployment_name) - .cloned() + async fn deployment_for( + &self, + deployment_name: &str, + lookup: DeploymentMetadataLookup, + ) -> Option { { - return Some(deployment); + let cache = self + .deployments + .lock() + .expect("Azure Foundry deployment cache poisoned"); + if cache.applies_to(lookup) { + return cache.deployments.get(deployment_name).cloned(); + } } - let (_, deployments) = self.fetch_deployments().await.ok()?; + let fetch_result = tokio::time::timeout( + Duration::from_secs(DEPLOYMENT_METADATA_TIMEOUT_SECS), + self.fetch_deployments(), + ) + .await + .ok() + .and_then(Result::ok); + let fetch_failed = fetch_result.is_none(); + let deployments = fetch_result + .map(|(_, deployments)| deployments) + .unwrap_or_default(); let deployment = deployments.get(deployment_name).cloned(); *self .deployments .lock() - .expect("Azure Foundry deployment cache poisoned") = deployments; + .expect("Azure Foundry deployment cache poisoned") = DeploymentCache { + deployments, + fetched_at: Some(Instant::now()), + fetch_failed, + }; deployment } } @@ -348,14 +391,24 @@ fn model_info_for_deployment(deployment_name: &str, model_name: &str) -> ModelIn "azure_foundry", &model_name.to_ascii_lowercase(), ) + }) + .or_else(|| { + let (base_model, effort) = extract_reasoning_effort(model_name); + effort.and_then(|_| { + crate::canonical::maybe_get_canonical_model("azure_foundry", &base_model).or_else( + || { + crate::canonical::maybe_get_canonical_model( + "azure_foundry", + &base_model.to_ascii_lowercase(), + ) + }, + ) + }) }); ModelInfo { name: deployment_name.to_string(), resolved_model: Some(model_name.to_string()), - context_limit: canonical - .as_ref() - .map(|model| model.limit.context) - .unwrap_or_else(|| ModelConfig::new(model_name).context_limit()), + context_limit: canonical.as_ref().map(|model| model.limit.context), input_token_cost: None, output_token_cost: None, currency: None, @@ -416,7 +469,11 @@ impl Provider for AzureFoundryProvider { *self .deployments .lock() - .expect("Azure Foundry deployment cache poisoned") = deployments; + .expect("Azure Foundry deployment cache poisoned") = DeploymentCache { + deployments, + fetched_at: Some(Instant::now()), + fetch_failed: false, + }; Ok(models) } @@ -442,7 +499,11 @@ impl Provider for AzureFoundryProvider { *self .deployments .lock() - .expect("Azure Foundry deployment cache poisoned") = deployments; + .expect("Azure Foundry deployment cache poisoned") = DeploymentCache { + deployments, + fetched_at: Some(Instant::now()), + fetch_failed: false, + }; Ok(model_info) } @@ -450,7 +511,7 @@ impl Provider for AzureFoundryProvider { let resolved_model = if let Some(model) = &self.maas_model { model.clone() } else if is_project_endpoint(&self.endpoint) { - self.deployment_for(model_name) + self.deployment_for(model_name, DeploymentMetadataLookup::ContextDiscovery) .await .map(|deployment| deployment.model_name) .unwrap_or_else(|| model_name.to_string()) @@ -460,14 +521,14 @@ impl Provider for AzureFoundryProvider { Ok(model_info_for_deployment(model_name, &resolved_model)) } - async fn get_context_limit(&self, model_config: &ModelConfig) -> Result { - if let Some(context_limit) = model_config.context_limit { - return Ok(context_limit); - } - Ok(self - .fetch_model_info(&model_config.model_name) - .await? - .context_limit) + async fn get_context_limit(&self, model: &str, override_limit: Option) -> usize { + goose_provider_types::context_limit::ContextLimitResolver::new(self.get_name()) + .resolve(model, override_limit, || async { + self.fetch_model_info(model) + .await + .map(|info| info.context_limit) + }) + .await } async fn stream( @@ -488,7 +549,8 @@ impl Provider for AzureFoundryProvider { .clone() .unwrap_or_else(|| model_config.model_name.clone()); let deployment = if is_project_endpoint(&self.endpoint) { - self.deployment_for(&wire_model).await + self.deployment_for(&wire_model, DeploymentMetadataLookup::InferenceRouting) + .await } else { None }; @@ -610,6 +672,30 @@ mod tests { ) } + #[test] + fn routing_retries_cached_deployment_failures() { + let cache = DeploymentCache { + fetched_at: Some(Instant::now()), + fetch_failed: true, + ..Default::default() + }; + + assert!(cache.applies_to(DeploymentMetadataLookup::ContextDiscovery)); + assert!(!cache.applies_to(DeploymentMetadataLookup::InferenceRouting)); + } + + #[test] + fn routing_reuses_successful_empty_deployment_inventory() { + let cache = DeploymentCache { + fetched_at: Some(Instant::now()), + fetch_failed: false, + ..Default::default() + }; + + assert!(cache.applies_to(DeploymentMetadataLookup::ContextDiscovery)); + assert!(cache.applies_to(DeploymentMetadataLookup::InferenceRouting)); + } + #[test] fn routing_matrix_uses_endpoint_publisher_and_underlying_model() { use InferenceRoute::*; @@ -859,7 +945,7 @@ mod tests { let info = model_info_for_deployment("production-chat", "gpt-5"); assert_eq!(info.name, "production-chat"); assert_eq!(info.resolved_model.as_deref(), Some("gpt-5")); - assert_eq!(info.context_limit, 400_000); + assert_eq!(info.context_limit, Some(400_000)); assert_eq!(info.input_token_cost, None); assert_eq!(info.output_token_cost, None); } @@ -868,7 +954,7 @@ mod tests { fn gpt_5_6_sol_uses_its_full_context_window() { let info = model_info_for_deployment("gpt-5.6-sol", "gpt-5.6-sol"); - assert_eq!(info.context_limit, 1_050_000); + assert_eq!(info.context_limit, Some(1_050_000)); assert!(info.reasoning); } @@ -936,10 +1022,7 @@ mod tests { let provider = project_provider(&server); assert_eq!( - provider - .get_context_limit(&ModelConfig::new("production-chat")) - .await - .unwrap(), + provider.get_context_limit("production-chat", None).await, 400_000 ); } @@ -963,16 +1046,24 @@ mod tests { let provider = project_provider(&server); let config = raw_model_config("gpt-5-high"); - assert_eq!(provider.get_context_limit(&config).await.unwrap(), 128_000); + assert_eq!( + provider.get_context_limit(&config.model_name, None).await, + 128_000 + ); } #[tokio::test] - async fn explicit_context_limit_overrides_deployment_metadata() { + async fn caller_override_precedes_deployment_metadata() { let server = MockServer::start().await; let provider = project_provider(&server); - let config = raw_model_config("gpt-5-high").with_context_limit(Some(64_000)); + let config = raw_model_config("gpt-5-high"); - assert_eq!(provider.get_context_limit(&config).await.unwrap(), 64_000); + assert_eq!( + provider + .get_context_limit(&config.model_name, Some(64_000)) + .await, + 64_000 + ); assert!(server.received_requests().await.unwrap().is_empty()); } @@ -1048,6 +1139,7 @@ mod tests { Mock::given(method("GET")) .and(path("/api/projects/test/deployments")) .respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": []}))) + .expect(1) .mount(&server) .await; Mock::given(method("POST")) @@ -1062,7 +1154,12 @@ mod tests { .mount(&server) .await; - project_provider(&server) + let provider = project_provider(&server); + assert_eq!( + provider.get_context_limit("gpt-5-high", None).await, + 400_000 + ); + provider .complete(&raw_model_config("gpt-5-high"), "system", &[], &[]) .await .unwrap(); diff --git a/crates/goose-providers/src/databricks.rs b/crates/goose-providers/src/databricks.rs index eb8a4380a..99d23c064 100644 --- a/crates/goose-providers/src/databricks.rs +++ b/crates/goose-providers/src/databricks.rs @@ -53,10 +53,24 @@ struct DatabricksUpstreamModel { #[derive(Debug, Clone)] struct CachedDatabricksEndpointInfo { - info: DatabricksEndpointInfo, + info: Option, fetched_at: Instant, } +#[derive(Clone, Copy, PartialEq, Eq)] +enum EndpointMetadataLookup { + ContextDiscovery, + InferenceRouting, +} + +impl CachedDatabricksEndpointInfo { + fn applies_to(&self, lookup: EndpointMetadataLookup) -> bool { + self.fetched_at.elapsed() < Duration::from_secs(DATABRICKS_ENDPOINT_METADATA_TTL_SECS) + && (lookup == EndpointMetadataLookup::ContextDiscovery || self.info.is_some()) + } +} + +const DATABRICKS_ENDPOINT_METADATA_TIMEOUT_SECS: u64 = 5; const DATABRICKS_ENDPOINT_METADATA_TTL_SECS: u64 = 60; static DATABRICKS_ENDPOINT_INFO_CACHE: LazyLock< Mutex>, @@ -424,6 +438,7 @@ impl DatabricksProvider { async fn resolve_endpoint_info_cached( &self, endpoint_name: &str, + lookup: EndpointMetadataLookup, ) -> Result { let cache_key = format!("{}:{}", self.host, endpoint_name); let cached = DATABRICKS_ENDPOINT_INFO_CACHE @@ -433,14 +448,22 @@ impl DatabricksProvider { .cloned(); if let Some(cached) = cached { - if cached.fetched_at.elapsed() - < Duration::from_secs(DATABRICKS_ENDPOINT_METADATA_TTL_SECS) - { - return Ok(cached.info); + if cached.applies_to(lookup) { + return cached.info.ok_or_else(|| { + ProviderError::RequestFailed( + "Databricks endpoint metadata is unavailable".to_string(), + ) + }); } } - let info = self.resolve_endpoint_info(endpoint_name).await?; + let info = tokio::time::timeout( + Duration::from_secs(DATABRICKS_ENDPOINT_METADATA_TIMEOUT_SECS), + self.resolve_endpoint_info(endpoint_name), + ) + .await + .ok() + .and_then(Result::ok); DATABRICKS_ENDPOINT_INFO_CACHE.lock().unwrap().insert( cache_key, CachedDatabricksEndpointInfo { @@ -448,14 +471,16 @@ impl DatabricksProvider { fetched_at: Instant::now(), }, ); - Ok(info) + info.ok_or_else(|| { + ProviderError::RequestFailed("Databricks endpoint metadata is unavailable".to_string()) + }) } fn model_info_from_endpoint(info: DatabricksEndpointInfo) -> ModelInfo { let context_model = info.upstream_model_name.as_deref().unwrap_or(&info.name); - let context_limit = ModelConfig::new(context_model) - .with_canonical_limits(DATABRICKS_PROVIDER_NAME) - .context_limit(); + let context_limit = + crate::canonical::maybe_get_canonical_model(DATABRICKS_PROVIDER_NAME, context_model) + .map(|model| model.limit.context); let reasoning = info .reasoning .unwrap_or_else(|| ModelConfig::new(context_model).is_reasoning_model()); @@ -521,6 +546,16 @@ impl Provider for DatabricksProvider { self.retry_config.clone() } + async fn get_context_limit(&self, model: &str, override_limit: Option) -> usize { + crate::context_limit::ContextLimitResolver::new(self.get_name()) + .resolve(model, override_limit, || async { + self.fetch_model_info(model) + .await + .map(|info| info.context_limit) + }) + .await + } + async fn refresh_credentials(&self) -> Result<(), ProviderError> { if let Some(refresh_hook) = &self.refresh_hook { refresh_hook(); @@ -543,7 +578,10 @@ impl Provider for DatabricksProvider { .and_then(|provider| provider()) .unwrap_or_default(); let (endpoint_name, _) = extract_reasoning_effort(&model_config.model_name); - let endpoint_info = self.resolve_endpoint_info_cached(&endpoint_name).await.ok(); + let endpoint_info = self + .resolve_endpoint_info_cached(&endpoint_name, EndpointMetadataLookup::InferenceRouting) + .await + .ok(); let effective_model_name = endpoint_info .as_ref() .and_then(|info| info.upstream_model_name.as_deref()) @@ -780,7 +818,9 @@ impl Provider for DatabricksProvider { async fn fetch_model_info(&self, model_name: &str) -> Result { let (endpoint_name, _) = extract_reasoning_effort(model_name); - let endpoint_info = self.resolve_endpoint_info_cached(&endpoint_name).await?; + let endpoint_info = self + .resolve_endpoint_info_cached(&endpoint_name, EndpointMetadataLookup::ContextDiscovery) + .await?; Ok(Self::model_info_from_endpoint(endpoint_info)) } @@ -796,6 +836,34 @@ impl Provider for DatabricksProvider { mod tests { use super::*; + #[test] + fn routing_retries_cached_metadata_failures() { + let cached = CachedDatabricksEndpointInfo { + info: None, + fetched_at: Instant::now(), + }; + + assert!(cached.applies_to(EndpointMetadataLookup::ContextDiscovery)); + assert!(!cached.applies_to(EndpointMetadataLookup::InferenceRouting)); + } + + #[test] + fn routing_reuses_cached_metadata_successes() { + let cached = CachedDatabricksEndpointInfo { + info: Some(DatabricksEndpointInfo { + name: "production-chat".to_string(), + upstream_model_name: Some("gpt-5".to_string()), + upstream_model_provider: Some("openai".to_string()), + reasoning: Some(true), + supports_responses_api: true, + }), + fetched_at: Instant::now(), + }; + + assert!(cached.applies_to(EndpointMetadataLookup::ContextDiscovery)); + assert!(cached.applies_to(EndpointMetadataLookup::InferenceRouting)); + } + #[test] fn endpoint_metadata_marks_reasoning_alias_from_external_model() { let endpoint = json!({ diff --git a/crates/goose-providers/src/lib.rs b/crates/goose-providers/src/lib.rs index fb75fac7b..642a8909e 100644 --- a/crates/goose-providers/src/lib.rs +++ b/crates/goose-providers/src/lib.rs @@ -6,8 +6,8 @@ pub mod databricks_auth; pub mod databricks_v2; pub mod google; pub use goose_provider_types::{ - base, cache_semantics, canonical, conversation, errors, formats, goose_mode, images, json, - model, permission, request_log, retry, thinking, utils, + base, cache_semantics, canonical, context_limit, conversation, errors, formats, goose_mode, + images, json, model, permission, request_log, retry, thinking, utils, }; pub mod declarative; pub mod http_status; diff --git a/crates/goose-providers/src/ollama.rs b/crates/goose-providers/src/ollama.rs index 2a90c0811..9020f9a1b 100644 --- a/crates/goose-providers/src/ollama.rs +++ b/crates/goose-providers/src/ollama.rs @@ -1,5 +1,5 @@ use super::api_client::ApiClient; -use super::base::{ConfigKey, MessageStream, Provider, ProviderMetadata}; +use super::base::{ConfigKey, MessageStream, ModelInfo, Provider, ProviderMetadata}; use super::openai_compatible::handle_status; use super::retry::{ProviderRetry, RetryConfig}; use crate::api_client::{AuthMethod, TlsConfig}; @@ -55,8 +55,8 @@ const OLLAMA_MAX_RETRY_INTERVAL_MS: u64 = 15_000; #[derive(Debug, Clone, serde::Serialize)] pub struct OllamaOptions { /// Explicit context window override from `GOOSE_INPUT_LIMIT`. - /// `None` when unset, zero, or invalid; the model's context limit is then - /// used as the fallback. + /// `None` when unset, zero, or invalid; `num_ctx` is then omitted so Ollama + /// uses its model default. pub input_limit: Option, /// Whether to keep `stream_options` in the request (`OLLAMA_STREAM_USAGE`, /// default `true`). @@ -82,7 +82,7 @@ pub struct OllamaProvider { #[serde(skip)] api_client: ApiClient, name: String, - custom_models: Option>, + custom_models: Option>, dynamic_models: Option, skip_canonical_filtering: bool, options: OllamaOptions, @@ -91,7 +91,7 @@ pub struct OllamaProvider { pub struct OllamaProviderBuilder { api_client: ApiClient, name: String, - custom_models: Option>, + custom_models: Option>, dynamic_models: Option, skip_canonical_filtering: bool, options: OllamaOptions, @@ -132,7 +132,7 @@ impl OllamaProviderBuilder { self } - pub fn custom_models(mut self, custom_models: Option>) -> Self { + pub fn custom_models(mut self, custom_models: Option>) -> Self { self.custom_models = custom_models; self } @@ -230,11 +230,11 @@ pub async fn fetch_ollama_model_names( Ok(Some(names)) } -fn resolve_ollama_num_ctx(options: &OllamaOptions, model_config: &ModelConfig) -> Option { - options.input_limit.or(model_config.context_limit) +fn resolve_ollama_num_ctx(options: &OllamaOptions) -> Option { + options.input_limit } -fn apply_ollama_options(payload: &mut Value, options: &OllamaOptions, model_config: &ModelConfig) { +fn apply_ollama_options(payload: &mut Value, options: &OllamaOptions, _model_config: &ModelConfig) { if let Some(obj) = payload.as_object_mut() { // Gate stream_options behind OLLAMA_STREAM_USAGE (default: true). // Older Ollama builds that don't support stream_options may stall before @@ -257,8 +257,8 @@ fn apply_ollama_options(payload: &mut Value, options: &OllamaOptions, model_conf } } - // Apply num_ctx from context limit settings. - if let Some(limit) = resolve_ollama_num_ctx(options, model_config) { + // Only an explicit GOOSE_INPUT_LIMIT overrides Ollama's num_ctx. + if let Some(limit) = resolve_ollama_num_ctx(options) { let options_value = obj.entry("options").or_insert_with(|| json!({})); if let Some(options_obj) = options_value.as_object_mut() { options_obj.insert("num_ctx".to_string(), json!(limit)); @@ -273,13 +273,7 @@ pub fn from_declarative_config( key_resolver: impl KeyResolver, ) -> Result { let custom_models = if !config.models.is_empty() { - Some( - config - .models - .iter() - .map(|m| m.name.clone()) - .collect::>(), - ) + Some(config.models.clone()) } else { None }; @@ -440,10 +434,27 @@ impl Provider for OllamaProvider { stream_ollama(response, self.options.chunk_timeout_secs, log) } + async fn get_context_limit(&self, model: &str, override_limit: Option) -> usize { + let configured_limits = self + .custom_models + .iter() + .flatten() + .filter_map(|model| model.context_limit.map(|limit| (model.name.clone(), limit))); + crate::context_limit::ContextLimitResolver::new(&self.name) + .with_configured_limits(configured_limits) + .resolve(model, override_limit, || async { + Ok(self.options.input_limit) + }) + .await + } + async fn fetch_supported_models(&self) -> Result, ProviderError> { if let Some(custom_models) = &self.custom_models { if self.dynamic_models == Some(false) { - return Ok(custom_models.clone()); + return Ok(custom_models + .iter() + .map(|model| model.name.clone()) + .collect()); } match self.fetch_models_from_api().await { @@ -454,7 +465,10 @@ impl Provider for OllamaProvider { self.name, e ); - return Ok(custom_models.clone()); + return Ok(custom_models + .iter() + .map(|model| model.name.clone()) + .collect()); } Err(e) => return Err(e), } @@ -581,7 +595,10 @@ mod tests { #[tokio::test] async fn fetch_supported_models_uses_static_models_when_dynamic_models_false() { let provider = from_declarative_config( - ollama_config(Some(false), vec![ModelInfo::new("static-model", 4096)]), + ollama_config( + Some(false), + vec![ModelInfo::new("static-model").with_context_limit(4096)], + ), None, crate::declarative::EnvKeyResolver, ) @@ -625,7 +642,7 @@ mod tests { let provider = from_declarative_config( ollama_config_with_base_url( None, - vec![ModelInfo::new("static-model", 4096)], + vec![ModelInfo::new("static-model").with_context_limit(4096)], &server.uri(), ), None, @@ -653,21 +670,11 @@ mod tests { } #[test] - fn test_apply_ollama_options_falls_back_to_context_limit() { + fn test_apply_ollama_options_ignores_context_management_limit() { let options = OllamaOptions::default(); let model_config = ModelConfig::new("qwen3").with_context_limit(Some(12_000)); let mut payload = json!({}); apply_ollama_options(&mut payload, &options, &model_config); - assert_eq!(payload["options"]["num_ctx"], 12_000); - } - - #[test] - fn test_apply_ollama_options_skips_when_no_limit() { - let options = OllamaOptions::default(); - let mut model_config = ModelConfig::new("qwen3"); - model_config.context_limit = None; - let mut payload = json!({}); - apply_ollama_options(&mut payload, &options, &model_config); assert!(payload.get("options").is_none()); } diff --git a/crates/goose-providers/src/openai.rs b/crates/goose-providers/src/openai.rs index 9a1379d28..c91d129c1 100644 --- a/crates/goose-providers/src/openai.rs +++ b/crates/goose-providers/src/openai.rs @@ -28,6 +28,7 @@ use reqwest::StatusCode; use serde_json::json; use std::collections::HashMap; use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; use crate::base::{MessageStream, ProviderDescriptor}; use crate::model::ModelConfig; @@ -38,6 +39,24 @@ pub const OPEN_AI_DEFAULT_BASE_PATH: &str = "v1/chat/completions"; pub const OPEN_AI_VERSIONLESS_BASE_PATH: &str = "chat/completions"; const OPEN_AI_DEFAULT_RESPONSES_PATH: &str = "v1/responses"; const OPEN_AI_DEFAULT_MODELS_PATH: &str = "v1/models"; +const N_CTX_PROBE_TIMEOUT: Duration = Duration::from_secs(5); +const N_CTX_FAILURE_TTL: Duration = Duration::from_secs(60); + +#[derive(Debug, Clone, Copy)] +enum CachedContextLimit { + Success(Option), + Failure(Instant), +} + +impl CachedContextLimit { + fn value(self) -> Option> { + match self { + Self::Success(limit) => Some(limit), + Self::Failure(fetched_at) if fetched_at.elapsed() < N_CTX_FAILURE_TTL => Some(None), + Self::Failure(_) => None, + } + } +} pub const OPEN_AI_DEFAULT_MODEL: &str = "gpt-4o"; pub const OPEN_AI_DEFAULT_FAST_MODEL: &str = "gpt-4o-mini"; pub const OPEN_AI_KNOWN_MODELS: &[(&str, usize)] = &[ @@ -149,7 +168,7 @@ pub struct OpenAiProvider { skip_canonical_filtering: bool, preserve_thinking_context: bool, #[serde(skip)] - n_ctx_cache: Arc>>>, + n_ctx_cache: Arc>>, } /// Builder for [`OpenAiProvider`]. @@ -572,17 +591,21 @@ impl OpenAiProvider { /// llama.cpp and Ollama expose the actual allocated context window in the /// non-standard `meta.n_ctx` field of `/v1/models`. Returns `None` when absent /// (e.g. real OpenAI). - async fn fetch_n_ctx_from_api(&self, model_name: &str) -> Option { + async fn fetch_n_ctx_from_api(&self, model_name: &str) -> Result, ProviderError> { let models_path = Self::map_base_path(&self.base_path, "models", OPEN_AI_DEFAULT_MODELS_PATH); - let response = self - .api_client - .request(&models_path) - .response_get() - .await - .ok()?; - let json = handle_response_openai_compat(response).await.ok()?; - parse_n_ctx_from_models(&json, model_name) + let response = self.api_client.request(&models_path).response_get().await?; + if response.status() == StatusCode::NOT_FOUND { + return Ok(None); + } + let json = handle_response_openai_compat(response).await.map_err(|error| { + if matches!(&error, ProviderError::RequestFailed(message) if message.contains("not valid JSON")) { + ProviderError::EndpointNotFound(error.to_string()) + } else { + error + } + })?; + Ok(parse_n_ctx_from_models(&json, model_name)) } } @@ -634,7 +657,7 @@ impl ProviderDescriptor for OpenAiProvider { fn metadata() -> ProviderMetadata { let models = OPEN_AI_KNOWN_MODELS .iter() - .map(|(name, limit)| ModelInfo::new(*name, *limit)) + .map(|(name, limit)| ModelInfo::new(*name).with_context_limit(*limit)) .collect(); ProviderMetadata::with_models( OPEN_AI_PROVIDER_NAME, @@ -692,41 +715,51 @@ impl Provider for OpenAiProvider { self.skip_canonical_filtering } - /// Resolve the effective context limit. When the config carries an explicit - /// limit (GOOSE_CONTEXT_LIMIT, a session override, or a known/canonical - /// value) it is used as-is. Otherwise probe `/v1/models`: llama.cpp and - /// Ollama report the real allocated window via the non-standard - /// `meta.n_ctx` field, which fixes auto-compaction for local servers that - /// would otherwise fall back to DEFAULT_CONTEXT_LIMIT. The probe is bounded - /// by a short timeout so a hung endpoint can't stall the caller. - async fn get_context_limit(&self, model_config: &ModelConfig) -> Result { - if let Some(limit) = model_config.context_limit { - return Ok(limit); - } + async fn get_context_limit(&self, model: &str, override_limit: Option) -> usize { + let configured_limits = self + .custom_models + .iter() + .flatten() + .filter_map(|model| model.context_limit.map(|limit| (model.name.clone(), limit))); + let resolver = goose_provider_types::context_limit::ContextLimitResolver::new(&self.name) + .with_configured_limits(configured_limits); - if let Some(cached) = self - .n_ctx_cache - .lock() - .ok() - .and_then(|cache| cache.get(&model_config.model_name).copied()) - { - return Ok(cached.unwrap_or_else(|| model_config.context_limit())); - } + resolver + .resolve(model, override_limit, || async { + if let Some(cached) = self + .n_ctx_cache + .lock() + .ok() + .and_then(|cache| cache.get(model).copied()) + .and_then(CachedContextLimit::value) + { + return Ok(cached); + } - const N_CTX_PROBE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); - let probed = tokio::time::timeout( - N_CTX_PROBE_TIMEOUT, - self.fetch_n_ctx_from_api(&model_config.model_name), - ) - .await - .ok() - .flatten(); + let probed = match tokio::time::timeout( + N_CTX_PROBE_TIMEOUT, + self.fetch_n_ctx_from_api(model), + ) + .await + { + Ok(Ok(limit)) => Ok(limit), + Ok(Err(error)) if error.is_endpoint_not_found() => Ok(None), + Ok(Err(error)) => Err(error), + Err(_) => Err(ProviderError::RequestFailed( + "Context-limit discovery timed out".into(), + )), + }; - if let Ok(mut cache) = self.n_ctx_cache.lock() { - cache.insert(model_config.model_name.clone(), probed); - } - - Ok(probed.unwrap_or_else(|| model_config.context_limit())) + if let Ok(mut cache) = self.n_ctx_cache.lock() { + let cached = match probed.as_ref() { + Ok(limit) => CachedContextLimit::Success(*limit), + Err(_) => CachedContextLimit::Failure(Instant::now()), + }; + cache.insert(model.to_string(), cached); + } + probed + }) + .await } async fn fetch_supported_models(&self) -> Result, ProviderError> { @@ -1343,7 +1376,7 @@ mod tests { description: None, api_key_env: String::new(), base_url: base_url.to_string(), - models: vec![crate::base::ModelInfo::new("test-model", 4096)], + models: vec![crate::base::ModelInfo::new("test-model").with_context_limit(4096)], headers: None, timeout_seconds: None, supports_streaming: None, @@ -1459,7 +1492,7 @@ mod tests { custom_models: Some( custom_models .into_iter() - .map(|model| ModelInfo::new(model, 4096)) + .map(|model| ModelInfo::new(model).with_context_limit(4096)) .collect(), ), dynamic_models: Some(true), @@ -1469,6 +1502,94 @@ mod tests { } } + #[tokio::test] + async fn context_limit_caches_failed_models_probe() { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/models")) + .respond_with(ResponseTemplate::new(500)) + .expect(1) + .mount(&server) + .await; + + let mut provider = make_provider_with_custom_models( + &server.uri(), + "v1/chat/completions", + vec!["other-model".to_string()], + ); + provider.custom_models = None; + + for _ in 0..2 { + assert_eq!( + provider.get_context_limit("unknown-model", None).await, + crate::model::DEFAULT_CONTEXT_LIMIT + ); + } + } + + #[tokio::test] + async fn context_limit_retries_expired_models_probe_failure() { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/models")) + .respond_with(ResponseTemplate::new(500)) + .expect(1) + .mount(&server) + .await; + + let mut provider = make_provider_with_custom_models( + &server.uri(), + "v1/chat/completions", + vec!["other-model".to_string()], + ); + provider.custom_models = None; + provider.n_ctx_cache.lock().unwrap().insert( + "unknown-model".to_string(), + CachedContextLimit::Failure(Instant::now() - N_CTX_FAILURE_TTL), + ); + + assert_eq!( + provider.get_context_limit("unknown-model", None).await, + crate::model::DEFAULT_CONTEXT_LIMIT + ); + } + + #[tokio::test] + async fn context_limit_caches_unsupported_models_endpoint() { + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/v1/models")) + .respond_with(ResponseTemplate::new(200).set_body_string("not json")) + .expect(1) + .mount(&server) + .await; + + let mut provider = make_provider_with_custom_models( + &server.uri(), + "v1/chat/completions", + vec!["other-model".to_string()], + ); + provider.custom_models = None; + + assert_eq!( + provider.get_context_limit("unknown-model", None).await, + crate::model::DEFAULT_CONTEXT_LIMIT + ); + assert_eq!( + provider.get_context_limit("unknown-model", None).await, + crate::model::DEFAULT_CONTEXT_LIMIT + ); + } + #[tokio::test] async fn fetch_models_treats_invalid_json_as_endpoint_not_found() { use wiremock::matchers::{method, path}; diff --git a/crates/goose-sdk/src/bindings.rs b/crates/goose-sdk/src/bindings.rs index ae530df94..889ea7dc7 100644 --- a/crates/goose-sdk/src/bindings.rs +++ b/crates/goose-sdk/src/bindings.rs @@ -389,7 +389,6 @@ pub struct ProviderModelConfig { impl ProviderModelConfig { fn to_goose_model_config(&self) -> Result { let mut config = ModelConfig::new(&self.model_name) - .with_context_limit(self.context_limit.map(|limit| limit.max(0) as usize)) .with_temperature(self.temperature) .with_max_tokens(self.max_tokens) .with_toolshim(self.toolshim) @@ -615,6 +614,20 @@ impl ProviderHandle { self.provider.get_name().to_string() } + async fn context_limit(&self, model: ProviderModelConfig) -> Result { + let normalized_model = ModelConfig::new(&model.model_name); + let override_limit = model + .context_limit + .and_then(|limit| (limit > 0).then_some(limit as usize)); + let provider = Arc::clone(&self.provider); + run_on_runtime(async move { + provider + .get_context_limit(&normalized_model.model_name, override_limit) + .await + }) + .await + } + async fn stream( &self, model: ProviderModelConfig, @@ -736,6 +749,10 @@ impl Provider { features } + pub async fn context_limit(&self, model: ProviderModelConfig) -> Result { + Ok(self.handle.context_limit(model).await? as u64) + } + pub async fn stream( &self, model: ProviderModelConfig, @@ -1125,6 +1142,43 @@ mod tests { } } + #[test] + fn context_limit_override_filters_nonpositive_values() { + let none = base_model_config(); + assert_eq!( + none.context_limit + .and_then(|limit| (limit > 0).then_some(limit as usize)), + None + ); + + let negative = ProviderModelConfig { + context_limit: Some(-1), + ..base_model_config() + }; + assert_eq!( + negative + .context_limit + .and_then(|limit| (limit > 0).then_some(limit as usize)), + None + ); + + let positive = ProviderModelConfig { + context_limit: Some(64_000), + ..base_model_config() + }; + assert_eq!( + positive + .context_limit + .and_then(|limit| (limit > 0).then_some(limit as usize)), + Some(64_000) + ); + } + + #[test] + fn model_config_normalizes_effort_suffix() { + assert_eq!(ModelConfig::new("gpt-5.4-xhigh").model_name, "gpt-5.4"); + } + #[test] fn model_config_rejects_invalid_request_params_json() { let config = ProviderModelConfig { diff --git a/crates/goose/src/acp/provider.rs b/crates/goose/src/acp/provider.rs index d9ecf9698..f6e14498a 100644 --- a/crates/goose/src/acp/provider.rs +++ b/crates/goose/src/acp/provider.rs @@ -285,9 +285,7 @@ pub struct AcpProvider { /// Failed or abandoned first prompts reset this so the next prompt can retry it. handoff_context_sent: Arc, /// Latest `size` reported by the ACP server in a `session/update` → - /// `usage_update` notification. 0 means no real update has arrived yet, - /// in which case `get_context_limit()` falls back to the supplied model - /// configuration's context limit. + /// `usage_update` notification. 0 means no real update has arrived yet. context_size: Arc, /// Config option id used to select the model, if this agent supports it. @@ -659,7 +657,9 @@ impl AcpProvider { } }; - let context_limit = self.get_context_limit(model_config).await.ok()?; + let context_limit = crate::context_limit::get_context_limit(self, &model_config.model_name) + .await + .ok()?; let budget = memo_token_budget(context_limit, prompt_token_cost(current_prompt, &counter)); build_handoff_context_memo(&messages[..last_user_index], budget, &counter) @@ -706,12 +706,13 @@ impl Provider for AcpProvider { Ok(()) } - async fn get_context_limit(&self, model_config: &ModelConfig) -> Result { - let size = self.context_size.load(Ordering::Relaxed); - if size > 0 { - return Ok(size as usize); - } - Ok(model_config.context_limit()) + async fn get_context_limit(&self, model: &str, override_limit: Option) -> usize { + goose_providers::context_limit::ContextLimitResolver::new(self.get_name()) + .resolve(model, override_limit, || async { + let size = self.context_size.load(Ordering::Relaxed); + Ok((size > 0).then_some(size as usize)) + }) + .await } async fn update_mode(&self, session_id: &str, mode: GooseMode) -> Result<(), ProviderError> { @@ -2888,12 +2889,35 @@ mod tests { async fn get_context_limit_surfaces_captured_context_size() { let (provider, model) = test_provider(); assert_eq!( - provider.get_context_limit(&model).await.unwrap(), + provider.get_context_limit(&model.model_name, None).await, goose_providers::model::DEFAULT_CONTEXT_LIMIT ); provider.context_size.store(200_000, Ordering::Relaxed); - assert_eq!(provider.get_context_limit(&model).await.unwrap(), 200_000); + assert_eq!( + provider.get_context_limit(&model.model_name, None).await, + 200_000 + ); + } + + #[tokio::test] + async fn handoff_budget_honors_global_context_limit_override() { + let _guard = env_lock::lock_env([("GOOSE_CONTEXT_LIMIT", Some("64"))]); + let (provider, model) = test_provider(); + provider.context_size.store(200_000, Ordering::Relaxed); + let messages = vec![ + Message::assistant().with_text("prior answer"), + Message::user().with_text("current request"), + ]; + let current_prompt = vec![ContentBlock::Text(TextContent::new("current request"))]; + + assert!( + provider + .bounded_handoff_memo(&model, &messages, ¤t_prompt) + .await + .is_none(), + "the global limit should leave too little room for a handoff memo" + ); } #[tokio::test] @@ -3064,6 +3088,7 @@ mod tests { #[tokio::test] async fn failed_handoff_send_consumes_the_claim() { + let _guard = env_lock::lock_env([("GOOSE_CONTEXT_LIMIT", None::<&str>)]); let (tx, rx) = mpsc::channel(1); drop(rx); let (provider, model) = test_provider_with_tx(Some(tx)); diff --git a/crates/goose/src/acp/response_builder.rs b/crates/goose/src/acp/response_builder.rs index 0a8dae16a..d2adf210d 100644 --- a/crates/goose/src/acp/response_builder.rs +++ b/crates/goose/src/acp/response_builder.rs @@ -463,18 +463,18 @@ pub(super) fn send_session_setup_notifications( cx: &ConnectionTo, session: &Session, totals: &SessionUsageTotals, + context_limit: usize, supports_goose_custom_notifications: bool, ) -> Result<(), agent_client_protocol::Error> { let session_id = SessionId::new(session.id.clone()); - if let Some(updates) = build_usage_updates(session, totals) { - if supports_goose_custom_notifications { - cx.send_notification(updates.custom)?; - } - cx.send_notification(SessionNotification::new( - session_id.clone(), - SessionUpdate::UsageUpdate(updates.standard), - ))?; + let updates = build_usage_updates(session, totals, context_limit); + if supports_goose_custom_notifications { + cx.send_notification(updates.custom)?; } + cx.send_notification(SessionNotification::new( + session_id.clone(), + SessionUpdate::UsageUpdate(updates.standard), + ))?; cx.send_notification(SessionNotification::new( session_id, SessionUpdate::AvailableCommandsUpdate(available_commands_update(&session.working_dir)), diff --git a/crates/goose/src/acp/server.rs b/crates/goose/src/acp/server.rs index e90f053cc..a789c5ade 100644 --- a/crates/goose/src/acp/server.rs +++ b/crates/goose/src/acp/server.rs @@ -755,14 +755,15 @@ pub(super) struct UsageUpdates { pub(super) fn build_usage_updates( session: &Session, totals: &SessionUsageTotals, -) -> Option { + context_limit: usize, +) -> UsageUpdates { let used = session.usage.total_tokens.unwrap_or(0).max(0) as u64; - let ctx_limit = session.model_config.as_ref()?.context_limit() as u64; + let ctx_limit = context_limit as u64; let accumulated_input_tokens = to_nonnegative_u64(totals.accumulated_usage.input_tokens).unwrap_or(0); let accumulated_output_tokens = to_nonnegative_u64(totals.accumulated_usage.output_tokens).unwrap_or(0); - Some(UsageUpdates { + UsageUpdates { custom: GooseSessionNotification { session_id: session.id.clone(), update: GooseSessionUpdate::UsageUpdate(SessionUsageUpdate { @@ -780,7 +781,7 @@ pub(super) fn build_usage_updates( } standard }, - }) + } } /// Resolve the cwd an existing session should be activated with: a @@ -851,39 +852,60 @@ impl GooseAcpAgent { .unwrap_or(false) } - pub(super) async fn notify_session_setup( - &self, - cx: &ConnectionTo, - session: &Session, - ) -> Result<(), agent_client_protocol::Error> { - let totals = self - .session_manager - .get_session_usage_totals(&session.id) - .await - .unwrap_or_default(); - send_session_setup_notifications( - cx, - session, - &totals, - self.supports_goose_custom_notifications(), - ) - } - pub(super) async fn prepare_session_setup_by_id( &self, session_id: &str, - ) -> Result<(Session, SessionUsageTotals), agent_client_protocol::Error> { - let session = self - .session_manager - .get_session(session_id, false) - .await - .internal_err_ctx("Failed to load session for setup notifications")?; - let totals = self - .session_manager - .get_session_usage_totals(session_id) - .await - .unwrap_or_default(); - Ok((session, totals)) + ) -> Result<(Session, SessionUsageTotals, usize), agent_client_protocol::Error> { + loop { + let session = self + .session_manager + .get_session(session_id, false) + .await + .internal_err_ctx("Failed to load session for setup notifications")?; + let model_name = session + .model_config + .as_ref() + .map(|model| model.model_name.clone()) + .ok_or_else(|| { + agent_client_protocol::Error::internal_error().data("Session has no model") + })?; + let provider_name = session.provider_name.clone(); + let agent = self.get_session_agent(session_id).await?; + let provider = agent + .provider() + .await + .internal_err_ctx("Failed to resolve session provider")?; + let context_limit = + crate::context_limit::get_context_limit(provider.as_ref(), &model_name) + .await + .internal_err_ctx("Failed to resolve context limit")?; + let session = self + .session_manager + .get_session(session_id, false) + .await + .internal_err_ctx("Failed to refresh session for setup notifications")?; + let current_provider = agent + .provider() + .await + .internal_err_ctx("Failed to refresh session provider")?; + let refreshed_model_name = session + .model_config + .as_ref() + .map(|model| model.model_name.as_str()); + if provider_name.as_deref() != Some(provider.get_name()) + || !Arc::ptr_eq(&provider, ¤t_provider) + || session.provider_name != provider_name + || refreshed_model_name != Some(model_name.as_str()) + { + continue; + } + let totals = self + .session_manager + .get_session_usage_totals(session_id) + .await + .unwrap_or_default(); + return Ok((session, totals, context_limit)); + } } pub(super) fn supports_recipe_param_requests(&self) -> bool { @@ -2245,18 +2267,25 @@ impl GooseAcpAgent { .get_session_usage_totals(&session_id) .await .unwrap_or_default(); - if let Some(updates) = build_usage_updates(&session, &totals) { - if self.supports_goose_custom_notifications() { - cx.send_notification(updates.custom)?; - } - // Standard ACP notification — emitted alongside the custom one for - // backwards compatibility. Remove once all known clients have - // migrated to `_goose/unstable/session/update`. - cx.send_notification(SessionNotification::new( - args.session_id.clone(), - SessionUpdate::UsageUpdate(updates.standard), - ))?; + let provider = agent + .provider() + .await + .internal_err_ctx("Failed to resolve session provider")?; + let model = session.model_config.as_ref().ok_or_else(|| { + agent_client_protocol::Error::internal_error().data("Session has no model") + })?; + let context_limit = + crate::context_limit::get_context_limit(provider.as_ref(), &model.model_name) + .await + .internal_err_ctx("Failed to resolve context limit")?; + let updates = build_usage_updates(&session, &totals, context_limit); + if self.supports_goose_custom_notifications() { + cx.send_notification(updates.custom)?; } + cx.send_notification(SessionNotification::new( + args.session_id.clone(), + SessionUpdate::UsageUpdate(updates.standard), + ))?; let stop_reason = prompt_stop_reason(was_cancelled, output_token_limit_reached); @@ -3374,8 +3403,7 @@ print(\"hello, world\") accumulated_usage: session.accumulated_usage, accumulated_cost: session.accumulated_cost, }; - let updates = - build_usage_updates(&session, &totals).expect("usage updates should be present"); + let updates = build_usage_updates(&session, &totals, 258_000); assert_eq!(updates.custom.session_id, "session-1"); let usage = match updates.custom.update { GooseSessionUpdate::UsageUpdate(usage) => usage, @@ -3387,15 +3415,6 @@ print(\"hello, world\") assert_eq!(updates.standard.size, 258_000); } - #[test] - fn test_build_usage_update_requires_model_config() { - let session = make_session_with_usage( - TokenUsage::new(Some(80), Some(40), Some(120)), - TokenUsage::default(), - ); - assert!(build_usage_updates(&session, &SessionUsageTotals::default()).is_none()); - } - #[test] fn test_goose_custom_notifications_capability_defaults_to_false() { let request = InitializeRequest::new(agent_client_protocol::schema::ProtocolVersion::V1); diff --git a/crates/goose/src/acp/server/dispatch.rs b/crates/goose/src/acp/server/dispatch.rs index 86ed61476..8623559eb 100644 --- a/crates/goose/src/acp/server/dispatch.rs +++ b/crates/goose/src/acp/server/dispatch.rs @@ -48,14 +48,15 @@ impl HandleDispatchFrom for GooseAcpHandler { match agent.on_new_session(&cx_clone, req).await { Ok(response) => { let session_id = response.session_id.0.to_string(); + responder.respond(response)?; let session_setup = agent.prepare_session_setup_by_id(&session_id).await; - responder.respond(response)?; - if let Err(error) = session_setup.and_then(|(session, totals)| { + if let Err(error) = session_setup.and_then(|(session, totals, context_limit)| { send_session_setup_notifications( &cx_clone, &session, &totals, + context_limit, agent.supports_goose_custom_notifications(), ) }) { @@ -85,6 +86,25 @@ impl HandleDispatchFrom for GooseAcpHandler { match agent.on_load_session(&cx_clone, req).await { Ok(response) => { responder.respond(response)?; + let session_setup = + agent.prepare_session_setup_by_id(&session_id).await; + if let Err(error) = session_setup.and_then( + |(session, totals, context_limit)| { + send_session_setup_notifications( + &cx_clone, + &session, + &totals, + context_limit, + agent.supports_goose_custom_notifications(), + ) + }, + ) { + tracing::warn!( + session_id = %session_id, + error = ?error, + "Failed to send ACP session setup notifications" + ); + } } Err(e) => { tracing::error!( @@ -420,14 +440,15 @@ impl HandleDispatchFrom for GooseAcpHandler { match agent.on_fork_session(&cx_spawn, req).await { Ok(response) => { let session_id = response.session_id.0.to_string(); + responder.respond(response)?; let session_setup = agent.prepare_session_setup_by_id(&session_id).await; - responder.respond(response)?; - if let Err(error) = session_setup.and_then(|(session, totals)| { + if let Err(error) = session_setup.and_then(|(session, totals, context_limit)| { send_session_setup_notifications( &cx_spawn, &session, &totals, + context_limit, agent.supports_goose_custom_notifications(), ) }) { diff --git a/crates/goose/src/acp/server/load_session.rs b/crates/goose/src/acp/server/load_session.rs index a11747f28..376fa55a8 100644 --- a/crates/goose/src/acp/server/load_session.rs +++ b/crates/goose/src/acp/server/load_session.rs @@ -353,8 +353,6 @@ impl GooseAcpAgent { ) .await?; - self.notify_session_setup(cx, &session).await?; - let mut response = LoadSessionResponse::new().modes(mode_state); if let Some(co) = config_options { response = response.config_options(co); diff --git a/crates/goose/src/acp/server/providers.rs b/crates/goose/src/acp/server/providers.rs index 86268d724..a22649d92 100644 --- a/crates/goose/src/acp/server/providers.rs +++ b/crates/goose/src/acp/server/providers.rs @@ -2,6 +2,7 @@ use super::*; use crate::config::declarative_providers; use crate::providers::inventory::ensure_refresh_identity_current; use crate::providers::provider_secrets; +use goose_providers::base::ModelInfo; use std::str::FromStr; const ACP_READINESS_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); @@ -378,6 +379,34 @@ fn custom_provider_headers(headers: HashMap) -> Option, + existing: &[ModelInfo], + catalog_provider_id: Option<&str>, +) -> Vec { + let catalog_models = catalog_provider_id + .and_then(crate::providers::catalog::get_provider_template) + .map(|template| template.models) + .unwrap_or_default(); + + names + .into_iter() + .map(|name| { + existing + .iter() + .find(|model| model.name == name) + .cloned() + .or_else(|| { + catalog_models + .iter() + .find(|model| model.id == name) + .map(|model| ModelInfo::new(&name).with_context_limit(model.context_limit)) + }) + .unwrap_or_else(|| ModelInfo::new(name)) + }) + .collect() +} + fn load_declarative_provider_for_client( provider_id: &str, ) -> Result { @@ -611,7 +640,11 @@ impl GooseAcpAgent { display_name: provider.display_name, api_url: provider.api_url, api_key: provider.api_key, - models: provider.models, + models: custom_provider_models( + provider.models, + &[], + provider.catalog_provider_id.as_deref(), + ), supports_streaming: provider.supports_streaming, headers: custom_provider_headers(provider.headers), requires_auth: provider.requires_auth, @@ -680,7 +713,11 @@ impl GooseAcpAgent { display_name: provider.display_name, api_url: provider.api_url, api_key: provider.api_key, - models: provider.models, + models: custom_provider_models( + provider.models, + &loaded.config.models, + provider.catalog_provider_id.as_deref(), + ), supports_streaming: provider.supports_streaming, headers: Some(provider.headers), requires_auth: provider.requires_auth, @@ -1183,7 +1220,9 @@ impl GooseAcpAgent { config_info.map(|info| CanonicalModelInfoDto { provider: req.provider.clone(), model: req.model.clone(), - context_limit: info.context_limit, + context_limit: info.context_limit.unwrap_or_else(|| { + ModelConfig::new(&req.model).context_limit() + }), // ModelInfo carries no max-output limit. max_output_tokens: None, // Configs deserialize a missing `reasoning` as false; keep diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 2232e5525..f8b31069d 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -893,10 +893,12 @@ impl Agent { Ok(v) => v, Err(_) => { let context_limit = match self.provider().await { - Ok(provider) => provider - .get_context_limit(&model_config) - .await - .unwrap_or_else(|_| model_config.context_limit()), + Ok(provider) => crate::context_limit::get_context_limit( + provider.as_ref(), + &model_config.model_name, + ) + .await + .unwrap_or(goose_providers::model::DEFAULT_CONTEXT_LIMIT), Err(_) => goose_providers::model::DEFAULT_CONTEXT_LIMIT, }; let compaction_threshold = Config::global() @@ -1908,10 +1910,9 @@ impl Agent { } }; - let context_limit = provider - .get_context_limit(&model_config) - .await - .unwrap_or_else(|_| model_config.context_limit()); + let context_limit = + crate::context_limit::get_context_limit(provider.as_ref(), &model_config.model_name) + .await?; let steer_queue = self.steer_queue(&session_id).await; let machine = self.create_state_machine( provider, @@ -3623,10 +3624,6 @@ impl Agent { ) -> Result<()> { let provider_name = provider.get_name().to_string(); - // Normalize against the provider entry so custom/declarative providers - // backfill `context_limit` from their known models before the config is - // persisted as the session source of truth; otherwise auto-compaction - // would fall back to DEFAULT_CONTEXT_LIMIT. let model_config = match crate::providers::get_from_registry(&provider_name).await { Ok(entry) => entry .normalize_model_config(model_config.clone()) diff --git a/crates/goose/src/agents/execute_commands.rs b/crates/goose/src/agents/execute_commands.rs index c4e923266..da675944e 100644 --- a/crates/goose/src/agents/execute_commands.rs +++ b/crates/goose/src/agents/execute_commands.rs @@ -255,10 +255,9 @@ impl Agent { async fn handle_status_command(&self, session_id: &str) -> Result> { let provider = self.provider().await?; let model_config = self.model_config_for_session(session_id).await?; - let context_limit = provider - .get_context_limit(&model_config) - .await - .unwrap_or_else(|_| model_config.context_limit()); + let context_limit = + crate::context_limit::get_context_limit(provider.as_ref(), &model_config.model_name) + .await?; let goose_mode = self.goose_mode().await; diff --git a/crates/goose/src/agents/moim.rs b/crates/goose/src/agents/moim.rs index 0b1e93a92..728e04915 100644 --- a/crates/goose/src/agents/moim.rs +++ b/crates/goose/src/agents/moim.rs @@ -50,12 +50,12 @@ pub(super) async fn compute_compaction_info( let context_limit = if let Some(model_config) = session_model_config.as_ref() { let provider = extension_manager.get_provider().lock().await.clone(); match provider { - Some(provider) => provider - .get_context_limit(model_config) - .await - .ok() - .or_else(|| Some(model_config.context_limit())), - None => Some(model_config.context_limit()), + Some(provider) => { + crate::context_limit::get_context_limit(provider.as_ref(), &model_config.model_name) + .await + .ok() + } + None => None, } } else { None @@ -95,12 +95,12 @@ pub async fn turn_context_message( let context_limit = if let Some(model_config) = session_model_config.as_ref() { let provider = extension_manager.get_provider().lock().await.clone(); match provider { - Some(provider) => provider - .get_context_limit(model_config) - .await - .ok() - .or_else(|| Some(model_config.context_limit())), - None => Some(model_config.context_limit()), + Some(provider) => { + crate::context_limit::get_context_limit(provider.as_ref(), &model_config.model_name) + .await + .ok() + } + None => None, } } else { None diff --git a/crates/goose/src/agents/platform_extensions/summon.rs b/crates/goose/src/agents/platform_extensions/summon.rs index da9f9fe00..41ccb205a 100644 --- a/crates/goose/src/agents/platform_extensions/summon.rs +++ b/crates/goose/src/agents/platform_extensions/summon.rs @@ -2860,6 +2860,7 @@ You review code."#; } #[tokio::test] + #[serial] async fn test_resolve_provider_reuses_unregistered_parent_provider() { let temp_dir = TempDir::new().unwrap(); let parent_provider: Arc = Arc::new( @@ -2979,13 +2980,11 @@ You review code."#; let parent = parent_config(); let overridden = goose_providers::model::ModelConfig::new(OVERRIDE_MODEL) .with_canonical_limits(PROVIDER); - assert_ne!(parent.context_limit, overridden.context_limit); assert_ne!(parent.reasoning, overridden.reasoning); let resolved = resolve_with_override(Some(OVERRIDE_MODEL), parent); assert_eq!(resolved.model_name, OVERRIDE_MODEL); - assert_eq!(resolved.context_limit, overridden.context_limit); assert_eq!(resolved.max_tokens, overridden.max_tokens); assert_eq!(resolved.reasoning, overridden.reasoning); } diff --git a/crates/goose/src/agents/state_machine/ops_llm.rs b/crates/goose/src/agents/state_machine/ops_llm.rs index ae030bb78..748564054 100644 --- a/crates/goose/src/agents/state_machine/ops_llm.rs +++ b/crates/goose/src/agents/state_machine/ops_llm.rs @@ -170,8 +170,8 @@ impl Provider for GooseInferenceProvider { }))) } - async fn get_context_limit(&self, model_config: &ModelConfig) -> Result { - self.inner.get_context_limit(model_config).await + async fn get_context_limit(&self, model: &str, override_limit: Option) -> usize { + self.inner.get_context_limit(model, override_limit).await } async fn fetch_model_info(&self, model_name: &str) -> Result { diff --git a/crates/goose/src/agents/state_machine/ops_status.rs b/crates/goose/src/agents/state_machine/ops_status.rs index f2e8e5066..b7241695d 100644 --- a/crates/goose/src/agents/state_machine/ops_status.rs +++ b/crates/goose/src/agents/state_machine/ops_status.rs @@ -43,11 +43,11 @@ impl Operation for StatusOperation { if command.command != "status" { return not_applicable(); } - let context_limit = self - .provider - .get_context_limit(&self.model_config) - .await - .unwrap_or_else(|_| self.model_config.context_limit()); + let context_limit = crate::context_limit::get_context_limit( + self.provider.as_ref(), + &self.model_config.model_name, + ) + .await?; let context_tokens = session.usage.total_tokens.unwrap_or(0); let lifetime_tokens = session.accumulated_usage.total_tokens.unwrap_or(0); let context_pct = if context_limit > 0 { diff --git a/crates/goose/src/agents/state_machine/tests/pipeline.rs b/crates/goose/src/agents/state_machine/tests/pipeline.rs index 0a8237aab..3dca2e206 100644 --- a/crates/goose/src/agents/state_machine/tests/pipeline.rs +++ b/crates/goose/src/agents/state_machine/tests/pipeline.rs @@ -60,11 +60,8 @@ impl Provider for FeatureProvider { .await } - async fn get_context_limit( - &self, - model_config: &ModelConfig, - ) -> Result { - self.inner.get_context_limit(model_config).await + async fn get_context_limit(&self, model: &str, override_limit: Option) -> usize { + self.inner.get_context_limit(model, override_limit).await } fn manages_own_context(&self) -> bool { diff --git a/crates/goose/src/agents/state_machine/tests/reconstruction_isolation_lifecycle.rs b/crates/goose/src/agents/state_machine/tests/reconstruction_isolation_lifecycle.rs index 1b1a5287c..58a2d048e 100644 --- a/crates/goose/src/agents/state_machine/tests/reconstruction_isolation_lifecycle.rs +++ b/crates/goose/src/agents/state_machine/tests/reconstruction_isolation_lifecycle.rs @@ -73,9 +73,7 @@ async fn reconstruction_and_session_isolation() -> Result<()> { .build() .expect("valid recipe"); pipeline.set_recipe(recipe).await?; - let model_config = ModelConfig::new("gpt-4o") - .with_canonical_limits("openai") - .with_context_limit(Some(100_000)); + let model_config = ModelConfig::new("gpt-4o").with_canonical_limits("openai"); let pipeline = pipeline .with_model_config(model_config) .await @@ -92,7 +90,7 @@ async fn reconstruction_and_session_isolation() -> Result<()> { .map(|config| config.model_name.as_str()), Some("gpt-4o") ); - assert_eq!(pipeline.context_limit(), 100_000); + assert_eq!(pipeline.context_limit(), 128_000); assert_eq!(restored.goose_mode, GooseMode::Chat); assert_eq!( restored.recipe.as_ref().map(|recipe| recipe.title.as_str()), diff --git a/crates/goose/src/config/declarative_providers.rs b/crates/goose/src/config/declarative_providers.rs index e32d9a81d..501c8e3a2 100644 --- a/crates/goose/src/config/declarative_providers.rs +++ b/crates/goose/src/config/declarative_providers.rs @@ -138,7 +138,7 @@ pub struct CreateCustomProviderParams { pub display_name: String, pub api_url: String, pub api_key: Option, - pub models: Vec, + pub models: Vec, pub supports_streaming: Option, pub headers: Option>, pub requires_auth: bool, @@ -154,7 +154,7 @@ pub struct UpdateCustomProviderParams { pub display_name: String, pub api_url: String, pub api_key: Option, - pub models: Vec, + pub models: Vec, pub supports_streaming: Option, pub headers: Option>, pub requires_auth: bool, @@ -183,11 +183,7 @@ pub fn create_custom_provider( String::new() }; - let model_infos: Vec = params - .models - .into_iter() - .map(|name| ModelInfo::new(name, 128000)) - .collect(); + let model_infos = params.models; let engine = ProviderEngine::from_str(¶ms.engine)?; let preserves_thinking = params @@ -257,16 +253,31 @@ pub fn update_custom_provider(params: UpdateCustomProviderParams) -> Result<()> }; if editable { - let model_infos: Vec = params + let model_infos = params .models .into_iter() - .map(|name| { - existing_config + .map(|mut model| { + if let Some(existing) = existing_config .models .iter() - .find(|existing| existing.name == name) - .cloned() - .unwrap_or_else(|| ModelInfo::new(name, 128000)) + .find(|existing| existing.name == model.name) + { + model.resolved_model = model.resolved_model.or(existing.resolved_model.clone()); + model.context_limit = model.context_limit.or(existing.context_limit); + model.input_token_cost = model.input_token_cost.or(existing.input_token_cost); + model.output_token_cost = + model.output_token_cost.or(existing.output_token_cost); + model.currency = model.currency.or(existing.currency.clone()); + model.supports_cache_control = model + .supports_cache_control + .or(existing.supports_cache_control); + model.reasoning |= existing.reasoning; + model.thinking_preservation_format = model + .thinking_preservation_format + .or(existing.thinking_preservation_format); + model.request_params = model.request_params.or(existing.request_params.clone()); + } + model }) .collect(); @@ -554,7 +565,7 @@ mod tests { models: vec![ModelInfo { name: "test/model".to_string(), resolved_model: None, - context_limit: 128_000, + context_limit: Some(128_000), input_token_cost: None, output_token_cost: None, currency: None, @@ -686,6 +697,59 @@ mod tests { } } + #[test] + fn custom_provider_update_preserves_model_metadata() { + let temp_dir = tempfile::tempdir().unwrap(); + let temp_root = temp_dir.path().display().to_string(); + let _guard = env_lock::lock_env([("GOOSE_PATH_ROOT", Some(temp_root.as_str()))]); + + let mut model = ModelInfo::with_cost("large-model", 1_048_576, 0.000002, 0.000006); + model.request_params = Some(HashMap::from([( + "temperature".to_string(), + serde_json::json!(0.25), + )])); + let created = create_custom_provider(CreateCustomProviderParams { + engine: "openai".to_string(), + display_name: "Large Context".to_string(), + api_url: "https://example.invalid/v1".to_string(), + api_key: None, + models: vec![model], + supports_streaming: Some(true), + headers: None, + requires_auth: false, + catalog_provider_id: None, + base_path: None, + preserves_thinking: None, + }) + .unwrap(); + + update_custom_provider(UpdateCustomProviderParams { + id: created.name.clone(), + engine: "openai".to_string(), + display_name: created.display_name.clone(), + api_url: created.base_url.clone(), + api_key: None, + models: vec![ModelInfo::new("large-model").with_context_limit(2_097_152)], + supports_streaming: Some(true), + headers: None, + requires_auth: false, + catalog_provider_id: None, + base_path: None, + preserves_thinking: None, + }) + .unwrap(); + + let loaded = load_provider(&created.name).unwrap(); + let model = &loaded.config.models[0]; + assert_eq!(model.context_limit, Some(2_097_152)); + assert_eq!(model.input_token_cost, Some(0.000002)); + assert_eq!(model.output_token_cost, Some(0.000006)); + assert_eq!( + model.request_params.as_ref().unwrap()["temperature"], + serde_json::json!(0.25) + ); + } + #[test] fn test_custom_openai_provider_missing_preserves_thinking_defaults_true() { let json = r#"{ @@ -781,7 +845,7 @@ mod tests { display_name: "Z.AI Updated".to_string(), api_url: "https://updated.example.invalid/v1/chat/completions".to_string(), api_key: None, - models: vec!["z-model".to_string()], + models: vec![ModelInfo::new("z-model")], supports_streaming: Some(true), headers: None, requires_auth: false, @@ -899,63 +963,4 @@ mod tests { let result = expand_env_vars("${TEST_EXPAND_OVERRIDE}/path", &env_vars).unwrap(); assert_eq!(result, "https://from-env.com/path"); } - - #[test] - fn test_update_custom_provider_preserves_model_pricing_and_context_limits() { - let temp_dir = tempfile::tempdir().unwrap(); - let temp_root = temp_dir.path().display().to_string(); - let _guard = env_lock::lock_env([("GOOSE_PATH_ROOT", Some(temp_root.as_str()))]); - - let custom_dir = custom_providers_dir(); - std::fs::create_dir_all(&custom_dir).unwrap(); - std::fs::write( - custom_dir.join("priced_provider.json"), - r#"{ - "name": "priced_provider", - "engine": "openai", - "display_name": "Priced", - "description": null, - "api_key_env": "", - "base_url": "https://example.invalid/v1", - "models": [ - { - "name": "kept-model", - "context_limit": 262144, - "input_token_cost": 0.000002, - "output_token_cost": 0.000006 - } - ], - "requires_auth": false - }"#, - ) - .unwrap(); - - update_custom_provider(UpdateCustomProviderParams { - id: "priced_provider".to_string(), - engine: "openai".to_string(), - display_name: "Renamed".to_string(), - api_url: "https://example.invalid/v1".to_string(), - api_key: None, - models: vec!["kept-model".to_string()], - supports_streaming: None, - headers: None, - requires_auth: false, - catalog_provider_id: None, - base_path: None, - preserves_thinking: None, - }) - .unwrap(); - - let loaded = load_provider("priced_provider").unwrap(); - let model = loaded - .config - .models - .iter() - .find(|m| m.name == "kept-model") - .expect("model survives update"); - assert_eq!(loaded.config.display_name, "Renamed"); - assert_eq!(model.context_limit, 262144); - assert_eq!(model.input_token_cost, Some(0.000002)); - assert_eq!(model.output_token_cost, Some(0.000006)); - } } diff --git a/crates/goose/src/context_limit.rs b/crates/goose/src/context_limit.rs new file mode 100644 index 000000000..8f4892d12 --- /dev/null +++ b/crates/goose/src/context_limit.rs @@ -0,0 +1,41 @@ +use anyhow::Result; +use goose_providers::context_limit::ContextLimitResolver; + +use crate::config::Config; +use crate::providers::base::Provider; + +pub async fn get_context_limit(provider: &dyn Provider, model: &str) -> Result { + let override_limit = Config::global().get_goose_context_limit()?; + Ok(provider.get_context_limit(model, override_limit).await) +} + +pub fn get_local_context_limit(provider_name: &str, model: &str) -> Result { + let override_limit = Config::global().get_goose_context_limit()?; + let mut configured_limits = Vec::new(); + + #[cfg(feature = "aws-providers")] + if provider_name == "aws_bedrock" { + if let Some(limit) = crate::providers::bedrock::local_context_limit(model) { + configured_limits.push((model.to_string(), limit)); + } + } + + #[cfg(feature = "local-inference")] + if provider_name == "local" { + if let Some(limit) = crate::providers::local_inference::local_context_limit(model) { + configured_limits.push((model.to_string(), limit)); + } + } + + configured_limits.extend( + crate::config::declarative_providers::load_provider(provider_name) + .ok() + .into_iter() + .flat_map(|loaded| loaded.config.models) + .filter_map(|model| model.context_limit.map(|limit| (model.name, limit))), + ); + + Ok(ContextLimitResolver::new(provider_name) + .with_configured_limits(configured_limits) + .resolve_local(model, override_limit)) +} diff --git a/crates/goose/src/context_mgmt/mod.rs b/crates/goose/src/context_mgmt/mod.rs index 4242e4805..fd2455935 100644 --- a/crates/goose/src/context_mgmt/mod.rs +++ b/crates/goose/src/context_mgmt/mod.rs @@ -244,10 +244,8 @@ pub async fn check_if_compaction_needed( .model_config .clone() .unwrap_or_else(|| ModelConfig::new("unknown")); - let context_limit = provider - .get_context_limit(&model_config) - .await - .unwrap_or_else(|_| model_config.context_limit()); + let context_limit = + crate::context_limit::get_context_limit(provider, &model_config.model_name).await?; let (current_tokens, _token_source) = match session.usage.total_tokens { Some(tokens) => (tokens as usize, "session metadata"), @@ -702,11 +700,8 @@ mod tests { Ok(stream_from_single_message(message, usage)) } - async fn get_context_limit( - &self, - _model_config: &ModelConfig, - ) -> Result { - Ok(self.config.context_limit()) + async fn get_context_limit(&self, _model: &str, override_limit: Option) -> usize { + override_limit.unwrap_or_else(|| self.config.context_limit()) } } diff --git a/crates/goose/src/lib.rs b/crates/goose/src/lib.rs index 76d9c0193..c3f32c4bd 100644 --- a/crates/goose/src/lib.rs +++ b/crates/goose/src/lib.rs @@ -10,6 +10,7 @@ pub mod agents; pub mod builtin_extension; pub mod checks; pub mod config; +pub mod context_limit; pub mod context_mgmt; pub mod conversation { pub use goose_providers::conversation::*; diff --git a/crates/goose/src/model_config.rs b/crates/goose/src/model_config.rs index 8f481af67..ed5461605 100644 --- a/crates/goose/src/model_config.rs +++ b/crates/goose/src/model_config.rs @@ -23,12 +23,11 @@ pub fn model_config_from_user_config_with_session_settings( model_name: impl AsRef, previous: Option<&ModelConfig>, request_params: Option>, - context_limit: Option, + _context_limit: Option, ) -> Result { let config = Config::global(); let model = base_model_config_from_user_config(provider_name, model_name.as_ref())?; let model = materialize_model_config_inner(model, provider_name, false)? - .with_context_limit(context_limit) .with_inherited_session_settings_from(previous, request_params) .with_default_thinking_effort(config.get_goose_thinking_effort()); @@ -63,9 +62,7 @@ fn materialize_model_config_inner( model = model.with_toolshim_model(get_goose_toolshim_model(config)?); } - model = model - .with_default_context_limit(config.get_goose_context_limit()?) - .with_default_max_tokens(config.get_goose_max_tokens()?); + model = model.with_default_max_tokens(config.get_goose_max_tokens()?); if include_default_thinking_effort { model = model.with_default_thinking_effort(config.get_goose_thinking_effort()); diff --git a/crates/goose/src/providers/anthropic_def.rs b/crates/goose/src/providers/anthropic_def.rs index 1cac2a71f..a2404a404 100644 --- a/crates/goose/src/providers/anthropic_def.rs +++ b/crates/goose/src/providers/anthropic_def.rs @@ -103,7 +103,9 @@ mod tests { .unwrap(); AnthropicProviderBuilder::new(api_client) .name("custom_anthropic") - .custom_models(custom_models) + .custom_models( + custom_models.map(|models| models.into_iter().map(ModelInfo::new).collect()), + ) .dynamic_models(dynamic_models) .build() } @@ -185,8 +187,10 @@ mod tests { #[test] fn from_custom_config_honors_explicit_timeout_seconds() { - let mut config = - base_declarative_config(vec![ModelInfo::new("m1".to_string(), 200000)], Some(false)); + let mut config = base_declarative_config( + vec![ModelInfo::new("m1").with_context_limit(200000)], + Some(false), + ); config.timeout_seconds = Some(120); assert_eq!(built_timeout(config), std::time::Duration::from_secs(120)); } @@ -195,8 +199,10 @@ mod tests { fn from_custom_config_defaults_timeout_when_unset() { // timeout_seconds: None in base config → 600s default, unchanged // behavior for providers that don't set the field. - let config = - base_declarative_config(vec![ModelInfo::new("m1".to_string(), 200000)], Some(false)); + let config = base_declarative_config( + vec![ModelInfo::new("m1").with_context_limit(200000)], + Some(false), + ); assert_eq!(built_timeout(config), std::time::Duration::from_secs(600)); } } diff --git a/crates/goose/src/providers/bedrock.rs b/crates/goose/src/providers/bedrock.rs index c925b0633..55f44c03c 100644 --- a/crates/goose/src/providers/bedrock.rs +++ b/crates/goose/src/providers/bedrock.rs @@ -141,6 +141,14 @@ const BEDROCK_MODEL_TABLE: &[BedrockModelEntry] = &[ }, ]; +pub(crate) fn local_context_limit(model: &str) -> Option { + BEDROCK_MODEL_TABLE + .iter() + .find(|entry| entry.name.eq_ignore_ascii_case(model)) + .and_then(|entry| entry.context_limit) + .map(|limit| limit as usize) +} + fn find_model_entry(name: &str) -> Option<&'static BedrockModelEntry> { // Direct lookup first (handles exact names like "google.gemma-4-31b") if let Some(entry) = BEDROCK_MODEL_TABLE.iter().find(|e| e.name == name) { @@ -835,7 +843,7 @@ impl goose_providers::base::ProviderDescriptor for BedrockProvider { .map(|entry| { entry.context_limit.map_or_else( || model_info_for_provider_model(BEDROCK_PROVIDER_NAME, entry.name), - |limit| ModelInfo::new(entry.name, limit as usize), + |limit| ModelInfo::new(entry.name).with_context_limit(limit as usize), ) }) .collect(); @@ -892,6 +900,18 @@ impl Provider for BedrockProvider { self.retry_config.clone() } + async fn get_context_limit(&self, model: &str, override_limit: Option) -> usize { + let configured_limits = BEDROCK_MODEL_TABLE.iter().filter_map(|entry| { + entry + .context_limit + .map(|limit| (entry.name.to_string(), limit as usize)) + }); + goose_providers::context_limit::ContextLimitResolver::new(&self.name) + .with_configured_limits(configured_limits) + .resolve(model, override_limit, || async { Ok(None) }) + .await + } + async fn fetch_supported_models(&self) -> Result, ProviderError> { Ok(BEDROCK_MODEL_TABLE .iter() @@ -1926,7 +1946,22 @@ mod tests { .iter() .find(|model| model.name == "google.gemma-4-31b") .unwrap(); - assert!(model.context_limit >= 262144); + assert!(model.context_limit.is_some_and(|limit| limit >= 262144)); + } + + #[tokio::test] + async fn test_gemma_context_limit_resolution() { + let (provider, _) = create_mock_provider_and_model("google.gemma-4-31b"); + assert_eq!( + provider.get_context_limit("google.gemma-4-31b", None).await, + 262_144 + ); + assert_eq!( + provider + .get_context_limit("google.gemma-4-31b", Some(64_000)) + .await, + 64_000 + ); } #[test] fn test_converse_model_not_mantle() { diff --git a/crates/goose/src/providers/canonical_cost.rs b/crates/goose/src/providers/canonical_cost.rs index fccfbdf48..2cbfc6a0a 100644 --- a/crates/goose/src/providers/canonical_cost.rs +++ b/crates/goose/src/providers/canonical_cost.rs @@ -219,13 +219,13 @@ mod tests { #[test] fn pricing_from_model_info_returns_none_without_prices() { - let info = ModelInfo::new("m", 1_000); + let info = ModelInfo::new("m").with_context_limit(1_000); assert!(pricing_from_model_info(&info).is_none()); } #[test] fn partial_pricing_is_rejected() { - let mut info = ModelInfo::new("m", 1_000); + let mut info = ModelInfo::new("m").with_context_limit(1_000); info.input_token_cost = Some(0.000002); assert!(pricing_from_model_info(&info).is_none()); } @@ -355,7 +355,7 @@ mod tests { #[test] fn missing_or_usd_currency_falls_back_to_the_dollar_symbol() { - let mut info = ModelInfo::new("m", 1_000); + let mut info = ModelInfo::new("m").with_context_limit(1_000); assert_eq!(display_currency(Some(&info)), "$"); assert_eq!(display_currency(None), "$"); diff --git a/crates/goose/src/providers/huggingface.rs b/crates/goose/src/providers/huggingface.rs index 6d47ceb21..6fa603de6 100644 --- a/crates/goose/src/providers/huggingface.rs +++ b/crates/goose/src/providers/huggingface.rs @@ -1,6 +1,6 @@ use super::api_client::{ApiClient, AuthMethod, AuthProvider}; use super::base::{ - ConfigKey, MessageStream, Provider, ProviderDef, ProviderMetadata, + ConfigKey, MessageStream, ModelInfo, Provider, ProviderDef, ProviderMetadata, DEFAULT_PROVIDER_TIMEOUT_SECS, }; use super::huggingface_auth; @@ -48,7 +48,7 @@ type EndpointParts = (String, String, QueryParams); pub struct HuggingFaceProvider { inner: OpenAiCompatibleProvider, - custom_models: Option>, + custom_models: Option>, dynamic_models: Option, } @@ -75,7 +75,7 @@ impl HuggingFaceProvider { config: DeclarativeProviderConfig, tls_config: Option, ) -> Result { - let custom_models = static_model_names(&config); + let custom_models = static_models(&config); if config.dynamic_models == Some(false) && custom_models.is_none() { return Err(anyhow!( "Provider '{}' has dynamic_models: false but no static models listed; \ @@ -133,10 +133,25 @@ impl Provider for HuggingFaceProvider { self.inner.get_name() } + async fn get_context_limit(&self, model: &str, override_limit: Option) -> usize { + let configured_limits = self + .custom_models + .iter() + .flatten() + .filter_map(|model| model.context_limit.map(|limit| (model.name.clone(), limit))); + goose_providers::context_limit::ContextLimitResolver::new(self.get_name()) + .with_configured_limits(configured_limits) + .resolve(model, override_limit, || async { Ok(None) }) + .await + } + async fn fetch_supported_models(&self) -> Result, ProviderError> { if let Some(custom_models) = &self.custom_models { if self.dynamic_models == Some(false) { - return Ok(custom_models.clone()); + return Ok(custom_models + .iter() + .map(|model| model.name.clone()) + .collect()); } match self.inner.fetch_supported_models().await { @@ -147,7 +162,10 @@ impl Provider for HuggingFaceProvider { self.inner.get_name(), e ); - return Ok(custom_models.clone()); + return Ok(custom_models + .iter() + .map(|model| model.name.clone()) + .collect()); } Err(e) => return Err(e), } @@ -247,14 +265,8 @@ fn configured_api_key(config: &DeclarativeProviderConfig) -> Result Option> { - (!config.models.is_empty()).then(|| { - config - .models - .iter() - .map(|model| model.name.clone()) - .collect() - }) +fn static_models(config: &DeclarativeProviderConfig) -> Option> { + (!config.models.is_empty()).then(|| config.models.clone()) } fn custom_auth_method(config: &DeclarativeProviderConfig) -> Result { @@ -439,8 +451,8 @@ mod tests { config.requires_auth = false; config.dynamic_models = Some(false); config.models = vec![ - ModelInfo::new("static-a".to_string(), 128000), - ModelInfo::new("static-b".to_string(), 128000), + ModelInfo::new("static-a").with_context_limit(128000), + ModelInfo::new("static-b").with_context_limit(128000), ]; let provider = HuggingFaceProvider::from_custom_config(config, None).unwrap(); @@ -449,6 +461,7 @@ mod tests { provider.fetch_supported_models().await.unwrap(), vec!["static-a".to_string(), "static-b".to_string()] ); + assert_eq!(provider.get_context_limit("static-a", None).await, 128_000); } #[test] diff --git a/crates/goose/src/providers/init.rs b/crates/goose/src/providers/init.rs index 91899aef3..5d9286a50 100644 --- a/crates/goose/src/providers/init.rs +++ b/crates/goose/src/providers/init.rs @@ -312,7 +312,6 @@ pub async fn create_with_named_model( mod tests { use super::*; use crate::config::paths::Paths; - use goose_providers::model::ModelConfig; use std::fs; #[tokio::test] @@ -445,24 +444,23 @@ mod tests { let inf_entry = get_from_registry("custom_inf") .await .expect("custom_inf entry should exist"); - let inf_config = inf_entry - .normalize_model_config( - crate::model_config::model_config_from_user_config("custom_inf", "kimi-k2.5") - .expect("custom_inf model config should resolve"), - ) - .expect("custom_inf model config should normalize"); - assert_eq!(inf_config.context_limit, Some(256_000)); + let provider = inf_entry + .create(vec![]) + .await + .expect("custom_inf provider should be created"); + assert_eq!(provider.get_context_limit("kimi-k2.5", None).await, 256_000); let zero_entry = get_from_registry("custom_zero") .await .expect("custom_zero entry should exist"); - let zero_config = zero_entry - .normalize_model_config( - crate::model_config::model_config_from_user_config("custom_zero", "zero-model") - .expect("custom_zero model config should resolve"), - ) - .expect("custom_zero model config should normalize"); - assert_eq!(zero_config.context_limit, None); + let zero_provider = zero_entry + .create(vec![]) + .await + .expect("custom_zero provider should be created"); + assert_eq!( + zero_provider.get_context_limit("zero-model", None).await, + goose_providers::model::DEFAULT_CONTEXT_LIMIT + ); std::env::remove_var("GOOSE_PATH_ROOT"); } @@ -482,10 +480,16 @@ mod tests { let openai = get_from_registry("openai") .await .expect("openai provider should be registered"); - let unknown = openai - .normalize_model_config(ModelConfig::new("totally-unknown-model")) - .expect("unknown model config should normalize"); - assert_eq!(unknown.context_limit(), 1_000_000); + let openai_provider = openai + .create(vec![]) + .await + .expect("openai provider should be created"); + assert_eq!( + openai_provider + .get_context_limit("totally-unknown-model", Some(1_000_000)) + .await, + 1_000_000 + ); let temp_dir = tempfile::tempdir().expect("tempdir should be created"); std::env::set_var("GOOSE_PATH_ROOT", temp_dir.path()); @@ -515,10 +519,16 @@ mod tests { let inf_entry = get_from_registry("custom_inf") .await .expect("custom_inf entry should exist"); - let inf_config = inf_entry - .normalize_model_config(ModelConfig::new("kimi-k2.5")) - .expect("custom_inf model config should normalize"); - assert_eq!(inf_config.context_limit(), 1_000_000); + let inf_provider = inf_entry + .create(vec![]) + .await + .expect("custom_inf provider should be created"); + assert_eq!( + inf_provider + .get_context_limit("kimi-k2.5", Some(1_000_000)) + .await, + 1_000_000 + ); std::env::remove_var("GOOSE_PATH_ROOT"); } diff --git a/crates/goose/src/providers/inventory/mod.rs b/crates/goose/src/providers/inventory/mod.rs index 75da082f2..66b9b6e87 100644 --- a/crates/goose/src/providers/inventory/mod.rs +++ b/crates/goose/src/providers/inventory/mod.rs @@ -1113,7 +1113,7 @@ fn configured_models_to_inventory( let mut result: Vec = Vec::new(); let mut seen_names: HashSet = HashSet::new(); for model in models { - let enriched = enriched_model(provider_family, &model.name, Some(model.context_limit)); + let enriched = enriched_model(provider_family, &model.name, model.context_limit); if seen_names.insert(enriched.name.clone()) { result.push(enriched); } @@ -1320,8 +1320,10 @@ mod tests { #[test] fn configured_models_use_canonical_enrichment() { - let models = - configured_models_to_inventory("anthropic", &[ModelInfo::new("claude-sonnet-4-5", 0)]); + let models = configured_models_to_inventory( + "anthropic", + &[ModelInfo::new("claude-sonnet-4-5").with_context_limit(0)], + ); assert_eq!(models.len(), 1); assert!(models[0].name.contains("Claude")); @@ -1349,7 +1351,7 @@ mod tests { #[test] fn inventory_uses_configured_models_before_first_successful_refresh() { - let configured_models = [ModelInfo::new("claude-sonnet-4-5", 0)]; + let configured_models = [ModelInfo::new("claude-sonnet-4-5").with_context_limit(0)]; let snapshot = InventorySnapshot { models: vec![], last_updated_at: None, @@ -1366,7 +1368,7 @@ mod tests { #[test] fn inventory_preserves_empty_models_after_successful_refresh() { - let configured_models = [ModelInfo::new("claude-sonnet-4-5", 0)]; + let configured_models = [ModelInfo::new("claude-sonnet-4-5").with_context_limit(0)]; let snapshot = InventorySnapshot { models: vec![], last_updated_at: Some(Utc::now()), @@ -1382,7 +1384,7 @@ mod tests { #[test] fn inventory_ignores_stale_snapshots_for_static_providers() { - let configured_models = [ModelInfo::new("gpt-5.6", 0)]; + let configured_models = [ModelInfo::new("gpt-5.6").with_context_limit(0)]; let snapshot = InventorySnapshot { models: vec![InventoryModel { id: "gpt-5.5".to_string(), diff --git a/crates/goose/src/providers/litellm.rs b/crates/goose/src/providers/litellm.rs index 2b0d3d6ea..285f220e9 100644 --- a/crates/goose/src/providers/litellm.rs +++ b/crates/goose/src/providers/litellm.rs @@ -7,6 +7,7 @@ use goose_providers::errors::ProviderError; use goose_providers::images::ImageFormat; use serde_json::Value; use std::collections::HashMap; +use std::time::{Duration, Instant}; use super::api_client::{ApiClient, AuthMethod}; use super::base::{ @@ -26,6 +27,15 @@ const LITELLM_DEFAULT_HOST: &str = "http://localhost:4000"; pub const LITELLM_DEFAULT_MODEL: &str = "gpt-4o-mini"; pub const LITELLM_DOC_URL: &str = "https://docs.litellm.ai/docs/"; +const MODEL_INFO_DISCOVERY_TIMEOUT: Duration = Duration::from_secs(5); +const MODEL_INFO_FAILURE_TTL: Duration = Duration::from_secs(60); + +#[derive(Debug)] +enum CachedModelInfo { + Success(Vec), + Failure(Instant), +} + #[derive(Debug, serde::Serialize)] pub struct LiteLLMProvider { #[serde(skip)] @@ -34,7 +44,7 @@ pub struct LiteLLMProvider { #[serde(skip)] name: String, #[serde(skip)] - cached_model_info: tokio::sync::OnceCell>, + cached_model_info: tokio::sync::Mutex>, } impl LiteLLMProvider { @@ -88,15 +98,41 @@ impl LiteLLMProvider { api_client, base_path, name: LITELLM_PROVIDER_NAME.to_string(), - cached_model_info: tokio::sync::OnceCell::new(), + cached_model_info: tokio::sync::Mutex::new(None), }) } - async fn get_or_fetch_models(&self) -> Result<&[ModelInfo], ProviderError> { - self.cached_model_info - .get_or_try_init(|| self.fetch_models_from_api()) - .await - .map(|v| v.as_slice()) + async fn get_or_fetch_models(&self) -> Result, ProviderError> { + let mut cache = self.cached_model_info.lock().await; + match cache.as_ref() { + Some(CachedModelInfo::Success(models)) => return Ok(models.clone()), + Some(CachedModelInfo::Failure(fetched_at)) + if fetched_at.elapsed() < MODEL_INFO_FAILURE_TTL => + { + return Err(ProviderError::RequestFailed( + "LiteLLM model metadata is unavailable".to_string(), + )); + } + Some(CachedModelInfo::Failure(_)) | None => {} + } + + match tokio::time::timeout(MODEL_INFO_DISCOVERY_TIMEOUT, self.fetch_models_from_api()).await + { + Ok(Ok(models)) => { + *cache = Some(CachedModelInfo::Success(models.clone())); + Ok(models) + } + Ok(Err(error)) => { + *cache = Some(CachedModelInfo::Failure(Instant::now())); + Err(error) + } + Err(_) => { + *cache = Some(CachedModelInfo::Failure(Instant::now())); + Err(ProviderError::RequestFailed( + "LiteLLM model metadata discovery timed out".to_string(), + )) + } + } } async fn fetch_models_from_api(&self) -> Result, ProviderError> { @@ -125,11 +161,13 @@ impl LiteLLMProvider { } let model_info = &model_data["model_info"]; - let context_length = - model_info["max_input_tokens"].as_u64().unwrap_or(128000) as usize; + let context_length = model_info["max_input_tokens"] + .as_u64() + .map(|limit| limit as usize); let supports_cache_control = model_info["supports_prompt_caching"].as_bool(); - let mut model_info_obj = ModelInfo::new(model_name, context_length); + let mut model_info_obj = + ModelInfo::new(model_name).with_optional_context_limit(context_length); model_info_obj.supports_cache_control = supports_cache_control; models.push(model_info_obj); } @@ -231,25 +269,17 @@ impl Provider for LiteLLMProvider { &self.name } - async fn get_context_limit(&self, model_config: &ModelConfig) -> Result { - if let Some(limit) = model_config.context_limit { - return Ok(limit); - } - - // The cache is populated lazily by the first stream() call (via - // supports_cache_control). On turn 1 this will be None and we fall - // back to DEFAULT_CONTEXT_LIMIT, which is fine — the conversation is - // too small to trigger compaction. From turn 2 onward the real limit - // from /model/info is used. - if let Some(models) = self.cached_model_info.get() { - if let Some(info) = models.iter().find(|m| m.name == model_config.model_name) { - if info.context_limit > 0 { - return Ok(info.context_limit); - } - } - } - - Ok(model_config.context_limit()) + async fn get_context_limit(&self, model: &str, override_limit: Option) -> usize { + goose_providers::context_limit::ContextLimitResolver::new(&self.name) + .resolve(model, override_limit, || async { + Ok(self + .get_or_fetch_models() + .await? + .iter() + .find(|info| info.name == model) + .and_then(|info| info.context_limit)) + }) + .await } async fn stream( @@ -311,3 +341,85 @@ fn parse_custom_headers(headers_str: String) -> HashMap { } headers } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn context_limit_negative_caches_failed_model_info() { + let provider = LiteLLMProvider { + api_client: ApiClient::new_with_tls( + "http://127.0.0.1:1".to_string(), + AuthMethod::NoAuth, + None, + ) + .unwrap(), + base_path: "v1/chat/completions".to_string(), + name: LITELLM_PROVIDER_NAME.to_string(), + cached_model_info: tokio::sync::Mutex::new(None), + }; + + assert_eq!( + provider.get_context_limit("unknown-model", None).await, + goose_providers::model::DEFAULT_CONTEXT_LIMIT + ); + assert!(matches!( + provider.cached_model_info.lock().await.as_ref(), + Some(CachedModelInfo::Failure(_)) + )); + assert_eq!( + provider.get_context_limit("unknown-model", None).await, + goose_providers::model::DEFAULT_CONTEXT_LIMIT + ); + } + + #[tokio::test] + async fn expired_failure_allows_model_info_retry() { + let provider = LiteLLMProvider { + api_client: ApiClient::new_with_tls( + "http://127.0.0.1:1".to_string(), + AuthMethod::NoAuth, + None, + ) + .unwrap(), + base_path: "v1/chat/completions".to_string(), + name: LITELLM_PROVIDER_NAME.to_string(), + cached_model_info: tokio::sync::Mutex::new(Some(CachedModelInfo::Failure( + Instant::now() - MODEL_INFO_FAILURE_TTL, + ))), + }; + + assert!(provider.get_or_fetch_models().await.is_err()); + assert!(matches!( + provider.cached_model_info.lock().await.as_ref(), + Some(CachedModelInfo::Failure(fetched_at)) + if fetched_at.elapsed() < MODEL_INFO_FAILURE_TTL + )); + } + + #[tokio::test] + async fn context_limit_uses_cached_model_info() { + let cached_model_info = + tokio::sync::Mutex::new(Some(CachedModelInfo::Success(vec![ModelInfo::new( + "cached-model", + ) + .with_context_limit(32_000)]))); + let provider = LiteLLMProvider { + api_client: ApiClient::new_with_tls( + "http://127.0.0.1:1".to_string(), + AuthMethod::NoAuth, + None, + ) + .unwrap(), + base_path: "v1/chat/completions".to_string(), + name: LITELLM_PROVIDER_NAME.to_string(), + cached_model_info, + }; + + assert_eq!( + provider.get_context_limit("cached-model", None).await, + 32_000 + ); + } +} diff --git a/crates/goose/src/providers/ollama_cloud.rs b/crates/goose/src/providers/ollama_cloud.rs index 8ac5771b3..7401c6970 100644 --- a/crates/goose/src/providers/ollama_cloud.rs +++ b/crates/goose/src/providers/ollama_cloud.rs @@ -5,7 +5,7 @@ use crate::conversation::message::Message; use anyhow::Result; use futures::future::BoxFuture; use goose_providers::api_client::{ApiClient, AuthMethod, TlsConfig}; -use goose_providers::base::ProviderDescriptor; +use goose_providers::base::{ModelInfo, ProviderDescriptor}; use goose_providers::errors::ProviderError; use goose_providers::model::ModelConfig; use goose_providers::ollama::fetch_ollama_model_names; @@ -13,19 +13,18 @@ use goose_providers::openai::OpenAiProvider; use rmcp::model::Tool; use serde_json::Value; use std::collections::HashMap; -use std::sync::{LazyLock, Mutex}; +use std::sync::Mutex; use tokio::sync::OnceCell; const OLLAMA_CLOUD_PROVIDER_NAME: &str = "ollama_cloud"; - -static SHOW_INFO_CACHE: LazyLock>> = - LazyLock::new(|| Mutex::new(HashMap::new())); +const SHOW_INFO_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); pub struct OllamaCloudProvider { inner: OpenAiProvider, ollama_api_client: ApiClient, model_names: OnceCell>, - custom_models: Option>, + context_limits: Mutex>>, + custom_models: Option>, dynamic_models: Option, } @@ -43,13 +42,7 @@ impl OllamaCloudProvider { crate::providers::openai_def::from_custom_config(config.clone(), tls_config.clone())?; let custom_models = if !config.models.is_empty() { - Some( - config - .models - .iter() - .map(|m| m.name.clone()) - .collect::>(), - ) + Some(config.models.clone()) } else { None }; @@ -68,6 +61,7 @@ impl OllamaCloudProvider { inner, ollama_api_client, model_names: OnceCell::new(), + context_limits: Mutex::new(HashMap::new()), custom_models, dynamic_models: config.dynamic_models, }) @@ -190,7 +184,10 @@ impl Provider for OllamaCloudProvider { async fn fetch_supported_models(&self) -> Result, ProviderError> { if let Some(custom_models) = &self.custom_models { if self.dynamic_models == Some(false) { - return Ok(custom_models.clone()); + return Ok(custom_models + .iter() + .map(|model| model.name.clone()) + .collect()); } match self.get_or_fetch_model_names().await { @@ -200,7 +197,10 @@ impl Provider for OllamaCloudProvider { "Ollama api/tags not available for provider '{}', using static model list", self.inner.get_name(), ); - return Ok(custom_models.clone()); + return Ok(custom_models + .iter() + .map(|model| model.name.clone()) + .collect()); } Err(e) => return Err(e), } @@ -209,31 +209,37 @@ impl Provider for OllamaCloudProvider { self.get_or_fetch_model_names().await } - async fn get_context_limit(&self, model_config: &ModelConfig) -> Result { - if let Some(limit) = model_config.context_limit { - return Ok(limit); - } + async fn get_context_limit(&self, model: &str, override_limit: Option) -> usize { + let configured_limits = self + .custom_models + .iter() + .flatten() + .filter_map(|model| model.context_limit.map(|limit| (model.name.clone(), limit))); + goose_providers::context_limit::ContextLimitResolver::new(self.get_name()) + .with_configured_limits(configured_limits) + .resolve(model, override_limit, || async { + if let Some(cached) = self + .context_limits + .lock() + .ok() + .and_then(|cache| cache.get(model).copied()) + { + return Ok(cached); + } - if let Some(cached) = SHOW_INFO_CACHE - .lock() - .ok() - .and_then(|cache| cache.get(&model_config.model_name).copied()) - { - return Ok(cached); - } - - if let Some(limit) = self - .fetch_context_limit_from_show(&model_config.model_name) + let limit = tokio::time::timeout( + SHOW_INFO_TIMEOUT, + self.fetch_context_limit_from_show(model), + ) + .await + .ok() + .flatten(); + if let Ok(mut cache) = self.context_limits.lock() { + cache.insert(model.to_string(), limit); + } + Ok(limit) + }) .await - { - if let Ok(mut cache) = SHOW_INFO_CACHE.lock() { - cache.insert(model_config.model_name.clone(), limit); - } - - return Ok(limit); - } - - Ok(model_config.context_limit()) } } @@ -299,7 +305,7 @@ mod tests { let provider = build_provider( server.uri(), Some(false), - vec![ModelInfo::new("static-model", 4096)], + vec![ModelInfo::new("static-model").with_context_limit(4096)], ); assert_eq!( @@ -314,7 +320,7 @@ mod tests { let provider = build_provider( server.uri(), None, - vec![ModelInfo::new("static-model", 4096)], + vec![ModelInfo::new("static-model").with_context_limit(4096)], ); assert_eq!( @@ -329,7 +335,7 @@ mod tests { let provider = build_provider( server.uri(), Some(true), - vec![ModelInfo::new("static-model", 4096)], + vec![ModelInfo::new("static-model").with_context_limit(4096)], ); let models = provider.fetch_supported_models().await.unwrap(); @@ -351,7 +357,9 @@ mod tests { let provider = build_provider(server.uri(), Some(true), vec![]); let model_config = ModelConfig::new("gemma3:4b"); - let limit = provider.get_context_limit(&model_config).await.unwrap(); + let limit = provider + .get_context_limit(&model_config.model_name, None) + .await; assert_eq!(limit, 131072); } @@ -361,18 +369,25 @@ mod tests { let provider = build_provider(server.uri(), Some(true), vec![]); let model_config = ModelConfig::new("qwen3-coder:480b"); - let limit = provider.get_context_limit(&model_config).await.unwrap(); + let limit = provider + .get_context_limit(&model_config.model_name, None) + .await; assert_eq!(limit, 262144); } #[tokio::test] - async fn get_context_limit_falls_back_on_missing_model_info() { + async fn get_context_limit_caches_missing_model_info() { let server = mock_show_server_no_model_info().await; let provider = build_provider(server.uri(), Some(true), vec![]); - let model_config = ModelConfig::new("unknown-model").with_context_limit(Some(8000)); - let limit = provider.get_context_limit(&model_config).await.unwrap(); - assert_eq!(limit, 8000); + for _ in 0..2 { + assert_eq!( + provider.get_context_limit("unknown-model", None).await, + goose_providers::model::DEFAULT_CONTEXT_LIMIT + ); + } + + assert_eq!(server.received_requests().await.unwrap().len(), 1); } fn build_provider( diff --git a/crates/goose/src/providers/provider_registry.rs b/crates/goose/src/providers/provider_registry.rs index 774dbb4a2..badb93f60 100644 --- a/crates/goose/src/providers/provider_registry.rs +++ b/crates/goose/src/providers/provider_registry.rs @@ -55,26 +55,8 @@ impl ProviderEntry { (self.inventory_configured)() } - /// Apply provider-specific normalization to a model config: materialize - /// global defaults and backfill `context_limit` from the provider's known - /// models when the canonical registry didn't already resolve one. Used by - /// the agent/session layer to resolve effective limits (e.g. for custom - /// providers that declare explicit context limits in their config). - pub fn normalize_model_config(&self, mut model: ModelConfig) -> Result { - model = crate::model_config::materialize_model_config(&self.metadata.name, model)?; - - if model.context_limit.is_none() { - if let Some(info) = self - .metadata - .known_models - .iter() - .find(|m| m.name.eq_ignore_ascii_case(&model.model_name) && m.context_limit > 0) - { - model.context_limit = Some(info.context_limit); - } - } - - Ok(model) + pub fn normalize_model_config(&self, model: ModelConfig) -> Result { + crate::model_config::materialize_model_config(&self.metadata.name, model) } pub async fn create_with_default_model( @@ -386,7 +368,7 @@ mod tests { description: None, api_key_env: String::new(), base_url: "https://router.huggingface.co/v1".to_string(), - models: vec![ModelInfo::new("test-model", 128_000)], + models: vec![ModelInfo::new("test-model").with_context_limit(128_000)], headers: None, timeout_seconds: None, supports_streaming: Some(true), diff --git a/crates/goose/src/providers/xai.rs b/crates/goose/src/providers/xai.rs index 4cab5df72..c775a471a 100644 --- a/crates/goose/src/providers/xai.rs +++ b/crates/goose/src/providers/xai.rs @@ -107,7 +107,7 @@ mod tests { .iter() .find(|model| model.name == "grok-4.5") .expect("grok-4.5 should be a known xAI model"); - assert_eq!(grok_4_5.context_limit, 500_000); + assert_eq!(grok_4_5.context_limit, Some(500_000)); assert!(grok_4_5.reasoning); let grok_4_20_non_reasoning = metadata diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index 3593e5589..1c3bbc56a 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -2806,6 +2806,18 @@ mod tests { assert_eq!(config.model_name, "gpt-5-high"); assert_eq!(config.thinking_effort(), None); + assert!(!json.contains("context_limit")); + } + + #[test] + fn legacy_session_context_limit_is_ignored() { + let config = deserialize_session_model_config( + Some("openai"), + r#"{"model_name":"gpt-4o","context_limit":64000,"temperature":null,"max_tokens":null,"toolshim":false,"toolshim_model":null}"#, + ) + .unwrap(); + + assert_eq!(config.context_limit, None); } #[test] diff --git a/crates/goose/tests/acp_custom_provider_methods_test.rs b/crates/goose/tests/acp_custom_provider_methods_test.rs index a73849d53..bdbb03ddd 100644 --- a/crates/goose/tests/acp_custom_provider_methods_test.rs +++ b/crates/goose/tests/acp_custom_provider_methods_test.rs @@ -25,25 +25,22 @@ fn write_secrets(config_dir: &std::path::Path, contents: &str) { #[test] #[serial] fn acp_catalog_and_custom_provider_methods_use_core_provider_store() { - let root = tempfile::tempdir().unwrap(); - let root_path = root.path().to_string_lossy().to_string(); let _env = env_lock::lock_env([ - ("GOOSE_PATH_ROOT", Some(root_path.as_str())), ("GOOSE_DISABLE_KEYRING", Some("1")), ("XAI_API_KEY", None), ("XAI_HOST", None), ("CUSTOM_STARK_ACP_PROVIDER_API_KEY", None), ]); - let config_dir = Paths::config_dir(); - write_config( - &config_dir, - "GOOSE_MODEL: gpt-4o\nGOOSE_PROVIDER: openai\nGOOSE_DISABLE_KEYRING: true\nXAI_HOST: https://api.x.ai/v1\n", - ); - write_secrets(&config_dir, "XAI_API_KEY: xai-configured-key\n"); - Config::global().invalidate_secrets_cache(); - run_test(async move { + let config_dir = Paths::config_dir(); + write_config( + &config_dir, + "GOOSE_MODEL: gpt-4o\nGOOSE_PROVIDER: openai\nGOOSE_DISABLE_KEYRING: true\nXAI_HOST: https://api.x.ai/v1\n", + ); + write_secrets(&config_dir, "XAI_API_KEY: xai-configured-key\n"); + Config::global().invalidate_secrets_cache(); + let openai = common_tests::fixtures::OpenAiFixture::new( vec![], Arc::new(EnforceSessionId::default()), diff --git a/crates/goose/tests/acp_fixtures/mod.rs b/crates/goose/tests/acp_fixtures/mod.rs index 28a4f1b2e..ec7e59bdf 100644 --- a/crates/goose/tests/acp_fixtures/mod.rs +++ b/crates/goose/tests/acp_fixtures/mod.rs @@ -480,7 +480,6 @@ pub fn to_notifications(updates: &[SessionUpdate]) -> Vec { } } SessionUpdate::Plan(_) => out.push(Notification::Plan), - SessionUpdate::AvailableCommandsUpdate(_) => out.push(Notification::AvailableCommands), SessionUpdate::CurrentModeUpdate(_) => out.push(Notification::CurrentMode), SessionUpdate::ConfigOptionUpdate(_) => out.push(Notification::ConfigOption), SessionUpdate::SessionInfoUpdate(update) => { diff --git a/documentation/docs/guides/environment-variables.md b/documentation/docs/guides/environment-variables.md index 7575bc7f5..c2735f71f 100644 --- a/documentation/docs/guides/environment-variables.md +++ b/documentation/docs/guides/environment-variables.md @@ -250,8 +250,7 @@ These variables allow you to override the default context window size (token lim | Variable | Purpose | Values | Default | |----------|---------|---------|---------| | `GOOSE_CONTEXT_LIMIT` | Override context limit for the main model | Integer (number of tokens) | Model-specific default or 128,000 | -| `GOOSE_INPUT_LIMIT` | Override input prompt limit for ollama requests (maps to `num_ctx`) | Integer (number of tokens) | Falls back to `GOOSE_CONTEXT_LIMIT` or model default | -| `GOOSE_PLANNER_CONTEXT_LIMIT` | Override context limit for the [planner model](/docs/guides/context-engineering/creating-plans) | Integer (number of tokens) | Falls back to `GOOSE_CONTEXT_LIMIT` or model default | +| `GOOSE_INPUT_LIMIT` | Override input prompt limit for ollama requests (maps to `num_ctx`) | Integer (number of tokens) | Unset; Ollama uses its model default | **Examples** @@ -260,9 +259,6 @@ These variables allow you to override the default context window size (token lim export GOOSE_CONTEXT_LIMIT=200000 # Override ollama input prompt limit export GOOSE_INPUT_LIMIT=32000 - -# Set context limit for planner -export GOOSE_PLANNER_CONTEXT_LIMIT=1000000 ``` For more details and examples, see [Model Context Limit Overrides](/docs/guides/sessions/smart-context-management#model-context-limit-overrides). diff --git a/documentation/docs/guides/sessions/smart-context-management.md b/documentation/docs/guides/sessions/smart-context-management.md index a8637e86d..6e4ba7013 100644 --- a/documentation/docs/guides/sessions/smart-context-management.md +++ b/documentation/docs/guides/sessions/smart-context-management.md @@ -295,7 +295,6 @@ Context limits are automatically detected based on your model name, but goose pr | Model | Description | Best For | Setting | |-------|-------------|----------|---------| | **Main** | Set context limit for the main model (also serves as fallback for other models) | LiteLLM proxies, custom models with non-standard names | `GOOSE_CONTEXT_LIMIT` | -| **Planner** | Set context for [planner models](/docs/guides/context-engineering/creating-plans) | Large planning tasks requiring extensive context | `GOOSE_PLANNER_CONTEXT_LIMIT` | :::info This setting only affects the displayed token usage and progress indicators. Actual context management is handled by your LLM, so you may experience more or less usage than the limit you set, regardless of what the display shows. @@ -310,10 +309,10 @@ This feature is particularly useful with: goose resolves context limits with the following precedence (highest to lowest): -1. Explicit context_limit in model configuration (if set programmatically) -2. Specific environment variable (e.g., `GOOSE_PLANNER_CONTEXT_LIMIT`) -3. Global environment variable (`GOOSE_CONTEXT_LIMIT`) -4. Model-specific default based on name pattern matching +1. Global environment variable (`GOOSE_CONTEXT_LIMIT`) +2. Explicit declarative or custom provider model configuration +3. Provider runtime discovery +4. Canonical model metadata 5. Global default (128,000 tokens) **Configuration** @@ -348,22 +347,6 @@ export GOOSE_MODEL="my-custom-gpt4-proxy" export GOOSE_CONTEXT_LIMIT=200000 # Override the 32k default ``` -2. Planner setup with a different context limit - -```bash -# Set a larger context window for planning -export GOOSE_PLANNER_MODEL="claude-opus-custom" -export GOOSE_PLANNER_CONTEXT_LIMIT=500000 -``` - -3. Planner with large context - -```bash -# Large context for complex planning -export GOOSE_PLANNER_MODEL="gpt-4-custom" -export GOOSE_PLANNER_CONTEXT_LIMIT=1000000 -``` - ## Credit Balance Monitoring goose monitors your API provider balance and warns you when credits are running low or exhausted. When this happens, you'll see an **Insufficient Credits** notification. diff --git a/documentation/src/data/gdk-api.json b/documentation/src/data/gdk-api.json index 7e05a287c..48d24b79e 100644 --- a/documentation/src/data/gdk-api.json +++ b/documentation/src/data/gdk-api.json @@ -862,6 +862,21 @@ "throws": null, "isAsync": false }, + { + "name": "context_limit", + "docs": "", + "params": [ + { + "name": "model", + "type": "ProviderModelConfig", + "default": null, + "docs": "" + } + ], + "returns": "u64", + "throws": "GooseError", + "isAsync": true + }, { "name": "stream", "docs": "", diff --git a/ui/desktop/src/acp/__tests__/gooseSessionNotifications.test.ts b/ui/desktop/src/acp/__tests__/gooseSessionNotifications.test.ts index be651ae8a..2d974b5e1 100644 --- a/ui/desktop/src/acp/__tests__/gooseSessionNotifications.test.ts +++ b/ui/desktop/src/acp/__tests__/gooseSessionNotifications.test.ts @@ -160,6 +160,7 @@ describe('applyGooseSessionNotification', () => { type: 'tokenState', tokenState: { totalTokens: 42, + contextLimit: 200, accumulatedInputTokens: 10, accumulatedOutputTokens: 15, accumulatedTotalTokens: 25, diff --git a/ui/desktop/src/acp/__tests__/providers.test.ts b/ui/desktop/src/acp/__tests__/providers.test.ts index 48feb4024..4cb0c45d4 100644 --- a/ui/desktop/src/acp/__tests__/providers.test.ts +++ b/ui/desktop/src/acp/__tests__/providers.test.ts @@ -219,7 +219,7 @@ describe('ACP providers', () => { expect(result.connectionChecked).toBe(true); expect(result.provider.metadata.known_models).toEqual([ - { name: 'claude-sonnet', context_limit: 0, reasoning: undefined }, + { name: 'claude-sonnet', context_limit: undefined, reasoning: undefined }, ]); }); @@ -268,7 +268,7 @@ describe('ACP providers', () => { }); expect(enabled.is_configured).toBe(true); expect(enabled.metadata.known_models).toEqual([ - { name: 'claude-sonnet', context_limit: 0, reasoning: undefined }, + { name: 'claude-sonnet', context_limit: undefined, reasoning: undefined }, ]); }); diff --git a/ui/desktop/src/acp/__tests__/sessionNotificationAdapter.test.ts b/ui/desktop/src/acp/__tests__/sessionNotificationAdapter.test.ts index af98c82dc..f3dd33158 100644 --- a/ui/desktop/src/acp/__tests__/sessionNotificationAdapter.test.ts +++ b/ui/desktop/src/acp/__tests__/sessionNotificationAdapter.test.ts @@ -872,6 +872,7 @@ describe('createAcpSessionNotificationAdapter', () => { type: 'tokenState', tokenState: { totalTokens: 42, + contextLimit: 200, accumulatedInputTokens: 10, accumulatedOutputTokens: 15, accumulatedTotalTokens: 25, diff --git a/ui/desktop/src/acp/adapter/gooseSessionNotifications.ts b/ui/desktop/src/acp/adapter/gooseSessionNotifications.ts index b2f11f471..0bcc3cc00 100644 --- a/ui/desktop/src/acp/adapter/gooseSessionNotifications.ts +++ b/ui/desktop/src/acp/adapter/gooseSessionNotifications.ts @@ -15,6 +15,7 @@ export function applyGooseSessionNotification( type: 'tokenState', tokenState: { totalTokens: update.used, + contextLimit: update.contextLimit, accumulatedInputTokens: update.accumulatedInputTokens, accumulatedOutputTokens: update.accumulatedOutputTokens, accumulatedTotalTokens: update.accumulatedInputTokens + update.accumulatedOutputTokens, diff --git a/ui/desktop/src/acp/providers.ts b/ui/desktop/src/acp/providers.ts index 87533b054..47e841ee9 100644 --- a/ui/desktop/src/acp/providers.ts +++ b/ui/desktop/src/acp/providers.ts @@ -71,7 +71,7 @@ function providerEntryToDetails(entry: ProviderInventoryEntryDto): ProviderDetai })), known_models: entry.models.map((model) => ({ name: model.id, - context_limit: model.contextLimit ?? 0, + context_limit: model.contextLimit ?? undefined, reasoning: model.reasoning ?? undefined, })), setup_steps: entry.setupSteps, diff --git a/ui/desktop/src/components/BaseChat.tsx b/ui/desktop/src/components/BaseChat.tsx index 500d7e3c1..48319a3df 100644 --- a/ui/desktop/src/components/BaseChat.tsx +++ b/ui/desktop/src/components/BaseChat.tsx @@ -522,6 +522,7 @@ export default function BaseChat({ initialValue={initialPrompt} setView={setView} totalTokens={tokenState?.totalTokens ?? session?.usage?.total_tokens ?? undefined} + contextLimit={tokenState?.contextLimit} accumulatedInputTokens={ tokenState?.accumulatedInputTokens ?? session?.accumulated_usage?.input_tokens ?? diff --git a/ui/desktop/src/components/ChatInput.tsx b/ui/desktop/src/components/ChatInput.tsx index 83ddea4f9..2f9831360 100644 --- a/ui/desktop/src/components/ChatInput.tsx +++ b/ui/desktop/src/components/ChatInput.tsx @@ -82,7 +82,7 @@ const removeQueuedMessage = (messages: QueuedMessage[], messageId: string): Queu const MAX_IMAGES_PER_MESSAGE = 10; -const TOKEN_LIMIT_DEFAULT = 128000; // fallback for custom models that the backend doesn't know about +const TOKEN_LIMIT_DEFAULT = 128000; // used before a session has a backend-resolved limit const getContextAlertType = (totalTokens: number, tokenLimit: number): AlertType => { const percentage = tokenLimit ? (totalTokens / tokenLimit) * 100 : 0; @@ -172,6 +172,7 @@ interface ChatInputProps { onFilesProcessed?: () => void; setView: (view: View) => void; totalTokens?: number; + contextLimit?: number; accumulatedInputTokens?: number; accumulatedOutputTokens?: number; accumulatedCost?: number | null; @@ -207,6 +208,7 @@ export default function ChatInput({ onFilesProcessed, setView, totalTokens, + contextLimit, accumulatedInputTokens, accumulatedOutputTokens, accumulatedCost, @@ -581,7 +583,11 @@ export default function ChatInput({ // Load providers and get current model's token limit const loadProviderDetails = async () => { try { - // Reset token limit loaded state + if (sessionId) { + setTokenLimit(0); + setIsTokenLimitLoaded(false); + return; + } setIsTokenLimitLoaded(false); // Use effective model/provider (includes overrides from in-session model changes), @@ -642,7 +648,20 @@ export default function ChatInput({ useEffect(() => { loadProviderDetails(); // eslint-disable-next-line react-hooks/exhaustive-deps - }, [effectiveModel, effectiveProvider, configModel, configProvider]); + }, [effectiveModel, effectiveProvider, configModel, configProvider, sessionId]); + + useEffect(() => { + if (contextLimit === undefined) { + if (sessionId) { + setTokenLimit(0); + setIsTokenLimitLoaded(false); + } + return; + } + + setTokenLimit(contextLimit); + setIsTokenLimitLoaded(true); + }, [contextLimit, sessionId]); // Handle token usage alerts useEffect(() => { diff --git a/ui/desktop/src/components/settings/providers/modal/subcomponents/forms/CustomProviderForm.tsx b/ui/desktop/src/components/settings/providers/modal/subcomponents/forms/CustomProviderForm.tsx index ea2ba77b3..5edfc6ce5 100644 --- a/ui/desktop/src/components/settings/providers/modal/subcomponents/forms/CustomProviderForm.tsx +++ b/ui/desktop/src/components/settings/providers/modal/subcomponents/forms/CustomProviderForm.tsx @@ -459,8 +459,8 @@ export default function CustomProviderForm({ const modelList = models .split(',') - .map((m) => m.trim()) - .filter((m) => m); + .map((name) => name.trim()) + .filter(Boolean); let allHeaders = [...headers]; diff --git a/ui/desktop/src/types/chat.ts b/ui/desktop/src/types/chat.ts index 32e007c74..6eed8824f 100644 --- a/ui/desktop/src/types/chat.ts +++ b/ui/desktop/src/types/chat.ts @@ -13,6 +13,7 @@ export type TokenState = { inputTokens: number; outputTokens: number; totalTokens: number; + contextLimit?: number; }; export interface ChatType { diff --git a/ui/desktop/src/types/providers.ts b/ui/desktop/src/types/providers.ts index 28d59fd76..912bb5708 100644 --- a/ui/desktop/src/types/providers.ts +++ b/ui/desktop/src/types/providers.ts @@ -13,7 +13,7 @@ export type ConfigKey = { }; export type ModelInfo = { - context_limit: number; + context_limit?: number | null; currency?: string | null; input_token_cost?: number | null; name: string;