agent: Fix streaming edit file tool inserting newlines when old_text ends with newline (#52661)

Bennet Bo Fenner created

Self-Review Checklist:

- [x] I've reviewed my own diff for quality, security, and reliability
- [x] Unsafe blocks (if any) have justifying comments
- [x] The content is consistent with the [UI/UX
checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)
- [x] Tests cover the new/changed behavior
- [x] Performance impact has been considered and is acceptable

Release Notes:

- N/A

Change summary

crates/agent/src/tools/streaming_edit_file_tool.rs |  39 +++++
crates/agent/src/tools/tool_edit_parser.rs         | 113 +++++++++++++--
2 files changed, 133 insertions(+), 19 deletions(-)

Detailed changes

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

@@ -3969,6 +3969,45 @@ mod tests {
         );
     }
 
+    #[gpui::test]
+    async fn test_streaming_edit_preserves_blank_line_after_trailing_newline_replacement(
+        cx: &mut TestAppContext,
+    ) {
+        let file_content = "before\ntarget\n\nafter\n";
+        let old_text = "target\n";
+        let new_text = "one\ntwo\ntarget\n";
+        let expected = "before\none\ntwo\ntarget\n\nafter\n";
+
+        let (tool, _project, _action_log, _fs, _thread) =
+            setup_test(cx, json!({"file.rs": file_content})).await;
+        let (sender, input) = ToolInput::<StreamingEditFileToolInput>::test();
+        let (event_stream, _receiver) = ToolCallEventStream::test();
+        let task = cx.update(|cx| tool.clone().run(input, event_stream, cx));
+
+        sender.send_final(json!({
+            "display_description": "description",
+            "path": "root/file.rs",
+            "mode": "edit",
+            "edits": [{"old_text": old_text, "new_text": new_text}]
+        }));
+
+        let result = task.await;
+
+        let StreamingEditFileToolOutput::Success {
+            new_text: final_text,
+            ..
+        } = result.unwrap()
+        else {
+            panic!("expected success");
+        };
+
+        pretty_assertions::assert_eq!(
+            final_text,
+            expected,
+            "Edit should preserve a single blank line before test_after"
+        );
+    }
+
     #[gpui::test]
     async fn test_streaming_reject_created_file_deletes_it(cx: &mut TestAppContext) {
         let (tool, _project, action_log, fs, _thread) = setup_test(cx, json!({"dir": {}})).await;

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

@@ -78,7 +78,7 @@ impl ToolEditParser {
                 if partial.new_text.is_some() {
                     // new_text appeared, so old_text is done — emit everything.
                     let start = state.old_text_emitted_len.min(old_text.len());
-                    let chunk = old_text[start..].to_string();
+                    let chunk = normalize_done_chunk(old_text[start..].to_string());
                     state.old_text_done = true;
                     state.old_text_emitted_len = old_text.len();
                     events.push(ToolEditEvent::OldTextChunk {
@@ -87,7 +87,8 @@ impl ToolEditParser {
                         done: true,
                     });
                 } else {
-                    let safe_end = safe_emit_end(old_text);
+                    let safe_end = safe_emit_end_for_edit_text(old_text);
+
                     if safe_end > state.old_text_emitted_len {
                         let chunk = old_text[state.old_text_emitted_len..safe_end].to_string();
                         state.old_text_emitted_len = safe_end;
@@ -104,7 +105,8 @@ impl ToolEditParser {
             if let Some(new_text) = &partial.new_text
                 && !state.new_text_done
             {
-                let safe_end = safe_emit_end(new_text);
+                let safe_end = safe_emit_end_for_edit_text(new_text);
+
                 if safe_end > state.new_text_emitted_len {
                     let chunk = new_text[state.new_text_emitted_len..safe_end].to_string();
                     state.new_text_emitted_len = safe_end;
@@ -160,7 +162,7 @@ impl ToolEditParser {
 
             if !state.old_text_done {
                 let start = state.old_text_emitted_len.min(edit.old_text.len());
-                let chunk = edit.old_text[start..].to_string();
+                let chunk = normalize_done_chunk(edit.old_text[start..].to_string());
                 state.old_text_done = true;
                 state.old_text_emitted_len = edit.old_text.len();
                 events.push(ToolEditEvent::OldTextChunk {
@@ -172,7 +174,7 @@ impl ToolEditParser {
 
             if !state.new_text_done {
                 let start = state.new_text_emitted_len.min(edit.new_text.len());
-                let chunk = edit.new_text[start..].to_string();
+                let chunk = normalize_done_chunk(edit.new_text[start..].to_string());
                 state.new_text_done = true;
                 state.new_text_emitted_len = edit.new_text.len();
                 events.push(ToolEditEvent::NewTextChunk {
@@ -252,6 +254,22 @@ fn safe_emit_end(text: &str) -> usize {
     }
 }
 
+fn safe_emit_end_for_edit_text(text: &str) -> usize {
+    let safe_end = safe_emit_end(text);
+    if safe_end > 0 && text.as_bytes()[safe_end - 1] == b'\n' {
+        safe_end - 1
+    } else {
+        safe_end
+    }
+}
+
+fn normalize_done_chunk(mut chunk: String) -> String {
+    if chunk.ends_with('\n') {
+        chunk.pop();
+    }
+    chunk
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -337,6 +355,69 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_done_chunks_strip_trailing_newline() {
+        let mut parser = ToolEditParser::default();
+
+        let events = parser.finalize_edits(&[Edit {
+            old_text: "before\n".into(),
+            new_text: "after\n".into(),
+        }]);
+        assert_eq!(
+            events.as_slice(),
+            &[
+                ToolEditEvent::OldTextChunk {
+                    edit_index: 0,
+                    chunk: "before".into(),
+                    done: true,
+                },
+                ToolEditEvent::NewTextChunk {
+                    edit_index: 0,
+                    chunk: "after".into(),
+                    done: true,
+                },
+            ]
+        );
+    }
+
+    #[test]
+    fn test_partial_edit_chunks_hold_back_trailing_newline() {
+        let mut parser = ToolEditParser::default();
+
+        let events = parser.push_edits(&[PartialEdit {
+            old_text: Some("before\n".into()),
+            new_text: Some("after\n".into()),
+        }]);
+        assert_eq!(
+            events.as_slice(),
+            &[
+                ToolEditEvent::OldTextChunk {
+                    edit_index: 0,
+                    chunk: "before".into(),
+                    done: true,
+                },
+                ToolEditEvent::NewTextChunk {
+                    edit_index: 0,
+                    chunk: "after".into(),
+                    done: false,
+                },
+            ]
+        );
+
+        let events = parser.finalize_edits(&[Edit {
+            old_text: "before\n".into(),
+            new_text: "after\n".into(),
+        }]);
+        assert_eq!(
+            events.as_slice(),
+            &[ToolEditEvent::NewTextChunk {
+                edit_index: 0,
+                chunk: "".into(),
+                done: true,
+            }]
+        );
+    }
+
     #[test]
     fn test_multiple_edits_sequential() {
         let mut parser = ToolEditParser::default();
@@ -858,22 +939,16 @@ mod tests {
         );
 
         // Next partial: the fixer corrects the escape to \n.
-        // The held-back byte was wrong, but we never emitted it. Now the
-        // correct newline at that position is emitted normally.
+        // Because edit text also holds back a trailing newline, nothing new
+        // is emitted yet.
         let events = parser.push_edits(&[PartialEdit {
             old_text: Some("hello,\n".into()),
             new_text: None,
         }]);
-        assert_eq!(
-            events.as_slice(),
-            &[ToolEditEvent::OldTextChunk {
-                edit_index: 0,
-                chunk: "\n".into(),
-                done: false,
-            }]
-        );
+        assert!(events.is_empty());
 
-        // Continue normally.
+        // Continue normally. The held-back newline is emitted together with the
+        // next content once it is no longer trailing.
         let events = parser.push_edits(&[PartialEdit {
             old_text: Some("hello,\nworld".into()),
             new_text: None,
@@ -882,7 +957,7 @@ mod tests {
             events.as_slice(),
             &[ToolEditEvent::OldTextChunk {
                 edit_index: 0,
-                chunk: "world".into(),
+                chunk: "\nworld".into(),
                 done: false,
             }]
         );
@@ -919,7 +994,7 @@ mod tests {
                 },
                 ToolEditEvent::NewTextChunk {
                     edit_index: 0,
-                    chunk: "LINE1\n".into(),
+                    chunk: "LINE1".into(),
                     done: false,
                 },
             ]
@@ -933,7 +1008,7 @@ mod tests {
             events.as_slice(),
             &[ToolEditEvent::NewTextChunk {
                 edit_index: 0,
-                chunk: "LINE2\nLINE3".into(),
+                chunk: "\nLINE2\nLINE3".into(),
                 done: false,
             }]
         );