Skip to main content

oxi_agent/tools/
memory_edit.rs

1//! `memory_edit` tool — update or delete a memory item.
2
3use async_trait::async_trait;
4use serde_json::{Value, json};
5
6use super::{AgentTool, AgentToolResult, ToolContext, ToolError};
7
8/// Tool that updates an existing memory item or deletes it when no new content
9/// is supplied.
10///
11/// Requires `ctx.memory` to be set; otherwise returns an error. When `content`
12/// is provided the item is updated via [`MemoryBackend::put`]; when `content`
13/// is absent the item is deleted via [`MemoryBackend::delete`].
14pub 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            // Update path: persist the new content.
76            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            // Delete path: content is absent.
93            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    /// Records every `put` and `delete` call.
112    #[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}