From 4353b61155dbb7e0ea3513278b0e61dfcefa6b28 Mon Sep 17 00:00:00 2001 From: Agus Zubiaga Date: Wed, 24 Sep 2025 10:10:52 -0300 Subject: [PATCH] zeta2: Compute smaller edits (#38786) The new cloud endpoint returns structured edits, but they may include more of the input excerpt than what we want to display in the preview, so we compute a smaller diff on the client side against the snapshot. Release Notes: - N/A --- Cargo.lock | 1 + crates/zeta2/Cargo.toml | 1 + crates/zeta2/src/zeta2.rs | 139 ++++++++++++++++++++++++++++++-------- 3 files changed, 114 insertions(+), 27 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 793c092e71a9e3717841832452e8758fde24064f..46de7e804b2b13ff775ac87953bb7c9593c39120 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -20587,6 +20587,7 @@ dependencies = [ "edit_prediction_context", "futures 0.3.31", "gpui", + "indoc", "language", "language_model", "log", diff --git a/crates/zeta2/Cargo.toml b/crates/zeta2/Cargo.toml index 11ca5fcfddd3e7a7c1085401e7575bc599aea91e..8ef1c5a64f7b64a6af6a0ce984e2f76e14eb5e77 100644 --- a/crates/zeta2/Cargo.toml +++ b/crates/zeta2/Cargo.toml @@ -22,6 +22,7 @@ edit_prediction.workspace = true edit_prediction_context.workspace = true futures.workspace = true gpui.workspace = true +indoc.workspace = true language.workspace = true language_model.workspace = true log.workspace = true diff --git a/crates/zeta2/src/zeta2.rs b/crates/zeta2/src/zeta2.rs index 47afc797d5d1d4dafca71f39ba15fbe6fd621c20..8986fb20740327f5598611d981b53569edcb559e 100644 --- a/crates/zeta2/src/zeta2.rs +++ b/crates/zeta2/src/zeta2.rs @@ -21,11 +21,13 @@ use gpui::{ }; use language::{ Anchor, Buffer, DiagnosticSet, LanguageServerId, OffsetRangeExt as _, ToOffset as _, ToPoint, + text_diff, }; use language::{BufferSnapshot, EditPreview}; use language_model::{LlmApiToken, RefreshLlmTokenListener}; use project::Project; use release_channel::AppVersion; +use std::borrow::Cow; use std::cmp; use std::collections::{HashMap, VecDeque, hash_map}; use std::path::PathBuf; @@ -438,7 +440,10 @@ impl Zeta { .ok(); } - anyhow::Ok(Some(response?)) + let (response, usage) = response?; + let edits = Self::compute_edits(&response.edits, &snapshot); + + anyhow::Ok(Some((response.request_id, edits, usage))) } }); @@ -446,9 +451,7 @@ impl Zeta { cx.spawn(async move |this, cx| { match request_task.await { - Ok(Some((response, usage))) => { - log::debug!("predicted edits: {:?}", &response.edits); - + Ok(Some((id, edits, usage))) => { if let Some(usage) = usage { this.update(cx, |this, cx| { this.user_store.update(cx, |user_store, cx| { @@ -459,28 +462,6 @@ impl Zeta { } // TODO telemetry: duration, etc - - // TODO produce smaller edits by diffing against snapshot first - // - // Cloud returns entire snippets/excerpts ranges as they were included - // in the request, but we should display smaller edits to the user. - // - // We can do this by computing a diff of each one against the snapshot. - // Similar to zeta::Zeta::compute_edits, but per edit. - let edits = response - .edits - .into_iter() - .map(|edit| { - // TODO edits to different files - ( - snapshot.anchor_before(edit.range.start) - ..snapshot.anchor_before(edit.range.end), - edit.content, - ) - }) - .collect::>() - .into(); - let Some((edits, snapshot, edit_preview_task)) = buffer.read_with(cx, |buffer, cx| { let new_snapshot = buffer.snapshot(); @@ -493,7 +474,7 @@ impl Zeta { }; Ok(Some(EditPrediction { - id: EditPredictionId(response.request_id), + id: EditPredictionId(id), edits, snapshot, edit_preview: edit_preview_task.await, @@ -604,6 +585,62 @@ impl Zeta { } } + fn compute_edits( + edits: &[predict_edits_v3::Edit], + snapshot: &BufferSnapshot, + ) -> Arc<[(Range, String)]> { + edits + .iter() + .flat_map(|edit| { + // TODO multi-file edits + let old_text = snapshot.text_for_range(edit.range.clone()); + + Self::compute_excerpt_edits( + old_text.collect::>(), + &edit.content, + edit.range.start, + &snapshot, + ) + }) + .collect::>() + .into() + } + + fn compute_excerpt_edits( + old_text: Cow, + new_text: &str, + offset: usize, + snapshot: &BufferSnapshot, + ) -> impl Iterator, String)> { + text_diff(&old_text, new_text) + .into_iter() + .map(move |(mut old_range, new_text)| { + old_range.start += offset; + old_range.end += offset; + + let prefix_len = common_prefix( + snapshot.chars_for_range(old_range.clone()), + new_text.chars(), + ); + old_range.start += prefix_len; + + let suffix_len = common_prefix( + snapshot.reversed_chars_for_range(old_range.clone()), + new_text[prefix_len..].chars().rev(), + ); + old_range.end = old_range.end.saturating_sub(suffix_len); + + let new_text = new_text[prefix_len..new_text.len() - suffix_len].to_string(); + let range = if old_range.is_empty() { + let anchor = snapshot.anchor_after(old_range.start); + anchor..anchor + } else { + snapshot.anchor_after(old_range.start)..snapshot.anchor_before(old_range.end) + }; + (range, new_text) + }) + } + fn gather_nearby_diagnostics( cursor_offset: usize, diagnostic_sets: &[(LanguageServerId, DiagnosticSet)], @@ -713,6 +750,13 @@ impl Zeta { } } +fn common_prefix, T2: Iterator>(a: T1, b: T2) -> usize { + a.zip(b) + .take_while(|(a, b)| a == b) + .map(|(a, _)| a.len_utf8()) + .sum() +} + #[derive(Error, Debug)] #[error( "You must update to Zed version {minimum_version} or higher to continue using edit predictions." @@ -1222,6 +1266,47 @@ fn interpolate( mod tests { use super::*; use gpui::TestAppContext; + use indoc::indoc; + + #[gpui::test] + async fn test_compute_edits(cx: &mut TestAppContext) { + let old = indoc! {r#" + fn main() { + let args = + println!("{}", args[1]) + } + "#}; + + let new = indoc! {r#" + fn main() { + let args = std::env::args(); + println!("{}", args[1]); + } + "#}; + + let buffer = cx.new(|cx| Buffer::local(old, cx)); + let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot()); + + // TODO cover more cases when multi-file is supported + let big_edits = vec![predict_edits_v3::Edit { + path: PathBuf::from("test.txt"), + range: 0..old.len(), + content: new.into(), + }]; + + let edits = Zeta::compute_edits(&big_edits, &snapshot); + assert_eq!(edits.len(), 2); + assert_eq!( + edits[0].0.to_point(&snapshot).start, + language::Point::new(1, 14) + ); + assert_eq!(edits[0].1, " std::env::args();"); + assert_eq!( + edits[1].0.to_point(&snapshot).start, + language::Point::new(2, 27) + ); + assert_eq!(edits[1].1, ";"); + } #[gpui::test] async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) {