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