edit predictions: Fix predictions bar disappearing while loading (#24582)

Release Notes:

- N/A

---------

Co-authored-by: Max <max@zed.dev>

crates/editor/src/editor.rs

@@ -190,6 +190,7 @@ pub const CODE_ACTIONS_DEBOUNCE_TIMEOUT: Duration = Duration::from_millis(250);
 pub(crate) const FORMAT_TIMEOUT: Duration = Duration::from_secs(2);
 pub(crate) const SCROLL_CENTER_TOP_BOTTOM_DEBOUNCE_TIMEOUT: Duration = Duration::from_secs(1);
 
+pub(crate) const EDIT_PREDICTION_KEY_CONTEXT: &str = "edit_prediction";
 pub(crate) const EDIT_PREDICTION_REQUIRES_MODIFIER_KEY_CONTEXT: &str =
     "edit_prediction_requires_modifier";
 
@@ -1488,13 +1489,13 @@ impl Editor {
         this
     }
 
-    pub fn mouse_menu_is_focused(&self, window: &mut Window, cx: &mut App) -> bool {
+    pub fn mouse_menu_is_focused(&self, window: &Window, cx: &App) -> bool {
         self.mouse_context_menu
             .as_ref()
             .is_some_and(|menu| menu.context_menu.focus_handle(cx).is_focused(window))
     }
 
-    fn key_context(&self, window: &mut Window, cx: &mut Context<Self>) -> KeyContext {
+    fn key_context(&self, window: &Window, cx: &App) -> KeyContext {
         let mut key_context = KeyContext::new_with_defaults();
         key_context.add("Editor");
         let mode = match self.mode {
@@ -1547,7 +1548,7 @@ impl Editor {
 
         if self.has_active_inline_completion() {
             key_context.add("copilot_suggestion");
-            key_context.add("edit_prediction");
+            key_context.add(EDIT_PREDICTION_KEY_CONTEXT);
 
             if showing_completions || self.edit_prediction_requires_modifier() {
                 key_context.add(EDIT_PREDICTION_REQUIRES_MODIFIER_KEY_CONTEXT);
@@ -1561,6 +1562,22 @@ impl Editor {
         key_context
     }
 
+    pub fn accept_edit_prediction_keybind(
+        &self,
+        window: &Window,
+        cx: &App,
+    ) -> AcceptEditPredictionBinding {
+        let mut context = self.key_context(window, cx);
+        context.add(EDIT_PREDICTION_KEY_CONTEXT);
+
+        AcceptEditPredictionBinding(
+            window
+                .bindings_for_action_in_context(&AcceptEditPrediction, context)
+                .into_iter()
+                .next(),
+        )
+    }
+
     pub fn new_file(
         workspace: &mut Workspace,
         _: &workspace::NewFile,
@@ -5128,8 +5145,7 @@ impl Editor {
         cx: &mut Context<Self>,
     ) {
         if self.show_edit_predictions_in_menu() {
-            let accept_binding =
+            let accept_binding = self.accept_edit_prediction_keybind(window, cx);
-                AcceptEditPredictionBinding::resolve(self.focus_handle(cx), window);
             if let Some(accept_keystroke) = accept_binding.keystroke() {
                 let was_previewing_inline_completion = self.previewing_inline_completion;
                 self.previewing_inline_completion = modifiers == accept_keystroke.modifiers
@@ -14408,7 +14424,8 @@ impl Editor {
         });
         supports
     }
-    pub fn is_focused(&self, window: &mut Window) -> bool {
+
+    pub fn is_focused(&self, window: &Window) -> bool {
         self.focus_handle.is_focused(window)
     }
 

crates/editor/src/element.rs

@@ -34,12 +34,12 @@ use gpui::{
     anchored, deferred, div, fill, linear_color_stop, linear_gradient, outline, point, px, quad,
     relative, size, svg, transparent_black, Action, AnyElement, App, AvailableSpace, Axis, Bounds,
     ClickEvent, ClipboardItem, ContentMask, Context, Corner, Corners, CursorStyle, DispatchPhase,
-    Edges, Element, ElementInputHandler, Entity, FocusHandle, Focusable as _, FontId,
+    Edges, Element, ElementInputHandler, Entity, Focusable as _, FontId, GlobalElementId, Hitbox,
-    GlobalElementId, Hitbox, Hsla, InteractiveElement, IntoElement, KeyBindingContextPredicate,
+    Hsla, InteractiveElement, IntoElement, KeyBindingContextPredicate, Keystroke, Length,
-    Keystroke, Length, ModifiersChangedEvent, MouseButton, MouseDownEvent, MouseMoveEvent,
+    ModifiersChangedEvent, MouseButton, MouseDownEvent, MouseMoveEvent, MouseUpEvent, PaintQuad,
-    MouseUpEvent, PaintQuad, ParentElement, Pixels, ScrollDelta, ScrollWheelEvent, ShapedLine,
+    ParentElement, Pixels, ScrollDelta, ScrollWheelEvent, ShapedLine, SharedString, Size,
-    SharedString, Size, StatefulInteractiveElement, Style, Styled, Subscription, TextRun,
+    StatefulInteractiveElement, Style, Styled, Subscription, TextRun, TextStyleRefinement,
-    TextStyleRefinement, WeakEntity, Window,
+    WeakEntity, Window,
 };
 use itertools::Itertools;
 use language::{
@@ -3167,10 +3167,8 @@ impl EditorElement {
                 );
 
                 let edit_prediction = if edit_prediction_popover_visible {
-                    let accept_binding =
-                        AcceptEditPredictionBinding::resolve(self.editor.focus_handle(cx), window);
-
                     self.editor.update(cx, move |editor, cx| {
+                        let accept_binding = editor.accept_edit_prediction_keybind(window, cx);
                         let mut element = editor.render_edit_prediction_cursor_popover(
                             min_width,
                             max_width,
@@ -3569,7 +3567,7 @@ impl EditorElement {
                         "Jump to Edit",
                         Some(IconName::ArrowUp),
                         previewing,
-                        self.editor.focus_handle(cx),
+                        editor,
                         window,
                         cx,
                     )?;
@@ -3582,7 +3580,7 @@ impl EditorElement {
                         "Jump to Edit",
                         Some(IconName::ArrowDown),
                         previewing,
-                        self.editor.focus_handle(cx),
+                        editor,
                         window,
                         cx,
                     )?;
@@ -3598,7 +3596,7 @@ impl EditorElement {
                         "Jump to Edit",
                         None,
                         previewing,
-                        self.editor.focus_handle(cx),
+                        editor,
                         window,
                         cx,
                     )?;
@@ -3657,10 +3655,16 @@ impl EditorElement {
                             target_display_point.row(),
                             editor_snapshot.line_len(target_display_point.row()),
                         );
-                        let (previewing_inline_completion, origin) =
+                        let (mut element, origin) = self.editor.update(cx, |editor, cx| {
-                            self.editor.update(cx, |editor, _cx| {
                             Some((
+                                inline_completion_accept_indicator(
+                                    "Accept",
+                                    None,
                                     editor.previewing_inline_completion,
+                                    editor,
+                                    window,
+                                    cx,
+                                )?,
                                 editor.display_to_pixel_point(
                                     target_line_end,
                                     editor_snapshot,
@@ -3669,15 +3673,6 @@ impl EditorElement {
                             ))
                         })?;
 
-                        let mut element = inline_completion_accept_indicator(
-                            "Accept",
-                            None,
-                            previewing_inline_completion,
-                            self.editor.focus_handle(cx),
-                            window,
-                            cx,
-                        )?;
-
                         element.prepaint_as_root(
                             text_bounds.origin + origin + point(PADDING_X, px(0.)),
                             AvailableSpace::min_size(),
@@ -5675,11 +5670,11 @@ fn inline_completion_accept_indicator(
     label: impl Into<SharedString>,
     icon: Option<IconName>,
     previewing: bool,
-    editor_focus_handle: FocusHandle,
+    editor: &Editor,
-    window: &Window,
+    window: &mut Window,
     cx: &App,
 ) -> Option<AnyElement> {
-    let accept_binding = AcceptEditPredictionBinding::resolve(editor_focus_handle, window);
+    let accept_binding = editor.accept_edit_prediction_keybind(window, cx);
     let accept_keystroke = accept_binding.keystroke()?;
 
     let accept_key = h_flex()
@@ -5728,18 +5723,9 @@ fn inline_completion_accept_indicator(
     )
 }
 
-pub struct AcceptEditPredictionBinding(Option<gpui::KeyBinding>);
+pub struct AcceptEditPredictionBinding(pub(crate) Option<gpui::KeyBinding>);
 
 impl AcceptEditPredictionBinding {
-    pub fn resolve(editor_focus_handle: FocusHandle, window: &Window) -> Self {
-        AcceptEditPredictionBinding(
-            window
-                .bindings_for_action_in(&AcceptEditPrediction, &editor_focus_handle)
-                .into_iter()
-                .next(),
-        )
-    }
-
     pub fn keystroke(&self) -> Option<&Keystroke> {
         if let Some(binding) = self.0.as_ref() {
             match &binding.keystrokes() {

crates/gpui/src/window.rs

@@ -3671,6 +3671,18 @@ impl Window {
         dispatch_tree.bindings_for_action(action, &context_stack)
     }
 
+    /// Returns the key bindings for the given action in the given context.
+    pub fn bindings_for_action_in_context(
+        &self,
+        action: &dyn Action,
+        context: KeyContext,
+    ) -> Vec<KeyBinding> {
+        let dispatch_tree = &self.rendered_frame.dispatch_tree;
+        dispatch_tree.bindings_for_action(action, &[context])
+    }
+
+    /// Returns a generic event listener that invokes the given listener with the view and context associated with the given view handle.
+
     /// Returns a generic event listener that invokes the given listener with the view and context associated with the given view handle.
     pub fn listener_for<V: Render, E>(
         &self,