From 0f20ee4eb871907defe7848f0d3e2203cfff057e Mon Sep 17 00:00:00 2001 From: Michelle Tilley Date: Tue, 21 Apr 2026 10:32:54 -0700 Subject: feat: AI tool rendering overhaul + edit_file tool (#3423) Overhaul of how AI tool calls are modeled, rendered, and displayed in the Atuin AI TUI. Fixes bugs in shell command output capture, implements the `edit_file` tool with full safety infrastructure, and adds a diff preview for edits. --- crates/atuin-ai/src/diff.rs | 294 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 294 insertions(+) create mode 100644 crates/atuin-ai/src/diff.rs (limited to 'crates/atuin-ai/src/diff.rs') diff --git a/crates/atuin-ai/src/diff.rs b/crates/atuin-ai/src/diff.rs new file mode 100644 index 00000000..663481c0 --- /dev/null +++ b/crates/atuin-ai/src/diff.rs @@ -0,0 +1,294 @@ +//! Structured diff computation for edit previews. +//! +//! Computes a line-level diff between old and new file content using +//! imara-diff's Histogram algorithm, producing structured hunks with +//! typed lines (Context, Added, Removed) suitable for TUI rendering. + +use imara_diff::{Algorithm, Diff, InternedInput}; + +/// Number of context lines to show around each change. +const CONTEXT_LINES: u32 = 3; + +/// A structured diff preview for a file edit, ready for rendering. +#[derive(Debug, Clone)] +pub(crate) struct EditPreview { + pub hunks: Vec, +} + +/// A contiguous group of diff lines (context + changes). +#[derive(Debug, Clone)] +pub(crate) struct DiffHunk { + /// 1-indexed line number of the first line in this hunk (in the original file). + pub before_start: u32, + /// 1-indexed line number of the first line in this hunk (in the new file). + pub after_start: u32, + pub lines: Vec, +} + +/// A single line in a diff hunk. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum DiffLine { + /// Unchanged line (shown for context). + Context(String), + /// Line added in the new version. + Added(String), + /// Line removed from the old version. + Removed(String), +} + +impl EditPreview { + /// Compute a structured diff between old and new file content. + /// + /// Uses the Histogram algorithm with line-level granularity and + /// indentation-aware postprocessing for readable output. + pub fn compute(old: &str, new: &str) -> Self { + let input = InternedInput::new(old, new); + let mut diff = Diff::compute(Algorithm::Histogram, &input); + diff.postprocess_lines(&input); + + let raw_hunks: Vec<_> = diff.hunks().collect(); + if raw_hunks.is_empty() { + return EditPreview { hunks: Vec::new() }; + } + + // Merge hunks that are within 2*CONTEXT_LINES of each other + // (same logic as unified diff format). + let mut merged_groups: Vec> = Vec::new(); + let mut current_group: Vec<&imara_diff::Hunk> = vec![&raw_hunks[0]]; + + for hunk in &raw_hunks[1..] { + let prev = current_group.last().unwrap(); + if hunk.before.start.saturating_sub(prev.before.end) <= 2 * CONTEXT_LINES { + current_group.push(hunk); + } else { + merged_groups.push(current_group); + current_group = vec![hunk]; + } + } + merged_groups.push(current_group); + + // Build structured hunks from merged groups + let hunks = merged_groups + .into_iter() + .map(|group| build_hunk(&group, &input)) + .collect(); + + EditPreview { hunks } + } + + /// The highest line number (from either file) that will be displayed. + /// Used to calculate gutter width. + pub fn max_line_number(&self) -> u32 { + self.hunks + .iter() + .map(|h| { + let mut before_pos = h.before_start; + let mut after_pos = h.after_start; + for line in &h.lines { + match line { + DiffLine::Context(_) => { + before_pos += 1; + after_pos += 1; + } + DiffLine::Removed(_) => before_pos += 1, + DiffLine::Added(_) => after_pos += 1, + } + } + before_pos.max(after_pos).saturating_sub(1) + }) + .max() + .unwrap_or(0) + } +} + +/// Build a single DiffHunk from a group of adjacent raw hunks. +fn build_hunk(group: &[&imara_diff::Hunk], input: &InternedInput<&str>) -> DiffHunk { + let first = group.first().unwrap(); + let last = group.last().unwrap(); + + let context_start = first.before.start.saturating_sub(CONTEXT_LINES); + let context_end = (last.before.end + CONTEXT_LINES).min(input.before.len() as u32); + + // The after-file position of context_start: same offset as before since + // context before the first change is identical in both files. + let after_context_start = first.after.start - (first.before.start - context_start); + + let mut lines = Vec::new(); + let mut pos = context_start; + + for hunk in group { + // Context lines before this hunk + for i in pos..hunk.before.start { + lines.push(DiffLine::Context(token_text(input, true, i))); + } + + // Removed lines + for i in hunk.before.start..hunk.before.end { + lines.push(DiffLine::Removed(token_text(input, true, i))); + } + + // Added lines + for i in hunk.after.start..hunk.after.end { + lines.push(DiffLine::Added(token_text(input, false, i))); + } + + pos = hunk.before.end; + } + + // Trailing context + for i in pos..context_end { + lines.push(DiffLine::Context(token_text(input, true, i))); + } + + DiffHunk { + before_start: context_start + 1, // 1-indexed + after_start: after_context_start + 1, // 1-indexed + lines, + } +} + +/// Extract the text content of a token, trimming the trailing newline +/// that imara-diff includes in line-based tokenization. +fn token_text(input: &InternedInput<&str>, is_before: bool, idx: u32) -> String { + let tokens = if is_before { + &input.before + } else { + &input.after + }; + let text = input.interner[tokens[idx as usize]]; + text.strip_suffix('\n') + .unwrap_or(text) + .strip_suffix('\r') + .unwrap_or(text.strip_suffix('\n').unwrap_or(text)) + .to_string() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn no_changes_produces_empty_preview() { + let preview = EditPreview::compute("hello\nworld\n", "hello\nworld\n"); + assert!(preview.hunks.is_empty()); + } + + #[test] + fn single_line_replacement() { + let old = "line1\nline2\nline3\n"; + let new = "line1\nchanged\nline3\n"; + let preview = EditPreview::compute(old, new); + + assert_eq!(preview.hunks.len(), 1); + let hunk = &preview.hunks[0]; + + // Should have: context(line1), removed(line2), added(changed), context(line3) + assert!(hunk.lines.contains(&DiffLine::Context("line1".into()))); + assert!(hunk.lines.contains(&DiffLine::Removed("line2".into()))); + assert!(hunk.lines.contains(&DiffLine::Added("changed".into()))); + assert!(hunk.lines.contains(&DiffLine::Context("line3".into()))); + } + + #[test] + fn addition_only() { + let old = "aaa\nbbb\n"; + let new = "aaa\nnew_line\nbbb\n"; + let preview = EditPreview::compute(old, new); + + assert_eq!(preview.hunks.len(), 1); + let hunk = &preview.hunks[0]; + assert!(hunk.lines.contains(&DiffLine::Added("new_line".into()))); + // Original lines are context + assert!(hunk.lines.contains(&DiffLine::Context("aaa".into()))); + assert!(hunk.lines.contains(&DiffLine::Context("bbb".into()))); + } + + #[test] + fn removal_only() { + let old = "aaa\nremove_me\nbbb\n"; + let new = "aaa\nbbb\n"; + let preview = EditPreview::compute(old, new); + + assert_eq!(preview.hunks.len(), 1); + let hunk = &preview.hunks[0]; + assert!(hunk.lines.contains(&DiffLine::Removed("remove_me".into()))); + } + + #[test] + fn distant_changes_produce_separate_hunks() { + // Two changes separated by more than 2*CONTEXT_LINES (6) lines + let old = "1\n2\n3\n4\n5\n6\n7\n8\n9\n10\n11\n12\n"; + let new = "1\nX\n3\n4\n5\n6\n7\n8\n9\n10\n11\nY\n"; + let preview = EditPreview::compute(old, new); + + assert_eq!(preview.hunks.len(), 2); + } + + #[test] + fn close_changes_merge_into_one_hunk() { + // Two changes separated by fewer than 2*CONTEXT_LINES lines + let old = "1\n2\n3\n4\n5\n"; + let new = "X\n2\n3\n4\nY\n"; + let preview = EditPreview::compute(old, new); + + assert_eq!(preview.hunks.len(), 1); + } + + #[test] + fn context_is_limited() { + // With CONTEXT_LINES=3, a change at line 10 shouldn't include line 1 + let mut old_lines: Vec<&str> = (1..=20).map(|_| "unchanged").collect(); + old_lines[9] = "target"; + let old = old_lines.join("\n") + "\n"; + let new = old.replace("target", "replaced"); + + let preview = EditPreview::compute(&old, &new); + assert_eq!(preview.hunks.len(), 1); + + // Should have at most 3 context lines before + 3 after + 1 removed + 1 added = 8 lines + assert!(preview.hunks[0].lines.len() <= 8); + } + + #[test] + fn max_line_number_reflects_file_position() { + let old = "a\nb\nc\n"; + let new = "a\nX\nc\n"; + let preview = EditPreview::compute(old, new); + // 3-line file, context + removed lines span positions 1-3 + assert_eq!(preview.max_line_number(), 3); + } + + #[test] + fn start_line_is_correct_for_later_changes() { + // Change at line 10 with 3 context lines → before_start = 7 + let mut lines: Vec = (1..=15).map(|i| format!("line{i}")).collect(); + let old = lines.join("\n") + "\n"; + lines[9] = "CHANGED".to_string(); + let new = lines.join("\n") + "\n"; + + let preview = EditPreview::compute(&old, &new); + assert_eq!(preview.hunks.len(), 1); + assert_eq!(preview.hunks[0].before_start, 7); // line 10 - 3 context = line 7 + assert_eq!(preview.hunks[0].after_start, 7); // same for a simple replacement + } + + #[test] + fn multiline_replacement() { + let old = "[section]\nkey1 = old1\nkey2 = old2\n[other]\n"; + let new = "[section]\nkey1 = new1\nkey2 = new2\n[other]\n"; + let preview = EditPreview::compute(old, new); + + assert_eq!(preview.hunks.len(), 1); + let hunk = &preview.hunks[0]; + assert!( + hunk.lines + .contains(&DiffLine::Removed("key1 = old1".into())) + ); + assert!( + hunk.lines + .contains(&DiffLine::Removed("key2 = old2".into())) + ); + assert!(hunk.lines.contains(&DiffLine::Added("key1 = new1".into()))); + assert!(hunk.lines.contains(&DiffLine::Added("key2 = new2".into()))); + } +} -- cgit v1.3.1