imp_core/tools/
memory.rs

1use async_trait::async_trait;
2use serde_json::json;
3
4use super::{Tool, ToolContext, ToolOutput};
5use crate::error::Result;
6use crate::memory::MemoryStore;
7use crate::storage;
8
9const DEFAULT_MEMORY_LIMIT: usize = 2200;
10const DEFAULT_USER_LIMIT: usize = 1400;
11
12pub struct MemoryTool;
13
14#[async_trait]
15impl Tool for MemoryTool {
16    fn name(&self) -> &str {
17        "memory"
18    }
19
20    fn label(&self) -> &str {
21        "Memory"
22    }
23
24    fn description(&self) -> &str {
25        "Manage persistent memory across sessions. Use to save environment facts, \
26         user preferences, and lessons learned. Target 'memory' for agent notes, \
27         'user' for user profile."
28    }
29
30    fn parameters(&self) -> serde_json::Value {
31        json!({
32            "type": "object",
33            "required": ["action", "target"],
34            "properties": {
35                "action": {
36                    "type": "string",
37                    "enum": ["add", "replace", "remove"],
38                    "description": "Action to perform"
39                },
40                "target": {
41                    "type": "string",
42                    "enum": ["memory", "user"],
43                    "description": "Which store: 'memory' for agent notes, 'user' for user profile"
44                },
45                "content": {
46                    "type": "string",
47                    "description": "Content to add or replacement text"
48                },
49                "old_text": {
50                    "type": "string",
51                    "description": "Unique substring identifying the entry to replace or remove"
52                }
53            }
54        })
55    }
56
57    fn is_readonly(&self) -> bool {
58        false
59    }
60
61    async fn execute(
62        &self,
63        _call_id: &str,
64        params: serde_json::Value,
65        _ctx: ToolContext,
66    ) -> Result<ToolOutput> {
67        let action = params["action"].as_str().unwrap_or("");
68        let target = params["target"].as_str().unwrap_or("");
69
70        if action.is_empty() {
71            return Ok(ToolOutput::error("Missing required parameter: action"));
72        }
73        if target.is_empty() {
74            return Ok(ToolOutput::error("Missing required parameter: target"));
75        }
76
77        let (path, char_limit) = match target {
78            "memory" => (storage::global_memory_path(), DEFAULT_MEMORY_LIMIT),
79            "user" => (storage::global_user_path(), DEFAULT_USER_LIMIT),
80            other => {
81                return Ok(ToolOutput::error(format!(
82                    "Unknown target \"{other}\". Use \"memory\" or \"user\"."
83                )));
84            }
85        };
86
87        let mut store = match MemoryStore::load(&path, char_limit) {
88            Ok(s) => s,
89            Err(e) => return Ok(ToolOutput::error(format!("Failed to load memory: {e}"))),
90        };
91
92        let result = match action {
93            "add" => {
94                let content = params["content"].as_str().unwrap_or("");
95                if content.is_empty() {
96                    return Ok(ToolOutput::error(
97                        "Missing required parameter: content (for 'add' action)",
98                    ));
99                }
100                store.add(content)?
101            }
102            "replace" => {
103                let old_text = params["old_text"].as_str().unwrap_or("");
104                let content = params["content"].as_str().unwrap_or("");
105                if old_text.is_empty() {
106                    return Ok(ToolOutput::error(
107                        "Missing required parameter: old_text (for 'replace' action)",
108                    ));
109                }
110                if content.is_empty() {
111                    return Ok(ToolOutput::error(
112                        "Missing required parameter: content (for 'replace' action)",
113                    ));
114                }
115                store.replace(old_text, content)?
116            }
117            "remove" => {
118                let old_text = params["old_text"].as_str().unwrap_or("");
119                if old_text.is_empty() {
120                    return Ok(ToolOutput::error(
121                        "Missing required parameter: old_text (for 'remove' action)",
122                    ));
123                }
124                store.remove(old_text)?
125            }
126            other => {
127                return Ok(ToolOutput::error(format!(
128                    "Unknown action \"{other}\". Use \"add\", \"replace\", or \"remove\"."
129                )));
130            }
131        };
132
133        let json_text = serde_json::to_string_pretty(&result.to_json()).unwrap_or_default();
134        if result.success {
135            Ok(ToolOutput::text(json_text))
136        } else {
137            Ok(ToolOutput::error(json_text))
138        }
139    }
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145    use crate::tools::ToolContext;
146    use std::sync::Arc;
147
148    fn test_ctx() -> ToolContext {
149        let (tx, _rx) = tokio::sync::mpsc::channel(16);
150        let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::channel(16);
151        let dir = std::env::temp_dir();
152        ToolContext {
153            cwd: dir,
154            cancelled: Arc::new(std::sync::atomic::AtomicBool::new(false)),
155            update_tx: tx,
156            command_tx: cmd_tx,
157            ui: Arc::new(crate::ui::NullInterface),
158            file_cache: Arc::new(crate::tools::FileCache::new()),
159            checkpoint_state: Arc::new(crate::tools::CheckpointState::new()),
160            file_tracker: Arc::new(std::sync::Mutex::new(crate::tools::FileTracker::new())),
161            anchor_store: Arc::new(crate::tools::AnchorStore::new()),
162            lua_tool_loader: None,
163            mode: crate::config::AgentMode::Full,
164            read_max_lines: 500,
165            turn_mana_review: Arc::new(std::sync::Mutex::new(
166                crate::mana_review::TurnManaReviewAccumulator::default(),
167            )),
168            config: Arc::new(crate::config::Config::default()),
169        }
170    }
171
172    #[tokio::test]
173    async fn memory_tool_validates_params() {
174        let tool = MemoryTool;
175
176        // Missing action
177        let r = tool
178            .execute("c1", json!({"target": "memory"}), test_ctx())
179            .await
180            .unwrap();
181        assert!(r.is_error);
182
183        // Missing target
184        let r = tool
185            .execute("c2", json!({"action": "add"}), test_ctx())
186            .await
187            .unwrap();
188        assert!(r.is_error);
189
190        // Missing content for add
191        let r = tool
192            .execute(
193                "c3",
194                json!({"action": "add", "target": "memory"}),
195                test_ctx(),
196            )
197            .await
198            .unwrap();
199        assert!(r.is_error);
200    }
201}
imp_core/tools/memory.rs

imp_core/tools/
memory.rs
