1use std::path::PathBuf;
2
3use async_trait::async_trait;
4use cookie_agent_engine::{
5 PreparedExecutor, PreparedSerializationKey, PreparedTool, SessionToolContext, ToolCall,
6 ToolError, ToolExecutionContext, ToolPreparationContext, ToolProvider, ToolSpec,
7};
8use cookie_agent_protocol::{PermissionAction, PersistedToolResult as ToolResult};
9use schemars::JsonSchema;
10use serde::{Deserialize, Serialize};
11
12use crate::{fs_cap, parse_args, prepared_operation, prepared_path_resources, schema};
13
14#[derive(Debug)]
15pub struct WriteTool {
16 workspace: PathBuf,
17}
18
19#[derive(Debug, Deserialize, JsonSchema, Serialize)]
20struct WriteArgs {
21 #[serde(
22 rename = "filePath",
23 deserialize_with = "crate::path_args::deserialize"
24 )]
25 file_path: String,
27 content: String,
28}
29
30struct WriteExecutor {
31 target: WriteTarget,
32 bytes: Vec<u8>,
33}
34
35enum WriteTarget {
36 Existing(fs_cap::PreparedExisting),
37 Absent(fs_cap::PreparedAbsent),
38 Chained(fs_cap::PreparedChained),
39}
40
41impl WriteTool {
42 #[must_use]
43 pub fn new(workspace: impl Into<PathBuf>) -> Self {
44 Self {
45 workspace: workspace.into(),
46 }
47 }
48
49 fn prepare_write(
55 &self,
56 ctx: &ToolPreparationContext,
57 args: WriteArgs,
58 target: fs_cap::PreparedTarget,
59 predecessor: Option<cookie_agent_protocol::Sha256Digest>,
60 ) -> Result<PreparedTool, ToolError> {
61 let mut binding = match &target {
62 fs_cap::PreparedTarget::Existing(existing) => {
63 if existing.directory {
64 return Err(ToolError::unsupported_security(
65 "write target is a directory",
66 ));
67 }
68 let mut bytes = existing.identity.canonical_bytes();
69 bytes.extend_from_slice(
70 predecessor
71 .clone()
72 .unwrap_or_else(|| existing.content_digest.clone())
73 .as_str()
74 .as_bytes(),
75 );
76 bytes
77 }
78 fs_cap::PreparedTarget::Absent(_) => b"expected-absent".to_vec(),
79 };
80 binding.extend_from_slice(
81 cookie_agent_protocol::Sha256Digest::of_bytes(args.content.as_bytes())
82 .as_str()
83 .as_bytes(),
84 );
85 binding.extend_from_slice(&target.manifest_bytes()?);
86 let display_path = match &target {
87 fs_cap::PreparedTarget::Existing(target) => &target.display_path,
88 fs_cap::PreparedTarget::Absent(target) => &target.display_path,
89 };
90 let (resources, policy_labels) = prepared_path_resources(
91 PermissionAction::Write,
92 "file",
93 display_path,
94 &self.workspace,
95 &binding,
96 )?;
97 let serialization_key = target.serialization_bytes()?;
98 let context = fs_cap::cwd_context_bytes(&ctx.cwd)?;
99 let operation = prepared_operation(
100 "write",
101 &args,
102 vec![(PermissionAction::Write, "replace")],
103 resources,
104 &context,
105 )?;
106 let normalized_arguments = serde_json::json!({
107 "filePath": display_path,
108 "content": args.content,
109 });
110 let executor_target = match (target, predecessor) {
111 (fs_cap::PreparedTarget::Existing(existing), Some(predecessor_digest)) => {
112 WriteTarget::Chained(fs_cap::PreparedChained {
113 original_digest: existing.content_digest.clone(),
114 predecessor_digest,
115 target: existing,
116 })
117 }
118 (fs_cap::PreparedTarget::Existing(existing), None) => WriteTarget::Existing(existing),
119 (fs_cap::PreparedTarget::Absent(absent), _) => WriteTarget::Absent(absent),
121 };
122 PreparedTool::new(
123 operation,
124 normalized_arguments,
125 Some(PreparedSerializationKey::new(serialization_key)),
126 Box::new(WriteExecutor {
127 target: executor_target,
128 bytes: args.content.into_bytes(),
129 }),
130 )?
131 .with_policy_labels(policy_labels)
132 }
133}
134
135impl Default for WriteTool {
136 fn default() -> Self {
137 Self::new(std::env::current_dir().expect("current directory"))
138 }
139}
140
141#[async_trait]
142impl ToolProvider for WriteTool {
143 fn provider_id(&self) -> &'static str {
144 "builtin.write"
145 }
146
147 fn tools_for_session(&self, _: &SessionToolContext) -> Result<Vec<ToolSpec>, ToolError> {
148 Ok(vec![ToolSpec {
149 output: Default::default(),
150 concurrency: cookie_agent_engine::ToolConcurrency::Parallel,
151 result_truncation: Default::default(),
152 name: "write".into(),
153 permission_name: Self::get_permission_name("write")?.into(),
154 description: "Atomically write an exact descriptor-bound target.".into(),
155 parameters: schema::<WriteArgs>(),
156 }])
157 }
158
159 fn get_permission_name(tool_name: &str) -> Result<&'static str, ToolError> {
160 match tool_name {
161 "write" => Ok("write"),
162 _ => Err(ToolError::execution("write provider received another tool")),
163 }
164 }
165
166 fn get_permission_resource(
167 &self,
168 name: &str,
169 arguments: &serde_json::Value,
170 ) -> Result<(&'static str, Option<String>), ToolError> {
171 let permission_name = Self::get_permission_name(name)?;
172 let args: WriteArgs = parse_args("write", arguments.clone())?;
173 if args.file_path.is_empty() {
174 return Err(ToolError::execution("filePath must not be empty"));
175 }
176 Ok((
177 permission_name,
178 Some(crate::permission_path_label(
179 &args.file_path,
180 &self.workspace,
181 )),
182 ))
183 }
184
185 fn get_display_argument(
186 &self,
187 name: &str,
188 arguments: &serde_json::Value,
189 ) -> Result<String, ToolError> {
190 let (_, Some(path)) = self.get_permission_resource(name, arguments)? else {
191 return Err(ToolError::execution("write permission resource is missing"));
192 };
193 Ok(crate::abbreviated_display_path(&path, &self.workspace))
194 }
195
196 async fn prepare(
197 &self,
198 ctx: ToolPreparationContext,
199 call: ToolCall,
200 ) -> Result<PreparedTool, ToolError> {
201 let args: WriteArgs = parse_args("write", call.arguments)?;
202 fs_cap::ensure_atomic_write_supported()?;
203 let target = fs_cap::prepare_target(&ctx.cwd, std::path::Path::new(&args.file_path))?;
204 self.prepare_write(&ctx, args, target, None)
205 }
206
207 async fn prepare_parallel(
208 &self,
209 ctx: ToolPreparationContext,
210 calls: Vec<ToolCall>,
211 ) -> Vec<Result<PreparedTool, ToolError>> {
212 if calls.len() == 1 {
213 let call = calls.into_iter().next().expect("single call batch");
214 return vec![self.prepare(ctx, call).await];
215 }
216 let mut pending = Vec::with_capacity(calls.len());
221 for call in calls {
222 pending.push(Some((|| {
223 let args: WriteArgs = parse_args("write", call.arguments)?;
224 fs_cap::ensure_atomic_write_supported()?;
225 let target =
226 fs_cap::prepare_target(&ctx.cwd, std::path::Path::new(&args.file_path))?;
227 Ok((args, target))
228 })()));
229 }
230 let mut groups: Vec<((u64, u64), Vec<usize>)> = Vec::new();
231 for (index, entry) in pending.iter().enumerate() {
232 let Some(Ok((_, fs_cap::PreparedTarget::Existing(existing)))) = entry else {
233 continue;
234 };
235 let key = (existing.identity.device, existing.identity.inode);
236 if let Some((_, indexes)) = groups.iter_mut().find(|(candidate, _)| *candidate == key) {
237 indexes.push(index);
238 } else {
239 groups.push((key, vec![index]));
240 }
241 }
242 let mut results: Vec<Option<Result<PreparedTool, ToolError>>> =
243 (0..pending.len()).map(|_| None).collect();
244 for (_, indexes) in groups {
245 let mut predecessor = None;
249 for index in indexes {
250 let (args, target) = pending[index]
251 .take()
252 .expect("grouped write pending")
253 .expect("grouped write prepared");
254 let output_digest =
255 cookie_agent_protocol::Sha256Digest::of_bytes(args.content.as_bytes());
256 let result = self.prepare_write(&ctx, args, target, predecessor.clone());
257 if result.is_ok() {
258 predecessor = Some(output_digest);
259 }
260 results[index] = Some(result);
261 }
262 }
263 for (index, entry) in pending.into_iter().enumerate() {
264 if results[index].is_none() {
265 results[index] = Some(
266 entry
267 .expect("ungrouped write pending")
268 .and_then(|(args, target)| self.prepare_write(&ctx, args, target, None)),
269 );
270 }
271 }
272 results
273 .into_iter()
274 .map(|result| result.expect("every write call is prepared"))
275 .collect()
276 }
277}
278
279#[async_trait]
280impl PreparedExecutor for WriteExecutor {
281 async fn revalidate(&self) -> Result<(), ToolError> {
282 match &self.target {
283 WriteTarget::Existing(target) => target.revalidate(),
284 WriteTarget::Absent(target) => target.revalidate(),
285 WriteTarget::Chained(target) => target.revalidate(),
286 }
287 }
288
289 async fn revalidate_for_execution(&self) -> Result<(), ToolError> {
290 match &self.target {
291 WriteTarget::Existing(target) => target.revalidate(),
292 WriteTarget::Absent(target) => target.revalidate(),
293 WriteTarget::Chained(target) => target.revalidate_for_execution(),
294 }
295 }
296
297 async fn execute(
298 self: Box<Self>,
299 context: ToolExecutionContext,
300 ) -> Result<cookie_agent_engine::ToolCompletion, ToolError> {
301 let result: Result<ToolResult, ToolError> = async move {
302 if context.cancellation.is_cancelled() {
303 return Err(ToolError::execution(
304 "prepared write cancelled before commit",
305 ));
306 }
307 let (path, outcome) = match &self.target {
308 WriteTarget::Existing(target) => {
309 let outcome = target.replace_atomically(&self.bytes)?;
310 (target.display_path.clone(), outcome)
311 }
312 WriteTarget::Absent(target) => {
313 let outcome = target.create_atomically(&self.bytes)?;
314 (target.display_path.clone(), outcome)
315 }
316 WriteTarget::Chained(target) => {
317 target.revalidate_for_execution()?;
320 let outcome = target.replace_atomically(&self.bytes)?;
321 (target.target.display_path.clone(), outcome)
322 }
323 };
324 Ok(ToolResult {
325display: None,
326retained_output: None,
327 title: crate::safe_title(format!("Wrote {}", path.display())),
328 output: format!("Wrote {} bytes to {}", self.bytes.len(), path.display()),
329 metadata: serde_json::json!({"bytes":self.bytes.len(),"cleanup_warning":outcome.cleanup_warning}),
330 truncation: None,
331 attachments: Vec::new(),
332 additional_messages: Vec::new(),
333 })
334}.await;
335 result.map(cookie_agent_engine::ToolCompletion::single)
336 }
337}
338
339#[cfg(test)]
340mod tests {
341 #[cfg(any(unix, windows))]
342 use cookie_agent_engine::ToolExecutionContext;
343 use cookie_agent_engine::{ToolCall, ToolError, ToolPreparationContext, ToolProvider};
344 use cookie_agent_protocol::{
345 OperationFingerprint, PermissionAction, RunId, SessionId, ToolCallId,
346 };
347
348 use super::WriteTool;
349
350 #[test]
351 fn permission_resource_is_the_file_path() {
352 let tool = WriteTool::new("/tmp");
353 assert_eq!(
354 tool.get_permission_resource(
355 "write",
356 &serde_json::json!({"filePath":"out.txt","content":"x"})
357 )
358 .expect("permission resource"),
359 ("write", Some("out.txt".into()))
360 );
361 assert!(matches!(
362 tool.get_permission_resource("write", &serde_json::json!({"content":"x"})),
363 Err(ToolError::Failed(_))
364 ));
365 }
366
367 #[test]
368 fn display_argument_abbreviates_workspace_and_home_paths() {
369 let workspace = tempfile::tempdir().expect("workspace");
370 let tool = WriteTool::new(workspace.path());
371 assert_eq!(
372 tool.get_display_argument(
373 "write",
374 &serde_json::json!({"filePath":"out.txt","content":"x"})
375 )
376 .expect("relative"),
377 "out.txt"
378 );
379 assert_eq!(
380 tool.get_display_argument(
381 "write",
382 &serde_json::json!({"filePath":workspace.path().join("out.txt"),"content":"x"})
383 )
384 .expect("workspace"),
385 "out.txt"
386 );
387 let home = cookie_agent_protocol::paths::home_dir().expect("home directory");
388 assert_eq!(
389 tool.get_display_argument(
390 "write",
391 &serde_json::json!({"filePath": home.join("notes.txt"),"content":"x"})
392 )
393 .expect("home"),
394 "~/notes.txt"
395 );
396 }
397
398 fn context(root: &std::path::Path) -> ToolPreparationContext {
399 ToolPreparationContext {
400 session: SessionId::new_v7(),
401 run: RunId::new_v7(),
402 cwd: root.to_owned(),
403 workspace_root: root.to_owned(),
404 turn_context: crate::test_turn_context(),
405 }
406 }
407
408 #[tokio::test]
409 async fn distinct_write_content_has_distinct_fingerprint() {
410 let root = tempfile::tempdir().expect("root");
411 let tool = WriteTool::new(root.path());
412 let first = tool
413 .prepare(
414 context(root.path()),
415 ToolCall {
416 id: ToolCallId::new_v7(),
417 name: "write".into(),
418 arguments: serde_json::json!({"filePath":"value.txt","content":"one"}),
419 },
420 )
421 .await
422 .expect("first");
423 let second = tool
424 .prepare(
425 context(root.path()),
426 ToolCall {
427 id: ToolCallId::new_v7(),
428 name: "write".into(),
429 arguments: serde_json::json!({"filePath":"value.txt","content":"two"}),
430 },
431 )
432 .await
433 .expect("second");
434 assert_ne!(
435 OperationFingerprint::from_prepared_operation(first.operation()),
436 OperationFingerprint::from_prepared_operation(second.operation())
437 );
438 }
439
440 #[tokio::test]
441 async fn missing_subtree_write_prepares_without_creating_components() {
442 let root = tempfile::tempdir().expect("root");
443 let prepared = WriteTool::new(root.path())
444 .prepare(
445 context(root.path()),
446 ToolCall {
447 id: ToolCallId::new_v7(),
448 name: "write".into(),
449 arguments: serde_json::json!({"filePath":"a/b/value.txt","content":"value"}),
450 },
451 )
452 .await
453 .expect("prepare");
454 assert!(!root.path().join("a").exists());
455 assert_eq!(
456 prepared.operation().resources()[0].capability,
457 PermissionAction::Write
458 );
459 }
460
461 #[cfg(unix)]
462 #[tokio::test]
463 async fn symlink_routes_keep_lexical_write_contract_and_preserve_links() {
464 use std::{fs, os::unix::fs::symlink};
465
466 let root = tempfile::tempdir().expect("root");
467 let workspace = root.path().join("workspace");
468 let external = root.path().join("external");
469 fs::create_dir(&workspace).expect("workspace");
470 fs::create_dir(workspace.join("routes")).expect("routes");
471 fs::create_dir(&external).expect("external");
472 for name in [
473 "relative-leaf.txt",
474 "absolute-leaf.txt",
475 "relative-ancestor.txt",
476 "absolute-ancestor.txt",
477 ] {
478 fs::write(external.join(name), "old").expect("fixture");
479 }
480 symlink(
481 "../../external/relative-leaf.txt",
482 workspace.join("routes/relative-leaf"),
483 )
484 .expect("relative leaf");
485 symlink(
486 external.join("absolute-leaf.txt"),
487 workspace.join("routes/absolute-leaf"),
488 )
489 .expect("absolute leaf");
490 symlink(
491 "absolute-ancestor",
492 workspace.join("routes/relative-ancestor"),
493 )
494 .expect("relative ancestor");
495 symlink(&external, workspace.join("routes/absolute-ancestor")).expect("absolute ancestor");
496
497 let tool = WriteTool::new(&workspace);
498 for (requested, destination, link) in [
499 (
500 "routes/relative-leaf",
501 external.join("relative-leaf.txt"),
502 workspace.join("routes/relative-leaf"),
503 ),
504 (
505 "routes/absolute-leaf",
506 external.join("absolute-leaf.txt"),
507 workspace.join("routes/absolute-leaf"),
508 ),
509 (
510 "routes/relative-ancestor/relative-ancestor.txt",
511 external.join("relative-ancestor.txt"),
512 workspace.join("routes/relative-ancestor"),
513 ),
514 (
515 "routes/absolute-ancestor/absolute-ancestor.txt",
516 external.join("absolute-ancestor.txt"),
517 workspace.join("routes/absolute-ancestor"),
518 ),
519 ] {
520 let original_link = fs::read_link(&link).expect("link target");
521 let prepared = tool
522 .prepare(
523 context(&workspace),
524 ToolCall {
525 id: ToolCallId::new_v7(),
526 name: "write".into(),
527 arguments: serde_json::json!({"filePath":requested,"content":"new"}),
528 },
529 )
530 .await
531 .expect("prepare symlink write");
532 assert_eq!(prepared.operation().resources().len(), 1);
533 assert_eq!(prepared.policy_labels(), [Some(requested.to_owned())]);
534 assert_eq!(
535 prepared.normalized_arguments()["filePath"],
536 workspace.join(requested).to_string_lossy().as_ref()
537 );
538 assert_eq!(
539 tool.get_permission_resource("write", prepared.normalized_arguments())
540 .expect("normalized permission resource"),
541 ("write", Some(requested.to_owned()))
542 );
543 crate::assert_workspace_rule_allows(
544 &prepared,
545 &workspace,
546 PermissionAction::Write,
547 "routes/*",
548 );
549 prepared
550 .execute_for_test(
551 ToolExecutionContext::for_test(
552 root.path().join("artifacts"),
553 crate::test_turn_context(),
554 )
555 .expect("execution context"),
556 )
557 .await
558 .expect("execute symlink write");
559 assert_eq!(fs::read_to_string(destination).unwrap(), "new");
560 assert_eq!(fs::read_link(link).unwrap(), original_link);
561 }
562 }
563
564 #[cfg(unix)]
565 #[tokio::test]
566 async fn write_follows_dangling_leaf_for_creation_without_replacing_link() {
567 use std::{fs, os::unix::fs::symlink};
568
569 let root = tempfile::tempdir().expect("root");
570 symlink("created.txt", root.path().join("route")).expect("dangling route");
571 let prepared = WriteTool::new(root.path())
572 .prepare(
573 context(root.path()),
574 ToolCall {
575 id: ToolCallId::new_v7(),
576 name: "write".into(),
577 arguments: serde_json::json!({"filePath":"route","content":"created"}),
578 },
579 )
580 .await
581 .expect("prepare dangling write");
582 prepared
583 .execute_for_test(
584 ToolExecutionContext::for_test(
585 root.path().join("artifacts"),
586 crate::test_turn_context(),
587 )
588 .expect("execution context"),
589 )
590 .await
591 .expect("execute dangling write");
592
593 assert_eq!(
594 fs::read_link(root.path().join("route")).unwrap(),
595 std::path::Path::new("created.txt")
596 );
597 assert_eq!(
598 fs::read_to_string(root.path().join("created.txt")).unwrap(),
599 "created"
600 );
601 }
602
603 #[cfg(unix)]
604 #[tokio::test]
605 async fn write_route_swap_is_operation_changed_and_routes_affect_fingerprint() {
606 use std::{fs, os::unix::fs::symlink};
607
608 let root = tempfile::tempdir().expect("root");
609 fs::write(root.path().join("value.txt"), "old").expect("fixture");
610 symlink("value.txt", root.path().join("left")).expect("left route");
611 symlink("value.txt", root.path().join("right")).expect("right route");
612 let tool = WriteTool::new(root.path());
613 let prepare = |path: &str| {
614 tool.prepare(
615 context(root.path()),
616 ToolCall {
617 id: ToolCallId::new_v7(),
618 name: "write".into(),
619 arguments: serde_json::json!({"filePath":path,"content":"new"}),
620 },
621 )
622 };
623 let left = prepare("left").await.expect("prepare left");
624 let right = prepare("right").await.expect("prepare right");
625 assert_ne!(
626 left.operation().resources()[0].binding_digest,
627 right.operation().resources()[0].binding_digest
628 );
629 assert_ne!(
630 OperationFingerprint::from_prepared_operation(left.operation()),
631 OperationFingerprint::from_prepared_operation(right.operation())
632 );
633
634 fs::remove_file(root.path().join("left")).expect("remove route");
635 symlink("value.txt", root.path().join("left")).expect("replace route");
636 let error = left
637 .execute_for_test(
638 ToolExecutionContext::for_test(
639 root.path().join("artifacts"),
640 crate::test_turn_context(),
641 )
642 .expect("execution context"),
643 )
644 .await
645 .expect_err("route swap must fail");
646 assert!(matches!(error, ToolError::OperationChanged(_)));
647 }
648
649 #[cfg(any(unix, windows))]
650 fn write_call(path: &str, content: &str) -> ToolCall {
651 ToolCall {
652 id: ToolCallId::new_v7(),
653 name: "write".into(),
654 arguments: serde_json::json!({"filePath": path, "content": content}),
655 }
656 }
657
658 #[cfg(any(unix, windows))]
659 fn execution_context(root: &std::path::Path) -> ToolExecutionContext {
660 ToolExecutionContext::for_test(root.join("artifacts"), crate::test_turn_context())
661 .expect("execution context")
662 }
663
664 #[cfg(any(unix, windows))]
665 #[tokio::test]
666 async fn parallel_same_file_writes_chain_in_call_order() {
667 use std::fs;
668
669 let root = tempfile::tempdir().expect("root");
670 fs::write(root.path().join("value.txt"), "old").expect("fixture");
671 let prepared = WriteTool::new(root.path())
672 .prepare_parallel(
673 context(root.path()),
674 vec![
675 write_call("value.txt", "one"),
676 write_call("value.txt", "two"),
677 ],
678 )
679 .await;
680 let mut prepared = prepared.into_iter();
681 let first = prepared.next().expect("first").expect("first write");
682 let second = prepared.next().expect("second").expect("second write");
683 assert!(prepared.next().is_none());
684 assert_ne!(
685 OperationFingerprint::from_prepared_operation(first.operation()),
686 OperationFingerprint::from_prepared_operation(second.operation())
687 );
688
689 let error = second
691 .revalidate_for_execution_for_test()
692 .await
693 .expect_err("predecessor has not run yet");
694 assert!(matches!(error, ToolError::OperationChanged(_)));
695
696 first
697 .execute_for_test(execution_context(root.path()))
698 .await
699 .expect("first write");
700 second
701 .execute_for_test(execution_context(root.path()))
702 .await
703 .expect("chained write");
704 assert_eq!(
706 fs::read_to_string(root.path().join("value.txt")).unwrap(),
707 "two"
708 );
709 }
710
711 #[cfg(any(unix, windows))]
712 #[tokio::test]
713 async fn absent_write_targets_are_not_chained() {
714 let root = tempfile::tempdir().expect("root");
715 let prepared = WriteTool::new(root.path())
716 .prepare_parallel(
717 context(root.path()),
718 vec![
719 write_call("created.txt", "one"),
720 write_call("created.txt", "two"),
721 ],
722 )
723 .await;
724 let mut prepared = prepared.into_iter();
725 let first = prepared.next().expect("first").expect("first write");
726 let second = prepared.next().expect("second").expect("second write");
727 first
728 .execute_for_test(execution_context(root.path()))
729 .await
730 .expect("first creation");
731 let error = second
732 .execute_for_test(execution_context(root.path()))
733 .await
734 .expect_err("second creation must observe the first");
735 assert!(matches!(error, ToolError::OperationChanged(_)));
736 }
737}