1use std::sync::Arc;
8
9use schemars::JsonSchema;
10use serde::Deserialize;
11use serde_json::Value;
12use sgr_agent_core::agent_tool::{Tool, ToolError, ToolOutput, parse_args};
13use sgr_agent_core::context::AgentContext;
14use sgr_agent_core::schema::json_schema_for;
15
16use crate::backend::FileBackend;
17use crate::helpers::backend_err;
18
19pub struct WriteTool<B: FileBackend>(pub Arc<B>);
20
21#[derive(Deserialize, JsonSchema)]
22struct WriteArgs {
23 path: String,
25 content: String,
27 #[serde(default)]
29 start_line: i32,
30 #[serde(default)]
32 end_line: i32,
33}
34
35fn maybe_repair_json(path: &str, content: &str) -> String {
38 if !path.ends_with(".json") {
39 return content.to_string();
40 }
41 match serde_json::from_str::<serde_json::Value>(content) {
42 Ok(_) => content.to_string(),
43 Err(_) => {
44 let opts = llm_json::RepairOptions::default();
45 llm_json::repair_json(content, &opts).unwrap_or_else(|_| content.to_string())
46 }
47 }
48}
49
50#[async_trait::async_trait]
51impl<B: FileBackend> Tool for WriteTool<B> {
52 fn name(&self) -> &str {
53 "write"
54 }
55 fn description(&self) -> &str {
56 "Write content to a file. Without start_line/end_line: overwrites entire file. \
57 With start_line and end_line: replaces only those lines (like sed). \
58 Example: start_line=5, end_line=7 replaces lines 5-7 with content. \
59 Use read with number=true first to see line numbers."
60 }
61 fn parameters_schema(&self) -> Value {
62 json_schema_for::<WriteArgs>()
63 }
64 async fn execute(&self, args: Value, _ctx: &mut AgentContext) -> Result<ToolOutput, ToolError> {
65 let a: WriteArgs = parse_args(&args)?;
66 let content = maybe_repair_json(&a.path, &a.content);
67
68 self.0
69 .write(&a.path, &content, a.start_line, a.end_line)
70 .await
71 .map_err(backend_err)?;
72
73 let msg = if a.start_line > 0 && a.end_line > 0 {
74 format!(
75 "Replaced lines {}-{} in {}",
76 a.start_line, a.end_line, a.path
77 )
78 } else if a.start_line > 0 {
79 format!("Replaced from line {} in {}", a.start_line, a.path)
80 } else {
81 format!("Written to {}", a.path)
82 };
83 Ok(ToolOutput::text(msg))
84 }
85}
86
87#[cfg(test)]
88mod tests {
89 use super::*;
90 use crate::mock_fs::MockFs;
91 use sgr_agent_core::agent_tool::Tool;
92
93 #[tokio::test]
94 async fn test_write_new_file() {
95 let fs = Arc::new(MockFs::new());
96 let tool = WriteTool(fs.clone());
97 let mut ctx = AgentContext::new();
98 let result = tool
99 .execute(
100 serde_json::json!({"path": "out.txt", "content": "hello"}),
101 &mut ctx,
102 )
103 .await
104 .unwrap();
105 assert!(result.content.contains("Written to out.txt"));
106 assert_eq!(fs.content("out.txt").unwrap(), "hello");
107 }
108
109 #[tokio::test]
110 async fn test_write_json_repair() {
111 let fs = Arc::new(MockFs::new());
112 let tool = WriteTool(fs.clone());
113 let mut ctx = AgentContext::new();
114 let result = tool
116 .execute(
117 serde_json::json!({"path": "data.json", "content": "{\"key\": \"value\""}),
118 &mut ctx,
119 )
120 .await
121 .unwrap();
122 assert!(result.content.contains("Written to data.json"));
123 let stored = fs.content("data.json").unwrap();
124 assert!(serde_json::from_str::<serde_json::Value>(&stored).is_ok());
126 }
127}