From 46b39f0077b4ad54421f1d48e59db5b1e5e02e29 Mon Sep 17 00:00:00 2001
From: Pranav Joglekar <pranav2000joglekar@gmail.com>
Date: Tue, 4 Nov 2025 21:00:37 +0530
Subject: [PATCH] editor: Clean up edit prediction preview when edit prediction
 is disabled (#40159)

Update the `editor::Editor.handle_modifiers_changed` method to ensure
that the `editor::Editor.update_edit_prediction_preview` method is
called even if edit prediction preview is disabled, if there's an active
edit prediction preview.

Without this change it was possible for users to get into a state where
holding the modifiers to show the prediction were part of the modifiers
used to disable edit prediction. When that keybinding was used, edit
prediction would be disabled, but the edit prediction preview would
remain as active, so the context menu for the editor would never be
shown again, as the editor would assume it was still showing the edit
prediction preview.

Closes #40056

Release Notes:

- Fixed a bug that could cause the completions menu to stop being show
when edit predictions were disabled

---------

Co-authored-by: dino <dinojoaocosta@gmail.com>
---
 crates/editor/src/edit_prediction_tests.rs | 64 +++++++++++++++++++++-
 crates/editor/src/editor.rs                |  9 ++-
 2 files changed, 69 insertions(+), 4 deletions(-)

diff --git a/crates/editor/src/edit_prediction_tests.rs b/crates/editor/src/edit_prediction_tests.rs
index 7d64dd9749c68cb0e436c1cfcb04e3458d052872..d897a670674cb23b075646c64f22e8b9bf0e4f90 100644
--- a/crates/editor/src/edit_prediction_tests.rs
+++ b/crates/editor/src/edit_prediction_tests.rs
@@ -1,12 +1,13 @@
 use edit_prediction::EditPredictionProvider;
-use gpui::{Entity, prelude::*};
+use gpui::{Entity, KeyBinding, Modifiers, prelude::*};
 use indoc::indoc;
 use multi_buffer::{Anchor, MultiBufferSnapshot, ToPoint};
 use std::ops::Range;
 use text::{Point, ToOffset};
 
 use crate::{
-    EditPrediction, editor_tests::init_test, test::editor_test_context::EditorTestContext,
+    AcceptEditPrediction, EditPrediction, MenuEditPredictionsPolicy, editor_tests::init_test,
+    test::editor_test_context::EditorTestContext,
 };
 
 #[gpui::test]
@@ -270,6 +271,63 @@ async fn test_edit_prediction_jump_disabled_for_non_zed_providers(cx: &mut gpui:
     });
 }
 
+#[gpui::test]
+async fn test_edit_prediction_preview_cleanup_on_toggle_off(cx: &mut gpui::TestAppContext) {
+    init_test(cx, |_| {});
+
+    // Bind `ctrl-shift-a` to accept the provided edit prediction. The actual key
+    // binding here doesn't matter, we simply need to confirm that holding the
+    // binding's modifiers triggers the edit prediction preview.
+    cx.update(|cx| cx.bind_keys([KeyBinding::new("ctrl-shift-a", AcceptEditPrediction, None)]));
+
+    let mut cx = EditorTestContext::new(cx).await;
+    let provider = cx.new(|_| FakeEditPredictionProvider::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.set_menu_edit_predictions_policy(MenuEditPredictionsPolicy::ByProvider);
+        editor.update_visible_edit_prediction(window, cx)
+    });
+
+    cx.editor(|editor, _, _| {
+        assert!(editor.has_active_edit_prediction());
+    });
+
+    // Simulate pressing the modifiers for `AcceptEditPrediction`, namely
+    // `ctrl-shift`, so that we can confirm that the edit prediction preview is
+    // activated.
+    let modifiers = Modifiers::control_shift();
+    cx.simulate_modifiers_change(modifiers);
+    cx.run_until_parked();
+
+    cx.editor(|editor, _, _| {
+        assert!(editor.edit_prediction_preview_is_active());
+    });
+
+    // Disable showing edit predictions without issuing a new modifiers changed
+    // event, to confirm that the edit prediction preview is still active.
+    cx.update_editor(|editor, window, cx| {
+        editor.set_show_edit_predictions(Some(false), window, cx);
+    });
+
+    cx.editor(|editor, _, _| {
+        assert!(!editor.has_active_edit_prediction());
+        assert!(editor.edit_prediction_preview_is_active());
+    });
+
+    // Now release the modifiers
+    // Simulate releasing all modifiers, ensuring that even with edit prediction
+    // disabled, the edit prediction preview is cleaned up.
+    cx.simulate_modifiers_change(Modifiers::none());
+    cx.run_until_parked();
+
+    cx.editor(|editor, _, _| {
+        assert!(!editor.edit_prediction_preview_is_active());
+    });
+}
+
 fn assert_editor_active_edit_completion(
     cx: &mut EditorTestContext,
     assert: impl FnOnce(MultiBufferSnapshot, &Vec<(Range<Anchor>, String)>),
@@ -395,7 +453,7 @@ impl EditPredictionProvider for FakeEditPredictionProvider {
     }
 
     fn show_completions_in_menu() -> bool {
-        false
+        true
     }
 
     fn supports_jump_to_edit() -> bool {
diff --git a/crates/editor/src/editor.rs b/crates/editor/src/editor.rs
index efd41e7465a0d375d957553d5c4cc283a78db276..9d2e2fc741c2ccce21a41fee4f9f32f20b1d33ab 100644
--- a/crates/editor/src/editor.rs
+++ b/crates/editor/src/editor.rs
@@ -7576,7 +7576,14 @@ impl Editor {
         window: &mut Window,
         cx: &mut Context<Self>,
     ) {
-        if self.show_edit_predictions_in_menu() {
+        // Ensure that the edit prediction preview is updated, even when not
+        // enabled, if there's an active edit prediction preview.
+        if self.show_edit_predictions_in_menu()
+            || matches!(
+                self.edit_prediction_preview,
+                EditPredictionPreview::Active { .. }
+            )
+        {
             self.update_edit_prediction_preview(&modifiers, window, cx);
         }