inline_assistant.rs

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