Show progress as the agent locates which range it needs to edit (#31582)

Release Notes:

- Improved latency when the agent starts streaming edits.

---------

Co-authored-by: Ben Brandt <benjamin.j.brandt@gmail.com>

crates/assistant_tools/src/edit_agent/edit_parser.rs

@@ -11,7 +11,7 @@ const END_TAGS: [&str; 3] = [OLD_TEXT_END_TAG, NEW_TEXT_END_TAG, EDITS_END_TAG];
 
 #[derive(Debug)]
 pub enum EditParserEvent {
-    OldText(String),
+    OldTextChunk { chunk: String, done: bool },
     NewTextChunk { chunk: String, done: bool },
 }
 
@@ -33,7 +33,7 @@ pub struct EditParser {
 #[derive(Debug, PartialEq)]
 enum EditParserState {
     Pending,
-    WithinOldText,
+    WithinOldText { start: bool },
     AfterOldText,
     WithinNewText { start: bool },
 }
@@ -56,20 +56,23 @@ impl EditParser {
                 EditParserState::Pending => {
                     if let Some(start) = self.buffer.find("<old_text>") {
                         self.buffer.drain(..start + "<old_text>".len());
-                        self.state = EditParserState::WithinOldText;
+                        self.state = EditParserState::WithinOldText { start: true };
                     } else {
                         break;
                     }
                 }
-                EditParserState::WithinOldText => {
-                    if let Some(tag_range) = self.find_end_tag() {
-                        let mut start = 0;
-                        if self.buffer.starts_with('\n') {
-                            start = 1;
+                EditParserState::WithinOldText { start } => {
+                    if !self.buffer.is_empty() {
+                        if *start && self.buffer.starts_with('\n') {
+                            self.buffer.remove(0);
                         }
-                        let mut old_text = self.buffer[start..tag_range.start].to_string();
-                        if old_text.ends_with('\n') {
-                            old_text.pop();
+                        *start = false;
+                    }
+
+                    if let Some(tag_range) = self.find_end_tag() {
+                        let mut chunk = self.buffer[..tag_range.start].to_string();
+                        if chunk.ends_with('\n') {
+                            chunk.pop();
                         }
 
                         self.metrics.tags += 1;
@@ -79,8 +82,14 @@ impl EditParser {
 
                         self.buffer.drain(..tag_range.end);
                         self.state = EditParserState::AfterOldText;
-                        edit_events.push(EditParserEvent::OldText(old_text));
+                        edit_events.push(EditParserEvent::OldTextChunk { chunk, done: true });
                     } else {
+                        if !self.ends_with_tag_prefix() {
+                            edit_events.push(EditParserEvent::OldTextChunk {
+                                chunk: mem::take(&mut self.buffer),
+                                done: false,
+                            });
+                        }
                         break;
                     }
                 }
@@ -115,11 +124,7 @@ impl EditParser {
                         self.state = EditParserState::Pending;
                         edit_events.push(EditParserEvent::NewTextChunk { chunk, done: true });
                     } else {
-                        let mut end_prefixes = END_TAGS
-                            .iter()
-                            .flat_map(|tag| (1..tag.len()).map(move |i| &tag[..i]))
-                            .chain(["\n"]);
-                        if end_prefixes.all(|prefix| !self.buffer.ends_with(&prefix)) {
+                        if !self.ends_with_tag_prefix() {
                             edit_events.push(EditParserEvent::NewTextChunk {
                                 chunk: mem::take(&mut self.buffer),
                                 done: false,
@@ -141,6 +146,14 @@ impl EditParser {
         Some(start_ix..start_ix + tag.len())
     }
 
+    fn ends_with_tag_prefix(&self) -> bool {
+        let mut end_prefixes = END_TAGS
+            .iter()
+            .flat_map(|tag| (1..tag.len()).map(move |i| &tag[..i]))
+            .chain(["\n"]);
+        end_prefixes.any(|prefix| self.buffer.ends_with(&prefix))
+    }
+
     pub fn finish(self) -> EditParserMetrics {
         self.metrics
     }
@@ -412,20 +425,28 @@ mod tests {
         chunk_indices.sort();
         chunk_indices.push(input.len());
 
+        let mut old_text = Some(String::new());
+        let mut new_text = None;
         let mut pending_edit = Edit::default();
         let mut edits = Vec::new();
         let mut last_ix = 0;
         for chunk_ix in chunk_indices {
             for event in parser.push(&input[last_ix..chunk_ix]) {
                 match event {
-                    EditParserEvent::OldText(old_text) => {
-                        pending_edit.old_text = old_text;
+                    EditParserEvent::OldTextChunk { chunk, done } => {
+                        old_text.as_mut().unwrap().push_str(&chunk);
+                        if done {
+                            pending_edit.old_text = old_text.take().unwrap();
+                            new_text = Some(String::new());
+                        }
                     }
                     EditParserEvent::NewTextChunk { chunk, done } => {
-                        pending_edit.new_text.push_str(&chunk);
+                        new_text.as_mut().unwrap().push_str(&chunk);
                         if done {
+                            pending_edit.new_text = new_text.take().unwrap();
                             edits.push(pending_edit);
                             pending_edit = Edit::default();
+                            old_text = Some(String::new());
                         }
                     }
                 }
@@ -433,8 +454,6 @@ mod tests {
             last_ix = chunk_ix;
         }
 
-        assert_eq!(pending_edit, Edit::default(), "unfinished edit");
-
         edits
     }
 }