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