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 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/// Typed arguments for [`WriteTool`].
23#[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;
34/// Threshold above which we switch from full-content to head/tail preview display
35const PREVIEW_THRESHOLD_LINES: usize = 20;
36
37/// WriteTool.
38pub struct WriteTool {
39    root_dir: Option<PathBuf>,
40}
41
42impl WriteTool {
43    /// Create with no explicit root (uses ToolContext.workspace_dir at runtime).
44    pub fn new() -> Self {
45        Self { root_dir: None }
46    }
47
48    /// Create with a specific working directory (overrides ToolContext).
49    pub fn with_cwd(cwd: PathBuf) -> Self {
50        Self {
51            root_dir: Some(cwd),
52        }
53    }
54
55    /// Build a human-readable preview of the content that was written.
56    /// For small files, shows everything. For large files, shows first/last few lines.
57    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    /// Core write implementation — runs inside the mutation queue lock.
83    async fn write_file_impl(
84        root_dir: &Path,
85        path: &str,
86        content: &str,
87        append: bool,
88    ) -> Result<String, ToolError> {
89        // Security: validate path with PathGuard
90        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        // Ensure parent directory exists (create if missing)
96        if let Some(parent) = file_path.parent() {
97            // Only try to create if the parent is non-empty (e.g. not just "")
98            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        // Check if file already existed before write (for reporting)
106        let existed = file_path.exists();
107
108        // Perform the write through the mutation queue for serialized access
109        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        // Build result with preview
151        let preview = Self::build_content_preview(content, total_lines);
152
153        // Truncate the preview if very large
154        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(&notify_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")); // middle should be omitted
299    }
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        // At threshold, should show full content
307        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        // Over threshold, should show head/tail
316        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        // Create initial file
359        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        // Write initial content
379        WriteTool::write_file_impl(Path::new("."), path_str, "line 1\n", false)
380            .await
381            .unwrap();
382
383        // Append to it
384        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        // First write
510        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        // Append
519        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"))); // append is optional
550    }
551}