Skip to main content

oxi_agent/tools/
write.rs

1/// Write file tool
2/// Supports:
3/// - Creating parent directories if they don't exist
4/// - Append mode (append=true)
5/// - Line count reporting
6/// - Diff-style output preview (first/last few lines for large files)
7/// - File mutation queue for serialized writes (concurrent safety)
8/// - Output truncation for very large content
9use 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
19/// Maximum number of lines to show in the diff-style output preview
20const PREVIEW_HEAD_LINES: usize = 5;
21const PREVIEW_TAIL_LINES: usize = 5;
22/// Threshold above which we switch from full-content to head/tail preview display
23const PREVIEW_THRESHOLD_LINES: usize = 20;
24
25/// WriteTool.
26pub struct WriteTool {
27    root_dir: Option<PathBuf>,
28}
29
30impl WriteTool {
31    /// Create with no explicit root (uses ToolContext.workspace_dir at runtime).
32    pub fn new() -> Self {
33        Self { root_dir: None }
34    }
35
36    /// Create with a specific working directory (overrides ToolContext).
37    pub fn with_cwd(cwd: PathBuf) -> Self {
38        Self {
39            root_dir: Some(cwd),
40        }
41    }
42
43    /// Build a human-readable preview of the content that was written.
44    /// For small files, shows everything. For large files, shows first/last few lines.
45    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    /// Core write implementation — runs inside the mutation queue lock.
71    async fn write_file_impl(
72        root_dir: &Path,
73        path: &str,
74        content: &str,
75        append: bool,
76    ) -> Result<String, ToolError> {
77        // Security: validate path with PathGuard
78        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        // Ensure parent directory exists (create if missing)
84        if let Some(parent) = file_path.parent() {
85            // Only try to create if the parent is non-empty (e.g. not just "")
86            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        // Check if file already existed before write (for reporting)
94        let existed = file_path.exists();
95
96        // Perform the write through the mutation queue for serialized access
97        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        // Build result with preview
139        let preview = Self::build_content_preview(content, total_lines);
140
141        // Truncate the preview if very large
142        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        // Use root_dir if set, else ctx.root()
234        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        // Notify LSP provider so diagnostics refresh after the
238        // file's contents change. Best-effort: a transient LSP
239        // error must not fail the write.
240        // Best-effort: resolve the path for LSP notification. The
241        // `write_file_impl` already does the same PathGuard
242        // validation; this is a read-only normalization for the
243        // background task only.
244        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(&notify_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")); // middle should be omitted
293    }
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        // At threshold, should show full content
301        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        // Over threshold, should show head/tail
310        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        // Create initial file
353        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        // Write initial content
373        WriteTool::write_file_impl(Path::new("."), path_str, "line 1\n", false)
374            .await
375            .unwrap();
376
377        // Append to it
378        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        // First write
504        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        // Append
513        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"))); // append is optional
544    }
545}