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