use std::sync::Arc; use async_trait::async_trait; use serde_json::json; use crate::storage::{MemoryRecord, MemoryRepository, MemoryUpsert, is_valid_namespace}; use crate::tools::traits::{Tool, ToolContext, ToolResult}; pub struct MemoryManageTool { memories: Arc, } impl MemoryManageTool { pub fn new(memories: Arc) -> Self { Self { memories } } } #[async_trait] impl Tool for MemoryManageTool { fn name(&self) -> &str { "memory_manage" } fn description(&self) -> &str { "Create, update, or delete long-term user memories in the configured memory repository. Supports actions: put, update, delete. Use memory_search as the default retrieval path before answering most requests, and use memory_search for all retrieval actions including search, get, and list. Only call this tool when you have determined that a high-value long-term memory should be created, overwritten, updated, or deleted. Memories are scoped to the current channel and sender, and record the originating session/message when available." } fn parameters_schema(&self) -> serde_json::Value { let namespaces = crate::storage::allowed_namespace_names(); json!({ "type": "object", "properties": { "action": { "type": "string", "enum": ["put", "update", "delete"], "description": "管理操作。put 用于创建或覆盖,update 用于修改已有记录,delete 用于删除。检索请使用 memory_search。" }, "namespace": { "type": "string", "enum": namespaces, "description": "记忆命名空间分类" }, "key": { "type": "string", "description": "命名空间内的记忆键名" }, "content": { "type": "string", "description": "put/update 时的记忆内容" } }, "required": ["action"] }) } async fn execute(&self, _args: serde_json::Value) -> anyhow::Result { Ok(error_result("memory_manage requires tool context")) } async fn execute_with_context( &self, context: &ToolContext, args: serde_json::Value, ) -> anyhow::Result { let action = match args.get("action").and_then(|value| value.as_str()) { Some(action) => action, None => return Ok(error_result("Missing required parameter: action")), }; let scope_key = match scope_key_from_context(context) { Ok(scope_key) => scope_key, Err(result) => return Ok(result), }; let namespace = args.get("namespace").and_then(|value| value.as_str()); let key = args.get("key").and_then(|value| value.as_str()); let payload = match action { "put" => { let input = match build_memory_upsert(context, &scope_key, &args, true) { Ok(input) => input, Err(result) => return Ok(result), }; memory_to_json(self.memories.put_memory(&input)?) } "update" => { let input = match build_memory_upsert(context, &scope_key, &args, false) { Ok(input) => input, Err(result) => return Ok(result), }; match self.memories.update_memory(&input)? { Some(memory) => memory_to_json(memory), None => { return Ok(error_result(&format!( "memory '{}.{}' not found", input.namespace, input.memory_key ))); } } } "delete" => { let namespace = match namespace { Some(namespace) => namespace, None => return Ok(error_result("Missing required parameter: namespace")), }; let key = match key { Some(key) => key, None => return Ok(error_result("Missing required parameter: key")), }; let deleted = self .memories .delete_memory("user", &scope_key, namespace, key)?; if !deleted { return Ok(error_result(&format!( "memory '{}.{}' not found", namespace, key ))); } json!({ "status": "deleted", "namespace": namespace, "key": key, }) } _ => return Ok(error_result("Unsupported action")), }; Ok(ToolResult { success: true, output: serde_json::to_string_pretty(&payload)?, error: None, }) } } fn build_memory_upsert( context: &ToolContext, scope_key: &str, args: &serde_json::Value, allow_put: bool, ) -> Result { let namespace = match args.get("namespace").and_then(|value| value.as_str()) { Some(namespace) => namespace, None => return Err(error_result("Missing required parameter: namespace")), }; // 验证 namespace 是否在允许列表中 if !is_valid_namespace(namespace) { let allowed = crate::storage::allowed_namespace_names().join(", "); return Err(error_result(&format!( "Invalid namespace '{}'. Allowed namespaces: {}", namespace, allowed ))); } let key = match args.get("key").and_then(|value| value.as_str()) { Some(key) => key, None => return Err(error_result("Missing required parameter: key")), }; let content = match args.get("content").and_then(|value| value.as_str()) { Some(content) => content, None => return Err(error_result("Missing required parameter: content")), }; let source_type = if context.message_id.is_some() { "message" } else if allow_put { "manual" } else { "session" }; Ok(MemoryUpsert { scope_kind: "user".to_string(), scope_key: scope_key.to_string(), namespace: namespace.to_string(), memory_key: key.to_string(), content: content.to_string(), source_type: source_type.to_string(), source_session_id: context.session_id.clone(), source_message_id: context.message_id.clone(), source_message_seq: context.message_seq, source_channel_name: context.channel_name.clone(), source_chat_id: context.chat_id.clone(), }) } fn scope_key_from_context(_context: &ToolContext) -> Result { Ok(crate::storage::GLOBAL_SCOPE_KEY.to_string()) } fn memory_to_json(memory: MemoryRecord) -> serde_json::Value { json!({ "id": memory.id, "scope_kind": memory.scope_kind, "scope_key": memory.scope_key, "namespace": memory.namespace, "key": memory.memory_key, "content": memory.content, "source_type": memory.source_type, "source_session_id": memory.source_session_id, "source_message_id": memory.source_message_id, "source_message_seq": memory.source_message_seq, "source_channel_name": memory.source_channel_name, "source_chat_id": memory.source_chat_id, "created_at": memory.created_at, "updated_at": memory.updated_at, }) } fn error_result(message: &str) -> ToolResult { ToolResult { success: false, output: String::new(), error: Some(message.to_string()), } } #[cfg(test)] mod tests { use super::*; use crate::storage::SessionStore; const TEST_CHANNEL: &str = "test-channel"; #[tokio::test] async fn test_memory_manage_put_returns_saved_memory() { let store = Arc::new(SessionStore::in_memory().unwrap()); let tool = MemoryManageTool::new(store); let context = ToolContext { channel_name: Some(TEST_CHANNEL.to_string()), chat_id: Some("chat-1".to_string()), session_id: Some(format!("{}:chat-1", TEST_CHANNEL)), message_id: Some("msg-1".to_string()), message_seq: Some(1), ..ToolContext::default() }; let put = tool .execute_with_context( &context, json!({ "action": "put", "namespace": "user", "key": "language", "content": "Rust" }), ) .await .unwrap(); assert!(put.success); assert!(put.output.contains("Rust")); assert!(put.output.contains("msg-1")); } #[tokio::test] async fn test_memory_manage_works_with_default_context() { let store = Arc::new(SessionStore::in_memory().unwrap()); let tool = MemoryManageTool::new(store); // scope_key 已全局统一为 "default",不再依赖 channel_name let result = tool .execute_with_context( &ToolContext::default(), json!({ "action": "put", "namespace": "user", "key": "language", "content": "Rust" }), ) .await .unwrap(); assert!(result.success); assert!(result.output.contains("Rust")); } #[tokio::test] async fn test_memory_manage_rejects_read_actions() { let store = Arc::new(SessionStore::in_memory().unwrap()); let tool = MemoryManageTool::new(store); let context = ToolContext { channel_name: Some(TEST_CHANNEL.to_string()), ..ToolContext::default() }; let result = tool .execute_with_context( &context, json!({ "action": "get", "namespace": "user", "key": "language" }), ) .await .unwrap(); assert!(!result.success); assert!(result.error.unwrap().contains("Unsupported action")); } }