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