From 73fd41d7c5350fd17cba699a6e68c56a948b8008 Mon Sep 17 00:00:00 2001 From: Jack Amadeo Date: Thu, 18 Jun 2026 14:29:58 -0400 Subject: [PATCH] move the provider trait into goose-providers (#9860) --- Cargo.lock | 3 + crates/goose-providers/Cargo.toml | 4 + crates/goose-providers/src/base.rs | 562 +++++++++++- .../src}/goose_mode.rs | 0 crates/goose-providers/src/lib.rs | 3 + .../src/permission.rs} | 0 .../src}/retry.rs | 14 +- crates/goose-server/src/openapi.rs | 5 +- crates/goose/src/config/mod.rs | 3 +- crates/goose/src/permission/mod.rs | 6 +- crates/goose/src/providers/base.rs | 804 +----------------- crates/goose/src/providers/mod.rs | 4 +- crates/goose/src/session/mod.rs | 1 + crates/goose/src/session/session_manager.rs | 25 +- crates/goose/src/session/session_naming.rs | 249 ++++++ 15 files changed, 855 insertions(+), 828 deletions(-) rename crates/{goose/src/config => goose-providers/src}/goose_mode.rs (100%) rename crates/{goose/src/permission/permission_confirmation.rs => goose-providers/src/permission.rs} (100%) rename crates/{goose/src/providers => goose-providers/src}/retry.rs (97%) create mode 100644 crates/goose/src/session/session_naming.rs diff --git a/Cargo.lock b/Cargo.lock index 18df6af4b..33523f476 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4984,16 +4984,19 @@ version = "1.38.0" dependencies = [ "anyhow", "async-stream", + "async-trait", "base64 0.22.1", "chrono", "env-lock", "futures", "once_cell", + "rand 0.8.6", "regex", "reqwest 0.13.4", "rmcp", "serde", "serde_json", + "strum 0.28.0", "tempfile", "test-case", "thiserror 2.0.18", diff --git a/crates/goose-providers/Cargo.toml b/crates/goose-providers/Cargo.toml index 54c92d34d..c054ea972 100644 --- a/crates/goose-providers/Cargo.toml +++ b/crates/goose-providers/Cargo.toml @@ -28,6 +28,10 @@ tracing = { workspace = true } unicode-normalization = { version = "0.1.22", default-features = false, features = ["std"] } utoipa = { workspace = true, features = ["chrono"] } uuid = { workspace = true, features = ["v4", "std"] } +async-trait = { workspace = true } +strum = { workspace = true } +tokio = { workspace = true } +rand = { workspace = true } [dev-dependencies] test-case = { workspace = true } diff --git a/crates/goose-providers/src/base.rs b/crates/goose-providers/src/base.rs index e063ddebc..b14579314 100644 --- a/crates/goose-providers/src/base.rs +++ b/crates/goose-providers/src/base.rs @@ -1,17 +1,563 @@ -use std::future::Future; +use async_trait::async_trait; +use futures::Stream; +use rmcp::model::Tool; +use serde::{Deserialize, Serialize}; +use std::pin::Pin; +use utoipa::ToSchema; -pub struct Error; +use crate::{ + canonical::{map_to_canonical_model, CanonicalModelRegistry}, + conversation::{ + message::{Message, MessageContent}, + token_usage::{ProviderUsage, Usage}, + }, + errors::ProviderError, + goose_mode::GooseMode, + model::ModelConfig, + permission::PermissionConfirmation, + retry::RetryConfig, +}; -pub struct Model { +/// Information about a model's capabilities +#[derive(Debug, Clone, Serialize, Deserialize, ToSchema, PartialEq)] +pub struct ModelInfo { + /// The name of the model pub name: String, + /// The underlying model resolved from provider metadata, when the configured model is an alias or endpoint. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub resolved_model: Option, + /// The maximum context length this model supports + pub context_limit: usize, + /// Cost per token for input in USD (optional) + pub input_token_cost: Option, + /// Cost per token for output in USD (optional) + pub output_token_cost: Option, + /// Currency for the costs (default: "$") + pub currency: Option, + /// Whether this model supports cache control + pub supports_cache_control: Option, + /// Whether this model supports reasoning/thinking controls + #[serde(default)] + pub reasoning: bool, } -pub struct StreamingRequest { - pub model: Model, +impl ModelInfo { + /// Create a new ModelInfo with just name and context limit + pub fn new(name: impl Into, context_limit: usize) -> Self { + Self { + name: name.into(), + resolved_model: None, + context_limit, + input_token_cost: None, + output_token_cost: None, + currency: None, + supports_cache_control: None, + reasoning: false, + } + } + + /// Create a new ModelInfo with cost information (per token) + pub fn with_cost( + name: impl Into, + context_limit: usize, + input_cost: f64, + output_cost: f64, + ) -> Self { + Self { + name: name.into(), + resolved_model: None, + context_limit, + input_token_cost: Some(input_cost), + output_token_cost: Some(output_cost), + currency: Some("$".to_string()), + supports_cache_control: None, + reasoning: false, + } + } } -pub struct StreamingResponse; +/// A message stream yields partial text content but complete tool calls, all within the Message object +/// So a message with text will contain potentially just a word of a longer response, but tool calls +/// messages will only be yielded once concatenated. +pub type MessageStream = Pin< + Box, Option), ProviderError>> + Send>, +>; -pub trait Provider { - fn stream(req: StreamingRequest) -> impl Future>; +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum PermissionRouting { + ActionRequired, + Noop, +} + +pub fn model_info_for_provider_model(provider_name: &str, model_name: &str) -> ModelInfo { + let registry = CanonicalModelRegistry::bundled().ok(); + let canonical = registry.as_ref().and_then(|registry| { + let canonical_id = map_to_canonical_model(provider_name, model_name, registry)?; + let (provider, model) = canonical_id.split_once('/')?; + registry.get(provider, model) + }); + + let reasoning = canonical + .as_ref() + .and_then(|model| model.reasoning) + .unwrap_or_else(|| ModelConfig::new_or_fail(model_name).is_reasoning_model()); + + ModelInfo { + name: model_name.to_string(), + resolved_model: None, + context_limit: ModelConfig::new_or_fail(model_name) + .with_canonical_limits(provider_name) + .context_limit(), + input_token_cost: None, + output_token_cost: None, + currency: None, + supports_cache_control: None, + reasoning, + } +} + +/// Collect all chunks from a MessageStream into a single Message and ProviderUsage +pub async fn collect_stream( + mut stream: MessageStream, +) -> Result<(Message, ProviderUsage), ProviderError> { + use futures::StreamExt; + + let mut final_message: Option = None; + let mut final_usage: Option = None; + + while let Some(result) = stream.next().await { + let (msg_opt, usage_opt) = result?; + + if let Some(msg) = msg_opt { + final_message = Some(match final_message { + Some(mut prev) => { + for new_content in msg.content { + match (&mut prev.content.last_mut(), &new_content) { + // Coalesce consecutive text blocks + ( + Some(MessageContent::Text(last_text)), + MessageContent::Text(new_text), + ) => { + last_text.text.push_str(&new_text.text); + } + _ => { + prev.content.push(new_content); + } + } + } + prev + } + None => msg, + }); + } + + if let Some(usage) = usage_opt { + final_usage = Some(usage); + } + } + + match final_message { + Some(msg) => { + let usage = final_usage + .unwrap_or_else(|| ProviderUsage::new("unknown".to_string(), Usage::default())); + Ok((msg, usage)) + } + None => Err(ProviderError::ExecutionError( + "Stream yielded no message".to_string(), + )), + } +} + +/// Base trait for AI providers (OpenAI, Anthropic, etc) +#[async_trait] +pub trait Provider: Send + Sync { + /// Get the name of this provider instance + fn get_name(&self) -> &str; + + /// Primary streaming method that all providers must implement. + /// + /// Note: Do not add `#[instrument]` here — the call sites (`complete` and + /// `stream_response_from_provider`) create the telemetry span so that + /// `session.id` is set once rather than in every provider. + async fn stream( + &self, + model_config: &ModelConfig, + session_id: &str, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result; + + /// Complete with a specific model config. + #[tracing::instrument( + skip(self, model_config, session_id, system, messages, tools), + fields(session.id = %session_id, gen_ai.request.model = %model_config.model_name) + )] + async fn complete( + &self, + model_config: &ModelConfig, + session_id: &str, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result<(Message, ProviderUsage), ProviderError> { + let stream = self + .stream(model_config, session_id, system, messages, tools) + .await?; + collect_stream(stream).await + } + + /// Try fast model first, fall back to regular model on failure. + async fn complete_fast( + &self, + session_id: &str, + system: &str, + messages: &[Message], + tools: &[Tool], + ) -> Result<(Message, ProviderUsage), ProviderError> { + let model_config = self.get_model_config(); + let fast_config = model_config.use_fast_model(); + + let result = self + .complete(&fast_config, session_id, system, messages, tools) + .await; + + match result { + Ok(response) => Ok(response), + Err(e) => { + if fast_config.model_name != model_config.model_name { + tracing::warn!( + "Fast model {} failed with error: {}. Falling back to regular model {}", + fast_config.model_name, + e, + model_config.model_name + ); + self.complete(&model_config, session_id, system, messages, tools) + .await + } else { + Err(e) + } + } + } + } + + /// Get the model config from the provider + fn get_model_config(&self) -> ModelConfig; + + fn retry_config(&self) -> RetryConfig { + RetryConfig::default() + } + + async fn fetch_supported_models(&self) -> Result, ProviderError> { + Ok(vec![]) + } + + async fn fetch_supported_model_info(&self) -> Result, ProviderError> { + Ok(self + .fetch_supported_models() + .await? + .iter() + .map(|model_name| model_info_for_provider_model(self.get_name(), model_name)) + .collect()) + } + + async fn fetch_model_info(&self, model_name: &str) -> Result { + Ok(model_info_for_provider_model(self.get_name(), model_name)) + } + + fn skip_canonical_filtering(&self) -> bool { + false + } + + /// Fetch inventory models filtered by canonical registry and usability. + async fn fetch_recommended_models(&self) -> Result, ProviderError> { + let all_models = self.fetch_supported_models().await?; + + if self.skip_canonical_filtering() { + return Ok(all_models); + } + + let registry = CanonicalModelRegistry::bundled().map_err(|e| { + ProviderError::ExecutionError(format!("Failed to load canonical registry: {}", e)) + })?; + + let provider_name = self.get_name(); + + // Get all text-capable models with their release dates + let mut models_with_dates: Vec<(String, Option)> = all_models + .iter() + .filter_map(|model| { + let canonical_id = map_to_canonical_model(provider_name, model, registry)?; + + let (provider, model_name) = canonical_id.split_once('/')?; + let canonical_model = registry.get(provider, model_name)?; + + if !canonical_model + .modalities + .input + .contains(&crate::canonical::Modality::Text) + { + return None; + } + + if !canonical_model.tool_call && !self.get_model_config().toolshim { + return None; + } + + let release_date = canonical_model.release_date.clone(); + + Some((model.clone(), release_date)) + }) + .collect(); + + // Sort by release date (most recent first), then alphabetically for models without dates + models_with_dates.sort_by(|a, b| match (&a.1, &b.1) { + (Some(date_a), Some(date_b)) => date_b.cmp(date_a), + (Some(_), None) => std::cmp::Ordering::Less, + (None, Some(_)) => std::cmp::Ordering::Greater, + (None, None) => a.0.cmp(&b.0), + }); + + let inventory_models: Vec = models_with_dates + .into_iter() + .map(|(name, _)| name) + .collect(); + + if inventory_models.is_empty() { + Ok(all_models) + } else { + Ok(inventory_models) + } + } + + async fn fetch_recommended_model_info(&self) -> Result, ProviderError> { + Ok(self + .fetch_recommended_models() + .await? + .iter() + .map(|model_name| model_info_for_provider_model(self.get_name(), model_name)) + .collect()) + } + + async fn map_to_canonical_model( + &self, + provider_model: &str, + ) -> Result, ProviderError> { + let registry = CanonicalModelRegistry::bundled().map_err(|e| { + ProviderError::ExecutionError(format!("Failed to load canonical registry: {}", e)) + })?; + + Ok(map_to_canonical_model( + self.get_name(), + provider_model, + registry, + )) + } + + fn supports_embeddings(&self) -> bool { + false + } + + /// Whether the provider manages its own conversation context (e.g. CLI + /// wrappers like Claude Code or Gemini CLI). When true, goose-side + /// context management such as tool-pair summarization is skipped because + /// the provider's internal state is the source of truth. + fn manages_own_context(&self) -> bool { + false + } + + async fn supports_cache_control(&self) -> bool { + false + } + + /// Create embeddings if supported. Default implementation returns an error. + async fn create_embeddings( + &self, + _session_id: &str, + _texts: Vec, + ) -> Result>, ProviderError> { + Err(ProviderError::ExecutionError( + "This provider does not support embeddings".to_string(), + )) + } + + /// Configure OAuth authentication for this provider + /// + /// This method is called when a provider has configuration keys marked with oauth_flow = true. + /// Providers that support OAuth should override this method to implement their specific OAuth flow. + /// + /// # Returns + /// * `Ok(())` if OAuth configuration succeeds and credentials are saved + /// * `Err(ProviderError)` if OAuth fails or is not supported by this provider + /// + /// # Default Implementation + /// The default implementation returns an error indicating OAuth is not supported. + async fn configure_oauth(&self) -> Result<(), ProviderError> { + Err(ProviderError::ExecutionError( + "OAuth configuration not supported by this provider".to_string(), + )) + } + + async fn refresh_credentials(&self) -> Result<(), ProviderError> { + Err(ProviderError::NotImplemented( + "credential refresh not supported by this provider".to_string(), + )) + } + + async fn update_mode(&self, _session_id: &str, _mode: GooseMode) -> Result<(), ProviderError> { + Ok(()) + } + + fn permission_routing(&self) -> PermissionRouting { + PermissionRouting::Noop + } + + async fn handle_permission_confirmation( + &self, + _request_id: &str, + _confirmation: &PermissionConfirmation, + ) -> bool { + false + } +} + +#[cfg(test)] +mod tests { + use super::*; + use test_case::test_case; + + fn content_from_str(s: String) -> MessageContent { + if let Some(img_data) = s.strip_prefix("*img:") { + MessageContent::image(format!("http://example.com/{}", img_data), "image/png") + } else if let Some(tool_name) = s.strip_prefix("*tool:") { + let tool_call = Ok( + rmcp::model::CallToolRequestParams::new(tool_name.to_string()) + .with_arguments(serde_json::Map::new()), + ); + MessageContent::tool_request(format!("tool_{}", tool_name), tool_call) + } else { + MessageContent::text(s) + } + } + + fn create_test_stream( + items: Vec, + ) -> impl Stream, Option), ProviderError>> { + use futures::stream; + stream::iter(items.into_iter().map(|item| { + let content = content_from_str(item); + let message = Message::new( + rmcp::model::Role::Assistant, + chrono::Utc::now().timestamp(), + vec![content], + ); + Ok((Some(message), None)) + })) + } + + fn content_to_strings(msg: &Message) -> Vec { + msg.content + .iter() + .map(|c| match c { + MessageContent::Text(t) => t.text.clone(), + MessageContent::Image(_) => "*img".to_string(), + MessageContent::ToolRequest(tr) => { + if let Ok(call) = &tr.tool_call { + format!("*tool:{}", call.name) + } else { + "*tool:error".to_string() + } + } + _ => "*other".to_string(), + }) + .collect() + } + + #[test_case( + vec!["Hello", " ", "world"], + vec!["Hello world"] + ; "consecutive text coalesces" + )] + #[test_case( + vec!["Hello", "*img:pic1", "world"], + vec!["Hello", "*img", "world"] + ; "non-text breaks coalescing" + )] + #[test_case( + vec!["A", "B", "*img:pic1", "C", "D", "*tool:read", "E", "F"], + vec!["AB", "*img", "CD", "*tool:read", "EF"] + ; "multiple text groups" + )] + #[test_case( + vec!["Text1", "*img:pic", "Text2"], + vec!["Text1", "*img", "Text2"] + ; "mixed content in chunk" + )] + #[tokio::test] + async fn test_collect_stream_coalescing(input_items: Vec<&str>, expected: Vec<&str>) { + let items: Vec = input_items.into_iter().map(|s| s.to_string()).collect(); + let stream = create_test_stream(items); + let (msg, _) = collect_stream(Box::pin(stream)).await.unwrap(); + assert_eq!(content_to_strings(&msg), expected); + } + + #[tokio::test] + async fn test_collect_stream_defaults_usage() { + let stream = create_test_stream(vec!["Hello".to_string()]); + let (msg, usage) = collect_stream(Box::pin(stream)).await.unwrap(); + assert_eq!(content_to_strings(&msg), vec!["Hello"]); + assert_eq!(usage.model, "unknown"); + } + + #[test] + fn test_model_info_creation() { + // Test direct ModelInfo creation + let info = ModelInfo { + name: "test-model".to_string(), + resolved_model: None, + context_limit: 1000, + input_token_cost: None, + output_token_cost: None, + currency: None, + supports_cache_control: None, + reasoning: false, + }; + assert_eq!(info.context_limit, 1000); + + // Test equality + let info2 = ModelInfo { + name: "test-model".to_string(), + resolved_model: None, + context_limit: 1000, + input_token_cost: None, + output_token_cost: None, + currency: None, + supports_cache_control: None, + reasoning: false, + }; + assert_eq!(info, info2); + + // Test inequality + let info3 = ModelInfo { + name: "test-model".to_string(), + resolved_model: None, + context_limit: 2000, + input_token_cost: None, + output_token_cost: None, + currency: None, + supports_cache_control: None, + reasoning: false, + }; + assert_ne!(info, info3); + } + + #[test] + 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.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/src/config/goose_mode.rs b/crates/goose-providers/src/goose_mode.rs similarity index 100% rename from crates/goose/src/config/goose_mode.rs rename to crates/goose-providers/src/goose_mode.rs diff --git a/crates/goose-providers/src/lib.rs b/crates/goose-providers/src/lib.rs index ae7f7e7aa..60d552f8b 100644 --- a/crates/goose-providers/src/lib.rs +++ b/crates/goose-providers/src/lib.rs @@ -3,9 +3,12 @@ pub mod canonical; pub mod conversation; pub mod errors; pub mod formats; +pub mod goose_mode; pub mod images; pub mod json; pub(crate) mod mcp_utils; pub mod model; +pub mod permission; +pub mod retry; pub mod thinking; pub mod utils; diff --git a/crates/goose/src/permission/permission_confirmation.rs b/crates/goose-providers/src/permission.rs similarity index 100% rename from crates/goose/src/permission/permission_confirmation.rs rename to crates/goose-providers/src/permission.rs diff --git a/crates/goose/src/providers/retry.rs b/crates/goose-providers/src/retry.rs similarity index 97% rename from crates/goose/src/providers/retry.rs rename to crates/goose-providers/src/retry.rs index 5cb12faa3..8660a9f1c 100644 --- a/crates/goose/src/providers/retry.rs +++ b/crates/goose-providers/src/retry.rs @@ -1,6 +1,6 @@ -use crate::providers::base::Provider; +use crate::base::Provider; +use crate::errors::ProviderError; use async_trait::async_trait; -use goose_providers::errors::ProviderError; use std::future::Future; use std::time::Duration; use tokio::time::sleep; @@ -13,16 +13,16 @@ pub const DEFAULT_MAX_RETRY_INTERVAL_MS: u64 = 30_000; #[derive(Debug, Clone)] pub struct RetryConfig { /// Maximum number of retry attempts - pub(crate) max_retries: usize, + pub max_retries: usize, /// Initial interval between retries in milliseconds - pub(crate) initial_interval_ms: u64, + pub initial_interval_ms: u64, /// Multiplier for backoff (exponential) - pub(crate) backoff_multiplier: f64, + pub backoff_multiplier: f64, /// Maximum interval between retries in milliseconds - pub(crate) max_interval_ms: u64, + pub max_interval_ms: u64, /// When true, only retry on transient errors (ServerError, NetworkError, /// RateLimitExceeded). RequestFailed (4xx client errors) will not be retried. - pub(crate) transient_only: bool, + pub transient_only: bool, } impl Default for RetryConfig { diff --git a/crates/goose-server/src/openapi.rs b/crates/goose-server/src/openapi.rs index 4610dc815..272f5a385 100644 --- a/crates/goose-server/src/openapi.rs +++ b/crates/goose-server/src/openapi.rs @@ -5,10 +5,11 @@ use goose::config::permission::PermissionLevel; use goose::config::ExtensionEntry; use goose::conversation::Conversation; use goose::download_manager::{DownloadProgress, DownloadStatus}; -use goose::permission::permission_confirmation::{Permission, PrincipalType}; use goose::providers::base::{ConfigKey, ModelInfo, ProviderMetadata, ProviderType}; use goose::session::{Session, SessionType, SystemInfo}; use goose_providers::model::ModelConfig; +use goose_providers::permission::Permission; +use goose_providers::permission::PrincipalType; use goose_providers::thinking::ThinkingEffort; use rmcp::model::{ Annotations, Content, EmbeddedResource, Icon, IconTheme, ImageContent, JsonObject, @@ -578,7 +579,7 @@ derive_utoipa!(IconTheme as IconThemeSchema); ThinkingEffort, super::routes::config_management::ProviderModelInfoQuery, Session, - goose::config::goose_mode::GooseMode, + goose_providers::goose_mode::GooseMode, SessionType, SystemInfo, Conversation, diff --git a/crates/goose/src/config/mod.rs b/crates/goose/src/config/mod.rs index ae6c46265..78bef69f0 100644 --- a/crates/goose/src/config/mod.rs +++ b/crates/goose/src/config/mod.rs @@ -2,7 +2,6 @@ pub mod base; pub mod declarative_providers; mod experiments; pub mod extensions; -pub mod goose_mode; mod migrations; pub mod paths; pub mod permission; @@ -21,7 +20,7 @@ pub use extensions::{ get_extension_by_name, get_warnings, is_extension_enabled, remove_extension, resolve_extensions_for_new_session, set_extension, set_extension_enabled, ExtensionEntry, }; -pub use goose_mode::GooseMode; +pub use goose_providers::goose_mode::GooseMode; pub use permission::PermissionManager; pub use signup_nanogpt::configure_nanogpt; pub use signup_openrouter::configure_openrouter; diff --git a/crates/goose/src/permission/mod.rs b/crates/goose/src/permission/mod.rs index dea80164d..bfe2caf86 100644 --- a/crates/goose/src/permission/mod.rs +++ b/crates/goose/src/permission/mod.rs @@ -1,8 +1,10 @@ -pub mod permission_confirmation; pub mod permission_inspector; pub mod permission_judge; pub mod permission_store; -pub use permission_confirmation::{Permission, PermissionConfirmation}; +pub use goose_providers::permission::{Permission, PermissionConfirmation}; +pub mod permission_confirmation { + pub use goose_providers::permission::PrincipalType; +} pub use permission_inspector::PermissionInspector; pub use permission_store::ToolPermissionStore; diff --git a/crates/goose/src/providers/base.rs b/crates/goose/src/providers/base.rs index 27868f5ad..f7ccf3018 100644 --- a/crates/goose/src/providers/base.rs +++ b/crates/goose/src/providers/base.rs @@ -1,12 +1,8 @@ use anyhow::Result; -use async_trait::async_trait; use futures::future::BoxFuture; -use futures::Stream; pub use goose_providers::conversation::token_usage::{ DraftStats, ProviderStats, ProviderUsage, Usage, }; -use goose_providers::errors::ProviderError; -use regex::Regex; use serde::{Deserialize, Serialize}; /// Default HTTP timeout for all provider API calls. @@ -14,175 +10,21 @@ use serde::{Deserialize, Serialize}; /// before giving up. Individual providers may override this via their own config key. pub const DEFAULT_PROVIDER_TIMEOUT_SECS: u64 = 600; -use super::canonical::{map_to_canonical_model, CanonicalModelRegistry}; -use super::retry::RetryConfig; use crate::config::base::ConfigValue; -use crate::config::{ExtensionConfig, GooseMode}; -use crate::conversation::message::{Message, MessageContent}; -use crate::conversation::Conversation; -use crate::permission::PermissionConfirmation; -use crate::utils::safe_truncate; +use crate::config::ExtensionConfig; +use goose_providers::conversation::message::Message; use goose_providers::model::ModelConfig; -use rmcp::model::Tool; use utoipa::ToSchema; use once_cell::sync::Lazy; use std::path::PathBuf; -use std::pin::Pin; -use std::sync::{LazyLock, Mutex}; +use std::sync::Mutex; + +pub use goose_providers::base::*; /// A global store for the current model being used, we use this as when a provider returns, it tells us the real model, not an alias pub static CURRENT_MODEL: Lazy>> = Lazy::new(|| Mutex::new(None)); -fn strip_xml_tags(text: &str) -> String { - static BLOCK_RE: LazyLock = LazyLock::new(|| { - Regex::new(r"(?s)<([a-zA-Z][a-zA-Z0-9_]*)[^>]*>.*?").unwrap() - }); - static TAG_RE: LazyLock = - LazyLock::new(|| Regex::new(r"]*>").unwrap()); - let pass1 = BLOCK_RE.replace_all(text, ""); - TAG_RE.replace_all(&pass1, "").into_owned() -} - -fn extract_short_title(text: &str) -> String { - let word_count = text.split_whitespace().count(); - if word_count <= 8 { - return text.to_string(); - } - - { - let mut results = Vec::new(); - let mut quote_char: Option = None; - let mut current = String::new(); - let mut prev_char: Option = None; - - for ch in text.chars() { - match quote_char { - None => { - if matches!(ch, '"' | '\'' | '`') { - let after_alnum = prev_char.map(|p| p.is_alphanumeric()).unwrap_or(false); - if !after_alnum { - quote_char = Some(ch); - current.clear(); - } - } - } - Some(q) => { - if ch == q { - let trimmed = current.trim().to_string(); - let wc = trimmed.split_whitespace().count(); - if (2..=8).contains(&wc) { - results.push(trimmed); - } - quote_char = None; - current.clear(); - } else { - current.push(ch); - } - } - } - prev_char = Some(ch); - } - - if let Some(title) = results.last() { - return title.clone(); - } - } - - if let Some(last) = text.lines().rev().find(|l| !l.trim().is_empty()) { - return last.trim().to_string(); - } - - text.to_string() -} - -pub static MSG_COUNT_FOR_SESSION_NAME_GENERATION: usize = 3; - -/// Information about a model's capabilities -#[derive(Debug, Clone, Serialize, Deserialize, ToSchema, PartialEq)] -pub struct ModelInfo { - /// The name of the model - pub name: String, - /// The underlying model resolved from provider metadata, when the configured model is an alias or endpoint. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub resolved_model: Option, - /// The maximum context length this model supports - pub context_limit: usize, - /// Cost per token for input in USD (optional) - pub input_token_cost: Option, - /// Cost per token for output in USD (optional) - pub output_token_cost: Option, - /// Currency for the costs (default: "$") - pub currency: Option, - /// Whether this model supports cache control - pub supports_cache_control: Option, - /// Whether this model supports reasoning/thinking controls - #[serde(default)] - pub reasoning: bool, -} - -impl ModelInfo { - /// Create a new ModelInfo with just name and context limit - pub fn new(name: impl Into, context_limit: usize) -> Self { - Self { - name: name.into(), - resolved_model: None, - context_limit, - input_token_cost: None, - output_token_cost: None, - currency: None, - supports_cache_control: None, - reasoning: false, - } - } - - /// Create a new ModelInfo with cost information (per token) - pub fn with_cost( - name: impl Into, - context_limit: usize, - input_cost: f64, - output_cost: f64, - ) -> Self { - Self { - name: name.into(), - resolved_model: None, - context_limit, - input_token_cost: Some(input_cost), - output_token_cost: Some(output_cost), - currency: Some("$".to_string()), - supports_cache_control: None, - reasoning: false, - } - } -} - -fn model_info_for_provider_model(provider_name: &str, model_name: &str) -> ModelInfo { - let registry = CanonicalModelRegistry::bundled().ok(); - let canonical = registry.as_ref().and_then(|registry| { - let canonical_id = map_to_canonical_model(provider_name, model_name, registry)?; - let (provider, model) = canonical_id.split_once('/')?; - registry.get(provider, model) - }); - - let reasoning = canonical - .as_ref() - .and_then(|model| model.reasoning) - .unwrap_or_else(|| ModelConfig::new_or_fail(model_name).is_reasoning_model()); - - ModelInfo { - name: model_name.to_string(), - resolved_model: None, - context_limit: ModelConfig::new_or_fail(model_name) - .with_canonical_limits(provider_name) - .context_limit(), - input_token_cost: None, - output_token_cost: None, - currency: None, - supports_cache_control: None, - reasoning, - } -} - #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, ToSchema)] pub enum ProviderType { Preferred, @@ -422,599 +264,15 @@ pub trait ProviderDef: Send + Sync { } } -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub enum PermissionRouting { - ActionRequired, - Noop, -} - -/// Base trait for AI providers (OpenAI, Anthropic, etc) -#[async_trait] -pub trait Provider: Send + Sync { - /// Get the name of this provider instance - fn get_name(&self) -> &str; - - /// Primary streaming method that all providers must implement. - /// - /// Note: Do not add `#[instrument]` here — the call sites (`complete` and - /// `stream_response_from_provider`) create the telemetry span so that - /// `session.id` is set once rather than in every provider. - async fn stream( - &self, - model_config: &ModelConfig, - session_id: &str, - system: &str, - messages: &[Message], - tools: &[Tool], - ) -> Result; - - /// Complete with a specific model config. - #[tracing::instrument( - skip(self, model_config, session_id, system, messages, tools), - fields(session.id = %session_id, gen_ai.request.model = %model_config.model_name) - )] - async fn complete( - &self, - model_config: &ModelConfig, - session_id: &str, - system: &str, - messages: &[Message], - tools: &[Tool], - ) -> Result<(Message, ProviderUsage), ProviderError> { - let stream = self - .stream(model_config, session_id, system, messages, tools) - .await?; - collect_stream(stream).await - } - - /// Try fast model first, fall back to regular model on failure. - async fn complete_fast( - &self, - session_id: &str, - system: &str, - messages: &[Message], - tools: &[Tool], - ) -> Result<(Message, ProviderUsage), ProviderError> { - let model_config = self.get_model_config(); - let fast_config = model_config.use_fast_model(); - - let result = self - .complete(&fast_config, session_id, system, messages, tools) - .await; - - match result { - Ok(response) => Ok(response), - Err(e) => { - if fast_config.model_name != model_config.model_name { - tracing::warn!( - "Fast model {} failed with error: {}. Falling back to regular model {}", - fast_config.model_name, - e, - model_config.model_name - ); - self.complete(&model_config, session_id, system, messages, tools) - .await - } else { - Err(e) - } - } - } - } - - /// Get the model config from the provider - fn get_model_config(&self) -> ModelConfig; - - fn retry_config(&self) -> RetryConfig { - RetryConfig::default() - } - - async fn fetch_supported_models(&self) -> Result, ProviderError> { - Ok(vec![]) - } - - async fn fetch_supported_model_info(&self) -> Result, ProviderError> { - Ok(self - .fetch_supported_models() - .await? - .iter() - .map(|model_name| model_info_for_provider_model(self.get_name(), model_name)) - .collect()) - } - - async fn fetch_model_info(&self, model_name: &str) -> Result { - Ok(model_info_for_provider_model(self.get_name(), model_name)) - } - - fn skip_canonical_filtering(&self) -> bool { - false - } - - /// Fetch inventory models filtered by canonical registry and usability. - async fn fetch_recommended_models(&self) -> Result, ProviderError> { - let all_models = self.fetch_supported_models().await?; - - if self.skip_canonical_filtering() { - return Ok(all_models); - } - - let registry = CanonicalModelRegistry::bundled().map_err(|e| { - ProviderError::ExecutionError(format!("Failed to load canonical registry: {}", e)) - })?; - - let provider_name = self.get_name(); - - // Get all text-capable models with their release dates - let mut models_with_dates: Vec<(String, Option)> = all_models - .iter() - .filter_map(|model| { - let canonical_id = map_to_canonical_model(provider_name, model, registry)?; - - let (provider, model_name) = canonical_id.split_once('/')?; - let canonical_model = registry.get(provider, model_name)?; - - if !canonical_model - .modalities - .input - .contains(&crate::providers::canonical::Modality::Text) - { - return None; - } - - if !canonical_model.tool_call && !self.get_model_config().toolshim { - return None; - } - - let release_date = canonical_model.release_date.clone(); - - Some((model.clone(), release_date)) - }) - .collect(); - - // Sort by release date (most recent first), then alphabetically for models without dates - models_with_dates.sort_by(|a, b| match (&a.1, &b.1) { - (Some(date_a), Some(date_b)) => date_b.cmp(date_a), - (Some(_), None) => std::cmp::Ordering::Less, - (None, Some(_)) => std::cmp::Ordering::Greater, - (None, None) => a.0.cmp(&b.0), - }); - - let inventory_models: Vec = models_with_dates - .into_iter() - .map(|(name, _)| name) - .collect(); - - if inventory_models.is_empty() { - Ok(all_models) - } else { - Ok(inventory_models) - } - } - - async fn fetch_recommended_model_info(&self) -> Result, ProviderError> { - Ok(self - .fetch_recommended_models() - .await? - .iter() - .map(|model_name| model_info_for_provider_model(self.get_name(), model_name)) - .collect()) - } - - async fn map_to_canonical_model( - &self, - provider_model: &str, - ) -> Result, ProviderError> { - let registry = CanonicalModelRegistry::bundled().map_err(|e| { - ProviderError::ExecutionError(format!("Failed to load canonical registry: {}", e)) - })?; - - Ok(map_to_canonical_model( - self.get_name(), - provider_model, - registry, - )) - } - - fn supports_embeddings(&self) -> bool { - false - } - - /// Whether the provider manages its own conversation context (e.g. CLI - /// wrappers like Claude Code or Gemini CLI). When true, goose-side - /// context management such as tool-pair summarization is skipped because - /// the provider's internal state is the source of truth. - fn manages_own_context(&self) -> bool { - false - } - - async fn supports_cache_control(&self) -> bool { - false - } - - /// Create embeddings if supported. Default implementation returns an error. - async fn create_embeddings( - &self, - _session_id: &str, - _texts: Vec, - ) -> Result>, ProviderError> { - Err(ProviderError::ExecutionError( - "This provider does not support embeddings".to_string(), - )) - } - - /// Returns the first 3 user messages as strings for session naming, - /// filtering out assistant-only content (e.g. preprompt blocks). - fn get_initial_user_messages(&self, messages: &Conversation) -> Vec { - messages - .iter() - .filter(|m| m.role == rmcp::model::Role::User) - .take(MSG_COUNT_FOR_SESSION_NAME_GENERATION) - .map(|m| { - m.content - .iter() - .filter_map(|c| c.filter_for_audience(rmcp::model::Role::User)) - .filter_map(|c| c.as_text().map(|s| s.to_string())) - .collect::>() - .join("\n") - }) - .collect() - } - - /// Extracts preprompt context (assistant-audience blocks) from the first user message. - /// These are content blocks visible to the assistant but not the user. - fn get_preprompt_context(&self, messages: &Conversation) -> String { - messages - .iter() - .filter(|m| m.role == rmcp::model::Role::User) - .take(1) - .flat_map(|m| m.content.iter()) - .filter_map(|c| { - // If this block is NOT visible to the user, it's preprompt/assistant-only content - if c.filter_for_audience(rmcp::model::Role::User).is_none() { - c.as_text().map(|s| s.to_string()) - } else { - None - } - }) - .collect::>() - .join("\n") - } - - /// Generate a session name/description based on the conversation history - /// Creates a prompt asking for a concise description in 4 words or less. - async fn generate_session_name( - &self, - session_id: &str, - messages: &Conversation, - ) -> Result { - let context = self.get_initial_user_messages(messages); - let preprompt_context = self.get_preprompt_context(messages); - let system = crate::prompt_template::render_template( - "session_name.md", - &std::collections::HashMap::::new(), - ) - .map_err(|e| ProviderError::ContextLengthExceeded(e.to_string()))?; - - use super::cli_common::{ - SESSION_NAME_BEGIN_MARKER, SESSION_NAME_END_MARKER, SESSION_NAME_SUFFIX, - }; - - let preprompt_section = if preprompt_context.is_empty() { - String::new() - } else { - format!( - "---BEGIN BACKGROUND CONTEXT (for understanding only, do NOT base the title on this)---\n{}\n---END BACKGROUND CONTEXT---\n\n", - preprompt_context - ) - }; - - let user_text = format!( - "{}{}\n{}\n{}\n\n{}", - preprompt_section, - SESSION_NAME_BEGIN_MARKER, - context.join("\n"), - SESSION_NAME_END_MARKER, - SESSION_NAME_SUFFIX, - ); - let message = Message::user().with_text(&user_text); - let result = self - .complete_fast(session_id, &system, &[message], &[]) - .await?; - - let raw: String = result - .0 - .content - .iter() - .filter_map(|c| c.as_text()) - .collect(); - let description = strip_xml_tags(&raw) - .split_whitespace() - .collect::>() - .join(" "); - - Ok(safe_truncate(&extract_short_title(&description), 100)) - } - - /// Configure OAuth authentication for this provider - /// - /// This method is called when a provider has configuration keys marked with oauth_flow = true. - /// Providers that support OAuth should override this method to implement their specific OAuth flow. - /// - /// # Returns - /// * `Ok(())` if OAuth configuration succeeds and credentials are saved - /// * `Err(ProviderError)` if OAuth fails or is not supported by this provider - /// - /// # Default Implementation - /// The default implementation returns an error indicating OAuth is not supported. - async fn configure_oauth(&self) -> Result<(), ProviderError> { - Err(ProviderError::ExecutionError( - "OAuth configuration not supported by this provider".to_string(), - )) - } - - async fn refresh_credentials(&self) -> Result<(), ProviderError> { - Err(ProviderError::NotImplemented( - "credential refresh not supported by this provider".to_string(), - )) - } - - async fn update_mode(&self, _session_id: &str, _mode: GooseMode) -> Result<(), ProviderError> { - Ok(()) - } - - fn permission_routing(&self) -> PermissionRouting { - PermissionRouting::Noop - } - - async fn handle_permission_confirmation( - &self, - _request_id: &str, - _confirmation: &PermissionConfirmation, - ) -> bool { - false - } -} - -/// A message stream yields partial text content but complete tool calls, all within the Message object -/// So a message with text will contain potentially just a word of a longer response, but tool calls -/// messages will only be yielded once concatenated. -pub type MessageStream = Pin< - Box, Option), ProviderError>> + Send>, ->; - pub fn stream_from_single_message(message: Message, usage: ProviderUsage) -> MessageStream { let stream = futures::stream::once(async move { Ok((Some(message), Some(usage))) }); Box::pin(stream) } -/// Collect all chunks from a MessageStream into a single Message and ProviderUsage -pub async fn collect_stream( - mut stream: MessageStream, -) -> Result<(Message, ProviderUsage), ProviderError> { - use futures::StreamExt; - - let mut final_message: Option = None; - let mut final_usage: Option = None; - - while let Some(result) = stream.next().await { - let (msg_opt, usage_opt) = result?; - - if let Some(msg) = msg_opt { - final_message = Some(match final_message { - Some(mut prev) => { - for new_content in msg.content { - match (&mut prev.content.last_mut(), &new_content) { - // Coalesce consecutive text blocks - ( - Some(MessageContent::Text(last_text)), - MessageContent::Text(new_text), - ) => { - last_text.text.push_str(&new_text.text); - } - _ => { - prev.content.push(new_content); - } - } - } - prev - } - None => msg, - }); - } - - if let Some(usage) = usage_opt { - final_usage = Some(usage); - } - } - - match final_message { - Some(msg) => { - let usage = final_usage - .unwrap_or_else(|| ProviderUsage::new("unknown".to_string(), Usage::default())); - Ok((msg, usage)) - } - None => Err(ProviderError::ExecutionError( - "Stream yielded no message".to_string(), - )), - } -} - #[cfg(test)] mod tests { use super::*; use std::collections::HashMap; - use test_case::test_case; - - #[test] - fn test_strip_xml_tags() { - assert_eq!(strip_xml_tags("reasoninganswer"), "answer"); - assert_eq!(strip_xml_tags("beforemidafter"), "beforeafter"); - assert_eq!(strip_xml_tags("xyz"), "z"); - assert_eq!(strip_xml_tags("no tags here"), "no tags here"); - assert_eq!(strip_xml_tags("a < b > c"), "a < b > c"); - assert_eq!(strip_xml_tags("überok"), "ok"); - assert_eq!(strip_xml_tags("日本語hello"), "hello"); - assert_eq!(strip_xml_tags(""), ""); - assert_eq!(strip_xml_tags("<>stuff"), "<>stuff"); - // attributes - assert_eq!( - strip_xml_tags(r#"reasoninganswer"#), - "answer" - ); - // self-closing tags - assert_eq!(strip_xml_tags("
self closing"), "self closing"); - // orphan closing tags - assert_eq!(strip_xml_tags("orphan tag"), "orphan tag"); - // multiline content - assert_eq!( - strip_xml_tags("\nline1\nline2\nresult"), - "result" - ); - } - - #[test] - fn test_extract_short_title() { - assert_eq!(extract_short_title("List files"), "List files"); - assert_eq!( - extract_short_title( - r#"blah blah blah blah blah blah blah blah blah "List files in folder""# - ), - "List files in folder" - ); - assert_eq!( - extract_short_title( - "blah blah blah blah blah blah blah blah blah `View current files`" - ), - "View current files" - ); - assert_eq!( - extract_short_title( - r#"stuff stuff stuff stuff stuff stuff stuff stuff "Abc title" "Zzz title""# - ), - "Zzz title" - ); - assert_eq!( - extract_short_title( - "long long long long long long long long long\nList files in folder" - ), - "List files in folder" - ); - assert_eq!( - extract_short_title( - r#"lots of words here and there and more and more "single" final line here"# - ), - "lots of words here and there and more and more \"single\" final line here" - ); - assert_eq!(extract_short_title("Hello world"), "Hello world"); - assert_eq!( - extract_short_title( - r#"1. Analyze the request. 2. The user's message says list files. 3. "List current folder files" fits perfectly. Result: List current folder files"# - ), - "List current folder files" - ); - assert_eq!( - extract_short_title( - r#"the user's phrasing is about listing files and the user's intent is clear. "List folder files" is best"# - ), - "List folder files" - ); - assert_eq!( - extract_short_title( - "lots of reasoning here about what to call it\nList current folder files" - ), - "List current folder files" - ); - } - - #[test] - fn test_usage_creation() { - let usage = Usage::new(Some(10), Some(20), Some(30)); - assert_eq!(usage.input_tokens, Some(10)); - assert_eq!(usage.output_tokens, Some(20)); - assert_eq!(usage.total_tokens, Some(30)); - } - - fn content_from_str(s: String) -> MessageContent { - if let Some(img_data) = s.strip_prefix("*img:") { - MessageContent::image(format!("http://example.com/{}", img_data), "image/png") - } else if let Some(tool_name) = s.strip_prefix("*tool:") { - let tool_call = Ok( - rmcp::model::CallToolRequestParams::new(tool_name.to_string()) - .with_arguments(serde_json::Map::new()), - ); - MessageContent::tool_request(format!("tool_{}", tool_name), tool_call) - } else { - MessageContent::text(s) - } - } - - fn create_test_stream( - items: Vec, - ) -> impl Stream, Option), ProviderError>> { - use futures::stream; - stream::iter(items.into_iter().map(|item| { - let content = content_from_str(item); - let message = Message::new( - rmcp::model::Role::Assistant, - chrono::Utc::now().timestamp(), - vec![content], - ); - Ok((Some(message), None)) - })) - } - - fn content_to_strings(msg: &Message) -> Vec { - msg.content - .iter() - .map(|c| match c { - MessageContent::Text(t) => t.text.clone(), - MessageContent::Image(_) => "*img".to_string(), - MessageContent::ToolRequest(tr) => { - if let Ok(call) = &tr.tool_call { - format!("*tool:{}", call.name) - } else { - "*tool:error".to_string() - } - } - _ => "*other".to_string(), - }) - .collect() - } - - #[test_case( - vec!["Hello", " ", "world"], - vec!["Hello world"] - ; "consecutive text coalesces" - )] - #[test_case( - vec!["Hello", "*img:pic1", "world"], - vec!["Hello", "*img", "world"] - ; "non-text breaks coalescing" - )] - #[test_case( - vec!["A", "B", "*img:pic1", "C", "D", "*tool:read", "E", "F"], - vec!["AB", "*img", "CD", "*tool:read", "EF"] - ; "multiple text groups" - )] - #[test_case( - vec!["Text1", "*img:pic", "Text2"], - vec!["Text1", "*img", "Text2"] - ; "mixed content in chunk" - )] - #[tokio::test] - async fn test_collect_stream_coalescing(input_items: Vec<&str>, expected: Vec<&str>) { - let items: Vec = input_items.into_iter().map(|s| s.to_string()).collect(); - let stream = create_test_stream(items); - let (msg, _) = collect_stream(Box::pin(stream)).await.unwrap(); - assert_eq!(content_to_strings(&msg), expected); - } - - #[tokio::test] - async fn test_collect_stream_defaults_usage() { - let stream = create_test_stream(vec!["Hello".to_string()]); - let (msg, usage) = collect_stream(Box::pin(stream)).await.unwrap(); - assert_eq!(content_to_strings(&msg), vec!["Hello"]); - assert_eq!(usage.model, "unknown"); - } #[test] fn test_provider_metadata_context_limits() { @@ -1048,56 +306,4 @@ mod tests { // unknown model should have default limit (128k) assert_eq!(*model_info.get("unknown-model").unwrap(), 128_000); } - - #[test] - fn test_model_info_creation() { - // Test direct ModelInfo creation - let info = ModelInfo { - name: "test-model".to_string(), - resolved_model: None, - context_limit: 1000, - input_token_cost: None, - output_token_cost: None, - currency: None, - supports_cache_control: None, - reasoning: false, - }; - assert_eq!(info.context_limit, 1000); - - // Test equality - let info2 = ModelInfo { - name: "test-model".to_string(), - resolved_model: None, - context_limit: 1000, - input_token_cost: None, - output_token_cost: None, - currency: None, - supports_cache_control: None, - reasoning: false, - }; - assert_eq!(info, info2); - - // Test inequality - let info3 = ModelInfo { - name: "test-model".to_string(), - resolved_model: None, - context_limit: 2000, - input_token_cost: None, - output_token_cost: None, - currency: None, - supports_cache_control: None, - reasoning: false, - }; - assert_ne!(info, info3); - } - - #[test] - 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.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/src/providers/mod.rs b/crates/goose/src/providers/mod.rs index d64215081..00fefee82 100644 --- a/crates/goose/src/providers/mod.rs +++ b/crates/goose/src/providers/mod.rs @@ -53,7 +53,9 @@ pub mod openrouter; pub mod pi_acp; pub mod provider_registry; pub mod provider_test; -mod retry; +mod retry { + pub use goose_providers::retry::*; +} #[cfg(feature = "aws-providers")] pub mod sagemaker_tgi; pub mod snowflake; diff --git a/crates/goose/src/session/mod.rs b/crates/goose/src/session/mod.rs index 58fdf8674..951c533bc 100644 --- a/crates/goose/src/session/mod.rs +++ b/crates/goose/src/session/mod.rs @@ -7,6 +7,7 @@ mod legacy; #[cfg(feature = "nostr")] pub mod nostr_share; pub mod session_manager; +mod session_naming; pub use diagnostics::{ config_path, generate_diagnostics, get_system_info, latest_llm_log_path, diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index e17a987e6..3ac0b6cfe 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -2,9 +2,12 @@ use crate::config::paths::Paths; use crate::config::GooseMode; use crate::conversation::message::Message; use crate::conversation::Conversation; -use crate::providers::base::{Provider, MSG_COUNT_FOR_SESSION_NAME_GENERATION}; +use crate::providers::base::Provider; use crate::recipe::Recipe; use crate::session::extension_data::ExtensionData; +use crate::session::session_naming::{ + generate_session_name, MSG_COUNT_FOR_SESSION_NAME_GENERATION, +}; use anyhow::Result; use chrono::{DateTime, Utc}; use goose_providers::model::ModelConfig; @@ -508,7 +511,7 @@ impl SessionManager { .count(); if user_message_count <= MSG_COUNT_FOR_SESSION_NAME_GENERATION { - let name = provider.generate_session_name(id, &conversation).await?; + let name = generate_session_name(provider.as_ref(), id, &conversation).await?; return Ok(Some(self.system_generated_name_update(id, name).await?)); } Ok(None) @@ -2042,6 +2045,9 @@ mod tests { use super::*; use crate::conversation::message::{Message, MessageContent}; use crate::providers::base::MessageStream; + use goose_providers::conversation::token_usage::ProviderUsage; + use goose_providers::errors::ProviderError; + use rmcp::model::Tool; use tempfile::TempDir; use test_case::test_case; @@ -2066,19 +2072,24 @@ mod tests { _messages: &[Message], _tools: &[rmcp::model::Tool], ) -> std::result::Result { - unimplemented!("session naming tests override generate_session_name") + unimplemented!("session naming calls complete_fast") } fn get_model_config(&self) -> ModelConfig { self.model_config.clone() } - async fn generate_session_name( + async fn complete_fast( &self, _session_id: &str, - _messages: &Conversation, - ) -> std::result::Result { - Ok(GENERATED_SESSION_NAME.to_string()) + _system: &str, + _messages: &[Message], + _tools: &[Tool], + ) -> Result<(Message, ProviderUsage), ProviderError> { + Ok(( + Message::assistant().with_text(GENERATED_SESSION_NAME), + ProviderUsage::new("test".to_string(), Default::default()), + )) } } diff --git a/crates/goose/src/session/session_naming.rs b/crates/goose/src/session/session_naming.rs new file mode 100644 index 000000000..74321016a --- /dev/null +++ b/crates/goose/src/session/session_naming.rs @@ -0,0 +1,249 @@ +use std::sync::LazyLock; + +use anyhow::Result; +use goose_providers::conversation::{message::Message, Conversation}; +use regex::Regex; + +use crate::{providers::base::Provider, utils::safe_truncate}; + +pub static MSG_COUNT_FOR_SESSION_NAME_GENERATION: usize = 3; + +fn strip_xml_tags(text: &str) -> String { + static BLOCK_RE: LazyLock = LazyLock::new(|| { + Regex::new(r"(?s)<([a-zA-Z][a-zA-Z0-9_]*)[^>]*>.*?").unwrap() + }); + static TAG_RE: LazyLock = + LazyLock::new(|| Regex::new(r"]*>").unwrap()); + let pass1 = BLOCK_RE.replace_all(text, ""); + TAG_RE.replace_all(&pass1, "").into_owned() +} + +fn extract_short_title(text: &str) -> String { + let word_count = text.split_whitespace().count(); + if word_count <= 8 { + return text.to_string(); + } + + { + let mut results = Vec::new(); + let mut quote_char: Option = None; + let mut current = String::new(); + let mut prev_char: Option = None; + + for ch in text.chars() { + match quote_char { + None => { + if matches!(ch, '"' | '\'' | '`') { + let after_alnum = prev_char.map(|p| p.is_alphanumeric()).unwrap_or(false); + if !after_alnum { + quote_char = Some(ch); + current.clear(); + } + } + } + Some(q) => { + if ch == q { + let trimmed = current.trim().to_string(); + let wc = trimmed.split_whitespace().count(); + if (2..=8).contains(&wc) { + results.push(trimmed); + } + quote_char = None; + current.clear(); + } else { + current.push(ch); + } + } + } + prev_char = Some(ch); + } + + if let Some(title) = results.last() { + return title.clone(); + } + } + + if let Some(last) = text.lines().rev().find(|l| !l.trim().is_empty()) { + return last.trim().to_string(); + } + + text.to_string() +} + +/// Returns the first 3 user messages as strings for session naming, +/// filtering out assistant-only content (e.g. preprompt blocks). +fn get_initial_user_messages(messages: &Conversation) -> Vec { + messages + .iter() + .filter(|m| m.role == rmcp::model::Role::User) + .take(MSG_COUNT_FOR_SESSION_NAME_GENERATION) + .map(|m| { + m.content + .iter() + .filter_map(|c| c.filter_for_audience(rmcp::model::Role::User)) + .filter_map(|c| c.as_text().map(|s| s.to_string())) + .collect::>() + .join("\n") + }) + .collect() +} + +/// Extracts preprompt context (assistant-audience blocks) from the first user message. +/// These are content blocks visible to the assistant but not the user. +fn get_preprompt_context(messages: &Conversation) -> String { + messages + .iter() + .filter(|m| m.role == rmcp::model::Role::User) + .take(1) + .flat_map(|m| m.content.iter()) + .filter_map(|c| { + // If this block is NOT visible to the user, it's preprompt/assistant-only content + if c.filter_for_audience(rmcp::model::Role::User).is_none() { + c.as_text().map(|s| s.to_string()) + } else { + None + } + }) + .collect::>() + .join("\n") +} + +/// Generate a session name/description based on the conversation history +/// Creates a prompt asking for a concise description in 4 words or less. +pub(crate) async fn generate_session_name( + provider: &dyn Provider, + session_id: &str, + messages: &Conversation, +) -> Result { + let context = get_initial_user_messages(messages); + let preprompt_context = get_preprompt_context(messages); + let system = crate::prompt_template::render_template( + "session_name.md", + &std::collections::HashMap::::new(), + )?; + + use crate::providers::cli_common::{ + SESSION_NAME_BEGIN_MARKER, SESSION_NAME_END_MARKER, SESSION_NAME_SUFFIX, + }; + + let preprompt_section = if preprompt_context.is_empty() { + String::new() + } else { + format!( + "---BEGIN BACKGROUND CONTEXT (for understanding only, do NOT base the title on this)---\n{}\n---END BACKGROUND CONTEXT---\n\n", + preprompt_context + ) + }; + + let user_text = format!( + "{}{}\n{}\n{}\n\n{}", + preprompt_section, + SESSION_NAME_BEGIN_MARKER, + context.join("\n"), + SESSION_NAME_END_MARKER, + SESSION_NAME_SUFFIX, + ); + let message = Message::user().with_text(&user_text); + let result = provider + .complete_fast(session_id, &system, &[message], &[]) + .await?; + + let raw: String = result + .0 + .content + .iter() + .filter_map(|c| c.as_text()) + .collect(); + let description = strip_xml_tags(&raw) + .split_whitespace() + .collect::>() + .join(" "); + + Ok(safe_truncate(&extract_short_title(&description), 100)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_strip_xml_tags() { + assert_eq!(strip_xml_tags("reasoninganswer"), "answer"); + assert_eq!(strip_xml_tags("beforemidafter"), "beforeafter"); + assert_eq!(strip_xml_tags("xyz"), "z"); + assert_eq!(strip_xml_tags("no tags here"), "no tags here"); + assert_eq!(strip_xml_tags("a < b > c"), "a < b > c"); + assert_eq!(strip_xml_tags("überok"), "ok"); + assert_eq!(strip_xml_tags("日本語hello"), "hello"); + assert_eq!(strip_xml_tags(""), ""); + assert_eq!(strip_xml_tags("<>stuff"), "<>stuff"); + // attributes + assert_eq!( + strip_xml_tags(r#"reasoninganswer"#), + "answer" + ); + // self-closing tags + assert_eq!(strip_xml_tags("
self closing"), "self closing"); + // orphan closing tags + assert_eq!(strip_xml_tags("orphan tag"), "orphan tag"); + // multiline content + assert_eq!( + strip_xml_tags("\nline1\nline2\nresult"), + "result" + ); + } + + #[test] + fn test_extract_short_title() { + assert_eq!(extract_short_title("List files"), "List files"); + assert_eq!( + extract_short_title( + r#"blah blah blah blah blah blah blah blah blah "List files in folder""# + ), + "List files in folder" + ); + assert_eq!( + extract_short_title( + "blah blah blah blah blah blah blah blah blah `View current files`" + ), + "View current files" + ); + assert_eq!( + extract_short_title( + r#"stuff stuff stuff stuff stuff stuff stuff stuff "Abc title" "Zzz title""# + ), + "Zzz title" + ); + assert_eq!( + extract_short_title( + "long long long long long long long long long\nList files in folder" + ), + "List files in folder" + ); + assert_eq!( + extract_short_title( + r#"lots of words here and there and more and more "single" final line here"# + ), + "lots of words here and there and more and more \"single\" final line here" + ); + assert_eq!(extract_short_title("Hello world"), "Hello world"); + assert_eq!( + extract_short_title( + r#"1. Analyze the request. 2. The user's message says list files. 3. "List current folder files" fits perfectly. Result: List current folder files"# + ), + "List current folder files" + ); + assert_eq!( + extract_short_title( + r#"the user's phrasing is about listing files and the user's intent is clear. "List folder files" is best"# + ), + "List folder files" + ); + assert_eq!( + extract_short_title( + "lots of reasoning here about what to call it\nList current folder files" + ), + "List current folder files" + ); + } +}