aegis_tools/
memory_search.rs

1use crate::registry::{Tool, ToolContext};
2use aegis_memory::MemoryGraph;
3use anyhow::Result;
4use async_trait::async_trait;
5use serde_json::{json, Value};
6use std::sync::{Arc, Mutex};
7
8pub struct MemorySearchTool {
9    pub graph: Arc<Mutex<MemoryGraph>>,
10}
11
12#[async_trait]
13impl Tool for MemorySearchTool {
14    fn name(&self) -> &str {
15        "memory_search"
16    }
17    fn description(&self) -> &str {
18        "Search, forget, or update agent memory. Actions: search (default), forget (delete matching memories), update (replace matching memory with new content)"
19    }
20    fn parameters(&self) -> Value {
21        json!({
22            "type": "object",
23            "properties": {
24                "query": { "type": "string", "description": "Keyword/text to search, or exact memory ID (mem-xxx) for precise deletion" },
25                "limit": { "type": "integer", "description": "Max results (default 5)" },
26                "action": { "type": "string", "enum": ["search", "forget", "update"], "description": "Action: search (default), forget (delete matches), update (replace with new_content)" },
27                "new_content": { "type": "string", "description": "New content to replace matched memory (required for action=update)" }
28            },
29            "required": ["query"]
30        })
31    }
32    async fn execute(&self, args: Value, _ctx: &ToolContext<'_>) -> Result<String> {
33        let query = args["query"].as_str().unwrap_or("");
34        let limit = args["limit"].as_u64().unwrap_or(5) as usize;
35        let action = args["action"].as_str().unwrap_or("search");
36
37        let mut graph = self.graph.lock().map_err(|e| anyhow::anyhow!("lock error: {e}"))?;
38
39        match action {
40            "forget" => {
41                // If query looks like an ID, delete precisely; otherwise search and batch delete
42                if query.starts_with("mem-") || (query.len() > 8 && query.chars().all(|c| c.is_ascii_hexdigit() || c == '-')) {
43                    if graph.forget(query) {
44                        Ok(format!("Deleted memory: {query}"))
45                    } else {
46                        Ok(format!("No memory found with id: {query}"))
47                    }
48                } else {
49                    let matches: Vec<String> = graph.search(query, limit)
50                        .iter()
51                        .map(|e| e.id.clone())
52                        .collect();
53                    if matches.is_empty() {
54                        return Ok("No matching memories found to forget.".to_string());
55                    }
56                    let count = matches.len();
57                    for id in &matches {
58                        graph.forget(id);
59                    }
60                    Ok(format!("Deleted {count} matching memories."))
61                }
62            }
63            "update" => {
64                let new_content = args["new_content"].as_str().unwrap_or("");
65                if new_content.is_empty() {
66                    return Ok("Error: new_content is required for action=update".to_string());
67                }
68                // Find best match, supersede it with new content
69                let matches: Vec<String> = graph.search(query, 1)
70                    .iter()
71                    .map(|e| e.id.clone())
72                    .collect();
73                if matches.is_empty() {
74                    // No existing memory to update — create new
75                    let id = format!(
76                        "mem-{:x}",
77                        std::time::SystemTime::now()
78                            .duration_since(std::time::UNIX_EPOCH)
79                            .map(|d| d.as_nanos())
80                            .unwrap_or(0)
81                    );
82                    graph.insert(aegis_memory::MemoryEntry::new(
83                        &id, new_content, aegis_memory::MemoryCategory::Fact, "user",
84                    ));
85                    Ok(format!("No existing memory matched; created new: {new_content}"))
86                } else {
87                    let old_id = &matches[0];
88                    // Deactivate old
89                    graph.deactivate(old_id);
90                    // Insert new
91                    let new_id = format!(
92                        "mem-{:x}",
93                        std::time::SystemTime::now()
94                            .duration_since(std::time::UNIX_EPOCH)
95                            .map(|d| d.as_nanos())
96                            .unwrap_or(0)
97                    );
98                    graph.insert(aegis_memory::MemoryEntry::new(
99                        &new_id, new_content, aegis_memory::MemoryCategory::Fact, "user",
100                    ));
101                    graph.supersede(old_id, &new_id);
102                    Ok(format!("Updated memory: superseded old with: {new_content}"))
103                }
104            }
105            _ => {
106                // Default: search
107                let results = graph.search(query, limit);
108                if results.is_empty() {
109                    return Ok("No memories found.".to_string());
110                }
111                let lines: Vec<String> = results
112                    .iter()
113                    .enumerate()
114                    .map(|(i, e)| format!("{}. [{}] (id:{}) {}", i + 1, e.category, e.id, e.content))
115                    .collect();
116                Ok(lines.join("\n"))
117            }
118        }
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125    use aegis_memory::{MemoryCategory, MemoryEntry};
126
127    #[tokio::test]
128    async fn test_memory_search_tool() {
129        let mut graph = MemoryGraph::new();
130        graph.insert(MemoryEntry::new("m1", "Rust is a systems language", MemoryCategory::Fact, "user"));
131
132        let tool = MemorySearchTool {
133            graph: Arc::new(Mutex::new(graph)),
134        };
135
136        let ctx = ToolContext {
137            cwd: std::path::PathBuf::from("/tmp"),
138            session_id: "test".to_string(),
139            approve_fn: &|_| true,
140            yolo: true,
141            identity: None,
142            sandbox_enabled: false,
143        };
144
145        let result = tool
146            .execute(json!({"query": "rust"}), &ctx)
147            .await
148            .unwrap();
149        assert!(result.contains("Rust is a systems language"));
150    }
151
152    #[tokio::test]
153    async fn test_memory_forget_by_query() {
154        let mut graph = MemoryGraph::new();
155        graph.insert(MemoryEntry::new("m1", "old password is abc123", MemoryCategory::Fact, "user"));
156        graph.insert(MemoryEntry::new("m2", "favorite color is blue", MemoryCategory::Fact, "user"));
157
158        let tool = MemorySearchTool {
159            graph: Arc::new(Mutex::new(graph)),
160        };
161        let ctx = ToolContext {
162            cwd: std::path::PathBuf::from("/tmp"),
163            session_id: "test".to_string(),
164            approve_fn: &|_| true,
165            yolo: true,
166            identity: None,
167            sandbox_enabled: false,
168        };
169
170        let result = tool.execute(json!({"query": "password", "action": "forget"}), &ctx).await.unwrap();
171        assert!(result.contains("Deleted 1"));
172
173        // Verify it's gone
174        let search = tool.execute(json!({"query": "password"}), &ctx).await.unwrap();
175        assert!(search.contains("No memories found"));
176    }
177
178    #[tokio::test]
179    async fn test_memory_update() {
180        let mut graph = MemoryGraph::new();
181        graph.insert(MemoryEntry::new("m1", "password is abc123", MemoryCategory::Fact, "user"));
182
183        let tool = MemorySearchTool {
184            graph: Arc::new(Mutex::new(graph)),
185        };
186        let ctx = ToolContext {
187            cwd: std::path::PathBuf::from("/tmp"),
188            session_id: "test".to_string(),
189            approve_fn: &|_| true,
190            yolo: true,
191            identity: None,
192            sandbox_enabled: false,
193        };
194
195        let result = tool.execute(json!({
196            "query": "password",
197            "action": "update",
198            "new_content": "password is xyz789"
199        }), &ctx).await.unwrap();
200        assert!(result.contains("Updated memory"));
201
202        // Verify new content is searchable
203        let search = tool.execute(json!({"query": "xyz789"}), &ctx).await.unwrap();
204        assert!(search.contains("xyz789"));
205    }
206}
aegis_tools/memory_search.rs

aegis_tools/
memory_search.rs
