edit_prediction_tests.rs

   1use edit_prediction_types::{
   2    EditPredictionDelegate, EditPredictionIconSet, PredictedCursorPosition,
   3};
   4use futures::StreamExt;
   5use gpui::{
   6    Entity, KeyBinding, KeybindingKeystroke, Keystroke, Modifiers, NoAction, Task, prelude::*,
   7};
   8use indoc::indoc;
   9use language::EditPredictionsMode;
  10use language::{Buffer, CodeLabel};
  11use multi_buffer::{Anchor, MultiBufferSnapshot, ToPoint};
  12use project::{Completion, CompletionResponse, CompletionSource};
  13use std::{
  14    ops::Range,
  15    rc::Rc,
  16    sync::{
  17        Arc,
  18        atomic::{self, AtomicUsize},
  19    },
  20};
  21use text::{Point, ToOffset};
  22use ui::prelude::*;
  23
  24use crate::{
  25    AcceptEditPrediction, CompletionContext, CompletionProvider, EditPrediction,
  26    EditPredictionKeybindAction, EditPredictionKeybindSurface, MenuEditPredictionsPolicy,
  27    ShowCompletions,
  28    editor_tests::{init_test, update_test_language_settings},
  29    test::{editor_lsp_test_context::EditorLspTestContext, editor_test_context::EditorTestContext},
  30};
  31use rpc::proto::PeerId;
  32use workspace::CollaboratorId;
  33
  34#[gpui::test]
  35async fn test_edit_prediction_insert(cx: &mut gpui::TestAppContext) {
  36    init_test(cx, |_| {});
  37
  38    let mut cx = EditorTestContext::new(cx).await;
  39    let provider = cx.new(|_| FakeEditPredictionDelegate::default());
  40    assign_editor_completion_provider(provider.clone(), &mut cx);
  41    cx.set_state("let absolute_zero_celsius = ˇ;");
  42
  43    propose_edits(&provider, vec![(28..28, "-273.15")], &mut cx);
  44    cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
  45
  46    assert_editor_active_edit_completion(&mut cx, |_, edits| {
  47        assert_eq!(edits.len(), 1);
  48        assert_eq!(edits[0].1.as_ref(), "-273.15");
  49    });
  50
  51    accept_completion(&mut cx);
  52
  53    cx.assert_editor_state("let absolute_zero_celsius = -273.15ˇ;")
  54}
  55
  56#[gpui::test]
  57async fn test_edit_prediction_cursor_position_inside_insertion(cx: &mut gpui::TestAppContext) {
  58    init_test(cx, |_| {
  59        eprintln!("");
  60    });
  61
  62    let mut cx = EditorTestContext::new(cx).await;
  63    let provider = cx.new(|_| FakeEditPredictionDelegate::default());
  64
  65    assign_editor_completion_provider(provider.clone(), &mut cx);
  66    // Buffer: "fn foo() {}" - we'll insert text and position cursor inside the insertion
  67    cx.set_state("fn foo() ˇ{}");
  68
  69    // Insert "bar()" at offset 9, with cursor at offset 2 within the insertion (after "ba")
  70    // This tests the case where cursor is inside newly inserted text
  71    propose_edits_with_cursor_position_in_insertion(
  72        &provider,
  73        vec![(9..9, "bar()")],
  74        9, // anchor at the insertion point
  75        2, // offset 2 within "bar()" puts cursor after "ba"
  76        &mut cx,
  77    );
  78    cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
  79
  80    assert_editor_active_edit_completion(&mut cx, |_, edits| {
  81        assert_eq!(edits.len(), 1);
  82        assert_eq!(edits[0].1.as_ref(), "bar()");
  83    });
  84
  85    accept_completion(&mut cx);
  86
  87    // Cursor should be inside the inserted text at "baˇr()"
  88    cx.assert_editor_state("fn foo() baˇr(){}");
  89}
  90
  91#[gpui::test]
  92async fn test_edit_prediction_cursor_position_outside_edit(cx: &mut gpui::TestAppContext) {
  93    init_test(cx, |_| {});
  94
  95    let mut cx = EditorTestContext::new(cx).await;
  96    let provider = cx.new(|_| FakeEditPredictionDelegate::default());
  97    assign_editor_completion_provider(provider.clone(), &mut cx);
  98    // Buffer: "let x = ;" with cursor before semicolon - we'll insert "42" and position cursor elsewhere
  99    cx.set_state("let x = ˇ;");
 100
 101    // Insert "42" at offset 8, but set cursor_position to offset 4 (the 'x')
 102    // This tests that cursor moves to the predicted position, not the end of the edit
 103    propose_edits_with_cursor_position(
 104        &provider,
 105        vec![(8..8, "42")],
 106        Some(4), // cursor at offset 4 (the 'x'), NOT at the edit location
 107        &mut cx,
 108    );
 109    cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
 110
 111    assert_editor_active_edit_completion(&mut cx, |_, edits| {
 112        assert_eq!(edits.len(), 1);
 113        assert_eq!(edits[0].1.as_ref(), "42");
 114    });
 115
 116    accept_completion(&mut cx);
 117
 118    // Cursor should be at offset 4 (the 'x'), not at the end of the inserted "42"
 119    cx.assert_editor_state("let ˇx = 42;");
 120}
 121
 122#[gpui::test]
 123async fn test_edit_prediction_cursor_position_fallback(cx: &mut gpui::TestAppContext) {
 124    init_test(cx, |_| {});
 125
 126    let mut cx = EditorTestContext::new(cx).await;
 127    let provider = cx.new(|_| FakeEditPredictionDelegate::default());
 128    assign_editor_completion_provider(provider.clone(), &mut cx);
 129    cx.set_state("let x = ˇ;");
 130
 131    // Propose an edit without a cursor position - should fall back to end of edit
 132    propose_edits(&provider, vec![(8..8, "42")], &mut cx);
 133    cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
 134
 135    accept_completion(&mut cx);
 136
 137    // Cursor should be at the end of the inserted text (default behavior)
 138    cx.assert_editor_state("let x = 42ˇ;")
 139}
 140
 141#[gpui::test]
 142async fn test_edit_prediction_modification(cx: &mut gpui::TestAppContext) {
 143    init_test(cx, |_| {});
 144
 145    let mut cx = EditorTestContext::new(cx).await;
 146    let provider = cx.new(|_| FakeEditPredictionDelegate::default());
 147    assign_editor_completion_provider(provider.clone(), &mut cx);
 148    cx.set_state("let pi = ˇ\"foo\";");
 149
 150    propose_edits(&provider, vec![(9..14, "3.14159")], &mut cx);
 151    cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
 152
 153    assert_editor_active_edit_completion(&mut cx, |_, edits| {
 154        assert_eq!(edits.len(), 1);
 155        assert_eq!(edits[0].1.as_ref(), "3.14159");
 156    });
 157
 158    accept_completion(&mut cx);
 159
 160    cx.assert_editor_state("let pi = 3.14159ˇ;")
 161}
 162
 163#[gpui::test]
 164async fn test_edit_prediction_jump_button(cx: &mut gpui::TestAppContext) {
 165    init_test(cx, |_| {});
 166
 167    let mut cx = EditorTestContext::new(cx).await;
 168    let provider = cx.new(|_| FakeEditPredictionDelegate::default());
 169    assign_editor_completion_provider(provider.clone(), &mut cx);
 170
 171    // Cursor is 2+ lines above the proposed edit
 172    cx.set_state(indoc! {"
 173        line 0
 174        line ˇ1
 175        line 2
 176        line 3
 177        line
 178    "});
 179
 180    propose_edits(
 181        &provider,
 182        vec![(Point::new(4, 3)..Point::new(4, 3), " 4")],
 183        &mut cx,
 184    );
 185
 186    cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
 187    assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
 188        assert_eq!(move_target.to_point(&snapshot), Point::new(4, 3));
 189    });
 190
 191    // When accepting, cursor is moved to the proposed location
 192    accept_completion(&mut cx);
 193    cx.assert_editor_state(indoc! {"
 194        line 0
 195        line 1
 196        line 2
 197        line 3
 198        linˇe
 199    "});
 200
 201    // Cursor is 2+ lines below the proposed edit
 202    cx.set_state(indoc! {"
 203        line 0
 204        line
 205        line 2
 206        line 3
 207        line ˇ4
 208    "});
 209
 210    propose_edits(
 211        &provider,
 212        vec![(Point::new(1, 3)..Point::new(1, 3), " 1")],
 213        &mut cx,
 214    );
 215
 216    cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
 217    assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
 218        assert_eq!(move_target.to_point(&snapshot), Point::new(1, 3));
 219    });
 220
 221    // When accepting, cursor is moved to the proposed location
 222    accept_completion(&mut cx);
 223    cx.assert_editor_state(indoc! {"
 224        line 0
 225        linˇe
 226        line 2
 227        line 3
 228        line 4
 229    "});
 230}
 231
 232#[gpui::test]
 233async fn test_edit_prediction_invalidation_range(cx: &mut gpui::TestAppContext) {
 234    init_test(cx, |_| {});
 235
 236    let mut cx = EditorTestContext::new(cx).await;
 237    let provider = cx.new(|_| FakeEditPredictionDelegate::default());
 238    assign_editor_completion_provider(provider.clone(), &mut cx);
 239
 240    // Cursor is 3+ lines above the proposed edit
 241    cx.set_state(indoc! {"
 242        line 0
 243        line ˇ1
 244        line 2
 245        line 3
 246        line 4
 247        line
 248    "});
 249    let edit_location = Point::new(5, 3);
 250
 251    propose_edits(
 252        &provider,
 253        vec![(edit_location..edit_location, " 5")],
 254        &mut cx,
 255    );
 256
 257    cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
 258    assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
 259        assert_eq!(move_target.to_point(&snapshot), edit_location);
 260    });
 261
 262    // If we move *towards* the completion, it stays active
 263    cx.set_selections_state(indoc! {"
 264        line 0
 265        line 1
 266        line ˇ2
 267        line 3
 268        line 4
 269        line
 270    "});
 271    assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
 272        assert_eq!(move_target.to_point(&snapshot), edit_location);
 273    });
 274
 275    // If we move *away* from the completion, it is discarded
 276    cx.set_selections_state(indoc! {"
 277        line ˇ0
 278        line 1
 279        line 2
 280        line 3
 281        line 4
 282        line
 283    "});
 284    cx.editor(|editor, _, _| {
 285        assert!(editor.active_edit_prediction.is_none());
 286    });
 287
 288    // Cursor is 3+ lines below the proposed edit
 289    cx.set_state(indoc! {"
 290        line
 291        line 1
 292        line 2
 293        line 3
 294        line ˇ4
 295        line 5
 296    "});
 297    let edit_location = Point::new(0, 3);
 298
 299    propose_edits(
 300        &provider,
 301        vec![(edit_location..edit_location, " 0")],
 302        &mut cx,
 303    );
 304
 305    cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
 306    assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
 307        assert_eq!(move_target.to_point(&snapshot), edit_location);
 308    });
 309
 310    // If we move *towards* the completion, it stays active
 311    cx.set_selections_state(indoc! {"
 312        line
 313        line 1
 314        line 2
 315        line ˇ3
 316        line 4
 317        line 5
 318    "});
 319    assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
 320        assert_eq!(move_target.to_point(&snapshot), edit_location);
 321    });
 322
 323    // If we move *away* from the completion, it is discarded
 324    cx.set_selections_state(indoc! {"
 325        line
 326        line 1
 327        line 2
 328        line 3
 329        line 4
 330        line ˇ5
 331    "});
 332    cx.editor(|editor, _, _| {
 333        assert!(editor.active_edit_prediction.is_none());
 334    });
 335}
 336
 337#[gpui::test]
 338async fn test_edit_prediction_jump_disabled_for_non_zed_providers(cx: &mut gpui::TestAppContext) {
 339    init_test(cx, |_| {});
 340
 341    let mut cx = EditorTestContext::new(cx).await;
 342    let provider = cx.new(|_| FakeNonZedEditPredictionDelegate::default());
 343    assign_editor_completion_provider_non_zed(provider.clone(), &mut cx);
 344
 345    // Cursor is 2+ lines above the proposed edit
 346    cx.set_state(indoc! {"
 347        line 0
 348        line ˇ1
 349        line 2
 350        line 3
 351        line
 352    "});
 353
 354    propose_edits_non_zed(
 355        &provider,
 356        vec![(Point::new(4, 3)..Point::new(4, 3), " 4")],
 357        &mut cx,
 358    );
 359
 360    cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
 361
 362    // For non-Zed providers, there should be no move completion (jump functionality disabled)
 363    cx.editor(|editor, _, _| {
 364        if let Some(completion_state) = &editor.active_edit_prediction {
 365            // Should be an Edit prediction, not a Move prediction
 366            match &completion_state.completion {
 367                EditPrediction::Edit { .. } => {
 368                    // This is expected for non-Zed providers
 369                }
 370                EditPrediction::MoveWithin { .. } | EditPrediction::MoveOutside { .. } => {
 371                    panic!(
 372                        "Non-Zed providers should not show Move predictions (jump functionality)"
 373                    );
 374                }
 375            }
 376        }
 377    });
 378}
 379
 380#[gpui::test]
 381async fn test_edit_prediction_refresh_suppressed_while_following(cx: &mut gpui::TestAppContext) {
 382    init_test(cx, |_| {});
 383
 384    let mut cx = EditorTestContext::new(cx).await;
 385    let provider = cx.new(|_| FakeEditPredictionDelegate::default());
 386    assign_editor_completion_provider(provider.clone(), &mut cx);
 387    cx.set_state("let x = ˇ;");
 388
 389    propose_edits(&provider, vec![(8..8, "42")], &mut cx);
 390
 391    cx.update_editor(|editor, window, cx| {
 392        editor.refresh_edit_prediction(false, false, window, cx);
 393        editor.update_visible_edit_prediction(window, cx);
 394    });
 395
 396    assert_eq!(
 397        provider.read_with(&cx.cx, |provider, _| {
 398            provider.refresh_count.load(atomic::Ordering::SeqCst)
 399        }),
 400        1
 401    );
 402    cx.editor(|editor, _, _| {
 403        assert!(editor.active_edit_prediction.is_some());
 404    });
 405
 406    cx.update_editor(|editor, window, cx| {
 407        editor.leader_id = Some(CollaboratorId::PeerId(PeerId::default()));
 408        editor.refresh_edit_prediction(false, false, window, cx);
 409    });
 410
 411    assert_eq!(
 412        provider.read_with(&cx.cx, |provider, _| {
 413            provider.refresh_count.load(atomic::Ordering::SeqCst)
 414        }),
 415        1
 416    );
 417    cx.editor(|editor, _, _| {
 418        assert!(editor.active_edit_prediction.is_none());
 419    });
 420
 421    cx.update_editor(|editor, window, cx| {
 422        editor.leader_id = None;
 423        editor.refresh_edit_prediction(false, false, window, cx);
 424    });
 425
 426    assert_eq!(
 427        provider.read_with(&cx.cx, |provider, _| {
 428            provider.refresh_count.load(atomic::Ordering::SeqCst)
 429        }),
 430        2
 431    );
 432}
 433
 434#[gpui::test]
 435async fn test_edit_prediction_preview_cleanup_on_toggle_off(cx: &mut gpui::TestAppContext) {
 436    init_test(cx, |_| {});
 437
 438    // Bind `ctrl-shift-a` to accept the provided edit prediction. The actual key
 439    // binding here doesn't matter, we simply need to confirm that holding the
 440    // binding's modifiers triggers the edit prediction preview.
 441    cx.update(|cx| cx.bind_keys([KeyBinding::new("ctrl-shift-a", AcceptEditPrediction, None)]));
 442
 443    let mut cx = EditorTestContext::new(cx).await;
 444    let provider = cx.new(|_| FakeEditPredictionDelegate::default());
 445    assign_editor_completion_provider(provider.clone(), &mut cx);
 446    cx.set_state("let x = ˇ;");
 447
 448    propose_edits(&provider, vec![(8..8, "42")], &mut cx);
 449    cx.update_editor(|editor, window, cx| {
 450        editor.set_menu_edit_predictions_policy(MenuEditPredictionsPolicy::ByProvider);
 451        editor.update_visible_edit_prediction(window, cx)
 452    });
 453
 454    cx.editor(|editor, _, _| {
 455        assert!(editor.has_active_edit_prediction());
 456    });
 457
 458    // Simulate pressing the modifiers for `AcceptEditPrediction`, namely
 459    // `ctrl-shift`, so that we can confirm that the edit prediction preview is
 460    // activated.
 461    let modifiers = Modifiers::control_shift();
 462    cx.simulate_modifiers_change(modifiers);
 463    cx.run_until_parked();
 464
 465    cx.editor(|editor, _, _| {
 466        assert!(editor.edit_prediction_preview_is_active());
 467    });
 468
 469    // Disable showing edit predictions without issuing a new modifiers changed
 470    // event, to confirm that the edit prediction preview is still active.
 471    cx.update_editor(|editor, window, cx| {
 472        editor.set_show_edit_predictions(Some(false), window, cx);
 473    });
 474
 475    cx.editor(|editor, _, _| {
 476        assert!(!editor.has_active_edit_prediction());
 477        assert!(editor.edit_prediction_preview_is_active());
 478    });
 479
 480    // Now release the modifiers
 481    // Simulate releasing all modifiers, ensuring that even with edit prediction
 482    // disabled, the edit prediction preview is cleaned up.
 483    cx.simulate_modifiers_change(Modifiers::none());
 484    cx.run_until_parked();
 485
 486    cx.editor(|editor, _, _| {
 487        assert!(!editor.edit_prediction_preview_is_active());
 488    });
 489}
 490
 491#[gpui::test]
 492async fn test_edit_prediction_preview_activates_when_prediction_arrives_with_modifier_held(
 493    cx: &mut gpui::TestAppContext,
 494) {
 495    init_test(cx, |_| {});
 496    load_default_keymap(cx);
 497    update_test_language_settings(cx, &|settings| {
 498        settings.edit_predictions.get_or_insert_default().mode = Some(EditPredictionsMode::Subtle);
 499    });
 500
 501    let mut cx = EditorTestContext::new(cx).await;
 502    let provider = cx.new(|_| FakeEditPredictionDelegate::default());
 503    assign_editor_completion_provider(provider.clone(), &mut cx);
 504    cx.set_state("let x = ˇ;");
 505
 506    cx.editor(|editor, _, _| {
 507        assert!(!editor.has_active_edit_prediction());
 508        assert!(!editor.edit_prediction_preview_is_active());
 509    });
 510
 511    let preview_modifiers = cx.update_editor(|editor, window, cx| {
 512        *editor
 513            .preview_edit_prediction_keystroke(window, cx)
 514            .unwrap()
 515            .modifiers()
 516    });
 517
 518    cx.simulate_modifiers_change(preview_modifiers);
 519    cx.run_until_parked();
 520
 521    cx.editor(|editor, _, _| {
 522        assert!(!editor.has_active_edit_prediction());
 523        assert!(editor.edit_prediction_preview_is_active());
 524    });
 525
 526    propose_edits(&provider, vec![(8..8, "42")], &mut cx);
 527    cx.update_editor(|editor, window, cx| {
 528        editor.set_menu_edit_predictions_policy(MenuEditPredictionsPolicy::ByProvider);
 529        editor.update_visible_edit_prediction(window, cx)
 530    });
 531
 532    cx.editor(|editor, _, _| {
 533        assert!(editor.has_active_edit_prediction());
 534        assert!(
 535            editor.edit_prediction_preview_is_active(),
 536            "prediction preview should activate immediately when the prediction arrives while the preview modifier is still held",
 537        );
 538    });
 539}
 540
 541#[gpui::test]
 542async fn test_edit_prediction_preview_does_not_hide_code_actions_on_modifier_press(
 543    cx: &mut gpui::TestAppContext,
 544) {
 545    init_test(cx, |_| {});
 546    update_test_language_settings(cx, &|settings| {
 547        settings.edit_predictions.get_or_insert_default().mode = Some(EditPredictionsMode::Subtle);
 548    });
 549    cx.update(|cx| {
 550        cx.bind_keys([KeyBinding::new(
 551            "ctrl-enter",
 552            AcceptEditPrediction,
 553            Some("Editor && edit_prediction && !showing_completions"),
 554        )]);
 555    });
 556
 557    let mut cx = EditorLspTestContext::new_rust(
 558        lsp::ServerCapabilities {
 559            code_action_provider: Some(lsp::CodeActionProviderCapability::Simple(true)),
 560            ..Default::default()
 561        },
 562        cx,
 563    )
 564    .await;
 565    cx.set_state(indoc! {"
 566        fn main() {
 567            let valueˇ = 1;
 568        }
 569    "});
 570
 571    let provider = cx.new(|_| FakeEditPredictionDelegate::default());
 572    cx.update_editor(|editor, window, cx| {
 573        editor.set_edit_prediction_provider(Some(provider.clone()), window, cx);
 574    });
 575
 576    let snapshot = cx.buffer_snapshot();
 577    let edit_position = snapshot.anchor_after(Point::new(1, 13));
 578    cx.update(|_, cx| {
 579        provider.update(cx, |provider, _| {
 580            provider.set_edit_prediction(Some(edit_prediction_types::EditPrediction::Local {
 581                id: None,
 582                edits: vec![(edit_position..edit_position, " + 1".into())],
 583                cursor_position: None,
 584                edit_preview: None,
 585            }))
 586        })
 587    });
 588    cx.update_editor(|editor, window, cx| {
 589        editor.set_menu_edit_predictions_policy(MenuEditPredictionsPolicy::ByProvider);
 590        editor.update_visible_edit_prediction(window, cx);
 591    });
 592    cx.update_editor(|editor, _, _| {
 593        assert!(editor.has_active_edit_prediction());
 594        assert!(editor.stale_edit_prediction_in_menu.is_none());
 595    });
 596
 597    let mut code_action_requests = cx.set_request_handler::<lsp::request::CodeActionRequest, _, _>(
 598        move |_, _, _| async move {
 599            Ok(Some(vec![lsp::CodeActionOrCommand::CodeAction(
 600                lsp::CodeAction {
 601                    title: "Inline value".to_string(),
 602                    kind: Some(lsp::CodeActionKind::QUICKFIX),
 603                    ..Default::default()
 604                },
 605            )]))
 606        },
 607    );
 608
 609    cx.update_editor(|editor, window, cx| {
 610        editor.toggle_code_actions(
 611            &crate::actions::ToggleCodeActions {
 612                deployed_from: None,
 613                quick_launch: false,
 614            },
 615            window,
 616            cx,
 617        );
 618    });
 619    code_action_requests.next().await;
 620    cx.run_until_parked();
 621    cx.condition(|editor, _| editor.context_menu_visible())
 622        .await;
 623
 624    cx.update_editor(|editor, _, _| {
 625        assert!(!editor.has_active_edit_prediction());
 626        assert!(editor.stale_edit_prediction_in_menu.is_some());
 627        assert!(editor.context_menu_visible());
 628        assert!(matches!(
 629            editor.context_menu.borrow().as_ref(),
 630            Some(crate::code_context_menus::CodeContextMenu::CodeActions(_))
 631        ));
 632        assert!(!editor.edit_prediction_preview_is_active());
 633    });
 634
 635    cx.simulate_modifiers_change(Modifiers::control());
 636    cx.run_until_parked();
 637
 638    cx.update_editor(|editor, _, _| {
 639        assert!(
 640            !editor.edit_prediction_preview_is_active(),
 641            "modifier-only press should not activate edit prediction preview while code actions are open"
 642        );
 643        assert!(
 644            editor.context_menu_visible(),
 645            "modifier-only press should not hide the code actions menu"
 646        );
 647        assert!(matches!(
 648            editor.context_menu.borrow().as_ref(),
 649            Some(crate::code_context_menus::CodeContextMenu::CodeActions(_))
 650        ));
 651    });
 652}
 653
 654#[gpui::test]
 655async fn test_edit_prediction_preview_supersedes_completions_menu(cx: &mut gpui::TestAppContext) {
 656    init_test(cx, |_| {});
 657    update_test_language_settings(cx, &|settings| {
 658        settings.edit_predictions.get_or_insert_default().mode = Some(EditPredictionsMode::Subtle);
 659    });
 660    cx.update(|cx| {
 661        cx.bind_keys([KeyBinding::new(
 662            "ctrl-enter",
 663            AcceptEditPrediction,
 664            Some("Editor && edit_prediction && showing_completions"),
 665        )]);
 666    });
 667
 668    let mut cx = EditorTestContext::new(cx).await;
 669    let provider = cx.new(|_| FakeEditPredictionDelegate::default());
 670    assign_editor_completion_provider(provider.clone(), &mut cx);
 671    assign_editor_completion_menu_provider(&mut cx);
 672    cx.set_state("let x = ˇ;");
 673
 674    propose_edits(&provider, vec![(8..8, "42")], &mut cx);
 675    cx.update_editor(|editor, window, cx| {
 676        editor.set_menu_edit_predictions_policy(MenuEditPredictionsPolicy::ByProvider);
 677        editor.update_visible_edit_prediction(window, cx);
 678    });
 679    cx.update_editor(|editor, window, cx| {
 680        editor.show_completions(&ShowCompletions, window, cx);
 681    });
 682    cx.run_until_parked();
 683
 684    cx.editor(|editor, _, _| {
 685        assert!(editor.has_active_edit_prediction());
 686        assert!(editor.context_menu_visible());
 687        assert!(matches!(
 688            editor.context_menu.borrow().as_ref(),
 689            Some(crate::code_context_menus::CodeContextMenu::Completions(_))
 690        ));
 691        assert!(!editor.edit_prediction_preview_is_active());
 692    });
 693
 694    cx.simulate_modifiers_change(Modifiers::control());
 695    cx.run_until_parked();
 696
 697    cx.editor(|editor, _, _| {
 698        assert!(editor.edit_prediction_preview_is_active());
 699        assert!(!editor.context_menu_visible());
 700        assert!(matches!(
 701            editor.context_menu.borrow().as_ref(),
 702            Some(crate::code_context_menus::CodeContextMenu::Completions(_))
 703        ));
 704    });
 705}
 706
 707fn load_default_keymap(cx: &mut gpui::TestAppContext) {
 708    cx.update(|cx| {
 709        cx.bind_keys(
 710            settings::KeymapFile::load_asset_allow_partial_failure(
 711                settings::DEFAULT_KEYMAP_PATH,
 712                cx,
 713            )
 714            .expect("failed to load default keymap"),
 715        );
 716    });
 717}
 718
 719#[gpui::test]
 720async fn test_inline_edit_prediction_keybind_selection_cases(cx: &mut gpui::TestAppContext) {
 721    enum InlineKeybindState {
 722        Normal,
 723        ShowingCompletions,
 724        InLeadingWhitespace,
 725        ShowingCompletionsAndLeadingWhitespace,
 726    }
 727
 728    enum ExpectedKeystroke {
 729        DefaultAccept,
 730        DefaultPreview,
 731        Literal(&'static str),
 732    }
 733
 734    struct InlineKeybindCase {
 735        name: &'static str,
 736        use_default_keymap: bool,
 737        mode: EditPredictionsMode,
 738        extra_bindings: Vec<KeyBinding>,
 739        state: InlineKeybindState,
 740        expected_accept_keystroke: ExpectedKeystroke,
 741        expected_preview_keystroke: ExpectedKeystroke,
 742        expected_displayed_keystroke: ExpectedKeystroke,
 743    }
 744
 745    init_test(cx, |_| {});
 746    load_default_keymap(cx);
 747    let mut default_cx = EditorTestContext::new(cx).await;
 748    let provider = default_cx.new(|_| FakeEditPredictionDelegate::default());
 749    assign_editor_completion_provider(provider.clone(), &mut default_cx);
 750    default_cx.set_state("let x = ˇ;");
 751    propose_edits(&provider, vec![(8..8, "42")], &mut default_cx);
 752    default_cx
 753        .update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
 754
 755    let (default_accept_keystroke, default_preview_keystroke) =
 756        default_cx.update_editor(|editor, window, cx| {
 757            let keybind_display = editor.edit_prediction_keybind_display(
 758                EditPredictionKeybindSurface::Inline,
 759                window,
 760                cx,
 761            );
 762            let accept_keystroke = keybind_display
 763                .accept_keystroke
 764                .as_ref()
 765                .expect("default inline edit prediction should have an accept binding")
 766                .clone();
 767            let preview_keystroke = keybind_display
 768                .preview_keystroke
 769                .as_ref()
 770                .expect("default inline edit prediction should have a preview binding")
 771                .clone();
 772            (accept_keystroke, preview_keystroke)
 773        });
 774
 775    let cases = [
 776        InlineKeybindCase {
 777            name: "default setup prefers tab over alt-tab for accept",
 778            use_default_keymap: true,
 779            mode: EditPredictionsMode::Eager,
 780            extra_bindings: Vec::new(),
 781            state: InlineKeybindState::Normal,
 782            expected_accept_keystroke: ExpectedKeystroke::DefaultAccept,
 783            expected_preview_keystroke: ExpectedKeystroke::DefaultPreview,
 784            expected_displayed_keystroke: ExpectedKeystroke::DefaultAccept,
 785        },
 786        InlineKeybindCase {
 787            name: "subtle mode displays preview binding inline",
 788            use_default_keymap: true,
 789            mode: EditPredictionsMode::Subtle,
 790            extra_bindings: Vec::new(),
 791            state: InlineKeybindState::Normal,
 792            expected_accept_keystroke: ExpectedKeystroke::DefaultPreview,
 793            expected_preview_keystroke: ExpectedKeystroke::DefaultPreview,
 794            expected_displayed_keystroke: ExpectedKeystroke::DefaultPreview,
 795        },
 796        InlineKeybindCase {
 797            name: "removing default tab binding still displays tab",
 798            use_default_keymap: true,
 799            mode: EditPredictionsMode::Eager,
 800            extra_bindings: vec![KeyBinding::new(
 801                "tab",
 802                NoAction,
 803                Some("Editor && edit_prediction && edit_prediction_mode == eager"),
 804            )],
 805            state: InlineKeybindState::Normal,
 806            expected_accept_keystroke: ExpectedKeystroke::DefaultPreview,
 807            expected_preview_keystroke: ExpectedKeystroke::DefaultPreview,
 808            expected_displayed_keystroke: ExpectedKeystroke::DefaultPreview,
 809        },
 810        InlineKeybindCase {
 811            name: "custom-only rebound accept key uses replacement key",
 812            use_default_keymap: true,
 813            mode: EditPredictionsMode::Eager,
 814            extra_bindings: vec![KeyBinding::new(
 815                "ctrl-enter",
 816                AcceptEditPrediction,
 817                Some("Editor && edit_prediction"),
 818            )],
 819            state: InlineKeybindState::Normal,
 820            expected_accept_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
 821            expected_preview_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
 822            expected_displayed_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
 823        },
 824        InlineKeybindCase {
 825            name: "showing completions restores conflict-context binding",
 826            use_default_keymap: true,
 827            mode: EditPredictionsMode::Eager,
 828            extra_bindings: vec![KeyBinding::new(
 829                "ctrl-enter",
 830                AcceptEditPrediction,
 831                Some("Editor && edit_prediction && showing_completions"),
 832            )],
 833            state: InlineKeybindState::ShowingCompletions,
 834            expected_accept_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
 835            expected_preview_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
 836            expected_displayed_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
 837        },
 838        InlineKeybindCase {
 839            name: "leading whitespace restores conflict-context binding",
 840            use_default_keymap: false,
 841            mode: EditPredictionsMode::Eager,
 842            extra_bindings: vec![KeyBinding::new(
 843                "ctrl-enter",
 844                AcceptEditPrediction,
 845                Some("Editor && edit_prediction && in_leading_whitespace"),
 846            )],
 847            state: InlineKeybindState::InLeadingWhitespace,
 848            expected_accept_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
 849            expected_preview_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
 850            expected_displayed_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
 851        },
 852        InlineKeybindCase {
 853            name: "showing completions and leading whitespace restore combined conflict binding",
 854            use_default_keymap: false,
 855            mode: EditPredictionsMode::Eager,
 856            extra_bindings: vec![KeyBinding::new(
 857                "ctrl-enter",
 858                AcceptEditPrediction,
 859                Some("Editor && edit_prediction && showing_completions && in_leading_whitespace"),
 860            )],
 861            state: InlineKeybindState::ShowingCompletionsAndLeadingWhitespace,
 862            expected_accept_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
 863            expected_preview_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
 864            expected_displayed_keystroke: ExpectedKeystroke::Literal("ctrl-enter"),
 865        },
 866    ];
 867
 868    for case in cases {
 869        init_test(cx, |_| {});
 870        if case.use_default_keymap {
 871            load_default_keymap(cx);
 872        }
 873        update_test_language_settings(cx, &|settings| {
 874            settings.edit_predictions.get_or_insert_default().mode = Some(case.mode);
 875        });
 876
 877        if !case.extra_bindings.is_empty() {
 878            cx.update(|cx| cx.bind_keys(case.extra_bindings.clone()));
 879        }
 880
 881        let mut cx = EditorTestContext::new(cx).await;
 882        let provider = cx.new(|_| FakeEditPredictionDelegate::default());
 883        assign_editor_completion_provider(provider.clone(), &mut cx);
 884
 885        match case.state {
 886            InlineKeybindState::Normal | InlineKeybindState::ShowingCompletions => {
 887                cx.set_state("let x = ˇ;");
 888            }
 889            InlineKeybindState::InLeadingWhitespace
 890            | InlineKeybindState::ShowingCompletionsAndLeadingWhitespace => {
 891                cx.set_state(indoc! {"
 892                    fn main() {
 893                        ˇ
 894                    }
 895                "});
 896            }
 897        }
 898
 899        propose_edits(&provider, vec![(8..8, "42")], &mut cx);
 900        cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
 901
 902        if matches!(
 903            case.state,
 904            InlineKeybindState::ShowingCompletions
 905                | InlineKeybindState::ShowingCompletionsAndLeadingWhitespace
 906        ) {
 907            assign_editor_completion_menu_provider(&mut cx);
 908            cx.update_editor(|editor, window, cx| {
 909                editor.show_completions(&ShowCompletions, window, cx);
 910            });
 911            cx.run_until_parked();
 912        }
 913
 914        cx.update_editor(|editor, window, cx| {
 915            assert!(
 916                editor.has_active_edit_prediction(),
 917                "case '{}' should have an active edit prediction",
 918                case.name
 919            );
 920
 921            let keybind_display = editor.edit_prediction_keybind_display(
 922                EditPredictionKeybindSurface::Inline,
 923                window,
 924                cx,
 925            );
 926            let accept_keystroke = keybind_display
 927                .accept_keystroke
 928                .as_ref()
 929                .unwrap_or_else(|| panic!("case '{}' should have an accept binding", case.name));
 930            let preview_keystroke = keybind_display
 931                .preview_keystroke
 932                .as_ref()
 933                .unwrap_or_else(|| panic!("case '{}' should have a preview binding", case.name));
 934            let displayed_keystroke = keybind_display
 935                .displayed_keystroke
 936                .as_ref()
 937                .unwrap_or_else(|| panic!("case '{}' should have a displayed binding", case.name));
 938
 939            let expected_accept_keystroke = match case.expected_accept_keystroke {
 940                ExpectedKeystroke::DefaultAccept => default_accept_keystroke.clone(),
 941                ExpectedKeystroke::DefaultPreview => default_preview_keystroke.clone(),
 942                ExpectedKeystroke::Literal(keystroke) => KeybindingKeystroke::from_keystroke(
 943                    Keystroke::parse(keystroke).expect("expected test keystroke to parse"),
 944                ),
 945            };
 946            let expected_preview_keystroke = match case.expected_preview_keystroke {
 947                ExpectedKeystroke::DefaultAccept => default_accept_keystroke.clone(),
 948                ExpectedKeystroke::DefaultPreview => default_preview_keystroke.clone(),
 949                ExpectedKeystroke::Literal(keystroke) => KeybindingKeystroke::from_keystroke(
 950                    Keystroke::parse(keystroke).expect("expected test keystroke to parse"),
 951                ),
 952            };
 953            let expected_displayed_keystroke = match case.expected_displayed_keystroke {
 954                ExpectedKeystroke::DefaultAccept => default_accept_keystroke.clone(),
 955                ExpectedKeystroke::DefaultPreview => default_preview_keystroke.clone(),
 956                ExpectedKeystroke::Literal(keystroke) => KeybindingKeystroke::from_keystroke(
 957                    Keystroke::parse(keystroke).expect("expected test keystroke to parse"),
 958                ),
 959            };
 960
 961            assert_eq!(
 962                accept_keystroke, &expected_accept_keystroke,
 963                "case '{}' selected the wrong accept binding",
 964                case.name
 965            );
 966            assert_eq!(
 967                preview_keystroke, &expected_preview_keystroke,
 968                "case '{}' selected the wrong preview binding",
 969                case.name
 970            );
 971            assert_eq!(
 972                displayed_keystroke, &expected_displayed_keystroke,
 973                "case '{}' selected the wrong displayed binding",
 974                case.name
 975            );
 976
 977            if matches!(case.mode, EditPredictionsMode::Subtle) {
 978                assert!(
 979                    editor.edit_prediction_requires_modifier(),
 980                    "case '{}' should require a modifier",
 981                    case.name
 982                );
 983            }
 984        });
 985    }
 986}
 987
 988#[gpui::test]
 989async fn test_tab_accepts_edit_prediction_over_completion(cx: &mut gpui::TestAppContext) {
 990    init_test(cx, |_| {});
 991    load_default_keymap(cx);
 992
 993    let mut cx = EditorTestContext::new(cx).await;
 994    let provider = cx.new(|_| FakeEditPredictionDelegate::default());
 995    assign_editor_completion_provider(provider.clone(), &mut cx);
 996    cx.set_state("let x = ˇ;");
 997
 998    propose_edits(&provider, vec![(8..8, "42")], &mut cx);
 999    cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
1000
1001    assert_editor_active_edit_completion(&mut cx, |_, edits| {
1002        assert_eq!(edits.len(), 1);
1003        assert_eq!(edits[0].1.as_ref(), "42");
1004    });
1005
1006    cx.simulate_keystroke("tab");
1007    cx.run_until_parked();
1008
1009    cx.assert_editor_state("let x = 42ˇ;");
1010}
1011
1012#[gpui::test]
1013async fn test_cursor_popover_edit_prediction_keybind_cases(cx: &mut gpui::TestAppContext) {
1014    enum CursorPopoverPredictionKind {
1015        SingleLine,
1016        MultiLine,
1017        SingleLineWithPreview,
1018        MultiLineWithPreview,
1019        DeleteSingleNewline,
1020        StaleSingleLineAfterMultiLine,
1021    }
1022
1023    struct CursorPopoverCase {
1024        name: &'static str,
1025        prediction_kind: CursorPopoverPredictionKind,
1026        expected_action: EditPredictionKeybindAction,
1027    }
1028
1029    let cases = [
1030        CursorPopoverCase {
1031            name: "single line prediction uses accept action",
1032            prediction_kind: CursorPopoverPredictionKind::SingleLine,
1033            expected_action: EditPredictionKeybindAction::Accept,
1034        },
1035        CursorPopoverCase {
1036            name: "multi line prediction uses preview action",
1037            prediction_kind: CursorPopoverPredictionKind::MultiLine,
1038            expected_action: EditPredictionKeybindAction::Preview,
1039        },
1040        CursorPopoverCase {
1041            name: "single line prediction with preview still uses accept action",
1042            prediction_kind: CursorPopoverPredictionKind::SingleLineWithPreview,
1043            expected_action: EditPredictionKeybindAction::Accept,
1044        },
1045        CursorPopoverCase {
1046            name: "multi line prediction with preview uses preview action",
1047            prediction_kind: CursorPopoverPredictionKind::MultiLineWithPreview,
1048            expected_action: EditPredictionKeybindAction::Preview,
1049        },
1050        CursorPopoverCase {
1051            name: "single line newline deletion uses accept action",
1052            prediction_kind: CursorPopoverPredictionKind::DeleteSingleNewline,
1053            expected_action: EditPredictionKeybindAction::Accept,
1054        },
1055        CursorPopoverCase {
1056            name: "stale multi line prediction does not force preview action",
1057            prediction_kind: CursorPopoverPredictionKind::StaleSingleLineAfterMultiLine,
1058            expected_action: EditPredictionKeybindAction::Accept,
1059        },
1060    ];
1061
1062    for case in cases {
1063        init_test(cx, |_| {});
1064        load_default_keymap(cx);
1065
1066        let mut cx = EditorTestContext::new(cx).await;
1067        let provider = cx.new(|_| FakeEditPredictionDelegate::default());
1068        assign_editor_completion_provider(provider.clone(), &mut cx);
1069
1070        match case.prediction_kind {
1071            CursorPopoverPredictionKind::SingleLine => {
1072                cx.set_state("let x = ˇ;");
1073                propose_edits(&provider, vec![(8..8, "42")], &mut cx);
1074                cx.update_editor(|editor, window, cx| {
1075                    editor.update_visible_edit_prediction(window, cx)
1076                });
1077            }
1078            CursorPopoverPredictionKind::MultiLine => {
1079                cx.set_state("let x = ˇ;");
1080                propose_edits(&provider, vec![(8..8, "42\n43")], &mut cx);
1081                cx.update_editor(|editor, window, cx| {
1082                    editor.update_visible_edit_prediction(window, cx)
1083                });
1084            }
1085            CursorPopoverPredictionKind::SingleLineWithPreview => {
1086                cx.set_state("let x = ˇ;");
1087                propose_edits_with_preview(&provider, vec![(8..8, "42")], &mut cx).await;
1088                cx.update_editor(|editor, window, cx| {
1089                    editor.update_visible_edit_prediction(window, cx)
1090                });
1091            }
1092            CursorPopoverPredictionKind::MultiLineWithPreview => {
1093                cx.set_state("let x = ˇ;");
1094                propose_edits_with_preview(&provider, vec![(8..8, "42\n43")], &mut cx).await;
1095                cx.update_editor(|editor, window, cx| {
1096                    editor.update_visible_edit_prediction(window, cx)
1097                });
1098            }
1099            CursorPopoverPredictionKind::DeleteSingleNewline => {
1100                cx.set_state(indoc! {"
1101                    fn main() {
1102                        let value = 1;
1103                        ˇprintln!(\"done\");
1104                    }
1105                "});
1106                propose_edits(
1107                    &provider,
1108                    vec![(Point::new(1, 18)..Point::new(2, 17), "")],
1109                    &mut cx,
1110                );
1111                cx.update_editor(|editor, window, cx| {
1112                    editor.update_visible_edit_prediction(window, cx)
1113                });
1114            }
1115            CursorPopoverPredictionKind::StaleSingleLineAfterMultiLine => {
1116                cx.set_state("let x = ˇ;");
1117                propose_edits(&provider, vec![(8..8, "42\n43")], &mut cx);
1118                cx.update_editor(|editor, window, cx| {
1119                    editor.update_visible_edit_prediction(window, cx)
1120                });
1121                cx.update_editor(|editor, _window, cx| {
1122                    assert!(editor.active_edit_prediction.is_some());
1123                    assert!(editor.stale_edit_prediction_in_menu.is_none());
1124                    editor.take_active_edit_prediction(true, cx);
1125                    assert!(editor.active_edit_prediction.is_none());
1126                    assert!(editor.stale_edit_prediction_in_menu.is_some());
1127                });
1128
1129                propose_edits(&provider, vec![(8..8, "42")], &mut cx);
1130                cx.update_editor(|editor, window, cx| {
1131                    editor.update_visible_edit_prediction(window, cx)
1132                });
1133            }
1134        }
1135
1136        cx.update_editor(|editor, window, cx| {
1137            assert!(
1138                editor.has_active_edit_prediction(),
1139                "case '{}' should have an active edit prediction",
1140                case.name
1141            );
1142
1143            let keybind_display = editor.edit_prediction_keybind_display(
1144                EditPredictionKeybindSurface::CursorPopoverExpanded,
1145                window,
1146                cx,
1147            );
1148            let accept_keystroke = keybind_display
1149                .accept_keystroke
1150                .as_ref()
1151                .unwrap_or_else(|| panic!("case '{}' should have an accept binding", case.name));
1152            let preview_keystroke = keybind_display
1153                .preview_keystroke
1154                .as_ref()
1155                .unwrap_or_else(|| panic!("case '{}' should have a preview binding", case.name));
1156
1157            assert_eq!(
1158                keybind_display.action, case.expected_action,
1159                "case '{}' selected the wrong cursor popover action",
1160                case.name
1161            );
1162            assert_eq!(
1163                accept_keystroke.key(),
1164                "tab",
1165                "case '{}' selected the wrong accept binding",
1166                case.name
1167            );
1168            assert!(
1169                preview_keystroke.modifiers().modified(),
1170                "case '{}' should use a modified preview binding",
1171                case.name
1172            );
1173
1174            if matches!(
1175                case.prediction_kind,
1176                CursorPopoverPredictionKind::StaleSingleLineAfterMultiLine
1177            ) {
1178                assert!(
1179                    editor.stale_edit_prediction_in_menu.is_none(),
1180                    "case '{}' should clear stale menu state",
1181                    case.name
1182                );
1183            }
1184        });
1185    }
1186}
1187
1188fn assert_editor_active_edit_completion(
1189    cx: &mut EditorTestContext,
1190    assert: impl FnOnce(MultiBufferSnapshot, &Vec<(Range<Anchor>, Arc<str>)>),
1191) {
1192    cx.editor(|editor, _, cx| {
1193        let completion_state = editor
1194            .active_edit_prediction
1195            .as_ref()
1196            .expect("editor has no active completion");
1197
1198        if let EditPrediction::Edit { edits, .. } = &completion_state.completion {
1199            assert(editor.buffer().read(cx).snapshot(cx), edits);
1200        } else {
1201            panic!("expected edit completion");
1202        }
1203    })
1204}
1205
1206fn assert_editor_active_move_completion(
1207    cx: &mut EditorTestContext,
1208    assert: impl FnOnce(MultiBufferSnapshot, Anchor),
1209) {
1210    cx.editor(|editor, _, cx| {
1211        let completion_state = editor
1212            .active_edit_prediction
1213            .as_ref()
1214            .expect("editor has no active completion");
1215
1216        if let EditPrediction::MoveWithin { target, .. } = &completion_state.completion {
1217            assert(editor.buffer().read(cx).snapshot(cx), *target);
1218        } else {
1219            panic!("expected move completion");
1220        }
1221    })
1222}
1223
1224#[gpui::test]
1225async fn test_cancel_clears_stale_edit_prediction_in_menu(cx: &mut gpui::TestAppContext) {
1226    init_test(cx, |_| {});
1227    load_default_keymap(cx);
1228
1229    let mut cx = EditorTestContext::new(cx).await;
1230    let provider = cx.new(|_| FakeEditPredictionDelegate::default());
1231    assign_editor_completion_provider(provider.clone(), &mut cx);
1232    cx.set_state("let x = ˇ;");
1233
1234    propose_edits(&provider, vec![(8..8, "42")], &mut cx);
1235    cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
1236
1237    cx.update_editor(|editor, _window, _cx| {
1238        assert!(editor.active_edit_prediction.is_some());
1239        assert!(editor.stale_edit_prediction_in_menu.is_none());
1240    });
1241
1242    cx.simulate_keystroke("escape");
1243    cx.run_until_parked();
1244
1245    cx.update_editor(|editor, _window, _cx| {
1246        assert!(editor.active_edit_prediction.is_none());
1247        assert!(editor.stale_edit_prediction_in_menu.is_none());
1248    });
1249}
1250
1251#[gpui::test]
1252async fn test_discard_clears_delegate_completion(cx: &mut gpui::TestAppContext) {
1253    init_test(cx, |_| {});
1254    load_default_keymap(cx);
1255
1256    let mut cx = EditorTestContext::new(cx).await;
1257    let provider = cx.new(|_| FakeEditPredictionDelegate::default());
1258    assign_editor_completion_provider(provider.clone(), &mut cx);
1259    cx.set_state("let x = ˇ;");
1260
1261    propose_edits(&provider, vec![(8..8, "42")], &mut cx);
1262    cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
1263
1264    cx.update_editor(|editor, _window, _cx| {
1265        assert!(editor.active_edit_prediction.is_some());
1266    });
1267
1268    // Dismiss the prediction — this must call discard() on the delegate,
1269    // which should clear self.completion.
1270    cx.simulate_keystroke("escape");
1271    cx.run_until_parked();
1272
1273    cx.update_editor(|editor, _window, _cx| {
1274        assert!(editor.active_edit_prediction.is_none());
1275    });
1276
1277    // update_visible_edit_prediction must NOT bring the prediction back,
1278    // because discard() cleared self.completion in the delegate.
1279    cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
1280
1281    cx.update_editor(|editor, _window, _cx| {
1282        assert!(
1283            editor.active_edit_prediction.is_none(),
1284            "prediction must not resurface after discard()"
1285        );
1286    });
1287}
1288
1289fn accept_completion(cx: &mut EditorTestContext) {
1290    cx.update_editor(|editor, window, cx| {
1291        editor.accept_edit_prediction(&crate::AcceptEditPrediction, window, cx)
1292    })
1293}
1294
1295fn propose_edits<T: ToOffset>(
1296    provider: &Entity<FakeEditPredictionDelegate>,
1297    edits: Vec<(Range<T>, &str)>,
1298    cx: &mut EditorTestContext,
1299) {
1300    propose_edits_with_cursor_position(provider, edits, None, cx);
1301}
1302
1303async fn propose_edits_with_preview<T: ToOffset + Clone>(
1304    provider: &Entity<FakeEditPredictionDelegate>,
1305    edits: Vec<(Range<T>, &str)>,
1306    cx: &mut EditorTestContext,
1307) {
1308    let snapshot = cx.buffer_snapshot();
1309    let edits = edits
1310        .into_iter()
1311        .map(|(range, text)| {
1312            let anchor_range =
1313                snapshot.anchor_after(range.start.clone())..snapshot.anchor_before(range.end);
1314            (anchor_range, Arc::<str>::from(text))
1315        })
1316        .collect::<Vec<_>>();
1317
1318    let preview_edits = edits
1319        .iter()
1320        .map(|(range, text)| (range.clone(), text.clone()))
1321        .collect::<Arc<[_]>>();
1322
1323    let edit_preview = cx
1324        .buffer(|buffer: &Buffer, app| buffer.preview_edits(preview_edits, app))
1325        .await;
1326
1327    let provider_edits = edits.into_iter().collect();
1328
1329    cx.update(|_, cx| {
1330        provider.update(cx, |provider, _| {
1331            provider.set_edit_prediction(Some(edit_prediction_types::EditPrediction::Local {
1332                id: None,
1333                edits: provider_edits,
1334                cursor_position: None,
1335                edit_preview: Some(edit_preview),
1336            }))
1337        })
1338    });
1339}
1340
1341fn propose_edits_with_cursor_position<T: ToOffset>(
1342    provider: &Entity<FakeEditPredictionDelegate>,
1343    edits: Vec<(Range<T>, &str)>,
1344    cursor_offset: Option<usize>,
1345    cx: &mut EditorTestContext,
1346) {
1347    let snapshot = cx.buffer_snapshot();
1348    let cursor_position = cursor_offset
1349        .map(|offset| PredictedCursorPosition::at_anchor(snapshot.anchor_after(offset)));
1350    let edits = edits.into_iter().map(|(range, text)| {
1351        let range = snapshot.anchor_after(range.start)..snapshot.anchor_before(range.end);
1352        (range, text.into())
1353    });
1354
1355    cx.update(|_, cx| {
1356        provider.update(cx, |provider, _| {
1357            provider.set_edit_prediction(Some(edit_prediction_types::EditPrediction::Local {
1358                id: None,
1359                edits: edits.collect(),
1360                cursor_position,
1361                edit_preview: None,
1362            }))
1363        })
1364    });
1365}
1366
1367fn propose_edits_with_cursor_position_in_insertion<T: ToOffset>(
1368    provider: &Entity<FakeEditPredictionDelegate>,
1369    edits: Vec<(Range<T>, &str)>,
1370    anchor_offset: usize,
1371    offset_within_insertion: usize,
1372    cx: &mut EditorTestContext,
1373) {
1374    let snapshot = cx.buffer_snapshot();
1375    // Use anchor_before (left bias) so the anchor stays at the insertion point
1376    // rather than moving past the inserted text
1377    let cursor_position = Some(PredictedCursorPosition::new(
1378        snapshot.anchor_before(anchor_offset),
1379        offset_within_insertion,
1380    ));
1381    let edits = edits.into_iter().map(|(range, text)| {
1382        let range = snapshot.anchor_after(range.start)..snapshot.anchor_before(range.end);
1383        (range, text.into())
1384    });
1385
1386    cx.update(|_, cx| {
1387        provider.update(cx, |provider, _| {
1388            provider.set_edit_prediction(Some(edit_prediction_types::EditPrediction::Local {
1389                id: None,
1390                edits: edits.collect(),
1391                cursor_position,
1392                edit_preview: None,
1393            }))
1394        })
1395    });
1396}
1397
1398fn assign_editor_completion_provider(
1399    provider: Entity<FakeEditPredictionDelegate>,
1400    cx: &mut EditorTestContext,
1401) {
1402    cx.update_editor(|editor, window, cx| {
1403        editor.set_edit_prediction_provider(Some(provider), window, cx);
1404    })
1405}
1406
1407fn assign_editor_completion_menu_provider(cx: &mut EditorTestContext) {
1408    cx.update_editor(|editor, _, _| {
1409        editor.set_completion_provider(Some(Rc::new(FakeCompletionMenuProvider)));
1410    });
1411}
1412
1413fn propose_edits_non_zed<T: ToOffset>(
1414    provider: &Entity<FakeNonZedEditPredictionDelegate>,
1415    edits: Vec<(Range<T>, &str)>,
1416    cx: &mut EditorTestContext,
1417) {
1418    let snapshot = cx.buffer_snapshot();
1419    let edits = edits.into_iter().map(|(range, text)| {
1420        let range = snapshot.anchor_after(range.start)..snapshot.anchor_before(range.end);
1421        (range, text.into())
1422    });
1423
1424    cx.update(|_, cx| {
1425        provider.update(cx, |provider, _| {
1426            provider.set_edit_prediction(Some(edit_prediction_types::EditPrediction::Local {
1427                id: None,
1428                edits: edits.collect(),
1429                cursor_position: None,
1430                edit_preview: None,
1431            }))
1432        })
1433    });
1434}
1435
1436fn assign_editor_completion_provider_non_zed(
1437    provider: Entity<FakeNonZedEditPredictionDelegate>,
1438    cx: &mut EditorTestContext,
1439) {
1440    cx.update_editor(|editor, window, cx| {
1441        editor.set_edit_prediction_provider(Some(provider), window, cx);
1442    })
1443}
1444
1445struct FakeCompletionMenuProvider;
1446
1447impl CompletionProvider for FakeCompletionMenuProvider {
1448    fn completions(
1449        &self,
1450        buffer: &Entity<Buffer>,
1451        _buffer_position: text::Anchor,
1452        _trigger: CompletionContext,
1453        _window: &mut Window,
1454        cx: &mut Context<crate::Editor>,
1455    ) -> Task<anyhow::Result<Vec<CompletionResponse>>> {
1456        let replace_range = text::Anchor::min_max_range_for_buffer(buffer.read(cx).remote_id());
1457        let completions = ["fake_completion", "fake_completion_2"]
1458            .into_iter()
1459            .map(|label| Completion {
1460                replace_range: replace_range.clone(),
1461                new_text: label.to_string(),
1462                label: CodeLabel::plain(label.to_string(), None),
1463                documentation: None,
1464                source: CompletionSource::Custom,
1465                icon_path: None,
1466                match_start: None,
1467                snippet_deduplication_key: None,
1468                insert_text_mode: None,
1469                confirm: None,
1470            })
1471            .collect();
1472
1473        Task::ready(Ok(vec![CompletionResponse {
1474            completions,
1475            display_options: Default::default(),
1476            is_incomplete: false,
1477        }]))
1478    }
1479
1480    fn is_completion_trigger(
1481        &self,
1482        _buffer: &Entity<Buffer>,
1483        _position: language::Anchor,
1484        _text: &str,
1485        _trigger_in_words: bool,
1486        _cx: &mut Context<crate::Editor>,
1487    ) -> bool {
1488        false
1489    }
1490
1491    fn filter_completions(&self) -> bool {
1492        false
1493    }
1494}
1495
1496#[derive(Default, Clone)]
1497pub struct FakeEditPredictionDelegate {
1498    pub completion: Option<edit_prediction_types::EditPrediction>,
1499    pub refresh_count: Arc<AtomicUsize>,
1500}
1501
1502impl FakeEditPredictionDelegate {
1503    pub fn set_edit_prediction(
1504        &mut self,
1505        completion: Option<edit_prediction_types::EditPrediction>,
1506    ) {
1507        self.completion = completion;
1508    }
1509}
1510
1511impl EditPredictionDelegate for FakeEditPredictionDelegate {
1512    fn name() -> &'static str {
1513        "fake-completion-provider"
1514    }
1515
1516    fn display_name() -> &'static str {
1517        "Fake Completion Provider"
1518    }
1519
1520    fn show_predictions_in_menu() -> bool {
1521        true
1522    }
1523
1524    fn supports_jump_to_edit() -> bool {
1525        true
1526    }
1527
1528    fn icons(&self, _cx: &gpui::App) -> EditPredictionIconSet {
1529        EditPredictionIconSet::new(IconName::ZedPredict)
1530    }
1531
1532    fn is_enabled(
1533        &self,
1534        _buffer: &gpui::Entity<language::Buffer>,
1535        _cursor_position: language::Anchor,
1536        _cx: &gpui::App,
1537    ) -> bool {
1538        true
1539    }
1540
1541    fn is_refreshing(&self, _cx: &gpui::App) -> bool {
1542        false
1543    }
1544
1545    fn refresh(
1546        &mut self,
1547        _buffer: gpui::Entity<language::Buffer>,
1548        _cursor_position: language::Anchor,
1549        _debounce: bool,
1550        _cx: &mut gpui::Context<Self>,
1551    ) {
1552        self.refresh_count.fetch_add(1, atomic::Ordering::SeqCst);
1553    }
1554
1555    fn accept(&mut self, _cx: &mut gpui::Context<Self>) {}
1556
1557    fn discard(
1558        &mut self,
1559        _reason: edit_prediction_types::EditPredictionDiscardReason,
1560        _cx: &mut gpui::Context<Self>,
1561    ) {
1562        self.completion.take();
1563    }
1564
1565    fn suggest<'a>(
1566        &mut self,
1567        _buffer: &gpui::Entity<language::Buffer>,
1568        _cursor_position: language::Anchor,
1569        _cx: &mut gpui::Context<Self>,
1570    ) -> Option<edit_prediction_types::EditPrediction> {
1571        self.completion.clone()
1572    }
1573}
1574
1575#[derive(Default, Clone)]
1576pub struct FakeNonZedEditPredictionDelegate {
1577    pub completion: Option<edit_prediction_types::EditPrediction>,
1578}
1579
1580impl FakeNonZedEditPredictionDelegate {
1581    pub fn set_edit_prediction(
1582        &mut self,
1583        completion: Option<edit_prediction_types::EditPrediction>,
1584    ) {
1585        self.completion = completion;
1586    }
1587}
1588
1589impl EditPredictionDelegate for FakeNonZedEditPredictionDelegate {
1590    fn name() -> &'static str {
1591        "fake-non-zed-provider"
1592    }
1593
1594    fn display_name() -> &'static str {
1595        "Fake Non-Zed Provider"
1596    }
1597
1598    fn show_predictions_in_menu() -> bool {
1599        false
1600    }
1601
1602    fn supports_jump_to_edit() -> bool {
1603        false
1604    }
1605
1606    fn icons(&self, _cx: &gpui::App) -> EditPredictionIconSet {
1607        EditPredictionIconSet::new(IconName::ZedPredict)
1608    }
1609
1610    fn is_enabled(
1611        &self,
1612        _buffer: &gpui::Entity<language::Buffer>,
1613        _cursor_position: language::Anchor,
1614        _cx: &gpui::App,
1615    ) -> bool {
1616        true
1617    }
1618
1619    fn is_refreshing(&self, _cx: &gpui::App) -> bool {
1620        false
1621    }
1622
1623    fn refresh(
1624        &mut self,
1625        _buffer: gpui::Entity<language::Buffer>,
1626        _cursor_position: language::Anchor,
1627        _debounce: bool,
1628        _cx: &mut gpui::Context<Self>,
1629    ) {
1630    }
1631
1632    fn accept(&mut self, _cx: &mut gpui::Context<Self>) {}
1633
1634    fn discard(
1635        &mut self,
1636        _reason: edit_prediction_types::EditPredictionDiscardReason,
1637        _cx: &mut gpui::Context<Self>,
1638    ) {
1639        self.completion.take();
1640    }
1641
1642    fn suggest<'a>(
1643        &mut self,
1644        _buffer: &gpui::Entity<language::Buffer>,
1645        _cursor_position: language::Anchor,
1646        _cx: &mut gpui::Context<Self>,
1647    ) -> Option<edit_prediction_types::EditPrediction> {
1648        self.completion.clone()
1649    }
1650}