ep: Fix applying patch to text without a trailing newline (#46471)

Oleksiy Syvokon and Agus Zubiaga created 16 hours ago

Release Notes:

- N/A

Co-authored-by: Agus Zubiaga <[e-mail]>

Change summary

crates/edit_prediction/src/zeta2.rs             |  7 ------
crates/edit_prediction_cli/src/format_prompt.rs | 22 +++++++++++++++---
2 files changed, 18 insertions(+), 11 deletions(-)

Detailed changes

crates/edit_prediction/src/zeta2.rs

@@ -277,10 +277,3 @@ pub(crate) fn edit_prediction_accepted(
     })
     .detach_and_log_err(cx);
 }
-
-#[cfg(feature = "cli-support")]
-pub fn zeta2_output_for_patch(input: &zeta_prompt::ZetaPromptInput, patch: &str) -> Result<String> {
-    let old_editable_region = &input.cursor_excerpt[input.editable_range_in_excerpt.clone()];
-    let new_editable_region = crate::udiff::apply_diff_to_string(patch, old_editable_region)?;
-    Ok(new_editable_region)
-}

crates/edit_prediction_cli/src/format_prompt.rs

@@ -7,11 +7,9 @@ use crate::{
     retrieve_context::run_context_retrieval,
 };
 use anyhow::{Context as _, Result};
-use edit_prediction::{
-    EditPredictionStore,
-    zeta2::{zeta2_output_for_patch, zeta2_prompt_input},
-};
+use edit_prediction::{EditPredictionStore, zeta2::zeta2_prompt_input};
 use gpui::{AsyncApp, Entity};
+use similar::DiffableStr;
 use std::fmt::Write as _;
 use std::sync::Arc;
 use zeta_prompt::format_zeta_prompt;
@@ -98,6 +96,22 @@ pub async fn run_format_prompt(
     Ok(())
 }
 
+pub fn zeta2_output_for_patch(input: &zeta_prompt::ZetaPromptInput, patch: &str) -> Result<String> {
+    let mut old_editable_region =
+        input.cursor_excerpt[input.editable_range_in_excerpt.clone()].to_string();
+
+    if !old_editable_region.ends_with_newline() {
+        old_editable_region.push('\n');
+    }
+
+    edit_prediction::udiff::apply_diff_to_string(patch, &old_editable_region).with_context(|| {
+        format!(
+            "Patch:\n```\n{}```\n\nEditable region:\n```\n{}```",
+            patch, old_editable_region
+        )
+    })
+}
+
 pub struct TeacherPrompt;
 
 impl TeacherPrompt {