From b5fc09e40515cc5a97a7c91de8fb6e5bebb9c447 Mon Sep 17 00:00:00 2001
From: Ben Kunkle <ben@zed.dev>
Date: Tue, 24 Mar 2026 13:09:02 -0500
Subject: [PATCH] Don't keep stale predictions in if they were discarded
 (#52334)

## 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
---
 crates/editor/src/edit_prediction_tests.rs | 29 +++++++++++++++++++++-
 crates/editor/src/editor.rs                | 18 ++++++++++----
 2 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/crates/editor/src/edit_prediction_tests.rs b/crates/editor/src/edit_prediction_tests.rs
index 684213e481762d7fb09a0bd6d8b7a0b9fc6d4a36..52939a9e5a8fd1a35a3a3c0bcd2a04b893bd6628 100644
--- a/crates/editor/src/edit_prediction_tests.rs
+++ b/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)
diff --git a/crates/editor/src/editor.rs b/crates/editor/src/editor.rs
index 3fb832b3ac7b2ebb6aa3f47b2dcd590d7eb081df..a6ad7219b59195bad61e73b44090d549b3ad7bf5 100644
--- a/crates/editor/src/editor.rs
+++ b/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();
     }