oxi_agent/tools/
memory_retain.rs

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

oxi_agent/tools/
memory_retain.rs
