diff --git a/crates/goose-cli/src/commands/mcp.rs b/crates/goose-cli/src/commands/mcp.rs index bda870e4b..579b8d2e5 100644 --- a/crates/goose-cli/src/commands/mcp.rs +++ b/crates/goose-cli/src/commands/mcp.rs @@ -1,6 +1,7 @@ use anyhow::Result; -use goose_mcp::NonDeveloperRouter; -use goose_mcp::{Developer2Router, DeveloperRouter, JetBrainsRouter}; +use goose_mcp::{ + Developer2Router, DeveloperRouter, JetBrainsRouter, MemoryRouter, NonDeveloperRouter, +}; use mcp_server::router::RouterService; use mcp_server::{BoundedService, ByteTransport, Server}; use tokio::io::{stdin, stdout}; @@ -16,6 +17,7 @@ pub async fn run_server(name: &str) -> Result<()> { "developer2" => Some(Box::new(RouterService(Developer2Router::new()))), "nondeveloper" => Some(Box::new(RouterService(NonDeveloperRouter::new()))), "jetbrains" => Some(Box::new(RouterService(JetBrainsRouter::new()))), + "memory" => Some(Box::new(RouterService(MemoryRouter::new()))), _ => None, }; diff --git a/crates/goose-mcp/examples/mcp.rs b/crates/goose-mcp/examples/mcp.rs index 15e401a06..072479892 100644 --- a/crates/goose-mcp/examples/mcp.rs +++ b/crates/goose-mcp/examples/mcp.rs @@ -1,6 +1,6 @@ // An example script to run an MCP server use anyhow::Result; -use goose_mcp::DeveloperRouter; +use goose_mcp::{DeveloperRouter, MemoryRouter}; use mcp_server::router::RouterService; use mcp_server::{ByteTransport, Server}; use tokio::io::{stdin, stdout}; @@ -10,7 +10,7 @@ use tracing_subscriber::{self, EnvFilter}; #[tokio::main] async fn main() -> Result<()> { // Set up file appender for logging - let file_appender = RollingFileAppender::new(Rotation::DAILY, "logs", "mcp-server.log"); + let file_appender = RollingFileAppender::new(Rotation::DAILY, "logs", "goose-mcp-example.log"); // Initialize the tracing subscriber with file and stdout logging tracing_subscriber::fmt() @@ -25,7 +25,7 @@ async fn main() -> Result<()> { tracing::info!("Starting MCP server"); // Create an instance of our counter router - let router = RouterService(DeveloperRouter::new()); + let router = RouterService(MemoryRouter::new()); // Create and run the server let server = Server::new(router); diff --git a/crates/goose-mcp/src/lib.rs b/crates/goose-mcp/src/lib.rs index ed1f397fe..a09e35ecc 100644 --- a/crates/goose-mcp/src/lib.rs +++ b/crates/goose-mcp/src/lib.rs @@ -1,9 +1,11 @@ mod developer; mod developer2; mod jetbrains; +mod memory; mod nondeveloper; pub use developer::DeveloperRouter; pub use developer2::Developer2Router; pub use jetbrains::JetBrainsRouter; +pub use memory::MemoryRouter; pub use nondeveloper::NonDeveloperRouter; diff --git a/crates/goose-mcp/src/memory/mod.rs b/crates/goose-mcp/src/memory/mod.rs new file mode 100644 index 000000000..f157c9f6b --- /dev/null +++ b/crates/goose-mcp/src/memory/mod.rs @@ -0,0 +1,550 @@ +use async_trait::async_trait; +use indoc::formatdoc; +use serde_json::{json, Value}; +use std::{ + collections::HashMap, + fs, + future::Future, + io::{self, Read, Write}, + path::PathBuf, + pin::Pin, + sync::{Arc, Mutex}, +}; +use tracing::info; +use url::Url; + +use mcp_core::{ + handler::{ResourceError, ToolError}, + protocol::ServerCapabilities, + resource::Resource, + tool::{Tool, ToolCall}, + Content, +}; +use mcp_server::router::CapabilitiesBuilder; +use mcp_server::Router; + +// MemoryRouter implementation +#[derive(Clone)] +pub struct MemoryRouter { + tools: Vec, + instructions: String, + active_resources: Arc>>, + global_memory_dir: PathBuf, + local_memory_dir: PathBuf, +} + +impl Default for MemoryRouter { + fn default() -> Self { + Self::new() + } +} + +impl MemoryRouter { + pub fn new() -> Self { + let remember_memory = Tool::new( + "remember_memory", + "Stores a memory with optional tags in a specified category", + json!({ + "type": "object", + "properties": { + "category": {"type": "string"}, + "data": {"type": "string"}, + "tags": {"type": "array", "items": {"type": "string"}}, + "is_global": {"type": "boolean"} + }, + "required": ["category", "data", "is_global"] + }), + ); + + let retrieve_memories = Tool::new( + "retrieve_memories", + "Retrieves all memories from a specified category", + json!({ + "type": "object", + "properties": { + "category": {"type": "string"}, + "is_global": {"type": "boolean"} + }, + "required": ["category", "is_global"] + }), + ); + + let remove_memory_category = Tool::new( + "remove_memory_category", + "Removes all memories within a specified category", + json!({ + "type": "object", + "properties": { + "category": {"type": "string"}, + "is_global": {"type": "boolean"} + }, + "required": ["category", "is_global"] + }), + ); + + let remove_specific_memory = Tool::new( + "remove_specific_memory", + "Removes a specific memory within a specified category", + json!({ + "type": "object", + "properties": { + "category": {"type": "string"}, + "memory_content": {"type": "string"}, + "is_global": {"type": "boolean"} + }, + "required": ["category", "memory_content", "is_global"] + }), + ); + + let instructions = formatdoc! {r#" + Memory Management System for Goose + This system allows storage and retrieval of categorized information with tagging support. It's designed to help + manage important information across sessions in a systematic and organized manner. + Capabilities: + 1. Store information in categories with optional tags for context-based retrieval. + 2. Search memories by content or specific tags to find relevant information. + 3. List all available memory categories for easy navigation. + 4. Remove entire categories of memories when they are no longer needed. + Interaction Protocol: + When important information is identified, such as: + - User-specific data (e.g., name, preferences) + - Project-related configurations + - Workflow descriptions + - Other critical settings + The protocol is: + 1. Identify the critical piece of information. + 2. Ask the user if they'd like to store it for later reference. + 3. Upon agreement: + - Suggest a relevant category like "personal" for user data or "development" for project preferences. + - Inquire about any specific tags they want to apply for easier lookup. + - Confirm the desired storage location: + - Local storage (.goose/memory) for project-specific details. + - Global storage (~/.config/goose/memory) for user-wide data. + Example Interaction for Storing Information: + User: "For this project, we use black for code formatting" + Assistant: "You've mentioned a development preference. Would you like to remember this for future conversations? + User: "Yes, please." + Assistant: "I'll store this in the 'development' category. Any specific tags to add? Suggestions: #formatting + #tools" + User: "Yes, use those tags." + Assistant: "Shall I store this locally for this project only, or globally for all projects?" + User: "Locally, please." + Assistant: *Stores the information under category="development", tags="formatting tools", scope="local"* + Retrieving Memories: + To access stored information, utilize the memory retrieval protocols: + - **Search by Category**: + - Provides all memories within the specified context. + - Use: `retrieve_memories(category="development", is_global=False)` + - **Filter by Tags**: + - Enables targeted retrieval based on specific tags. + - Use: Provide tag filters to refine search. + The Protocol is: + 1. Confirm what kind of information the user seeks by category or keyword. + 2. Suggest categories or relevant tags based on the user's request. + 3. Use the retrieve function to access relevant memory entries. + 4. Present a summary of findings, offering detailed exploration upon request. + Example Interaction for Retrieving Information: + User: "What configuration do we use for code formatting?" + Assistant: "Let me check the 'development' category for any related memories. Searching using #formatting tag." + Assistant: *Executes retrieval: `retrieve_memories(category="development", is_global=False)`* + Assistant: "We have 'black' configured for code formatting, specific to this project. Would you like further + details?" + Memory Overview: + - Categories can include a wide range of topics, structured to keep information grouped logically. + - Tags enable quick filtering and identification of specific entries. + Operational Guidelines: + - Always confirm with the user before saving information. + - Propose suitable categories and tag suggestions. + - Discuss storage scope thoroughly to align with user needs. + - Acknowledge the user about what is stored and where, for transparency and ease of future retrieval. + "#}; + + // Check for .goose/memory in current directory + let local_memory_dir = std::env::current_dir() + .unwrap() + .join(".goose") + .join("memory"); + // Check for .config/goose/memory in user's home directory + let global_memory_dir = dirs::home_dir() + .map(|home| home.join(".config/goose/memory")) + .unwrap_or_else(|| PathBuf::from(".config/goose/memory")); + fs::create_dir_all(&global_memory_dir).unwrap(); + fs::create_dir_all(&local_memory_dir).unwrap(); + + Self { + tools: vec![ + remember_memory, + retrieve_memories, + remove_memory_category, + remove_specific_memory, + ], + instructions, + active_resources: Arc::new(Mutex::new(HashMap::new())), + global_memory_dir, + local_memory_dir, + } + } + + fn get_memory_file(&self, category: &str, is_global: bool) -> PathBuf { + // Defaults to local memory if no is_global flag is provided + let base_dir = if is_global { + &self.global_memory_dir + } else { + &self.local_memory_dir + }; + base_dir.join(format!("{}.txt", category)) + } + + pub fn retrieve_all(&self, is_global: bool) -> io::Result>> { + let base_dir = if is_global { + &self.global_memory_dir + } else { + &self.local_memory_dir + }; + let mut memories = HashMap::new(); + if base_dir.exists() { + for entry in fs::read_dir(base_dir)? { + let entry = entry?; + if entry.file_type()?.is_file() { + let category = entry.file_name().to_string_lossy().replace(".txt", ""); + let category_memories = self.retrieve(&category, is_global)?; + memories.insert( + category, + category_memories.into_iter().flat_map(|(_, v)| v).collect(), + ); + } + } + } + Ok(memories) + } + + pub fn remember( + &self, + _context: &str, + category: &str, + data: &str, + tags: &[&str], + is_global: bool, + ) -> io::Result<()> { + let memory_file_path = self.get_memory_file(category, is_global); + let uri = Url::from_file_path(&memory_file_path).unwrap().to_string(); + + let mut file = fs::OpenOptions::new() + .append(true) + .create(true) + .open(&memory_file_path)?; + if !tags.is_empty() { + writeln!(file, "# {}", tags.join(" "))?; + } + writeln!(file, "{}\n", data)?; + + // Create and store the resource + let resource = Resource::new(uri.clone(), Some("text".to_string()), None).unwrap(); + self.add_active_resource(uri, resource); + + Ok(()) + } + + pub fn retrieve( + &self, + category: &str, + is_global: bool, + ) -> io::Result>> { + let memory_file_path = self.get_memory_file(category, is_global); + if !memory_file_path.exists() { + return Ok(HashMap::new()); + } + + let uri = Url::from_file_path(&memory_file_path).unwrap().to_string(); + + let mut file = fs::File::open(memory_file_path)?; + let mut content = String::new(); + file.read_to_string(&mut content)?; + + let mut memories = HashMap::new(); + for entry in content.split("\n\n") { + let mut lines = entry.lines(); + if let Some(first_line) = lines.next() { + if let Some(stripped) = first_line.strip_prefix('#') { + let tags = stripped + .split_whitespace() + .map(String::from) + .collect::>(); + memories.insert(tags.join(" "), lines.map(String::from).collect()); + } else { + let entry_data: Vec = std::iter::once(first_line.to_string()) + .chain(lines.map(String::from)) + .collect(); + memories + .entry("untagged".to_string()) + .or_insert_with(Vec::new) + .extend(entry_data); + } + } + } + + // Update resource + if let Some(resource) = self.active_resources.lock().unwrap().get_mut(&uri) { + resource.update_timestamp(); + } + + Ok(memories) + } + + pub fn remove_specific_memory( + &self, + category: &str, + memory_content: &str, + is_global: bool, + ) -> io::Result<()> { + let memory_file_path = self.get_memory_file(category, is_global); + if !memory_file_path.exists() { + return Ok(()); + } + + let uri = Url::from_file_path(&memory_file_path).unwrap().to_string(); + + let mut file = fs::File::open(&memory_file_path)?; + let mut content = String::new(); + file.read_to_string(&mut content)?; + + let memories: Vec<&str> = content.split("\n\n").collect(); + let new_content: Vec = memories + .into_iter() + .filter(|entry| !entry.contains(memory_content)) + .map(|s| s.to_string()) + .collect(); + + fs::write(memory_file_path, new_content.join("\n\n"))?; + + // Update resource + if let Some(resource) = self.active_resources.lock().unwrap().get_mut(&uri) { + resource.update_timestamp(); + } + + Ok(()) + } + + pub fn clear_memory(&self, category: &str, is_global: bool) -> io::Result<()> { + let memory_file_path = self.get_memory_file(category, is_global); + let uri = Url::from_file_path(&memory_file_path).unwrap().to_string(); + if memory_file_path.exists() { + fs::remove_file(memory_file_path)?; + } + + // Remove resource from active resources + self.active_resources.lock().unwrap().remove(&uri); + + Ok(()) + } + + async fn execute_tool_call(&self, tool_call: ToolCall) -> Result { + match tool_call.name.as_str() { + "remember_memory" => { + let args = MemoryArgs::from_value(&tool_call.arguments)?; + let data = args + .data + .as_deref() + .filter(|d| !d.is_empty()) + .ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidInput, + "Data must exist when remembering a memory", + ) + })?; + self.remember("context", args.category, data, &args.tags, args.is_global)?; + Ok(format!("Stored memory in category: {}", args.category)) + } + "retrieve_memories" => { + let args = MemoryArgs::from_value(&tool_call.arguments)?; + let memories = self.retrieve(&args.category, args.is_global)?; + Ok(format!("Retrieved memories: {:?}", memories)) + } + "remove_memory_category" => { + let args = MemoryArgs::from_value(&tool_call.arguments)?; + self.clear_memory(&args.category, args.is_global)?; + Ok(format!("Cleared memories in category: {}", args.category)) + } + "remove_specific_memory" => { + let args = MemoryArgs::from_value(&tool_call.arguments)?; + let memory_content = tool_call.arguments["memory_content"].as_str().unwrap(); + self.remove_specific_memory(&args.category, memory_content, args.is_global)?; + Ok(format!( + "Removed specific memory from category: {}", + args.category + )) + } + _ => Err(io::Error::new(io::ErrorKind::InvalidInput, "Unknown tool")), + } + } + + fn add_active_resource(&self, uri: String, resource: Resource) { + self.active_resources.lock().unwrap().insert(uri, resource); + } + + async fn read_memory_resources(&self, uri: &str) -> Result { + // Ensure the resource exists in the active resources map + let active_resources = self.active_resources.lock().unwrap(); + let resource = active_resources + .get(uri) + .ok_or_else(|| ResourceError::NotFound(format!("Resource '{}' not found", uri)))?; + + let url = + Url::parse(uri).map_err(|e| ResourceError::NotFound(format!("Invalid URI: {}", e)))?; + + // Read content based on scheme and mime_type + match url.scheme() { + "file" => { + let path = url + .to_file_path() + .map_err(|_| ResourceError::NotFound("Invalid file path in URI".into()))?; + + // Ensure file exists + if !path.exists() { + return Err(ResourceError::NotFound(format!( + "File does not exist: {}", + path.display() + ))); + } + + match resource.mime_type.as_str() { + "text" => { + // Read the file as UTF-8 text + fs::read_to_string(&path).map_err(|e| { + ResourceError::ExecutionError(format!("Failed to read file: {}", e)) + }) + } + mime_type => Err(ResourceError::ExecutionError(format!( + "Unsupported mime type: {}", + mime_type + ))), + } + } + scheme => Err(ResourceError::NotFound(format!( + "Unsupported URI scheme: {}", + scheme + ))), + } + } +} + +#[async_trait] +impl Router for MemoryRouter { + fn name(&self) -> String { + "memory".to_string() + } + + fn instructions(&self) -> String { + self.instructions.clone() + } + + fn capabilities(&self) -> ServerCapabilities { + CapabilitiesBuilder::new() + .with_tools(false) + .with_resources(false, false) + .build() + } + + fn list_tools(&self) -> Vec { + self.tools.clone() + } + + fn call_tool( + &self, + tool_name: &str, + arguments: Value, + ) -> Pin, ToolError>> + Send + 'static>> { + let this = self.clone(); + let tool_name = tool_name.to_string(); + + Box::pin(async move { + let tool_call = ToolCall { + name: tool_name, + arguments, + }; + match this.execute_tool_call(tool_call).await { + Ok(result) => Ok(vec![Content::text(result)]), + Err(err) => Err(ToolError::ExecutionError(err.to_string())), + } + }) + } + + fn list_resources(&self) -> Vec { + let resources = self + .active_resources + .lock() + .unwrap() + .values() + .cloned() + .collect(); + info!("Listing resources: {:?}", resources); + resources + } + + fn read_resource( + &self, + uri: &str, + ) -> Pin> + Send + 'static>> { + let this = self.clone(); + let uri = uri.to_string(); + info!("Reading resource: {}", uri); + Box::pin(async move { + match this.read_memory_resources(&uri).await { + Ok(content) => Ok(content), + Err(e) => Err(e), + } + }) + } +} + +#[derive(Debug)] +struct MemoryArgs<'a> { + category: &'a str, + data: Option<&'a str>, + tags: Vec<&'a str>, + is_global: bool, +} + +impl<'a> MemoryArgs<'a> { + // Category is required, data is optional, tags are optional, is_global is optional + fn from_value(args: &'a Value) -> Result { + let category = args["category"].as_str().ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidInput, "Category must be a string") + })?; + + if category.is_empty() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "Category must be a string", + )); + } + + let data = args.get("data").and_then(|d| d.as_str()); + + let tags = match &args["tags"] { + Value::Array(arr) => arr.iter().filter_map(|v| v.as_str()).collect(), + Value::String(s) => vec![s.as_str()], + _ => Vec::new(), + }; + + let is_global = match &args.get("is_global") { + // Default to false if no is_global flag is provided + Some(Value::Bool(b)) => *b, + Some(Value::String(s)) => s.to_lowercase() == "true", + None => false, + _ => { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "is_global must be a boolean or string 'true'/'false'", + )) + } + }; + + Ok(Self { + category, + data, + tags, + is_global, + }) + } +} diff --git a/crates/goose-server/src/commands/mcp.rs b/crates/goose-server/src/commands/mcp.rs index dead113d9..0615342d9 100644 --- a/crates/goose-server/src/commands/mcp.rs +++ b/crates/goose-server/src/commands/mcp.rs @@ -1,6 +1,7 @@ use anyhow::Result; -use goose_mcp::NonDeveloperRouter; -use goose_mcp::{Developer2Router, DeveloperRouter, JetBrainsRouter}; +use goose_mcp::{ + Developer2Router, DeveloperRouter, JetBrainsRouter, MemoryRouter, NonDeveloperRouter, +}; use mcp_server::router::RouterService; use mcp_server::{BoundedService, ByteTransport, Server}; use tokio::io::{stdin, stdout}; @@ -15,6 +16,7 @@ pub async fn run(name: &str) -> Result<()> { "developer2" => Some(Box::new(RouterService(Developer2Router::new()))), "nondeveloper" => Some(Box::new(RouterService(NonDeveloperRouter::new()))), "jetbrains" => Some(Box::new(RouterService(JetBrainsRouter::new()))), + "memory" => Some(Box::new(RouterService(MemoryRouter::new()))), _ => None, };