1use super::{AgentTool, AgentToolResult, MemoryItem, ToolContext, ToolError};
4use crate::tools::typed::TypedTool;
5use async_trait::async_trait;
6use schemars::JsonSchema;
7use serde::Deserialize;
8use serde_json::{Value, json};
9
10#[allow(dead_code)]
12const DEFAULT_LIMIT: usize = 5;
13const MAX_LIMIT: usize = 20;
15
16#[derive(Deserialize, JsonSchema)]
19pub struct MemoryRecallArgs {
20 query: String,
21 #[serde(default = "default_mem_limit")]
22 limit: u64,
23}
24
25fn default_mem_limit() -> u64 {
26 5
27}
28
29pub struct MemoryRecallTool;
32
33#[async_trait]
34impl AgentTool for MemoryRecallTool {
35 fn name(&self) -> &str {
36 "memory_recall"
37 }
38
39 fn label(&self) -> &str {
40 "Memory Recall"
41 }
42
43 fn description(&self) -> &str {
44 "Search long-term memory for information relevant to a query. \
45 Returns the most relevant stored memories (facts, preferences, \
46 context, summaries)."
47 }
48
49 fn essential(&self) -> bool {
50 false
51 }
52
53 fn parameters_schema(&self) -> Value {
54 json!({
55 "type": "object",
56 "properties": {
57 "query": {
58 "type": "string",
59 "description": "What to search for in memory."
60 },
61 "limit": {
62 "type": "integer",
63 "minimum": 1,
64 "maximum": 20,
65 "default": 5,
66 "description": "Maximum number of results to return."
67 }
68 },
69 "required": ["query"]
70 })
71 }
72
73 async fn execute(
74 &self,
75 _tool_call_id: &str,
76 params: Value,
77 _signal: Option<tokio::sync::oneshot::Receiver<()>>,
78 ctx: &ToolContext,
79 ) -> Result<AgentToolResult, ToolError> {
80 let args: MemoryRecallArgs =
81 serde_json::from_value(params).map_err(|e| format!("invalid params: {e}"))?;
82 self.execute_typed(_tool_call_id, args, _signal, ctx).await
83 }
84}
85
86#[async_trait]
87impl TypedTool for MemoryRecallTool {
88 type Args = MemoryRecallArgs;
89
90 async fn execute_typed(
91 &self,
92 _tool_call_id: &str,
93 args: Self::Args,
94 _signal: Option<tokio::sync::oneshot::Receiver<()>>,
95 ctx: &ToolContext,
96 ) -> Result<AgentToolResult, ToolError> {
97 let backend = ctx.memory.as_ref().ok_or("Memory not configured")?;
98 let limit = (args.limit as usize).clamp(1, MAX_LIMIT);
99 let results = backend.search(&args.query, limit).await?;
100 Ok(AgentToolResult::success(format_results(&results)))
101 }
102}
103
104fn format_results(items: &[MemoryItem]) -> String {
106 if items.is_empty() {
107 return "No matching memories found.".to_string();
108 }
109 let mut out = format!(
110 "Found {} memor{}:\n\n",
111 items.len(),
112 if items.len() == 1 { "y" } else { "ies" }
113 );
114 for (i, item) in items.iter().enumerate() {
115 out.push_str(&format!("{}. [{}] {}\n", i + 1, item.kind, item.content));
116 }
117 out
118}
119
120#[cfg(test)]
121mod tests {
122 use super::*;
123 use crate::tools::MemoryBackend;
124 use parking_lot::Mutex;
125 use std::future::Future;
126 use std::pin::Pin;
127 use std::sync::Arc;
128
129 #[derive(Debug)]
131 struct MockMemory {
132 items: Vec<MemoryItem>,
133 last_k: Mutex<Option<usize>>,
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 Box::pin(async move { Ok("mem-1".to_string()) })
144 }
145
146 fn search<'a>(
147 &'a self,
148 _query: &'a str,
149 k: usize,
150 ) -> Pin<Box<dyn Future<Output = Result<Vec<MemoryItem>, ToolError>> + Send + 'a>> {
151 *self.last_k.lock() = Some(k);
152 let items: Vec<MemoryItem> = self.items.iter().take(k).cloned().collect();
153 Box::pin(async move { Ok(items) })
154 }
155
156 fn list<'a>(
157 &'a self,
158 _subject: &'a str,
159 ) -> Pin<Box<dyn Future<Output = Result<Vec<MemoryItem>, ToolError>> + Send + 'a>> {
160 Box::pin(async move { Ok(vec![]) })
161 }
162
163 fn delete<'a>(
164 &'a self,
165 _id: &'a str,
166 ) -> Pin<Box<dyn Future<Output = Result<(), ToolError>> + Send + 'a>> {
167 Box::pin(async move { Ok(()) })
168 }
169 }
170
171 fn make_item(id: &str, kind: &str, content: &str) -> MemoryItem {
172 MemoryItem {
173 id: id.into(),
174 kind: kind.into(),
175 content: content.into(),
176 subject: "s".into(),
177 }
178 }
179
180 #[tokio::test]
181 async fn recall_returns_formatted_results() {
182 let mock = Arc::new(MockMemory {
183 items: vec![
184 make_item("1", "fact", "Rust is fast"),
185 make_item("2", "preference", "Likes dark mode"),
186 ],
187 last_k: Mutex::new(None),
188 });
189 let ctx = ToolContext::default().with_memory(mock.clone());
190 let result = MemoryRecallTool
191 .execute("c1", json!({"query": "rust", "limit": 5}), None, &ctx)
192 .await
193 .unwrap();
194 assert!(result.success);
195 assert!(result.output.contains("[fact] Rust is fast"));
196 assert!(result.output.contains("[preference] Likes dark mode"));
197 assert_eq!(*mock.last_k.lock(), Some(5));
198 }
199
200 #[tokio::test]
201 async fn recall_reports_empty_results() {
202 let mock = Arc::new(MockMemory {
203 items: vec![],
204 last_k: Mutex::new(None),
205 });
206 let ctx = ToolContext::default().with_memory(mock);
207 let result = MemoryRecallTool
208 .execute("c1", json!({"query": "nothing"}), None, &ctx)
209 .await
210 .unwrap();
211 assert!(result.success);
212 assert_eq!(result.output, "No matching memories found.");
213 }
214
215 #[tokio::test]
216 async fn recall_uses_default_limit() {
217 let mock = Arc::new(MockMemory {
218 items: vec![],
219 last_k: Mutex::new(None),
220 });
221 let ctx = ToolContext::default().with_memory(mock.clone());
222 MemoryRecallTool
223 .execute("c1", json!({"query": "x"}), None, &ctx)
224 .await
225 .unwrap();
226 assert_eq!(*mock.last_k.lock(), Some(DEFAULT_LIMIT));
227 }
228
229 #[tokio::test]
230 async fn recall_clamps_oversized_limit() {
231 let mock = Arc::new(MockMemory {
232 items: vec![],
233 last_k: Mutex::new(None),
234 });
235 let ctx = ToolContext::default().with_memory(mock.clone());
236 MemoryRecallTool
237 .execute("c1", json!({"query": "x", "limit": 100}), None, &ctx)
238 .await
239 .unwrap();
240 assert_eq!(*mock.last_k.lock(), Some(MAX_LIMIT));
241 }
242
243 #[tokio::test]
244 async fn recall_clamps_zero_limit() {
245 let mock = Arc::new(MockMemory {
246 items: vec![],
247 last_k: Mutex::new(None),
248 });
249 let ctx = ToolContext::default().with_memory(mock.clone());
250 MemoryRecallTool
251 .execute("c1", json!({"query": "x", "limit": 0}), None, &ctx)
252 .await
253 .unwrap();
254 assert_eq!(*mock.last_k.lock(), Some(1));
255 }
256
257 #[tokio::test]
258 async fn recall_errors_when_memory_not_configured() {
259 let ctx = ToolContext::default();
260 let err = MemoryRecallTool
261 .execute("c1", json!({"query": "x"}), None, &ctx)
262 .await
263 .unwrap_err();
264 assert_eq!(err, "Memory not configured");
265 }
266
267 #[tokio::test]
268 async fn recall_rejects_missing_query() {
269 let mock = Arc::new(MockMemory {
270 items: vec![],
271 last_k: Mutex::new(None),
272 });
273 let ctx = ToolContext::default().with_memory(mock);
274 let err = MemoryRecallTool
275 .execute("c1", json!({"limit": 3}), None, &ctx)
276 .await
277 .unwrap_err();
278 assert!(err.contains("query"));
279 }
280}