From 02ee88bdf099e0b4efec6fe77553ff87f43e9bd7 Mon Sep 17 00:00:00 2001 From: Martin Ye Date: Wed, 8 Jul 2026 11:12:31 -0700 Subject: [PATCH] agent: Fix edit_file corrupting indentation when old_text omits first-line indent --- crates/agent/src/tools/edit_file_tool.rs | 71 ++++++++ crates/agent/src/tools/edit_session.rs | 29 +++- .../agent/src/tools/edit_session/reindent.rs | 161 ++++++++++++++++-- .../edit_session/streaming_fuzzy_matcher.rs | 153 ++++++++++++++--- 4 files changed, 377 insertions(+), 37 deletions(-) diff --git a/crates/agent/src/tools/edit_file_tool.rs b/crates/agent/src/tools/edit_file_tool.rs index 8cf5610531d..4734dd6b36c 100644 --- a/crates/agent/src/tools/edit_file_tool.rs +++ b/crates/agent/src/tools/edit_file_tool.rs @@ -324,6 +324,77 @@ mod tests { assert_eq!(new_text, "line 1\nmodified line 2\nline 3\n"); } + #[gpui::test] + async fn test_streaming_edit_first_line_missing_indent(cx: &mut TestAppContext) { + // Reproduces https://github.com/zed-industries/zed/issues/60302: the + // first line of the multi-line `old_text` omits its leading + // indentation while subsequent lines include theirs, so the indent + // delta computed from the first line must not be applied to the + // following lines. `old_text` also omits the `self.extra` line, so + // the query lines don't correspond one-to-one to the matched buffer + // rows and the indent pairing must follow the fuzzy match's + // alignment instead of assuming equal line counts. + let content = concat!( + "class Outer:\n", + " def method(self):\n", + " self.kept = \"unchanged\"\n", + " self.target_a = \"before\"\n", + " self.extra = \"row\"\n", + " self.target_b = \"before\"\n", + " self.target_c = \"before\"\n", + " self.target_d = \"before\"\n", + " self.kept_2 = \"unchanged\"\n", + ); + let (edit_tool, _project, _action_log, _fs, _thread) = + setup_test(cx, json!({"file.py": content})).await; + let result = cx + .update(|cx| { + edit_tool.clone().run( + ToolInput::resolved(EditFileToolInput { + path: "root/file.py".into(), + edits: vec![Edit { + old_text: concat!( + "self.target_a = \"before\"\n", + " self.target_b = \"before\"\n", + " self.target_c = \"before\"\n", + " self.target_d = \"before\"", + ) + .into(), + new_text: concat!( + "self.target_a = \"after\"\n", + " self.target_b = \"after\"\n", + " self.target_c = \"after\"\n", + " self.target_d = \"after\"", + ) + .into(), + }], + }), + ToolCallEventStream::test().0, + cx, + ) + }) + .await; + + let EditFileToolOutput::Success { new_text, .. } = result.unwrap() else { + panic!("expected success"); + }; + // The matched range includes the `self.extra` row, so it is replaced + // along with the rest of the match. + assert_eq!( + new_text, + concat!( + "class Outer:\n", + " def method(self):\n", + " self.kept = \"unchanged\"\n", + " self.target_a = \"after\"\n", + " self.target_b = \"after\"\n", + " self.target_c = \"after\"\n", + " self.target_d = \"after\"\n", + " self.kept_2 = \"unchanged\"\n", + ) + ); + } + #[gpui::test] async fn test_streaming_edit_multiple_edits(cx: &mut TestAppContext) { let (edit_tool, _project, _action_log, _fs, _thread) = setup_test( diff --git a/crates/agent/src/tools/edit_session.rs b/crates/agent/src/tools/edit_session.rs index b6f8c0f1cfc..d47a8712333 100644 --- a/crates/agent/src/tools/edit_session.rs +++ b/crates/agent/src/tools/edit_session.rs @@ -16,7 +16,7 @@ use language::{Buffer, BufferEditSource, BufferEvent, LanguageRegistry}; use language_model::LanguageModelToolResultContent; use project::lsp_store::{FormatTrigger, LspFormatTarget}; use project::{AgentLocation, Project, ProjectPath}; -use reindent::{Reindenter, compute_indent_delta}; +use reindent::{Reindenter, compute_indent_delta, compute_rest_indent_delta}; use schemars::JsonSchema; use serde::{Deserialize, Serialize, de::DeserializeOwned}; use std::ops::Range; @@ -527,15 +527,34 @@ impl EditPipeline { ); let buffer_indent = snapshot.line_indent_for_row(line); + let query_lines = matcher.query_lines(); let query_indent = text::LineIndent::from_iter( - matcher - .query_lines() + query_lines .first() .map(|s| s.as_str()) .unwrap_or("") .chars(), ); - let indent_delta = compute_indent_delta(buffer_indent, query_indent); + let first_line_delta = compute_indent_delta(buffer_indent, query_indent); + + // Query row 0 is excluded: its delta is `first_line_delta`, + // which intentionally differs when the model stripped the + // first line's indentation. + let rest_delta = compute_rest_indent_delta( + first_line_delta, + matcher + .line_pairs(&range) + .unwrap_or(&[]) + .iter() + .filter(|(query_row, _)| *query_row != 0) + .filter_map(|(query_row, buffer_row)| { + let query_line = query_lines.get(*query_row as usize)?; + Some(( + snapshot.line_indent_for_row(*buffer_row), + text::LineIndent::from_iter(query_line.chars()), + )) + }), + ); let old_text_in_buffer = snapshot.text_for_range(range.clone()).collect::(); @@ -551,7 +570,7 @@ impl EditPipeline { self.current_edit = Some(EditPipelineEntry::StreamingNewText { streaming_diff: StreamingDiff::new(old_text_in_buffer), edit_cursor: range.start, - reindenter: Reindenter::new(indent_delta), + reindenter: Reindenter::with_deltas(first_line_delta, rest_delta), original_snapshot: text_snapshot, }); diff --git a/crates/agent/src/tools/edit_session/reindent.rs b/crates/agent/src/tools/edit_session/reindent.rs index 7f08749e475..ed4066e30d2 100644 --- a/crates/agent/src/tools/edit_session/reindent.rs +++ b/crates/agent/src/tools/edit_session/reindent.rs @@ -1,7 +1,7 @@ use language::LineIndent; use std::{cmp, iter}; -#[derive(Copy, Clone, Debug)] +#[derive(Copy, Clone, Debug, PartialEq, Eq)] pub enum IndentDelta { Spaces(isize), Tabs(isize), @@ -31,20 +31,60 @@ pub fn compute_indent_delta(buffer_indent: LineIndent, query_indent: LineIndent) } } +/// Computes the indent delta for the lines after the first, given per-line +/// `(buffer, query)` indents for those lines. +/// +/// When the remaining lines agree on a consistent delta, that delta is +/// returned even if it differs from `first_line_delta`. This handles queries +/// where only the first line's indentation was stripped. When the remaining +/// lines are inconsistent (or all blank), falls back to `first_line_delta`, +/// preserving the uniform re-indentation behavior. +pub fn compute_rest_indent_delta( + first_line_delta: IndentDelta, + indent_pairs: impl IntoIterator, +) -> IndentDelta { + let mut rest_delta = None; + for (buffer_indent, query_indent) in indent_pairs { + if buffer_indent.line_blank || query_indent.line_blank { + continue; + } + let delta = compute_indent_delta(buffer_indent, query_indent); + match rest_delta { + None => rest_delta = Some(delta), + Some(existing) if existing == delta => {} + Some(_) => return first_line_delta, + } + } + rest_delta.unwrap_or(first_line_delta) +} + /// Synchronous re-indentation adapter. Buffers incomplete lines and applies /// an `IndentDelta` to each line's leading whitespace before emitting it. +/// +/// Models sometimes omit the leading indentation only on the first line of +/// `old_text`/`new_text` (e.g. when copying from mid-line context), so the +/// first line and the remaining lines can require different deltas. pub struct Reindenter { - delta: IndentDelta, + first_line_delta: IndentDelta, + rest_delta: IndentDelta, buffer: String, in_leading_whitespace: bool, + on_first_line: bool, } impl Reindenter { - pub fn new(delta: IndentDelta) -> Self { + #[cfg(test)] + fn uniform(delta: IndentDelta) -> Self { + Self::with_deltas(delta, delta) + } + + pub fn with_deltas(first_line_delta: IndentDelta, rest_delta: IndentDelta) -> Self { Self { - delta, + first_line_delta, + rest_delta, buffer: String::new(), in_leading_whitespace: true, + on_first_line: true, } } @@ -70,14 +110,19 @@ impl Reindenter { None => (self.buffer.len(), true), }; let line = &self.buffer[start_ix..line_end]; + let delta = if self.on_first_line { + self.first_line_delta + } else { + self.rest_delta + }; if self.in_leading_whitespace { - if let Some(non_whitespace_ix) = line.find(|c| self.delta.character() != c) { + if let Some(non_whitespace_ix) = line.find(|c| delta.character() != c) { // We found a non-whitespace character, adjust indentation // based on the delta. let new_indent_len = - cmp::max(0, non_whitespace_ix as isize + self.delta.len()) as usize; - indented.extend(iter::repeat(self.delta.character()).take(new_indent_len)); + cmp::max(0, non_whitespace_ix as isize + delta.len()) as usize; + indented.extend(iter::repeat(delta.character()).take(new_indent_len)); indented.push_str(&line[non_whitespace_ix..]); self.in_leading_whitespace = false; } else if is_pending_line && !is_final { @@ -97,6 +142,7 @@ impl Reindenter { break; } else { self.in_leading_whitespace = true; + self.on_first_line = false; indented.push('\n'); start_ix = line_end + 1; } @@ -116,7 +162,7 @@ mod tests { #[test] fn test_indent_single_chunk() { - let mut r = Reindenter::new(IndentDelta::Spaces(2)); + let mut r = Reindenter::uniform(IndentDelta::Spaces(2)); let out = r.push(" abc\n def\n ghi"); // All three lines are emitted: "ghi" starts with spaces but // contains non-whitespace, so it's processed immediately. @@ -127,7 +173,7 @@ mod tests { #[test] fn test_outdent_tabs() { - let mut r = Reindenter::new(IndentDelta::Tabs(-2)); + let mut r = Reindenter::uniform(IndentDelta::Tabs(-2)); let out = r.push("\t\t\t\tabc\n\t\tdef\n\t\t\t\t\t\tghi"); assert_eq!(out, "\t\tabc\ndef\n\t\t\t\tghi"); let out = r.finish(); @@ -136,7 +182,7 @@ mod tests { #[test] fn test_incremental_chunks() { - let mut r = Reindenter::new(IndentDelta::Spaces(2)); + let mut r = Reindenter::uniform(IndentDelta::Spaces(2)); // Feed " ab" — the `a` is non-whitespace, so the line is // processed immediately even without a trailing newline. let out = r.push(" ab"); @@ -151,7 +197,7 @@ mod tests { #[test] fn test_zero_delta() { - let mut r = Reindenter::new(IndentDelta::Spaces(0)); + let mut r = Reindenter::uniform(IndentDelta::Spaces(0)); let out = r.push(" hello\n world\n"); assert_eq!(out, " hello\n world\n"); let out = r.finish(); @@ -160,7 +206,7 @@ mod tests { #[test] fn test_clamp_negative_indent() { - let mut r = Reindenter::new(IndentDelta::Spaces(-10)); + let mut r = Reindenter::uniform(IndentDelta::Spaces(-10)); let out = r.push(" abc\n"); // max(0, 2 - 10) = 0, so no leading spaces. assert_eq!(out, "abc\n"); @@ -170,7 +216,7 @@ mod tests { #[test] fn test_whitespace_only_lines() { - let mut r = Reindenter::new(IndentDelta::Spaces(2)); + let mut r = Reindenter::uniform(IndentDelta::Spaces(2)); let out = r.push(" \n code\n"); // First line is all whitespace — emitted verbatim. Second line is indented. assert_eq!(out, " \n code\n"); @@ -178,6 +224,95 @@ mod tests { assert_eq!(out, ""); } + #[test] + fn test_distinct_first_line_delta() { + // First line's indentation was stripped in the query (delta +8), + // while the remaining lines are already correct (delta 0). Chunks + // split mid-line and mid-indentation to exercise the streaming path, + // and the blank line is passed through verbatim. + let mut r = Reindenter::with_deltas(IndentDelta::Spaces(8), IndentDelta::Spaces(0)); + let mut out = r.push("self.target_a = "); + out.push_str(&r.push("\"after\"\n ")); + out.push_str(&r.push(" self.target_b = \"after\"\n")); + out.push_str(&r.push("\n self.target_c = \"after\"")); + out.push_str(&r.finish()); + assert_eq!( + out, + concat!( + " self.target_a = \"after\"\n", + " self.target_b = \"after\"\n", + "\n", + " self.target_c = \"after\"", + ) + ); + } + + fn line_indent(text: &str) -> LineIndent { + LineIndent::from_iter(text.chars()) + } + + #[test] + fn test_compute_rest_indent_delta() { + let first_line_delta = IndentDelta::Spaces(8); + + // Remaining lines that agree on a delta override the first-line + // delta, and blank lines are skipped when forming the consensus. + assert_eq!( + compute_rest_indent_delta( + first_line_delta, + vec![ + (line_indent(" b"), line_indent(" b")), + (line_indent(""), line_indent("")), + (line_indent(" c"), line_indent(" c")), + ], + ), + IndentDelta::Spaces(0) + ); + assert_eq!( + compute_rest_indent_delta( + first_line_delta, + vec![ + (line_indent(" b"), line_indent(" b")), + (line_indent(" "), line_indent("")), + (line_indent(" c"), line_indent(" c")), + ], + ), + IndentDelta::Spaces(4) + ); + assert_eq!( + compute_rest_indent_delta( + first_line_delta, + vec![(line_indent("\t\tb"), line_indent("\tb"))], + ), + IndentDelta::Tabs(1) + ); + + // Inconsistent remaining lines fall back to the first-line delta... + assert_eq!( + compute_rest_indent_delta( + first_line_delta, + vec![ + (line_indent(" b"), line_indent(" b")), + (line_indent(" c"), line_indent(" c")), + ], + ), + first_line_delta + ); + + // ...and so do all-blank and empty pairings. + assert_eq!( + compute_rest_indent_delta( + first_line_delta, + vec![(line_indent(" "), line_indent(""))], + ), + first_line_delta + ); + assert_eq!( + compute_rest_indent_delta(first_line_delta, vec![]), + first_line_delta + ); + } + #[test] fn test_compute_indent_delta_spaces() { let buffer = LineIndent { diff --git a/crates/agent/src/tools/edit_session/streaming_fuzzy_matcher.rs b/crates/agent/src/tools/edit_session/streaming_fuzzy_matcher.rs index e6a56099a29..3515275f70a 100644 --- a/crates/agent/src/tools/edit_session/streaming_fuzzy_matcher.rs +++ b/crates/agent/src/tools/edit_session/streaming_fuzzy_matcher.rs @@ -12,10 +12,18 @@ pub struct StreamingFuzzyMatcher { query_lines: Vec, line_hint: Option, incomplete_line: String, - matches: Vec>, + matches: Vec, matrix: SearchMatrix, } +/// A match candidate: the matched byte range plus the 0-based +/// `(query_row, buffer_row)` line pairs the search aligned to produce it. +#[derive(Clone, Debug)] +struct SearchMatch { + range: Range, + line_pairs: Vec<(u32, u32)>, +} + impl StreamingFuzzyMatcher { pub fn new(snapshot: TextBufferSnapshot) -> Self { let buffer_line_count = snapshot.max_point().row as usize + 1; @@ -34,6 +42,23 @@ impl StreamingFuzzyMatcher { &self.query_lines } + /// Returns the 0-based `(query_row, buffer_row)` line pairs that the + /// search aligned for the match with the given range. Lines that were + /// skipped on either side of the alignment are absent. + pub fn line_pairs(&self, range: &Range) -> Option<&[(u32, u32)]> { + self.matches + .iter() + .find(|search_match| search_match.range == *range) + .map(|search_match| search_match.line_pairs.as_slice()) + } + + fn match_ranges(&self) -> Vec> { + self.matches + .iter() + .map(|search_match| search_match.range.clone()) + .collect() + } + /// Push a new chunk of text and get the best match found so far. /// /// This method accumulates text chunks and processes complete lines. @@ -62,7 +87,11 @@ impl StreamingFuzzyMatcher { } let best_match = self.select_best_match(); - best_match.or_else(|| self.matches.first().cloned()) + best_match.or_else(|| { + self.matches + .first() + .map(|search_match| search_match.range.clone()) + }) } /// Finish processing and return the final best match(es). @@ -72,26 +101,35 @@ impl StreamingFuzzyMatcher { pub fn finish(&mut self) -> Vec> { // Process any remaining incomplete line if !self.incomplete_line.is_empty() { - if self.matches.len() == 1 { - let range = &mut self.matches[0]; + if let [only_match] = self.matches.as_mut_slice() { + let range = &mut only_match.range; if range.end < self.snapshot.len() && self .snapshot .contains_str_at(range.end + 1, &self.incomplete_line) { range.end += 1 + self.incomplete_line.len(); - return self.matches.clone(); + // Record the line and its alignment so that `query_lines` + // and `line_pairs` stay in sync with the lines covered by + // the returned range. + let extended_row = self.snapshot.offset_to_point(range.end).row; + self.query_lines + .push(std::mem::take(&mut self.incomplete_line)); + only_match + .line_pairs + .push(((self.query_lines.len() - 1) as u32, extended_row)); + return self.match_ranges(); } } - self.query_lines.push(self.incomplete_line.clone()); - self.incomplete_line.clear(); + self.query_lines + .push(std::mem::take(&mut self.incomplete_line)); self.matches = self.resolve_location_fuzzy(); } - self.matches.clone() + self.match_ranges() } - fn resolve_location_fuzzy(&mut self) -> Vec> { + fn resolve_location_fuzzy(&mut self) -> Vec { let new_query_line_count = self.query_lines.len(); let old_query_line_count = self.matrix.rows.saturating_sub(1); if new_query_line_count == old_query_line_count { @@ -167,7 +205,7 @@ impl StreamingFuzzyMatcher { // Find ranges for the matches let mut valid_matches = Vec::new(); for &buffer_row_end in &matches_with_best_cost { - let mut matched_lines = 0; + let mut line_pairs = Vec::new(); let mut query_row = new_query_line_count; let mut buffer_row_start = buffer_row_end; while query_row > 0 && buffer_row_start > 0 { @@ -176,7 +214,7 @@ impl StreamingFuzzyMatcher { SearchDirection::Diagonal => { query_row -= 1; buffer_row_start -= 1; - matched_lines += 1; + line_pairs.push((query_row as u32, buffer_row_start)); } SearchDirection::Up => { query_row -= 1; @@ -186,9 +224,10 @@ impl StreamingFuzzyMatcher { } } } + line_pairs.reverse(); let matched_buffer_row_count = buffer_row_end - buffer_row_start; - let matched_ratio = matched_lines as f32 + let matched_ratio = line_pairs.len() as f32 / (matched_buffer_row_count as f32).max(new_query_line_count as f32); if matched_ratio >= 0.8 { let buffer_start_ix = self @@ -198,11 +237,14 @@ impl StreamingFuzzyMatcher { buffer_row_end - 1, self.snapshot.line_len(buffer_row_end - 1), )); - valid_matches.push((buffer_row_start, buffer_start_ix..buffer_end_ix)); + valid_matches.push(SearchMatch { + range: buffer_start_ix..buffer_end_ix, + line_pairs, + }); } } - valid_matches.into_iter().map(|(_, range)| range).collect() + valid_matches } /// Return the best match with starting position close enough to line_hint. @@ -216,8 +258,8 @@ impl StreamingFuzzyMatcher { return None; } - if self.matches.len() == 1 { - return self.matches.first().cloned(); + if let [only_match] = self.matches.as_slice() { + return Some(only_match.range.clone()); } let Some(line_hint) = self.line_hint else { @@ -228,14 +270,14 @@ impl StreamingFuzzyMatcher { let mut best_match = None; let mut best_distance = u32::MAX; - for range in &self.matches { - let start_point = self.snapshot.offset_to_point(range.start); + for search_match in &self.matches { + let start_point = self.snapshot.offset_to_point(search_match.range.start); let start_line = start_point.row; let distance = start_line.abs_diff(line_hint); if distance <= LINE_HINT_TOLERANCE && distance < best_distance { best_distance = distance; - best_match = Some(range.clone()); + best_match = Some(search_match.range.clone()); } } @@ -831,6 +873,79 @@ mod tests { } } + #[test] + fn test_line_pairs_skip_unmatched_buffer_line() { + let text = indoc! {r#" + class Outer: + def method(self): + self.kept = "unchanged" + self.target_a = "before" + self.extra = "row" + self.target_b = "before" + self.target_c = "before" + self.target_d = "before" + self.kept_2 = "unchanged" + "#}; + let buffer = TextBuffer::new( + ReplicaId::LOCAL, + BufferId::new(1).unwrap(), + text.to_string(), + ); + let mut matcher = StreamingFuzzyMatcher::new(buffer.snapshot().clone()); + + // The query omits the `self.extra` row that sits between the matched + // buffer lines. + matcher.push( + concat!( + " self.target_a = \"before\"\n", + " self.target_b = \"before\"\n", + " self.target_c = \"before\"\n", + " self.target_d = \"before\"\n", + ), + None, + ); + let matches = matcher.finish(); + + assert_eq!(matches.len(), 1); + assert_eq!( + matcher.line_pairs(&matches[0]), + Some(&[(0, 3), (1, 5), (2, 6), (3, 7)][..]) + ); + } + + #[test] + fn test_line_pairs_include_extended_incomplete_line() { + let text = indoc! {r#" + fn on_query_change(&mut self, cx: &mut Context) { + self.filter(cx); + } + + + + fn render_search(&self, cx: &mut Context) -> Div { + div() + } + "#}; + let buffer = TextBuffer::new( + ReplicaId::LOCAL, + BufferId::new(1).unwrap(), + text.to_string(), + ); + let mut matcher = StreamingFuzzyMatcher::new(buffer.snapshot().clone()); + + // The last query line is incomplete and gets appended to the match by + // `finish` via verbatim comparison rather than the fuzzy search. + matcher.push("}\n\n\n\nfn render_search", None); + let matches = matcher.finish(); + + assert_eq!(matches.len(), 1); + assert_eq!( + matcher.line_pairs(&matches[0]), + Some(&[(0, 2), (1, 3), (2, 4), (3, 5), (4, 6)][..]) + ); + assert_eq!(matcher.query_lines().len(), 5); + } + fn to_random_chunks(rng: &mut StdRng, input: &str) -> Vec { let chunk_count = rng.random_range(1..=cmp::min(input.len(), 50)); let mut chunk_indices = (0..input.len()).choose_multiple(rng, chunk_count);