1use async_trait::async_trait;
4use serde_json::{Value, json};
5
6use super::{AgentTool, AgentToolResult, ToolContext, ToolError};
7
8const VALID_KINDS: [&str; 4] = ["fact", "preference", "context", "summary"];
10
11pub 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 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 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 #[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}