PicoBot/src/tools/memory_manage.rs

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<dyn MemoryRepository>,
}

impl MemoryManageTool {
    pub fn new(memories: Arc<dyn MemoryRepository>) -> 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<ToolResult> {
        Ok(error_result("memory_manage requires tool context"))
    }

    async fn execute_with_context(
        &self,
        context: &ToolContext,
        args: serde_json::Value,
    ) -> anyhow::Result<ToolResult> {
        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<MemoryUpsert, ToolResult> {
    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<String, ToolResult> {
    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"));
    }
}