Fix issue in `StreamingEditFileTool` with incomplete last line (#51747)

Bennet Bo Fenner and Antonio Scandurra created

<img width="927" height="368" alt="image"
src="https://github.com/user-attachments/assets/05de37f7-3034-4060-96b6-49c426170f39"
/>

Before you mark this PR as ready for review, make sure that you have:
- [x] Added a solid test coverage and/or screenshots from doing manual
testing
- [x] Done a self-review taking into account security and performance
aspects
- [x] Aligned any UI changes with the [UI
checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)

Release Notes:

- N/A

---------

Co-authored-by: Antonio Scandurra <me@as-cii.com>

Change summary

crates/agent/src/edit_agent/streaming_fuzzy_matcher.rs | 60 ++++++++++++
crates/agent/src/tools/streaming_edit_file_tool.rs     | 53 ++++++++++
2 files changed, 112 insertions(+), 1 deletion(-)

Detailed changes

crates/agent/src/edit_agent/streaming_fuzzy_matcher.rs 🔗

@@ -72,6 +72,18 @@ impl StreamingFuzzyMatcher {
     pub fn finish(&mut self) -> Vec<Range<usize>> {
         // Process any remaining incomplete line
         if !self.incomplete_line.is_empty() {
+            if self.matches.len() == 1 {
+                let range = &mut self.matches[0];
+                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();
+                }
+            }
+
             self.query_lines.push(self.incomplete_line.clone());
             self.incomplete_line.clear();
             self.matches = self.resolve_location_fuzzy();
@@ -722,6 +734,54 @@ mod tests {
         );
     }
 
+    #[gpui::test]
+    fn test_prefix_of_last_line_resolves_to_correct_range() {
+        let text = indoc! {r#"
+            fn on_query_change(&mut self, cx: &mut Context<Self>) {
+                self.filter(cx);
+            }
+
+
+
+            fn render_search(&self, cx: &mut Context<Self>) -> Div {
+                div()
+            }
+        "#};
+
+        let buffer = TextBuffer::new(
+            ReplicaId::LOCAL,
+            BufferId::new(1).unwrap(),
+            text.to_string(),
+        );
+        let snapshot = buffer.snapshot();
+
+        // Query with a partial last line.
+        let query = "}\n\n\n\nfn render_search";
+
+        let mut matcher = StreamingFuzzyMatcher::new(snapshot.clone());
+        matcher.push(query, None);
+        let matches = matcher.finish();
+
+        // The match should include the line containing "fn render_search".
+        let matched_text = matches
+            .first()
+            .map(|range| snapshot.text_for_range(range.clone()).collect::<String>());
+
+        assert!(
+            matches.len() == 1,
+            "Expected exactly one match, got {}: {:?}",
+            matches.len(),
+            matched_text,
+        );
+
+        let matched_text = matched_text.unwrap();
+        pretty_assertions::assert_eq!(
+            matched_text,
+            "}\n\n\n\nfn render_search",
+            "Match should include the render_search line",
+        );
+    }
+
     #[track_caller]
     fn assert_location_resolution(text_with_expected_range: &str, query: &str, rng: &mut StdRng) {
         let (text, expected_ranges) = marked_text_ranges(text_with_expected_range, false);

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

@@ -116,7 +116,6 @@ pub struct Edit {
     /// The exact text to find in the file. This will be matched using fuzzy matching
     /// to handle minor differences in whitespace or formatting.
     ///
-    /// Always include complete lines. Do not start or end mid-line.
     /// Be minimal with replacements:
     /// - For unique lines, include only those lines
     /// - For non-unique lines, include enough context to identify them
@@ -3916,6 +3915,58 @@ mod tests {
         assert_eq!(new_text, "new_content");
     }
 
+    #[gpui::test]
+    async fn test_streaming_edit_partial_last_line(cx: &mut TestAppContext) {
+        let file_content = indoc::indoc! {r#"
+            fn on_query_change(&mut self, cx: &mut Context<Self>) {
+                self.filter(cx);
+            }
+
+
+
+            fn render_search(&self, cx: &mut Context<Self>) -> Div {
+                div()
+            }
+        "#}
+        .to_string();
+
+        let (tool, _project, _action_log, _fs, _thread) =
+            setup_test(cx, json!({"file.rs": file_content})).await;
+
+        // The model sends old_text with a PARTIAL last line.
+        let old_text = "}\n\n\n\nfn render_search";
+        let new_text = "}\n\nfn render_search";
+
+        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": "Remove extra blank lines",
+            "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");
+        };
+
+        // The edit should reduce 3 blank lines to 1 blank line before
+        // fn render_search, without duplicating the function signature.
+        let expected = file_content.replace("}\n\n\n\nfn render_search", "}\n\nfn render_search");
+        pretty_assertions::assert_eq!(
+            final_text,
+            expected,
+            "Edit should only remove blank lines before render_search"
+        );
+    }
+
     #[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;