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