telos_agent/tools/builtin/
memory.rs

1use async_trait::async_trait;
2use serde_json::{Value, json};
3use std::sync::{Arc, Mutex};
4
5use crate::error::AgentError;
6use crate::knowledge::memory::MemoryEntry;
7use crate::knowledge::memory::MemoryStore;
8use crate::tools::api::{PermissionDecision, Tool, ToolContext, ToolDefinition, ToolOutput};
9
10mod maintenance;
11pub use maintenance::MemoryMaintenanceTool;
12
13// ── MemoryRead ───────────────────────────────────────────
14
15pub struct MemoryReadTool {
16    store: Arc<Mutex<MemoryStore>>,
17}
18
19impl MemoryReadTool {
20    pub fn new(store: Arc<Mutex<MemoryStore>>) -> Self {
21        Self { store }
22    }
23}
24
25#[async_trait]
26impl Tool for MemoryReadTool {
27    fn definition(&self) -> ToolDefinition {
28        ToolDefinition {
29            name: "MemoryRead".into(),
30            description: "Read a memory entry by name. Returns the full content including body."
31                .into(),
32            input_schema: json!({"type":"object","properties":{"name":{"type":"string","description":"Name of the memory to read"}},"required":["name"]}),
33        }
34    }
35    async fn check_permission(
36        &self,
37        _: &Value,
38        _: &ToolContext,
39    ) -> Result<PermissionDecision, AgentError> {
40        Ok(PermissionDecision::Allow)
41    }
42    fn is_concurrency_safe(&self, _: &Value) -> bool {
43        true
44    }
45
46    async fn invoke(&self, args: Value, _: ToolContext) -> Result<ToolOutput, AgentError> {
47        let name = args
48            .get("name")
49            .and_then(|v| v.as_str())
50            .ok_or_else(|| AgentError::Validation("missing `name`".into()))?
51            .to_string();
52        let store = self.store.clone();
53        tokio::task::spawn_blocking(move || {
54            let mut store = store.lock().map_err(|e| AgentError::ToolExecution {
55                tool: "MemoryRead".into(),
56                message: format!("memory store poisoned: {e}"),
57            })?;
58            match store.read(&name) {
59                Some(entry) => {
60                    store.record_use(&name).map_err(|e| AgentError::ToolExecution {
61                        tool: "MemoryRead".into(),
62                        message: e.to_string(),
63                    })?;
64                    let mut value = serde_json::to_value(&entry)
65                        .unwrap_or(json!({"error":"serialization failed"}));
66                    value["body"] = json!(entry.body);
67                    Ok(ToolOutput::json(value))
68                }
69                None => {
70                    Ok(ToolOutput::json(json!({"error": format!("memory '{}' not found", name)})))
71                }
72            }
73        })
74        .await
75        .map_err(|e| AgentError::ToolExecution {
76            tool: "MemoryRead".into(),
77            message: format!("memory task panicked: {e}"),
78        })?
79    }
80}
81
82// ── MemoryWrite ──────────────────────────────────────────
83
84pub struct MemoryWriteTool {
85    store: Arc<Mutex<MemoryStore>>,
86}
87
88impl MemoryWriteTool {
89    pub fn new(store: Arc<Mutex<MemoryStore>>) -> Self {
90        Self { store }
91    }
92}
93
94#[async_trait]
95impl Tool for MemoryWriteTool {
96    fn definition(&self) -> ToolDefinition {
97        ToolDefinition {
98            name: "MemoryWrite".into(),
99            description:
100                "Write a new memory entry. Categories: script, command, pattern, fact, workflow."
101                    .into(),
102            input_schema: json!({"type":"object","properties":{
103                "name":{"type":"string"},
104                "description":{"type":"string"},
105                "category":{"type":"string","enum":["script","command","pattern","fact","workflow"]},
106                "body":{"type":"string","description":"Markdown body content"},
107                "status":{"type":"string","enum":["working","needs_fix","deprecated"],"default":"working"},
108                "confidence":{"type":"string"},
109                "tags":{"type":"array","items":{"type":"string"},"default":[]},
110                "related":{"type":"array","items":{"type":"string"},"default":[]}
111            },"required":["name","description","category","body"]}),
112        }
113    }
114    async fn check_permission(
115        &self,
116        _: &Value,
117        _: &ToolContext,
118    ) -> Result<PermissionDecision, AgentError> {
119        Ok(PermissionDecision::Allow)
120    }
121
122    async fn invoke(&self, args: Value, _: ToolContext) -> Result<ToolOutput, AgentError> {
123        let name = args
124            .get("name")
125            .and_then(|v| v.as_str())
126            .ok_or(AgentError::Validation("missing name".into()))?
127            .to_string();
128        let desc = args.get("description").and_then(|v| v.as_str()).unwrap_or("").to_string();
129        let cat_str = args
130            .get("category")
131            .and_then(|v| v.as_str())
132            .ok_or(AgentError::Validation("missing category".into()))?
133            .to_string();
134        let body = args.get("body").and_then(|v| v.as_str()).unwrap_or("").to_string();
135        let status_str = args.get("status").and_then(|v| v.as_str()).unwrap_or("working");
136        let status = match status_str {
137            "working" => crate::knowledge::memory::MemoryStatus::Working,
138            "needs_fix" => crate::knowledge::memory::MemoryStatus::NeedsFix,
139            "deprecated" => crate::knowledge::memory::MemoryStatus::Deprecated,
140            _ => return Err(AgentError::Validation(format!("unknown status: {status_str}"))),
141        };
142        let confidence = args.get("confidence").and_then(|v| v.as_str()).map(String::from);
143        let category = match cat_str.as_str() {
144            "script" => crate::knowledge::memory::MemoryCategory::Script,
145            "command" => crate::knowledge::memory::MemoryCategory::Command,
146            "pattern" => crate::knowledge::memory::MemoryCategory::Pattern,
147            "fact" => crate::knowledge::memory::MemoryCategory::Fact,
148            "workflow" => crate::knowledge::memory::MemoryCategory::Workflow,
149            _ => return Err(AgentError::Validation(format!("unknown category: {cat_str}"))),
150        };
151        let tags: Vec<String> = args
152            .get("tags")
153            .and_then(|v| v.as_array())
154            .map(|a| a.iter().filter_map(|v| v.as_str().map(String::from)).collect())
155            .unwrap_or_default();
156        let related: Vec<String> = args
157            .get("related")
158            .and_then(|v| v.as_array())
159            .map(|a| a.iter().filter_map(|v| v.as_str().map(String::from)).collect())
160            .unwrap_or_default();
161        let now = chrono_now();
162
163        let entry = MemoryEntry {
164            name,
165            description: desc,
166            category,
167            tags,
168            created: now.clone(),
169            updated: now,
170            status,
171            times_used: 0,
172            confidence,
173            related,
174            source_session: None,
175            body,
176        };
177
178        let store = self.store.clone();
179        tokio::task::spawn_blocking(move || {
180            let entry_name = entry.name.clone();
181            let mut store = store.lock().map_err(|e| AgentError::ToolExecution {
182                tool: "MemoryWrite".into(),
183                message: format!("memory store poisoned: {e}"),
184            })?;
185            store.upsert(entry).map_err(|e| AgentError::ToolExecution {
186                tool: "MemoryWrite".into(),
187                message: e.to_string(),
188            })?;
189            Ok(ToolOutput::json(json!({"status": "written", "name": entry_name})))
190        })
191        .await
192        .map_err(|e| AgentError::ToolExecution {
193            tool: "MemoryWrite".into(),
194            message: format!("memory task panicked: {e}"),
195        })?
196    }
197}
198
199// ── MemoryGrep ───────────────────────────────────────────
200
201pub struct MemoryGrepTool {
202    store: Arc<Mutex<MemoryStore>>,
203}
204
205impl MemoryGrepTool {
206    pub fn new(store: Arc<Mutex<MemoryStore>>) -> Self {
207        Self { store }
208    }
209}
210
211#[async_trait]
212impl Tool for MemoryGrepTool {
213    fn definition(&self) -> ToolDefinition {
214        ToolDefinition {
215            name: "MemoryGrep".into(),
216            description:
217                "Search memories by keyword. Matches against name, description, tags, and body."
218                    .into(),
219            input_schema: json!({"type":"object","properties":{"query":{"type":"string"}},"required":["query"]}),
220        }
221    }
222    async fn check_permission(
223        &self,
224        _: &Value,
225        _: &ToolContext,
226    ) -> Result<PermissionDecision, AgentError> {
227        Ok(PermissionDecision::Allow)
228    }
229    fn is_concurrency_safe(&self, _: &Value) -> bool {
230        true
231    }
232
233    async fn invoke(&self, args: Value, _: ToolContext) -> Result<ToolOutput, AgentError> {
234        let query = args
235            .get("query")
236            .and_then(|v| v.as_str())
237            .ok_or(AgentError::Validation("missing query".into()))?
238            .to_string();
239        let store = self.store.clone();
240        tokio::task::spawn_blocking(move || {
241            let store = store.lock().map_err(|e| AgentError::ToolExecution {
242                tool: "MemoryGrep".into(),
243                message: format!("memory store poisoned: {e}"),
244            })?;
245            let results = store.search(&query);
246            let summary: Vec<Value> = results.iter().map(|e| json!({"name": e.name, "description": e.description, "category": format!("{:?}", e.category), "tags": e.tags})).collect();
247            Ok(ToolOutput::json(json!({"results": summary, "count": summary.len()})))
248        })
249        .await
250        .map_err(|e| AgentError::ToolExecution {
251            tool: "MemoryGrep".into(),
252            message: format!("memory task panicked: {e}"),
253        })?
254    }
255}
256
257// ── MemoryEdit ───────────────────────────────────────────
258
259pub struct MemoryEditTool {
260    store: Arc<Mutex<MemoryStore>>,
261}
262
263impl MemoryEditTool {
264    pub fn new(store: Arc<Mutex<MemoryStore>>) -> Self {
265        Self { store }
266    }
267}
268
269#[async_trait]
270impl Tool for MemoryEditTool {
271    fn definition(&self) -> ToolDefinition {
272        ToolDefinition {
273            name: "MemoryEdit".into(),
274            description: "Edit a memory entry's body content. Replaces the full body.".into(),
275            input_schema: json!({"type":"object","properties":{"name":{"type":"string"},"body":{"type":"string","description":"New body content"}},"required":["name","body"]}),
276        }
277    }
278    async fn check_permission(
279        &self,
280        _: &Value,
281        _: &ToolContext,
282    ) -> Result<PermissionDecision, AgentError> {
283        Ok(PermissionDecision::Allow)
284    }
285
286    async fn invoke(&self, args: Value, _: ToolContext) -> Result<ToolOutput, AgentError> {
287        let name = args
288            .get("name")
289            .and_then(|v| v.as_str())
290            .ok_or(AgentError::Validation("missing name".into()))?
291            .to_string();
292        let body = args.get("body").and_then(|v| v.as_str()).unwrap_or("").to_string();
293        let store = self.store.clone();
294        tokio::task::spawn_blocking(move || {
295            let mut store = store.lock().map_err(|e| AgentError::ToolExecution {
296                tool: "MemoryEdit".into(),
297                message: format!("memory store poisoned: {e}"),
298            })?;
299            if let Some(mut entry) = store.read(&name) {
300                entry.body = body;
301                entry.updated = chrono_now();
302                store.write(entry).map_err(|e| AgentError::ToolExecution {
303                    tool: "MemoryEdit".into(),
304                    message: e.to_string(),
305                })?;
306                Ok(ToolOutput::json(json!({"status": "updated", "name": name})))
307            } else {
308                Ok(ToolOutput::json(json!({"error": format!("memory '{}' not found", name)})))
309            }
310        })
311        .await
312        .map_err(|e| AgentError::ToolExecution {
313            tool: "MemoryEdit".into(),
314            message: format!("memory task panicked: {e}"),
315        })?
316    }
317}
318
319// ── MemoryStatus ─────────────────────────────────────────
320
321pub struct MemoryStatusTool {
322    store: Arc<Mutex<MemoryStore>>,
323}
324
325impl MemoryStatusTool {
326    pub fn new(store: Arc<Mutex<MemoryStore>>) -> Self {
327        Self { store }
328    }
329}
330
331#[async_trait]
332impl Tool for MemoryStatusTool {
333    fn definition(&self) -> ToolDefinition {
334        ToolDefinition {
335            name: "MemoryStatus".into(),
336            description: "Update the status of a memory entry: working, needs_fix, or deprecated."
337                .into(),
338            input_schema: json!({"type":"object","properties":{"name":{"type":"string"},"status":{"type":"string","enum":["working","needs_fix","deprecated"]}},"required":["name","status"]}),
339        }
340    }
341    async fn check_permission(
342        &self,
343        _: &Value,
344        _: &ToolContext,
345    ) -> Result<PermissionDecision, AgentError> {
346        Ok(PermissionDecision::Allow)
347    }
348
349    async fn invoke(&self, args: Value, _: ToolContext) -> Result<ToolOutput, AgentError> {
350        let name = args
351            .get("name")
352            .and_then(|v| v.as_str())
353            .ok_or(AgentError::Validation("missing name".into()))?
354            .to_string();
355        let status_str = args
356            .get("status")
357            .and_then(|v| v.as_str())
358            .ok_or(AgentError::Validation("missing status".into()))?
359            .to_string();
360        let status = match status_str.as_str() {
361            "working" => crate::knowledge::memory::MemoryStatus::Working,
362            "needs_fix" => crate::knowledge::memory::MemoryStatus::NeedsFix,
363            "deprecated" => crate::knowledge::memory::MemoryStatus::Deprecated,
364            _ => return Err(AgentError::Validation(format!("unknown status: {status_str}"))),
365        };
366        let store = self.store.clone();
367        tokio::task::spawn_blocking(move || {
368            let mut store = store.lock().map_err(|e| AgentError::ToolExecution {
369                tool: "MemoryStatus".into(),
370                message: format!("memory store poisoned: {e}"),
371            })?;
372            store.update_status(&name, status).map_err(|e| AgentError::ToolExecution {
373                tool: "MemoryStatus".into(),
374                message: e.to_string(),
375            })?;
376            Ok(ToolOutput::json(
377                json!({"status": "updated", "name": name, "new_status": status_str}),
378            ))
379        })
380        .await
381        .map_err(|e| AgentError::ToolExecution {
382            tool: "MemoryStatus".into(),
383            message: format!("memory task panicked: {e}"),
384        })?
385    }
386}
387
388fn is_leap_year(y: u64) -> bool {
389    (y.is_multiple_of(4) && !y.is_multiple_of(100)) || y.is_multiple_of(400)
390}
391
392fn chrono_now() -> String {
393    // Simple date string without chrono dep.
394    // Walk forward from Unix epoch to find the current year, then month.
395    let now =
396        std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default();
397    let mut days = now.as_secs() / 86400;
398    let mut year: u64 = 1970;
399    loop {
400        let days_in_year: u64 = if is_leap_year(year) { 366 } else { 365 };
401        if days < days_in_year {
402            break;
403        }
404        days -= days_in_year;
405        year += 1;
406    }
407    let ml: [u64; 12] = if is_leap_year(year) {
408        [31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
409    } else {
410        [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
411    };
412    let (mut rem, mut m) = (days, 1usize);
413    for &l in &ml {
414        if rem < l {
415            break;
416        }
417        rem -= l;
418        m += 1;
419    }
420    let d = rem + 1;
421    format!("{year}-{m:02}-{d:02}")
422}
telos_agent/tools/builtin/memory.rs

telos_agent/tools/builtin/
memory.rs
