inline_assistant.rs

   1use std::cmp;
   2use std::mem;
   3use std::ops::Range;
   4use std::rc::Rc;
   5use std::sync::Arc;
   6
   7use anyhow::{Context as _, Result};
   8use assistant_settings::AssistantSettings;
   9use client::telemetry::Telemetry;
  10use collections::{HashMap, HashSet, VecDeque, hash_map};
  11use editor::{
  12    Anchor, AnchorRangeExt, CodeActionProvider, Editor, EditorEvent, ExcerptId, ExcerptRange,
  13    GutterDimensions, MultiBuffer, MultiBufferSnapshot, ToOffset as _, ToPoint,
  14    actions::SelectAll,
  15    display_map::{
  16        BlockContext, BlockPlacement, BlockProperties, BlockStyle, CustomBlockId, RenderBlock,
  17        ToDisplayPoint,
  18    },
  19};
  20use feature_flags::{Assistant2FeatureFlag, FeatureFlagViewExt as _};
  21use fs::Fs;
  22use gpui::{
  23    App, Context, Entity, Focusable, Global, HighlightStyle, Subscription, Task, UpdateGlobal,
  24    WeakEntity, Window, point,
  25};
  26use language::{Buffer, Point, Selection, TransactionId};
  27use language_model::{LanguageModelRegistry, report_assistant_event};
  28use multi_buffer::MultiBufferRow;
  29use parking_lot::Mutex;
  30use project::LspAction;
  31use project::{CodeAction, ProjectTransaction};
  32use prompt_store::PromptBuilder;
  33use settings::{Settings, SettingsStore};
  34use telemetry_events::{AssistantEventData, AssistantKind, AssistantPhase};
  35use terminal_view::{TerminalView, terminal_panel::TerminalPanel};
  36use text::{OffsetRangeExt, ToPoint as _};
  37use ui::prelude::*;
  38use util::RangeExt;
  39use util::ResultExt;
  40use workspace::{ItemHandle, Toast, Workspace, dock::Panel, notifications::NotificationId};
  41use zed_actions::agent::OpenConfiguration;
  42
  43use crate::AssistantPanel;
  44use crate::buffer_codegen::{BufferCodegen, CodegenAlternative, CodegenEvent};
  45use crate::context_store::ContextStore;
  46use crate::inline_prompt_editor::{CodegenStatus, InlineAssistId, PromptEditor, PromptEditorEvent};
  47use crate::terminal_inline_assistant::TerminalInlineAssistant;
  48use crate::thread_store::ThreadStore;
  49
  50pub fn init(
  51    fs: Arc<dyn Fs>,
  52    prompt_builder: Arc<PromptBuilder>,
  53    telemetry: Arc<Telemetry>,
  54    cx: &mut App,
  55) {
  56    cx.set_global(InlineAssistant::new(fs, prompt_builder, telemetry));
  57    cx.observe_new(|_workspace: &mut Workspace, window, cx| {
  58        let Some(window) = window else {
  59            return;
  60        };
  61        let workspace = cx.entity().clone();
  62        InlineAssistant::update_global(cx, |inline_assistant, cx| {
  63            inline_assistant.register_workspace(&workspace, window, cx)
  64        });
  65
  66        cx.observe_flag::<Assistant2FeatureFlag, _>(window, {
  67            |is_assistant2_enabled, _workspace, _window, cx| {
  68                InlineAssistant::update_global(cx, |inline_assistant, _cx| {
  69                    inline_assistant.is_assistant2_enabled = is_assistant2_enabled;
  70                });
  71            }
  72        })
  73        .detach();
  74    })
  75    .detach();
  76}
  77
  78const PROMPT_HISTORY_MAX_LEN: usize = 20;
  79
  80enum InlineAssistTarget {
  81    Editor(Entity<Editor>),
  82    Terminal(Entity<TerminalView>),
  83}
  84
  85pub struct InlineAssistant {
  86    next_assist_id: InlineAssistId,
  87    next_assist_group_id: InlineAssistGroupId,
  88    assists: HashMap<InlineAssistId, InlineAssist>,
  89    assists_by_editor: HashMap<WeakEntity<Editor>, EditorInlineAssists>,
  90    assist_groups: HashMap<InlineAssistGroupId, InlineAssistGroup>,
  91    confirmed_assists: HashMap<InlineAssistId, Entity<CodegenAlternative>>,
  92    prompt_history: VecDeque<String>,
  93    prompt_builder: Arc<PromptBuilder>,
  94    telemetry: Arc<Telemetry>,
  95    fs: Arc<dyn Fs>,
  96    is_assistant2_enabled: bool,
  97}
  98
  99impl Global for InlineAssistant {}
 100
 101impl InlineAssistant {
 102    pub fn new(
 103        fs: Arc<dyn Fs>,
 104        prompt_builder: Arc<PromptBuilder>,
 105        telemetry: Arc<Telemetry>,
 106    ) -> Self {
 107        Self {
 108            next_assist_id: InlineAssistId::default(),
 109            next_assist_group_id: InlineAssistGroupId::default(),
 110            assists: HashMap::default(),
 111            assists_by_editor: HashMap::default(),
 112            assist_groups: HashMap::default(),
 113            confirmed_assists: HashMap::default(),
 114            prompt_history: VecDeque::default(),
 115            prompt_builder,
 116            telemetry,
 117            fs,
 118            is_assistant2_enabled: false,
 119        }
 120    }
 121
 122    pub fn register_workspace(
 123        &mut self,
 124        workspace: &Entity<Workspace>,
 125        window: &mut Window,
 126        cx: &mut App,
 127    ) {
 128        window
 129            .subscribe(workspace, cx, |workspace, event, window, cx| {
 130                Self::update_global(cx, |this, cx| {
 131                    this.handle_workspace_event(workspace, event, window, cx)
 132                });
 133            })
 134            .detach();
 135
 136        let workspace = workspace.downgrade();
 137        cx.observe_global::<SettingsStore>(move |cx| {
 138            let Some(workspace) = workspace.upgrade() else {
 139                return;
 140            };
 141            let Some(terminal_panel) = workspace.read(cx).panel::<TerminalPanel>(cx) else {
 142                return;
 143            };
 144            let enabled = AssistantSettings::get_global(cx).enabled;
 145            terminal_panel.update(cx, |terminal_panel, cx| {
 146                terminal_panel.set_assistant_enabled(enabled, cx)
 147            });
 148        })
 149        .detach();
 150    }
 151
 152    fn handle_workspace_event(
 153        &mut self,
 154        workspace: Entity<Workspace>,
 155        event: &workspace::Event,
 156        window: &mut Window,
 157        cx: &mut App,
 158    ) {
 159        match event {
 160            workspace::Event::UserSavedItem { item, .. } => {
 161                // When the user manually saves an editor, automatically accepts all finished transformations.
 162                if let Some(editor) = item.upgrade().and_then(|item| item.act_as::<Editor>(cx)) {
 163                    if let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) {
 164                        for assist_id in editor_assists.assist_ids.clone() {
 165                            let assist = &self.assists[&assist_id];
 166                            if let CodegenStatus::Done = assist.codegen.read(cx).status(cx) {
 167                                self.finish_assist(assist_id, false, window, cx)
 168                            }
 169                        }
 170                    }
 171                }
 172            }
 173            workspace::Event::ItemAdded { item } => {
 174                self.register_workspace_item(&workspace, item.as_ref(), window, cx);
 175            }
 176            _ => (),
 177        }
 178    }
 179
 180    fn register_workspace_item(
 181        &mut self,
 182        workspace: &Entity<Workspace>,
 183        item: &dyn ItemHandle,
 184        window: &mut Window,
 185        cx: &mut App,
 186    ) {
 187        let is_assistant2_enabled = self.is_assistant2_enabled;
 188
 189        if let Some(editor) = item.act_as::<Editor>(cx) {
 190            editor.update(cx, |editor, cx| {
 191                if is_assistant2_enabled {
 192                    let thread_store = workspace
 193                        .read(cx)
 194                        .panel::<AssistantPanel>(cx)
 195                        .map(|assistant_panel| assistant_panel.read(cx).thread_store().downgrade());
 196
 197                    editor.add_code_action_provider(
 198                        Rc::new(AssistantCodeActionProvider {
 199                            editor: cx.entity().downgrade(),
 200                            workspace: workspace.downgrade(),
 201                            thread_store,
 202                        }),
 203                        window,
 204                        cx,
 205                    );
 206
 207                    // Remove the Assistant1 code action provider, as it still might be registered.
 208                    editor.remove_code_action_provider("assistant".into(), window, cx);
 209                } else {
 210                    editor.remove_code_action_provider(
 211                        ASSISTANT_CODE_ACTION_PROVIDER_ID.into(),
 212                        window,
 213                        cx,
 214                    );
 215                }
 216            });
 217        }
 218    }
 219
 220    pub fn inline_assist(
 221        workspace: &mut Workspace,
 222        _action: &zed_actions::assistant::InlineAssist,
 223        window: &mut Window,
 224        cx: &mut Context<Workspace>,
 225    ) {
 226        let settings = AssistantSettings::get_global(cx);
 227        if !settings.enabled {
 228            return;
 229        }
 230
 231        let Some(inline_assist_target) = Self::resolve_inline_assist_target(
 232            workspace,
 233            workspace.panel::<AssistantPanel>(cx),
 234            window,
 235            cx,
 236        ) else {
 237            return;
 238        };
 239
 240        let is_authenticated = || {
 241            LanguageModelRegistry::read_global(cx)
 242                .inline_assistant_model()
 243                .map_or(false, |model| model.provider.is_authenticated(cx))
 244        };
 245
 246        let thread_store = workspace
 247            .panel::<AssistantPanel>(cx)
 248            .map(|assistant_panel| assistant_panel.read(cx).thread_store().downgrade());
 249
 250        let handle_assist =
 251            |window: &mut Window, cx: &mut Context<Workspace>| match inline_assist_target {
 252                InlineAssistTarget::Editor(active_editor) => {
 253                    InlineAssistant::update_global(cx, |assistant, cx| {
 254                        assistant.assist(
 255                            &active_editor,
 256                            cx.entity().downgrade(),
 257                            thread_store,
 258                            window,
 259                            cx,
 260                        )
 261                    })
 262                }
 263                InlineAssistTarget::Terminal(active_terminal) => {
 264                    TerminalInlineAssistant::update_global(cx, |assistant, cx| {
 265                        assistant.assist(&active_terminal, cx.entity(), thread_store, window, cx)
 266                    })
 267                }
 268            };
 269
 270        if is_authenticated() {
 271            handle_assist(window, cx);
 272        } else {
 273            cx.spawn_in(window, async move |_workspace, cx| {
 274                let Some(task) = cx.update(|_, cx| {
 275                    LanguageModelRegistry::read_global(cx)
 276                        .inline_assistant_model()
 277                        .map_or(None, |model| Some(model.provider.authenticate(cx)))
 278                })?
 279                else {
 280                    let answer = cx
 281                        .prompt(
 282                            gpui::PromptLevel::Warning,
 283                            "No language model provider configured",
 284                            None,
 285                            &["Configure", "Cancel"],
 286                        )
 287                        .await
 288                        .ok();
 289                    if let Some(answer) = answer {
 290                        if answer == 0 {
 291                            cx.update(|window, cx| {
 292                                window.dispatch_action(Box::new(OpenConfiguration), cx)
 293                            })
 294                            .ok();
 295                        }
 296                    }
 297                    return Ok(());
 298                };
 299                task.await?;
 300
 301                anyhow::Ok(())
 302            })
 303            .detach_and_log_err(cx);
 304
 305            if is_authenticated() {
 306                handle_assist(window, cx);
 307            }
 308        }
 309    }
 310
 311    pub fn assist(
 312        &mut self,
 313        editor: &Entity<Editor>,
 314        workspace: WeakEntity<Workspace>,
 315        thread_store: Option<WeakEntity<ThreadStore>>,
 316        window: &mut Window,
 317        cx: &mut App,
 318    ) {
 319        let Some(project) = workspace
 320            .upgrade()
 321            .map(|workspace| workspace.read(cx).project().downgrade())
 322        else {
 323            return;
 324        };
 325
 326        let (snapshot, initial_selections) = editor.update(cx, |editor, cx| {
 327            (
 328                editor.snapshot(window, cx),
 329                editor.selections.all::<Point>(cx),
 330            )
 331        });
 332
 333        let mut selections = Vec::<Selection<Point>>::new();
 334        let mut newest_selection = None;
 335        for mut selection in initial_selections {
 336            if selection.end > selection.start {
 337                selection.start.column = 0;
 338                // If the selection ends at the start of the line, we don't want to include it.
 339                if selection.end.column == 0 {
 340                    selection.end.row -= 1;
 341                }
 342                selection.end.column = snapshot
 343                    .buffer_snapshot
 344                    .line_len(MultiBufferRow(selection.end.row));
 345            } else if let Some(fold) =
 346                snapshot.crease_for_buffer_row(MultiBufferRow(selection.end.row))
 347            {
 348                selection.start = fold.range().start;
 349                selection.end = fold.range().end;
 350                if MultiBufferRow(selection.end.row) < snapshot.buffer_snapshot.max_row() {
 351                    let chars = snapshot
 352                        .buffer_snapshot
 353                        .chars_at(Point::new(selection.end.row + 1, 0));
 354
 355                    for c in chars {
 356                        if c == '\n' {
 357                            break;
 358                        }
 359                        if c.is_whitespace() {
 360                            continue;
 361                        }
 362                        if snapshot
 363                            .language_at(selection.end)
 364                            .is_some_and(|language| language.config().brackets.is_closing_brace(c))
 365                        {
 366                            selection.end.row += 1;
 367                            selection.end.column = snapshot
 368                                .buffer_snapshot
 369                                .line_len(MultiBufferRow(selection.end.row));
 370                        }
 371                    }
 372                }
 373            }
 374
 375            if let Some(prev_selection) = selections.last_mut() {
 376                if selection.start <= prev_selection.end {
 377                    prev_selection.end = selection.end;
 378                    continue;
 379                }
 380            }
 381
 382            let latest_selection = newest_selection.get_or_insert_with(|| selection.clone());
 383            if selection.id > latest_selection.id {
 384                *latest_selection = selection.clone();
 385            }
 386            selections.push(selection);
 387        }
 388        let snapshot = &snapshot.buffer_snapshot;
 389        let newest_selection = newest_selection.unwrap();
 390
 391        let mut codegen_ranges = Vec::new();
 392        for (buffer, buffer_range, excerpt_id) in
 393            snapshot.ranges_to_buffer_ranges(selections.iter().map(|selection| {
 394                snapshot.anchor_before(selection.start)..snapshot.anchor_after(selection.end)
 395            }))
 396        {
 397            let anchor_range = Anchor::range_in_buffer(
 398                excerpt_id,
 399                buffer.remote_id(),
 400                buffer.anchor_before(buffer_range.start)..buffer.anchor_after(buffer_range.end),
 401            );
 402
 403            codegen_ranges.push(anchor_range);
 404
 405            if let Some(model) = LanguageModelRegistry::read_global(cx).inline_assistant_model() {
 406                self.telemetry.report_assistant_event(AssistantEventData {
 407                    conversation_id: None,
 408                    kind: AssistantKind::Inline,
 409                    phase: AssistantPhase::Invoked,
 410                    message_id: None,
 411                    model: model.model.telemetry_id(),
 412                    model_provider: model.provider.id().to_string(),
 413                    response_latency: None,
 414                    error_message: None,
 415                    language_name: buffer.language().map(|language| language.name().to_proto()),
 416                });
 417            }
 418        }
 419
 420        let assist_group_id = self.next_assist_group_id.post_inc();
 421        let prompt_buffer =
 422            cx.new(|cx| MultiBuffer::singleton(cx.new(|cx| Buffer::local(String::new(), cx)), cx));
 423
 424        let mut assists = Vec::new();
 425        let mut assist_to_focus = None;
 426        for range in codegen_ranges {
 427            let assist_id = self.next_assist_id.post_inc();
 428            let context_store =
 429                cx.new(|_cx| ContextStore::new(project.clone(), thread_store.clone()));
 430            let codegen = cx.new(|cx| {
 431                BufferCodegen::new(
 432                    editor.read(cx).buffer().clone(),
 433                    range.clone(),
 434                    None,
 435                    context_store.clone(),
 436                    self.telemetry.clone(),
 437                    self.prompt_builder.clone(),
 438                    cx,
 439                )
 440            });
 441
 442            let gutter_dimensions = Arc::new(Mutex::new(GutterDimensions::default()));
 443            let prompt_editor = cx.new(|cx| {
 444                PromptEditor::new_buffer(
 445                    assist_id,
 446                    gutter_dimensions.clone(),
 447                    self.prompt_history.clone(),
 448                    prompt_buffer.clone(),
 449                    codegen.clone(),
 450                    self.fs.clone(),
 451                    context_store,
 452                    workspace.clone(),
 453                    thread_store.clone(),
 454                    window,
 455                    cx,
 456                )
 457            });
 458
 459            if assist_to_focus.is_none() {
 460                let focus_assist = if newest_selection.reversed {
 461                    range.start.to_point(&snapshot) == newest_selection.start
 462                } else {
 463                    range.end.to_point(&snapshot) == newest_selection.end
 464                };
 465                if focus_assist {
 466                    assist_to_focus = Some(assist_id);
 467                }
 468            }
 469
 470            let [prompt_block_id, end_block_id] =
 471                self.insert_assist_blocks(editor, &range, &prompt_editor, cx);
 472
 473            assists.push((
 474                assist_id,
 475                range,
 476                prompt_editor,
 477                prompt_block_id,
 478                end_block_id,
 479            ));
 480        }
 481
 482        let editor_assists = self
 483            .assists_by_editor
 484            .entry(editor.downgrade())
 485            .or_insert_with(|| EditorInlineAssists::new(&editor, window, cx));
 486        let mut assist_group = InlineAssistGroup::new();
 487        for (assist_id, range, prompt_editor, prompt_block_id, end_block_id) in assists {
 488            let codegen = prompt_editor.read(cx).codegen().clone();
 489
 490            self.assists.insert(
 491                assist_id,
 492                InlineAssist::new(
 493                    assist_id,
 494                    assist_group_id,
 495                    editor,
 496                    &prompt_editor,
 497                    prompt_block_id,
 498                    end_block_id,
 499                    range,
 500                    codegen,
 501                    workspace.clone(),
 502                    window,
 503                    cx,
 504                ),
 505            );
 506            assist_group.assist_ids.push(assist_id);
 507            editor_assists.assist_ids.push(assist_id);
 508        }
 509        self.assist_groups.insert(assist_group_id, assist_group);
 510
 511        if let Some(assist_id) = assist_to_focus {
 512            self.focus_assist(assist_id, window, cx);
 513        }
 514    }
 515
 516    pub fn suggest_assist(
 517        &mut self,
 518        editor: &Entity<Editor>,
 519        mut range: Range<Anchor>,
 520        initial_prompt: String,
 521        initial_transaction_id: Option<TransactionId>,
 522        focus: bool,
 523        workspace: Entity<Workspace>,
 524        thread_store: Option<WeakEntity<ThreadStore>>,
 525        window: &mut Window,
 526        cx: &mut App,
 527    ) -> InlineAssistId {
 528        let assist_group_id = self.next_assist_group_id.post_inc();
 529        let prompt_buffer = cx.new(|cx| Buffer::local(&initial_prompt, cx));
 530        let prompt_buffer = cx.new(|cx| MultiBuffer::singleton(prompt_buffer, cx));
 531
 532        let assist_id = self.next_assist_id.post_inc();
 533
 534        let buffer = editor.read(cx).buffer().clone();
 535        {
 536            let snapshot = buffer.read(cx).read(cx);
 537            range.start = range.start.bias_left(&snapshot);
 538            range.end = range.end.bias_right(&snapshot);
 539        }
 540
 541        let project = workspace.read(cx).project().downgrade();
 542        let context_store = cx.new(|_cx| ContextStore::new(project, thread_store.clone()));
 543
 544        let codegen = cx.new(|cx| {
 545            BufferCodegen::new(
 546                editor.read(cx).buffer().clone(),
 547                range.clone(),
 548                initial_transaction_id,
 549                context_store.clone(),
 550                self.telemetry.clone(),
 551                self.prompt_builder.clone(),
 552                cx,
 553            )
 554        });
 555
 556        let gutter_dimensions = Arc::new(Mutex::new(GutterDimensions::default()));
 557        let prompt_editor = cx.new(|cx| {
 558            PromptEditor::new_buffer(
 559                assist_id,
 560                gutter_dimensions.clone(),
 561                self.prompt_history.clone(),
 562                prompt_buffer.clone(),
 563                codegen.clone(),
 564                self.fs.clone(),
 565                context_store,
 566                workspace.downgrade(),
 567                thread_store,
 568                window,
 569                cx,
 570            )
 571        });
 572
 573        let [prompt_block_id, end_block_id] =
 574            self.insert_assist_blocks(editor, &range, &prompt_editor, cx);
 575
 576        let editor_assists = self
 577            .assists_by_editor
 578            .entry(editor.downgrade())
 579            .or_insert_with(|| EditorInlineAssists::new(&editor, window, cx));
 580
 581        let mut assist_group = InlineAssistGroup::new();
 582        self.assists.insert(
 583            assist_id,
 584            InlineAssist::new(
 585                assist_id,
 586                assist_group_id,
 587                editor,
 588                &prompt_editor,
 589                prompt_block_id,
 590                end_block_id,
 591                range,
 592                codegen.clone(),
 593                workspace.downgrade(),
 594                window,
 595                cx,
 596            ),
 597        );
 598        assist_group.assist_ids.push(assist_id);
 599        editor_assists.assist_ids.push(assist_id);
 600        self.assist_groups.insert(assist_group_id, assist_group);
 601
 602        if focus {
 603            self.focus_assist(assist_id, window, cx);
 604        }
 605
 606        assist_id
 607    }
 608
 609    fn insert_assist_blocks(
 610        &self,
 611        editor: &Entity<Editor>,
 612        range: &Range<Anchor>,
 613        prompt_editor: &Entity<PromptEditor<BufferCodegen>>,
 614        cx: &mut App,
 615    ) -> [CustomBlockId; 2] {
 616        let prompt_editor_height = prompt_editor.update(cx, |prompt_editor, cx| {
 617            prompt_editor
 618                .editor
 619                .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1 + 2)
 620        });
 621        let assist_blocks = vec![
 622            BlockProperties {
 623                style: BlockStyle::Sticky,
 624                placement: BlockPlacement::Above(range.start),
 625                height: Some(prompt_editor_height),
 626                render: build_assist_editor_renderer(prompt_editor),
 627                priority: 0,
 628            },
 629            BlockProperties {
 630                style: BlockStyle::Sticky,
 631                placement: BlockPlacement::Below(range.end),
 632                height: None,
 633                render: Arc::new(|cx| {
 634                    v_flex()
 635                        .h_full()
 636                        .w_full()
 637                        .border_t_1()
 638                        .border_color(cx.theme().status().info_border)
 639                        .into_any_element()
 640                }),
 641                priority: 0,
 642            },
 643        ];
 644
 645        editor.update(cx, |editor, cx| {
 646            let block_ids = editor.insert_blocks(assist_blocks, None, cx);
 647            [block_ids[0], block_ids[1]]
 648        })
 649    }
 650
 651    fn handle_prompt_editor_focus_in(&mut self, assist_id: InlineAssistId, cx: &mut App) {
 652        let assist = &self.assists[&assist_id];
 653        let Some(decorations) = assist.decorations.as_ref() else {
 654            return;
 655        };
 656        let assist_group = self.assist_groups.get_mut(&assist.group_id).unwrap();
 657        let editor_assists = self.assists_by_editor.get_mut(&assist.editor).unwrap();
 658
 659        assist_group.active_assist_id = Some(assist_id);
 660        if assist_group.linked {
 661            for assist_id in &assist_group.assist_ids {
 662                if let Some(decorations) = self.assists[assist_id].decorations.as_ref() {
 663                    decorations.prompt_editor.update(cx, |prompt_editor, cx| {
 664                        prompt_editor.set_show_cursor_when_unfocused(true, cx)
 665                    });
 666                }
 667            }
 668        }
 669
 670        assist
 671            .editor
 672            .update(cx, |editor, cx| {
 673                let scroll_top = editor.scroll_position(cx).y;
 674                let scroll_bottom = scroll_top + editor.visible_line_count().unwrap_or(0.);
 675                let prompt_row = editor
 676                    .row_for_block(decorations.prompt_block_id, cx)
 677                    .unwrap()
 678                    .0 as f32;
 679
 680                if (scroll_top..scroll_bottom).contains(&prompt_row) {
 681                    editor_assists.scroll_lock = Some(InlineAssistScrollLock {
 682                        assist_id,
 683                        distance_from_top: prompt_row - scroll_top,
 684                    });
 685                } else {
 686                    editor_assists.scroll_lock = None;
 687                }
 688            })
 689            .ok();
 690    }
 691
 692    fn handle_prompt_editor_focus_out(&mut self, assist_id: InlineAssistId, cx: &mut App) {
 693        let assist = &self.assists[&assist_id];
 694        let assist_group = self.assist_groups.get_mut(&assist.group_id).unwrap();
 695        if assist_group.active_assist_id == Some(assist_id) {
 696            assist_group.active_assist_id = None;
 697            if assist_group.linked {
 698                for assist_id in &assist_group.assist_ids {
 699                    if let Some(decorations) = self.assists[assist_id].decorations.as_ref() {
 700                        decorations.prompt_editor.update(cx, |prompt_editor, cx| {
 701                            prompt_editor.set_show_cursor_when_unfocused(false, cx)
 702                        });
 703                    }
 704                }
 705            }
 706        }
 707    }
 708
 709    fn handle_prompt_editor_event(
 710        &mut self,
 711        prompt_editor: Entity<PromptEditor<BufferCodegen>>,
 712        event: &PromptEditorEvent,
 713        window: &mut Window,
 714        cx: &mut App,
 715    ) {
 716        let assist_id = prompt_editor.read(cx).id();
 717        match event {
 718            PromptEditorEvent::StartRequested => {
 719                self.start_assist(assist_id, window, cx);
 720            }
 721            PromptEditorEvent::StopRequested => {
 722                self.stop_assist(assist_id, cx);
 723            }
 724            PromptEditorEvent::ConfirmRequested { execute: _ } => {
 725                self.finish_assist(assist_id, false, window, cx);
 726            }
 727            PromptEditorEvent::CancelRequested => {
 728                self.finish_assist(assist_id, true, window, cx);
 729            }
 730            PromptEditorEvent::DismissRequested => {
 731                self.dismiss_assist(assist_id, window, cx);
 732            }
 733            PromptEditorEvent::Resized { .. } => {
 734                // This only matters for the terminal inline assistant
 735            }
 736        }
 737    }
 738
 739    fn handle_editor_newline(&mut self, editor: Entity<Editor>, window: &mut Window, cx: &mut App) {
 740        let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
 741            return;
 742        };
 743
 744        if editor.read(cx).selections.count() == 1 {
 745            let (selection, buffer) = editor.update(cx, |editor, cx| {
 746                (
 747                    editor.selections.newest::<usize>(cx),
 748                    editor.buffer().read(cx).snapshot(cx),
 749                )
 750            });
 751            for assist_id in &editor_assists.assist_ids {
 752                let assist = &self.assists[assist_id];
 753                let assist_range = assist.range.to_offset(&buffer);
 754                if assist_range.contains(&selection.start) && assist_range.contains(&selection.end)
 755                {
 756                    if matches!(assist.codegen.read(cx).status(cx), CodegenStatus::Pending) {
 757                        self.dismiss_assist(*assist_id, window, cx);
 758                    } else {
 759                        self.finish_assist(*assist_id, false, window, cx);
 760                    }
 761
 762                    return;
 763                }
 764            }
 765        }
 766
 767        cx.propagate();
 768    }
 769
 770    fn handle_editor_cancel(&mut self, editor: Entity<Editor>, window: &mut Window, cx: &mut App) {
 771        let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
 772            return;
 773        };
 774
 775        if editor.read(cx).selections.count() == 1 {
 776            let (selection, buffer) = editor.update(cx, |editor, cx| {
 777                (
 778                    editor.selections.newest::<usize>(cx),
 779                    editor.buffer().read(cx).snapshot(cx),
 780                )
 781            });
 782            let mut closest_assist_fallback = None;
 783            for assist_id in &editor_assists.assist_ids {
 784                let assist = &self.assists[assist_id];
 785                let assist_range = assist.range.to_offset(&buffer);
 786                if assist.decorations.is_some() {
 787                    if assist_range.contains(&selection.start)
 788                        && assist_range.contains(&selection.end)
 789                    {
 790                        self.focus_assist(*assist_id, window, cx);
 791                        return;
 792                    } else {
 793                        let distance_from_selection = assist_range
 794                            .start
 795                            .abs_diff(selection.start)
 796                            .min(assist_range.start.abs_diff(selection.end))
 797                            + assist_range
 798                                .end
 799                                .abs_diff(selection.start)
 800                                .min(assist_range.end.abs_diff(selection.end));
 801                        match closest_assist_fallback {
 802                            Some((_, old_distance)) => {
 803                                if distance_from_selection < old_distance {
 804                                    closest_assist_fallback =
 805                                        Some((assist_id, distance_from_selection));
 806                                }
 807                            }
 808                            None => {
 809                                closest_assist_fallback = Some((assist_id, distance_from_selection))
 810                            }
 811                        }
 812                    }
 813                }
 814            }
 815
 816            if let Some((&assist_id, _)) = closest_assist_fallback {
 817                self.focus_assist(assist_id, window, cx);
 818            }
 819        }
 820
 821        cx.propagate();
 822    }
 823
 824    fn handle_editor_release(
 825        &mut self,
 826        editor: WeakEntity<Editor>,
 827        window: &mut Window,
 828        cx: &mut App,
 829    ) {
 830        if let Some(editor_assists) = self.assists_by_editor.get_mut(&editor) {
 831            for assist_id in editor_assists.assist_ids.clone() {
 832                self.finish_assist(assist_id, true, window, cx);
 833            }
 834        }
 835    }
 836
 837    fn handle_editor_change(&mut self, editor: Entity<Editor>, window: &mut Window, cx: &mut App) {
 838        let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
 839            return;
 840        };
 841        let Some(scroll_lock) = editor_assists.scroll_lock.as_ref() else {
 842            return;
 843        };
 844        let assist = &self.assists[&scroll_lock.assist_id];
 845        let Some(decorations) = assist.decorations.as_ref() else {
 846            return;
 847        };
 848
 849        editor.update(cx, |editor, cx| {
 850            let scroll_position = editor.scroll_position(cx);
 851            let target_scroll_top = editor
 852                .row_for_block(decorations.prompt_block_id, cx)
 853                .unwrap()
 854                .0 as f32
 855                - scroll_lock.distance_from_top;
 856            if target_scroll_top != scroll_position.y {
 857                editor.set_scroll_position(point(scroll_position.x, target_scroll_top), window, cx);
 858            }
 859        });
 860    }
 861
 862    fn handle_editor_event(
 863        &mut self,
 864        editor: Entity<Editor>,
 865        event: &EditorEvent,
 866        window: &mut Window,
 867        cx: &mut App,
 868    ) {
 869        let Some(editor_assists) = self.assists_by_editor.get_mut(&editor.downgrade()) else {
 870            return;
 871        };
 872
 873        match event {
 874            EditorEvent::Edited { transaction_id } => {
 875                let buffer = editor.read(cx).buffer().read(cx);
 876                let edited_ranges =
 877                    buffer.edited_ranges_for_transaction::<usize>(*transaction_id, cx);
 878                let snapshot = buffer.snapshot(cx);
 879
 880                for assist_id in editor_assists.assist_ids.clone() {
 881                    let assist = &self.assists[&assist_id];
 882                    if matches!(
 883                        assist.codegen.read(cx).status(cx),
 884                        CodegenStatus::Error(_) | CodegenStatus::Done
 885                    ) {
 886                        let assist_range = assist.range.to_offset(&snapshot);
 887                        if edited_ranges
 888                            .iter()
 889                            .any(|range| range.overlaps(&assist_range))
 890                        {
 891                            self.finish_assist(assist_id, false, window, cx);
 892                        }
 893                    }
 894                }
 895            }
 896            EditorEvent::ScrollPositionChanged { .. } => {
 897                if let Some(scroll_lock) = editor_assists.scroll_lock.as_ref() {
 898                    let assist = &self.assists[&scroll_lock.assist_id];
 899                    if let Some(decorations) = assist.decorations.as_ref() {
 900                        let distance_from_top = editor.update(cx, |editor, cx| {
 901                            let scroll_top = editor.scroll_position(cx).y;
 902                            let prompt_row = editor
 903                                .row_for_block(decorations.prompt_block_id, cx)
 904                                .unwrap()
 905                                .0 as f32;
 906                            prompt_row - scroll_top
 907                        });
 908
 909                        if distance_from_top != scroll_lock.distance_from_top {
 910                            editor_assists.scroll_lock = None;
 911                        }
 912                    }
 913                }
 914            }
 915            EditorEvent::SelectionsChanged { .. } => {
 916                for assist_id in editor_assists.assist_ids.clone() {
 917                    let assist = &self.assists[&assist_id];
 918                    if let Some(decorations) = assist.decorations.as_ref() {
 919                        if decorations
 920                            .prompt_editor
 921                            .focus_handle(cx)
 922                            .is_focused(window)
 923                        {
 924                            return;
 925                        }
 926                    }
 927                }
 928
 929                editor_assists.scroll_lock = None;
 930            }
 931            _ => {}
 932        }
 933    }
 934
 935    pub fn finish_assist(
 936        &mut self,
 937        assist_id: InlineAssistId,
 938        undo: bool,
 939        window: &mut Window,
 940        cx: &mut App,
 941    ) {
 942        if let Some(assist) = self.assists.get(&assist_id) {
 943            let assist_group_id = assist.group_id;
 944            if self.assist_groups[&assist_group_id].linked {
 945                for assist_id in self.unlink_assist_group(assist_group_id, window, cx) {
 946                    self.finish_assist(assist_id, undo, window, cx);
 947                }
 948                return;
 949            }
 950        }
 951
 952        self.dismiss_assist(assist_id, window, cx);
 953
 954        if let Some(assist) = self.assists.remove(&assist_id) {
 955            if let hash_map::Entry::Occupied(mut entry) = self.assist_groups.entry(assist.group_id)
 956            {
 957                entry.get_mut().assist_ids.retain(|id| *id != assist_id);
 958                if entry.get().assist_ids.is_empty() {
 959                    entry.remove();
 960                }
 961            }
 962
 963            if let hash_map::Entry::Occupied(mut entry) =
 964                self.assists_by_editor.entry(assist.editor.clone())
 965            {
 966                entry.get_mut().assist_ids.retain(|id| *id != assist_id);
 967                if entry.get().assist_ids.is_empty() {
 968                    entry.remove();
 969                    if let Some(editor) = assist.editor.upgrade() {
 970                        self.update_editor_highlights(&editor, cx);
 971                    }
 972                } else {
 973                    entry.get().highlight_updates.send(()).ok();
 974                }
 975            }
 976
 977            let active_alternative = assist.codegen.read(cx).active_alternative().clone();
 978            let message_id = active_alternative.read(cx).message_id.clone();
 979
 980            if let Some(model) = LanguageModelRegistry::read_global(cx).inline_assistant_model() {
 981                let language_name = assist.editor.upgrade().and_then(|editor| {
 982                    let multibuffer = editor.read(cx).buffer().read(cx);
 983                    let snapshot = multibuffer.snapshot(cx);
 984                    let ranges = snapshot.range_to_buffer_ranges(assist.range.clone());
 985                    ranges
 986                        .first()
 987                        .and_then(|(buffer, _, _)| buffer.language())
 988                        .map(|language| language.name())
 989                });
 990                report_assistant_event(
 991                    AssistantEventData {
 992                        conversation_id: None,
 993                        kind: AssistantKind::Inline,
 994                        message_id,
 995                        phase: if undo {
 996                            AssistantPhase::Rejected
 997                        } else {
 998                            AssistantPhase::Accepted
 999                        },
1000                        model: model.model.telemetry_id(),
1001                        model_provider: model.model.provider_id().to_string(),
1002                        response_latency: None,
1003                        error_message: None,
1004                        language_name: language_name.map(|name| name.to_proto()),
1005                    },
1006                    Some(self.telemetry.clone()),
1007                    cx.http_client(),
1008                    model.model.api_key(cx),
1009                    cx.background_executor(),
1010                );
1011            }
1012
1013            if undo {
1014                assist.codegen.update(cx, |codegen, cx| codegen.undo(cx));
1015            } else {
1016                self.confirmed_assists.insert(assist_id, active_alternative);
1017            }
1018        }
1019    }
1020
1021    fn dismiss_assist(
1022        &mut self,
1023        assist_id: InlineAssistId,
1024        window: &mut Window,
1025        cx: &mut App,
1026    ) -> bool {
1027        let Some(assist) = self.assists.get_mut(&assist_id) else {
1028            return false;
1029        };
1030        let Some(editor) = assist.editor.upgrade() else {
1031            return false;
1032        };
1033        let Some(decorations) = assist.decorations.take() else {
1034            return false;
1035        };
1036
1037        editor.update(cx, |editor, cx| {
1038            let mut to_remove = decorations.removed_line_block_ids;
1039            to_remove.insert(decorations.prompt_block_id);
1040            to_remove.insert(decorations.end_block_id);
1041            editor.remove_blocks(to_remove, None, cx);
1042        });
1043
1044        if decorations
1045            .prompt_editor
1046            .focus_handle(cx)
1047            .contains_focused(window, cx)
1048        {
1049            self.focus_next_assist(assist_id, window, cx);
1050        }
1051
1052        if let Some(editor_assists) = self.assists_by_editor.get_mut(&editor.downgrade()) {
1053            if editor_assists
1054                .scroll_lock
1055                .as_ref()
1056                .map_or(false, |lock| lock.assist_id == assist_id)
1057            {
1058                editor_assists.scroll_lock = None;
1059            }
1060            editor_assists.highlight_updates.send(()).ok();
1061        }
1062
1063        true
1064    }
1065
1066    fn focus_next_assist(&mut self, assist_id: InlineAssistId, window: &mut Window, cx: &mut App) {
1067        let Some(assist) = self.assists.get(&assist_id) else {
1068            return;
1069        };
1070
1071        let assist_group = &self.assist_groups[&assist.group_id];
1072        let assist_ix = assist_group
1073            .assist_ids
1074            .iter()
1075            .position(|id| *id == assist_id)
1076            .unwrap();
1077        let assist_ids = assist_group
1078            .assist_ids
1079            .iter()
1080            .skip(assist_ix + 1)
1081            .chain(assist_group.assist_ids.iter().take(assist_ix));
1082
1083        for assist_id in assist_ids {
1084            let assist = &self.assists[assist_id];
1085            if assist.decorations.is_some() {
1086                self.focus_assist(*assist_id, window, cx);
1087                return;
1088            }
1089        }
1090
1091        assist
1092            .editor
1093            .update(cx, |editor, cx| window.focus(&editor.focus_handle(cx)))
1094            .ok();
1095    }
1096
1097    fn focus_assist(&mut self, assist_id: InlineAssistId, window: &mut Window, cx: &mut App) {
1098        let Some(assist) = self.assists.get(&assist_id) else {
1099            return;
1100        };
1101
1102        if let Some(decorations) = assist.decorations.as_ref() {
1103            decorations.prompt_editor.update(cx, |prompt_editor, cx| {
1104                prompt_editor.editor.update(cx, |editor, cx| {
1105                    window.focus(&editor.focus_handle(cx));
1106                    editor.select_all(&SelectAll, window, cx);
1107                })
1108            });
1109        }
1110
1111        self.scroll_to_assist(assist_id, window, cx);
1112    }
1113
1114    pub fn scroll_to_assist(
1115        &mut self,
1116        assist_id: InlineAssistId,
1117        window: &mut Window,
1118        cx: &mut App,
1119    ) {
1120        let Some(assist) = self.assists.get(&assist_id) else {
1121            return;
1122        };
1123        let Some(editor) = assist.editor.upgrade() else {
1124            return;
1125        };
1126
1127        let position = assist.range.start;
1128        editor.update(cx, |editor, cx| {
1129            editor.change_selections(None, window, cx, |selections| {
1130                selections.select_anchor_ranges([position..position])
1131            });
1132
1133            let mut scroll_target_top;
1134            let mut scroll_target_bottom;
1135            if let Some(decorations) = assist.decorations.as_ref() {
1136                scroll_target_top = editor
1137                    .row_for_block(decorations.prompt_block_id, cx)
1138                    .unwrap()
1139                    .0 as f32;
1140                scroll_target_bottom = editor
1141                    .row_for_block(decorations.end_block_id, cx)
1142                    .unwrap()
1143                    .0 as f32;
1144            } else {
1145                let snapshot = editor.snapshot(window, cx);
1146                let start_row = assist
1147                    .range
1148                    .start
1149                    .to_display_point(&snapshot.display_snapshot)
1150                    .row();
1151                scroll_target_top = start_row.0 as f32;
1152                scroll_target_bottom = scroll_target_top + 1.;
1153            }
1154            scroll_target_top -= editor.vertical_scroll_margin() as f32;
1155            scroll_target_bottom += editor.vertical_scroll_margin() as f32;
1156
1157            let height_in_lines = editor.visible_line_count().unwrap_or(0.);
1158            let scroll_top = editor.scroll_position(cx).y;
1159            let scroll_bottom = scroll_top + height_in_lines;
1160
1161            if scroll_target_top < scroll_top {
1162                editor.set_scroll_position(point(0., scroll_target_top), window, cx);
1163            } else if scroll_target_bottom > scroll_bottom {
1164                if (scroll_target_bottom - scroll_target_top) <= height_in_lines {
1165                    editor.set_scroll_position(
1166                        point(0., scroll_target_bottom - height_in_lines),
1167                        window,
1168                        cx,
1169                    );
1170                } else {
1171                    editor.set_scroll_position(point(0., scroll_target_top), window, cx);
1172                }
1173            }
1174        });
1175    }
1176
1177    fn unlink_assist_group(
1178        &mut self,
1179        assist_group_id: InlineAssistGroupId,
1180        window: &mut Window,
1181        cx: &mut App,
1182    ) -> Vec<InlineAssistId> {
1183        let assist_group = self.assist_groups.get_mut(&assist_group_id).unwrap();
1184        assist_group.linked = false;
1185        for assist_id in &assist_group.assist_ids {
1186            let assist = self.assists.get_mut(assist_id).unwrap();
1187            if let Some(editor_decorations) = assist.decorations.as_ref() {
1188                editor_decorations
1189                    .prompt_editor
1190                    .update(cx, |prompt_editor, cx| prompt_editor.unlink(window, cx));
1191            }
1192        }
1193        assist_group.assist_ids.clone()
1194    }
1195
1196    pub fn start_assist(&mut self, assist_id: InlineAssistId, window: &mut Window, cx: &mut App) {
1197        let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
1198            assist
1199        } else {
1200            return;
1201        };
1202
1203        let assist_group_id = assist.group_id;
1204        if self.assist_groups[&assist_group_id].linked {
1205            for assist_id in self.unlink_assist_group(assist_group_id, window, cx) {
1206                self.start_assist(assist_id, window, cx);
1207            }
1208            return;
1209        }
1210
1211        let Some(user_prompt) = assist.user_prompt(cx) else {
1212            return;
1213        };
1214
1215        self.prompt_history.retain(|prompt| *prompt != user_prompt);
1216        self.prompt_history.push_back(user_prompt.clone());
1217        if self.prompt_history.len() > PROMPT_HISTORY_MAX_LEN {
1218            self.prompt_history.pop_front();
1219        }
1220
1221        assist
1222            .codegen
1223            .update(cx, |codegen, cx| codegen.start(user_prompt, cx))
1224            .log_err();
1225    }
1226
1227    pub fn stop_assist(&mut self, assist_id: InlineAssistId, cx: &mut App) {
1228        let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
1229            assist
1230        } else {
1231            return;
1232        };
1233
1234        assist.codegen.update(cx, |codegen, cx| codegen.stop(cx));
1235    }
1236
1237    fn update_editor_highlights(&self, editor: &Entity<Editor>, cx: &mut App) {
1238        let mut gutter_pending_ranges = Vec::new();
1239        let mut gutter_transformed_ranges = Vec::new();
1240        let mut foreground_ranges = Vec::new();
1241        let mut inserted_row_ranges = Vec::new();
1242        let empty_assist_ids = Vec::new();
1243        let assist_ids = self
1244            .assists_by_editor
1245            .get(&editor.downgrade())
1246            .map_or(&empty_assist_ids, |editor_assists| {
1247                &editor_assists.assist_ids
1248            });
1249
1250        for assist_id in assist_ids {
1251            if let Some(assist) = self.assists.get(assist_id) {
1252                let codegen = assist.codegen.read(cx);
1253                let buffer = codegen.buffer(cx).read(cx).read(cx);
1254                foreground_ranges.extend(codegen.last_equal_ranges(cx).iter().cloned());
1255
1256                let pending_range =
1257                    codegen.edit_position(cx).unwrap_or(assist.range.start)..assist.range.end;
1258                if pending_range.end.to_offset(&buffer) > pending_range.start.to_offset(&buffer) {
1259                    gutter_pending_ranges.push(pending_range);
1260                }
1261
1262                if let Some(edit_position) = codegen.edit_position(cx) {
1263                    let edited_range = assist.range.start..edit_position;
1264                    if edited_range.end.to_offset(&buffer) > edited_range.start.to_offset(&buffer) {
1265                        gutter_transformed_ranges.push(edited_range);
1266                    }
1267                }
1268
1269                if assist.decorations.is_some() {
1270                    inserted_row_ranges
1271                        .extend(codegen.diff(cx).inserted_row_ranges.iter().cloned());
1272                }
1273            }
1274        }
1275
1276        let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
1277        merge_ranges(&mut foreground_ranges, &snapshot);
1278        merge_ranges(&mut gutter_pending_ranges, &snapshot);
1279        merge_ranges(&mut gutter_transformed_ranges, &snapshot);
1280        editor.update(cx, |editor, cx| {
1281            enum GutterPendingRange {}
1282            if gutter_pending_ranges.is_empty() {
1283                editor.clear_gutter_highlights::<GutterPendingRange>(cx);
1284            } else {
1285                editor.highlight_gutter::<GutterPendingRange>(
1286                    &gutter_pending_ranges,
1287                    |cx| cx.theme().status().info_background,
1288                    cx,
1289                )
1290            }
1291
1292            enum GutterTransformedRange {}
1293            if gutter_transformed_ranges.is_empty() {
1294                editor.clear_gutter_highlights::<GutterTransformedRange>(cx);
1295            } else {
1296                editor.highlight_gutter::<GutterTransformedRange>(
1297                    &gutter_transformed_ranges,
1298                    |cx| cx.theme().status().info,
1299                    cx,
1300                )
1301            }
1302
1303            if foreground_ranges.is_empty() {
1304                editor.clear_highlights::<InlineAssist>(cx);
1305            } else {
1306                editor.highlight_text::<InlineAssist>(
1307                    foreground_ranges,
1308                    HighlightStyle {
1309                        fade_out: Some(0.6),
1310                        ..Default::default()
1311                    },
1312                    cx,
1313                );
1314            }
1315
1316            editor.clear_row_highlights::<InlineAssist>();
1317            for row_range in inserted_row_ranges {
1318                editor.highlight_rows::<InlineAssist>(
1319                    row_range,
1320                    cx.theme().status().info_background,
1321                    false,
1322                    cx,
1323                );
1324            }
1325        });
1326    }
1327
1328    fn update_editor_blocks(
1329        &mut self,
1330        editor: &Entity<Editor>,
1331        assist_id: InlineAssistId,
1332        window: &mut Window,
1333        cx: &mut App,
1334    ) {
1335        let Some(assist) = self.assists.get_mut(&assist_id) else {
1336            return;
1337        };
1338        let Some(decorations) = assist.decorations.as_mut() else {
1339            return;
1340        };
1341
1342        let codegen = assist.codegen.read(cx);
1343        let old_snapshot = codegen.snapshot(cx);
1344        let old_buffer = codegen.old_buffer(cx);
1345        let deleted_row_ranges = codegen.diff(cx).deleted_row_ranges.clone();
1346
1347        editor.update(cx, |editor, cx| {
1348            let old_blocks = mem::take(&mut decorations.removed_line_block_ids);
1349            editor.remove_blocks(old_blocks, None, cx);
1350
1351            let mut new_blocks = Vec::new();
1352            for (new_row, old_row_range) in deleted_row_ranges {
1353                let (_, buffer_start) = old_snapshot
1354                    .point_to_buffer_offset(Point::new(*old_row_range.start(), 0))
1355                    .unwrap();
1356                let (_, buffer_end) = old_snapshot
1357                    .point_to_buffer_offset(Point::new(
1358                        *old_row_range.end(),
1359                        old_snapshot.line_len(MultiBufferRow(*old_row_range.end())),
1360                    ))
1361                    .unwrap();
1362
1363                let deleted_lines_editor = cx.new(|cx| {
1364                    let multi_buffer =
1365                        cx.new(|_| MultiBuffer::without_headers(language::Capability::ReadOnly));
1366                    multi_buffer.update(cx, |multi_buffer, cx| {
1367                        multi_buffer.push_excerpts(
1368                            old_buffer.clone(),
1369                            Some(ExcerptRange::new(buffer_start..buffer_end)),
1370                            cx,
1371                        );
1372                    });
1373
1374                    enum DeletedLines {}
1375                    let mut editor = Editor::for_multibuffer(multi_buffer, None, window, cx);
1376                    editor.set_soft_wrap_mode(language::language_settings::SoftWrap::None, cx);
1377                    editor.set_show_wrap_guides(false, cx);
1378                    editor.set_show_gutter(false, cx);
1379                    editor.scroll_manager.set_forbid_vertical_scroll(true);
1380                    editor.set_show_scrollbars(false, cx);
1381                    editor.set_read_only(true);
1382                    editor.set_show_edit_predictions(Some(false), window, cx);
1383                    editor.highlight_rows::<DeletedLines>(
1384                        Anchor::min()..Anchor::max(),
1385                        cx.theme().status().deleted_background,
1386                        false,
1387                        cx,
1388                    );
1389                    editor
1390                });
1391
1392                let height =
1393                    deleted_lines_editor.update(cx, |editor, cx| editor.max_point(cx).row().0 + 1);
1394                new_blocks.push(BlockProperties {
1395                    placement: BlockPlacement::Above(new_row),
1396                    height: Some(height),
1397                    style: BlockStyle::Flex,
1398                    render: Arc::new(move |cx| {
1399                        div()
1400                            .block_mouse_down()
1401                            .bg(cx.theme().status().deleted_background)
1402                            .size_full()
1403                            .h(height as f32 * cx.window.line_height())
1404                            .pl(cx.gutter_dimensions.full_width())
1405                            .child(deleted_lines_editor.clone())
1406                            .into_any_element()
1407                    }),
1408                    priority: 0,
1409                });
1410            }
1411
1412            decorations.removed_line_block_ids = editor
1413                .insert_blocks(new_blocks, None, cx)
1414                .into_iter()
1415                .collect();
1416        })
1417    }
1418
1419    fn resolve_inline_assist_target(
1420        workspace: &mut Workspace,
1421        assistant_panel: Option<Entity<AssistantPanel>>,
1422        window: &mut Window,
1423        cx: &mut App,
1424    ) -> Option<InlineAssistTarget> {
1425        if let Some(terminal_panel) = workspace.panel::<TerminalPanel>(cx) {
1426            if terminal_panel
1427                .read(cx)
1428                .focus_handle(cx)
1429                .contains_focused(window, cx)
1430            {
1431                if let Some(terminal_view) = terminal_panel.read(cx).pane().and_then(|pane| {
1432                    pane.read(cx)
1433                        .active_item()
1434                        .and_then(|t| t.downcast::<TerminalView>())
1435                }) {
1436                    return Some(InlineAssistTarget::Terminal(terminal_view));
1437                }
1438            }
1439        }
1440
1441        let context_editor = assistant_panel
1442            .and_then(|panel| panel.read(cx).active_context_editor())
1443            .and_then(|editor| {
1444                let editor = &editor.read(cx).editor().clone();
1445                if editor.read(cx).is_focused(window) {
1446                    Some(editor.clone())
1447                } else {
1448                    None
1449                }
1450            });
1451
1452        if let Some(context_editor) = context_editor {
1453            Some(InlineAssistTarget::Editor(context_editor))
1454        } else if let Some(workspace_editor) = workspace
1455            .active_item(cx)
1456            .and_then(|item| item.act_as::<Editor>(cx))
1457        {
1458            Some(InlineAssistTarget::Editor(workspace_editor))
1459        } else if let Some(terminal_view) = workspace
1460            .active_item(cx)
1461            .and_then(|item| item.act_as::<TerminalView>(cx))
1462        {
1463            Some(InlineAssistTarget::Terminal(terminal_view))
1464        } else {
1465            None
1466        }
1467    }
1468}
1469
1470struct EditorInlineAssists {
1471    assist_ids: Vec<InlineAssistId>,
1472    scroll_lock: Option<InlineAssistScrollLock>,
1473    highlight_updates: async_watch::Sender<()>,
1474    _update_highlights: Task<Result<()>>,
1475    _subscriptions: Vec<gpui::Subscription>,
1476}
1477
1478struct InlineAssistScrollLock {
1479    assist_id: InlineAssistId,
1480    distance_from_top: f32,
1481}
1482
1483impl EditorInlineAssists {
1484    fn new(editor: &Entity<Editor>, window: &mut Window, cx: &mut App) -> Self {
1485        let (highlight_updates_tx, mut highlight_updates_rx) = async_watch::channel(());
1486        Self {
1487            assist_ids: Vec::new(),
1488            scroll_lock: None,
1489            highlight_updates: highlight_updates_tx,
1490            _update_highlights: cx.spawn({
1491                let editor = editor.downgrade();
1492                async move |cx| {
1493                    while let Ok(()) = highlight_updates_rx.changed().await {
1494                        let editor = editor.upgrade().context("editor was dropped")?;
1495                        cx.update_global(|assistant: &mut InlineAssistant, cx| {
1496                            assistant.update_editor_highlights(&editor, cx);
1497                        })?;
1498                    }
1499                    Ok(())
1500                }
1501            }),
1502            _subscriptions: vec![
1503                cx.observe_release_in(editor, window, {
1504                    let editor = editor.downgrade();
1505                    |_, window, cx| {
1506                        InlineAssistant::update_global(cx, |this, cx| {
1507                            this.handle_editor_release(editor, window, cx);
1508                        })
1509                    }
1510                }),
1511                window.observe(editor, cx, move |editor, window, cx| {
1512                    InlineAssistant::update_global(cx, |this, cx| {
1513                        this.handle_editor_change(editor, window, cx)
1514                    })
1515                }),
1516                window.subscribe(editor, cx, move |editor, event, window, cx| {
1517                    InlineAssistant::update_global(cx, |this, cx| {
1518                        this.handle_editor_event(editor, event, window, cx)
1519                    })
1520                }),
1521                editor.update(cx, |editor, cx| {
1522                    let editor_handle = cx.entity().downgrade();
1523                    editor.register_action(move |_: &editor::actions::Newline, window, cx| {
1524                        InlineAssistant::update_global(cx, |this, cx| {
1525                            if let Some(editor) = editor_handle.upgrade() {
1526                                this.handle_editor_newline(editor, window, cx)
1527                            }
1528                        })
1529                    })
1530                }),
1531                editor.update(cx, |editor, cx| {
1532                    let editor_handle = cx.entity().downgrade();
1533                    editor.register_action(move |_: &editor::actions::Cancel, window, cx| {
1534                        InlineAssistant::update_global(cx, |this, cx| {
1535                            if let Some(editor) = editor_handle.upgrade() {
1536                                this.handle_editor_cancel(editor, window, cx)
1537                            }
1538                        })
1539                    })
1540                }),
1541            ],
1542        }
1543    }
1544}
1545
1546struct InlineAssistGroup {
1547    assist_ids: Vec<InlineAssistId>,
1548    linked: bool,
1549    active_assist_id: Option<InlineAssistId>,
1550}
1551
1552impl InlineAssistGroup {
1553    fn new() -> Self {
1554        Self {
1555            assist_ids: Vec::new(),
1556            linked: true,
1557            active_assist_id: None,
1558        }
1559    }
1560}
1561
1562fn build_assist_editor_renderer(editor: &Entity<PromptEditor<BufferCodegen>>) -> RenderBlock {
1563    let editor = editor.clone();
1564
1565    Arc::new(move |cx: &mut BlockContext| {
1566        let gutter_dimensions = editor.read(cx).gutter_dimensions();
1567
1568        *gutter_dimensions.lock() = *cx.gutter_dimensions;
1569        editor.clone().into_any_element()
1570    })
1571}
1572
1573#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
1574struct InlineAssistGroupId(usize);
1575
1576impl InlineAssistGroupId {
1577    fn post_inc(&mut self) -> InlineAssistGroupId {
1578        let id = *self;
1579        self.0 += 1;
1580        id
1581    }
1582}
1583
1584pub struct InlineAssist {
1585    group_id: InlineAssistGroupId,
1586    range: Range<Anchor>,
1587    editor: WeakEntity<Editor>,
1588    decorations: Option<InlineAssistDecorations>,
1589    codegen: Entity<BufferCodegen>,
1590    _subscriptions: Vec<Subscription>,
1591    workspace: WeakEntity<Workspace>,
1592}
1593
1594impl InlineAssist {
1595    fn new(
1596        assist_id: InlineAssistId,
1597        group_id: InlineAssistGroupId,
1598        editor: &Entity<Editor>,
1599        prompt_editor: &Entity<PromptEditor<BufferCodegen>>,
1600        prompt_block_id: CustomBlockId,
1601        end_block_id: CustomBlockId,
1602        range: Range<Anchor>,
1603        codegen: Entity<BufferCodegen>,
1604        workspace: WeakEntity<Workspace>,
1605        window: &mut Window,
1606        cx: &mut App,
1607    ) -> Self {
1608        let prompt_editor_focus_handle = prompt_editor.focus_handle(cx);
1609        InlineAssist {
1610            group_id,
1611            editor: editor.downgrade(),
1612            decorations: Some(InlineAssistDecorations {
1613                prompt_block_id,
1614                prompt_editor: prompt_editor.clone(),
1615                removed_line_block_ids: HashSet::default(),
1616                end_block_id,
1617            }),
1618            range,
1619            codegen: codegen.clone(),
1620            workspace: workspace.clone(),
1621            _subscriptions: vec![
1622                window.on_focus_in(&prompt_editor_focus_handle, cx, move |_, cx| {
1623                    InlineAssistant::update_global(cx, |this, cx| {
1624                        this.handle_prompt_editor_focus_in(assist_id, cx)
1625                    })
1626                }),
1627                window.on_focus_out(&prompt_editor_focus_handle, cx, move |_, _, cx| {
1628                    InlineAssistant::update_global(cx, |this, cx| {
1629                        this.handle_prompt_editor_focus_out(assist_id, cx)
1630                    })
1631                }),
1632                window.subscribe(prompt_editor, cx, |prompt_editor, event, window, cx| {
1633                    InlineAssistant::update_global(cx, |this, cx| {
1634                        this.handle_prompt_editor_event(prompt_editor, event, window, cx)
1635                    })
1636                }),
1637                window.observe(&codegen, cx, {
1638                    let editor = editor.downgrade();
1639                    move |_, window, cx| {
1640                        if let Some(editor) = editor.upgrade() {
1641                            InlineAssistant::update_global(cx, |this, cx| {
1642                                if let Some(editor_assists) =
1643                                    this.assists_by_editor.get(&editor.downgrade())
1644                                {
1645                                    editor_assists.highlight_updates.send(()).ok();
1646                                }
1647
1648                                this.update_editor_blocks(&editor, assist_id, window, cx);
1649                            })
1650                        }
1651                    }
1652                }),
1653                window.subscribe(&codegen, cx, move |codegen, event, window, cx| {
1654                    InlineAssistant::update_global(cx, |this, cx| match event {
1655                        CodegenEvent::Undone => this.finish_assist(assist_id, false, window, cx),
1656                        CodegenEvent::Finished => {
1657                            let assist = if let Some(assist) = this.assists.get(&assist_id) {
1658                                assist
1659                            } else {
1660                                return;
1661                            };
1662
1663                            if let CodegenStatus::Error(error) = codegen.read(cx).status(cx) {
1664                                if assist.decorations.is_none() {
1665                                    if let Some(workspace) = assist.workspace.upgrade() {
1666                                        let error = format!("Inline assistant error: {}", error);
1667                                        workspace.update(cx, |workspace, cx| {
1668                                            struct InlineAssistantError;
1669
1670                                            let id =
1671                                                NotificationId::composite::<InlineAssistantError>(
1672                                                    assist_id.0,
1673                                                );
1674
1675                                            workspace.show_toast(Toast::new(id, error), cx);
1676                                        })
1677                                    }
1678                                }
1679                            }
1680
1681                            if assist.decorations.is_none() {
1682                                this.finish_assist(assist_id, false, window, cx);
1683                            }
1684                        }
1685                    })
1686                }),
1687            ],
1688        }
1689    }
1690
1691    fn user_prompt(&self, cx: &App) -> Option<String> {
1692        let decorations = self.decorations.as_ref()?;
1693        Some(decorations.prompt_editor.read(cx).prompt(cx))
1694    }
1695}
1696
1697struct InlineAssistDecorations {
1698    prompt_block_id: CustomBlockId,
1699    prompt_editor: Entity<PromptEditor<BufferCodegen>>,
1700    removed_line_block_ids: HashSet<CustomBlockId>,
1701    end_block_id: CustomBlockId,
1702}
1703
1704struct AssistantCodeActionProvider {
1705    editor: WeakEntity<Editor>,
1706    workspace: WeakEntity<Workspace>,
1707    thread_store: Option<WeakEntity<ThreadStore>>,
1708}
1709
1710const ASSISTANT_CODE_ACTION_PROVIDER_ID: &str = "assistant2";
1711
1712impl CodeActionProvider for AssistantCodeActionProvider {
1713    fn id(&self) -> Arc<str> {
1714        ASSISTANT_CODE_ACTION_PROVIDER_ID.into()
1715    }
1716
1717    fn code_actions(
1718        &self,
1719        buffer: &Entity<Buffer>,
1720        range: Range<text::Anchor>,
1721        _: &mut Window,
1722        cx: &mut App,
1723    ) -> Task<Result<Vec<CodeAction>>> {
1724        if !AssistantSettings::get_global(cx).enabled {
1725            return Task::ready(Ok(Vec::new()));
1726        }
1727
1728        let snapshot = buffer.read(cx).snapshot();
1729        let mut range = range.to_point(&snapshot);
1730
1731        // Expand the range to line boundaries.
1732        range.start.column = 0;
1733        range.end.column = snapshot.line_len(range.end.row);
1734
1735        let mut has_diagnostics = false;
1736        for diagnostic in snapshot.diagnostics_in_range::<_, Point>(range.clone(), false) {
1737            range.start = cmp::min(range.start, diagnostic.range.start);
1738            range.end = cmp::max(range.end, diagnostic.range.end);
1739            has_diagnostics = true;
1740        }
1741        if has_diagnostics {
1742            if let Some(symbols_containing_start) = snapshot.symbols_containing(range.start, None) {
1743                if let Some(symbol) = symbols_containing_start.last() {
1744                    range.start = cmp::min(range.start, symbol.range.start.to_point(&snapshot));
1745                    range.end = cmp::max(range.end, symbol.range.end.to_point(&snapshot));
1746                }
1747            }
1748
1749            if let Some(symbols_containing_end) = snapshot.symbols_containing(range.end, None) {
1750                if let Some(symbol) = symbols_containing_end.last() {
1751                    range.start = cmp::min(range.start, symbol.range.start.to_point(&snapshot));
1752                    range.end = cmp::max(range.end, symbol.range.end.to_point(&snapshot));
1753                }
1754            }
1755
1756            Task::ready(Ok(vec![CodeAction {
1757                server_id: language::LanguageServerId(0),
1758                range: snapshot.anchor_before(range.start)..snapshot.anchor_after(range.end),
1759                lsp_action: LspAction::Action(Box::new(lsp::CodeAction {
1760                    title: "Fix with Assistant".into(),
1761                    ..Default::default()
1762                })),
1763                resolved: true,
1764            }]))
1765        } else {
1766            Task::ready(Ok(Vec::new()))
1767        }
1768    }
1769
1770    fn apply_code_action(
1771        &self,
1772        buffer: Entity<Buffer>,
1773        action: CodeAction,
1774        excerpt_id: ExcerptId,
1775        _push_to_history: bool,
1776        window: &mut Window,
1777        cx: &mut App,
1778    ) -> Task<Result<ProjectTransaction>> {
1779        let editor = self.editor.clone();
1780        let workspace = self.workspace.clone();
1781        let thread_store = self.thread_store.clone();
1782        window.spawn(cx, async move |cx| {
1783            let workspace = workspace.upgrade().context("workspace was released")?;
1784            let editor = editor.upgrade().context("editor was released")?;
1785            let range = editor
1786                .update(cx, |editor, cx| {
1787                    editor.buffer().update(cx, |multibuffer, cx| {
1788                        let buffer = buffer.read(cx);
1789                        let multibuffer_snapshot = multibuffer.read(cx);
1790
1791                        let old_context_range =
1792                            multibuffer_snapshot.context_range_for_excerpt(excerpt_id)?;
1793                        let mut new_context_range = old_context_range.clone();
1794                        if action
1795                            .range
1796                            .start
1797                            .cmp(&old_context_range.start, buffer)
1798                            .is_lt()
1799                        {
1800                            new_context_range.start = action.range.start;
1801                        }
1802                        if action.range.end.cmp(&old_context_range.end, buffer).is_gt() {
1803                            new_context_range.end = action.range.end;
1804                        }
1805                        drop(multibuffer_snapshot);
1806
1807                        if new_context_range != old_context_range {
1808                            multibuffer.resize_excerpt(excerpt_id, new_context_range, cx);
1809                        }
1810
1811                        let multibuffer_snapshot = multibuffer.read(cx);
1812                        Some(
1813                            multibuffer_snapshot
1814                                .anchor_in_excerpt(excerpt_id, action.range.start)?
1815                                ..multibuffer_snapshot
1816                                    .anchor_in_excerpt(excerpt_id, action.range.end)?,
1817                        )
1818                    })
1819                })?
1820                .context("invalid range")?;
1821
1822            cx.update_global(|assistant: &mut InlineAssistant, window, cx| {
1823                let assist_id = assistant.suggest_assist(
1824                    &editor,
1825                    range,
1826                    "Fix Diagnostics".into(),
1827                    None,
1828                    true,
1829                    workspace,
1830                    thread_store,
1831                    window,
1832                    cx,
1833                );
1834                assistant.start_assist(assist_id, window, cx);
1835            })?;
1836
1837            Ok(ProjectTransaction::default())
1838        })
1839    }
1840}
1841
1842fn merge_ranges(ranges: &mut Vec<Range<Anchor>>, buffer: &MultiBufferSnapshot) {
1843    ranges.sort_unstable_by(|a, b| {
1844        a.start
1845            .cmp(&b.start, buffer)
1846            .then_with(|| b.end.cmp(&a.end, buffer))
1847    });
1848
1849    let mut ix = 0;
1850    while ix + 1 < ranges.len() {
1851        let b = ranges[ix + 1].clone();
1852        let a = &mut ranges[ix];
1853        if a.end.cmp(&b.start, buffer).is_gt() {
1854            if a.end.cmp(&b.end, buffer).is_lt() {
1855                a.end = b.end;
1856            }
1857            ranges.remove(ix + 1);
1858        } else {
1859            ix += 1;
1860        }
1861    }
1862}