tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

move the provider trait into goose-providers (#9860)

Browse Source
This commit is contained in:
Jack Amadeo
2026-06-18 14:29:58 -04:00
committed by GitHub
parent c72c6531aa
commit 73fd41d7c5
15 changed files with 855 additions and 828 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+3
View File
@@ -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",
crates/goose-providers/Cargo.toml
+4
View File
@@ -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 }
crates/goose-providers/src/base.rs
+554 -8
View File
@@ -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<String>,
/// The maximum context length this model supports
pub context_limit: usize,
/// Cost per token for input in USD (optional)
pub input_token_cost: Option<f64>,
/// Cost per token for output in USD (optional)
pub output_token_cost: Option<f64>,
/// Currency for the costs (default: "$")
pub currency: Option<String>,
/// Whether this model supports cache control
pub supports_cache_control: Option<bool>,
/// 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<String>, 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<String>,
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<dyn Stream<Item = Result<(Option<Message>, Option<ProviderUsage>), ProviderError>> + Send>,
>;
pub trait Provider {
fn stream(req: StreamingRequest) -> impl Future<Output = Result<StreamingResponse, Error>>;
#[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<Message> = None;
let mut final_usage: Option<ProviderUsage> = 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<MessageStream, ProviderError>;
/// 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<Vec<String>, ProviderError> {
Ok(vec![])
}
async fn fetch_supported_model_info(&self) -> Result<Vec<ModelInfo>, 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<ModelInfo, ProviderError> {
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<Vec<String>, 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<String>)> = 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<String> = 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<Vec<ModelInfo>, 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<Option<String>, 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<String>,
) -> Result<Vec<Vec<f32>>, 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<String>,
) -> impl Stream<Item = Result<(Option<Message>, Option<ProviderUsage>), 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<String> {
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<String> = 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()));
}
}
crates/goose/src/config/goose_mode.rs → crates/goose-providers/src/goose_mode.rs
View File
crates/goose-providers/src/lib.rs
+3
View File
@@ -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;
crates/goose/src/permission/permission_confirmation.rs → crates/goose-providers/src/permission.rs
View File
crates/goose/src/providers/retry.rs → crates/goose-providers/src/retry.rs
+7 -7
View File
@@ -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 {
crates/goose-server/src/openapi.rs
+3 -2
View File
@@ -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,
crates/goose/src/config/mod.rs
+1 -2
View File
@@ -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;
crates/goose/src/permission/mod.rs
+4 -2
View File
@@ -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;
crates/goose/src/providers/base.rs
+5 -799
View File
@@ -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<Mutex<Option<String>>> = Lazy::new(|| Mutex::new(None));
fn strip_xml_tags(text: &str) -> String {
static BLOCK_RE: LazyLock<Regex> = LazyLock::new(|| {
Regex::new(r"(?s)<([a-zA-Z][a-zA-Z0-9_]*)[^>]*>.*?</[a-zA-Z][a-zA-Z0-9_]*>").unwrap()
});
static TAG_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"</?[a-zA-Z][a-zA-Z0-9_]*[^>]*>").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<char> = None;
let mut current = String::new();
let mut prev_char: Option<char> = 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<String>,
/// The maximum context length this model supports
pub context_limit: usize,
/// Cost per token for input in USD (optional)
pub input_token_cost: Option<f64>,
/// Cost per token for output in USD (optional)
pub output_token_cost: Option<f64>,
/// Currency for the costs (default: "$")
pub currency: Option<String>,
/// Whether this model supports cache control
pub supports_cache_control: Option<bool>,
/// 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<String>, 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<String>,
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<MessageStream, ProviderError>;
/// 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<Vec<String>, ProviderError> {
Ok(vec![])
}
async fn fetch_supported_model_info(&self) -> Result<Vec<ModelInfo>, 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<ModelInfo, ProviderError> {
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<Vec<String>, 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<String>)> = 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<String> = 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<Vec<ModelInfo>, 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<Option<String>, 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<String>,
) -> Result<Vec<Vec<f32>>, 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<String> {
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::<Vec<_>>()
.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::<Vec<_>>()
.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<String, ProviderError> {
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::<String, String>::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::<Vec<_>>()
.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<dyn Stream<Item = Result<(Option<Message>, Option<ProviderUsage>), 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<Message> = None;
let mut final_usage: Option<ProviderUsage> = 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("<think>reasoning</think>answer"), "answer");
assert_eq!(strip_xml_tags("before<t>mid</t>after"), "beforeafter");
assert_eq!(strip_xml_tags("<a>x</a><b>y</b>z"), "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("<think>über</think>ok"), "ok");
assert_eq!(strip_xml_tags("<think>日本語</think>hello"), "hello");
assert_eq!(strip_xml_tags(""), "");
assert_eq!(strip_xml_tags("<>stuff</>"), "<>stuff</>");
// attributes
assert_eq!(
strip_xml_tags(r#"<think class="deep">reasoning</think>answer"#),
"answer"
);
// self-closing tags
assert_eq!(strip_xml_tags("<br/>self closing"), "self closing");
// orphan closing tags
assert_eq!(strip_xml_tags("orphan </think> tag"), "orphan tag");
// multiline content
assert_eq!(
strip_xml_tags("<think>\nline1\nline2\n</think>result"),
"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<String>,
) -> impl Stream<Item = Result<(Option<Message>, Option<ProviderUsage>), 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<String> {
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<String> = 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()));
}
}
crates/goose/src/providers/mod.rs
+3 -1
View File
@@ -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;
crates/goose/src/session/mod.rs
+1
View File
@@ -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,
crates/goose/src/session/session_manager.rs
+18 -7
View File
@@ -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<MessageStream, goose_providers::errors::ProviderError> {
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<String, goose_providers::errors::ProviderError> {
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()),
))
}
}
crates/goose/src/session/session_naming.rs
+249
View File
@@ -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<Regex> = LazyLock::new(|| {
Regex::new(r"(?s)<([a-zA-Z][a-zA-Z0-9_]*)[^>]*>.*?</[a-zA-Z][a-zA-Z0-9_]*>").unwrap()
});
static TAG_RE: LazyLock<Regex> =
LazyLock::new(|| Regex::new(r"</?[a-zA-Z][a-zA-Z0-9_]*[^>]*>").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<char> = None;
let mut current = String::new();
let mut prev_char: Option<char> = 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<String> {
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::<Vec<_>>()
.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::<Vec<_>>()
.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<String> {
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::<String, String>::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::<Vec<_>>()
.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("<think>reasoning</think>answer"), "answer");
assert_eq!(strip_xml_tags("before<t>mid</t>after"), "beforeafter");
assert_eq!(strip_xml_tags("<a>x</a><b>y</b>z"), "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("<think>über</think>ok"), "ok");
assert_eq!(strip_xml_tags("<think>日本語</think>hello"), "hello");
assert_eq!(strip_xml_tags(""), "");
assert_eq!(strip_xml_tags("<>stuff</>"), "<>stuff</>");
// attributes
assert_eq!(
strip_xml_tags(r#"<think class="deep">reasoning</think>answer"#),
"answer"
);
// self-closing tags
assert_eq!(strip_xml_tags("<br/>self closing"), "self closing");
// orphan closing tags
assert_eq!(strip_xml_tags("orphan </think> tag"), "orphan tag");
// multiline content
assert_eq!(
strip_xml_tags("<think>\nline1\nline2\n</think>result"),
"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"
);
}
}