swink-agent 0.13.2

//! Built-in tool for making surgical find-and-replace edits to a file.

use std::io::Write as _;
use std::ops::Range;
use std::path::{Path, PathBuf};

use schemars::JsonSchema;
use serde::Deserialize;
use serde_json::Value;
use sha2::{Digest as _, Sha256};
use tokio_util::sync::CancellationToken;

use super::path::resolve_existing_path;
use crate::tool::{AgentTool, AgentToolResult, ToolFuture, validated_schema_for};
use crate::types::ContentBlock;

/// Built-in tool for making precise, surgical edits to a file.
///
/// Supports multiple edits per call, atomic writes, stale-read detection,
/// whitespace-normalised matching, and line-number-based disambiguation.
pub struct EditFileTool {
    schema: Value,
    execution_root: Option<PathBuf>,
}

impl EditFileTool {
    /// Create a new `EditFileTool`.
    #[must_use]
    pub fn new() -> Self {
        Self {
            schema: validated_schema_for::<Params>(),
            execution_root: None,
        }
    }

    /// Set the working directory used to resolve relative file paths.
    #[must_use]
    pub fn with_execution_root(mut self, root: impl Into<PathBuf>) -> Self {
        self.execution_root = Some(root.into());
        self
    }
}

impl Default for EditFileTool {
    fn default() -> Self {
        Self::new()
    }
}

/// One find-and-replace operation.
#[derive(Deserialize, JsonSchema)]
#[schemars(deny_unknown_fields)]
struct EditOp {
    /// Text to find in the file.  Exact match is tried first; if that fails a
    /// line-by-line match that ignores trailing whitespace is attempted.
    old_string: String,
    /// Replacement text.
    new_string: String,
    /// When `true`, every occurrence is replaced.  When `false` (the default)
    /// exactly one occurrence must exist, or `line_hint` must be provided.
    #[serde(default)]
    replace_all: bool,
    /// 1-based line number of the desired occurrence.  Used to pick among
    /// multiple matches when `replace_all` is `false`.
    line_hint: Option<u32>,
}

#[derive(Deserialize, JsonSchema)]
#[schemars(deny_unknown_fields)]
struct Params {
    /// Absolute path to the file to edit.
    path: String,
    /// Edits to apply in order (top-to-bottom).
    edits: Vec<EditOp>,
    /// SHA-256 hex digest of the file content as previously read.  When
    /// provided the edit is rejected if the file has changed since.
    expected_hash: Option<String>,
}

// ---------------------------------------------------------------------------
// Matching helpers
// ---------------------------------------------------------------------------

/// Compute the SHA-256 hex digest of `data`.
fn sha256_hex(data: &[u8]) -> String {
    Sha256::digest(data)
        .iter()
        .fold(String::with_capacity(64), |mut s, b| {
            use std::fmt::Write as _;
            let _ = write!(s, "{b:02x}");
            s
        })
}

/// Return `(byte_start, line_content_without_newline)` for every line.
///
/// Splits on `'\n'`; the `'\n'` itself is not included in the line slice.
/// Windows `\r\n` files: the `\r` will appear as trailing content in each
/// slice, which is stripped by [`str::trim_end`] during normalised matching.
fn line_spans(s: &str) -> Vec<(usize, &str)> {
    let mut spans = Vec::new();
    let mut pos = 0;
    for line in s.split('\n') {
        spans.push((pos, line));
        pos += line.len() + 1; // +1 for the '\n'
    }
    spans
}

/// Find all non-overlapping exact byte ranges of `pattern` in `content`.
fn find_exact(content: &str, pattern: &str) -> Vec<Range<usize>> {
    if pattern.is_empty() {
        return Vec::new();
    }
    let mut ranges = Vec::new();
    let mut start = 0;
    while let Some(pos) = content[start..].find(pattern) {
        let abs = start + pos;
        ranges.push(abs..abs + pattern.len());
        start = abs + pattern.len();
    }
    ranges
}

