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