Fix issues with applying editor history

Max Brunsfeld , Piotr Osiewicz , Ben Kunkle , and Agus Zubiaga created 16 hours ago

Co-authored-by: Piotr Osiewicz <[e-mail]>
Co-authored-by: Ben Kunkle <[e-mail]>
Co-authored-by: Agus Zubiaga <[e-mail]>

Change summary

crates/zeta2/src/zeta2.rs      |  5 +++++
crates/zeta_cli/src/example.rs | 27 ++++++++++++++++-----------
crates/zeta_cli/src/main.rs    |  7 ++++---
3 files changed, 25 insertions(+), 14 deletions(-)

Detailed changes

crates/zeta2/src/zeta2.rs

@@ -135,6 +135,7 @@ impl ContextMode {
     }
 }
 
+#[derive(Debug)]
 pub enum ZetaDebugInfo {
     ContextRetrievalStarted(ZetaContextRetrievalStartedDebugInfo),
     SearchQueriesGenerated(ZetaSearchQueryDebugInfo),
@@ -144,17 +145,20 @@ pub enum ZetaDebugInfo {
     EditPredicted(ZetaEditPredictionDebugInfo),
 }
 
+#[derive(Debug)]
 pub struct ZetaContextRetrievalStartedDebugInfo {
     pub project: Entity<Project>,
     pub timestamp: Instant,
     pub search_prompt: String,
 }
 
+#[derive(Debug)]
 pub struct ZetaContextRetrievalDebugInfo {
     pub project: Entity<Project>,
     pub timestamp: Instant,
 }
 
+#[derive(Debug)]
 pub struct ZetaEditPredictionDebugInfo {
     pub request: predict_edits_v3::PredictEditsRequest,
     pub retrieval_time: TimeDelta,
@@ -164,6 +168,7 @@ pub struct ZetaEditPredictionDebugInfo {
     pub response_rx: oneshot::Receiver<Result<predict_edits_v3::PredictEditsResponse, String>>,
 }
 
+#[derive(Debug)]
 pub struct ZetaSearchQueryDebugInfo {
     pub project: Entity<Project>,
     pub timestamp: Instant,

crates/zeta_cli/src/example.rs

@@ -13,6 +13,7 @@ use clap::ValueEnum;
 use collections::HashSet;
 use futures::AsyncWriteExt as _;
 use gpui::{AsyncApp, Entity, http_client::Url};
+use language::Buffer;
 use project::{Project, ProjectPath};
 use pulldown_cmark::CowStr;
 use serde::{Deserialize, Serialize};
@@ -331,15 +332,16 @@ impl NamedExample {
         }
     }
 
+    #[must_use]
     pub async fn apply_edit_history(
         &self,
         project: &Entity<Project>,
         cx: &mut AsyncApp,
-    ) -> Result<()> {
+    ) -> Result<HashSet<Entity<Buffer>>> {
         use cloud_llm_client::udiff::DiffLine;
         use std::fmt::Write;
 
-        #[derive(Default)]
+        #[derive(Debug, Default)]
         struct Edit {
             context: String,
             deletion_start: Option<usize>,
@@ -359,7 +361,13 @@ impl NamedExample {
 
         while let Some(diff_line) = diff_lines.next() {
             match diff_line {
-                DiffLine::OldPath { path } => old_path = Some(path),
+                DiffLine::OldPath { path } => {
+                    mem::take(&mut pending);
+                    old_path = Some(path)
+                }
+                DiffLine::HunkHeader(_) => {
+                    mem::take(&mut pending);
+                }
                 DiffLine::NewPath { path } => {
                     if old_path.is_none() {
                         anyhow::bail!(
@@ -382,7 +390,7 @@ impl NamedExample {
 
                     writeln!(&mut pending.addition, "{add}")?;
                 }
-                DiffLine::HunkHeader(_) | DiffLine::Garbage => {}
+                DiffLine::Garbage => {}
             }
 
             let commit_pending = match diff_lines.peek() {
@@ -396,7 +404,7 @@ impl NamedExample {
                 Some(DiffLine::Deletion(_)) => {
                     // start a new cluster if we have any additions specifically
                     // if we only have deletions, we continue to aggregate them
-                    pending.addition.is_empty()
+                    !pending.addition.is_empty()
                 }
                 _ => false,
             };
@@ -404,10 +412,6 @@ impl NamedExample {
             if commit_pending {
                 let edit = mem::take(&mut pending);
 
-                if edit.addition.is_empty() || edit.deletion_start.is_none() {
-                    return anyhow::Ok(());
-                }
-
                 let Some(old_path) = old_path.as_deref() else {
                     anyhow::bail!("Missing old path (`---`) header")
                 };
@@ -444,6 +448,7 @@ impl NamedExample {
                 // TODO is it worth using project search?
                 buffer.update(cx, |buffer, cx| {
                     let text = buffer.text();
+                    // todo! check there's only one
                     if let Some(context_offset) = text.find(&edit.context) {
                         let end = context_offset + edit.context.len();
                         let start = if let Some(deletion_start) = edit.deletion_start {
@@ -456,13 +461,13 @@ impl NamedExample {
 
                         anyhow::Ok(())
                     } else {
-                        anyhow::bail!("Failed to match context");
+                        anyhow::bail!("Failed to match context:\n{}", edit.context);
                     }
                 })??;
             }
         }
 
-        anyhow::Ok(())
+        anyhow::Ok(open_buffers)
     }
 }
 

crates/zeta_cli/src/main.rs

@@ -372,6 +372,8 @@ async fn zeta2_predict(
         })?
         .await;
 
+    let _edited_buffers = example.apply_edit_history(&project, cx).await?;
+
     let cursor_path = RelPath::new(&example.example.cursor_path, PathStyle::Posix)?.into_arc();
 
     let cursor_buffer = project
@@ -420,15 +422,12 @@ async fn zeta2_predict(
         zeta.register_buffer(&cursor_buffer, &project, cx);
     })?;
 
-    example.apply_edit_history(&project, cx).await?;
-
     let (prediction_task, mut debug_rx) = zeta.update(cx, |zeta, cx| {
         let receiver = zeta.debug_info();
         let prediction_task = zeta.request_prediction(&project, &cursor_buffer, cursor_anchor, cx);
         (prediction_task, receiver)
     })?;
 
-    prediction_task.await.context("No prediction")?;
     let mut response = None;
 
     let mut excerpts_text = String::new();
@@ -456,6 +455,8 @@ async fn zeta2_predict(
         }
     }
 
+    prediction_task.await.context("No prediction")?;
+
     println!("## Excerpts\n");
     println!("{excerpts_text}");