/// Find all non-overlapping byte ranges in `content` that match `pattern`
/// line-by-line, ignoring trailing whitespace on each line.
///
/// Leading and trailing blank lines in `pattern` are stripped before
/// matching.  The returned ranges refer to byte positions in the original
/// (un-normalised) `content`.
fn find_normalized(content: &str, pattern: &str) -> Vec<Range<usize>> {
    let pattern = pattern.trim_matches('\n');
    if pattern.is_empty() {
        return Vec::new();
    }
    let pattern_lines: Vec<&str> = pattern.split('\n').collect();
    let spans = line_spans(content);
    let n = pattern_lines.len();

    if n > spans.len() {
        return Vec::new();
    }

    let mut ranges = Vec::new();
    let mut i = 0;
    while i + n <= spans.len() {
        let all_match = pattern_lines
            .iter()
            .enumerate()
            .all(|(j, &pl)| spans[i + j].1.trim_end() == pl.trim_end());

        if all_match {
            let byte_start = spans[i].0;
            let last = &spans[i + n - 1];
            let byte_end = last.0 + last.1.len();
            ranges.push(byte_start..byte_end);
            i += n; // skip past the match so occurrences don't overlap
        } else {
            i += 1;
        }
    }
    ranges
}

/// Return the 1-based line number of the character at `byte_pos`.
fn line_number_at(content: &str, byte_pos: usize) -> usize {
    content[..byte_pos].chars().filter(|&c| c == '\n').count() + 1
}

/// Replace all `ranges` in `content` with `replacement`.
///
/// `ranges` must be sorted ascending and non-overlapping.
fn replace_ranges(content: &str, ranges: &[Range<usize>], replacement: &str) -> String {
    let mut out = String::with_capacity(content.len());
    let mut cursor = 0;
    for r in ranges {
        out.push_str(&content[cursor..r.start]);
        out.push_str(replacement);
        cursor = r.end;
    }
    out.push_str(&content[cursor..]);
    out
}

/// Apply a single [`EditOp`] to `content`, returning the modified string or
/// an error message.
fn apply_op(content: &str, op: &EditOp) -> Result<String, String> {
    if op.old_string.is_empty() {
        return Err("old_string must not be empty".to_owned());
    }

    // Prefer exact match; fall back to whitespace-normalised line matching.
    let candidates: Vec<Range<usize>> = {
        let exact = find_exact(content, &op.old_string);
        if exact.is_empty() {
            let norm = find_normalized(content, &op.old_string);
            if norm.is_empty() {
                return Err(format!(
                    "old_string not found (tried exact and whitespace-normalised match):\n{}",
                    op.old_string
                ));
            }
            norm
        } else {
            exact
        }
    };

    if op.replace_all {
        return Ok(replace_ranges(content, &candidates, &op.new_string));
    }

    match candidates.len() {
        0 => unreachable!("candidates is non-empty at this point"),
        1 => Ok(replace_ranges(content, &candidates, &op.new_string)),
        n => op.line_hint.map_or_else(
            || {
                Err(format!(
                    "old_string matched {n} times; set replace_all to replace every \
                     occurrence, or provide line_hint to select one"
                ))
            },
            |hint| {
                let best = candidates
                    .iter()
                    .min_by_key(|r| {
                        let line =
                            i64::try_from(line_number_at(content, r.start)).unwrap_or(i64::MAX);
                        (line - i64::from(hint)).abs()
                    })
                    .expect("candidates is non-empty");
                Ok(replace_ranges(
                    content,
                    std::slice::from_ref(best),
                    &op.new_string,
                ))
            },
        ),
    }
}

// ---------------------------------------------------------------------------
// Locked read-modify-write
// ---------------------------------------------------------------------------

/// Apply edits and write the result while holding the shared per-target lock.
fn edit_file_locked(
    path: &Path,
    expected_hash: Option<&str>,
    edits: &[EditOp],
) -> Result<(String, String), String> {
    crate::atomic_fs::with_target_lock(path, || Ok(edit_file_unlocked(path, expected_hash, edits)))
        .map_err(|error| format!("failed to lock {}: {error}", path.display()))?
}

