Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 71 additions & 0 deletions crates/agent/src/tools/edit_file_tool.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
29 changes: 24 additions & 5 deletions crates/agent/src/tools/edit_session.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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::<String>();

Expand All @@ -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,
});

Expand Down
161 changes: 148 additions & 13 deletions crates/agent/src/tools/edit_session/reindent.rs
Original file line number Diff line number Diff line change
@@ -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),
Expand Down Expand Up @@ -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<Item = (LineIndent, LineIndent)>,
) -> 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,
}
}

Expand All @@ -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 {
Expand All @@ -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;
}
Expand All @@ -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.
Expand All @@ -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();
Expand All @@ -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");
Expand All @@ -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();
Expand All @@ -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");
Expand All @@ -170,14 +216,103 @@ 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");
let out = r.finish();
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 {
Expand Down
Loading
Loading