Don't keep stale predictions in if they were discarded (#52334)

Ben Kunkle created 2 weeks ago

## Context

We're seeing issues where discarding a prediction (e.g. `editor::Cancel`
while a prediction is being shown) seems to just hide it temporarily.
This is becoming a larger issue with #51842 as we now expect people to
be dismissing predictions in order to insert literal tabs in certain
contexts.
The logic changed in this PR is part of the problem, but the model
generating the same prediction multiple times is also likely
contributing. That will be solved as-needed later.


## Self-Review Checklist

<!-- Check before requesting review: -->
- [x] I've reviewed my own diff for quality, security, and reliability
- [x] Unsafe blocks (if any) have justifying comments
- [x] The content is consistent with the [UI/UX
checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)
- [x] Tests cover the new/changed behavior
- [x] Performance impact has been considered and is acceptable

Release Notes:

- Fixed an issue where dismissing an edit prediction would not fully
discard it, causing it to re-appear

Change summary

crates/editor/src/edit_prediction_tests.rs | 29 +++++++++++++++++++++++
crates/editor/src/editor.rs                | 18 ++++++++++----
2 files changed, 41 insertions(+), 6 deletions(-)

Detailed changes

crates/editor/src/edit_prediction_tests.rs

@@ -954,7 +954,7 @@ async fn test_cursor_popover_edit_prediction_keybind_cases(cx: &mut gpui::TestAp
                 cx.update_editor(|editor, _window, cx| {
                     assert!(editor.active_edit_prediction.is_some());
                     assert!(editor.stale_edit_prediction_in_menu.is_none());
-                    editor.take_active_edit_prediction(cx);
+                    editor.take_active_edit_prediction(true, cx);
                     assert!(editor.active_edit_prediction.is_none());
                     assert!(editor.stale_edit_prediction_in_menu.is_some());
                 });
@@ -1054,6 +1054,33 @@ fn assert_editor_active_move_completion(
     })
 }
 
+#[gpui::test]
+async fn test_cancel_clears_stale_edit_prediction_in_menu(cx: &mut gpui::TestAppContext) {
+    init_test(cx, |_| {});
+    load_default_keymap(cx);
+
+    let mut cx = EditorTestContext::new(cx).await;
+    let provider = cx.new(|_| FakeEditPredictionDelegate::default());
+    assign_editor_completion_provider(provider.clone(), &mut cx);
+    cx.set_state("let x = ˇ;");
+
+    propose_edits(&provider, vec![(8..8, "42")], &mut cx);
+    cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
+
+    cx.update_editor(|editor, _window, _cx| {
+        assert!(editor.active_edit_prediction.is_some());
+        assert!(editor.stale_edit_prediction_in_menu.is_none());
+    });
+
+    cx.simulate_keystroke("escape");
+    cx.run_until_parked();
+
+    cx.update_editor(|editor, _window, _cx| {
+        assert!(editor.active_edit_prediction.is_none());
+        assert!(editor.stale_edit_prediction_in_menu.is_none());
+    });
+}
+
 fn accept_completion(cx: &mut EditorTestContext) {
     cx.update_editor(|editor, window, cx| {
         editor.accept_edit_prediction(&crate::AcceptEditPrediction, window, cx)

crates/editor/src/editor.rs

@@ -8427,7 +8427,7 @@ impl Editor {
             provider.discard(reason, cx);
         }
 
-        self.take_active_edit_prediction(cx)
+        self.take_active_edit_prediction(reason == EditPredictionDiscardReason::Ignored, cx)
     }
 
     fn report_edit_prediction_event(&self, id: Option<SharedString>, accepted: bool, cx: &App) {
@@ -8495,14 +8495,22 @@ impl Editor {
         self.active_edit_prediction.is_some()
     }
 
-    fn take_active_edit_prediction(&mut self, cx: &mut Context<Self>) -> bool {
+    fn take_active_edit_prediction(
+        &mut self,
+        preserve_stale_in_menu: bool,
+        cx: &mut Context<Self>,
+    ) -> bool {
         let Some(active_edit_prediction) = self.active_edit_prediction.take() else {
+            if !preserve_stale_in_menu {
+                self.stale_edit_prediction_in_menu = None;
+            }
             return false;
         };
 
         self.splice_inlays(&active_edit_prediction.inlay_ids, Default::default(), cx);
         self.clear_highlights(HighlightKey::EditPredictionHighlight, cx);
-        self.stale_edit_prediction_in_menu = Some(active_edit_prediction);
+        self.stale_edit_prediction_in_menu =
+            preserve_stale_in_menu.then_some(active_edit_prediction);
         true
     }
 
@@ -8715,7 +8723,7 @@ impl Editor {
             return None;
         }
 
-        self.take_active_edit_prediction(cx);
+        self.take_active_edit_prediction(true, cx);
         let Some(provider) = self.edit_prediction_provider() else {
             self.edit_prediction_settings = EditPredictionSettings::Disabled;
             return None;
@@ -25205,7 +25213,7 @@ impl Editor {
         {
             self.hide_context_menu(window, cx);
         }
-        self.take_active_edit_prediction(cx);
+        self.take_active_edit_prediction(true, cx);
         cx.emit(EditorEvent::Blurred);
         cx.notify();
     }