1use async_trait::async_trait;
4use serde_json::{Value, json};
5
6use super::{AgentTool, AgentToolResult, ToolContext, ToolError};
7
8pub struct MemoryEditTool;
15
16#[async_trait]
17impl AgentTool for MemoryEditTool {
18 fn name(&self) -> &str {
19 "memory_edit"
20 }
21
22 fn label(&self) -> &str {
23 "Memory Edit"
24 }
25
26 fn description(&self) -> &str {
27 "Update or delete a memory item. \
28 Provide a new `content` to update the item, or omit it to delete."
29 }
30
31 fn essential(&self) -> bool {
32 false
33 }
34
35 fn parameters_schema(&self) -> Value {
36 json!({
37 "type": "object",
38 "properties": {
39 "id": {
40 "type": "string",
41 "description": "The ID of the memory item to edit."
42 },
43 "content": {
44 "type": "string",
45 "description": "New content for the memory item. If omitted, the item is deleted."
46 },
47 "kind": {
48 "type": "string",
49 "description": "Category of the memory (only used when updating)."
50 },
51 "subject": {
52 "type": "string",
53 "description": "Subject scope for the memory (only used when updating)."
54 }
55 },
56 "required": ["id"]
57 })
58 }
59
60 async fn execute(
61 &self,
62 _tool_call_id: &str,
63 params: Value,
64 _signal: Option<tokio::sync::oneshot::Receiver<()>>,
65 ctx: &ToolContext,
66 ) -> Result<AgentToolResult, ToolError> {
67 let backend = ctx.memory.as_ref().ok_or("Memory not configured")?;
68
69 let id = params
70 .get("id")
71 .and_then(|v| v.as_str())
72 .ok_or("Missing required parameter: id")?;
73
74 if let Some(content) = params.get("content").and_then(|v| v.as_str()) {
75 let kind = params
77 .get("kind")
78 .and_then(|v| v.as_str())
79 .unwrap_or("fact");
80 let subject = params
81 .get("subject")
82 .and_then(|v| v.as_str())
83 .or(ctx.session_id.as_deref())
84 .unwrap_or("default");
85
86 let new_id = backend.put(content, kind, subject).await?;
87 Ok(AgentToolResult::success(format!(
88 "Updated memory item (old id: {}, new id: {}).",
89 id, new_id
90 )))
91 } else {
92 backend.delete(id).await?;
94 Ok(AgentToolResult::success(format!(
95 "Deleted memory item (id: {}).",
96 id
97 )))
98 }
99 }
100}
101
102#[cfg(test)]
103mod tests {
104 use super::*;
105 use crate::tools::MemoryBackend;
106 use parking_lot::Mutex;
107 use std::future::Future;
108 use std::pin::Pin;
109 use std::sync::Arc;
110
111 #[derive(Debug)]
113 struct MockMemory {
114 puts: Mutex<Vec<(String, String, String)>>,
115 deletes: Mutex<Vec<String>>,
116 }
117
118 impl MockMemory {
119 fn new() -> Self {
120 Self {
121 puts: Mutex::new(vec![]),
122 deletes: Mutex::new(vec![]),
123 }
124 }
125 }
126
127 impl MemoryBackend for MockMemory {
128 fn put<'a>(
129 &'a self,
130 content: &'a str,
131 kind: &'a str,
132 subject: &'a str,
133 ) -> Pin<Box<dyn Future<Output = Result<String, ToolError>> + Send + 'a>> {
134 self.puts
135 .lock()
136 .push((content.into(), kind.into(), subject.into()));
137 Box::pin(async move { Ok("new-id".to_string()) })
138 }
139
140 fn search<'a>(
141 &'a self,
142 _query: &'a str,
143 _k: usize,
144 ) -> Pin<
145 Box<dyn Future<Output = Result<Vec<crate::tools::MemoryItem>, ToolError>> + Send + 'a>,
146 > {
147 Box::pin(async move { Ok(vec![]) })
148 }
149
150 fn list<'a>(
151 &'a self,
152 _subject: &'a str,
153 ) -> Pin<
154 Box<dyn Future<Output = Result<Vec<crate::tools::MemoryItem>, ToolError>> + Send + 'a>,
155 > {
156 Box::pin(async move { Ok(vec![]) })
157 }
158
159 fn delete<'a>(
160 &'a self,
161 id: &'a str,
162 ) -> Pin<Box<dyn Future<Output = Result<(), ToolError>> + Send + 'a>> {
163 self.deletes.lock().push(id.into());
164 Box::pin(async move { Ok(()) })
165 }
166 }
167
168 #[tokio::test]
169 async fn edit_update_calls_put_with_correct_args() {
170 let mock = Arc::new(MockMemory::new());
171 let ctx = ToolContext::default()
172 .with_session("sess-42")
173 .with_memory(mock.clone());
174 let result = MemoryEditTool
175 .execute(
176 "c1",
177 json!({"id": "mem-1", "content": "updated", "kind": "fact"}),
178 None,
179 &ctx,
180 )
181 .await
182 .unwrap();
183 assert!(result.success);
184 assert_eq!(
185 result.output,
186 "Updated memory item (old id: mem-1, new id: new-id)."
187 );
188 let puts = mock.puts.lock();
189 assert_eq!(puts.len(), 1);
190 assert_eq!(puts[0].0, "updated");
191 assert_eq!(puts[0].1, "fact");
192 assert_eq!(puts[0].2, "sess-42");
193 assert_eq!(mock.deletes.lock().len(), 0);
194 }
195
196 #[tokio::test]
197 async fn edit_delete_calls_delete() {
198 let mock = Arc::new(MockMemory::new());
199 let ctx = ToolContext::default().with_memory(mock.clone());
200 let result = MemoryEditTool
201 .execute("c1", json!({"id": "mem-1"}), None, &ctx)
202 .await
203 .unwrap();
204 assert!(result.success);
205 assert_eq!(result.output, "Deleted memory item (id: mem-1).");
206 assert_eq!(mock.deletes.lock().len(), 1);
207 assert_eq!(mock.deletes.lock()[0], "mem-1");
208 assert_eq!(mock.puts.lock().len(), 0);
209 }
210
211 #[tokio::test]
212 async fn edit_update_defaults_kind_to_fact() {
213 let mock = Arc::new(MockMemory::new());
214 let ctx = ToolContext::default().with_memory(mock.clone());
215 MemoryEditTool
216 .execute("c1", json!({"id": "mem-1", "content": "x"}), None, &ctx)
217 .await
218 .unwrap();
219 assert_eq!(mock.puts.lock()[0].1, "fact");
220 }
221
222 #[tokio::test]
223 async fn edit_update_uses_default_subject_without_session() {
224 let mock = Arc::new(MockMemory::new());
225 let ctx = ToolContext::default().with_memory(mock.clone());
226 MemoryEditTool
227 .execute("c1", json!({"id": "mem-1", "content": "x"}), None, &ctx)
228 .await
229 .unwrap();
230 assert_eq!(mock.puts.lock()[0].2, "default");
231 }
232
233 #[tokio::test]
234 async fn edit_update_uses_explicit_subject() {
235 let mock = Arc::new(MockMemory::new());
236 let ctx = ToolContext::default()
237 .with_session("sess-42")
238 .with_memory(mock.clone());
239 MemoryEditTool
240 .execute(
241 "c1",
242 json!({"id": "mem-1", "content": "x", "subject": "custom-scope"}),
243 None,
244 &ctx,
245 )
246 .await
247 .unwrap();
248 assert_eq!(mock.puts.lock()[0].2, "custom-scope");
249 }
250
251 #[tokio::test]
252 async fn edit_errors_when_memory_not_configured() {
253 let ctx = ToolContext::default();
254 let err = MemoryEditTool
255 .execute("c1", json!({"id": "mem-1"}), None, &ctx)
256 .await
257 .unwrap_err();
258 assert_eq!(err, "Memory not configured");
259 }
260
261 #[tokio::test]
262 async fn edit_rejects_missing_id() {
263 let mock = Arc::new(MockMemory::new());
264 let ctx = ToolContext::default().with_memory(mock.clone());
265 let err = MemoryEditTool
266 .execute("c1", json!({"content": "x"}), None, &ctx)
267 .await
268 .unwrap_err();
269 assert!(err.contains("id"));
270 }
271}