Skip to main content

oxicode_agent/tools/
memory_retain.rs

1//! `memory_retain` tool — persist a memory item to the backend.
2
3use async_trait::async_trait;
4use serde_json::{Value, json};
5
6use super::{AgentTool, AgentToolResult, ToolContext, ToolError};
7
8/// Valid memory kinds accepted by [`MemoryRetainTool`].
9const VALID_KINDS: [&str; 4] = ["fact", "preference", "context", "summary"];
10
11/// Tool that persists a memory item (content, kind, importance) to the
12/// configured `MemoryBackend`.
13///
14/// Requires `ctx.memory` to be set; otherwise returns an error. The memory
15/// is scoped to the current session (`ctx.session_id`), falling back to
16/// `"default"` when no session id is available.
17pub struct MemoryRetainTool;
18
19#[async_trait]
20impl AgentTool for MemoryRetainTool {
21    fn name(&self) -> &str {
22        "memory_retain"
23    }
24
25    fn label(&self) -> &str {
26        "Memory Retain"
27    }
28
29    fn description(&self) -> &str {
30        "Store a piece of information to long-term memory for later recall. \
31         Use for facts, preferences, context, or summaries worth remembering \
32         across sessions."
33    }
34
35    fn essential(&self) -> bool {
36        false
37    }
38
39    fn parameters_schema(&self) -> Value {
40        json!({
41            "type": "object",
42            "properties": {
43                "content": {
44                    "type": "string",
45                    "description": "The text to remember."
46                },
47                "kind": {
48                    "type": "string",
49                    "enum": ["fact", "preference", "context", "summary"],
50                    "default": "fact",
51                    "description": "Category of the memory."
52                },
53                "importance": {
54                    "type": "number",
55                    "minimum": 0.0,
56                    "maximum": 1.0,
57                    "default": 0.5,
58                    "description": "How important this memory is (0–1)."
59                }
60            },
61            "required": ["content"]
62        })
63    }
64
65    async fn execute(
66        &self,
67        _tool_call_id: &str,
68        params: Value,
69        _signal: Option<tokio::sync::oneshot::Receiver<()>>,
70        ctx: &ToolContext,
71    ) -> Result<AgentToolResult, ToolError> {
72        let backend = ctx.memory.as_ref().ok_or("Memory not configured")?;
73
74        let content = params
75            .get("content")
76            .and_then(|v| v.as_str())
77            .ok_or("Missing required parameter: content")?;
78
79        let kind = params
80            .get("kind")
81            .and_then(|v| v.as_str())
82            .unwrap_or("fact");
83        if !VALID_KINDS.contains(&kind) {
84            return Err(format!(
85                "Invalid kind '{}': expected one of {:?}",
86                kind, VALID_KINDS
87            ));
88        }
89
90        // `importance` is validated for forward-compatibility; the current
91        // `MemoryBackend::put` signature does not persist it.
92        if let Some(importance) = params.get("importance").and_then(|v| v.as_f64())
93            && !(0.0..=1.0).contains(&importance)
94        {
95            return Err(format!(
96                "importance must be between 0 and 1, got {}",
97                importance
98            ));
99        }
100
101        let subject = ctx.session_id.as_deref().unwrap_or("default");
102        let id = backend.put(content, kind, subject).await?;
103
104        // Plan §5.d: surface the Brain ID + scope the durable
105        // authority acknowledged. The tool must NOT claim success
106        // until the ledger append is acknowledged (which `put` does).
107        Ok(AgentToolResult::success(format!(
108            "Retained [{}] (Brain id: {}) to scope '{}'.",
109            kind, id, subject
110        )))
111    }
112}
113#[cfg(test)]
114mod tests {
115    use super::*;
116    use crate::tools::MemoryBackend;
117    use parking_lot::Mutex;
118    use std::future::Future;
119    use std::pin::Pin;
120    use std::sync::Arc;
121
122    /// Records every `put` call; the remaining trait methods are stubbed.
123    #[derive(Debug)]
124    struct MockMemory {
125        puts: Mutex<Vec<(String, String, String)>>,
126    }
127
128    impl MockMemory {
129        fn new() -> Self {
130            Self {
131                puts: Mutex::new(vec![]),
132            }
133        }
134    }
135
136    impl MemoryBackend for MockMemory {
137        fn put<'a>(
138            &'a self,
139            content: &'a str,
140            kind: &'a str,
141            subject: &'a str,
142        ) -> Pin<Box<dyn Future<Output = Result<String, ToolError>> + Send + 'a>> {
143            self.puts
144                .lock()
145                .push((content.into(), kind.into(), subject.into()));
146            Box::pin(async move { Ok("mem-1".to_string()) })
147        }
148
149        fn search<'a>(
150            &'a self,
151            _query: &'a str,
152            _k: usize,
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 list<'a>(
160            &'a self,
161            _subject: &'a str,
162        ) -> Pin<
163            Box<dyn Future<Output = Result<Vec<crate::tools::MemoryItem>, ToolError>> + Send + 'a>,
164        > {
165            Box::pin(async move { Ok(vec![]) })
166        }
167
168        fn delete<'a>(
169            &'a self,
170            _id: &'a str,
171        ) -> Pin<Box<dyn Future<Output = Result<(), ToolError>> + Send + 'a>> {
172            Box::pin(async move { Ok(()) })
173        }
174    }
175
176    #[tokio::test]
177    async fn retain_calls_put_with_correct_args() {
178        let mock = Arc::new(MockMemory::new());
179        let ctx = ToolContext::default()
180            .with_session("sess-42")
181            .with_memory(mock.clone());
182        let result = MemoryRetainTool
183            .execute(
184                "c1",
185                json!({"content": "hello", "kind": "fact", "importance": 0.9}),
186                None,
187                &ctx,
188            )
189            .await
190            .unwrap();
191        assert!(result.success);
192        assert_eq!(
193            result.output,
194            "Retained [fact] (Brain id: mem-1) to scope 'sess-42'."
195        );
196        let puts = mock.puts.lock();
197        assert_eq!(puts.len(), 1);
198        assert_eq!(puts[0].0, "hello");
199        assert_eq!(puts[0].1, "fact");
200        assert_eq!(puts[0].2, "sess-42");
201    }
202
203    #[tokio::test]
204    async fn retain_defaults_kind_to_fact() {
205        let mock = Arc::new(MockMemory::new());
206        let ctx = ToolContext::default().with_memory(mock.clone());
207        let result = MemoryRetainTool
208            .execute("c1", json!({"content": "x"}), None, &ctx)
209            .await
210            .unwrap();
211        assert_eq!(
212            result.output,
213            "Retained [fact] (Brain id: mem-1) to scope 'default'."
214        );
215        assert_eq!(mock.puts.lock()[0].1, "fact");
216    }
217
218    #[tokio::test]
219    async fn retain_uses_default_subject_without_session() {
220        let mock = Arc::new(MockMemory::new());
221        let ctx = ToolContext::default().with_memory(mock.clone());
222        MemoryRetainTool
223            .execute("c1", json!({"content": "x"}), None, &ctx)
224            .await
225            .unwrap();
226        assert_eq!(mock.puts.lock()[0].2, "default");
227    }
228
229    #[tokio::test]
230    async fn retain_errors_when_memory_not_configured() {
231        let ctx = ToolContext::default();
232        let err = MemoryRetainTool
233            .execute("c1", json!({"content": "x"}), None, &ctx)
234            .await
235            .unwrap_err();
236        assert_eq!(err, "Memory not configured");
237    }
238
239    #[tokio::test]
240    async fn retain_rejects_invalid_kind() {
241        let mock = Arc::new(MockMemory::new());
242        let ctx = ToolContext::default().with_memory(mock.clone());
243        let err = MemoryRetainTool
244            .execute("c1", json!({"content": "x", "kind": "bogus"}), None, &ctx)
245            .await
246            .unwrap_err();
247        assert!(err.contains("Invalid kind"));
248    }
249
250    #[tokio::test]
251    async fn retain_rejects_out_of_range_importance() {
252        let mock = Arc::new(MockMemory::new());
253        let ctx = ToolContext::default().with_memory(mock.clone());
254        let err = MemoryRetainTool
255            .execute("c1", json!({"content": "x", "importance": 1.5}), None, &ctx)
256            .await
257            .unwrap_err();
258        assert!(err.contains("importance"));
259    }
260
261    #[tokio::test]
262    async fn retain_rejects_missing_content() {
263        let mock = Arc::new(MockMemory::new());
264        let ctx = ToolContext::default().with_memory(mock.clone());
265        let err = MemoryRetainTool
266            .execute("c1", json!({"kind": "fact"}), None, &ctx)
267            .await
268            .unwrap_err();
269        assert!(err.contains("content"));
270    }
271}