diff --git a/crates/goose-providers/src/cli.rs b/crates/goose-providers/src/cli.rs new file mode 100644 index 000000000000..2636212ee51c --- /dev/null +++ b/crates/goose-providers/src/cli.rs @@ -0,0 +1,155 @@ +use crate::errors::ProviderError; +use goose_types::Usage; +use serde_json::Value; + +pub const SESSION_NAME_BEGIN_MARKER: &str = "---BEGIN USER MESSAGES---"; +pub const SESSION_NAME_END_MARKER: &str = "---END USER MESSAGES---"; +pub const SESSION_NAME_SUFFIX: &str = "Generate a short title for the above messages."; + +pub fn extract_usage_tokens(usage_info: &Value) -> Usage { + let get = |key: &str| { + usage_info + .get(key) + .and_then(|value| value.as_i64()) + .and_then(|value| i32::try_from(value).ok()) + }; + Usage::new( + get("input_tokens"), + get("output_tokens"), + get("total_tokens"), + ) +} + +pub fn error_from_event(provider_name: &str, parsed: &Value) -> ProviderError { + let error_msg = parsed + .get("error") + .and_then(|error| error.as_str()) + .or_else(|| parsed.get("message").and_then(|message| message.as_str())) + .unwrap_or("Unknown error"); + if error_msg.contains("context window exceeded") { + ProviderError::ContextLengthExceeded(error_msg.to_string()) + } else { + ProviderError::RequestFailed(format!("{provider_name} error: {error_msg}")) + } +} + +pub fn is_session_description_request(system: &str) -> bool { + system.contains("four words or less") || system.contains("4 words or less") +} + +pub fn strip_session_name_prompt_wrapper(text: &str) -> &str { + let text = text + .rfind(SESSION_NAME_BEGIN_MARKER) + .and_then(|idx| text.get(idx..)) + .unwrap_or(text); + let stripped = text + .strip_prefix(SESSION_NAME_BEGIN_MARKER) + .unwrap_or(text) + .trim_start_matches(['\n', '\r']); + let full_suffix = format!("{}\n\n{}", SESSION_NAME_END_MARKER, SESSION_NAME_SUFFIX); + stripped + .strip_suffix(&full_suffix) + .or_else(|| stripped.strip_suffix(SESSION_NAME_END_MARKER)) + .unwrap_or(stripped) + .trim() +} + +pub fn simple_session_description_from_text(text: &str) -> String { + let stripped = strip_session_name_prompt_wrapper(text); + let description = stripped + .split_whitespace() + .take(4) + .collect::>() + .join(" "); + if description.is_empty() { + "Simple task".to_string() + } else { + safe_truncate(&description, 100) + } +} + +fn safe_truncate(text: &str, max_chars: usize) -> String { + if text.chars().count() <= max_chars { + text.to_string() + } else { + let truncated: String = text.chars().take(max_chars.saturating_sub(3)).collect(); + format!("{}...", truncated) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn extracts_usage_tokens() { + let usage = extract_usage_tokens(&json!({ + "input_tokens": 10, + "output_tokens": 20, + "total_tokens": 30, + "ignored": 40, + })); + assert_eq!(usage.input_tokens, Some(10)); + assert_eq!(usage.output_tokens, Some(20)); + assert_eq!(usage.total_tokens, Some(30)); + } + + #[test] + fn event_error_maps_context_window_errors() { + let error = error_from_event( + "Claude CLI", + &json!({"error": "context window exceeded: too many tokens"}), + ); + assert!(matches!(error, ProviderError::ContextLengthExceeded(_))); + } + + #[test] + fn event_error_maps_generic_errors() { + let error = error_from_event("Claude CLI", &json!({"message": "boom"})); + assert_eq!( + error, + ProviderError::RequestFailed("Claude CLI error: boom".to_string()) + ); + } + + #[test] + fn detects_session_description_requests() { + assert!(is_session_description_request( + "Please use four words or less" + )); + assert!(is_session_description_request("Please use 4 words or less")); + assert!(!is_session_description_request( + "Please summarize this session" + )); + } + + #[test] + fn strips_session_name_prompt_wrapper() { + let text = format!( + "{}\nList files in the repo\n{}\n\n{}", + SESSION_NAME_BEGIN_MARKER, SESSION_NAME_END_MARKER, SESSION_NAME_SUFFIX + ); + assert_eq!( + strip_session_name_prompt_wrapper(&text), + "List files in the repo" + ); + } + + #[test] + fn simple_session_description_uses_first_four_words() { + let text = format!( + "{}\nList files in the repo please\n{}\n\n{}", + SESSION_NAME_BEGIN_MARKER, SESSION_NAME_END_MARKER, SESSION_NAME_SUFFIX + ); + assert_eq!( + simple_session_description_from_text(&text), + "List files in the" + ); + } + + #[test] + fn simple_session_description_defaults_when_empty() { + assert_eq!(simple_session_description_from_text(""), "Simple task"); + } +} diff --git a/crates/goose-providers/src/lib.rs b/crates/goose-providers/src/lib.rs index f6c6efe042bc..d622f3ba40d5 100644 --- a/crates/goose-providers/src/lib.rs +++ b/crates/goose-providers/src/lib.rs @@ -1,6 +1,7 @@ pub mod api_client; pub mod avian; pub mod canonical; +pub mod cli; pub mod errors; pub mod function_name; pub mod google_response; diff --git a/crates/goose/src/providers/cli_common.rs b/crates/goose/src/providers/cli_common.rs index f2531ec9f5a3..0abf096238ad 100644 --- a/crates/goose/src/providers/cli_common.rs +++ b/crates/goose/src/providers/cli_common.rs @@ -1,45 +1,13 @@ -use serde_json::Value; - -use super::base::{ProviderUsage, Usage}; +use super::base::ProviderUsage; use super::errors::ProviderError; use crate::conversation::message::{Message, MessageContent}; -use crate::utils::safe_truncate; use rmcp::model::Role; -pub(crate) fn extract_usage_tokens(usage_info: &Value) -> Usage { - let get = |key: &str| { - usage_info - .get(key) - .and_then(|v| v.as_i64()) - .and_then(|v| i32::try_from(v).ok()) - }; - Usage::new( - get("input_tokens"), - get("output_tokens"), - get("total_tokens"), - ) -} - -pub(crate) fn error_from_event(provider_name: &str, parsed: &Value) -> ProviderError { - let error_msg = parsed - .get("error") - .and_then(|e| e.as_str()) - .or_else(|| parsed.get("message").and_then(|m| m.as_str())) - .unwrap_or("Unknown error"); - if error_msg.contains("context window exceeded") { - ProviderError::ContextLengthExceeded(error_msg.to_string()) - } else { - ProviderError::RequestFailed(format!("{provider_name} error: {error_msg}")) - } -} - -pub(crate) const SESSION_NAME_BEGIN_MARKER: &str = "---BEGIN USER MESSAGES---"; -pub(crate) const SESSION_NAME_END_MARKER: &str = "---END USER MESSAGES---"; -pub(crate) const SESSION_NAME_SUFFIX: &str = "Generate a short title for the above messages."; - -pub(crate) fn is_session_description_request(system: &str) -> bool { - system.contains("four words or less") || system.contains("4 words or less") -} +pub(crate) use goose_providers::cli::{ + error_from_event, extract_usage_tokens, is_session_description_request, + simple_session_description_from_text, SESSION_NAME_BEGIN_MARKER, SESSION_NAME_END_MARKER, + SESSION_NAME_SUFFIX, +}; pub(crate) fn generate_simple_session_description( model_name: &str, @@ -47,42 +15,14 @@ pub(crate) fn generate_simple_session_description( ) -> Result<(Message, ProviderUsage), ProviderError> { let description = messages .iter() - .find(|m| m.role == Role::User) - .and_then(|m| { - m.content.iter().find_map(|c| match c { + .find(|message| message.role == Role::User) + .and_then(|message| { + message.content.iter().find_map(|content| match content { MessageContent::Text(text_content) => Some(&text_content.text), _ => None, }) }) - .map(|text| { - // Strip the wrapper added by generate_session_name so we get - // the actual user content. First strip the optional background context section. - let text = text - .rfind(SESSION_NAME_BEGIN_MARKER) - .and_then(|idx| text.get(idx..)) - .unwrap_or(text); - let stripped = text - .strip_prefix(SESSION_NAME_BEGIN_MARKER) - .unwrap_or(text) - .trim_start_matches(['\n', '\r']); - let full_suffix = format!("{}\n\n{}", SESSION_NAME_END_MARKER, SESSION_NAME_SUFFIX); - let stripped = stripped - .strip_suffix(&full_suffix) - .or_else(|| stripped.strip_suffix(SESSION_NAME_END_MARKER)) - .unwrap_or(stripped) - .trim(); - - let desc: String = stripped - .split_whitespace() - .take(4) - .collect::>() - .join(" "); - if desc.is_empty() { - "Simple task".to_string() - } else { - safe_truncate(&desc, 100) - } - }) + .map(|text| simple_session_description_from_text(text)) .unwrap_or_else(|| "Simple task".to_string()); tracing::debug!( @@ -98,6 +38,6 @@ pub(crate) fn generate_simple_session_description( Ok(( message, - ProviderUsage::new(model_name.to_string(), Usage::default()), + ProviderUsage::new(model_name.to_string(), Default::default()), )) }