Skip to main content

tycode_core/file/modify/
write_file.rs

1use crate::chat::events::{ToolExecutionResult, ToolRequest as ToolRequestEvent, ToolRequestType};
2use crate::file::access::FileAccessManager;
3use crate::file::manager::FileModificationManager;
4use crate::tools::r#trait::{
5    ContinuationPreference, FileModification, FileOperation, ToolCallHandle, ToolCategory,
6    ToolExecutor, ToolOutput, ToolRequest,
7};
8use crate::tools::ToolName;
9use anyhow::Result;
10use serde_json::{json, Value};
11use std::path::PathBuf;
12
13#[derive(Clone)]
14pub struct WriteFileTool {
15    file_manager: FileAccessManager,
16}
17
18impl WriteFileTool {
19    pub fn tool_name() -> ToolName {
20        ToolName::new("write_file")
21    }
22
23    pub fn new(workspace_roots: Vec<PathBuf>) -> anyhow::Result<Self> {
24        let file_manager = FileAccessManager::new(workspace_roots)?;
25        Ok(Self { file_manager })
26    }
27}
28
29struct WriteFileHandle {
30    modification: FileModification,
31    tool_use_id: String,
32    file_manager: FileAccessManager,
33}
34
35#[async_trait::async_trait(?Send)]
36impl ToolCallHandle for WriteFileHandle {
37    fn tool_request(&self) -> ToolRequestEvent {
38        ToolRequestEvent {
39            tool_call_id: self.tool_use_id.clone(),
40            tool_name: "write_file".to_string(),
41            tool_type: ToolRequestType::ModifyFile {
42                file_path: self.modification.path.to_string_lossy().to_string(),
43                before: self
44                    .modification
45                    .original_content
46                    .clone()
47                    .unwrap_or_default(),
48                after: self.modification.new_content.clone().unwrap_or_default(),
49            },
50        }
51    }
52
53    async fn execute(self: Box<Self>) -> ToolOutput {
54        let manager = FileModificationManager::new(self.file_manager.clone());
55        match manager.apply_modification(self.modification).await {
56            Ok(stats) => ToolOutput::Result {
57                content: json!({
58                    "success": true,
59                    "lines_added": stats.lines_added,
60                    "lines_removed": stats.lines_removed
61                })
62                .to_string(),
63                is_error: false,
64                continuation: ContinuationPreference::Continue,
65                ui_result: ToolExecutionResult::ModifyFile {
66                    lines_added: stats.lines_added,
67                    lines_removed: stats.lines_removed,
68                },
69            },
70            Err(e) => {
71                let msg = format!("{e:?}");
72                ToolOutput::Result {
73                    content: msg.clone(),
74                    is_error: true,
75                    continuation: ContinuationPreference::Continue,
76                    ui_result: ToolExecutionResult::Error {
77                        short_message: if msg.len() > 100 {
78                            format!("{}...", &msg[..97])
79                        } else {
80                            msg.clone()
81                        },
82                        detailed_message: msg,
83                    },
84                }
85            }
86        }
87    }
88}
89
90#[async_trait::async_trait(?Send)]
91impl ToolExecutor for WriteFileTool {
92    fn name(&self) -> String {
93        "write_file".to_string()
94    }
95
96    fn description(&self) -> String {
97        "Create a new file or completely overwrite an existing file".to_string()
98    }
99
100    fn input_schema(&self) -> Value {
101        json!({
102            "type": "object",
103            "properties": {
104                "file_path": {
105                    "type": "string",
106                    "description": "Path where the file should be created"
107                },
108                "content": {
109                    "type": "string",
110                    "description": "Complete content to write to the file"
111                }
112            },
113            "required": ["file_path", "content"]
114        })
115    }
116
117    fn category(&self) -> ToolCategory {
118        ToolCategory::Execution
119    }
120
121    async fn process(&self, request: &ToolRequest) -> Result<Box<dyn ToolCallHandle>> {
122        let file_path = request
123            .arguments
124            .get("file_path")
125            .and_then(|v| v.as_str())
126            .ok_or_else(|| anyhow::anyhow!("Missing required parameter: file_path"))?;
127
128        let content = request.arguments
129            .get("content")
130            .and_then(|v| v.as_str())
131            .ok_or_else(|| anyhow::anyhow!("Missing required parameter: content. Sometimes this can happen if you hit a token limit; try writing a smaller file"))?;
132
133        let original_content = self.file_manager.read_file(file_path).await.ok();
134        let operation = if original_content.is_some() {
135            FileOperation::Update
136        } else {
137            FileOperation::Create
138        };
139
140        let modification = FileModification {
141            path: PathBuf::from(file_path),
142            operation,
143            original_content,
144            new_content: Some(content.to_string()),
145            warning: None,
146        };
147
148        Ok(Box::new(WriteFileHandle {
149            modification,
150            tool_use_id: request.tool_use_id.clone(),
151            file_manager: self.file_manager.clone(),
152        }))
153    }
154}