microclaw_engine/tools/
structured_memory.rs

1use async_trait::async_trait;
2use serde_json::json;
3use std::sync::Arc;
4use tracing::info;
5
6use crate::internal::storage::db::Database;
7use crate::internal::storage::db::Memory;
8use crate::memory_backend::MemoryBackend;
9use microclaw_core::llm_types::ToolDefinition;
10
11use super::{auth_context_from_input, authorize_chat_access, schema_object, Tool, ToolResult};
12
13// ── Search ────────────────────────────────────────────────────────────────────
14
15pub struct StructuredMemorySearchTool {
16    memory_backend: Arc<MemoryBackend>,
17}
18
19impl StructuredMemorySearchTool {
20    pub fn new(db: Arc<Database>, memory_backend: Arc<MemoryBackend>) -> Self {
21        let _ = db;
22        Self { memory_backend }
23    }
24
25    fn filter_visible_memories(chat_id: i64, memories: Vec<Memory>) -> Vec<Memory> {
26        memories
27            .into_iter()
28            .filter(|m| m.chat_id.is_none() || m.chat_id == Some(chat_id))
29            .collect()
30    }
31}
32
33#[async_trait]
34impl Tool for StructuredMemorySearchTool {
35    fn name(&self) -> &str {
36        "structured_memory_search"
37    }
38
39    fn definition(&self) -> ToolDefinition {
40        ToolDefinition {
41            name: "structured_memory_search".into(),
42            description: "Search structured memories extracted from past conversations. Returns memories whose content contains the query string. Leave query empty to list recent visible memories.".into(),
43            input_schema: schema_object(
44                json!({
45                    "query": {
46                        "type": "string",
47                        "description": "Keyword(s) to search for in memory content; leave empty to list recent visible memories"
48                    },
49                    "limit": {
50                        "type": "integer",
51                        "description": "Maximum number of results to return (default 10, max 50)"
52                    },
53                    "include_archived": {
54                        "type": "boolean",
55                        "description": "Whether to include archived memories in results (default false)"
56                    }
57                }),
58                &[],
59            ),
60        }
61    }
62
63    async fn execute(&self, input: serde_json::Value) -> ToolResult {
64        let query = input
65            .get("query")
66            .and_then(|v| v.as_str())
67            .map(|q| q.trim().to_string())
68            .unwrap_or_default();
69        let limit = input
70            .get("limit")
71            .and_then(|v| v.as_u64())
72            .map(|n| n.min(50) as usize)
73            .unwrap_or(10);
74        let include_archived = input
75            .get("include_archived")
76            .and_then(|v| v.as_bool())
77            .unwrap_or(false);
78
79        let chat_id = auth_context_from_input(&input)
80            .map(|a| a.caller_chat_id)
81            .unwrap_or(0);
82
83        info!(
84            "structured_memory_search: query={query:?} chat_id={chat_id} limit={limit} include_archived={include_archived}"
85        );
86
87        let result = if query.is_empty() {
88            let mut memories = match self
89                .memory_backend
90                .get_all_memories_for_chat(Some(chat_id))
91                .await
92            {
93                Ok(m) => m,
94                Err(e) => return ToolResult::error(format!("Search failed: {e}")),
95            };
96            let mut global = match self.memory_backend.get_all_memories_for_chat(None).await {
97                Ok(m) => m,
98                Err(e) => return ToolResult::error(format!("Search failed: {e}")),
99            };
100            memories.append(&mut global);
101            memories.sort_by(|a, b| b.updated_at.cmp(&a.updated_at));
102            if !include_archived {
103                memories.retain(|m| !m.is_archived);
104            }
105            memories.truncate(limit);
106            Ok(memories)
107        } else {
108            self.memory_backend
109                .search_memories_with_options(chat_id, &query, limit, include_archived, true)
110                .await
111        };
112
113        match result {
114            Ok(memories) if memories.is_empty() => {
115                let message = if query.is_empty() {
116                    "No visible memories found.".to_string()
117                } else {
118                    "No memories found matching that query.".to_string()
119                };
120                ToolResult::success(message)
121            }
122            Ok(memories) => {
123                let original_count = memories.len();
124                let memories = Self::filter_visible_memories(chat_id, memories);
125                if memories.is_empty() {
126                    let message = if query.is_empty() {
127                        "No visible memories found.".to_string()
128                    } else {
129                        "No visible memories found matching that query.".to_string()
130                    };
131                    return ToolResult::success(message);
132                }
133                let filtered_count = original_count.saturating_sub(memories.len());
134                if filtered_count > 0 {
135                    info!(
136                        "structured_memory_search: filtered {} cross-chat memories for chat_id={}",
137                        filtered_count, chat_id
138                    );
139                }
140                let lines: Vec<String> = memories
141                    .iter()
142                    .map(|m| {
143                        let scope = if m.chat_id.is_none() {
144                            "global"
145                        } else {
146                            "chat"
147                        };
148                        format!("[id={}] [{}] [{}] {}", m.id, m.category, scope, m.content)
149                    })
150                    .collect();
151                ToolResult::success(lines.join("\n"))
152            }
153            Err(e) => ToolResult::error(format!("Search failed: {e}")),
154        }
155    }
156}
157
158// ── Delete ────────────────────────────────────────────────────────────────────
159
160pub struct StructuredMemoryDeleteTool {
161    memory_backend: Arc<MemoryBackend>,
162}
163
164impl StructuredMemoryDeleteTool {
165    pub fn new(db: Arc<Database>, memory_backend: Arc<MemoryBackend>) -> Self {
166        let _ = db;
167        Self { memory_backend }
168    }
169}
170
171#[async_trait]
172impl Tool for StructuredMemoryDeleteTool {
173    fn name(&self) -> &str {
174        "structured_memory_delete"
175    }
176
177    fn definition(&self) -> ToolDefinition {
178        ToolDefinition {
179            name: "structured_memory_delete".into(),
180            description: "Archive a structured memory by its id (soft delete). Use structured_memory_search first to find the id. You can only archive memories that belong to the current chat or global memories if you are a control chat.".into(),
181            input_schema: schema_object(
182                json!({
183                    "id": {
184                        "type": "integer",
185                        "description": "The id of the memory to delete"
186                    }
187                }),
188                &["id"],
189            ),
190        }
191    }
192
193    async fn execute(&self, input: serde_json::Value) -> ToolResult {
194        let id = match input.get("id").and_then(|v| v.as_i64()) {
195            Some(id) => id,
196            None => return ToolResult::error("Missing 'id' parameter".into()),
197        };
198
199        // Load memory first to check ownership
200        let mem = match self.memory_backend.get_memory_by_id(id).await {
201            Ok(Some(m)) => m,
202            Ok(None) => return ToolResult::error(format!("Memory id={id} not found")),
203            Err(e) => return ToolResult::error(format!("DB error: {e}")),
204        };
205
206        // Authorize: caller must own the chat or be a control chat; global memories only by control
207        if let Some(auth) = auth_context_from_input(&input) {
208            match mem.chat_id {
209                Some(mem_chat_id) => {
210                    if let Err(e) = authorize_chat_access(&input, mem_chat_id) {
211                        return ToolResult::error(format!(
212                            "{e} (memory id={id}, owner_chat_id={mem_chat_id})"
213                        ));
214                    }
215                }
216                None => {
217                    // Global memory — requires control chat
218                    if !auth.is_control_chat() {
219                        return ToolResult::error(format!(
220                            "Permission denied: only control chats can delete global memories (caller: {}, memory id={id}, owner_scope=global)",
221                            auth.caller_chat_id
222                        ));
223                    }
224                }
225            }
226        }
227
228        info!("structured_memory_delete: id={id}");
229
230        match self.memory_backend.archive_memory(id).await {
231            Ok(true) => ToolResult::success(format!("Memory id={id} archived.")),
232            Ok(false) => ToolResult::error(format!("Memory id={id} not found")),
233            Err(e) => ToolResult::error(format!("Delete failed: {e}")),
234        }
235    }
236}
237
238// ── Update ────────────────────────────────────────────────────────────────────
239
240pub struct StructuredMemoryUpdateTool {
241    memory_backend: Arc<MemoryBackend>,
242}
243
244impl StructuredMemoryUpdateTool {
245    pub fn new(db: Arc<Database>, memory_backend: Arc<MemoryBackend>) -> Self {
246        let _ = db;
247        Self { memory_backend }
248    }
249}
250
251#[async_trait]
252impl Tool for StructuredMemoryUpdateTool {
253    fn name(&self) -> &str {
254        "structured_memory_update"
255    }
256
257    fn definition(&self) -> ToolDefinition {
258        ToolDefinition {
259            name: "structured_memory_update".into(),
260            description: "Update the content or category of an existing structured memory. Use this to correct outdated or wrong memories instead of creating a duplicate.".into(),
261            input_schema: schema_object(
262                json!({
263                    "id": {
264                        "type": "integer",
265                        "description": "The id of the memory to update"
266                    },
267                    "content": {
268                        "type": "string",
269                        "description": "New content for the memory (max 300 characters)"
270                    },
271                    "category": {
272                        "type": "string",
273                        "description": "Category: PROFILE, KNOWLEDGE, or EVENT",
274                        "enum": ["PROFILE", "KNOWLEDGE", "EVENT"]
275                    }
276                }),
277                &["id", "content"],
278            ),
279        }
280    }
281
282    async fn execute(&self, input: serde_json::Value) -> ToolResult {
283        let id = match input.get("id").and_then(|v| v.as_i64()) {
284            Some(id) => id,
285            None => return ToolResult::error("Missing 'id' parameter".into()),
286        };
287        let content = match input.get("content").and_then(|v| v.as_str()) {
288            Some(c) if !c.trim().is_empty() => c.trim().to_string(),
289            _ => return ToolResult::error("Missing or empty 'content' parameter".into()),
290        };
291        if content.len() > 300 {
292            return ToolResult::error("Content exceeds 300 character limit".into());
293        }
294
295        // Load memory first to check ownership and get current category
296        let mem = match self.memory_backend.get_memory_by_id(id).await {
297            Ok(Some(m)) => m,
298            Ok(None) => return ToolResult::error(format!("Memory id={id} not found")),
299            Err(e) => return ToolResult::error(format!("DB error: {e}")),
300        };
301
302        // Authorize same as delete
303        if let Some(auth) = auth_context_from_input(&input) {
304            match mem.chat_id {
305                Some(mem_chat_id) => {
306                    if let Err(e) = authorize_chat_access(&input, mem_chat_id) {
307                        return ToolResult::error(format!(
308                            "{e} (memory id={id}, owner_chat_id={mem_chat_id})"
309                        ));
310                    }
311                }
312                None => {
313                    if !auth.is_control_chat() {
314                        return ToolResult::error(format!(
315                            "Permission denied: only control chats can update global memories (caller: {}, memory id={id}, owner_scope=global)",
316                            auth.caller_chat_id
317                        ));
318                    }
319                }
320            }
321        }
322
323        let category = input
324            .get("category")
325            .and_then(|v| v.as_str())
326            .unwrap_or(&mem.category)
327            .to_string();
328
329        let valid_categories = ["PROFILE", "KNOWLEDGE", "EVENT"];
330        if !valid_categories.contains(&category.as_str()) {
331            return ToolResult::error(format!(
332                "Invalid category '{category}'. Must be one of: PROFILE, KNOWLEDGE, EVENT"
333            ));
334        }
335
336        info!("structured_memory_update: id={id}");
337
338        match self
339            .memory_backend
340            .update_memory_content(id, &content, &category)
341            .await
342        {
343            Ok(true) => ToolResult::success(format!("Memory id={id} updated.")),
344            Ok(false) => ToolResult::error(format!("Memory id={id} not found")),
345            Err(e) => ToolResult::error(format!("Update failed: {e}")),
346        }
347    }
348}
349
350#[cfg(test)]
351mod tests {
352    use super::*;
353    use serde_json::json;
354
355    fn test_db() -> Arc<Database> {
356        let dir = std::env::temp_dir().join(format!("mc_smem_{}", uuid::Uuid::new_v4()));
357        std::fs::create_dir_all(&dir).unwrap();
358        Arc::new(Database::new(dir.to_str().unwrap()).unwrap())
359    }
360
361    fn test_backend(db: Arc<Database>) -> Arc<MemoryBackend> {
362        Arc::new(MemoryBackend::local_only(db))
363    }
364
365    #[tokio::test]
366    async fn test_search_returns_results() {
367        let db = test_db();
368        db.insert_memory(Some(100), "User loves Rust programming", "PROFILE")
369            .unwrap();
370        db.insert_memory(Some(100), "User likes coffee", "PROFILE")
371            .unwrap();
372        let tool = StructuredMemorySearchTool::new(db.clone(), test_backend(db));
373        let result = tool
374            .execute(json!({
375                "query": "rust",
376                "__microclaw_auth": {"caller_chat_id": 100, "control_chat_ids": []}
377            }))
378            .await;
379        assert!(!result.is_error);
380        assert!(result.content.contains("Rust"));
381        assert!(!result.content.contains("coffee"));
382    }
383
384    #[tokio::test]
385    async fn test_search_empty_query_lists_no_visible_memories() {
386        let db = test_db();
387        let tool = StructuredMemorySearchTool::new(db.clone(), test_backend(db));
388        let result = tool.execute(json!({"query": "  "})).await;
389        assert!(!result.is_error);
390        assert!(result.content.contains("No visible memories"));
391    }
392
393    #[tokio::test]
394    async fn test_delete_own_chat_memory() {
395        let db = test_db();
396        let id = db.insert_memory(Some(100), "to delete", "EVENT").unwrap();
397        let tool = StructuredMemoryDeleteTool::new(db.clone(), test_backend(db.clone()));
398        let result = tool
399            .execute(json!({
400                "id": id,
401                "__microclaw_auth": {"caller_chat_id": 100, "control_chat_ids": []}
402            }))
403            .await;
404        assert!(!result.is_error, "{}", result.content);
405        let mem = db.get_memory_by_id(id).unwrap().unwrap();
406        assert!(mem.is_archived);
407    }
408
409    #[tokio::test]
410    async fn test_delete_other_chat_denied() {
411        let db = test_db();
412        let id = db.insert_memory(Some(200), "other chat", "EVENT").unwrap();
413        let tool = StructuredMemoryDeleteTool::new(db.clone(), test_backend(db));
414        let result = tool
415            .execute(json!({
416                "id": id,
417                "__microclaw_auth": {"caller_chat_id": 100, "control_chat_ids": []}
418            }))
419            .await;
420        assert!(result.is_error);
421        assert!(result.content.contains("Permission denied"));
422    }
423
424    #[tokio::test]
425    async fn test_search_empty_query_lists_recent_visible_memories() {
426        let db = test_db();
427        db.insert_memory(Some(100), "chat memory", "PROFILE")
428            .unwrap();
429        db.insert_memory(None, "global memory", "KNOWLEDGE")
430            .unwrap();
431        let tool = StructuredMemorySearchTool::new(db.clone(), test_backend(db));
432        let result = tool
433            .execute(json!({
434                "limit": 10,
435                "__microclaw_auth": {"caller_chat_id": 100, "control_chat_ids": []}
436            }))
437            .await;
438        assert!(!result.is_error, "{}", result.content);
439        assert!(result.content.contains("chat memory"));
440        assert!(result.content.contains("global memory"));
441    }
442
443    #[tokio::test]
444    async fn test_update_memory() {
445        let db = test_db();
446        let id = db
447            .insert_memory(Some(100), "User lives in Tokyo", "PROFILE")
448            .unwrap();
449        let tool = StructuredMemoryUpdateTool::new(db.clone(), test_backend(db.clone()));
450        let result = tool
451            .execute(json!({
452                "id": id,
453                "content": "User lives in Osaka",
454                "__microclaw_auth": {"caller_chat_id": 100, "control_chat_ids": []}
455            }))
456            .await;
457        assert!(!result.is_error, "{}", result.content);
458        let mem = db.get_memory_by_id(id).unwrap().unwrap();
459        assert_eq!(mem.content, "User lives in Osaka");
460    }
461
462    #[tokio::test]
463    async fn test_update_content_too_long() {
464        let db = test_db();
465        let id = db.insert_memory(Some(100), "short", "EVENT").unwrap();
466        let tool = StructuredMemoryUpdateTool::new(db.clone(), test_backend(db));
467        let long = "x".repeat(301);
468        let result = tool
469            .execute(json!({
470                "id": id,
471                "content": long,
472                "__microclaw_auth": {"caller_chat_id": 100, "control_chat_ids": []}
473            }))
474            .await;
475        assert!(result.is_error);
476        assert!(result.content.contains("300 character"));
477    }
478}
microclaw_engine/tools/structured_memory.rs

microclaw_engine/tools/
structured_memory.rs
