Skip to main content

zeroclaw_tools/
memory_purge.rs

1use async_trait::async_trait;
2use serde_json::json;
3use std::sync::Arc;
4use zeroclaw_api::tool::{Tool, ToolResult};
5use zeroclaw_config::policy::SecurityPolicy;
6use zeroclaw_config::policy::ToolOperation;
7use zeroclaw_memory::Memory;
8
9/// Let the agent bulk-delete memories by namespace or session
10pub struct MemoryPurgeTool {
11    memory: Arc<dyn Memory>,
12    security: Arc<SecurityPolicy>,
13}
14
15impl MemoryPurgeTool {
16    pub fn new(memory: Arc<dyn Memory>, security: Arc<SecurityPolicy>) -> Self {
17        Self { memory, security }
18    }
19}
20
21#[async_trait]
22impl Tool for MemoryPurgeTool {
23    fn name(&self) -> &str {
24        "memory_purge"
25    }
26
27    fn description(&self) -> &str {
28        "Remove all memories in a namespace or session. Use to bulk-delete per-tenant or per-conversation data. Returns the number of deleted entries. WARNING: This operation cannot be undone."
29    }
30
31    fn parameters_schema(&self) -> serde_json::Value {
32        json!({
33            "type": "object",
34            "properties": {
35                "namespace": {
36                    "type": "string",
37                    "description": "The namespace to purge. Deletes all memories whose namespace field equals this value."
38                },
39                "session_id": {
40                    "type": "string",
41                    "description": "The session ID to purge. Deletes all memories in this session."
42                }
43            },
44            "minProperties": 1
45        })
46    }
47
48    async fn execute(&self, args: serde_json::Value) -> anyhow::Result<ToolResult> {
49        let namespace = args.get("namespace").and_then(|v| v.as_str());
50        let session_id = args.get("session_id").and_then(|v| v.as_str());
51
52        if namespace.is_none() && session_id.is_none() {
53            ::zeroclaw_log::record!(
54                WARN,
55                ::zeroclaw_log::Event::new(module_path!(), ::zeroclaw_log::Action::Reject)
56                    .with_outcome(::zeroclaw_log::EventOutcome::Failure)
57                    .with_attrs(::serde_json::json!({"missing": "namespace_or_session_id"})),
58                "memory_purge: must provide namespace or session_id"
59            );
60            return Err(anyhow::Error::msg(
61                "Must provide either 'namespace' or 'session_id' parameter",
62            ));
63        }
64
65        if let Err(error) = self
66            .security
67            .enforce_tool_operation(ToolOperation::Act, "memory_purge")
68        {
69            return Ok(ToolResult {
70                success: false,
71                output: String::new(),
72                error: Some(error),
73            });
74        }
75
76        let mut total_purged = 0;
77        let mut output_parts = Vec::new();
78
79        if let Some(ns) = namespace {
80            match self.memory.purge_namespace(ns).await {
81                Ok(count) => {
82                    total_purged += count;
83                    output_parts.push(format!("Purged {count} memories from namespace '{ns}'"));
84                }
85                Err(e) => {
86                    return Ok(ToolResult {
87                        success: false,
88                        output: String::new(),
89                        error: Some(format!("Failed to purge namespace: {e}")),
90                    });
91                }
92            }
93        }
94
95        if let Some(sid) = session_id {
96            match self.memory.purge_session(sid).await {
97                Ok(count) => {
98                    total_purged += count;
99                    output_parts.push(format!("Purged {count} memories from session '{sid}'"));
100                }
101                Err(e) => {
102                    return Ok(ToolResult {
103                        success: false,
104                        output: String::new(),
105                        error: Some(format!("Failed to purge session: {e}")),
106                    });
107                }
108            }
109        }
110
111        Ok(ToolResult {
112            success: true,
113            output: if output_parts.is_empty() {
114                format!("Purged {total_purged} memories")
115            } else {
116                output_parts.join("; ")
117            },
118            error: None,
119        })
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126    use tempfile::TempDir;
127    use zeroclaw_config::autonomy::AutonomyLevel;
128    use zeroclaw_config::policy::SecurityPolicy;
129    use zeroclaw_memory::{MemoryCategory, MemoryEntry, SqliteMemory};
130
131    fn test_security() -> Arc<SecurityPolicy> {
132        Arc::new(SecurityPolicy::default())
133    }
134
135    fn test_mem() -> (TempDir, Arc<dyn Memory>) {
136        let tmp = TempDir::new().unwrap();
137        let mem = SqliteMemory::new("test", tmp.path()).unwrap();
138        (tmp, Arc::new(mem))
139    }
140
141    #[test]
142    fn name_and_schema() {
143        let (_tmp, mem) = test_mem();
144        let tool = MemoryPurgeTool::new(mem, test_security());
145        assert_eq!(tool.name(), "memory_purge");
146        assert!(tool.parameters_schema()["properties"]["namespace"].is_object());
147        assert!(tool.parameters_schema()["properties"]["session_id"].is_object());
148    }
149
150    #[tokio::test]
151    async fn purge_namespace_removes_only_all_matching_memories() {
152        let (_tmp, mem) = test_mem();
153
154        mem.store_with_metadata("a", "data", MemoryCategory::Core, None, Some("ns1"), None)
155            .await
156            .unwrap();
157        mem.store_with_metadata("b", "data", MemoryCategory::Core, None, Some("ns2"), None)
158            .await
159            .unwrap();
160
161        let in_ns1 =
162            |entries: &[MemoryEntry]| entries.iter().filter(|e| e.namespace == "ns1").count();
163
164        let before = mem.list(None, None).await.unwrap();
165        let tool = MemoryPurgeTool::new(mem.clone(), test_security());
166        let result = tool.execute(json!({"namespace": "ns1"})).await.unwrap();
167        let after = mem.list(None, None).await.unwrap();
168
169        assert!(result.success);
170        assert_eq!(in_ns1(&after), 0);
171        assert_eq!(after.len() - in_ns1(&after), before.len() - in_ns1(&before));
172    }
173
174    #[tokio::test]
175    async fn purge_session_removes_all_memories() {
176        let (_tmp, mem) = test_mem();
177        mem.store("a1", "data1", MemoryCategory::Core, Some("sess-x"))
178            .await
179            .unwrap();
180        mem.store("a2", "data2", MemoryCategory::Core, Some("sess-x"))
181            .await
182            .unwrap();
183        mem.store("b1", "data3", MemoryCategory::Core, Some("sess-y"))
184            .await
185            .unwrap();
186
187        let tool = MemoryPurgeTool::new(mem.clone(), test_security());
188        let result = tool.execute(json!({"session_id": "sess-x"})).await.unwrap();
189        assert!(result.success);
190        assert!(result.output.contains("2 memories"));
191
192        assert_eq!(mem.count().await.unwrap(), 1);
193    }
194
195    #[tokio::test]
196    async fn purge_namespace_nonexistent_is_noop() {
197        let (_tmp, mem) = test_mem();
198        mem.store("a", "data", MemoryCategory::Core, None)
199            .await
200            .unwrap();
201
202        let tool = MemoryPurgeTool::new(mem.clone(), test_security());
203        let result = tool
204            .execute(json!({"namespace": "nonexistent"}))
205            .await
206            .unwrap();
207        assert!(result.success);
208        assert!(result.output.contains("0 memories"));
209
210        assert_eq!(mem.count().await.unwrap(), 1);
211    }
212
213    #[tokio::test]
214    async fn purge_session_nonexistent_is_noop() {
215        let (_tmp, mem) = test_mem();
216        mem.store("a", "data", MemoryCategory::Core, Some("sess"))
217            .await
218            .unwrap();
219
220        let tool = MemoryPurgeTool::new(mem.clone(), test_security());
221        let result = tool
222            .execute(json!({"session_id": "nonexistent"}))
223            .await
224            .unwrap();
225        assert!(result.success);
226        assert!(result.output.contains("0 memories"));
227
228        assert_eq!(mem.count().await.unwrap(), 1);
229    }
230
231    #[tokio::test]
232    async fn purge_missing_parameter() {
233        let (_tmp, mem) = test_mem();
234        let tool = MemoryPurgeTool::new(mem, test_security());
235        let result = tool.execute(json!({})).await;
236        assert!(result.is_err());
237    }
238
239    #[tokio::test]
240    async fn purge_blocked_in_readonly_mode() {
241        let (_tmp, mem) = test_mem();
242        mem.store_with_metadata(
243            "a",
244            "data",
245            MemoryCategory::Core,
246            None,
247            Some("test-ns"),
248            None,
249        )
250        .await
251        .unwrap();
252        let readonly = Arc::new(SecurityPolicy {
253            autonomy: AutonomyLevel::ReadOnly,
254            ..SecurityPolicy::default()
255        });
256        let tool = MemoryPurgeTool::new(mem.clone(), readonly);
257        let result = tool.execute(json!({"namespace": "test-ns"})).await.unwrap();
258        assert!(!result.success);
259        assert!(
260            result
261                .error
262                .as_deref()
263                .unwrap_or("")
264                .contains("read-only mode")
265        );
266        assert_eq!(mem.count().await.unwrap(), 1);
267    }
268
269    #[tokio::test]
270    async fn purge_blocked_when_rate_limited() {
271        let (_tmp, mem) = test_mem();
272        mem.store_with_metadata(
273            "a",
274            "data",
275            MemoryCategory::Core,
276            None,
277            Some("test-ns"),
278            None,
279        )
280        .await
281        .unwrap();
282        let limited = Arc::new(SecurityPolicy {
283            max_actions_per_hour: 0,
284            ..SecurityPolicy::default()
285        });
286        let tool = MemoryPurgeTool::new(mem.clone(), limited);
287        let result = tool.execute(json!({"namespace": "test-ns"})).await.unwrap();
288        assert!(!result.success);
289        assert!(
290            result
291                .error
292                .as_deref()
293                .unwrap_or("")
294                .contains("Rate limit exceeded")
295        );
296        assert_eq!(mem.count().await.unwrap(), 1);
297    }
298}