Skip to main content

limit_cli/tools/
file.rs

1use async_trait::async_trait;
2use limit_agent::error::AgentError;
3use limit_agent::Tool;
4use serde_json::Value;
5use std::fs;
6use std::path::Path;
7
8const MAX_FILE_SIZE: u64 = 50 * 1024 * 1024; // 50MB
9
10pub struct FileReadTool;
11
12impl FileReadTool {
13    pub fn new() -> Self {
14        FileReadTool
15    }
16}
17
18impl Default for FileReadTool {
19    fn default() -> Self {
20        Self::new()
21    }
22}
23
24#[async_trait]
25impl Tool for FileReadTool {
26    fn name(&self) -> &str {
27        "file_read"
28    }
29
30    async fn execute(&self, args: Value) -> Result<Value, AgentError> {
31        let path: String = serde_json::from_value(args["path"].clone())
32            .map_err(|e| AgentError::ToolError(format!("Invalid path argument: {}", e)))?;
33
34        let path_obj = Path::new(&path);
35
36        // Check if file exists
37        if !path_obj.exists() {
38            return Err(AgentError::ToolError(format!("File not found: {}", path)));
39        }
40
41        // Check if it's a file
42        if !path_obj.is_file() {
43            return Err(AgentError::ToolError(format!(
44                "Path is not a file: {}",
45                path
46            )));
47        }
48
49        // Check file size
50        let metadata = fs::metadata(&path)
51            .map_err(|e| AgentError::IoError(format!("Failed to read file metadata: {}", e)))?;
52
53        if metadata.len() > MAX_FILE_SIZE {
54            return Err(AgentError::ToolError(format!(
55                "File too large: {} bytes (max: {} bytes)",
56                metadata.len(),
57                MAX_FILE_SIZE
58            )));
59        }
60
61        // Read file
62        let content = fs::read_to_string(&path)
63            .map_err(|e| AgentError::IoError(format!("Failed to read file: {}", e)))?;
64
65        // Check for binary content (null bytes)
66        if content.contains('\0') {
67            return Err(AgentError::ToolError(format!(
68                "Binary file detected: {}",
69                path
70            )));
71        }
72
73        Ok(serde_json::json!({
74            "content": content,
75            "size": metadata.len()
76        }))
77    }
78}
79
80pub struct FileWriteTool;
81
82impl FileWriteTool {
83    pub fn new() -> Self {
84        FileWriteTool
85    }
86}
87
88impl Default for FileWriteTool {
89    fn default() -> Self {
90        Self::new()
91    }
92}
93
94#[async_trait]
95impl Tool for FileWriteTool {
96    fn name(&self) -> &str {
97        "file_write"
98    }
99
100    async fn execute(&self, args: Value) -> Result<Value, AgentError> {
101        let path: String = serde_json::from_value(args["path"].clone())
102            .map_err(|e| AgentError::ToolError(format!("Invalid path argument: {}", e)))?;
103
104        let content: String = serde_json::from_value(args["content"].clone())
105            .map_err(|e| AgentError::ToolError(format!("Invalid content argument: {}", e)))?;
106
107        let path_obj = Path::new(&path);
108
109        // Create parent directories if they don't exist
110        if let Some(parent) = path_obj.parent() {
111            if !parent.exists() {
112                fs::create_dir_all(parent).map_err(|e| {
113                    AgentError::IoError(format!("Failed to create directories: {}", e))
114                })?;
115            }
116        }
117
118        // Write file
119        fs::write(&path, &content)
120            .map_err(|e| AgentError::IoError(format!("Failed to write file: {}", e)))?;
121
122        Ok(serde_json::json!({
123            "success": true,
124            "path": path,
125            "size": content.len()
126        }))
127    }
128}
129
130pub struct FileEditTool;
131
132impl FileEditTool {
133    pub fn new() -> Self {
134        FileEditTool
135    }
136}
137
138impl Default for FileEditTool {
139    fn default() -> Self {
140        Self::new()
141    }
142}
143
144#[async_trait]
145impl Tool for FileEditTool {
146    fn name(&self) -> &str {
147        "file_edit"
148    }
149
150    async fn execute(&self, args: Value) -> Result<Value, AgentError> {
151        let path: String = serde_json::from_value(args["path"].clone())
152            .map_err(|e| AgentError::ToolError(format!("Invalid path argument: {}", e)))?;
153
154        let old_text: String = serde_json::from_value(args["old_text"].clone())
155            .map_err(|e| AgentError::ToolError(format!("Invalid old_text argument: {}", e)))?;
156
157        let new_text: String = serde_json::from_value(args["new_text"].clone())
158            .map_err(|e| AgentError::ToolError(format!("Invalid new_text argument: {}", e)))?;
159
160        let path_obj = Path::new(&path);
161
162        // Check if file exists
163        if !path_obj.exists() {
164            return Err(AgentError::ToolError(format!("File not found: {}", path)));
165        }
166
167        // Read file content
168        let current_content = fs::read_to_string(&path)
169            .map_err(|e| AgentError::IoError(format!("Failed to read file: {}", e)))?;
170
171        // Check file size
172        let metadata = fs::metadata(&path)
173            .map_err(|e| AgentError::IoError(format!("Failed to read file metadata: {}", e)))?;
174
175        if metadata.len() > MAX_FILE_SIZE {
176            return Err(AgentError::ToolError(format!(
177                "File too large: {} bytes (max: {} bytes)",
178                metadata.len(),
179                MAX_FILE_SIZE
180            )));
181        }
182
183        // Check for binary content
184        if current_content.contains('\0') {
185            return Err(AgentError::ToolError(format!(
186                "Binary file detected: {}",
187                path
188            )));
189        }
190
191        // Check if old_text exists in file
192        if !current_content.contains(&old_text) {
193            return Err(AgentError::ToolError(
194                "old_text not found in file".to_string(),
195            ));
196        }
197
198        // Replace old_text with new_text
199        let new_content = current_content.replace(&old_text, &new_text);
200
201        // Write modified content back
202        fs::write(&path, new_content)
203            .map_err(|e| AgentError::IoError(format!("Failed to write file: {}", e)))?;
204
205        Ok(serde_json::json!({
206            "success": true,
207            "path": path,
208            "replacements": current_content.matches(&old_text).count()
209        }))
210    }
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216    use std::io::Write;
217    use tempfile::NamedTempFile;
218
219    #[tokio::test]
220    async fn test_file_read_tool_name() {
221        let tool = FileReadTool::new();
222        assert_eq!(tool.name(), "file_read");
223    }
224
225    #[tokio::test]
226    async fn test_file_read_tool_execute() {
227        let mut temp_file = NamedTempFile::new().unwrap();
228        writeln!(temp_file, "Hello, World!").unwrap();
229        writeln!(temp_file, "This is a test.").unwrap();
230
231        let tool = FileReadTool::new();
232        let args = serde_json::json!({
233            "path": temp_file.path().to_str().unwrap()
234        });
235
236        let result = tool.execute(args).await.unwrap();
237        assert!(result["content"].is_string());
238        assert!(result["size"].is_u64());
239        assert!(result["content"]
240            .as_str()
241            .unwrap()
242            .contains("Hello, World!"));
243    }
244
245    #[tokio::test]
246    async fn test_file_read_tool_file_not_found() {
247        let tool = FileReadTool::new();
248        let args = serde_json::json!({
249            "path": "/nonexistent/file.txt"
250        });
251
252        let result = tool.execute(args).await;
253        assert!(result.is_err());
254        assert!(result.unwrap_err().to_string().contains("File not found"));
255    }
256
257    #[tokio::test]
258    async fn test_file_read_tool_invalid_path() {
259        let tool = FileReadTool::new();
260        let args = serde_json::json!({}); // Missing path
261
262        let result = tool.execute(args).await;
263        assert!(result.is_err());
264    }
265
266    #[tokio::test]
267    async fn test_file_write_tool_name() {
268        let tool = FileWriteTool::new();
269        assert_eq!(tool.name(), "file_write");
270    }
271
272    #[tokio::test]
273    async fn test_file_write_tool_execute() {
274        let temp_file = NamedTempFile::new().unwrap();
275        let tool = FileWriteTool::new();
276        let content = "Hello from test!";
277
278        let args = serde_json::json!({
279            "path": temp_file.path().to_str().unwrap(),
280            "content": content
281        });
282
283        let result = tool.execute(args).await.unwrap();
284        assert_eq!(result["success"], true);
285        assert_eq!(result["size"], content.len());
286
287        // Verify content was written
288        let written = fs::read_to_string(temp_file.path()).unwrap();
289        assert_eq!(written, content);
290    }
291
292    #[tokio::test]
293    async fn test_file_write_tool_create_dirs() {
294        let temp_dir = tempfile::tempdir().unwrap();
295        let nested_path = temp_dir.path().join("nested/dir/file.txt");
296
297        let tool = FileWriteTool::new();
298        let args = serde_json::json!({
299            "path": nested_path.to_str().unwrap(),
300            "content": "Test content"
301        });
302
303        let result = tool.execute(args).await.unwrap();
304        assert_eq!(result["success"], true);
305        assert!(nested_path.exists());
306    }
307
308    #[tokio::test]
309    async fn test_file_edit_tool_name() {
310        let tool = FileEditTool::new();
311        assert_eq!(tool.name(), "file_edit");
312    }
313
314    #[tokio::test]
315    async fn test_file_edit_tool_execute() {
316        let mut temp_file = NamedTempFile::new().unwrap();
317        writeln!(temp_file, "Hello, World!").unwrap();
318        writeln!(temp_file, "Goodbye, World!").unwrap();
319
320        let tool = FileEditTool::new();
321        let args = serde_json::json!({
322            "path": temp_file.path().to_str().unwrap(),
323            "old_text": "Hello, World!",
324            "new_text": "Hello, Rust!"
325        });
326
327        let result = tool.execute(args).await.unwrap();
328        assert_eq!(result["success"], true);
329        assert_eq!(result["replacements"], 1);
330
331        // Verify edit
332        let content = fs::read_to_string(temp_file.path()).unwrap();
333        assert!(content.contains("Hello, Rust!"));
334        assert!(!content.contains("Hello, World!"));
335    }
336
337    #[tokio::test]
338    async fn test_file_edit_tool_old_text_not_found() {
339        let mut temp_file = NamedTempFile::new().unwrap();
340        writeln!(temp_file, "Hello, World!").unwrap();
341
342        let tool = FileEditTool::new();
343        let args = serde_json::json!({
344            "path": temp_file.path().to_str().unwrap(),
345            "old_text": "Nonexistent text",
346            "new_text": "Replacement"
347        });
348
349        let result = tool.execute(args).await;
350        assert!(result.is_err());
351        assert!(result
352            .unwrap_err()
353            .to_string()
354            .contains("old_text not found"));
355    }
356
357    #[tokio::test]
358    async fn test_file_read_tool_binary_detection() {
359        let mut temp_file = NamedTempFile::new().unwrap();
360        // Write binary content with null byte
361        temp_file.write_all(b"Hello\x00World").unwrap();
362
363        let tool = FileReadTool::new();
364        let args = serde_json::json!({
365            "path": temp_file.path().to_str().unwrap()
366        });
367
368        let result = tool.execute(args).await;
369        assert!(result.is_err());
370        assert!(result.unwrap_err().to_string().contains("Binary file"));
371    }
372}