Polish edit predictions (#24732)

Release Notes:

- N/A

---------

Co-authored-by: Antonio Scandurra <me@as-cii.com>
Co-authored-by: as-cii <as-cii@zed.dev>
Co-authored-by: Danilo Leal <daniloleal09@gmail.com>

assets/settings/default.json

@@ -792,11 +792,11 @@
     ],
     // When to show edit predictions previews in buffer.
     // This setting takes two possible values:
-    // 1. Display inline when holding modifier key (alt by default).
+    // 1. Display inline when there are no language server completions available.
-    //     "mode": "auto"
-    // 2. Display inline when there are no language server completions available.
     //     "mode": "eager_preview"
-    "mode": "auto"
+    // 2. Display inline when holding modifier key (alt by default).
+    //     "mode": "auto"
+    "mode": "eager_preview"
   },
   // Settings specific to journaling
   "journal": {

crates/editor/src/editor.rs

@@ -5648,6 +5648,77 @@ impl Editor {
         }
     }
 
+    fn render_edit_prediction_accept_keybind(&self, window: &mut Window, cx: &App) -> Option<Div> {
+        let accept_binding = self.accept_edit_prediction_keybind(window, cx);
+        let accept_keystroke = accept_binding.keystroke()?;
+        let colors = cx.theme().colors();
+        let accent_color = colors.text_accent;
+        let editor_bg_color = colors.editor_background;
+        let bg_color = editor_bg_color.blend(accent_color.opacity(0.1));
+
+        h_flex()
+            .px_0p5()
+            .gap_1()
+            .bg(bg_color)
+            .font(theme::ThemeSettings::get_global(cx).buffer_font.clone())
+            .text_size(TextSize::XSmall.rems(cx))
+            .when(!self.edit_prediction_preview_is_active(), |parent| {
+                parent.children(ui::render_modifiers(
+                    &accept_keystroke.modifiers,
+                    PlatformStyle::platform(),
+                    Some(if accept_keystroke.modifiers == window.modifiers() {
+                        Color::Accent
+                    } else {
+                        Color::Muted
+                    }),
+                    Some(IconSize::XSmall.rems().into()),
+                    false,
+                ))
+            })
+            .child(accept_keystroke.key.clone())
+            .into()
+    }
+
+    fn render_edit_prediction_line_popover(
+        &self,
+        label: impl Into<SharedString>,
+        icon: Option<IconName>,
+        window: &mut Window,
+        cx: &App,
+    ) -> Option<Div> {
+        let bg_color = Self::edit_prediction_line_popover_bg_color(cx);
+
+        let padding_right = if icon.is_some() { px(4.) } else { px(8.) };
+
+        let result = h_flex()
+            .gap_1()
+            .border_1()
+            .rounded_lg()
+            .shadow_sm()
+            .bg(bg_color)
+            .border_color(cx.theme().colors().text_accent.opacity(0.4))
+            .py_0p5()
+            .pl_1()
+            .pr(padding_right)
+            .children(self.render_edit_prediction_accept_keybind(window, cx))
+            .child(Label::new(label).size(LabelSize::Small))
+            .when_some(icon, |element, icon| {
+                element.child(
+                    div()
+                        .mt(px(1.5))
+                        .child(Icon::new(icon).size(IconSize::Small)),
+                )
+            });
+
+        Some(result)
+    }
+
+    fn edit_prediction_line_popover_bg_color(cx: &App) -> Hsla {
+        let accent_color = cx.theme().colors().text_accent;
+        let editor_bg_color = cx.theme().colors().editor_background;
+        editor_bg_color.blend(accent_color.opacity(0.1))
+    }
+
     #[allow(clippy::too_many_arguments)]
     fn render_edit_prediction_cursor_popover(
         &self,
@@ -5788,18 +5859,26 @@ impl Editor {
                 .min_w(min_width)
                 .max_w(max_width)
                 .flex_1()
-                .px_2()
                 .elevation_2(cx)
                 .border_color(cx.theme().colors().border)
-                .child(div().py_1().overflow_hidden().child(completion))
+                .child(
+                    div()
+                        .flex_1()
+                        .py_1()
+                        .px_2()
+                        .overflow_hidden()
+                        .child(completion),
+                )
                 .child(
                     h_flex()
                         .h_full()
                         .border_l_1()
+                        .rounded_r_lg()
                         .border_color(cx.theme().colors().border)
+                        .bg(Self::edit_prediction_line_popover_bg_color(cx))
                         .gap_1()
                         .py_1()
-                        .pl_2()
+                        .px_2()
                         .child(
                             h_flex()
                                 .font(theme::ThemeSettings::get_global(cx).buffer_font.clone())
@@ -14548,6 +14627,7 @@ impl Editor {
         }
 
         self.hide_context_menu(window, cx);
+        self.discard_inline_completion(false, cx);
         cx.emit(EditorEvent::Blurred);
         cx.notify();
     }

crates/editor/src/element.rs

@@ -3583,14 +3583,14 @@ impl EditorElement {
                     }
 
                     if target_display_point.row() < visible_row_range.start {
-                        let mut element = inline_completion_accept_indicator(
+                        let mut element = editor
-                            "Scroll",
+                            .render_edit_prediction_line_popover(
-                            Some(IconName::ArrowUp),
+                                "Scroll",
-                            editor,
+                                Some(IconName::ArrowUp),
-                            window,
+                                window,
-                            cx,
+                                cx,
-                        )?
+                            )?
-                        .into_any();
+                            .into_any();
 
                         element.layout_as_root(AvailableSpace::min_size(), window, cx);
 
@@ -3608,14 +3608,14 @@ impl EditorElement {
                         element.prepaint_at(origin, window, cx);
                         return Some(element);
                     } else if target_display_point.row() >= visible_row_range.end {
-                        let mut element = inline_completion_accept_indicator(
+                        let mut element = editor
-                            "Scroll",
+                            .render_edit_prediction_line_popover(
-                            Some(IconName::ArrowDown),
+                                "Scroll",
-                            editor,
+                                Some(IconName::ArrowDown),
-                            window,
+                                window,
-                            cx,
+                                cx,
-                        )?
+                            )?
-                        .into_any();
+                            .into_any();
 
                         let size = element.layout_as_root(AvailableSpace::min_size(), window, cx);
 
@@ -3640,12 +3640,11 @@ impl EditorElement {
 
                         let mut element = v_flex()
                             .child(
-                                inline_completion_accept_indicator(
+                                editor
-                                    "Jump", None, editor, window, cx,
+                                    .render_edit_prediction_line_popover("Jump", None, window, cx)?
-                                )?
+                                    .rounded_br(px(0.))
-                                .rounded_br(px(0.))
+                                    .rounded_tr(px(0.))
-                                .rounded_tr(px(0.))
+                                    .border_r_2(),
-                                .border_r_2(),
                             )
                             .child(
                                 div()
@@ -3680,28 +3679,30 @@ impl EditorElement {
                 }
 
                 if target_display_point.row().as_f32() < scroll_top {
-                    let mut element = inline_completion_accept_indicator(
+                    let mut element = editor
-                        "Jump to Edit",
+                        .render_edit_prediction_line_popover(
-                        Some(IconName::ArrowUp),
+                            "Jump to Edit",
-                        editor,
+                            Some(IconName::ArrowUp),
-                        window,
+                            window,
-                        cx,
+                            cx,
-                    )?
+                        )?
-                    .into_any();
+                        .into_any();
+
                     let size = element.layout_as_root(AvailableSpace::min_size(), window, cx);
                     let offset = point((text_bounds.size.width - size.width) / 2., PADDING_Y);
 
                     element.prepaint_at(text_bounds.origin + offset, window, cx);
                     Some(element)
                 } else if (target_display_point.row().as_f32() + 1.) > scroll_bottom {
-                    let mut element = inline_completion_accept_indicator(
+                    let mut element = editor
-                        "Jump to Edit",
+                        .render_edit_prediction_line_popover(
-                        Some(IconName::ArrowDown),
+                            "Jump to Edit",
-                        editor,
+                            Some(IconName::ArrowDown),
-                        window,
+                            window,
-                        cx,
+                            cx,
-                    )?
+                        )?
-                    .into_any();
+                        .into_any();
+
                     let size = element.layout_as_root(AvailableSpace::min_size(), window, cx);
                     let offset = point(
                         (text_bounds.size.width - size.width) / 2.,
@@ -3711,14 +3712,9 @@ impl EditorElement {
                     element.prepaint_at(text_bounds.origin + offset, window, cx);
                     Some(element)
                 } else {
-                    let mut element = inline_completion_accept_indicator(
+                    let mut element = editor
-                        "Jump to Edit",
+                        .render_edit_prediction_line_popover("Jump to Edit", None, window, cx)?
-                        None,
+                        .into_any();
-                        editor,
-                        window,
-                        cx,
-                    )?
-                    .into_any();
                     let target_line_end = DisplayPoint::new(
                         target_display_point.row(),
                         editor_snapshot.line_len(target_display_point.row()),
@@ -3776,10 +3772,11 @@ impl EditorElement {
                         );
                         let (mut element, origin) = self.editor.update(cx, |editor, cx| {
                             Some((
-                                inline_completion_accept_indicator(
+                                editor
-                                    "Accept", None, editor, window, cx,
+                                    .render_edit_prediction_line_popover(
-                                )?
+                                        "Accept", None, window, cx,
-                                .into_any(),
+                                    )?
+                                    .into_any(),
                                 editor.display_to_pixel_point(
                                     target_line_end,
                                     editor_snapshot,
@@ -3808,6 +3805,37 @@ impl EditorElement {
                     cx,
                 );
 
+                let styled_text = highlighted_edits.to_styled_text(&style.text);
+
+                const ACCEPT_INDICATOR_HEIGHT: Pixels = px(24.);
+
+                let mut element = v_flex()
+                    .items_end()
+                    .shadow_sm()
+                    .child(
+                        h_flex()
+                            .h(ACCEPT_INDICATOR_HEIGHT)
+                            .mb(px(-1.))
+                            .px_1p5()
+                            .gap_1()
+                            .bg(Editor::edit_prediction_line_popover_bg_color(cx))
+                            .border_1()
+                            .border_b_0()
+                            .border_color(cx.theme().colors().border)
+                            .rounded_t_lg()
+                            .children(editor.render_edit_prediction_accept_keybind(window, cx)),
+                    )
+                    .child(
+                        div()
+                            .bg(cx.theme().colors().editor_background)
+                            .border_1()
+                            .border_color(cx.theme().colors().border)
+                            .rounded_lg()
+                            .rounded_tr(Pixels::ZERO)
+                            .child(styled_text),
+                    )
+                    .into_any();
+
                 let line_count = highlighted_edits.text.lines().count();
 
                 let longest_row =
@@ -3827,16 +3855,6 @@ impl EditorElement {
                     .width
                 };
 
-                let styled_text = highlighted_edits.to_styled_text(&style.text);
-
-                let mut element = div()
-                    .bg(cx.theme().colors().editor_background)
-                    .border_1()
-                    .border_color(cx.theme().colors().border)
-                    .rounded_md()
-                    .child(styled_text)
-                    .into_any();
-
                 let viewport_bounds = Bounds::new(Default::default(), window.viewport_size())
                     .extend(Edges {
                         right: -Self::SCROLLBAR_WIDTH,
@@ -3853,7 +3871,7 @@ impl EditorElement {
                 let is_fully_visible = x_after_longest < text_bounds.right()
                     && x_after_longest + element_bounds.width < viewport_bounds.right();
 
-                let origin = if is_fully_visible {
+                let mut origin = if is_fully_visible {
                     point(
                         x_after_longest,
                         text_bounds.origin.y + edit_start.row().as_f32() * line_height
@@ -3898,6 +3916,8 @@ impl EditorElement {
                         )
                 };
 
+                origin.y -= ACCEPT_INDICATOR_HEIGHT;
+
                 window.defer_draw(element, origin, 1);
 
                 // Do not return an element, since it will already be drawn due to defer_draw.
@@ -5796,63 +5816,6 @@ fn header_jump_data(
     }
 }
 
-fn inline_completion_accept_indicator(
-    label: impl Into<SharedString>,
-    icon: Option<IconName>,
-    editor: &Editor,
-    window: &mut Window,
-    cx: &App,
-) -> Option<Div> {
-    let accept_binding = editor.accept_edit_prediction_keybind(window, cx);
-    let accept_keystroke = accept_binding.keystroke()?;
-
-    let accept_key = h_flex()
-        .px_0p5()
-        .font(theme::ThemeSettings::get_global(cx).buffer_font.clone())
-        .text_size(TextSize::XSmall.rems(cx))
-        .text_color(cx.theme().colors().text)
-        .gap_1()
-        .when(!editor.edit_prediction_preview_is_active(), |parent| {
-            parent.children(ui::render_modifiers(
-                &accept_keystroke.modifiers,
-                PlatformStyle::platform(),
-                Some(Color::Default),
-                None,
-                false,
-            ))
-        })
-        .child(accept_keystroke.key.clone());
-
-    let colors = cx.theme().colors();
-
-    let accent_color = colors.text_accent;
-    let editor_bg_color = colors.editor_background;
-    let bg_color = editor_bg_color.blend(accent_color.opacity(0.2));
-    let padding_right = if icon.is_some() { px(4.) } else { px(8.) };
-
-    let result = h_flex()
-        .gap_1()
-        .border_1()
-        .rounded_md()
-        .shadow_sm()
-        .bg(bg_color)
-        .border_color(colors.text_accent.opacity(0.8))
-        .py_0p5()
-        .pl_1()
-        .pr(padding_right)
-        .child(accept_key)
-        .child(Label::new(label).size(LabelSize::Small))
-        .when_some(icon, |element, icon| {
-            element.child(
-                div()
-                    .mt(px(1.5))
-                    .child(Icon::new(icon).size(IconSize::Small)),
-            )
-        });
-
-    Some(result)
-}
-
 pub struct AcceptEditPredictionBinding(pub(crate) Option<gpui::KeyBinding>);
 
 impl AcceptEditPredictionBinding {

crates/inline_completion_button/src/inline_completion_button.rs

@@ -573,37 +573,41 @@ impl InlineCompletionButton {
             language::EditPredictionsMode::Auto => false,
             language::EditPredictionsMode::EagerPreview => true,
         };
-        menu = menu.separator().toggleable_entry(
+        menu = if cx.is_staff() {
-            "Eager Preview Mode",
+            menu.separator().toggleable_entry(
-            is_eager_preview_enabled,
+                "Eager Preview Mode",
-            IconPosition::Start,
+                is_eager_preview_enabled,
-            None,
+                IconPosition::Start,
-            {
+                None,
-                let fs = fs.clone();
+                {
-                move |_window, cx| {
+                    let fs = fs.clone();
-                    update_settings_file::<AllLanguageSettings>(
+                    move |_window, cx| {
-                        fs.clone(),
+                        update_settings_file::<AllLanguageSettings>(
-                        cx,
+                            fs.clone(),
-                        move |settings, _cx| {
+                            cx,
-                            let new_mode = match is_eager_preview_enabled {
+                            move |settings, _cx| {
-                                true => language::EditPredictionsMode::Auto,
+                                let new_mode = match is_eager_preview_enabled {
-                                false => language::EditPredictionsMode::EagerPreview,
+                                    true => language::EditPredictionsMode::Auto,
-                            };
+                                    false => language::EditPredictionsMode::EagerPreview,
+                                };
 
-                            if let Some(edit_predictions) = settings.edit_predictions.as_mut() {
+                                if let Some(edit_predictions) = settings.edit_predictions.as_mut() {
-                                edit_predictions.mode = new_mode;
+                                    edit_predictions.mode = new_mode;
-                            } else {
+                                } else {
-                                settings.edit_predictions =
+                                    settings.edit_predictions =
-                                    Some(language_settings::EditPredictionSettingsContent {
+                                        Some(language_settings::EditPredictionSettingsContent {
-                                        mode: new_mode,
+                                            mode: new_mode,
-                                        ..Default::default()
+                                            ..Default::default()
-                                    });
+                                        });
-                            }
+                                }
-                        },
+                            },
-                    );
+                        );
-                }
+                    }
-            },
+                },
-        );
+            )
+        } else {
+            menu
+        };
 
         if let Some(editor_focus_handle) = self.editor_focus_handle.clone() {
             menu = menu

crates/language/src/language_settings.rs

@@ -237,9 +237,9 @@ pub struct EditPredictionSettings {
 pub enum EditPredictionsMode {
     /// If provider supports it, display inline when holding modifier key (e.g., alt).
     /// Otherwise, eager preview is used.
-    #[default]
     Auto,
     /// Display inline when there are no language server completions available.
+    #[default]
     EagerPreview,
 }
 

crates/zeta/src/zeta.rs

@@ -27,7 +27,10 @@ use gpui::{
 };
 use http_client::{HttpClient, Method};
 use input_excerpt::excerpt_for_cursor_position;
-use language::{Anchor, Buffer, BufferSnapshot, EditPreview, OffsetRangeExt, ToOffset, ToPoint};
+use language::{
+    Anchor, Buffer, BufferSnapshot, CharClassifier, CharKind, EditPreview, OffsetRangeExt,
+    ToOffset, ToPoint,
+};
 use language_models::LlmApiToken;
 use postage::watch;
 use project::Project;
@@ -57,9 +60,9 @@ const EDITABLE_REGION_END_MARKER: &'static str = "<|editable_region_end|>";
 const BUFFER_CHANGE_GROUPING_INTERVAL: Duration = Duration::from_secs(1);
 const ZED_PREDICT_DATA_COLLECTION_CHOICE: &str = "zed_predict_data_collection_choice";
 
-const MAX_CONTEXT_TOKENS: usize = 100;
+const MAX_CONTEXT_TOKENS: usize = 150;
-const MAX_REWRITE_TOKENS: usize = 300;
+const MAX_REWRITE_TOKENS: usize = 350;
-const MAX_EVENT_TOKENS: usize = 400;
+const MAX_EVENT_TOKENS: usize = 500;
 
 /// Maximum number of events to track.
 const MAX_EVENT_COUNT: usize = 16;
@@ -834,8 +837,34 @@ and then another
         offset: usize,
         snapshot: &BufferSnapshot,
     ) -> Vec<(Range<Anchor>, String)> {
-        let diff = similar::TextDiff::from_words(old_text.as_str(), new_text);
+        fn tokenize(text: &str) -> Vec<&str> {
+            let classifier = CharClassifier::new(None).for_completion(true);
+            let mut chars = text.chars().peekable();
+            let mut prev_ch = chars.peek().copied();
+            let mut tokens = Vec::new();
+            let mut start = 0;
+            let mut end = 0;
+            while let Some(ch) = chars.next() {
+                let prev_kind = prev_ch.map(|ch| classifier.kind(ch));
+                let kind = classifier.kind(ch);
+                if Some(kind) != prev_kind || (kind == CharKind::Punctuation && Some(ch) != prev_ch)
+                {
+                    tokens.push(&text[start..end]);
+                    start = end;
+                }
+                end += ch.len_utf8();
+                prev_ch = Some(ch);
+            }
+            tokens.push(&text[start..end]);
+            tokens
+        }
 
+        let old_tokens = tokenize(&old_text);
+        let new_tokens = tokenize(new_text);
+
+        let diff = similar::TextDiffConfig::default()
+            .algorithm(similar::Algorithm::Patience)
+            .diff_slices(&old_tokens, &new_tokens);
         let mut edits: Vec<(Range<usize>, String)> = Vec::new();
         let mut old_start = offset;
         for change in diff.iter_all_changes() {
@@ -1705,6 +1734,70 @@ mod tests {
         })
     }
 
+    #[gpui::test]
+    async fn test_clean_up_diff(cx: &mut TestAppContext) {
+        cx.update(|cx| {
+            let settings_store = SettingsStore::test(cx);
+            cx.set_global(settings_store);
+            client::init_settings(cx);
+        });
+
+        let edits = edits_for_prediction(
+            indoc! {"
+                fn main() {
+                    let word_1 = \"lorem\";
+                    let range = word.len()..word.len();
+                }
+            "},
+            indoc! {"
+                <|editable_region_start|>
+                fn main() {
+                    let word_1 = \"lorem\";
+                    let range = word_1.len()..word_1.len();
+                }
+
+                <|editable_region_end|>
+            "},
+            cx,
+        )
+        .await;
+        assert_eq!(
+            edits,
+            [
+                (Point::new(2, 20)..Point::new(2, 20), "_1".to_string()),
+                (Point::new(2, 32)..Point::new(2, 32), "_1".to_string()),
+            ]
+        );
+
+        let edits = edits_for_prediction(
+            indoc! {"
+                fn main() {
+                    let story = \"the quick\"
+                }
+            "},
+            indoc! {"
+                <|editable_region_start|>
+                fn main() {
+                    let story = \"the quick brown fox jumps over the lazy dog\";
+                }
+
+                <|editable_region_end|>
+            "},
+            cx,
+        )
+        .await;
+        assert_eq!(
+            edits,
+            [
+                (
+                    Point::new(1, 26)..Point::new(1, 26),
+                    " brown fox jumps over the lazy dog".to_string()
+                ),
+                (Point::new(1, 27)..Point::new(1, 27), ";".to_string()),
+            ]
+        );
+    }
+
     #[gpui::test]
     async fn test_inline_completion_end_of_buffer(cx: &mut TestAppContext) {
         cx.update(|cx| {
@@ -1768,6 +1861,58 @@ mod tests {
         );
     }
 
+    async fn edits_for_prediction(
+        buffer_content: &str,
+        completion_response: &str,
+        cx: &mut TestAppContext,
+    ) -> Vec<(Range<Point>, String)> {
+        let completion_response = completion_response.to_string();
+        let http_client = FakeHttpClient::create(move |_| {
+            let completion = completion_response.clone();
+            async move {
+                Ok(http_client::Response::builder()
+                    .status(200)
+                    .body(
+                        serde_json::to_string(&PredictEditsResponse {
+                            request_id: Uuid::new_v4(),
+                            output_excerpt: completion,
+                        })
+                        .unwrap()
+                        .into(),
+                    )
+                    .unwrap())
+            }
+        });
+
+        let client = cx.update(|cx| Client::new(Arc::new(FakeSystemClock::new()), http_client, cx));
+        cx.update(|cx| {
+            RefreshLlmTokenListener::register(client.clone(), cx);
+        });
+        let server = FakeServer::for_client(42, &client, cx).await;
+        let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
+        let zeta = cx.new(|cx| Zeta::new(client, user_store, cx));
+
+        let buffer = cx.new(|cx| Buffer::local(buffer_content, cx));
+        let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot());
+        let cursor = buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(1, 0)));
+        let completion_task = zeta.update(cx, |zeta, cx| {
+            zeta.request_completion(None, &buffer, cursor, false, cx)
+        });
+
+        let token_request = server.receive::<proto::GetLlmToken>().await.unwrap();
+        server.respond(
+            token_request.receipt(),
+            proto::GetLlmTokenResponse { token: "".into() },
+        );
+
+        let completion = completion_task.await.unwrap().unwrap();
+        completion
+            .edits
+            .into_iter()
+            .map(|(old_range, new_text)| (old_range.to_point(&snapshot), new_text.clone()))
+            .collect::<Vec<_>>()
+    }
+
     fn to_completion_edits(
         iterator: impl IntoIterator<Item = (Range<usize>, String)>,
         buffer: &Entity<Buffer>,