1use async_trait::async_trait;
2use bamboo_agent_core::{Tool, ToolCtx, ToolError, ToolOutcome, ToolResult};
3use serde::Deserialize;
4use serde_json::json;
5use std::path::Path;
6
7use super::read_tracker::{BaselineAdvance, ReadState};
8use super::{content_diagnostics, file_change, read_tracker};
9
10#[derive(Debug, Deserialize)]
11struct WriteArgs {
12 file_path: String,
13 content: String,
14}
15
16pub struct WriteTool;
17
18impl WriteTool {
19 pub fn new() -> Self {
20 Self
21 }
22}
23
24impl Default for WriteTool {
25 fn default() -> Self {
26 Self::new()
27 }
28}
29
30#[async_trait]
31impl Tool for WriteTool {
32 fn name(&self) -> &str {
33 "Write"
34 }
35
36 fn description(&self) -> &str {
37 "Write a local file (create or replace full content). IMPORTANT: for existing files, call Read first in this session or Write will fail."
38 }
39
40 fn parameters_schema(&self) -> serde_json::Value {
41 json!({
42 "type": "object",
43 "properties": {
44 "file_path": {
45 "type": "string",
46 "description": "The absolute path to the file to write"
47 },
48 "content": {
49 "type": "string",
50 "description": "The content to write to the file"
51 }
52 },
53 "required": ["file_path", "content"],
54 "additionalProperties": false
55 })
56 }
57
58 async fn invoke(
59 &self,
60 args: serde_json::Value,
61 ctx: ToolCtx,
62 ) -> Result<ToolOutcome, ToolError> {
63 let parsed: WriteArgs = serde_json::from_value(args)
64 .map_err(|e| ToolError::InvalidArguments(format!("Invalid Write args: {}", e)))?;
65
66 let file_path = parsed.file_path.trim();
67 let path = Path::new(file_path);
68
69 if !path.is_absolute() {
70 return Err(ToolError::InvalidArguments(
71 "file_path must be an absolute path".to_string(),
72 ));
73 }
74
75 let session_id = ctx.session_id().map(str::to_owned);
76 let target_existed = tokio::fs::try_exists(path)
77 .await
78 .map_err(|e| ToolError::Execution(format!("Failed to inspect target file: {}", e)))?;
79 let validated_read = if target_existed {
80 if let Some(session_id) = session_id.as_deref() {
81 match read_tracker::read_if_fresh(session_id, file_path).await {
82 Ok(validated) => Some(validated),
83 Err(ReadState::Unread) => {
84 return Err(ToolError::Execution(
85 "Write requires reading the target file first via Read".to_string(),
86 ));
87 }
88 Err(ReadState::Stale) => {
89 return Err(ToolError::Execution(
90 "Target file changed after last Read; call Read again before Write"
91 .to_string(),
92 ));
93 }
94 Err(ReadState::Fresh) => {
95 unreachable!("Fresh is returned as a validated read")
96 }
97 }
98 } else {
99 None
100 }
101 } else {
102 None
103 };
104 let new_file_slot =
105 if let (Some(session_id), false) = (session_id.as_deref(), target_existed) {
106 Some(read_tracker::capture_write_slot(session_id, file_path).await)
107 } else {
108 None
109 };
110
111 let previous_bytes = if let Some(validated) = validated_read.as_ref() {
112 Some(validated.bytes().to_vec())
113 } else if target_existed {
114 file_change::read_existing_bytes(path).await?
115 } else {
116 None
117 };
118 let checkpoint = file_change::create_checkpoint(path, previous_bytes.as_deref()).await?;
119 let next_content = parsed.content;
120
121 let write_expectation = if let Some(validated) = validated_read.as_ref() {
122 file_change::AtomicWriteExpectation::Exact(validated.bytes())
123 } else if session_id.is_some() && !target_existed {
124 file_change::AtomicWriteExpectation::Missing
125 } else {
126 file_change::AtomicWriteExpectation::Unchecked
127 };
128 file_change::atomic_write_text_with_expectation(path, &next_content, write_expectation)
129 .await?;
130
131 let mutation_slot = validated_read
132 .as_ref()
133 .map(|validated| validated.slot())
134 .or(new_file_slot.as_ref());
135 if let Some(slot) = mutation_slot {
136 if read_tracker::advance_after_verified_write(file_path, slot, next_content.as_bytes())
137 .await
138 == BaselineAdvance::Conflict
139 {
140 return Err(ToolError::Execution(
141 "Write committed, but the target changed before it could be verified; call Read again"
142 .to_string(),
143 ));
144 }
145 }
146
147 let previous_text = file_change::bytes_to_lossy_text(previous_bytes.as_deref());
148 let mut payload = file_change::build_file_change_payload_value(
149 "Write",
150 path,
151 format!("Wrote file: {}", file_path),
152 checkpoint,
153 &previous_text,
154 &next_content,
155 );
156 content_diagnostics::attach_file_diagnostics(&mut payload, path, &next_content);
157
158 Ok(ToolOutcome::Completed(ToolResult {
159 success: true,
160 result: payload.to_string(),
161 display_preference: Some("Default".to_string()),
162 images: Vec::new(),
163 }))
164 }
165}
166
167#[cfg(test)]
168mod tests {
169 use super::*;
170 use crate::tools::ReadTool;
171 use serde_json::json;
172
173 fn ctx(session_id: &str) -> ToolCtx {
174 ToolCtx {
175 session_id: Some(std::sync::Arc::from(session_id)),
176 tool_call_id: std::sync::Arc::from("call_1"),
177 event_tx: None,
178 available_tool_schemas: std::sync::Arc::from(Vec::new()),
179 bypass_permissions: false,
180 auto_approve_permissions: false,
181 plan_read_only: false,
182 can_async_resume: false,
183 async_completion_sink: None,
184 bash_completion_sink: None,
185 }
186 }
187
188 #[tokio::test]
189 async fn write_requires_fresh_read_for_existing_files() {
190 let file = tempfile::NamedTempFile::new().unwrap();
191 tokio::fs::write(file.path(), "v1").await.unwrap();
192 let write_tool = WriteTool::new();
193 let read_tool = ReadTool::new();
194
195 let denied = write_tool
196 .invoke(
197 json!({"file_path": file.path(), "content": "v2"}),
198 ctx("session_a"),
199 )
200 .await;
201 assert!(matches!(denied, Err(ToolError::Execution(_))));
202
203 let _ = read_tool
204 .invoke(json!({"file_path": file.path()}), ctx("session_a"))
205 .await
206 .unwrap();
207
208 tokio::fs::write(file.path(), "external change")
209 .await
210 .unwrap();
211
212 let stale = write_tool
213 .invoke(
214 json!({"file_path": file.path(), "content": "v3"}),
215 ctx("session_a"),
216 )
217 .await;
218 assert!(matches!(stale, Err(ToolError::Execution(msg)) if msg.contains("changed")));
219
220 let _ = read_tool
221 .invoke(json!({"file_path": file.path()}), ctx("session_a"))
222 .await
223 .unwrap();
224 let out = write_tool
225 .invoke(
226 json!({"file_path": file.path(), "content": "final"}),
227 ctx("session_a"),
228 )
229 .await
230 .unwrap();
231 let ToolOutcome::Completed(ok) = out else {
232 panic!("expected Completed")
233 };
234 assert!(ok.success);
235 }
236
237 #[tokio::test]
238 async fn read_write_write_succeeds_without_an_external_change() {
239 let file = tempfile::NamedTempFile::new().unwrap();
240 tokio::fs::write(file.path(), "v1").await.unwrap();
241 let session = format!("write-twice-{}", uuid::Uuid::new_v4());
242 let read_tool = ReadTool::new();
243 let write_tool = WriteTool::new();
244
245 read_tool
246 .invoke(json!({"file_path": file.path()}), ctx(&session))
247 .await
248 .unwrap();
249 write_tool
250 .invoke(
251 json!({"file_path": file.path(), "content": "v2"}),
252 ctx(&session),
253 )
254 .await
255 .unwrap();
256 write_tool
257 .invoke(
258 json!({"file_path": file.path(), "content": "v3"}),
259 ctx(&session),
260 )
261 .await
262 .unwrap();
263
264 assert_eq!(tokio::fs::read_to_string(file.path()).await.unwrap(), "v3");
265 }
266
267 #[tokio::test]
268 async fn write_rejects_external_change_after_a_successful_write() {
269 let file = tempfile::NamedTempFile::new().unwrap();
270 tokio::fs::write(file.path(), "aa").await.unwrap();
271 let session = format!("write-external-{}", uuid::Uuid::new_v4());
272 let read_tool = ReadTool::new();
273 let write_tool = WriteTool::new();
274
275 read_tool
276 .invoke(json!({"file_path": file.path()}), ctx(&session))
277 .await
278 .unwrap();
279 write_tool
280 .invoke(
281 json!({"file_path": file.path(), "content": "bb"}),
282 ctx(&session),
283 )
284 .await
285 .unwrap();
286
287 tokio::fs::write(file.path(), "cc").await.unwrap();
288 let stale = write_tool
289 .invoke(
290 json!({"file_path": file.path(), "content": "dd"}),
291 ctx(&session),
292 )
293 .await;
294
295 assert!(matches!(stale, Err(ToolError::Execution(message)) if message.contains("changed")));
296 assert_eq!(tokio::fs::read_to_string(file.path()).await.unwrap(), "cc");
297 }
298
299 #[tokio::test]
300 async fn stale_write_failure_does_not_advance_the_baseline() {
301 let file = tempfile::NamedTempFile::new().unwrap();
302 tokio::fs::write(file.path(), "v1").await.unwrap();
303 let session = format!("write-stale-failure-{}", uuid::Uuid::new_v4());
304 let read_tool = ReadTool::new();
305 let write_tool = WriteTool::new();
306
307 read_tool
308 .invoke(json!({"file_path": file.path()}), ctx(&session))
309 .await
310 .unwrap();
311 tokio::fs::write(file.path(), "external").await.unwrap();
312
313 for intended in ["first-attempt", "second-attempt"] {
314 let stale = write_tool
315 .invoke(
316 json!({"file_path": file.path(), "content": intended}),
317 ctx(&session),
318 )
319 .await;
320 assert!(
321 matches!(stale, Err(ToolError::Execution(message)) if message.contains("changed"))
322 );
323 }
324 assert_eq!(
325 tokio::fs::read_to_string(file.path()).await.unwrap(),
326 "external"
327 );
328 }
329
330 #[tokio::test]
331 async fn consecutive_writes_to_a_new_file_use_the_verified_first_write() {
332 let dir = tempfile::tempdir().unwrap();
333 let path = dir.path().join("new.txt");
334 let session = format!("write-new-twice-{}", uuid::Uuid::new_v4());
335 let write_tool = WriteTool::new();
336
337 write_tool
338 .invoke(
339 json!({"file_path": path, "content": "first"}),
340 ctx(&session),
341 )
342 .await
343 .unwrap();
344 write_tool
345 .invoke(
346 json!({"file_path": path, "content": "second"}),
347 ctx(&session),
348 )
349 .await
350 .unwrap();
351
352 assert_eq!(tokio::fs::read_to_string(path).await.unwrap(), "second");
353 }
354
355 #[tokio::test]
356 async fn concurrent_read_of_new_file_is_idempotent_for_first_write() {
357 let dir = tempfile::tempdir().unwrap();
358 let path = dir.path().join("new-concurrent.txt");
359 let path_str = path.to_string_lossy().into_owned();
360 let session = format!("write-new-concurrent-{}", uuid::Uuid::new_v4());
361 let (advance_reached, resume_advance) =
362 read_tracker::pause_next_advance_for_test(&session, &path_str).await;
363
364 let writer_path = path.clone();
365 let writer_session = session.clone();
366 let writer = tokio::spawn(async move {
367 WriteTool::new()
368 .invoke(
369 json!({"file_path": writer_path, "content": "first"}),
370 ctx(&writer_session),
371 )
372 .await
373 });
374
375 tokio::time::timeout(
376 std::time::Duration::from_secs(5),
377 advance_reached.notified(),
378 )
379 .await
380 .expect("Write did not reach post-write baseline advancement");
381 ReadTool::new()
382 .invoke(json!({"file_path": path}), ctx(&session))
383 .await
384 .unwrap();
385 resume_advance.notify_one();
386
387 let first = tokio::time::timeout(std::time::Duration::from_secs(5), writer)
388 .await
389 .expect("Write did not resume")
390 .unwrap()
391 .unwrap();
392 assert!(matches!(first, ToolOutcome::Completed(result) if result.success));
393
394 WriteTool::new()
395 .invoke(
396 json!({"file_path": path, "content": "second"}),
397 ctx(&session),
398 )
399 .await
400 .unwrap();
401 assert_eq!(tokio::fs::read_to_string(path).await.unwrap(), "second");
402 }
403
404 #[tokio::test]
405 async fn committed_postverify_conflict_is_clear_and_leaves_baseline_stale() {
406 let file = tempfile::NamedTempFile::new().unwrap();
407 tokio::fs::write(file.path(), "before").await.unwrap();
408 let path = file.path().to_path_buf();
409 let path_str = path.to_string_lossy().into_owned();
410 let session = format!("write-postverify-conflict-{}", uuid::Uuid::new_v4());
411
412 ReadTool::new()
413 .invoke(json!({"file_path": path}), ctx(&session))
414 .await
415 .unwrap();
416 let (advance_reached, resume_advance) =
417 read_tracker::pause_next_advance_for_test(&session, &path_str).await;
418
419 let writer_path = path.clone();
420 let writer_session = session.clone();
421 let writer = tokio::spawn(async move {
422 WriteTool::new()
423 .invoke(
424 json!({"file_path": writer_path, "content": "intended"}),
425 ctx(&writer_session),
426 )
427 .await
428 });
429
430 tokio::time::timeout(
431 std::time::Duration::from_secs(5),
432 advance_reached.notified(),
433 )
434 .await
435 .expect("Write did not reach post-write baseline advancement");
436 tokio::fs::write(&path, "other").await.unwrap();
437 ReadTool::new()
438 .invoke(json!({"file_path": path}), ctx(&session))
439 .await
440 .unwrap();
441 tokio::fs::write(&path, "intended").await.unwrap();
442 resume_advance.notify_one();
443
444 let result = tokio::time::timeout(std::time::Duration::from_secs(5), writer)
445 .await
446 .expect("Write did not resume")
447 .unwrap();
448 assert!(
449 matches!(result, Err(ToolError::Execution(message)) if message.contains("Write committed"))
450 );
451 assert_eq!(
452 read_tracker::read_state(&session, &path_str).await,
453 ReadState::Stale
454 );
455 assert_eq!(tokio::fs::read_to_string(path).await.unwrap(), "intended");
456 }
457
458 #[cfg(unix)]
459 #[tokio::test]
460 async fn write_rejects_symlinked_path_components() {
461 use std::os::unix::fs::symlink;
462 let dir = tempfile::tempdir().unwrap();
463 let real = dir.path().join("real");
464 let link = dir.path().join("link");
465 tokio::fs::create_dir_all(&real).await.unwrap();
466 symlink(&real, &link).unwrap();
467
468 let write_tool = WriteTool::new();
469 let result = write_tool
470 .invoke(
471 json!({
472 "file_path": link.join("test.txt"),
473 "content": "hello"
474 }),
475 ToolCtx::none("t"),
476 )
477 .await;
478 assert!(matches!(result, Err(ToolError::Execution(msg)) if msg.contains("symlinked")));
479 }
480
481 #[tokio::test]
482 async fn write_includes_json_diagnostics_for_invalid_content() {
483 let file = tempfile::Builder::new().suffix(".json").tempfile().unwrap();
484 let write_tool = WriteTool::new();
485
486 let out = write_tool
487 .invoke(
488 json!({
489 "file_path": file.path(),
490 "content": "{"
491 }),
492 ToolCtx::none("t"),
493 )
494 .await
495 .unwrap();
496 let ToolOutcome::Completed(result) = out else {
497 panic!("expected Completed")
498 };
499
500 let payload: serde_json::Value = serde_json::from_str(&result.result).unwrap();
501 assert_eq!(payload["diagnostics"]["format"], "json");
502 assert_eq!(payload["diagnostics"]["valid"], false);
503 }
504}