ai-harness/crates/harness-tools/src/edit/mod.rs

mod replacers;

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex as StdMutex};

use async_trait::async_trait;
use harness_core::tool::{Tool, ToolCtx, ToolError, ToolOutput};
use schemars::JsonSchema;
use serde::Deserialize;
use similar::{ChangeTag, TextDiff};
use tokio::sync::Mutex as AsyncMutex;

use crate::paths;

#[derive(Debug, Deserialize, JsonSchema)]
struct EditParams {
    file_path: String,
    old_string: String,
    new_string: String,
    replace_all: Option<bool>,
}

/// Per-file `tokio::Mutex` lock map so concurrent edits to different files proceed in
/// parallel while edits to the *same* file serialize (docs/05-tools.md).
pub struct EditTool {
    locks: StdMutex<HashMap<PathBuf, Arc<AsyncMutex<()>>>>,
}

impl EditTool {
    pub fn new() -> Self {
        Self {
            locks: StdMutex::new(HashMap::new()),
        }
    }

    fn lock_for(&self, path: &Path) -> Arc<AsyncMutex<()>> {
        self.locks
            .lock()
            .unwrap()
            .entry(path.to_path_buf())
            .or_insert_with(|| Arc::new(AsyncMutex::new(())))
            .clone()
    }
}

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

fn normalize_line_endings(text: &str) -> String {
    text.replace("\r\n", "\n")
}

fn detect_line_ending(text: &str) -> &'static str {
    if text.contains("\r\n") {
        "\r\n"
    } else {
        "\n"
    }
}

fn convert_line_ending(text: &str, ending: &str) -> String {
    if ending == "\n" {
        text.to_string()
    } else {
        text.replace('\n', "\r\n")
    }
}

const UTF8_BOM: [u8; 3] = [0xEF, 0xBB, 0xBF];

fn strip_bom(bytes: &[u8]) -> (bool, &[u8]) {
    if bytes.starts_with(&UTF8_BOM) {
        (true, &bytes[3..])
    } else {
        (false, bytes)
    }
}

fn unified_diff(before: &str, after: &str, file_path: &str) -> String {
    TextDiff::from_lines(before, after)
        .unified_diff()
        .header(&format!("a/{file_path}"), &format!("b/{file_path}"))
        .to_string()
}

fn diff_stats(before: &str, after: &str) -> (usize, usize) {
    let diff = TextDiff::from_lines(before, after);
    let mut added = 0;
    let mut removed = 0;
    for change in diff.iter_all_changes() {
        match change.tag() {
            ChangeTag::Insert => added += 1,
            ChangeTag::Delete => removed += 1,
            ChangeTag::Equal => {}
        }
    }
    (added, removed)
}

#[async_trait]
impl Tool for EditTool {
    fn name(&self) -> &str {
        "edit"
    }

    fn description(&self) -> &str {
        "Replaces an exact span of text in a file (fuzzy-matched against whitespace/indentation drift)."
    }

    fn parameters(&self) -> serde_json::Value {
        serde_json::to_value(schemars::schema_for!(EditParams)).unwrap()
    }

