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; pub 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 if !path_obj.exists() {
38 return Err(AgentError::ToolError(format!("File not found: {}", path)));
39 }
40
41 if !path_obj.is_file() {
43 return Err(AgentError::ToolError(format!(
44 "Path is not a file: {}",
45 path
46 )));
47 }
48
49 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 let content = fs::read_to_string(&path)
63 .map_err(|e| AgentError::IoError(format!("Failed to read file: {}", e)))?;
64
65 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 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 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 if !path_obj.exists() {
164 return Err(AgentError::ToolError(format!("File not found: {}", path)));
165 }
166
167 let current_content = fs::read_to_string(&path)
169 .map_err(|e| AgentError::IoError(format!("Failed to read file: {}", e)))?;
170
171 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 if current_content.contains('\0') {
185 return Err(AgentError::ToolError(format!(
186 "Binary file detected: {}",
187 path
188 )));
189 }
190
191 if !current_content.contains(&old_text) {
193 return Err(AgentError::ToolError(
194 "old_text not found in file".to_string(),
195 ));
196 }
197
198 let new_content = current_content.replace(&old_text, &new_text);
200
201 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!({}); 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 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 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 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}