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
9pub 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}