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