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