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 if !path_obj.exists() {
39 return Err(AgentError::ToolError(format!("File not found: {}", path)));
40 }
41
42 if !path_obj.is_file() {
44 return Err(AgentError::ToolError(format!(
45 "Path is not a file: {}",
46 path
47 )));
48 }
49
50 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 let content = fs::read_to_string(&path)
64 .map_err(|e| AgentError::IoError(format!("Failed to read file: {}", e)))?;
65
66 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 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 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 if !path_obj.exists() {
187 return Err(AgentError::ToolError(format!("File not found: {}", path)));
188 }
189
190 let current_content = fs::read_to_string(&path)
192 .map_err(|e| AgentError::IoError(format!("Failed to read file: {}", e)))?;
193
194 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 if current_content.contains('\0') {
208 return Err(AgentError::ToolError(format!(
209 "Binary file detected: {}",
210 path
211 )));
212 }
213
214 if !current_content.contains(&old_text) {
216 return Err(AgentError::ToolError(
217 "old_text not found in file".to_string(),
218 ));
219 }
220
221 let new_content = current_content.replace(&old_text, &new_text);
223
224 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!({}); 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 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 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 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}