agent: Simplify EditPipeline state in StreamingEditFileTool (#50792)

Bennet Bo Fenner created

Release Notes:

- N/A

Change summary

crates/agent/src/tools/streaming_edit_file_tool.rs | 72 ++++-----------
1 file changed, 21 insertions(+), 51 deletions(-)

Detailed changes

crates/agent/src/tools/streaming_edit_file_tool.rs 🔗

@@ -435,7 +435,7 @@ pub struct EditSession {
 }
 
 struct EditPipeline {
-    edits: Vec<EditPipelineEntry>,
+    current_edit: Option<EditPipelineEntry>,
     content_written: bool,
 }
 
@@ -449,26 +449,20 @@ enum EditPipelineEntry {
         reindenter: Reindenter,
         original_snapshot: text::BufferSnapshot,
     },
-    Done,
 }
 
 impl EditPipeline {
     fn new() -> Self {
         Self {
-            edits: Vec::new(),
+            current_edit: None,
             content_written: false,
         }
     }
 
-    fn ensure_resolving_old_text(
-        &mut self,
-        edit_index: usize,
-        buffer: &Entity<Buffer>,
-        cx: &mut AsyncApp,
-    ) {
-        while self.edits.len() <= edit_index {
+    fn ensure_resolving_old_text(&mut self, buffer: &Entity<Buffer>, cx: &mut AsyncApp) {
+        if self.current_edit.is_none() {
             let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.text_snapshot());
-            self.edits.push(EditPipelineEntry::ResolvingOldText {
+            self.current_edit = Some(EditPipelineEntry::ResolvingOldText {
                 matcher: StreamingFuzzyMatcher::new(snapshot),
             });
         }
@@ -703,15 +697,12 @@ impl EditSession {
                 }
 
                 ToolEditEvent::OldTextChunk {
-                    edit_index,
-                    chunk,
-                    done: false,
+                    chunk, done: false, ..
                 } => {
-                    self.pipeline
-                        .ensure_resolving_old_text(*edit_index, &self.buffer, cx);
+                    self.pipeline.ensure_resolving_old_text(&self.buffer, cx);
 
-                    if let EditPipelineEntry::ResolvingOldText { matcher } =
-                        &mut self.pipeline.edits[*edit_index]
+                    if let Some(EditPipelineEntry::ResolvingOldText { matcher }) =
+                        &mut self.pipeline.current_edit
                         && !chunk.is_empty()
                     {
                         if let Some(match_range) = matcher.push(chunk, None) {
@@ -734,11 +725,10 @@ impl EditSession {
                     chunk,
                     done: true,
                 } => {
-                    self.pipeline
-                        .ensure_resolving_old_text(*edit_index, &self.buffer, cx);
+                    self.pipeline.ensure_resolving_old_text(&self.buffer, cx);
 
-                    let EditPipelineEntry::ResolvingOldText { matcher } =
-                        &mut self.pipeline.edits[*edit_index]
+                    let Some(EditPipelineEntry::ResolvingOldText { matcher }) =
+                        &mut self.pipeline.current_edit
                     else {
                         continue;
                     };
@@ -763,13 +753,7 @@ impl EditSession {
                         ]),
                     );
 
-                    let EditPipelineEntry::ResolvingOldText { matcher } =
-                        &self.pipeline.edits[*edit_index]
-                    else {
-                        continue;
-                    };
-                    let buffer_indent =
-                        snapshot.line_indent_for_row(snapshot.offset_to_point(range.start).row);
+                    let buffer_indent = snapshot.line_indent_for_row(line);
                     let query_indent = text::LineIndent::from_iter(
                         matcher
                             .query_lines()
@@ -786,12 +770,12 @@ impl EditSession {
                     let text_snapshot = self
                         .buffer
                         .read_with(cx, |buffer, _cx| buffer.text_snapshot());
-                    self.pipeline.edits[*edit_index] = EditPipelineEntry::StreamingNewText {
+                    self.pipeline.current_edit = Some(EditPipelineEntry::StreamingNewText {
                         streaming_diff: StreamingDiff::new(old_text_in_buffer),
                         edit_cursor: range.start,
                         reindenter: Reindenter::new(indent_delta),
                         original_snapshot: text_snapshot,
-                    };
+                    });
 
                     cx.update(|cx| {
                         let position = self.buffer.read(cx).anchor_before(range.end);
@@ -800,20 +784,15 @@ impl EditSession {
                 }
 
                 ToolEditEvent::NewTextChunk {
-                    edit_index,
-                    chunk,
-                    done: false,
+                    chunk, done: false, ..
                 } => {
-                    if *edit_index >= self.pipeline.edits.len() {
-                        continue;
-                    }
-                    let EditPipelineEntry::StreamingNewText {
+                    let Some(EditPipelineEntry::StreamingNewText {
                         streaming_diff,
                         edit_cursor,
                         reindenter,
                         original_snapshot,
                         ..
-                    } = &mut self.pipeline.edits[*edit_index]
+                    }) = &mut self.pipeline.current_edit
                     else {
                         continue;
                     };
@@ -840,23 +819,14 @@ impl EditSession {
                 }
 
                 ToolEditEvent::NewTextChunk {
-                    edit_index,
-                    chunk,
-                    done: true,
+                    chunk, done: true, ..
                 } => {
-                    if *edit_index >= self.pipeline.edits.len() {
-                        continue;
-                    }
-
-                    let EditPipelineEntry::StreamingNewText {
+                    let Some(EditPipelineEntry::StreamingNewText {
                         mut streaming_diff,
                         mut edit_cursor,
                         mut reindenter,
                         original_snapshot,
-                    } = std::mem::replace(
-                        &mut self.pipeline.edits[*edit_index],
-                        EditPipelineEntry::Done,
-                    )
+                    }) = self.pipeline.current_edit.take()
                     else {
                         continue;
                     };