ep: Apply diff to editable region only and edit history fixes (#44737)

Oleksiy Syvokon , Max Brunsfeld , and Agus Zubiaga created 2 weeks ago

Release Notes:

- N/A

---------

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

Change summary

crates/edit_prediction/src/edit_prediction.rs   |  3 +
crates/edit_prediction/src/zeta2.rs             | 21 ++++++++++--------
crates/edit_prediction_cli/src/format_prompt.rs |  6 ++--
3 files changed, 17 insertions(+), 13 deletions(-)

Detailed changes

crates/edit_prediction/src/edit_prediction.rs

@@ -586,10 +586,11 @@ impl EditPredictionStore {
     pub fn edit_history_for_project(
         &self,
         project: &Entity<Project>,
+        cx: &App,
     ) -> Vec<Arc<zeta_prompt::Event>> {
         self.projects
             .get(&project.entity_id())
-            .map(|project_state| project_state.events.iter().cloned().collect())
+            .map(|project_state| project_state.events(cx))
             .unwrap_or_default()
     }
 

crates/edit_prediction/src/zeta2.rs

@@ -228,13 +228,16 @@ pub fn zeta2_prompt_input(
 }
 
 #[cfg(feature = "cli-support")]
-pub fn zeta2_output_for_patch(input: &zeta_prompt::ZetaPromptInput, patch: &str) -> String {
-    eprintln!("{}", patch);
-    eprintln!("---------------------");
-    eprintln!("{}", input.cursor_excerpt);
-    crate::udiff::apply_diff_to_string(
-        patch,
-        &input.cursor_excerpt[input.editable_range_in_excerpt.clone()],
-    )
-    .unwrap()
+pub fn zeta2_output_for_patch(input: &zeta_prompt::ZetaPromptInput, patch: &str) -> Result<String> {
+    let text = &input.cursor_excerpt;
+    let editable_region = input.editable_range_in_excerpt.clone();
+    let old_prefix = &text[..editable_region.start];
+    let old_suffix = &text[editable_region.end..];
+
+    let new = crate::udiff::apply_diff_to_string(patch, text)?;
+    if !new.starts_with(old_prefix) || !new.ends_with(old_suffix) {
+        anyhow::bail!("Patch shouldn't affect text outside of editable region");
+    }
+
+    Ok(new[editable_region.start..new.len() - old_suffix.len()].to_string())
 }

crates/edit_prediction_cli/src/format_prompt.rs

@@ -44,7 +44,7 @@ pub async fn run_format_prompt(
             let state = example.state.as_ref().context("state must be set")?;
             let snapshot = state.buffer.read_with(&cx, |buffer, _| buffer.snapshot())?;
             let project = state.project.clone();
-            let (_, input) = ep_store.update(&mut cx, |ep_store, _cx| {
+            let (_, input) = ep_store.update(&mut cx, |ep_store, cx| {
                 anyhow::Ok(zeta2_prompt_input(
                     &snapshot,
                     example
@@ -53,7 +53,7 @@ pub async fn run_format_prompt(
                         .context("context must be set")?
                         .files
                         .clone(),
-                    ep_store.edit_history_for_project(&project),
+                    ep_store.edit_history_for_project(&project, cx),
                     example.cursor_path.clone(),
                     example
                         .buffer
@@ -63,7 +63,7 @@ pub async fn run_format_prompt(
                 ))
             })??;
             let prompt = format_zeta_prompt(&input);
-            let expected_output = zeta2_output_for_patch(&input, &example.expected_patch.clone());
+            let expected_output = zeta2_output_for_patch(&input, &example.expected_patch.clone())?;
             example.prompt = Some(ExamplePrompt {
                 input: prompt,
                 expected_output,