    async fn execute(
        &self,
        input: serde_json::Value,
        ctx: ToolCtx,
    ) -> Result<ToolOutput, ToolError> {
        let params: EditParams =
            serde_json::from_value(input).map_err(|e| ToolError::Invalid(e.to_string()))?;

        if params.old_string == params.new_string {
            return Err(ToolError::Invalid(
                "No changes to apply: oldString and newString are identical.".into(),
            ));
        }

        let path = paths::resolve(&ctx.cwd, &params.file_path);
        let pattern = paths::relative_pattern(&ctx.cwd, &path);
        let lock = self.lock_for(&path);
        let _guard = lock.lock().await;

        let metadata = tokio::fs::metadata(&path).await.ok();
        if metadata.as_ref().is_some_and(|m| m.is_dir()) {
            return Err(ToolError::Invalid(format!(
                "Path is a directory, not a file: {}",
                path.display()
            )));
        }

        let (content_old, content_new, had_bom) = if params.old_string.is_empty() {
            if metadata.is_some() {
                return Err(ToolError::Invalid(
                    "oldString cannot be empty when editing an existing file. Provide the exact text to replace, or use write for an intentional full-file replacement.".into(),
                ));
            }
            (String::new(), params.new_string.clone(), false)
        } else {
            if metadata.is_none() {
                return Err(ToolError::Invalid(format!(
                    "File {} not found",
                    path.display()
                )));
            }
            let bytes = tokio::fs::read(&path)
                .await
                .map_err(|e| ToolError::Other(format!("{}: {e}", path.display())))?;
            let (had_bom, stripped) = strip_bom(&bytes);
            let text = String::from_utf8_lossy(stripped).into_owned();

            let ending = detect_line_ending(&text);
            let old = convert_line_ending(&normalize_line_endings(&params.old_string), ending);
            let new_ = convert_line_ending(&normalize_line_endings(&params.new_string), ending);
            let replaced =
                replacers::replace(&text, &old, &new_, params.replace_all.unwrap_or(false))
                    .map_err(|e| ToolError::Invalid(e.to_string()))?;
            (text, replaced, had_bom)
        };

        let diff = unified_diff(&content_old, &content_new, &params.file_path);
        ctx.ask
            .ask(
                "edit",
                pattern.clone(),
                "*",
                serde_json::json!({"file_path": params.file_path, "diff": diff}),
            )
            .await?;

        if let Some(parent) = path.parent() {
            tokio::fs::create_dir_all(parent)
                .await
                .map_err(|e| ToolError::Other(format!("{}: {e}", parent.display())))?;
        }
        let mut out_bytes = Vec::with_capacity(content_new.len() + 3);
        if had_bom {
            out_bytes.extend_from_slice(&UTF8_BOM);
        }
        out_bytes.extend_from_slice(content_new.as_bytes());
        tokio::fs::write(&path, &out_bytes)
            .await
            .map_err(|e| ToolError::Other(format!("{}: {e}", path.display())))?;

        let (added, removed) = diff_stats(&content_old, &content_new);
        let mut output = ToolOutput::new(pattern, "Edit applied successfully.".to_string());
        output.metadata = serde_json::json!({"diff": diff, "added": added, "removed": removed});
        Ok(output)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use harness_core::event::EventBus;
    use harness_core::permission::{spawn_auto_approve, PermissionService};
    use harness_core::tool::{MetadataSink, PermissionHandle};
    use harness_core::types::SessionId;
    use std::sync::Mutex;
    use tokio_util::sync::CancellationToken;

    fn ctx(cwd: PathBuf) -> ToolCtx {
        let bus = EventBus::new();
        let service = Arc::new(PermissionService::new(bus.clone()));
        spawn_auto_approve(bus, service.clone());
        let (metadata, _rx) = MetadataSink::channel();
        ToolCtx {
            session_id: SessionId::new(),
            message_id: harness_core::types::MessageId::new(),
            call_id: "call_1".into(),
            data_dir: cwd.join("tool-output"),
            cwd,
            cancel: CancellationToken::new(),
            ask: PermissionHandle::new(
                service,
                SessionId::new(),
                Vec::new(),
                Arc::new(Mutex::new(Vec::new())),
                CancellationToken::new(),
            ),
            metadata,
        }
    }

    #[tokio::test]
    async fn creates_new_file_when_old_string_is_empty() {
        let dir = tempfile::tempdir().unwrap();
        let result = EditTool::new()
            .execute(
                serde_json::json!({"file_path": "newfile.txt", "old_string": "", "new_string": "new content"}),
                ctx(dir.path().to_path_buf()),
            )
            .await
            .unwrap();
        assert!(result.metadata["diff"]
            .as_str()
            .unwrap()
            .contains("new content"));
        assert_eq!(
            std::fs::read_to_string(dir.path().join("newfile.txt")).unwrap(),
            "new content"
        );
    }

    #[tokio::test]
    async fn rejects_empty_old_string_on_existing_file() {
        let dir = tempfile::tempdir().unwrap();
        std::fs::write(dir.path().join("existing.txt"), "original").unwrap();
        let err = EditTool::new()
            .execute(
                serde_json::json!({"file_path": "existing.txt", "old_string": "", "new_string": "x"}),
                ctx(dir.path().to_path_buf()),
            )
            .await
            .unwrap_err();
        assert!(
            matches!(&err, ToolError::Invalid(msg) if msg.contains("oldString cannot be empty"))
        );
        assert_eq!(
            std::fs::read_to_string(dir.path().join("existing.txt")).unwrap(),
            "original"
        );
    }

    #[tokio::test]
    async fn creates_new_file_with_nested_directories() {
        let dir = tempfile::tempdir().unwrap();
        EditTool::new()
            .execute(
                serde_json::json!({"file_path": "nested/dir/file.txt", "old_string": "", "new_string": "nested file"}),
                ctx(dir.path().to_path_buf()),
            )
            .await
            .unwrap();
        assert_eq!(
            std::fs::read_to_string(dir.path().join("nested/dir/file.txt")).unwrap(),
            "nested file"
        );
    }

    #[tokio::test]
    async fn replaces_text_in_existing_file() {
        let dir = tempfile::tempdir().unwrap();
        std::fs::write(dir.path().join("existing.txt"), "old content here").unwrap();
        let result = EditTool::new()
            .execute(
                serde_json::json!({"file_path": "existing.txt", "old_string": "old content", "new_string": "new content"}),
                ctx(dir.path().to_path_buf()),
            )
            .await
            .unwrap();
        assert!(result.output.contains("Edit applied successfully"));
        assert_eq!(
            std::fs::read_to_string(dir.path().join("existing.txt")).unwrap(),
            "new content here"
        );
    }

    #[tokio::test]
    async fn preserves_bom_and_only_edits_visible_content() {
        let dir = tempfile::tempdir().unwrap();
        let mut original = UTF8_BOM.to_vec();
        original.extend_from_slice(b"using System;\nclass Test {}\n");
        std::fs::write(dir.path().join("existing.cs"), &original).unwrap();

        let result = EditTool::new()
            .execute(
                serde_json::json!({"file_path": "existing.cs", "old_string": "using System;", "new_string": "using Up;"}),
                ctx(dir.path().to_path_buf()),
            )
            .await
            .unwrap();
        assert!(result.metadata["diff"]
            .as_str()
            .unwrap()
            .contains("-using System;"));
        assert!(result.metadata["diff"]
            .as_str()
            .unwrap()
            .contains("+using Up;"));

        let content = std::fs::read(dir.path().join("existing.cs")).unwrap();
        assert_eq!(&content[..3], &UTF8_BOM);
        assert_eq!(&content[3..], b"using Up;\nclass Test {}\n");
    }

    #[tokio::test]
    async fn errors_when_file_does_not_exist() {
        let dir = tempfile::tempdir().unwrap();
        let err = EditTool::new()
            .execute(
                serde_json::json!({"file_path": "nonexistent.txt", "old_string": "old", "new_string": "new"}),
                ctx(dir.path().to_path_buf()),
            )
            .await
            .unwrap_err();
        assert!(matches!(&err, ToolError::Invalid(msg) if msg.contains("not found")));
    }

    #[tokio::test]
    async fn errors_when_old_equals_new() {
        let dir = tempfile::tempdir().unwrap();
        std::fs::write(dir.path().join("file.txt"), "content").unwrap();
        let err = EditTool::new()
            .execute(
                serde_json::json!({"file_path": "file.txt", "old_string": "same", "new_string": "same"}),
                ctx(dir.path().to_path_buf()),
            )
            .await
            .unwrap_err();
        assert!(matches!(&err, ToolError::Invalid(msg) if msg.contains("identical")));
    }

    #[tokio::test]
    async fn errors_when_path_is_a_directory() {
        let dir = tempfile::tempdir().unwrap();
        std::fs::create_dir(dir.path().join("adir")).unwrap();
        let err = EditTool::new()
            .execute(
                serde_json::json!({"file_path": "adir", "old_string": "old", "new_string": "new"}),
                ctx(dir.path().to_path_buf()),
            )
            .await
            .unwrap_err();
        assert!(matches!(&err, ToolError::Invalid(msg) if msg.contains("directory")));
    }

    #[tokio::test]
    async fn replaces_all_occurrences_with_replace_all_option() {
        let dir = tempfile::tempdir().unwrap();
        std::fs::write(dir.path().join("file.txt"), "foo bar foo baz foo").unwrap();
        EditTool::new()
            .execute(
                serde_json::json!({"file_path": "file.txt", "old_string": "foo", "new_string": "qux", "replace_all": true}),
                ctx(dir.path().to_path_buf()),
            )
            .await
            .unwrap();
        assert_eq!(
            std::fs::read_to_string(dir.path().join("file.txt")).unwrap(),
            "qux bar qux baz qux"
        );
    }

    #[tokio::test]
    async fn handles_crlf_line_endings() {
        let dir = tempfile::tempdir().unwrap();
        std::fs::write(dir.path().join("file.txt"), "line1\r\nold\r\nline3").unwrap();
        EditTool::new()
            .execute(
                serde_json::json!({"file_path": "file.txt", "old_string": "old", "new_string": "new"}),
                ctx(dir.path().to_path_buf()),
            )
            .await
            .unwrap();
        assert_eq!(
            std::fs::read_to_string(dir.path().join("file.txt")).unwrap(),
            "line1\r\nnew\r\nline3"
        );
    }

    #[tokio::test]
    async fn tracks_file_diff_statistics() {
        let dir = tempfile::tempdir().unwrap();
        std::fs::write(dir.path().join("file.txt"), "line1\nline2\nline3").unwrap();
        let result = EditTool::new()
            .execute(
                serde_json::json!({"file_path": "file.txt", "old_string": "line2", "new_string": "new line a\nnew line b"}),
                ctx(dir.path().to_path_buf()),
            )
            .await
            .unwrap();
        assert!(result.metadata["added"].as_u64().unwrap() > 0);
    }

    #[tokio::test]
    async fn concurrent_edits_to_the_same_file_serialize_without_losing_either_change() {
        let dir = tempfile::tempdir().unwrap();
        std::fs::write(
            dir.path().join("file.txt"),
            "top = 0\nmiddle = keep\nbottom = 0\n",
        )
        .unwrap();
        let tool = Arc::new(EditTool::new());

        let t1 = tool.clone();
        let cwd1 = dir.path().to_path_buf();
        let h1 = tokio::spawn(async move {
            t1.execute(
                serde_json::json!({"file_path": "file.txt", "old_string": "top = 0", "new_string": "top = 1"}),
                ctx(cwd1),
            )
            .await
        });

        let t2 = tool.clone();
        let cwd2 = dir.path().to_path_buf();
        let h2 = tokio::spawn(async move {
            t2.execute(
                serde_json::json!({"file_path": "file.txt", "old_string": "bottom = 0", "new_string": "bottom = 2"}),
                ctx(cwd2),
            )
            .await
        });

        h1.await.unwrap().unwrap();
        h2.await.unwrap().unwrap();

        assert_eq!(
            std::fs::read_to_string(dir.path().join("file.txt")).unwrap(),
            "top = 1\nmiddle = keep\nbottom = 2\n"
        );
    }
}