fn edit_file_unlocked(
    path: &Path,
    expected_hash: Option<&str>,
    edits: &[EditOp],
) -> Result<(String, String), String> {
    let raw_bytes = std::fs::read(path)
        .map_err(|error| format!("failed to read {}: {error}", path.display()))?;

    let original = std::str::from_utf8(&raw_bytes)
        .map_err(|_| format!("{} is not valid UTF-8", path.display()))?
        .to_owned();

    if let Some(expected) = expected_hash {
        let actual = sha256_hex(&raw_bytes);
        if actual != expected.to_ascii_lowercase() {
            return Err(format!(
                "{} has changed since it was last read (hash mismatch); \
                 re-read the file before editing",
                path.display()
            ));
        }
    }

    let mut content = original.clone();
    for (i, op) in edits.iter().enumerate() {
        content = apply_op(&content, op).map_err(|msg| format!("edit {}: {msg}", i + 1))?;
    }

    crate::atomic_fs::atomic_write_unlocked(path, |writer| writer.write_all(content.as_bytes()))
        .map_err(|error| format!("failed to write {}: {error}", path.display()))?;

    Ok((original, content))
}

// ---------------------------------------------------------------------------
// AgentTool impl
// ---------------------------------------------------------------------------

#[allow(clippy::unnecessary_literal_bound)]
impl AgentTool for EditFileTool {
    fn name(&self) -> &str {
        "edit_file"
    }

    fn label(&self) -> &str {
        "Edit File"
    }

    fn description(&self) -> &str {
        "Apply one or more surgical find-and-replace edits to a file. \
         Edits are applied top-to-bottom. Trailing whitespace is ignored \
         during matching when an exact match is not found. The write is \
         atomic: the file is never left in a partially-written state."
    }

    fn parameters_schema(&self) -> &Value {
        &self.schema
    }

    fn requires_approval(&self) -> bool {
        true
    }

    fn execution_root(&self) -> Option<&Path> {
        self.execution_root.as_deref()
    }

    fn execute(
        &self,
        _tool_call_id: &str,
        params: Value,
        cancellation_token: CancellationToken,
        _on_update: Option<Box<dyn Fn(AgentToolResult) + Send + Sync>>,
        _state: std::sync::Arc<std::sync::RwLock<crate::SessionState>>,
        _credential: Option<crate::credential::ResolvedCredential>,
    ) -> ToolFuture<'_> {
        Box::pin(async move {
            let parsed: Params = match serde_json::from_value(params) {
                Ok(p) => p,
                Err(e) => return AgentToolResult::error(format!("invalid parameters: {e}")),
            };

            if cancellation_token.is_cancelled() {
                return AgentToolResult::error("cancelled");
            }

            let path =
                match resolve_existing_path(&parsed.path, self.execution_root.as_deref()).await {
                    Ok(path) => path,
                    Err(error) => return AgentToolResult::error(error),
                };

            if cancellation_token.is_cancelled() {
                return AgentToolResult::error("cancelled");
            }

            let path = if self.execution_root.is_none() {
                match tokio::fs::canonicalize(&path).await {
                    Ok(path) => path,
                    Err(error) => {
                        return AgentToolResult::error(format!(
                            "failed to resolve path {}: {error}",
                            path.display()
                        ));
                    }
                }
            } else {
                path
            };

            if parsed.edits.is_empty() {
                return AgentToolResult::text("no edits specified; file unchanged");
            }

            let n = parsed.edits.len();
            let expected_hash = parsed.expected_hash;
            let edits = parsed.edits;
            let edit_path = path.clone();
            let edit_result = tokio::task::spawn_blocking(move || {
                edit_file_locked(&edit_path, expected_hash.as_deref(), &edits)
            })
            .await;
            let (original, content) = match edit_result {
                Ok(Ok(result)) => result,
                Ok(Err(error)) => return AgentToolResult::error(error),
                Err(error) => {
                    return AgentToolResult::error(format!("edit task failed: {error}"));
                }
            };

            AgentToolResult {
                content: vec![ContentBlock::Text {
                    text: format!(
                        "Applied {} edit{} to {}",
                        n,
                        if n == 1 { "" } else { "s" },
                        path.display()
                    ),
                }],
                details: serde_json::json!({
                    "path": path,
                    "edits_applied": n,
                    "old_content": original,
                    "new_content": content,
                }),
                is_error: false,
                transfer_signal: None,
            }
        })
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
#[path = "edit_file_tests.rs"]
mod tests;