pi/core/tools/
write.rs

1//! Write tool: create or overwrite a file under the mutation queue.
2//!
3//! Ports `.references/pi/packages/coding-agent/src/core/tools/write.ts`.
4//! Success text reports a JS-style UTF-16 code-unit length as "bytes".
5
6use std::path::{Path, PathBuf};
7use std::sync::Arc;
8
9use futures::FutureExt as _;
10use futures::future::BoxFuture;
11use pi_agent::{AgentTool, AgentToolResult, ToolError, ToolUpdates};
12use pi_ai::types::{TextContent, ToolResultContent};
13use schemars::JsonSchema;
14use serde::{Deserialize, Serialize};
15use serde_json::{Map, Value};
16use tokio_util::sync::CancellationToken;
17
18use super::{MutationQueueError, PathResolveError, resolve_to_cwd, with_file_mutation_queue};
19
20/// TypeBox-compatible write arguments (fixture `write.json`).
21#[derive(Clone, Debug, Deserialize, JsonSchema, PartialEq, Serialize)]
22pub struct WriteToolInput {
23    /// Path to the file to write (relative or absolute).
24    #[schemars(description = "Path to the file to write (relative or absolute)")]
25    pub path: String,
26    /// Content to write to the file.
27    #[schemars(description = "Content to write to the file")]
28    pub content: String,
29}
30
31/// Write tool has no structured `details` payload.
32pub type WriteToolDetails = ();
33
34/// Options for [`WriteTool`].
35#[derive(Clone, Debug)]
36pub struct WriteToolOptions {
37    /// Working directory used to resolve relative paths.
38    pub cwd: PathBuf,
39}
40
41impl WriteToolOptions {
42    /// Builds options for `cwd`.
43    #[must_use]
44    pub fn new(cwd: impl Into<PathBuf>) -> Self {
45        Self { cwd: cwd.into() }
46    }
47}
48
49/// Agent tool that writes UTF-8 file contents under the mutation queue.
50#[derive(Clone, Debug)]
51pub struct WriteTool {
52    cwd: PathBuf,
53    parameters: Value,
54    description: String,
55}
56
57impl WriteTool {
58    /// Creates a write tool rooted at `cwd`.
59    #[must_use]
60    pub fn new(cwd: impl Into<PathBuf>) -> Self {
61        Self::with_options(WriteToolOptions::new(cwd))
62    }
63
64    /// Creates a write tool from explicit options.
65    #[must_use]
66    pub fn with_options(options: WriteToolOptions) -> Self {
67        Self {
68            cwd: options.cwd,
69            parameters: write_parameters_schema(),
70            description: "Write content to a file. Creates the file if it doesn't exist, overwrites if it does. Automatically creates parent directories.".to_owned(),
71        }
72    }
73
74    /// Returns the JSON Schema for write arguments (normalized `TypeBox` shape).
75    #[must_use]
76    pub fn parameters_schema() -> Value {
77        write_parameters_schema()
78    }
79
80    /// Validates raw tool arguments into [`WriteToolInput`].
81    ///
82    /// # Errors
83    ///
84    /// Returns [`ToolError`] when required fields are missing or mistyped.
85    pub fn parse_input(args: &Map<String, Value>) -> Result<WriteToolInput, ToolError> {
86        serde_json::from_value(Value::Object(args.clone()))
87            .map_err(|error| ToolError::new(format!("Write tool input is invalid. {error}")))
88    }
89
90    /// JS `String.prototype.length` (UTF-16 code units).
91    #[must_use]
92    pub fn utf16_length(content: &str) -> usize {
93        content.encode_utf16().count()
94    }
95
96    /// Formats the success text using the caller's path string and JS length.
97    #[must_use]
98    pub fn success_text(path: &str, content: &str) -> String {
99        format!(
100            "Successfully wrote {} bytes to {path}",
101            Self::utf16_length(content)
102        )
103    }
104}
105
106impl AgentTool for WriteTool {
107    fn name(&self) -> &'static str {
108        "write"
109    }
110
111    fn label(&self) -> &'static str {
112        "write"
113    }
114
115    fn description(&self) -> &str {
116        &self.description
117    }
118
119    fn parameters(&self) -> &Value {
120        &self.parameters
121    }
122
123    fn validate_arguments(
124        &self,
125        args: &Map<String, Value>,
126    ) -> Result<Map<String, Value>, ToolError> {
127        let _ = Self::parse_input(args)?;
128        Ok(args.clone())
129    }
130
131    fn execute(
132        &self,
133        _tool_call_id: &str,
134        args: Map<String, Value>,
135        cancel: CancellationToken,
136        _updates: ToolUpdates,
137    ) -> BoxFuture<'static, Result<AgentToolResult, ToolError>> {
138        let cwd = self.cwd.clone();
139        async move {
140            throw_if_cancelled(&cancel)?;
141            let input = WriteTool::parse_input(&args)?;
142            let absolute_path = resolve_to_cwd(&input.path, cwd.to_string_lossy().as_ref())
143                .map_err(|error| path_error(&error))?;
144            let absolute = PathBuf::from(&absolute_path);
145            let parent = absolute
146                .parent()
147                .map_or_else(|| PathBuf::from("."), Path::to_path_buf);
148            let path_for_message = input.path.clone();
149            let content = input.content;
150            let cancel_for_queue = cancel.clone();
151
152            with_file_mutation_queue(&absolute, || {
153                let parent = parent.clone();
154                let absolute = absolute.clone();
155                let content = content.clone();
156                let cancel = cancel_for_queue.clone();
157                async move {
158                    apply_write_mutation_with_commit_hooks(
159                        &parent,
160                        &absolute,
161                        &path_for_message,
162                        &content,
163                        &cancel,
164                        || {},
165                        || {},
166                    )
167                    .await
168                }
169            })
170            .await
171            .map_err(|error| mutation_error(&error))?
172        }
173        .boxed()
174    }
175}
176
177async fn apply_write_mutation_with_commit_hooks<BeforeCommit, AfterCommit>(
178    parent: &Path,
179    absolute: &Path,
180    path_for_message: &str,
181    content: &str,
182    cancel: &CancellationToken,
183    before_commit: BeforeCommit,
184    after_commit: AfterCommit,
185) -> Result<AgentToolResult, ToolError>
186where
187    BeforeCommit: FnOnce() + Send,
188    AfterCommit: FnOnce() + Send,
189{
190    throw_if_cancelled(cancel)?;
191
192    tokio::fs::create_dir_all(parent).await.map_err(|error| {
193        ToolError::new(format!(
194            "Could not create parent directories for {}: {error}",
195            absolute.display()
196        ))
197    })?;
198
199    let result = AgentToolResult {
200        content: vec![ToolResultContent::Text(TextContent::new(
201            WriteTool::success_text(path_for_message, content),
202        ))],
203        details: Value::Null,
204        added_tool_names: None,
205        terminate: None,
206    };
207
208    before_commit();
209    throw_if_cancelled(cancel)?;
210    tokio::fs::write(absolute, content.as_bytes())
211        .await
212        .map_err(|error| {
213            ToolError::new(format!(
214                "Could not write file {}: {error}",
215                absolute.display()
216            ))
217        })?;
218    after_commit();
219
220    // A successful write is the durable commit point. Cancellation observed
221    // after it must not turn the committed write into a reported failure.
222    Ok(result)
223}
224
225fn write_parameters_schema() -> Value {
226    normalize_tool_schema(schemars::schema_for!(WriteToolInput))
227}
228
229fn normalize_tool_schema(schema: schemars::Schema) -> Value {
230    let mut value = serde_json::to_value(schema).unwrap_or_else(|_| Value::Object(Map::new()));
231    if let Value::Object(map) = &mut value {
232        map.remove("$schema");
233        map.remove("title");
234        map.remove("description");
235        map.remove("additionalProperties");
236        normalize_schema_node(map);
237    }
238    value
239}
240
241fn normalize_schema_node(map: &mut Map<String, Value>) {
242    map.remove("format");
243    if let Some(Value::Array(types)) = map.get("type").cloned() {
244        let non_null: Vec<Value> = types
245            .into_iter()
246            .filter(|t| t.as_str() != Some("null"))
247            .collect();
248        if non_null.len() == 1 {
249            map.insert("type".to_owned(), non_null[0].clone());
250        } else if !non_null.is_empty() {
251            map.insert("type".to_owned(), Value::Array(non_null));
252        }
253    }
254    let keys: Vec<String> = map.keys().cloned().collect();
255    for key in keys {
256        match map.get_mut(&key) {
257            Some(Value::Object(child)) => normalize_schema_node(child),
258            Some(Value::Array(items)) => {
259                for item in items {
260                    if let Value::Object(child) = item {
261                        normalize_schema_node(child);
262                    }
263                }
264            }
265            _ => {}
266        }
267    }
268}
269
270fn throw_if_cancelled(cancel: &CancellationToken) -> Result<(), ToolError> {
271    if cancel.is_cancelled() {
272        Err(ToolError::new("Operation aborted"))
273    } else {
274        Ok(())
275    }
276}
277
278fn path_error(error: &PathResolveError) -> ToolError {
279    ToolError::new(error.to_string())
280}
281
282fn mutation_error(error: &MutationQueueError) -> ToolError {
283    ToolError::new(error.to_string())
284}
285
286/// Builds an [`Arc<dyn AgentTool>`] write tool for `cwd`.
287#[must_use]
288pub fn create_write_tool(cwd: impl Into<PathBuf>) -> Arc<dyn AgentTool> {
289    Arc::new(WriteTool::new(cwd))
290}
291
292#[cfg(test)]
293mod tests {
294    use super::*;
295    use std::sync::Arc;
296    use std::time::Duration;
297
298    use serde_json::json;
299    use tempfile::tempdir;
300    use tokio::sync::Barrier;
301
302    fn fixture_schema() -> Result<Value, serde_json::Error> {
303        let text = include_str!("../../../tests/fixtures/tool-schemas/write.json");
304        serde_json::from_str(text)
305    }
306
307    #[test]
308    fn schema_matches_typebox_fixture() -> Result<(), Box<dyn std::error::Error>> {
309        let schema = WriteTool::parameters_schema();
310        assert_eq!(schema, fixture_schema()?);
311        Ok(())
312    }
313
314    #[test]
315    fn utf16_length_counts_surrogate_pairs() {
316        // "😀" is one scalar, two UTF-16 code units; JS length is 2.
317        assert_eq!(WriteTool::utf16_length("😀"), 2);
318        assert_eq!(WriteTool::utf16_length("abc"), 3);
319        assert_eq!(
320            WriteTool::success_text("a.txt", "😀"),
321            "Successfully wrote 2 bytes to a.txt"
322        );
323    }
324
325    #[tokio::test]
326    async fn writes_create_parent_and_overwrite() -> Result<(), Box<dyn std::error::Error>> {
327        let dir = tempdir()?;
328        let tool = WriteTool::new(dir.path());
329        let nested = "sub/dir/file.txt";
330        let args = json_map(json!({
331            "path": nested,
332            "content": "hello"
333        }))?;
334        let result = tool
335            .execute("1", args, CancellationToken::new(), ToolUpdates::noop())
336            .await?;
337        let text = text_of(&result);
338        assert_eq!(text, "Successfully wrote 5 bytes to sub/dir/file.txt");
339        let path = dir.path().join(nested);
340        assert_eq!(tokio::fs::read_to_string(&path).await?, "hello");
341
342        let args = json_map(json!({
343            "path": nested,
344            "content": "overwrite"
345        }))?;
346        let result = tool
347            .execute("2", args, CancellationToken::new(), ToolUpdates::noop())
348            .await?;
349        assert_eq!(
350            text_of(&result),
351            "Successfully wrote 9 bytes to sub/dir/file.txt"
352        );
353        assert_eq!(tokio::fs::read_to_string(&path).await?, "overwrite");
354        Ok(())
355    }
356
357    #[tokio::test]
358    async fn cancellation_before_work_aborts() -> Result<(), Box<dyn std::error::Error>> {
359        let dir = tempdir()?;
360        let tool = WriteTool::new(dir.path());
361        let cancel = CancellationToken::new();
362        cancel.cancel();
363        let result = tool
364            .execute(
365                "1",
366                json_map(json!({"path": "a.txt", "content": "x"}))?,
367                cancel,
368                ToolUpdates::noop(),
369            )
370            .await;
371        let Err(err) = result else {
372            return Err("expected cancellation error".into());
373        };
374        assert_eq!(err.message(), "Operation aborted");
375        Ok(())
376    }
377
378    #[tokio::test]
379    async fn cancellation_before_write_commit_aborts_without_mutating()
380    -> Result<(), Box<dyn std::error::Error>> {
381        let dir = tempdir()?;
382        let path = dir.path().join("pre-commit.txt");
383        tokio::fs::write(&path, "before").await?;
384        let cancel = CancellationToken::new();
385        let cancel_at_boundary = cancel.clone();
386
387        let result = apply_write_mutation_with_commit_hooks(
388            dir.path(),
389            &path,
390            "pre-commit.txt",
391            "after",
392            &cancel,
393            move || cancel_at_boundary.cancel(),
394            || {},
395        )
396        .await;
397        let Err(error) = result else {
398            return Err("pre-commit cancellation unexpectedly succeeded".into());
399        };
400
401        assert_eq!(error.message(), "Operation aborted");
402        assert_eq!(tokio::fs::read_to_string(&path).await?, "before");
403        Ok(())
404    }
405
406    #[tokio::test]
407    async fn cancellation_after_write_commit_reports_success()
408    -> Result<(), Box<dyn std::error::Error>> {
409        let dir = tempdir()?;
410        let path = dir.path().join("post-commit.txt");
411        let cancel = CancellationToken::new();
412        let cancel_at_boundary = cancel.clone();
413
414        let result = apply_write_mutation_with_commit_hooks(
415            dir.path(),
416            &path,
417            "post-commit.txt",
418            "committed",
419            &cancel,
420            || {},
421            move || cancel_at_boundary.cancel(),
422        )
423        .await?;
424
425        assert!(cancel.is_cancelled());
426        assert_eq!(
427            text_of(&result),
428            "Successfully wrote 9 bytes to post-commit.txt"
429        );
430        assert_eq!(tokio::fs::read_to_string(&path).await?, "committed");
431        Ok(())
432    }
433
434    #[tokio::test]
435    async fn serialized_writes_same_path() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
436        let dir = tempdir()?;
437        let path = dir.path().join("serial.txt");
438        let barrier = Arc::new(Barrier::new(2));
439        let order = Arc::new(std::sync::Mutex::new(Vec::new()));
440
441        let start_a = barrier.clone();
442        let order_a = order.clone();
443        let path_a = path.clone();
444        let a = tokio::spawn(async move {
445            with_file_mutation_queue(&path_a, || {
446                let start_a = start_a.clone();
447                let order_a = order_a.clone();
448                let path_a = path_a.clone();
449                async move {
450                    start_a.wait().await;
451                    order_a
452                        .lock()
453                        .map_err(|_| std::io::Error::other("lock poisoned"))?
454                        .push("a-start");
455                    tokio::time::sleep(Duration::from_millis(40)).await;
456                    tokio::fs::write(&path_a, b"a").await?;
457                    order_a
458                        .lock()
459                        .map_err(|_| std::io::Error::other("lock poisoned"))?
460                        .push("a-end");
461                    Ok::<(), std::io::Error>(())
462                }
463            })
464            .await
465            .map_err(|error| std::io::Error::other(error.to_string()))?
466        });
467
468        let start_b = barrier;
469        let order_b = order.clone();
470        let path_b = path.clone();
471        let b = tokio::spawn(async move {
472            // Ensure both tasks race to acquire.
473            start_b.wait().await;
474            with_file_mutation_queue(&path_b, || {
475                let order_b = order_b.clone();
476                let path_b = path_b.clone();
477                async move {
478                    order_b
479                        .lock()
480                        .map_err(|_| std::io::Error::other("lock poisoned"))?
481                        .push("b-start");
482                    tokio::fs::write(&path_b, b"b").await?;
483                    order_b
484                        .lock()
485                        .map_err(|_| std::io::Error::other("lock poisoned"))?
486                        .push("b-end");
487                    Ok::<(), std::io::Error>(())
488                }
489            })
490            .await
491            .map_err(|error| std::io::Error::other(error.to_string()))?
492        });
493
494        a.await??;
495        b.await??;
496        let observed = order
497            .lock()
498            .map_err(|_| std::io::Error::other("lock poisoned"))?
499            .clone();
500        // FIFO: a fully finishes before b starts, or b fully finishes before a
501        // starts — never interleaved start/end across the critical section.
502        assert!(
503            observed == ["a-start", "a-end", "b-start", "b-end"]
504                || observed == ["b-start", "b-end", "a-start", "a-end"],
505            "unexpected order: {observed:?}"
506        );
507        Ok(())
508    }
509
510    fn json_map(value: Value) -> Result<Map<String, Value>, &'static str> {
511        if let Value::Object(map) = value {
512            Ok(map)
513        } else {
514            Err("expected JSON object")
515        }
516    }
517
518    fn text_of(result: &AgentToolResult) -> String {
519        match result.content.first() {
520            Some(ToolResultContent::Text(text)) => text.text.to_string(),
521            _ => String::new(),
522        }
523    }
524}
pi/core/tools/write.rs

pi/core/tools/
write.rs
