Skip to main content

imp_core/tools/
multi_edit.rs

1use async_trait::async_trait;
2use imp_llm::truncate_chars_with_suffix;
3use serde_json::json;
4
5use super::edit::apply_edit;
6use super::{generate_diff, suggest_similar_files, Tool, ToolContext, ToolOutput};
7use crate::error::Result;
8
9pub struct MultiEditTool;
10
11#[async_trait]
12impl Tool for MultiEditTool {
13    fn name(&self) -> &str {
14        "multi_edit"
15    }
16    fn label(&self) -> &str {
17        "Multi Edit"
18    }
19    fn description(&self) -> &str {
20        "Legacy compatibility shim for multi-edit transactions. Prefer the canonical edit tool with edits[]."
21    }
22    fn parameters(&self) -> serde_json::Value {
23        json!({
24            "type": "object",
25            "properties": {
26                "path": { "type": "string", "description": "Default path to edit; may be omitted when each edit includes its own path" },
27                "dry_run": { "type": "boolean", "description": "Validate and return combined diff without writing files" },
28                "edits": {
29                    "type": "array",
30                    "items": {
31                        "type": "object",
32                        "properties": {
33                            "path": { "type": "string", "description": "Optional per-edit path for multi-file transactions" },
34                            "old_text": { "type": "string" },
35                            "new_text": { "type": "string" }
36                        },
37                        "required": ["old_text", "new_text"]
38                    },
39                    "description": "Array of {old_text, new_text, path?} edits validated before any file is written"
40                }
41            },
42            "required": ["edits"]
43        })
44    }
45    fn is_readonly(&self) -> bool {
46        false
47    }
48
49    async fn execute(
50        &self,
51        _call_id: &str,
52        params: serde_json::Value,
53        ctx: ToolContext,
54    ) -> Result<ToolOutput> {
55        let raw_path = params["path"].as_str().unwrap_or("");
56        let dry_run = params
57            .get("dry_run")
58            .and_then(|v| v.as_bool())
59            .or_else(|| params.get("dryRun").and_then(|v| v.as_bool()))
60            .unwrap_or(false);
61        let edits = match params["edits"].as_array() {
62            Some(e) if !e.is_empty() => e,
63            _ => return Ok(ToolOutput::error("Missing or empty edits array")),
64        };
65
66        let mut edits_by_path: std::collections::BTreeMap<String, Vec<&serde_json::Value>> =
67            std::collections::BTreeMap::new();
68        for edit in edits {
69            let edit_path = edit["path"].as_str().unwrap_or(raw_path);
70            if edit_path.is_empty() {
71                return Ok(ToolOutput::error(
72                    "Missing required parameter: path (top-level path or per-edit path)",
73                ));
74            }
75            edits_by_path
76                .entry(edit_path.to_string())
77                .or_default()
78                .push(edit);
79        }
80
81        let mut prepared = Vec::new();
82        let mut any_fuzzy = false;
83        let mut total_edits = 0usize;
84        let mut warnings = Vec::new();
85
86        for (edit_path, file_edits) in edits_by_path {
87            let path = super::resolve_path(&ctx.cwd, &edit_path);
88            if let Err(error) = ctx.check_write_path(&path) {
89                return Ok(ToolOutput::error(error));
90            }
91            if !path.exists() {
92                let suggestions = suggest_similar_files(&ctx.cwd, &edit_path);
93                let mut msg = format!("File not found: {}", path.display());
94                if !suggestions.is_empty() {
95                    msg.push_str("\n\nDid you mean:");
96                    for s in &suggestions {
97                        msg.push_str(&format!("\n  {s}"));
98                    }
99                }
100                return Ok(ToolOutput::error(msg));
101            }
102
103            if let Some(warning) = tracker_warning(&ctx, &path) {
104                warnings.push(warning);
105            }
106
107            let raw_content = tokio::fs::read_to_string(&path).await?;
108            let original = raw_content.replace("\r\n", "\n");
109            let has_crlf = raw_content.contains("\r\n");
110
111            if let Err(error) = reject_overlapping_exact_edits(&edit_path, &original, &file_edits) {
112                return Ok(ToolOutput::error(error.to_string()));
113            }
114
115            let mut current = original.clone();
116            for (i, edit) in file_edits.iter().enumerate() {
117                let old_text = edit
118                    .get("old_text")
119                    .and_then(|v| v.as_str())
120                    .or_else(|| edit.get("oldText").and_then(|v| v.as_str()))
121                    .unwrap_or("")
122                    .replace("\r\n", "\n");
123                let new_text = edit
124                    .get("new_text")
125                    .and_then(|v| v.as_str())
126                    .or_else(|| edit.get("newText").and_then(|v| v.as_str()))
127                    .unwrap_or("")
128                    .replace("\r\n", "\n");
129                if old_text.is_empty() {
130                    return Ok(ToolOutput::error(format!(
131                        "Edit {} in {edit_path}: missing old_text",
132                        i + 1
133                    )));
134                }
135                match apply_edit(&current, &old_text, &new_text) {
136                    Ok((new_content, was_fuzzy)) => {
137                        any_fuzzy |= was_fuzzy;
138                        current = new_content;
139                    }
140                    Err(_) => {
141                        return Ok(ToolOutput::error(format!(
142                            "Edit {} of {} failed in {edit_path}: could not find old_text in file (after applying previous edits).\nold_text starts with: {:?}",
143                            i + 1,
144                            file_edits.len(),
145                            truncate_chars_with_suffix(&old_text, 80, "")
146                        )));
147                    }
148                }
149            }
150
151            total_edits += file_edits.len();
152            let diff = generate_diff(&edit_path, &original, &current);
153            let final_content = if has_crlf {
154                current.replace('\n', "\r\n")
155            } else {
156                current.clone()
157            };
158            prepared.push(PreparedEditFile {
159                input_path: edit_path,
160                path,
161                final_content,
162                diff,
163                edit_count: file_edits.len(),
164            });
165        }
166
167        let touched_paths = prepared
168            .iter()
169            .map(|prepared| prepared.path.clone())
170            .collect::<Vec<_>>();
171        if !dry_run {
172            ctx.checkpoint_state
173                .snapshot_paths(&touched_paths, Some("multi_edit transaction".to_string()))?;
174            for prepared in &prepared {
175                tokio::fs::write(&prepared.path, &prepared.final_content).await?;
176                if let Ok(mut tracker) = ctx.file_tracker.lock() {
177                    tracker.record_read(&prepared.path);
178                }
179            }
180        }
181
182        let combined_diff = prepared
183            .iter()
184            .map(|prepared| prepared.diff.as_str())
185            .collect::<Vec<_>>()
186            .join("\n\n");
187        let mut msg = format!(
188            "Validated {} edits across {} file(s) as one transaction",
189            total_edits,
190            prepared.len()
191        );
192        if dry_run {
193            msg.push_str(" (dry run: no changes written)");
194        } else {
195            msg.push_str(" and applied them");
196        }
197        msg.push_str("\n\n");
198        msg.push_str(&combined_diff);
199        if any_fuzzy {
200            msg.push_str("\n(some edits used fuzzy matching)");
201        }
202        for warning in &warnings {
203            msg.push('\n');
204            msg.push_str(warning);
205        }
206
207        Ok(ToolOutput {
208            content: vec![imp_llm::ContentBlock::Text { text: msg }],
209            details: json!({
210                "transaction": true,
211                "dry_run": dry_run,
212                "files": prepared.iter().map(|prepared| json!({
213                    "path": prepared.path.display().to_string(),
214                    "input_path": prepared.input_path,
215                    "edit_count": prepared.edit_count,
216                })).collect::<Vec<_>>(),
217                "edit_count": total_edits,
218                "edits_applied": if dry_run { 0 } else { total_edits },
219                "fuzzy_match": any_fuzzy,
220                "checkpoint_created": !dry_run,
221            }),
222            is_error: false,
223        })
224    }
225}
226
227struct PreparedEditFile {
228    input_path: String,
229    path: std::path::PathBuf,
230    final_content: String,
231    diff: String,
232    edit_count: usize,
233}
234
235fn tracker_warning(ctx: &ToolContext, path: &std::path::Path) -> Option<String> {
236    let tracker = ctx.file_tracker.lock().ok()?;
237    if !tracker.was_read(path) {
238        Some(format!(
239            "Warning: editing {} without reading it first. Consider reading to verify current content.",
240            path.display()
241        ))
242    } else if tracker.is_stale(path) {
243        Some(format!(
244            "Warning: {} was modified externally since last read. Re-read to verify current content.",
245            path.display()
246        ))
247    } else {
248        None
249    }
250}
251
252fn reject_overlapping_exact_edits(
253    edit_path: &str,
254    original: &str,
255    file_edits: &[&serde_json::Value],
256) -> Result<()> {
257    let mut exact_ranges = Vec::new();
258    for (i, edit) in file_edits.iter().enumerate() {
259        let old_text = edit
260            .get("old_text")
261            .and_then(|v| v.as_str())
262            .or_else(|| edit.get("oldText").and_then(|v| v.as_str()))
263            .unwrap_or("")
264            .replace("\r\n", "\n");
265        if old_text.is_empty() {
266            continue;
267        }
268        if let Some(pos) = original.find(&old_text) {
269            exact_ranges.push((pos, pos + old_text.len(), i + 1));
270        }
271    }
272    exact_ranges.sort_by_key(|(start, _, _)| *start);
273    for pair in exact_ranges.windows(2) {
274        let (_, prev_end, prev_idx) = pair[0];
275        let (next_start, _, next_idx) = pair[1];
276        if next_start < prev_end {
277            return Err(crate::error::Error::Tool(format!(
278                "Overlapping edits rejected in {edit_path}: edit {prev_idx} overlaps edit {next_idx}. No changes made."
279            )));
280        }
281    }
282    Ok(())
283}
284
285#[cfg(test)]
286mod tests {
287    use super::*;
288    use crate::tools::ToolContext;
289    use std::sync::Arc;
290
291    fn test_ctx(dir: &std::path::Path) -> ToolContext {
292        let (tx, _rx) = tokio::sync::mpsc::channel(16);
293        let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::channel(16);
294        ToolContext {
295            cwd: dir.to_path_buf(),
296            cancelled: Arc::new(std::sync::atomic::AtomicBool::new(false)),
297            update_tx: tx,
298            command_tx: cmd_tx,
299            ui: Arc::new(crate::ui::NullInterface),
300            file_cache: Arc::new(crate::tools::FileCache::new()),
301            checkpoint_state: Arc::new(crate::tools::CheckpointState::new()),
302            file_tracker: Arc::new(std::sync::Mutex::new(crate::tools::FileTracker::new())),
303            anchor_store: Arc::new(crate::tools::AnchorStore::new()),
304            lua_tool_loader: None,
305            mode: crate::config::AgentMode::Full,
306            read_max_lines: 500,
307            turn_mana_review: Arc::new(std::sync::Mutex::new(
308                crate::mana_review::TurnManaReviewAccumulator::default(),
309            )),
310            config: Arc::new(crate::config::Config::default()),
311            run_policy: Default::default(),
312            supporting_provenance: Vec::new(),
313        }
314    }
315
316    #[tokio::test]
317    async fn multi_edit_sequential() {
318        let dir = tempfile::tempdir().unwrap();
319        let file = dir.path().join("seq.txt");
320        std::fs::write(&file, "aaa\nbbb\nccc\n").unwrap();
321
322        let result = MultiEditTool
323            .execute(
324                "c1",
325                json!({
326                    "path": "seq.txt",
327                    "edits": [
328                        {"oldText": "aaa", "newText": "AAA"},
329                        {"oldText": "bbb", "newText": "BBB"}
330                    ]
331                }),
332                test_ctx(dir.path()),
333            )
334            .await
335            .unwrap();
336
337        assert!(!result.is_error);
338        let written = std::fs::read_to_string(&file).unwrap();
339        assert!(written.contains("AAA"));
340        assert!(written.contains("BBB"));
341        assert!(written.contains("ccc"));
342        assert_eq!(result.details["transaction"], true);
343    }
344
345    #[tokio::test]
346    async fn multi_edit_atomic_rollback() {
347        let dir = tempfile::tempdir().unwrap();
348        let file = dir.path().join("atomic.txt");
349        std::fs::write(&file, "foo\nbar\nbaz\n").unwrap();
350
351        let result = MultiEditTool
352            .execute(
353                "c2",
354                json!({
355                    "path": "atomic.txt",
356                    "edits": [
357                        {"oldText": "foo", "newText": "FOO"},
358                        {"oldText": "nonexistent", "newText": "X"}
359                    ]
360                }),
361                test_ctx(dir.path()),
362            )
363            .await
364            .unwrap();
365
366        assert!(result.is_error);
367        assert_eq!(std::fs::read_to_string(&file).unwrap(), "foo\nbar\nbaz\n");
368    }
369
370    #[tokio::test]
371    async fn multi_edit_sees_previous_results() {
372        let dir = tempfile::tempdir().unwrap();
373        let file = dir.path().join("chain.txt");
374        std::fs::write(&file, "hello world\n").unwrap();
375
376        let result = MultiEditTool
377            .execute(
378                "c3",
379                json!({
380                    "path": "chain.txt",
381                    "edits": [
382                        {"oldText": "hello", "newText": "goodbye"},
383                        {"oldText": "goodbye world", "newText": "farewell"}
384                    ]
385                }),
386                test_ctx(dir.path()),
387            )
388            .await
389            .unwrap();
390
391        assert!(!result.is_error);
392        assert_eq!(std::fs::read_to_string(&file).unwrap(), "farewell\n");
393    }
394
395    #[tokio::test]
396    async fn multi_edit_creates_checkpoint_snapshot() {
397        let dir = tempfile::tempdir().unwrap();
398        let file = dir.path().join("checkpoint.txt");
399        std::fs::write(&file, "foo\nbar\n").unwrap();
400
401        let ctx = test_ctx(dir.path());
402        let checkpoint_state = ctx.checkpoint_state.clone();
403        let result = MultiEditTool
404            .execute(
405                "c-checkpoint",
406                json!({
407                    "path": "checkpoint.txt",
408                    "edits": [
409                        {"oldText": "foo", "newText": "FOO"},
410                        {"oldText": "bar", "newText": "BAR"}
411                    ]
412                }),
413                ctx,
414            )
415            .await
416            .unwrap();
417
418        assert!(!result.is_error);
419        assert_eq!(
420            checkpoint_state.original(&file).as_deref(),
421            Some("foo\nbar\n")
422        );
423        assert_eq!(checkpoint_state.checkpoints().len(), 1);
424        assert_eq!(result.details["checkpoint_created"], true);
425    }
426
427    #[tokio::test]
428    async fn multi_edit_empty_edits_error() {
429        let dir = tempfile::tempdir().unwrap();
430        let file = dir.path().join("empty_edits.txt");
431        std::fs::write(&file, "content\n").unwrap();
432
433        let result = MultiEditTool
434            .execute(
435                "c5",
436                json!({"path": "empty_edits.txt", "edits": []}),
437                test_ctx(dir.path()),
438            )
439            .await
440            .unwrap();
441
442        assert!(result.is_error);
443    }
444
445    #[tokio::test]
446    async fn multi_edit_missing_path_error() {
447        let dir = tempfile::tempdir().unwrap();
448
449        let result = MultiEditTool
450            .execute(
451                "c6",
452                json!({"edits": [{"oldText": "a", "newText": "b"}]}),
453                test_ctx(dir.path()),
454            )
455            .await
456            .unwrap();
457
458        assert!(result.is_error);
459    }
460
461    #[tokio::test]
462    async fn multi_edit_chained_three_edits() {
463        let dir = tempfile::tempdir().unwrap();
464        let file = dir.path().join("chain3.txt");
465        std::fs::write(&file, "apple banana cherry\n").unwrap();
466
467        let result = MultiEditTool
468            .execute(
469                "c7",
470                json!({
471                    "path": "chain3.txt",
472                    "edits": [
473                        {"oldText": "apple", "newText": "APPLE"},
474                        {"oldText": "APPLE banana", "newText": "FRUIT"},
475                        {"oldText": "cherry", "newText": "CHERRY"}
476                    ]
477                }),
478                test_ctx(dir.path()),
479            )
480            .await
481            .unwrap();
482
483        assert!(!result.is_error);
484        assert_eq!(std::fs::read_to_string(&file).unwrap(), "FRUIT CHERRY\n");
485    }
486
487    #[tokio::test]
488    async fn multi_edit_combined_diff() {
489        let dir = tempfile::tempdir().unwrap();
490        let file = dir.path().join("diff.txt");
491        std::fs::write(&file, "alpha\nbeta\ngamma\n").unwrap();
492
493        let result = MultiEditTool
494            .execute(
495                "c4",
496                json!({
497                    "path": "diff.txt",
498                    "edits": [
499                        {"oldText": "alpha", "newText": "ALPHA"},
500                        {"oldText": "gamma", "newText": "GAMMA"}
501                    ]
502                }),
503                test_ctx(dir.path()),
504            )
505            .await
506            .unwrap();
507
508        assert!(!result.is_error);
509        let text = result.text_content().unwrap();
510        assert!(text.contains("ALPHA"));
511        assert!(text.contains("GAMMA"));
512    }
513
514    #[tokio::test]
515    async fn multi_edit_can_edit_two_files_transactionally() {
516        let dir = tempfile::tempdir().unwrap();
517        let one = dir.path().join("one.txt");
518        let two = dir.path().join("two.txt");
519        std::fs::write(&one, "alpha\n").unwrap();
520        std::fs::write(&two, "beta\n").unwrap();
521
522        let result = MultiEditTool
523            .execute(
524                "c-multi-file",
525                json!({
526                    "edits": [
527                        {"path": "one.txt", "oldText": "alpha", "newText": "ALPHA"},
528                        {"path": "two.txt", "oldText": "beta", "newText": "BETA"}
529                    ]
530                }),
531                test_ctx(dir.path()),
532            )
533            .await
534            .unwrap();
535
536        assert!(!result.is_error);
537        assert_eq!(std::fs::read_to_string(&one).unwrap(), "ALPHA\n");
538        assert_eq!(std::fs::read_to_string(&two).unwrap(), "BETA\n");
539        assert_eq!(result.details["files"].as_array().unwrap().len(), 2);
540    }
541
542    #[tokio::test]
543    async fn multi_edit_rejects_overlaps_without_writing() {
544        let dir = tempfile::tempdir().unwrap();
545        let file = dir.path().join("overlap.txt");
546        std::fs::write(&file, "abcdef\n").unwrap();
547
548        let result = MultiEditTool
549            .execute(
550                "c-overlap",
551                json!({
552                    "path": "overlap.txt",
553                    "edits": [
554                        {"oldText": "abc", "newText": "ABC"},
555                        {"oldText": "bc", "newText": "BC"}
556                    ]
557                }),
558                test_ctx(dir.path()),
559            )
560            .await
561            .unwrap();
562
563        assert!(result.is_error);
564        assert!(result.text_content().unwrap().contains("Overlapping edits"));
565        assert_eq!(std::fs::read_to_string(&file).unwrap(), "abcdef\n");
566    }
567
568    #[tokio::test]
569    async fn multi_edit_dry_run_writes_nothing() {
570        let dir = tempfile::tempdir().unwrap();
571        let file = dir.path().join("dry.txt");
572        std::fs::write(&file, "alpha\n").unwrap();
573        let ctx = test_ctx(dir.path());
574
575        let result = MultiEditTool
576            .execute(
577                "c-dry",
578                json!({
579                    "path": "dry.txt",
580                    "dryRun": true,
581                    "edits": [{"oldText": "alpha", "newText": "ALPHA"}]
582                }),
583                ctx.clone(),
584            )
585            .await
586            .unwrap();
587
588        assert!(!result.is_error);
589        assert_eq!(std::fs::read_to_string(&file).unwrap(), "alpha\n");
590        assert!(ctx.checkpoint_state.checkpoints().is_empty());
591        assert_eq!(result.details["dry_run"], true);
592    }
593}