inline_assistant.rs

   1use std::cmp;
   2use std::mem;
   3use std::ops::Range;
   4use std::rc::Rc;
   5use std::sync::Arc;
   6
   7use crate::context::load_context;
   8use crate::mention_set::MentionSet;
   9use crate::{
  10    AgentPanel,
  11    buffer_codegen::{BufferCodegen, CodegenAlternative, CodegenEvent},
  12    inline_prompt_editor::{CodegenStatus, InlineAssistId, PromptEditor, PromptEditorEvent},
  13    terminal_inline_assistant::TerminalInlineAssistant,
  14};
  15use agent::HistoryStore;
  16use agent_settings::AgentSettings;
  17use anyhow::{Context as _, Result};
  18use client::telemetry::Telemetry;
  19use collections::{HashMap, HashSet, VecDeque, hash_map};
  20use editor::EditorSnapshot;
  21use editor::MultiBufferOffset;
  22use editor::RowExt;
  23use editor::SelectionEffects;
  24use editor::scroll::ScrollOffset;
  25use editor::{
  26    Anchor, AnchorRangeExt, CodeActionProvider, Editor, EditorEvent, ExcerptId, ExcerptRange,
  27    MultiBuffer, MultiBufferSnapshot, ToOffset as _, ToPoint,
  28    actions::SelectAll,
  29    display_map::{
  30        BlockContext, BlockPlacement, BlockProperties, BlockStyle, CustomBlockId, EditorMargins,
  31        RenderBlock, ToDisplayPoint,
  32    },
  33};
  34use fs::Fs;
  35use futures::{FutureExt, channel::mpsc};
  36use gpui::{
  37    App, Context, Entity, Focusable, Global, HighlightStyle, Subscription, Task, UpdateGlobal,
  38    WeakEntity, Window, point,
  39};
  40use language::{Buffer, Point, Selection, TransactionId};
  41use language_model::{
  42    ConfigurationError, ConfiguredModel, LanguageModelRegistry, report_assistant_event,
  43};
  44use multi_buffer::MultiBufferRow;
  45use parking_lot::Mutex;
  46use project::{CodeAction, DisableAiSettings, LspAction, Project, ProjectTransaction};
  47use prompt_store::{PromptBuilder, PromptStore};
  48use settings::{Settings, SettingsStore};
  49use telemetry_events::{AssistantEventData, AssistantKind, AssistantPhase};
  50use terminal_view::{TerminalView, terminal_panel::TerminalPanel};
  51use text::{OffsetRangeExt, ToPoint as _};
  52use ui::prelude::*;
  53use util::{RangeExt, ResultExt, maybe};
  54use workspace::{ItemHandle, Toast, Workspace, dock::Panel, notifications::NotificationId};
  55use zed_actions::agent::OpenSettings;
  56
  57pub fn init(
  58    fs: Arc<dyn Fs>,
  59    prompt_builder: Arc<PromptBuilder>,
  60    telemetry: Arc<Telemetry>,
  61    cx: &mut App,
  62) {
  63    cx.set_global(InlineAssistant::new(fs, prompt_builder, telemetry));
  64
  65    cx.observe_global::<SettingsStore>(|cx| {
  66        if DisableAiSettings::get_global(cx).disable_ai {
  67            // Hide any active inline assist UI when AI is disabled
  68            InlineAssistant::update_global(cx, |assistant, cx| {
  69                assistant.cancel_all_active_completions(cx);
  70            });
  71        }
  72    })
  73    .detach();
  74
  75    cx.observe_new(|_workspace: &mut Workspace, window, cx| {
  76        let Some(window) = window else {
  77            return;
  78        };
  79        let workspace = cx.entity();
  80        InlineAssistant::update_global(cx, |inline_assistant, cx| {
  81            inline_assistant.register_workspace(&workspace, window, cx)
  82        });
  83    })
  84    .detach();
  85}
  86
  87const PROMPT_HISTORY_MAX_LEN: usize = 20;
  88
  89enum InlineAssistTarget {
  90    Editor(Entity<Editor>),
  91    Terminal(Entity<TerminalView>),
  92}
  93
  94pub struct InlineAssistant {
  95    next_assist_id: InlineAssistId,
  96    next_assist_group_id: InlineAssistGroupId,
  97    assists: HashMap<InlineAssistId, InlineAssist>,
  98    assists_by_editor: HashMap<WeakEntity<Editor>, EditorInlineAssists>,
  99    assist_groups: HashMap<InlineAssistGroupId, InlineAssistGroup>,
 100    confirmed_assists: HashMap<InlineAssistId, Entity<CodegenAlternative>>,
 101    prompt_history: VecDeque<String>,
 102    prompt_builder: Arc<PromptBuilder>,
 103    telemetry: Arc<Telemetry>,
 104    fs: Arc<dyn Fs>,
 105    _inline_assistant_completions: Option<mpsc::UnboundedSender<anyhow::Result<InlineAssistId>>>,
 106}
 107
 108impl Global for InlineAssistant {}
 109
 110impl InlineAssistant {
 111    pub fn new(
 112        fs: Arc<dyn Fs>,
 113        prompt_builder: Arc<PromptBuilder>,
 114        telemetry: Arc<Telemetry>,
 115    ) -> Self {
 116        Self {
 117            next_assist_id: InlineAssistId::default(),
 118            next_assist_group_id: InlineAssistGroupId::default(),
 119            assists: HashMap::default(),
 120            assists_by_editor: HashMap::default(),
 121            assist_groups: HashMap::default(),
 122            confirmed_assists: HashMap::default(),
 123            prompt_history: VecDeque::default(),
 124            prompt_builder,
 125            telemetry,
 126            fs,
 127            _inline_assistant_completions: None,
 128        }
 129    }
 130
 131    #[cfg(any(test, feature = "test-support"))]
 132    pub fn set_completion_receiver(
 133        &mut self,
 134        sender: mpsc::UnboundedSender<anyhow::Result<InlineAssistId>>,
 135    ) {
 136        self._inline_assistant_completions = Some(sender);
 137    }
 138
 139    pub fn register_workspace(
 140        &mut self,
 141        workspace: &Entity<Workspace>,
 142        window: &mut Window,
 143        cx: &mut App,
 144    ) {
 145        window
 146            .subscribe(workspace, cx, |workspace, event, window, cx| {
 147                Self::update_global(cx, |this, cx| {
 148                    this.handle_workspace_event(workspace, event, window, cx)
 149                });
 150            })
 151            .detach();
 152
 153        let workspace = workspace.downgrade();
 154        cx.observe_global::<SettingsStore>(move |cx| {
 155            let Some(workspace) = workspace.upgrade() else {
 156                return;
 157            };
 158            let Some(terminal_panel) = workspace.read(cx).panel::<TerminalPanel>(cx) else {
 159                return;
 160            };
 161            let enabled = AgentSettings::get_global(cx).enabled(cx);
 162            terminal_panel.update(cx, |terminal_panel, cx| {
 163                terminal_panel.set_assistant_enabled(enabled, cx)
 164            });
 165        })
 166        .detach();
 167    }
 168
 169    /// Hides all active inline assists when AI is disabled
 170    pub fn cancel_all_active_completions(&mut self, cx: &mut App) {
 171        // Cancel all active completions in editors
 172        for (editor_handle, _) in self.assists_by_editor.iter() {
 173            if let Some(editor) = editor_handle.upgrade() {
 174                let windows = cx.windows();
 175                if !windows.is_empty() {
 176                    let window = windows[0];
 177                    let _ = window.update(cx, |_, window, cx| {
 178                        editor.update(cx, |editor, cx| {
 179                            if editor.has_active_edit_prediction() {
 180                                editor.cancel(&Default::default(), window, cx);
 181                            }
 182                        });
 183                    });
 184                }
 185            }
 186        }
 187    }
 188
 189    fn handle_workspace_event(
 190        &mut self,
 191        workspace: Entity<Workspace>,
 192        event: &workspace::Event,
 193        window: &mut Window,
 194        cx: &mut App,
 195    ) {
 196        match event {
 197            workspace::Event::UserSavedItem { item, .. } => {
 198                // When the user manually saves an editor, automatically accepts all finished transformations.
 199                if let Some(editor) = item.upgrade().and_then(|item| item.act_as::<Editor>(cx))
 200                    && let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade())
 201                {
 202                    for assist_id in editor_assists.assist_ids.clone() {
 203                        let assist = &self.assists[&assist_id];
 204                        if let CodegenStatus::Done = assist.codegen.read(cx).status(cx) {
 205                            self.finish_assist(assist_id, false, window, cx)
 206                        }
 207                    }
 208                }
 209            }
 210            workspace::Event::ItemAdded { item } => {
 211                self.register_workspace_item(&workspace, item.as_ref(), window, cx);
 212            }
 213            _ => (),
 214        }
 215    }
 216
 217    fn register_workspace_item(
 218        &mut self,
 219        workspace: &Entity<Workspace>,
 220        item: &dyn ItemHandle,
 221        window: &mut Window,
 222        cx: &mut App,
 223    ) {
 224        let is_ai_enabled = !DisableAiSettings::get_global(cx).disable_ai;
 225
 226        if let Some(editor) = item.act_as::<Editor>(cx) {
 227            editor.update(cx, |editor, cx| {
 228                if is_ai_enabled {
 229                    editor.add_code_action_provider(
 230                        Rc::new(AssistantCodeActionProvider {
 231                            editor: cx.entity().downgrade(),
 232                            workspace: workspace.downgrade(),
 233                        }),
 234                        window,
 235                        cx,
 236                    );
 237
 238                    if DisableAiSettings::get_global(cx).disable_ai {
 239                        // Cancel any active edit predictions
 240                        if editor.has_active_edit_prediction() {
 241                            editor.cancel(&Default::default(), window, cx);
 242                        }
 243                    }
 244                } else {
 245                    editor.remove_code_action_provider(
 246                        ASSISTANT_CODE_ACTION_PROVIDER_ID.into(),
 247                        window,
 248                        cx,
 249                    );
 250                }
 251            });
 252        }
 253    }
 254
 255    pub fn inline_assist(
 256        workspace: &mut Workspace,
 257        action: &zed_actions::assistant::InlineAssist,
 258        window: &mut Window,
 259        cx: &mut Context<Workspace>,
 260    ) {
 261        if !AgentSettings::get_global(cx).enabled(cx) {
 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 = agent_panel.thread_store().clone();
 286
 287        let handle_assist =
 288            |window: &mut Window, cx: &mut Context<Workspace>| match inline_assist_target {
 289                InlineAssistTarget::Editor(active_editor) => {
 290                    InlineAssistant::update_global(cx, |assistant, cx| {
 291                        assistant.assist(
 292                            &active_editor,
 293                            cx.entity().downgrade(),
 294                            workspace.project().downgrade(),
 295                            thread_store,
 296                            prompt_store,
 297                            action.prompt.clone(),
 298                            window,
 299                            cx,
 300                        );
 301                    })
 302                }
 303                InlineAssistTarget::Terminal(active_terminal) => {
 304                    TerminalInlineAssistant::update_global(cx, |assistant, cx| {
 305                        assistant.assist(
 306                            &active_terminal,
 307                            cx.entity().downgrade(),
 308                            workspace.project().downgrade(),
 309                            thread_store,
 310                            prompt_store,
 311                            action.prompt.clone(),
 312                            window,
 313                            cx,
 314                        );
 315                    });
 316                }
 317            };
 318
 319        if let Some(error) = configuration_error() {
 320            if let ConfigurationError::ProviderNotAuthenticated(provider) = error {
 321                cx.spawn(async move |_, cx| {
 322                    cx.update(|cx| provider.authenticate(cx))?.await?;
 323                    anyhow::Ok(())
 324                })
 325                .detach_and_log_err(cx);
 326
 327                if configuration_error().is_none() {
 328                    handle_assist(window, cx);
 329                }
 330            } else {
 331                cx.spawn_in(window, async move |_, cx| {
 332                    let answer = cx
 333                        .prompt(
 334                            gpui::PromptLevel::Warning,
 335                            &error.to_string(),
 336                            None,
 337                            &["Configure", "Cancel"],
 338                        )
 339                        .await
 340                        .ok();
 341                    if let Some(answer) = answer
 342                        && answer == 0
 343                    {
 344                        cx.update(|window, cx| window.dispatch_action(Box::new(OpenSettings), cx))
 345                            .ok();
 346                    }
 347                    anyhow::Ok(())
 348                })
 349                .detach_and_log_err(cx);
 350            }
 351        } else {
 352            handle_assist(window, cx);
 353        }
 354    }
 355
 356    fn codegen_ranges(
 357        &mut self,
 358        editor: &Entity<Editor>,
 359        snapshot: &EditorSnapshot,
 360        window: &mut Window,
 361        cx: &mut App,
 362    ) -> Option<(Vec<Range<Anchor>>, Selection<Point>)> {
 363        let (initial_selections, newest_selection) = editor.update(cx, |editor, _| {
 364            (
 365                editor.selections.all::<Point>(&snapshot.display_snapshot),
 366                editor
 367                    .selections
 368                    .newest::<Point>(&snapshot.display_snapshot),
 369            )
 370        });
 371
 372        // Check if there is already an inline assistant that contains the
 373        // newest selection, if there is, focus it
 374        if let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) {
 375            for assist_id in &editor_assists.assist_ids {
 376                let assist = &self.assists[assist_id];
 377                let range = assist.range.to_point(&snapshot.buffer_snapshot());
 378                if range.start.row <= newest_selection.start.row
 379                    && newest_selection.end.row <= range.end.row
 380                {
 381                    self.focus_assist(*assist_id, window, cx);
 382                    return None;
 383                }
 384            }
 385        }
 386
 387        let mut selections = Vec::<Selection<Point>>::new();
 388        let mut newest_selection = None;
 389        for mut selection in initial_selections {
 390            if selection.end > selection.start {
 391                selection.start.column = 0;
 392                // If the selection ends at the start of the line, we don't want to include it.
 393                if selection.end.column == 0 {
 394                    selection.end.row -= 1;
 395                }
 396                selection.end.column = snapshot
 397                    .buffer_snapshot()
 398                    .line_len(MultiBufferRow(selection.end.row));
 399            } else if let Some(fold) =
 400                snapshot.crease_for_buffer_row(MultiBufferRow(selection.end.row))
 401            {
 402                selection.start = fold.range().start;
 403                selection.end = fold.range().end;
 404                if MultiBufferRow(selection.end.row) < snapshot.buffer_snapshot().max_row() {
 405                    let chars = snapshot
 406                        .buffer_snapshot()
 407                        .chars_at(Point::new(selection.end.row + 1, 0));
 408
 409                    for c in chars {
 410                        if c == '\n' {
 411                            break;
 412                        }
 413                        if c.is_whitespace() {
 414                            continue;
 415                        }
 416                        if snapshot
 417                            .language_at(selection.end)
 418                            .is_some_and(|language| language.config().brackets.is_closing_brace(c))
 419                        {
 420                            selection.end.row += 1;
 421                            selection.end.column = snapshot
 422                                .buffer_snapshot()
 423                                .line_len(MultiBufferRow(selection.end.row));
 424                        }
 425                    }
 426                }
 427            }
 428
 429            if let Some(prev_selection) = selections.last_mut()
 430                && selection.start <= prev_selection.end
 431            {
 432                prev_selection.end = selection.end;
 433                continue;
 434            }
 435
 436            let latest_selection = newest_selection.get_or_insert_with(|| selection.clone());
 437            if selection.id > latest_selection.id {
 438                *latest_selection = selection.clone();
 439            }
 440            selections.push(selection);
 441        }
 442        let snapshot = &snapshot.buffer_snapshot();
 443        let newest_selection = newest_selection.unwrap();
 444
 445        let mut codegen_ranges = Vec::new();
 446        for (buffer, buffer_range, excerpt_id) in
 447            snapshot.ranges_to_buffer_ranges(selections.iter().map(|selection| {
 448                snapshot.anchor_before(selection.start)..snapshot.anchor_after(selection.end)
 449            }))
 450        {
 451            let anchor_range = Anchor::range_in_buffer(
 452                excerpt_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        project: WeakEntity<Project>,
 481        thread_store: Entity<HistoryStore>,
 482        prompt_store: Option<Entity<PromptStore>>,
 483        initial_prompt: Option<String>,
 484        window: &mut Window,
 485        codegen_ranges: &[Range<Anchor>],
 486        newest_selection: Option<Selection<Point>>,
 487        initial_transaction_id: Option<TransactionId>,
 488        cx: &mut App,
 489    ) -> Option<InlineAssistId> {
 490        let snapshot = editor.update(cx, |editor, cx| editor.snapshot(window, cx));
 491
 492        let assist_group_id = self.next_assist_group_id.post_inc();
 493        let prompt_buffer = cx.new(|cx| {
 494            MultiBuffer::singleton(
 495                cx.new(|cx| Buffer::local(initial_prompt.unwrap_or_default(), cx)),
 496                cx,
 497            )
 498        });
 499
 500        let mut assists = Vec::new();
 501        let mut assist_to_focus = None;
 502
 503        for range in codegen_ranges {
 504            let assist_id = self.next_assist_id.post_inc();
 505            let codegen = cx.new(|cx| {
 506                BufferCodegen::new(
 507                    editor.read(cx).buffer().clone(),
 508                    range.clone(),
 509                    initial_transaction_id,
 510                    self.telemetry.clone(),
 511                    self.prompt_builder.clone(),
 512                    cx,
 513                )
 514            });
 515
 516            let editor_margins = Arc::new(Mutex::new(EditorMargins::default()));
 517            let prompt_editor = cx.new(|cx| {
 518                PromptEditor::new_buffer(
 519                    assist_id,
 520                    editor_margins,
 521                    self.prompt_history.clone(),
 522                    prompt_buffer.clone(),
 523                    codegen.clone(),
 524                    self.fs.clone(),
 525                    thread_store.clone(),
 526                    prompt_store.clone(),
 527                    project.clone(),
 528                    workspace.clone(),
 529                    window,
 530                    cx,
 531                )
 532            });
 533
 534            if let Some(newest_selection) = newest_selection.as_ref()
 535                && assist_to_focus.is_none()
 536            {
 537                let focus_assist = if newest_selection.reversed {
 538                    range.start.to_point(&snapshot) == newest_selection.start
 539                } else {
 540                    range.end.to_point(&snapshot) == newest_selection.end
 541                };
 542                if focus_assist {
 543                    assist_to_focus = Some(assist_id);
 544                }
 545            }
 546
 547            let [prompt_block_id, end_block_id] =
 548                self.insert_assist_blocks(editor, &range, &prompt_editor, cx);
 549
 550            assists.push((
 551                assist_id,
 552                range.clone(),
 553                prompt_editor,
 554                prompt_block_id,
 555                end_block_id,
 556            ));
 557        }
 558
 559        let editor_assists = self
 560            .assists_by_editor
 561            .entry(editor.downgrade())
 562            .or_insert_with(|| EditorInlineAssists::new(editor, window, cx));
 563
 564        let assist_to_focus = if let Some(focus_id) = assist_to_focus {
 565            Some(focus_id)
 566        } else if assists.len() >= 1 {
 567            Some(assists[0].0)
 568        } else {
 569            None
 570        };
 571
 572        let mut assist_group = InlineAssistGroup::new();
 573        for (assist_id, range, prompt_editor, prompt_block_id, end_block_id) in assists {
 574            let codegen = prompt_editor.read(cx).codegen().clone();
 575
 576            self.assists.insert(
 577                assist_id,
 578                InlineAssist::new(
 579                    assist_id,
 580                    assist_group_id,
 581                    editor,
 582                    &prompt_editor,
 583                    prompt_block_id,
 584                    end_block_id,
 585                    range,
 586                    codegen,
 587                    workspace.clone(),
 588                    window,
 589                    cx,
 590                ),
 591            );
 592            assist_group.assist_ids.push(assist_id);
 593            editor_assists.assist_ids.push(assist_id);
 594        }
 595
 596        self.assist_groups.insert(assist_group_id, assist_group);
 597
 598        assist_to_focus
 599    }
 600
 601    pub fn assist(
 602        &mut self,
 603        editor: &Entity<Editor>,
 604        workspace: WeakEntity<Workspace>,
 605        project: WeakEntity<Project>,
 606        thread_store: Entity<HistoryStore>,
 607        prompt_store: Option<Entity<PromptStore>>,
 608        initial_prompt: Option<String>,
 609        window: &mut Window,
 610        cx: &mut App,
 611    ) -> Option<InlineAssistId> {
 612        let snapshot = editor.update(cx, |editor, cx| editor.snapshot(window, cx));
 613
 614        let Some((codegen_ranges, newest_selection)) =
 615            self.codegen_ranges(editor, &snapshot, window, cx)
 616        else {
 617            return None;
 618        };
 619
 620        let assist_to_focus = self.batch_assist(
 621            editor,
 622            workspace,
 623            project,
 624            thread_store,
 625            prompt_store,
 626            initial_prompt,
 627            window,
 628            &codegen_ranges,
 629            Some(newest_selection),
 630            None,
 631            cx,
 632        );
 633
 634        if let Some(assist_id) = assist_to_focus {
 635            self.focus_assist(assist_id, window, cx);
 636        }
 637
 638        assist_to_focus
 639    }
 640
 641    pub fn suggest_assist(
 642        &mut self,
 643        editor: &Entity<Editor>,
 644        mut range: Range<Anchor>,
 645        initial_prompt: String,
 646        initial_transaction_id: Option<TransactionId>,
 647        focus: bool,
 648        workspace: Entity<Workspace>,
 649        thread_store: Entity<HistoryStore>,
 650        prompt_store: Option<Entity<PromptStore>>,
 651        window: &mut Window,
 652        cx: &mut App,
 653    ) -> InlineAssistId {
 654        let buffer = editor.read(cx).buffer().clone();
 655        {
 656            let snapshot = buffer.read(cx).read(cx);
 657            range.start = range.start.bias_left(&snapshot);
 658            range.end = range.end.bias_right(&snapshot);
 659        }
 660
 661        let project = workspace.read(cx).project().downgrade();
 662
 663        let assist_id = self
 664            .batch_assist(
 665                editor,
 666                workspace.downgrade(),
 667                project,
 668                thread_store,
 669                prompt_store,
 670                Some(initial_prompt),
 671                window,
 672                &[range],
 673                None,
 674                initial_transaction_id,
 675                cx,
 676            )
 677            .expect("batch_assist returns an id if there's only one range");
 678
 679        if focus {
 680            self.focus_assist(assist_id, window, cx);
 681        }
 682
 683        assist_id
 684    }
 685
 686    fn insert_assist_blocks(
 687        &self,
 688        editor: &Entity<Editor>,
 689        range: &Range<Anchor>,
 690        prompt_editor: &Entity<PromptEditor<BufferCodegen>>,
 691        cx: &mut App,
 692    ) -> [CustomBlockId; 2] {
 693        let prompt_editor_height = prompt_editor.update(cx, |prompt_editor, cx| {
 694            prompt_editor
 695                .editor
 696                .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1 + 2)
 697        });
 698        let assist_blocks = vec![
 699            BlockProperties {
 700                style: BlockStyle::Sticky,
 701                placement: BlockPlacement::Above(range.start),
 702                height: Some(prompt_editor_height),
 703                render: build_assist_editor_renderer(prompt_editor),
 704                priority: 0,
 705            },
 706            BlockProperties {
 707                style: BlockStyle::Sticky,
 708                placement: BlockPlacement::Below(range.end),
 709                height: None,
 710                render: Arc::new(|cx| {
 711                    v_flex()
 712                        .h_full()
 713                        .w_full()
 714                        .border_t_1()
 715                        .border_color(cx.theme().status().info_border)
 716                        .into_any_element()
 717                }),
 718                priority: 0,
 719            },
 720        ];
 721
 722        editor.update(cx, |editor, cx| {
 723            let block_ids = editor.insert_blocks(assist_blocks, None, cx);
 724            [block_ids[0], block_ids[1]]
 725        })
 726    }
 727
 728    fn handle_prompt_editor_focus_in(&mut self, assist_id: InlineAssistId, cx: &mut App) {
 729        let assist = &self.assists[&assist_id];
 730        let Some(decorations) = assist.decorations.as_ref() else {
 731            return;
 732        };
 733        let assist_group = self.assist_groups.get_mut(&assist.group_id).unwrap();
 734        let editor_assists = self.assists_by_editor.get_mut(&assist.editor).unwrap();
 735
 736        assist_group.active_assist_id = Some(assist_id);
 737        if assist_group.linked {
 738            for assist_id in &assist_group.assist_ids {
 739                if let Some(decorations) = self.assists[assist_id].decorations.as_ref() {
 740                    decorations.prompt_editor.update(cx, |prompt_editor, cx| {
 741                        prompt_editor.set_show_cursor_when_unfocused(true, cx)
 742                    });
 743                }
 744            }
 745        }
 746
 747        assist
 748            .editor
 749            .update(cx, |editor, cx| {
 750                let scroll_top = editor.scroll_position(cx).y;
 751                let scroll_bottom = scroll_top + editor.visible_line_count().unwrap_or(0.);
 752                editor_assists.scroll_lock = editor
 753                    .row_for_block(decorations.prompt_block_id, cx)
 754                    .map(|row| row.as_f64())
 755                    .filter(|prompt_row| (scroll_top..scroll_bottom).contains(&prompt_row))
 756                    .map(|prompt_row| InlineAssistScrollLock {
 757                        assist_id,
 758                        distance_from_top: prompt_row - scroll_top,
 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
 817                        .selections
 818                        .newest::<MultiBufferOffset>(&editor.display_snapshot(cx)),
 819                    editor.buffer().read(cx).snapshot(cx),
 820                )
 821            });
 822            for assist_id in &editor_assists.assist_ids {
 823                let assist = &self.assists[assist_id];
 824                let assist_range = assist.range.to_offset(&buffer);
 825                if assist_range.contains(&selection.start) && assist_range.contains(&selection.end)
 826                {
 827                    if matches!(assist.codegen.read(cx).status(cx), CodegenStatus::Pending) {
 828                        self.dismiss_assist(*assist_id, window, cx);
 829                    } else {
 830                        self.finish_assist(*assist_id, false, window, cx);
 831                    }
 832
 833                    return;
 834                }
 835            }
 836        }
 837
 838        cx.propagate();
 839    }
 840
 841    fn handle_editor_cancel(&mut self, editor: Entity<Editor>, window: &mut Window, cx: &mut App) {
 842        let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
 843            return;
 844        };
 845
 846        if editor.read(cx).selections.count() == 1 {
 847            let (selection, buffer) = editor.update(cx, |editor, cx| {
 848                (
 849                    editor
 850                        .selections
 851                        .newest::<MultiBufferOffset>(&editor.display_snapshot(cx)),
 852                    editor.buffer().read(cx).snapshot(cx),
 853                )
 854            });
 855            let mut closest_assist_fallback = None;
 856            for assist_id in &editor_assists.assist_ids {
 857                let assist = &self.assists[assist_id];
 858                let assist_range = assist.range.to_offset(&buffer);
 859                if assist.decorations.is_some() {
 860                    if assist_range.contains(&selection.start)
 861                        && assist_range.contains(&selection.end)
 862                    {
 863                        self.focus_assist(*assist_id, window, cx);
 864                        return;
 865                    } else {
 866                        let distance_from_selection = assist_range
 867                            .start
 868                            .0
 869                            .abs_diff(selection.start.0)
 870                            .min(assist_range.start.0.abs_diff(selection.end.0))
 871                            + assist_range
 872                                .end
 873                                .0
 874                                .abs_diff(selection.start.0)
 875                                .min(assist_range.end.0.abs_diff(selection.end.0));
 876                        match closest_assist_fallback {
 877                            Some((_, old_distance)) => {
 878                                if distance_from_selection < old_distance {
 879                                    closest_assist_fallback =
 880                                        Some((assist_id, distance_from_selection));
 881                                }
 882                            }
 883                            None => {
 884                                closest_assist_fallback = Some((assist_id, distance_from_selection))
 885                            }
 886                        }
 887                    }
 888                }
 889            }
 890
 891            if let Some((&assist_id, _)) = closest_assist_fallback {
 892                self.focus_assist(assist_id, window, cx);
 893            }
 894        }
 895
 896        cx.propagate();
 897    }
 898
 899    fn handle_editor_release(
 900        &mut self,
 901        editor: WeakEntity<Editor>,
 902        window: &mut Window,
 903        cx: &mut App,
 904    ) {
 905        if let Some(editor_assists) = self.assists_by_editor.get_mut(&editor) {
 906            for assist_id in editor_assists.assist_ids.clone() {
 907                self.finish_assist(assist_id, true, window, cx);
 908            }
 909        }
 910    }
 911
 912    fn handle_editor_change(&mut self, editor: Entity<Editor>, window: &mut Window, cx: &mut App) {
 913        let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
 914            return;
 915        };
 916        let Some(scroll_lock) = editor_assists.scroll_lock.as_ref() else {
 917            return;
 918        };
 919        let assist = &self.assists[&scroll_lock.assist_id];
 920        let Some(decorations) = assist.decorations.as_ref() else {
 921            return;
 922        };
 923
 924        editor.update(cx, |editor, cx| {
 925            let scroll_position = editor.scroll_position(cx);
 926            let target_scroll_top = editor
 927                .row_for_block(decorations.prompt_block_id, cx)?
 928                .as_f64()
 929                - scroll_lock.distance_from_top;
 930            if target_scroll_top != scroll_position.y {
 931                editor.set_scroll_position(point(scroll_position.x, target_scroll_top), window, cx);
 932            }
 933            Some(())
 934        });
 935    }
 936
 937    fn handle_editor_event(
 938        &mut self,
 939        editor: Entity<Editor>,
 940        event: &EditorEvent,
 941        window: &mut Window,
 942        cx: &mut App,
 943    ) {
 944        let Some(editor_assists) = self.assists_by_editor.get_mut(&editor.downgrade()) else {
 945            return;
 946        };
 947
 948        match event {
 949            EditorEvent::Edited { transaction_id } => {
 950                let buffer = editor.read(cx).buffer().read(cx);
 951                let edited_ranges =
 952                    buffer.edited_ranges_for_transaction::<MultiBufferOffset>(*transaction_id, cx);
 953                let snapshot = buffer.snapshot(cx);
 954
 955                for assist_id in editor_assists.assist_ids.clone() {
 956                    let assist = &self.assists[&assist_id];
 957                    if matches!(
 958                        assist.codegen.read(cx).status(cx),
 959                        CodegenStatus::Error(_) | CodegenStatus::Done
 960                    ) {
 961                        let assist_range = assist.range.to_offset(&snapshot);
 962                        if edited_ranges
 963                            .iter()
 964                            .any(|range| range.overlaps(&assist_range))
 965                        {
 966                            self.finish_assist(assist_id, false, window, cx);
 967                        }
 968                    }
 969                }
 970            }
 971            EditorEvent::ScrollPositionChanged { .. } => {
 972                if let Some(scroll_lock) = editor_assists.scroll_lock.as_ref() {
 973                    let assist = &self.assists[&scroll_lock.assist_id];
 974                    if let Some(decorations) = assist.decorations.as_ref() {
 975                        let distance_from_top = editor.update(cx, |editor, cx| {
 976                            let scroll_top = editor.scroll_position(cx).y;
 977                            let prompt_row = editor
 978                                .row_for_block(decorations.prompt_block_id, cx)?
 979                                .0 as ScrollOffset;
 980                            Some(prompt_row - scroll_top)
 981                        });
 982
 983                        if distance_from_top.is_none_or(|distance_from_top| {
 984                            distance_from_top != scroll_lock.distance_from_top
 985                        }) {
 986                            editor_assists.scroll_lock = None;
 987                        }
 988                    }
 989                }
 990            }
 991            EditorEvent::SelectionsChanged { .. } => {
 992                for assist_id in editor_assists.assist_ids.clone() {
 993                    let assist = &self.assists[&assist_id];
 994                    if let Some(decorations) = assist.decorations.as_ref()
 995                        && decorations
 996                            .prompt_editor
 997                            .focus_handle(cx)
 998                            .is_focused(window)
 999                    {
1000                        return;
1001                    }
1002                }
1003
1004                editor_assists.scroll_lock = None;
1005            }
1006            _ => {}
1007        }
1008    }
1009
1010    pub fn finish_assist(
1011        &mut self,
1012        assist_id: InlineAssistId,
1013        undo: bool,
1014        window: &mut Window,
1015        cx: &mut App,
1016    ) {
1017        if let Some(assist) = self.assists.get(&assist_id) {
1018            let assist_group_id = assist.group_id;
1019            if self.assist_groups[&assist_group_id].linked {
1020                for assist_id in self.unlink_assist_group(assist_group_id, window, cx) {
1021                    self.finish_assist(assist_id, undo, window, cx);
1022                }
1023                return;
1024            }
1025        }
1026
1027        self.dismiss_assist(assist_id, window, cx);
1028
1029        if let Some(assist) = self.assists.remove(&assist_id) {
1030            if let hash_map::Entry::Occupied(mut entry) = self.assist_groups.entry(assist.group_id)
1031            {
1032                entry.get_mut().assist_ids.retain(|id| *id != assist_id);
1033                if entry.get().assist_ids.is_empty() {
1034                    entry.remove();
1035                }
1036            }
1037
1038            if let hash_map::Entry::Occupied(mut entry) =
1039                self.assists_by_editor.entry(assist.editor.clone())
1040            {
1041                entry.get_mut().assist_ids.retain(|id| *id != assist_id);
1042                if entry.get().assist_ids.is_empty() {
1043                    entry.remove();
1044                    if let Some(editor) = assist.editor.upgrade() {
1045                        self.update_editor_highlights(&editor, cx);
1046                    }
1047                } else {
1048                    entry.get_mut().highlight_updates.send(()).ok();
1049                }
1050            }
1051
1052            let active_alternative = assist.codegen.read(cx).active_alternative().clone();
1053            let message_id = active_alternative.read(cx).message_id.clone();
1054
1055            if let Some(model) = LanguageModelRegistry::read_global(cx).inline_assistant_model() {
1056                let language_name = assist.editor.upgrade().and_then(|editor| {
1057                    let multibuffer = editor.read(cx).buffer().read(cx);
1058                    let snapshot = multibuffer.snapshot(cx);
1059                    let ranges = snapshot.range_to_buffer_ranges(assist.range.clone());
1060                    ranges
1061                        .first()
1062                        .and_then(|(buffer, _, _)| buffer.language())
1063                        .map(|language| language.name())
1064                });
1065                report_assistant_event(
1066                    AssistantEventData {
1067                        conversation_id: None,
1068                        kind: AssistantKind::Inline,
1069                        message_id,
1070                        phase: if undo {
1071                            AssistantPhase::Rejected
1072                        } else {
1073                            AssistantPhase::Accepted
1074                        },
1075                        model: model.model.telemetry_id(),
1076                        model_provider: model.model.provider_id().to_string(),
1077                        response_latency: None,
1078                        error_message: None,
1079                        language_name: language_name.map(|name| name.to_proto()),
1080                    },
1081                    Some(self.telemetry.clone()),
1082                    cx.http_client(),
1083                    model.model.api_key(cx),
1084                    cx.background_executor(),
1085                );
1086            }
1087
1088            if undo {
1089                assist.codegen.update(cx, |codegen, cx| codegen.undo(cx));
1090            } else {
1091                self.confirmed_assists.insert(assist_id, active_alternative);
1092            }
1093        }
1094    }
1095
1096    fn dismiss_assist(
1097        &mut self,
1098        assist_id: InlineAssistId,
1099        window: &mut Window,
1100        cx: &mut App,
1101    ) -> bool {
1102        let Some(assist) = self.assists.get_mut(&assist_id) else {
1103            return false;
1104        };
1105        let Some(editor) = assist.editor.upgrade() else {
1106            return false;
1107        };
1108        let Some(decorations) = assist.decorations.take() else {
1109            return false;
1110        };
1111
1112        editor.update(cx, |editor, cx| {
1113            let mut to_remove = decorations.removed_line_block_ids;
1114            to_remove.insert(decorations.prompt_block_id);
1115            to_remove.insert(decorations.end_block_id);
1116            editor.remove_blocks(to_remove, None, cx);
1117        });
1118
1119        if decorations
1120            .prompt_editor
1121            .focus_handle(cx)
1122            .contains_focused(window, cx)
1123        {
1124            self.focus_next_assist(assist_id, window, cx);
1125        }
1126
1127        if let Some(editor_assists) = self.assists_by_editor.get_mut(&editor.downgrade()) {
1128            if editor_assists
1129                .scroll_lock
1130                .as_ref()
1131                .is_some_and(|lock| lock.assist_id == assist_id)
1132            {
1133                editor_assists.scroll_lock = None;
1134            }
1135            editor_assists.highlight_updates.send(()).ok();
1136        }
1137
1138        true
1139    }
1140
1141    fn focus_next_assist(&mut self, assist_id: InlineAssistId, window: &mut Window, cx: &mut App) {
1142        let Some(assist) = self.assists.get(&assist_id) else {
1143            return;
1144        };
1145
1146        let assist_group = &self.assist_groups[&assist.group_id];
1147        let assist_ix = assist_group
1148            .assist_ids
1149            .iter()
1150            .position(|id| *id == assist_id)
1151            .unwrap();
1152        let assist_ids = assist_group
1153            .assist_ids
1154            .iter()
1155            .skip(assist_ix + 1)
1156            .chain(assist_group.assist_ids.iter().take(assist_ix));
1157
1158        for assist_id in assist_ids {
1159            let assist = &self.assists[assist_id];
1160            if assist.decorations.is_some() {
1161                self.focus_assist(*assist_id, window, cx);
1162                return;
1163            }
1164        }
1165
1166        assist
1167            .editor
1168            .update(cx, |editor, cx| window.focus(&editor.focus_handle(cx)))
1169            .ok();
1170    }
1171
1172    fn focus_assist(&mut self, assist_id: InlineAssistId, window: &mut Window, cx: &mut App) {
1173        let Some(assist) = self.assists.get(&assist_id) else {
1174            return;
1175        };
1176
1177        if let Some(decorations) = assist.decorations.as_ref() {
1178            decorations.prompt_editor.update(cx, |prompt_editor, cx| {
1179                prompt_editor.editor.update(cx, |editor, cx| {
1180                    window.focus(&editor.focus_handle(cx));
1181                    editor.select_all(&SelectAll, window, cx);
1182                })
1183            });
1184        }
1185
1186        self.scroll_to_assist(assist_id, window, cx);
1187    }
1188
1189    pub fn scroll_to_assist(
1190        &mut self,
1191        assist_id: InlineAssistId,
1192        window: &mut Window,
1193        cx: &mut App,
1194    ) {
1195        let Some(assist) = self.assists.get(&assist_id) else {
1196            return;
1197        };
1198        let Some(editor) = assist.editor.upgrade() else {
1199            return;
1200        };
1201
1202        let position = assist.range.start;
1203        editor.update(cx, |editor, cx| {
1204            editor.change_selections(SelectionEffects::no_scroll(), window, cx, |selections| {
1205                selections.select_anchor_ranges([position..position])
1206            });
1207
1208            let mut scroll_target_range = None;
1209            if let Some(decorations) = assist.decorations.as_ref() {
1210                scroll_target_range = maybe!({
1211                    let top = editor.row_for_block(decorations.prompt_block_id, cx)?.0 as f64;
1212                    let bottom = editor.row_for_block(decorations.end_block_id, cx)?.0 as f64;
1213                    Some((top, bottom))
1214                });
1215                if scroll_target_range.is_none() {
1216                    log::error!("bug: failed to find blocks for scrolling to inline assist");
1217                }
1218            }
1219            let scroll_target_range = scroll_target_range.unwrap_or_else(|| {
1220                let snapshot = editor.snapshot(window, cx);
1221                let start_row = assist
1222                    .range
1223                    .start
1224                    .to_display_point(&snapshot.display_snapshot)
1225                    .row();
1226                let top = start_row.0 as ScrollOffset;
1227                let bottom = top + 1.0;
1228                (top, bottom)
1229            });
1230            let mut scroll_target_top = scroll_target_range.0;
1231            let mut scroll_target_bottom = scroll_target_range.1;
1232
1233            scroll_target_top -= editor.vertical_scroll_margin() as ScrollOffset;
1234            scroll_target_bottom += editor.vertical_scroll_margin() as ScrollOffset;
1235
1236            let height_in_lines = editor.visible_line_count().unwrap_or(0.);
1237            let scroll_top = editor.scroll_position(cx).y;
1238            let scroll_bottom = scroll_top + height_in_lines;
1239
1240            if scroll_target_top < scroll_top {
1241                editor.set_scroll_position(point(0., scroll_target_top), window, cx);
1242            } else if scroll_target_bottom > scroll_bottom {
1243                if (scroll_target_bottom - scroll_target_top) <= height_in_lines {
1244                    editor.set_scroll_position(
1245                        point(0., scroll_target_bottom - height_in_lines),
1246                        window,
1247                        cx,
1248                    );
1249                } else {
1250                    editor.set_scroll_position(point(0., scroll_target_top), window, cx);
1251                }
1252            }
1253        });
1254    }
1255
1256    fn unlink_assist_group(
1257        &mut self,
1258        assist_group_id: InlineAssistGroupId,
1259        window: &mut Window,
1260        cx: &mut App,
1261    ) -> Vec<InlineAssistId> {
1262        let assist_group = self.assist_groups.get_mut(&assist_group_id).unwrap();
1263        assist_group.linked = false;
1264
1265        for assist_id in &assist_group.assist_ids {
1266            let assist = self.assists.get_mut(assist_id).unwrap();
1267            if let Some(editor_decorations) = assist.decorations.as_ref() {
1268                editor_decorations
1269                    .prompt_editor
1270                    .update(cx, |prompt_editor, cx| prompt_editor.unlink(window, cx));
1271            }
1272        }
1273        assist_group.assist_ids.clone()
1274    }
1275
1276    pub fn start_assist(&mut self, assist_id: InlineAssistId, window: &mut Window, cx: &mut App) {
1277        let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
1278            assist
1279        } else {
1280            return;
1281        };
1282
1283        let assist_group_id = assist.group_id;
1284        if self.assist_groups[&assist_group_id].linked {
1285            for assist_id in self.unlink_assist_group(assist_group_id, window, cx) {
1286                self.start_assist(assist_id, window, cx);
1287            }
1288            return;
1289        }
1290
1291        let Some((user_prompt, mention_set)) = assist.user_prompt(cx).zip(assist.mention_set(cx))
1292        else {
1293            return;
1294        };
1295
1296        self.prompt_history.retain(|prompt| *prompt != user_prompt);
1297        self.prompt_history.push_back(user_prompt.clone());
1298        if self.prompt_history.len() > PROMPT_HISTORY_MAX_LEN {
1299            self.prompt_history.pop_front();
1300        }
1301
1302        let Some(ConfiguredModel { model, .. }) =
1303            LanguageModelRegistry::read_global(cx).inline_assistant_model()
1304        else {
1305            return;
1306        };
1307
1308        let context_task = load_context(&mention_set, cx).shared();
1309        assist
1310            .codegen
1311            .update(cx, |codegen, cx| {
1312                codegen.start(model, user_prompt, context_task, cx)
1313            })
1314            .log_err();
1315    }
1316
1317    pub fn stop_assist(&mut self, assist_id: InlineAssistId, cx: &mut App) {
1318        let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
1319            assist
1320        } else {
1321            return;
1322        };
1323
1324        assist.codegen.update(cx, |codegen, cx| codegen.stop(cx));
1325    }
1326
1327    fn update_editor_highlights(&self, editor: &Entity<Editor>, cx: &mut App) {
1328        let mut gutter_pending_ranges = Vec::new();
1329        let mut gutter_transformed_ranges = Vec::new();
1330        let mut foreground_ranges = Vec::new();
1331        let mut inserted_row_ranges = Vec::new();
1332        let empty_assist_ids = Vec::new();
1333        let assist_ids = self
1334            .assists_by_editor
1335            .get(&editor.downgrade())
1336            .map_or(&empty_assist_ids, |editor_assists| {
1337                &editor_assists.assist_ids
1338            });
1339
1340        for assist_id in assist_ids {
1341            if let Some(assist) = self.assists.get(assist_id) {
1342                let codegen = assist.codegen.read(cx);
1343                let buffer = codegen.buffer(cx).read(cx).read(cx);
1344                foreground_ranges.extend(codegen.last_equal_ranges(cx).iter().cloned());
1345
1346                let pending_range =
1347                    codegen.edit_position(cx).unwrap_or(assist.range.start)..assist.range.end;
1348                if pending_range.end.to_offset(&buffer) > pending_range.start.to_offset(&buffer) {
1349                    gutter_pending_ranges.push(pending_range);
1350                }
1351
1352                if let Some(edit_position) = codegen.edit_position(cx) {
1353                    let edited_range = assist.range.start..edit_position;
1354                    if edited_range.end.to_offset(&buffer) > edited_range.start.to_offset(&buffer) {
1355                        gutter_transformed_ranges.push(edited_range);
1356                    }
1357                }
1358
1359                if assist.decorations.is_some() {
1360                    inserted_row_ranges
1361                        .extend(codegen.diff(cx).inserted_row_ranges.iter().cloned());
1362                }
1363            }
1364        }
1365
1366        let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
1367        merge_ranges(&mut foreground_ranges, &snapshot);
1368        merge_ranges(&mut gutter_pending_ranges, &snapshot);
1369        merge_ranges(&mut gutter_transformed_ranges, &snapshot);
1370        editor.update(cx, |editor, cx| {
1371            enum GutterPendingRange {}
1372            if gutter_pending_ranges.is_empty() {
1373                editor.clear_gutter_highlights::<GutterPendingRange>(cx);
1374            } else {
1375                editor.highlight_gutter::<GutterPendingRange>(
1376                    gutter_pending_ranges,
1377                    |cx| cx.theme().status().info_background,
1378                    cx,
1379                )
1380            }
1381
1382            enum GutterTransformedRange {}
1383            if gutter_transformed_ranges.is_empty() {
1384                editor.clear_gutter_highlights::<GutterTransformedRange>(cx);
1385            } else {
1386                editor.highlight_gutter::<GutterTransformedRange>(
1387                    gutter_transformed_ranges,
1388                    |cx| cx.theme().status().info,
1389                    cx,
1390                )
1391            }
1392
1393            if foreground_ranges.is_empty() {
1394                editor.clear_highlights::<InlineAssist>(cx);
1395            } else {
1396                editor.highlight_text::<InlineAssist>(
1397                    foreground_ranges,
1398                    HighlightStyle {
1399                        fade_out: Some(0.6),
1400                        ..Default::default()
1401                    },
1402                    cx,
1403                );
1404            }
1405
1406            editor.clear_row_highlights::<InlineAssist>();
1407            for row_range in inserted_row_ranges {
1408                editor.highlight_rows::<InlineAssist>(
1409                    row_range,
1410                    cx.theme().status().info_background,
1411                    Default::default(),
1412                    cx,
1413                );
1414            }
1415        });
1416    }
1417
1418    fn update_editor_blocks(
1419        &mut self,
1420        editor: &Entity<Editor>,
1421        assist_id: InlineAssistId,
1422        window: &mut Window,
1423        cx: &mut App,
1424    ) {
1425        let Some(assist) = self.assists.get_mut(&assist_id) else {
1426            return;
1427        };
1428        let Some(decorations) = assist.decorations.as_mut() else {
1429            return;
1430        };
1431
1432        let codegen = assist.codegen.read(cx);
1433        let old_snapshot = codegen.snapshot(cx);
1434        let old_buffer = codegen.old_buffer(cx);
1435        let deleted_row_ranges = codegen.diff(cx).deleted_row_ranges.clone();
1436
1437        editor.update(cx, |editor, cx| {
1438            let old_blocks = mem::take(&mut decorations.removed_line_block_ids);
1439            editor.remove_blocks(old_blocks, None, cx);
1440
1441            let mut new_blocks = Vec::new();
1442            for (new_row, old_row_range) in deleted_row_ranges {
1443                let (_, buffer_start) = old_snapshot
1444                    .point_to_buffer_offset(Point::new(*old_row_range.start(), 0))
1445                    .unwrap();
1446                let (_, buffer_end) = old_snapshot
1447                    .point_to_buffer_offset(Point::new(
1448                        *old_row_range.end(),
1449                        old_snapshot.line_len(MultiBufferRow(*old_row_range.end())),
1450                    ))
1451                    .unwrap();
1452
1453                let deleted_lines_editor = cx.new(|cx| {
1454                    let multi_buffer =
1455                        cx.new(|_| MultiBuffer::without_headers(language::Capability::ReadOnly));
1456                    multi_buffer.update(cx, |multi_buffer, cx| {
1457                        multi_buffer.push_excerpts(
1458                            old_buffer.clone(),
1459                            // todo(lw): buffer_start and buffer_end might come from different snapshots!
1460                            Some(ExcerptRange::new(buffer_start..buffer_end)),
1461                            cx,
1462                        );
1463                    });
1464
1465                    enum DeletedLines {}
1466                    let mut editor = Editor::for_multibuffer(multi_buffer, None, window, cx);
1467                    editor.disable_scrollbars_and_minimap(window, cx);
1468                    editor.set_soft_wrap_mode(language::language_settings::SoftWrap::None, cx);
1469                    editor.set_show_wrap_guides(false, cx);
1470                    editor.set_show_gutter(false, cx);
1471                    editor.set_offset_content(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                                #[cfg(any(test, feature = "test-support"))]
1756                                if let Some(sender) = &mut this._inline_assistant_completions {
1757                                    sender
1758                                        .unbounded_send(Err(anyhow::anyhow!(
1759                                            "Inline assistant error: {}",
1760                                            error
1761                                        )))
1762                                        .ok();
1763                                }
1764
1765                                let error = format!("Inline assistant error: {}", error);
1766                                workspace.update(cx, |workspace, cx| {
1767                                    struct InlineAssistantError;
1768
1769                                    let id = NotificationId::composite::<InlineAssistantError>(
1770                                        assist_id.0,
1771                                    );
1772
1773                                    workspace.show_toast(Toast::new(id, error), cx);
1774                                })
1775                            } else {
1776                                #[cfg(any(test, feature = "test-support"))]
1777                                if let Some(sender) = &mut this._inline_assistant_completions {
1778                                    sender.unbounded_send(Ok(assist_id)).ok();
1779                                }
1780                            }
1781
1782                            if assist.decorations.is_none() {
1783                                this.finish_assist(assist_id, false, window, cx);
1784                            }
1785                        }
1786                    })
1787                }),
1788            ],
1789        }
1790    }
1791
1792    fn user_prompt(&self, cx: &App) -> Option<String> {
1793        let decorations = self.decorations.as_ref()?;
1794        Some(decorations.prompt_editor.read(cx).prompt(cx))
1795    }
1796
1797    fn mention_set(&self, cx: &App) -> Option<Entity<MentionSet>> {
1798        let decorations = self.decorations.as_ref()?;
1799        Some(decorations.prompt_editor.read(cx).mention_set().clone())
1800    }
1801}
1802
1803struct InlineAssistDecorations {
1804    prompt_block_id: CustomBlockId,
1805    prompt_editor: Entity<PromptEditor<BufferCodegen>>,
1806    removed_line_block_ids: HashSet<CustomBlockId>,
1807    end_block_id: CustomBlockId,
1808}
1809
1810struct AssistantCodeActionProvider {
1811    editor: WeakEntity<Editor>,
1812    workspace: WeakEntity<Workspace>,
1813}
1814
1815const ASSISTANT_CODE_ACTION_PROVIDER_ID: &str = "assistant";
1816
1817impl CodeActionProvider for AssistantCodeActionProvider {
1818    fn id(&self) -> Arc<str> {
1819        ASSISTANT_CODE_ACTION_PROVIDER_ID.into()
1820    }
1821
1822    fn code_actions(
1823        &self,
1824        buffer: &Entity<Buffer>,
1825        range: Range<text::Anchor>,
1826        _: &mut Window,
1827        cx: &mut App,
1828    ) -> Task<Result<Vec<CodeAction>>> {
1829        if !AgentSettings::get_global(cx).enabled(cx) {
1830            return Task::ready(Ok(Vec::new()));
1831        }
1832
1833        let snapshot = buffer.read(cx).snapshot();
1834        let mut range = range.to_point(&snapshot);
1835
1836        // Expand the range to line boundaries.
1837        range.start.column = 0;
1838        range.end.column = snapshot.line_len(range.end.row);
1839
1840        let mut has_diagnostics = false;
1841        for diagnostic in snapshot.diagnostics_in_range::<_, Point>(range.clone(), false) {
1842            range.start = cmp::min(range.start, diagnostic.range.start);
1843            range.end = cmp::max(range.end, diagnostic.range.end);
1844            has_diagnostics = true;
1845        }
1846        if has_diagnostics {
1847            let symbols_containing_start = snapshot.symbols_containing(range.start, None);
1848            if let Some(symbol) = symbols_containing_start.last() {
1849                range.start = cmp::min(range.start, symbol.range.start.to_point(&snapshot));
1850                range.end = cmp::max(range.end, symbol.range.end.to_point(&snapshot));
1851            }
1852            let symbols_containing_end = snapshot.symbols_containing(range.end, None);
1853            if let Some(symbol) = symbols_containing_end.last() {
1854                range.start = cmp::min(range.start, symbol.range.start.to_point(&snapshot));
1855                range.end = cmp::max(range.end, symbol.range.end.to_point(&snapshot));
1856            }
1857
1858            Task::ready(Ok(vec![CodeAction {
1859                server_id: language::LanguageServerId(0),
1860                range: snapshot.anchor_before(range.start)..snapshot.anchor_after(range.end),
1861                lsp_action: LspAction::Action(Box::new(lsp::CodeAction {
1862                    title: "Fix with Assistant".into(),
1863                    ..Default::default()
1864                })),
1865                resolved: true,
1866            }]))
1867        } else {
1868            Task::ready(Ok(Vec::new()))
1869        }
1870    }
1871
1872    fn apply_code_action(
1873        &self,
1874        buffer: Entity<Buffer>,
1875        action: CodeAction,
1876        excerpt_id: ExcerptId,
1877        _push_to_history: bool,
1878        window: &mut Window,
1879        cx: &mut App,
1880    ) -> Task<Result<ProjectTransaction>> {
1881        let editor = self.editor.clone();
1882        let workspace = self.workspace.clone();
1883        let prompt_store = PromptStore::global(cx);
1884        window.spawn(cx, async move |cx| {
1885            let workspace = workspace.upgrade().context("workspace was released")?;
1886            let thread_store = cx.update(|_window, cx| {
1887                anyhow::Ok(
1888                    workspace
1889                        .read(cx)
1890                        .panel::<AgentPanel>(cx)
1891                        .context("missing agent panel")?
1892                        .read(cx)
1893                        .thread_store()
1894                        .clone(),
1895                )
1896            })??;
1897            let editor = editor.upgrade().context("editor was released")?;
1898            let range = editor
1899                .update(cx, |editor, cx| {
1900                    editor.buffer().update(cx, |multibuffer, cx| {
1901                        let buffer = buffer.read(cx);
1902                        let multibuffer_snapshot = multibuffer.read(cx);
1903
1904                        let old_context_range =
1905                            multibuffer_snapshot.context_range_for_excerpt(excerpt_id)?;
1906                        let mut new_context_range = old_context_range.clone();
1907                        if action
1908                            .range
1909                            .start
1910                            .cmp(&old_context_range.start, buffer)
1911                            .is_lt()
1912                        {
1913                            new_context_range.start = action.range.start;
1914                        }
1915                        if action.range.end.cmp(&old_context_range.end, buffer).is_gt() {
1916                            new_context_range.end = action.range.end;
1917                        }
1918                        drop(multibuffer_snapshot);
1919
1920                        if new_context_range != old_context_range {
1921                            multibuffer.resize_excerpt(excerpt_id, new_context_range, cx);
1922                        }
1923
1924                        let multibuffer_snapshot = multibuffer.read(cx);
1925                        multibuffer_snapshot.anchor_range_in_excerpt(excerpt_id, action.range)
1926                    })
1927                })?
1928                .context("invalid range")?;
1929
1930            let prompt_store = prompt_store.await.ok();
1931            cx.update_global(|assistant: &mut InlineAssistant, window, cx| {
1932                let assist_id = assistant.suggest_assist(
1933                    &editor,
1934                    range,
1935                    "Fix Diagnostics".into(),
1936                    None,
1937                    true,
1938                    workspace,
1939                    thread_store,
1940                    prompt_store,
1941                    window,
1942                    cx,
1943                );
1944                assistant.start_assist(assist_id, window, cx);
1945            })?;
1946
1947            Ok(ProjectTransaction::default())
1948        })
1949    }
1950}
1951
1952fn merge_ranges(ranges: &mut Vec<Range<Anchor>>, buffer: &MultiBufferSnapshot) {
1953    ranges.sort_unstable_by(|a, b| {
1954        a.start
1955            .cmp(&b.start, buffer)
1956            .then_with(|| b.end.cmp(&a.end, buffer))
1957    });
1958
1959    let mut ix = 0;
1960    while ix + 1 < ranges.len() {
1961        let b = ranges[ix + 1].clone();
1962        let a = &mut ranges[ix];
1963        if a.end.cmp(&b.start, buffer).is_gt() {
1964            if a.end.cmp(&b.end, buffer).is_lt() {
1965                a.end = b.end;
1966            }
1967            ranges.remove(ix + 1);
1968        } else {
1969            ix += 1;
1970        }
1971    }
1972}
1973
1974#[cfg(any(test, feature = "test-support"))]
1975pub mod test {
1976    use std::sync::Arc;
1977
1978    use agent::HistoryStore;
1979    use assistant_text_thread::TextThreadStore;
1980    use client::{Client, UserStore};
1981    use editor::{Editor, MultiBuffer, MultiBufferOffset};
1982    use fs::FakeFs;
1983    use futures::channel::mpsc;
1984    use gpui::{AppContext, TestAppContext, UpdateGlobal as _};
1985    use language::Buffer;
1986    use language_model::LanguageModelRegistry;
1987    use project::Project;
1988    use prompt_store::PromptBuilder;
1989    use smol::stream::StreamExt as _;
1990    use util::test::marked_text_ranges;
1991    use workspace::Workspace;
1992
1993    use crate::InlineAssistant;
1994
1995    pub fn run_inline_assistant_test<SetupF, TestF>(
1996        base_buffer: String,
1997        prompt: String,
1998        setup: SetupF,
1999        test: TestF,
2000        cx: &mut TestAppContext,
2001    ) -> String
2002    where
2003        SetupF: FnOnce(&mut gpui::VisualTestContext),
2004        TestF: FnOnce(&mut gpui::VisualTestContext),
2005    {
2006        let fs = FakeFs::new(cx.executor());
2007        let app_state = cx.update(|cx| workspace::AppState::test(cx));
2008        let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
2009        let http = Arc::new(reqwest_client::ReqwestClient::user_agent("agent tests").unwrap());
2010        let client = cx.update(|cx| {
2011            cx.set_http_client(http);
2012            Client::production(cx)
2013        });
2014        let mut inline_assistant =
2015            InlineAssistant::new(fs.clone(), prompt_builder, client.telemetry().clone());
2016
2017        let (tx, mut completion_rx) = mpsc::unbounded();
2018        inline_assistant.set_completion_receiver(tx);
2019
2020        // Initialize settings and client
2021        cx.update(|cx| {
2022            gpui_tokio::init(cx);
2023            settings::init(cx);
2024            client::init(&client, cx);
2025            workspace::init(app_state.clone(), cx);
2026            let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
2027            language_model::init(client.clone(), cx);
2028            language_models::init(user_store, client.clone(), cx);
2029
2030            cx.set_global(inline_assistant);
2031        });
2032
2033        let project = cx
2034            .executor()
2035            .block_test(async { Project::test(fs.clone(), [], cx).await });
2036
2037        // Create workspace with window
2038        let (workspace, cx) = cx.add_window_view(|window, cx| {
2039            window.activate_window();
2040            Workspace::new(None, project.clone(), app_state.clone(), window, cx)
2041        });
2042
2043        setup(cx);
2044
2045        let (_editor, buffer) = cx.update(|window, cx| {
2046            let buffer = cx.new(|cx| Buffer::local("", cx));
2047            let multibuffer = cx.new(|cx| MultiBuffer::singleton(buffer.clone(), cx));
2048            let editor = cx.new(|cx| Editor::for_multibuffer(multibuffer, None, window, cx));
2049            editor.update(cx, |editor, cx| {
2050                let (unmarked_text, selection_ranges) = marked_text_ranges(&base_buffer, true);
2051                editor.set_text(unmarked_text, window, cx);
2052                editor.change_selections(Default::default(), window, cx, |s| {
2053                    s.select_ranges(
2054                        selection_ranges.into_iter().map(|range| {
2055                            MultiBufferOffset(range.start)..MultiBufferOffset(range.end)
2056                        }),
2057                    )
2058                })
2059            });
2060
2061            let text_thread_store = cx.new(|cx| TextThreadStore::fake(project.clone(), cx));
2062            let history_store = cx.new(|cx| HistoryStore::new(text_thread_store, cx));
2063
2064            // Add editor to workspace
2065            workspace.update(cx, |workspace, cx| {
2066                workspace.add_item_to_active_pane(Box::new(editor.clone()), None, true, window, cx);
2067            });
2068
2069            // Call assist method
2070            InlineAssistant::update_global(cx, |inline_assistant, cx| {
2071                let assist_id = inline_assistant
2072                    .assist(
2073                        &editor,
2074                        workspace.downgrade(),
2075                        project.downgrade(),
2076                        history_store, // thread_store
2077                        None,          // prompt_store
2078                        Some(prompt),
2079                        window,
2080                        cx,
2081                    )
2082                    .unwrap();
2083
2084                inline_assistant.start_assist(assist_id, window, cx);
2085            });
2086
2087            (editor, buffer)
2088        });
2089
2090        cx.run_until_parked();
2091
2092        test(cx);
2093
2094        cx.executor()
2095            .block_test(async { completion_rx.next().await });
2096
2097        buffer.read_with(cx, |buffer, _| buffer.text())
2098    }
2099
2100    #[allow(unused)]
2101    pub fn test_inline_assistant(
2102        base_buffer: &'static str,
2103        llm_output: &'static str,
2104        cx: &mut TestAppContext,
2105    ) -> String {
2106        run_inline_assistant_test(
2107            base_buffer.to_string(),
2108            "Prompt doesn't matter because we're using a fake model".to_string(),
2109            |cx| {
2110                cx.update(|_, cx| LanguageModelRegistry::test(cx));
2111            },
2112            |cx| {
2113                let fake_model = cx.update(|_, cx| {
2114                    LanguageModelRegistry::global(cx)
2115                        .update(cx, |registry, _| registry.fake_model())
2116                });
2117                let fake = fake_model.as_fake();
2118
2119                // let fake = fake_model;
2120                fake.send_last_completion_stream_text_chunk(llm_output.to_string());
2121                fake.end_last_completion_stream();
2122
2123                // Run again to process the model's response
2124                cx.run_until_parked();
2125            },
2126            cx,
2127        )
2128    }
2129}