Skip to main content

oxi_agent/tools/
memory_edit.rs

1//! `memory_edit` tool — update or delete a memory item.
2
3use 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/// Tool that updates an existing memory item or deletes it when no new content
12/// is supplied.
13///
14/// Requires `ctx.memory` to be set; otherwise returns an error. When `content`
15/// is provided the item is updated via `MemoryBackend::put`; when `content`
16/// is absent the item is deleted via `MemoryBackend::delete`.
17#[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    /// Records every `put` and `delete` call.
129    #[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}