diff --git a/crates/goose/src/agents/chatrecall_extension.rs b/crates/goose/src/agents/chatrecall_extension.rs new file mode 100644 index 0000000000..6176d4d42e --- /dev/null +++ b/crates/goose/src/agents/chatrecall_extension.rs @@ -0,0 +1,352 @@ +use crate::agents::extension::PlatformExtensionContext; +use crate::agents::mcp_client::{Error, McpClientTrait}; +use crate::session::SessionManager; +use anyhow::Result; +use async_trait::async_trait; +use indoc::indoc; +use rmcp::model::{ + CallToolResult, Content, GetPromptResult, Implementation, InitializeResult, JsonObject, + ListPromptsResult, ListResourcesResult, ListToolsResult, ProtocolVersion, ReadResourceResult, + ServerCapabilities, ServerNotification, Tool, ToolAnnotations, ToolsCapability, +}; +use schemars::{schema_for, JsonSchema}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +pub static EXTENSION_NAME: &str = "chatrecall"; + +/// Parameters for the chatrecall tool +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct ChatRecallParams { + /// Search keywords. Use multiple related terms/synonyms (e.g., 'database postgres sql'). Mutually exclusive with session_id. + #[serde(skip_serializing_if = "Option::is_none")] + query: Option, + /// Session ID to load. Returns first/last 3 messages. Mutually exclusive with query. + #[serde(skip_serializing_if = "Option::is_none")] + session_id: Option, + /// Max results (default: 10, max: 50). Search mode only. + #[serde(skip_serializing_if = "Option::is_none")] + limit: Option, + /// ISO 8601 date (e.g., '2025-10-01T00:00:00Z'). Search mode only. + #[serde(skip_serializing_if = "Option::is_none")] + after_date: Option, + /// ISO 8601 date (e.g., '2025-10-15T23:59:59Z'). Search mode only. + #[serde(skip_serializing_if = "Option::is_none")] + before_date: Option, +} + +pub struct ChatRecallClient { + info: InitializeResult, + context: PlatformExtensionContext, +} + +impl ChatRecallClient { + pub fn new(context: PlatformExtensionContext) -> Result { + let info = InitializeResult { + protocol_version: ProtocolVersion::V_2025_03_26, + capabilities: ServerCapabilities { + tools: Some(ToolsCapability { + list_changed: Some(false), + }), + resources: None, + prompts: None, + completions: None, + experimental: None, + logging: None, + }, + server_info: Implementation { + name: EXTENSION_NAME.to_string(), + title: Some("Chat Recall".to_string()), + version: "1.0.0".to_string(), + icons: None, + website_url: None, + }, + instructions: Some(indoc! {r#" + Chat Recall + + Search past conversations and load session summaries when the user expects some memory or context. + + Two modes: + - Search mode: Use query with keywords/synonyms to find relevant messages + - Load mode: Use session_id to get first and last messages of a specific session + "#}.to_string()), + }; + + Ok(Self { info, context }) + } + + #[allow(clippy::too_many_lines)] + async fn handle_chatrecall( + &self, + arguments: Option, + ) -> Result, String> { + let arguments = arguments.ok_or("Missing arguments")?; + + let session_id = arguments + .get("session_id") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + if let Some(sid) = session_id { + // LOAD MODE: Get session summary (first and last few messages) + match SessionManager::get_session(&sid, true).await { + Ok(loaded_session) => { + let conversation = loaded_session.conversation.as_ref(); + + if conversation.is_none() { + return Ok(vec![Content::text(format!( + "Session {} has no conversation.", + sid + ))]); + } + + let msgs = conversation.unwrap().messages(); + let total = msgs.len(); + + if total == 0 { + return Ok(vec![Content::text(format!( + "Session {} has no messages.", + sid + ))]); + } + + let mut output = format!( + "Session: {} (ID: {})\nWorking Dir: {}\nTotal Messages: {}\n\n", + loaded_session.name, + sid, + loaded_session.working_dir.display(), + total + ); + + // Show first 3 messages + let first_count = std::cmp::min(3, total); + output.push_str("--- First Few Messages ---\n\n"); + for (idx, msg) in msgs.iter().take(first_count).enumerate() { + output.push_str(&format!("{}. [{:?}] ", idx + 1, msg.role)); + for content in &msg.content { + if let Some(text) = content.as_text() { + output.push_str(text); + output.push('\n'); + } + } + output.push('\n'); + } + + // Show last 3 messages (if different from first) + if total > first_count { + output.push_str("--- Last Few Messages ---\n\n"); + let last_count = std::cmp::min(3, total); + let skip_count = total.saturating_sub(last_count); + for (idx, msg) in msgs.iter().skip(skip_count).enumerate() { + output.push_str(&format!( + "{}. [{:?}] ", + skip_count + idx + 1, + msg.role + )); + for content in &msg.content { + if let Some(text) = content.as_text() { + output.push_str(text); + output.push('\n'); + } + } + output.push('\n'); + } + } + + Ok(vec![Content::text(output)]) + } + Err(e) => Err(format!("Failed to load session: {}", e)), + } + } else { + // SEARCH MODE: Search across all sessions + let query = arguments + .get("query") + .and_then(|v| v.as_str()) + .ok_or("Missing required parameter: query or session_id")? + .to_string(); + + let limit = arguments + .get("limit") + .and_then(|v| v.as_i64()) + .map(|l| l as usize) + .unwrap_or(10) + .min(50); + + let after_date = arguments + .get("after_date") + .and_then(|v| v.as_str()) + .and_then(|s| chrono::DateTime::parse_from_rfc3339(s).ok()) + .map(|dt| dt.with_timezone(&chrono::Utc)); + + let before_date = arguments + .get("before_date") + .and_then(|v| v.as_str()) + .and_then(|s| chrono::DateTime::parse_from_rfc3339(s).ok()) + .map(|dt| dt.with_timezone(&chrono::Utc)); + + // Exclude current session from results to avoid self-referential loops + let exclude_session_id = self.context.session_id.clone(); + + match SessionManager::search_chat_history( + &query, + Some(limit), + after_date, + before_date, + exclude_session_id, + ) + .await + { + Ok(results) => { + let formatted_results = if results.total_matches == 0 { + format!("No results found for query: '{}'", query) + } else { + let mut output = format!( + "Found {} matching message(s) across {} session(s) for query: '{}'\n\n", + results.total_matches, + results.results.len(), + query + ); + for (idx, result) in results.results.iter().enumerate() { + output.push_str(&format!( + "{}. Session: {} (ID: {})\n Working Dir: {}\n Last Activity: {}\n Showing {} of {} total message(s) in session:\n\n", + idx + 1, + result.session_description, + result.session_id, + result.session_working_dir, + result.last_activity.format("%Y-%m-%d"), + result.messages.len(), + result.total_messages_in_session + )); + + for (msg_idx, message) in result.messages.iter().enumerate() { + output.push_str(&format!( + " {}.{} [{}]\n {}\n\n", + idx + 1, + msg_idx + 1, + message.role, + message + .content + .lines() + .map(|line| format!(" {}", line)) + .collect::>() + .join("\n") + )); + } + } + output + }; + Ok(vec![Content::text(formatted_results)]) + } + Err(e) => Err(format!("Chat recall failed: {}", e)), + } + } + } + + fn get_tools() -> Vec { + // Generate JSON schema from the ChatRecallParams struct + let schema = schema_for!(ChatRecallParams); + let schema_value = + serde_json::to_value(schema).expect("Failed to serialize ChatRecallParams schema"); + + let input_schema = schema_value + .as_object() + .expect("Schema should be an object") + .clone(); + + vec![Tool::new( + "chatrecall".to_string(), + indoc! {r#" + Search past chat or load session summaries. Use when it is clear user expects some memory or context. + + search mode (query): Use multiple keywords/synonyms. Returns messages grouped by session, ordered by recency. Supports date filters. + load mode (session_id): Returns first/last 3 messages of a session. + "#} + .to_string(), + input_schema, + ) + .annotate(ToolAnnotations { + title: Some("Recall past conversations".to_string()), + read_only_hint: Some(true), + destructive_hint: Some(false), + idempotent_hint: Some(true), + open_world_hint: Some(false), + })] + } +} + +#[async_trait] +impl McpClientTrait for ChatRecallClient { + async fn list_resources( + &self, + _next_cursor: Option, + _cancellation_token: CancellationToken, + ) -> Result { + Err(Error::TransportClosed) + } + + async fn read_resource( + &self, + _uri: &str, + _cancellation_token: CancellationToken, + ) -> Result { + Err(Error::TransportClosed) + } + + async fn list_tools( + &self, + _next_cursor: Option, + _cancellation_token: CancellationToken, + ) -> Result { + Ok(ListToolsResult { + tools: Self::get_tools(), + next_cursor: None, + }) + } + + async fn call_tool( + &self, + name: &str, + arguments: Option, + _cancellation_token: CancellationToken, + ) -> Result { + let content = match name { + "chatrecall" => self.handle_chatrecall(arguments).await, + _ => Err(format!("Unknown tool: {}", name)), + }; + + match content { + Ok(content) => Ok(CallToolResult::success(content)), + Err(error) => Ok(CallToolResult::error(vec![Content::text(format!( + "Error: {}", + error + ))])), + } + } + + async fn list_prompts( + &self, + _next_cursor: Option, + _cancellation_token: CancellationToken, + ) -> Result { + Err(Error::TransportClosed) + } + + async fn get_prompt( + &self, + _name: &str, + _arguments: Value, + _cancellation_token: CancellationToken, + ) -> Result { + Err(Error::TransportClosed) + } + + async fn subscribe(&self) -> mpsc::Receiver { + mpsc::channel(1).1 + } + + fn get_info(&self) -> Option<&InitializeResult> { + Some(&self.info) + } +} diff --git a/crates/goose/src/agents/extension.rs b/crates/goose/src/agents/extension.rs index ab41a7cdea..284a98cf5b 100644 --- a/crates/goose/src/agents/extension.rs +++ b/crates/goose/src/agents/extension.rs @@ -1,3 +1,4 @@ +use crate::agents::chatrecall_extension; use crate::agents::extension_manager_extension; use crate::agents::todo_extension; use std::collections::HashMap; @@ -50,6 +51,19 @@ pub static PLATFORM_EXTENSIONS: Lazy }, ); + map.insert( + chatrecall_extension::EXTENSION_NAME, + PlatformExtensionDef { + name: chatrecall_extension::EXTENSION_NAME, + description: + "Search past conversations and load session summaries for contextual memory", + default_enabled: false, + client_factory: |ctx| { + Box::new(chatrecall_extension::ChatRecallClient::new(ctx).unwrap()) + }, + }, + ); + map.insert( "extensionmanager", PlatformExtensionDef { diff --git a/crates/goose/src/agents/mod.rs b/crates/goose/src/agents/mod.rs index 7876510d58..5fa9f0ce10 100644 --- a/crates/goose/src/agents/mod.rs +++ b/crates/goose/src/agents/mod.rs @@ -1,4 +1,5 @@ mod agent; +pub(crate) mod chatrecall_extension; pub mod extension; pub mod extension_malware_check; pub mod extension_manager; diff --git a/crates/goose/src/session/chat_history_search.rs b/crates/goose/src/session/chat_history_search.rs new file mode 100644 index 0000000000..f3436701ee --- /dev/null +++ b/crates/goose/src/session/chat_history_search.rs @@ -0,0 +1,286 @@ +use crate::conversation::message::MessageContent; +use anyhow::Result; +use chrono::{DateTime, Utc}; +use serde::Serialize; +use sqlx::{Pool, Sqlite}; +use std::collections::HashMap; + +#[derive(Debug, Clone, Serialize)] +pub struct ChatRecallResult { + pub session_id: String, + pub session_description: String, + pub session_working_dir: String, + pub last_activity: DateTime, + pub total_messages_in_session: usize, + pub messages: Vec, +} + +#[derive(Debug, Clone, Serialize)] +pub struct ChatRecallMessage { + pub role: String, + pub content: String, + pub timestamp: DateTime, +} + +#[derive(Debug, Serialize)] +pub struct ChatRecallResults { + pub results: Vec, + pub total_matches: usize, +} + +type SqlQueryRow = ( + String, + String, + String, + DateTime, + String, + String, + DateTime, +); + +type SessionMessageGroup = ( + String, + String, + DateTime, + Vec<(String, String, DateTime)>, +); + +pub struct ChatHistorySearch<'a> { + pool: &'a Pool, + query: &'a str, + limit: usize, + after_date: Option>, + before_date: Option>, + exclude_session_id: Option, +} + +impl<'a> ChatHistorySearch<'a> { + pub fn new( + pool: &'a Pool, + query: &'a str, + limit: Option, + after_date: Option>, + before_date: Option>, + exclude_session_id: Option, + ) -> Self { + Self { + pool, + query, + limit: limit.unwrap_or(10), + after_date, + before_date, + exclude_session_id, + } + } + + pub async fn execute(self) -> Result { + let keywords = self.parse_keywords(); + if keywords.is_empty() { + return Ok(ChatRecallResults { + results: vec![], + total_matches: 0, + }); + } + + let rows = self.fetch_rows(&keywords).await?; + let session_messages = self.process_rows(rows); + let session_totals = self.get_session_totals(&session_messages).await?; + let results = self.convert_to_results(session_messages, session_totals); + + Ok(results) + } + + async fn fetch_rows(&self, keywords: &[String]) -> Result> { + let sql = self.build_sql(keywords); + let mut query_builder = sqlx::query_as::<_, SqlQueryRow>(&sql); + + for keyword in keywords { + query_builder = query_builder.bind(keyword); + } + + if let Some(exclude_id) = &self.exclude_session_id { + query_builder = query_builder.bind(exclude_id); + } + + if let Some(after) = self.after_date { + query_builder = query_builder.bind(after); + } + if let Some(before) = self.before_date { + query_builder = query_builder.bind(before); + } + + query_builder = query_builder.bind(self.limit as i64); + + Ok(query_builder.fetch_all(self.pool).await?) + } + + fn parse_keywords(&self) -> Vec { + self.query + .split_whitespace() + .map(|word| format!("%{}%", word.to_lowercase())) + .collect() + } + + fn build_sql(&self, keywords: &[String]) -> String { + let mut sql = String::from( + r#" + SELECT + s.id as session_id, + s.description as session_description, + s.working_dir as session_working_dir, + s.created_at as session_created_at, + m.role, + m.content_json, + m.timestamp + FROM messages m + INNER JOIN sessions s ON m.session_id = s.id + WHERE EXISTS ( + SELECT 1 FROM json_each(m.content_json) + WHERE json_extract(value, '$.type') = 'text' + AND ( + "#, + ); + + for (i, _) in keywords.iter().enumerate() { + if i > 0 { + sql.push_str(" OR "); + } + sql.push_str("LOWER(json_extract(value, '$.text')) LIKE ?"); + } + + sql.push_str( + r#" + ) + ) + "#, + ); + + if self.exclude_session_id.is_some() { + sql.push_str(" AND s.id != ?"); + } + + if self.after_date.is_some() { + sql.push_str(" AND m.timestamp >= ?"); + } + if self.before_date.is_some() { + sql.push_str(" AND m.timestamp <= ?"); + } + + sql.push_str(" ORDER BY m.timestamp DESC LIMIT ?"); + + sql + } + + fn process_rows(&self, rows: Vec) -> HashMap { + let mut session_messages: HashMap = HashMap::new(); + + for ( + session_id, + session_description, + session_working_dir, + session_created_at, + role, + content_json, + timestamp, + ) in rows + { + if let Ok(content_vec) = serde_json::from_str::>(&content_json) { + let text_parts = Self::extract_text_content(content_vec); + + if !text_parts.is_empty() { + let entry = session_messages.entry(session_id.clone()).or_insert(( + session_description.clone(), + session_working_dir.clone(), + session_created_at, + Vec::new(), + )); + entry + .3 + .push((role.clone(), text_parts.join("\n"), timestamp)); + } + } + } + + session_messages + } + + fn extract_text_content(content_vec: Vec) -> Vec { + content_vec + .into_iter() + .filter_map(|content| match content { + MessageContent::Text(ref tc) => Some(tc.text.clone()), + MessageContent::ToolRequest(ref tr) => { + Some(format!("[Tool: {}]", tr.to_readable_string())) + } + MessageContent::ToolResponse(_) => Some("[Tool Response]".to_string()), + MessageContent::Thinking(ref t) => Some(format!("[Thinking: {}]", t.thinking)), + _ => None, + }) + .collect() + } + + async fn get_session_totals( + &self, + session_messages: &HashMap, + ) -> Result> { + let mut session_totals: HashMap = HashMap::new(); + for session_id in session_messages.keys() { + let count: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM messages WHERE session_id = ?") + .bind(session_id) + .fetch_one(self.pool) + .await + .unwrap_or(0); + session_totals.insert(session_id.clone(), count as usize); + } + Ok(session_totals) + } + + fn convert_to_results( + &self, + session_messages: HashMap, + session_totals: HashMap, + ) -> ChatRecallResults { + let mut results: Vec = session_messages + .into_iter() + .map( + |(session_id, (description, working_dir, _created_at, messages))| { + let message_vec: Vec = messages + .into_iter() + .map(|(role, content, timestamp)| ChatRecallMessage { + role, + content, + timestamp, + }) + .collect(); + + let last_activity = message_vec + .iter() + .map(|m| m.timestamp) + .max() + .unwrap_or_else(chrono::Utc::now); + + let total_messages_in_session = + session_totals.get(&session_id).copied().unwrap_or(0); + + ChatRecallResult { + session_id, + session_description: description, + session_working_dir: working_dir, + last_activity, + total_messages_in_session, + messages: message_vec, + } + }, + ) + .collect(); + + results.sort_by(|a, b| b.last_activity.cmp(&a.last_activity)); + + let total_matches = results.iter().map(|r| r.messages.len()).sum(); + ChatRecallResults { + results, + total_matches, + } + } +} diff --git a/crates/goose/src/session/mod.rs b/crates/goose/src/session/mod.rs index 502f754d8e..e15579bb83 100644 --- a/crates/goose/src/session/mod.rs +++ b/crates/goose/src/session/mod.rs @@ -1,3 +1,4 @@ +mod chat_history_search; mod diagnostics; pub mod extension_data; mod legacy; diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index 1ec7a29603..ce97c8ead3 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -263,6 +263,19 @@ impl SessionManager { Ok(()) } } + + pub async fn search_chat_history( + query: &str, + limit: Option, + after_date: Option>, + before_date: Option>, + exclude_session_id: Option, + ) -> Result { + Self::instance() + .await? + .search_chat_history(query, limit, after_date, before_date, exclude_session_id) + .await + } } pub struct SessionStorage { @@ -1023,6 +1036,28 @@ impl SessionStorage { self.get_session(&session.id, true).await } + + async fn search_chat_history( + &self, + query: &str, + limit: Option, + after_date: Option>, + before_date: Option>, + exclude_session_id: Option, + ) -> Result { + use crate::session::chat_history_search::ChatHistorySearch; + + ChatHistorySearch::new( + &self.pool, + query, + limit, + after_date, + before_date, + exclude_session_id, + ) + .execute() + .await + } } #[cfg(test)]