1use super::file_mutation_queue::global_mutation_queue;
10use super::path_security::PathGuard;
11use super::truncate::{self, TruncationOptions};
12use super::{AgentTool, AgentToolResult, ToolContext, ToolError};
13use async_trait::async_trait;
14use serde_json::{Value, json};
15use std::path::{Path, PathBuf};
16use tokio::fs;
17use tokio::sync::oneshot;
18
19const PREVIEW_HEAD_LINES: usize = 5;
21const PREVIEW_TAIL_LINES: usize = 5;
22const PREVIEW_THRESHOLD_LINES: usize = 20;
24
25pub struct WriteTool {
27 root_dir: Option<PathBuf>,
28}
29
30impl WriteTool {
31 pub fn new() -> Self {
33 Self { root_dir: None }
34 }
35
36 pub fn with_cwd(cwd: PathBuf) -> Self {
38 Self {
39 root_dir: Some(cwd),
40 }
41 }
42
43 fn build_content_preview(content: &str, total_lines: usize) -> String {
46 if total_lines <= PREVIEW_THRESHOLD_LINES {
47 return content.to_string();
48 }
49
50 let lines: Vec<&str> = content.lines().collect();
51 let head: Vec<&str> = lines.iter().copied().take(PREVIEW_HEAD_LINES).collect();
52 let tail: Vec<&str> = lines
53 .iter()
54 .copied()
55 .rev()
56 .take(PREVIEW_TAIL_LINES)
57 .rev()
58 .collect();
59
60 let omitted = total_lines - PREVIEW_HEAD_LINES - PREVIEW_TAIL_LINES;
61
62 format!(
63 "{}\n\n... [{} lines omitted] ...\n\n{}",
64 head.join("\n"),
65 omitted,
66 tail.join("\n")
67 )
68 }
69
70 async fn write_file_impl(
72 root_dir: &Path,
73 path: &str,
74 content: &str,
75 append: bool,
76 ) -> Result<String, ToolError> {
77 let guard = PathGuard::new(root_dir);
79 let file_path = guard
80 .validate_traversal(Path::new(path))
81 .map_err(|e| e.to_string())?;
82
83 if let Some(parent) = file_path.parent() {
85 if !parent.as_os_str().is_empty() {
87 fs::create_dir_all(parent)
88 .await
89 .map_err(|e| format!("Cannot create parent directory: {}", e))?;
90 }
91 }
92
93 let existed = file_path.exists();
95
96 let content_owned = content.to_string();
98 let result = global_mutation_queue()
99 .with_queue(&file_path, || async {
100 if append {
101 let mut file = tokio::fs::OpenOptions::new()
102 .create(true)
103 .append(true)
104 .open(&file_path)
105 .await
106 .map_err(|e| format!("Cannot open file for append: {}", e))?;
107 use tokio::io::AsyncWriteExt;
108 file.write_all(content_owned.as_bytes())
109 .await
110 .map_err(|e| format!("Cannot write file: {}", e))?;
111 file.flush()
112 .await
113 .map_err(|e| format!("Cannot flush file: {}", e))?;
114 } else {
115 fs::write(&file_path, &content_owned)
116 .await
117 .map_err(|e| format!("Cannot write file: {}", e))?;
118 }
119 Ok::<(), ToolError>(())
120 })
121 .await;
122
123 result?;
124
125 let total_lines = content.lines().count();
126 let total_bytes = content.len();
127 let action = if append { "Appended" } else { "Wrote" };
128 let status = if existed && !append {
129 " (overwritten)"
130 } else if append && existed {
131 " (appended)"
132 } else if !existed {
133 " (new file)"
134 } else {
135 ""
136 };
137
138 let preview = Self::build_content_preview(content, total_lines);
140
141 let truncation_opts = TruncationOptions {
143 max_lines: Some(50),
144 max_bytes: Some(4 * 1024),
145 };
146 let truncated = truncate::truncate_head(&preview, &truncation_opts);
147
148 let mut msg = format!(
149 "{} {} lines ({} bytes) to {}{}\n",
150 action, total_lines, total_bytes, path, status
151 );
152
153 msg.push_str(&format!("--- Content Preview ---\n{}", truncated.content));
154
155 if truncated.truncated {
156 msg.push_str(&format!(
157 "\n[Output truncated: {} total lines, {} total bytes]",
158 truncated.total_lines, truncated.total_bytes
159 ));
160 }
161
162 Ok(msg)
163 }
164}
165
166impl Default for WriteTool {
167 fn default() -> Self {
168 Self::new()
169 }
170}
171
172#[async_trait]
173impl AgentTool for WriteTool {
174 fn name(&self) -> &str {
175 "write"
176 }
177
178 fn label(&self) -> &str {
179 "Write File"
180 }
181
182 fn essential(&self) -> bool {
183 true
184 }
185 fn description(&self) -> &str {
186 "Write content to a file, creating parent directories as needed. Existing files will be overwritten. Use append=true to append to existing files."
187 }
188
189 fn parameters_schema(&self) -> Value {
190 json!({
191 "type": "object",
192 "properties": {
193 "path": {
194 "type": "string",
195 "description": "The path to the file to write"
196 },
197 "content": {
198 "type": "string",
199 "description": "The content to write to the file"
200 },
201 "append": {
202 "type": "boolean",
203 "description": "If true, append to existing file instead of overwriting",
204 "default": false
205 }
206 },
207 "required": ["path", "content"]
208 })
209 }
210
211 async fn execute(
212 &self,
213 _tool_call_id: &str,
214 params: Value,
215 _signal: Option<oneshot::Receiver<()>>,
216 ctx: &ToolContext,
217 ) -> Result<AgentToolResult, ToolError> {
218 let path = params
219 .get("path")
220 .and_then(|v| v.as_str())
221 .ok_or_else(|| "Missing required parameter: path".to_string())?;
222
223 let content = params
224 .get("content")
225 .and_then(|v| v.as_str())
226 .ok_or_else(|| "Missing required parameter: content".to_string())?;
227
228 let append = params
229 .get("append")
230 .and_then(|v| v.as_bool())
231 .unwrap_or(false);
232
233 let root = self.root_dir.as_deref().unwrap_or(ctx.root());
235
236 let write_result = Self::write_file_impl(root, path, content, append).await;
237 let notify_path = std::path::Path::new(path).to_path_buf();
245 let notify_abs = if notify_path.is_absolute() {
246 notify_path
247 } else {
248 std::path::Path::new(root).join(¬ify_path)
249 };
250 if write_result.is_ok()
251 && let Some(provider) = ctx.lsp.as_ref()
252 {
253 let provider_clone = provider.clone();
254 let abs_path_clone = notify_abs.clone();
255 let content_owned = content.to_string();
256 tokio::spawn(async move {
257 provider_clone
258 .notify_file_changed(&abs_path_clone, &content_owned)
259 .await;
260 });
261 }
262 match write_result {
263 Ok(msg) => Ok(AgentToolResult::success(msg)),
264 Err(e) => Ok(AgentToolResult::error(e)),
265 }
266 }
267}
268
269#[cfg(test)]
270mod tests {
271 use super::*;
272 use tempfile::TempDir;
273
274 #[test]
275 fn test_build_content_preview_small() {
276 let content = "line1\nline2\nline3";
277 let preview = WriteTool::build_content_preview(content, 3);
278 assert_eq!(preview, content);
279 }
280
281 #[test]
282 fn test_build_content_preview_large() {
283 let lines: Vec<String> = (1..=30).map(|i| format!("line {}", i)).collect();
284 let content = lines.join("\n");
285 let preview = WriteTool::build_content_preview(&content, 30);
286
287 assert!(preview.contains("line 1"));
288 assert!(preview.contains("line 5"));
289 assert!(preview.contains("line 26"));
290 assert!(preview.contains("line 30"));
291 assert!(preview.contains("lines omitted"));
292 assert!(!preview.contains("line 10")); }
294
295 #[test]
296 fn test_build_content_preview_exact_threshold() {
297 let lines: Vec<String> = (1..=20).map(|i| format!("line {}", i)).collect();
298 let content = lines.join("\n");
299 let preview = WriteTool::build_content_preview(&content, 20);
300 assert_eq!(preview, content);
302 }
303
304 #[test]
305 fn test_build_content_preview_one_over_threshold() {
306 let lines: Vec<String> = (1..=21).map(|i| format!("line {}", i)).collect();
307 let content = lines.join("\n");
308 let preview = WriteTool::build_content_preview(&content, 21);
309 assert!(preview.contains("lines omitted"));
311 }
312
313 #[tokio::test]
314 async fn test_write_new_file() {
315 let tmp = TempDir::new().unwrap();
316 let path = tmp.path().join("test.txt");
317 let path_str = path.to_str().unwrap();
318
319 let result =
320 WriteTool::write_file_impl(Path::new("."), path_str, "hello world\nline 2", false)
321 .await;
322 assert!(result.is_ok());
323
324 let written = std::fs::read_to_string(&path).unwrap();
325 assert_eq!(written, "hello world\nline 2");
326
327 let msg = result.unwrap();
328 assert!(msg.contains("2 lines"));
329 assert!(msg.contains("new file"));
330 }
331
332 #[tokio::test]
333 async fn test_write_creates_parent_dirs() {
334 let tmp = TempDir::new().unwrap();
335 let path = tmp.path().join("a/b/c/test.txt");
336 let path_str = path.to_str().unwrap();
337
338 let result =
339 WriteTool::write_file_impl(Path::new("."), path_str, "deep nested", false).await;
340 assert!(result.is_ok());
341
342 let written = std::fs::read_to_string(&path).unwrap();
343 assert_eq!(written, "deep nested");
344 }
345
346 #[tokio::test]
347 async fn test_write_overwrites_existing() {
348 let tmp = TempDir::new().unwrap();
349 let path = tmp.path().join("test.txt");
350 let path_str = path.to_str().unwrap();
351
352 std::fs::write(&path, "old content").unwrap();
354
355 let result =
356 WriteTool::write_file_impl(Path::new("."), path_str, "new content", false).await;
357 assert!(result.is_ok());
358
359 let written = std::fs::read_to_string(&path).unwrap();
360 assert_eq!(written, "new content");
361
362 let msg = result.unwrap();
363 assert!(msg.contains("overwritten"));
364 }
365
366 #[tokio::test]
367 async fn test_write_append_mode() {
368 let tmp = TempDir::new().unwrap();
369 let path = tmp.path().join("test.txt");
370 let path_str = path.to_str().unwrap();
371
372 WriteTool::write_file_impl(Path::new("."), path_str, "line 1\n", false)
374 .await
375 .unwrap();
376
377 let result = WriteTool::write_file_impl(Path::new("."), path_str, "line 2\n", true).await;
379 assert!(result.is_ok());
380
381 let written = std::fs::read_to_string(&path).unwrap();
382 assert_eq!(written, "line 1\nline 2\n");
383
384 let msg = result.unwrap();
385 assert!(msg.contains("Appended"));
386 }
387
388 #[tokio::test]
389 async fn test_write_append_to_nonexistent() {
390 let tmp = TempDir::new().unwrap();
391 let path = tmp.path().join("new.txt");
392 let path_str = path.to_str().unwrap();
393
394 let result =
395 WriteTool::write_file_impl(Path::new("."), path_str, "appended content", true).await;
396 assert!(result.is_ok());
397
398 let written = std::fs::read_to_string(&path).unwrap();
399 assert_eq!(written, "appended content");
400 }
401
402 #[tokio::test]
403 async fn test_write_path_traversal_blocked() {
404 let result =
405 WriteTool::write_file_impl(Path::new("."), "../../etc/passwd", "hack", false).await;
406 assert!(result.is_err());
407 assert!(result.unwrap_err().contains("Path traversal"));
408 }
409
410 #[tokio::test]
411 async fn test_write_empty_content() {
412 let tmp = TempDir::new().unwrap();
413 let path = tmp.path().join("empty.txt");
414 let path_str = path.to_str().unwrap();
415
416 let result = WriteTool::write_file_impl(Path::new("."), path_str, "", false).await;
417 assert!(result.is_ok());
418
419 let written = std::fs::read_to_string(&path).unwrap();
420 assert_eq!(written, "");
421
422 let msg = result.unwrap();
423 assert!(msg.contains("0 lines"));
424 }
425
426 #[tokio::test]
427 async fn test_write_large_file_has_preview() {
428 let tmp = TempDir::new().unwrap();
429 let path = tmp.path().join("large.txt");
430 let path_str = path.to_str().unwrap();
431
432 let lines: Vec<String> = (1..=100).map(|i| format!("line {}", i)).collect();
433 let content = lines.join("\n");
434
435 let result = WriteTool::write_file_impl(Path::new("."), path_str, &content, false).await;
436 assert!(result.is_ok());
437
438 let msg = result.unwrap();
439 assert!(msg.contains("100 lines"));
440 assert!(msg.contains("Content Preview"));
441 }
442
443 #[tokio::test]
444 async fn test_execute_via_tool_trait() {
445 let tmp = TempDir::new().unwrap();
446 let path = tmp.path().join("trait_test.txt");
447 let path_str = path.to_str().unwrap().to_string();
448
449 let tool = WriteTool::new();
450 let params = json!({
451 "path": path_str,
452 "content": "via trait"
453 });
454
455 let result = tool
456 .execute("test-id", params, None, &ToolContext::default())
457 .await;
458 assert!(result.is_ok());
459 let tool_result = result.unwrap();
460 assert!(tool_result.success);
461 assert!(tool_result.output.contains("via trait"));
462
463 let written = std::fs::read_to_string(&path).unwrap();
464 assert_eq!(written, "via trait");
465 }
466
467 #[tokio::test]
468 async fn test_execute_missing_path_param() {
469 let tool = WriteTool::new();
470 let params = json!({
471 "content": "no path"
472 });
473
474 let result = tool
475 .execute("test-id", params, None, &ToolContext::default())
476 .await;
477 assert!(result.is_err());
478 assert!(result.unwrap_err().contains("path"));
479 }
480
481 #[tokio::test]
482 async fn test_execute_missing_content_param() {
483 let tool = WriteTool::new();
484 let params = json!({
485 "path": "/tmp/test.txt"
486 });
487
488 let result = tool
489 .execute("test-id", params, None, &ToolContext::default())
490 .await;
491 assert!(result.is_err());
492 assert!(result.unwrap_err().contains("content"));
493 }
494
495 #[tokio::test]
496 async fn test_execute_append_via_trait() {
497 let tmp = TempDir::new().unwrap();
498 let path = tmp.path().join("append_trait.txt");
499 let path_str = path.to_str().unwrap().to_string();
500
501 let tool = WriteTool::new();
502
503 let params = json!({
505 "path": &path_str,
506 "content": "first "
507 });
508 tool.execute("test-id-1", params, None, &ToolContext::default())
509 .await
510 .unwrap();
511
512 let params = json!({
514 "path": &path_str,
515 "content": "second",
516 "append": true
517 });
518 let result = tool
519 .execute("test-id-2", params, None, &ToolContext::default())
520 .await
521 .unwrap();
522 assert!(result.success);
523 assert!(result.output.contains("Appended"));
524
525 let written = std::fs::read_to_string(&path).unwrap();
526 assert_eq!(written, "first second");
527 }
528
529 #[test]
530 fn test_default_impl() {
531 let tool = WriteTool::default();
532 assert_eq!(tool.name(), "write");
533 assert_eq!(tool.label(), "Write File");
534 }
535
536 #[test]
537 fn test_parameters_schema_required_fields() {
538 let tool = WriteTool::new();
539 let schema = tool.parameters_schema();
540 let required = schema.get("required").unwrap().as_array().unwrap();
541 assert!(required.contains(&json!("path")));
542 assert!(required.contains(&json!("content")));
543 assert!(!required.contains(&json!("append"))); }
545}