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                && let Some(fold) =
 392                    snapshot.crease_for_buffer_row(MultiBufferRow(selection.end.row))
 393            {
 394                selection.start = fold.range().start;
 395                selection.end = fold.range().end;
 396                if MultiBufferRow(selection.end.row) < snapshot.buffer_snapshot().max_row() {
 397                    let chars = snapshot
 398                        .buffer_snapshot()
 399                        .chars_at(Point::new(selection.end.row + 1, 0));
 400
 401                    for c in chars {
 402                        if c == '\n' {
 403                            break;
 404                        }
 405                        if c.is_whitespace() {
 406                            continue;
 407                        }
 408                        if snapshot
 409                            .language_at(selection.end)
 410                            .is_some_and(|language| language.config().brackets.is_closing_brace(c))
 411                        {
 412                            selection.end.row += 1;
 413                            selection.end.column = snapshot
 414                                .buffer_snapshot()
 415                                .line_len(MultiBufferRow(selection.end.row));
 416                        }
 417                    }
 418                }
 419            } else {
 420                selection.start.column = 0;
 421                // If the selection ends at the start of the line, we don't want to include it.
 422                if selection.end.column == 0 && selection.start.row != selection.end.row {
 423                    selection.end.row -= 1;
 424                }
 425                selection.end.column = snapshot
 426                    .buffer_snapshot()
 427                    .line_len(MultiBufferRow(selection.end.row));
 428            }
 429
 430            if let Some(prev_selection) = selections.last_mut()
 431                && selection.start <= prev_selection.end
 432            {
 433                prev_selection.end = selection.end;
 434                continue;
 435            }
 436
 437            let latest_selection = newest_selection.get_or_insert_with(|| selection.clone());
 438            if selection.id > latest_selection.id {
 439                *latest_selection = selection.clone();
 440            }
 441            selections.push(selection);
 442        }
 443        let snapshot = &snapshot.buffer_snapshot();
 444        let newest_selection = newest_selection.unwrap();
 445
 446        let mut codegen_ranges = Vec::new();
 447        for (buffer, buffer_range, excerpt_id) in
 448            snapshot.ranges_to_buffer_ranges(selections.iter().map(|selection| {
 449                snapshot.anchor_before(selection.start)..snapshot.anchor_after(selection.end)
 450            }))
 451        {
 452            let anchor_range = Anchor::range_in_buffer(
 453                excerpt_id,
 454                buffer.anchor_before(buffer_range.start)..buffer.anchor_after(buffer_range.end),
 455            );
 456
 457            codegen_ranges.push(anchor_range);
 458
 459            if let Some(model) = LanguageModelRegistry::read_global(cx).inline_assistant_model() {
 460                self.telemetry.report_assistant_event(AssistantEventData {
 461                    conversation_id: None,
 462                    kind: AssistantKind::Inline,
 463                    phase: AssistantPhase::Invoked,
 464                    message_id: None,
 465                    model: model.model.telemetry_id(),
 466                    model_provider: model.provider.id().to_string(),
 467                    response_latency: None,
 468                    error_message: None,
 469                    language_name: buffer.language().map(|language| language.name().to_proto()),
 470                });
 471            }
 472        }
 473
 474        Some((codegen_ranges, newest_selection))
 475    }
 476
 477    fn batch_assist(
 478        &mut self,
 479        editor: &Entity<Editor>,
 480        workspace: WeakEntity<Workspace>,
 481        project: WeakEntity<Project>,
 482        thread_store: Entity<HistoryStore>,
 483        prompt_store: Option<Entity<PromptStore>>,
 484        initial_prompt: Option<String>,
 485        window: &mut Window,
 486        codegen_ranges: &[Range<Anchor>],
 487        newest_selection: Option<Selection<Point>>,
 488        initial_transaction_id: Option<TransactionId>,
 489        cx: &mut App,
 490    ) -> Option<InlineAssistId> {
 491        let snapshot = editor.update(cx, |editor, cx| editor.snapshot(window, cx));
 492
 493        let assist_group_id = self.next_assist_group_id.post_inc();
 494        let prompt_buffer = cx.new(|cx| {
 495            MultiBuffer::singleton(
 496                cx.new(|cx| Buffer::local(initial_prompt.unwrap_or_default(), cx)),
 497                cx,
 498            )
 499        });
 500
 501        let mut assists = Vec::new();
 502        let mut assist_to_focus = None;
 503
 504        for range in codegen_ranges {
 505            let assist_id = self.next_assist_id.post_inc();
 506            let codegen = cx.new(|cx| {
 507                BufferCodegen::new(
 508                    editor.read(cx).buffer().clone(),
 509                    range.clone(),
 510                    initial_transaction_id,
 511                    self.telemetry.clone(),
 512                    self.prompt_builder.clone(),
 513                    cx,
 514                )
 515            });
 516
 517            let editor_margins = Arc::new(Mutex::new(EditorMargins::default()));
 518            let prompt_editor = cx.new(|cx| {
 519                PromptEditor::new_buffer(
 520                    assist_id,
 521                    editor_margins,
 522                    self.prompt_history.clone(),
 523                    prompt_buffer.clone(),
 524                    codegen.clone(),
 525                    self.fs.clone(),
 526                    thread_store.clone(),
 527                    prompt_store.clone(),
 528                    project.clone(),
 529                    workspace.clone(),
 530                    window,
 531                    cx,
 532                )
 533            });
 534
 535            if let Some(newest_selection) = newest_selection.as_ref()
 536                && assist_to_focus.is_none()
 537            {
 538                let focus_assist = if newest_selection.reversed {
 539                    range.start.to_point(&snapshot) == newest_selection.start
 540                } else {
 541                    range.end.to_point(&snapshot) == newest_selection.end
 542                };
 543                if focus_assist {
 544                    assist_to_focus = Some(assist_id);
 545                }
 546            }
 547
 548            let [prompt_block_id, tool_description_block_id, end_block_id] =
 549                self.insert_assist_blocks(&editor, &range, &prompt_editor, cx);
 550
 551            assists.push((
 552                assist_id,
 553                range.clone(),
 554                prompt_editor,
 555                prompt_block_id,
 556                tool_description_block_id,
 557                end_block_id,
 558            ));
 559        }
 560
 561        let editor_assists = self
 562            .assists_by_editor
 563            .entry(editor.downgrade())
 564            .or_insert_with(|| EditorInlineAssists::new(editor, window, cx));
 565
 566        let assist_to_focus = if let Some(focus_id) = assist_to_focus {
 567            Some(focus_id)
 568        } else if assists.len() >= 1 {
 569            Some(assists[0].0)
 570        } else {
 571            None
 572        };
 573
 574        let mut assist_group = InlineAssistGroup::new();
 575        for (
 576            assist_id,
 577            range,
 578            prompt_editor,
 579            prompt_block_id,
 580            tool_description_block_id,
 581            end_block_id,
 582        ) in assists
 583        {
 584            let codegen = prompt_editor.read(cx).codegen().clone();
 585
 586            self.assists.insert(
 587                assist_id,
 588                InlineAssist::new(
 589                    assist_id,
 590                    assist_group_id,
 591                    editor,
 592                    &prompt_editor,
 593                    prompt_block_id,
 594                    tool_description_block_id,
 595                    end_block_id,
 596                    range,
 597                    codegen,
 598                    workspace.clone(),
 599                    window,
 600                    cx,
 601                ),
 602            );
 603            assist_group.assist_ids.push(assist_id);
 604            editor_assists.assist_ids.push(assist_id);
 605        }
 606
 607        self.assist_groups.insert(assist_group_id, assist_group);
 608
 609        assist_to_focus
 610    }
 611
 612    pub fn assist(
 613        &mut self,
 614        editor: &Entity<Editor>,
 615        workspace: WeakEntity<Workspace>,
 616        project: WeakEntity<Project>,
 617        thread_store: Entity<HistoryStore>,
 618        prompt_store: Option<Entity<PromptStore>>,
 619        initial_prompt: Option<String>,
 620        window: &mut Window,
 621        cx: &mut App,
 622    ) -> Option<InlineAssistId> {
 623        let snapshot = editor.update(cx, |editor, cx| editor.snapshot(window, cx));
 624
 625        let Some((codegen_ranges, newest_selection)) =
 626            self.codegen_ranges(editor, &snapshot, window, cx)
 627        else {
 628            return None;
 629        };
 630
 631        let assist_to_focus = self.batch_assist(
 632            editor,
 633            workspace,
 634            project,
 635            thread_store,
 636            prompt_store,
 637            initial_prompt,
 638            window,
 639            &codegen_ranges,
 640            Some(newest_selection),
 641            None,
 642            cx,
 643        );
 644
 645        if let Some(assist_id) = assist_to_focus {
 646            self.focus_assist(assist_id, window, cx);
 647        }
 648
 649        assist_to_focus
 650    }
 651
 652    pub fn suggest_assist(
 653        &mut self,
 654        editor: &Entity<Editor>,
 655        mut range: Range<Anchor>,
 656        initial_prompt: String,
 657        initial_transaction_id: Option<TransactionId>,
 658        focus: bool,
 659        workspace: Entity<Workspace>,
 660        thread_store: Entity<HistoryStore>,
 661        prompt_store: Option<Entity<PromptStore>>,
 662        window: &mut Window,
 663        cx: &mut App,
 664    ) -> InlineAssistId {
 665        let buffer = editor.read(cx).buffer().clone();
 666        {
 667            let snapshot = buffer.read(cx).read(cx);
 668            range.start = range.start.bias_left(&snapshot);
 669            range.end = range.end.bias_right(&snapshot);
 670        }
 671
 672        let project = workspace.read(cx).project().downgrade();
 673
 674        let assist_id = self
 675            .batch_assist(
 676                editor,
 677                workspace.downgrade(),
 678                project,
 679                thread_store,
 680                prompt_store,
 681                Some(initial_prompt),
 682                window,
 683                &[range],
 684                None,
 685                initial_transaction_id,
 686                cx,
 687            )
 688            .expect("batch_assist returns an id if there's only one range");
 689
 690        if focus {
 691            self.focus_assist(assist_id, window, cx);
 692        }
 693
 694        assist_id
 695    }
 696
 697    fn insert_assist_blocks(
 698        &self,
 699        editor: &Entity<Editor>,
 700        range: &Range<Anchor>,
 701        prompt_editor: &Entity<PromptEditor<BufferCodegen>>,
 702        cx: &mut App,
 703    ) -> [CustomBlockId; 3] {
 704        let prompt_editor_height = prompt_editor.update(cx, |prompt_editor, cx| {
 705            prompt_editor
 706                .editor
 707                .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1 + 2)
 708        });
 709        let assist_blocks = vec![
 710            BlockProperties {
 711                style: BlockStyle::Sticky,
 712                placement: BlockPlacement::Above(range.start),
 713                height: Some(prompt_editor_height),
 714                render: build_assist_editor_renderer(prompt_editor),
 715                priority: 0,
 716            },
 717            // Placeholder for tool description - will be updated dynamically
 718            BlockProperties {
 719                style: BlockStyle::Flex,
 720                placement: BlockPlacement::Below(range.end),
 721                height: Some(0),
 722                render: Arc::new(|_cx| div().into_any_element()),
 723                priority: 0,
 724            },
 725            BlockProperties {
 726                style: BlockStyle::Sticky,
 727                placement: BlockPlacement::Below(range.end),
 728                height: None,
 729                render: Arc::new(|cx| {
 730                    v_flex()
 731                        .h_full()
 732                        .w_full()
 733                        .border_t_1()
 734                        .border_color(cx.theme().status().info_border)
 735                        .into_any_element()
 736                }),
 737                priority: 0,
 738            },
 739        ];
 740
 741        editor.update(cx, |editor, cx| {
 742            let block_ids = editor.insert_blocks(assist_blocks, None, cx);
 743            [block_ids[0], block_ids[1], block_ids[2]]
 744        })
 745    }
 746
 747    fn handle_prompt_editor_focus_in(&mut self, assist_id: InlineAssistId, cx: &mut App) {
 748        let assist = &self.assists[&assist_id];
 749        let Some(decorations) = assist.decorations.as_ref() else {
 750            return;
 751        };
 752        let assist_group = self.assist_groups.get_mut(&assist.group_id).unwrap();
 753        let editor_assists = self.assists_by_editor.get_mut(&assist.editor).unwrap();
 754
 755        assist_group.active_assist_id = Some(assist_id);
 756        if assist_group.linked {
 757            for assist_id in &assist_group.assist_ids {
 758                if let Some(decorations) = self.assists[assist_id].decorations.as_ref() {
 759                    decorations.prompt_editor.update(cx, |prompt_editor, cx| {
 760                        prompt_editor.set_show_cursor_when_unfocused(true, cx)
 761                    });
 762                }
 763            }
 764        }
 765
 766        assist
 767            .editor
 768            .update(cx, |editor, cx| {
 769                let scroll_top = editor.scroll_position(cx).y;
 770                let scroll_bottom = scroll_top + editor.visible_line_count().unwrap_or(0.);
 771                editor_assists.scroll_lock = editor
 772                    .row_for_block(decorations.prompt_block_id, cx)
 773                    .map(|row| row.as_f64())
 774                    .filter(|prompt_row| (scroll_top..scroll_bottom).contains(&prompt_row))
 775                    .map(|prompt_row| InlineAssistScrollLock {
 776                        assist_id,
 777                        distance_from_top: prompt_row - scroll_top,
 778                    });
 779            })
 780            .ok();
 781    }
 782
 783    fn handle_prompt_editor_focus_out(&mut self, assist_id: InlineAssistId, cx: &mut App) {
 784        let assist = &self.assists[&assist_id];
 785        let assist_group = self.assist_groups.get_mut(&assist.group_id).unwrap();
 786        if assist_group.active_assist_id == Some(assist_id) {
 787            assist_group.active_assist_id = None;
 788            if assist_group.linked {
 789                for assist_id in &assist_group.assist_ids {
 790                    if let Some(decorations) = self.assists[assist_id].decorations.as_ref() {
 791                        decorations.prompt_editor.update(cx, |prompt_editor, cx| {
 792                            prompt_editor.set_show_cursor_when_unfocused(false, cx)
 793                        });
 794                    }
 795                }
 796            }
 797        }
 798    }
 799
 800    fn handle_prompt_editor_event(
 801        &mut self,
 802        prompt_editor: Entity<PromptEditor<BufferCodegen>>,
 803        event: &PromptEditorEvent,
 804        window: &mut Window,
 805        cx: &mut App,
 806    ) {
 807        let assist_id = prompt_editor.read(cx).id();
 808        match event {
 809            PromptEditorEvent::StartRequested => {
 810                self.start_assist(assist_id, window, cx);
 811            }
 812            PromptEditorEvent::StopRequested => {
 813                self.stop_assist(assist_id, cx);
 814            }
 815            PromptEditorEvent::ConfirmRequested { execute: _ } => {
 816                self.finish_assist(assist_id, false, window, cx);
 817            }
 818            PromptEditorEvent::CancelRequested => {
 819                self.finish_assist(assist_id, true, window, cx);
 820            }
 821            PromptEditorEvent::Resized { .. } => {
 822                // This only matters for the terminal inline assistant
 823            }
 824        }
 825    }
 826
 827    fn handle_editor_newline(&mut self, editor: Entity<Editor>, window: &mut Window, cx: &mut App) {
 828        let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
 829            return;
 830        };
 831
 832        if editor.read(cx).selections.count() == 1 {
 833            let (selection, buffer) = editor.update(cx, |editor, cx| {
 834                (
 835                    editor
 836                        .selections
 837                        .newest::<MultiBufferOffset>(&editor.display_snapshot(cx)),
 838                    editor.buffer().read(cx).snapshot(cx),
 839                )
 840            });
 841            for assist_id in &editor_assists.assist_ids {
 842                let assist = &self.assists[assist_id];
 843                let assist_range = assist.range.to_offset(&buffer);
 844                if assist_range.contains(&selection.start) && assist_range.contains(&selection.end)
 845                {
 846                    if matches!(assist.codegen.read(cx).status(cx), CodegenStatus::Pending) {
 847                        self.dismiss_assist(*assist_id, window, cx);
 848                    } else {
 849                        self.finish_assist(*assist_id, false, window, cx);
 850                    }
 851
 852                    return;
 853                }
 854            }
 855        }
 856
 857        cx.propagate();
 858    }
 859
 860    fn handle_editor_cancel(&mut self, editor: Entity<Editor>, window: &mut Window, cx: &mut App) {
 861        let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
 862            return;
 863        };
 864
 865        if editor.read(cx).selections.count() == 1 {
 866            let (selection, buffer) = editor.update(cx, |editor, cx| {
 867                (
 868                    editor
 869                        .selections
 870                        .newest::<MultiBufferOffset>(&editor.display_snapshot(cx)),
 871                    editor.buffer().read(cx).snapshot(cx),
 872                )
 873            });
 874            let mut closest_assist_fallback = None;
 875            for assist_id in &editor_assists.assist_ids {
 876                let assist = &self.assists[assist_id];
 877                let assist_range = assist.range.to_offset(&buffer);
 878                if assist.decorations.is_some() {
 879                    if assist_range.contains(&selection.start)
 880                        && assist_range.contains(&selection.end)
 881                    {
 882                        self.focus_assist(*assist_id, window, cx);
 883                        return;
 884                    } else {
 885                        let distance_from_selection = assist_range
 886                            .start
 887                            .0
 888                            .abs_diff(selection.start.0)
 889                            .min(assist_range.start.0.abs_diff(selection.end.0))
 890                            + assist_range
 891                                .end
 892                                .0
 893                                .abs_diff(selection.start.0)
 894                                .min(assist_range.end.0.abs_diff(selection.end.0));
 895                        match closest_assist_fallback {
 896                            Some((_, old_distance)) => {
 897                                if distance_from_selection < old_distance {
 898                                    closest_assist_fallback =
 899                                        Some((assist_id, distance_from_selection));
 900                                }
 901                            }
 902                            None => {
 903                                closest_assist_fallback = Some((assist_id, distance_from_selection))
 904                            }
 905                        }
 906                    }
 907                }
 908            }
 909
 910            if let Some((&assist_id, _)) = closest_assist_fallback {
 911                self.focus_assist(assist_id, window, cx);
 912            }
 913        }
 914
 915        cx.propagate();
 916    }
 917
 918    fn handle_editor_release(
 919        &mut self,
 920        editor: WeakEntity<Editor>,
 921        window: &mut Window,
 922        cx: &mut App,
 923    ) {
 924        if let Some(editor_assists) = self.assists_by_editor.get_mut(&editor) {
 925            for assist_id in editor_assists.assist_ids.clone() {
 926                self.finish_assist(assist_id, true, window, cx);
 927            }
 928        }
 929    }
 930
 931    fn handle_editor_change(&mut self, editor: Entity<Editor>, window: &mut Window, cx: &mut App) {
 932        let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
 933            return;
 934        };
 935        let Some(scroll_lock) = editor_assists.scroll_lock.as_ref() else {
 936            return;
 937        };
 938        let assist = &self.assists[&scroll_lock.assist_id];
 939        let Some(decorations) = assist.decorations.as_ref() else {
 940            return;
 941        };
 942
 943        editor.update(cx, |editor, cx| {
 944            let scroll_position = editor.scroll_position(cx);
 945            let target_scroll_top = editor
 946                .row_for_block(decorations.prompt_block_id, cx)?
 947                .as_f64()
 948                - scroll_lock.distance_from_top;
 949            if target_scroll_top != scroll_position.y {
 950                editor.set_scroll_position(point(scroll_position.x, target_scroll_top), window, cx);
 951            }
 952            Some(())
 953        });
 954    }
 955
 956    fn handle_editor_event(
 957        &mut self,
 958        editor: Entity<Editor>,
 959        event: &EditorEvent,
 960        window: &mut Window,
 961        cx: &mut App,
 962    ) {
 963        let Some(editor_assists) = self.assists_by_editor.get_mut(&editor.downgrade()) else {
 964            return;
 965        };
 966
 967        match event {
 968            EditorEvent::Edited { transaction_id } => {
 969                let buffer = editor.read(cx).buffer().read(cx);
 970                let edited_ranges =
 971                    buffer.edited_ranges_for_transaction::<MultiBufferOffset>(*transaction_id, cx);
 972                let snapshot = buffer.snapshot(cx);
 973
 974                for assist_id in editor_assists.assist_ids.clone() {
 975                    let assist = &self.assists[&assist_id];
 976                    if matches!(
 977                        assist.codegen.read(cx).status(cx),
 978                        CodegenStatus::Error(_) | CodegenStatus::Done
 979                    ) {
 980                        let assist_range = assist.range.to_offset(&snapshot);
 981                        if edited_ranges
 982                            .iter()
 983                            .any(|range| range.overlaps(&assist_range))
 984                        {
 985                            self.finish_assist(assist_id, false, window, cx);
 986                        }
 987                    }
 988                }
 989            }
 990            EditorEvent::ScrollPositionChanged { .. } => {
 991                if let Some(scroll_lock) = editor_assists.scroll_lock.as_ref() {
 992                    let assist = &self.assists[&scroll_lock.assist_id];
 993                    if let Some(decorations) = assist.decorations.as_ref() {
 994                        let distance_from_top = editor.update(cx, |editor, cx| {
 995                            let scroll_top = editor.scroll_position(cx).y;
 996                            let prompt_row = editor
 997                                .row_for_block(decorations.prompt_block_id, cx)?
 998                                .0 as ScrollOffset;
 999                            Some(prompt_row - scroll_top)
1000                        });
1001
1002                        if distance_from_top.is_none_or(|distance_from_top| {
1003                            distance_from_top != scroll_lock.distance_from_top
1004                        }) {
1005                            editor_assists.scroll_lock = None;
1006                        }
1007                    }
1008                }
1009            }
1010            EditorEvent::SelectionsChanged { .. } => {
1011                for assist_id in editor_assists.assist_ids.clone() {
1012                    let assist = &self.assists[&assist_id];
1013                    if let Some(decorations) = assist.decorations.as_ref()
1014                        && decorations
1015                            .prompt_editor
1016                            .focus_handle(cx)
1017                            .is_focused(window)
1018                    {
1019                        return;
1020                    }
1021                }
1022
1023                editor_assists.scroll_lock = None;
1024            }
1025            _ => {}
1026        }
1027    }
1028
1029    pub fn finish_assist(
1030        &mut self,
1031        assist_id: InlineAssistId,
1032        undo: bool,
1033        window: &mut Window,
1034        cx: &mut App,
1035    ) {
1036        if let Some(assist) = self.assists.get(&assist_id) {
1037            let assist_group_id = assist.group_id;
1038            if self.assist_groups[&assist_group_id].linked {
1039                for assist_id in self.unlink_assist_group(assist_group_id, window, cx) {
1040                    self.finish_assist(assist_id, undo, window, cx);
1041                }
1042                return;
1043            }
1044        }
1045
1046        self.dismiss_assist(assist_id, window, cx);
1047
1048        if let Some(assist) = self.assists.remove(&assist_id) {
1049            if let hash_map::Entry::Occupied(mut entry) = self.assist_groups.entry(assist.group_id)
1050            {
1051                entry.get_mut().assist_ids.retain(|id| *id != assist_id);
1052                if entry.get().assist_ids.is_empty() {
1053                    entry.remove();
1054                }
1055            }
1056
1057            if let hash_map::Entry::Occupied(mut entry) =
1058                self.assists_by_editor.entry(assist.editor.clone())
1059            {
1060                entry.get_mut().assist_ids.retain(|id| *id != assist_id);
1061                if entry.get().assist_ids.is_empty() {
1062                    entry.remove();
1063                    if let Some(editor) = assist.editor.upgrade() {
1064                        self.update_editor_highlights(&editor, cx);
1065                    }
1066                } else {
1067                    entry.get_mut().highlight_updates.send(()).ok();
1068                }
1069            }
1070
1071            let active_alternative = assist.codegen.read(cx).active_alternative().clone();
1072            let message_id = active_alternative.read(cx).message_id.clone();
1073
1074            if let Some(model) = LanguageModelRegistry::read_global(cx).inline_assistant_model() {
1075                let language_name = assist.editor.upgrade().and_then(|editor| {
1076                    let multibuffer = editor.read(cx).buffer().read(cx);
1077                    let snapshot = multibuffer.snapshot(cx);
1078                    let ranges = snapshot.range_to_buffer_ranges(assist.range.clone());
1079                    ranges
1080                        .first()
1081                        .and_then(|(buffer, _, _)| buffer.language())
1082                        .map(|language| language.name())
1083                });
1084                report_assistant_event(
1085                    AssistantEventData {
1086                        conversation_id: None,
1087                        kind: AssistantKind::Inline,
1088                        message_id,
1089                        phase: if undo {
1090                            AssistantPhase::Rejected
1091                        } else {
1092                            AssistantPhase::Accepted
1093                        },
1094                        model: model.model.telemetry_id(),
1095                        model_provider: model.model.provider_id().to_string(),
1096                        response_latency: None,
1097                        error_message: None,
1098                        language_name: language_name.map(|name| name.to_proto()),
1099                    },
1100                    Some(self.telemetry.clone()),
1101                    cx.http_client(),
1102                    model.model.api_key(cx),
1103                    cx.background_executor(),
1104                );
1105            }
1106
1107            if undo {
1108                assist.codegen.update(cx, |codegen, cx| codegen.undo(cx));
1109            } else {
1110                self.confirmed_assists.insert(assist_id, active_alternative);
1111            }
1112        }
1113    }
1114
1115    fn dismiss_assist(
1116        &mut self,
1117        assist_id: InlineAssistId,
1118        window: &mut Window,
1119        cx: &mut App,
1120    ) -> bool {
1121        let Some(assist) = self.assists.get_mut(&assist_id) else {
1122            return false;
1123        };
1124        let Some(editor) = assist.editor.upgrade() else {
1125            return false;
1126        };
1127        let Some(decorations) = assist.decorations.take() else {
1128            return false;
1129        };
1130
1131        editor.update(cx, |editor, cx| {
1132            let mut to_remove = decorations.removed_line_block_ids;
1133            to_remove.insert(decorations.prompt_block_id);
1134            to_remove.insert(decorations.end_block_id);
1135            if let Some(tool_description_block_id) = decorations.model_explanation {
1136                to_remove.insert(tool_description_block_id);
1137            }
1138            editor.remove_blocks(to_remove, None, cx);
1139        });
1140
1141        if decorations
1142            .prompt_editor
1143            .focus_handle(cx)
1144            .contains_focused(window, cx)
1145        {
1146            self.focus_next_assist(assist_id, window, cx);
1147        }
1148
1149        if let Some(editor_assists) = self.assists_by_editor.get_mut(&editor.downgrade()) {
1150            if editor_assists
1151                .scroll_lock
1152                .as_ref()
1153                .is_some_and(|lock| lock.assist_id == assist_id)
1154            {
1155                editor_assists.scroll_lock = None;
1156            }
1157            editor_assists.highlight_updates.send(()).ok();
1158        }
1159
1160        true
1161    }
1162
1163    fn focus_next_assist(&mut self, assist_id: InlineAssistId, window: &mut Window, cx: &mut App) {
1164        let Some(assist) = self.assists.get(&assist_id) else {
1165            return;
1166        };
1167
1168        let assist_group = &self.assist_groups[&assist.group_id];
1169        let assist_ix = assist_group
1170            .assist_ids
1171            .iter()
1172            .position(|id| *id == assist_id)
1173            .unwrap();
1174        let assist_ids = assist_group
1175            .assist_ids
1176            .iter()
1177            .skip(assist_ix + 1)
1178            .chain(assist_group.assist_ids.iter().take(assist_ix));
1179
1180        for assist_id in assist_ids {
1181            let assist = &self.assists[assist_id];
1182            if assist.decorations.is_some() {
1183                self.focus_assist(*assist_id, window, cx);
1184                return;
1185            }
1186        }
1187
1188        assist
1189            .editor
1190            .update(cx, |editor, cx| window.focus(&editor.focus_handle(cx)))
1191            .ok();
1192    }
1193
1194    fn focus_assist(&mut self, assist_id: InlineAssistId, window: &mut Window, cx: &mut App) {
1195        let Some(assist) = self.assists.get(&assist_id) else {
1196            return;
1197        };
1198
1199        if let Some(decorations) = assist.decorations.as_ref() {
1200            decorations.prompt_editor.update(cx, |prompt_editor, cx| {
1201                prompt_editor.editor.update(cx, |editor, cx| {
1202                    window.focus(&editor.focus_handle(cx));
1203                    editor.select_all(&SelectAll, window, cx);
1204                })
1205            });
1206        }
1207
1208        self.scroll_to_assist(assist_id, window, cx);
1209    }
1210
1211    pub fn scroll_to_assist(
1212        &mut self,
1213        assist_id: InlineAssistId,
1214        window: &mut Window,
1215        cx: &mut App,
1216    ) {
1217        let Some(assist) = self.assists.get(&assist_id) else {
1218            return;
1219        };
1220        let Some(editor) = assist.editor.upgrade() else {
1221            return;
1222        };
1223
1224        let position = assist.range.start;
1225        editor.update(cx, |editor, cx| {
1226            editor.change_selections(SelectionEffects::no_scroll(), window, cx, |selections| {
1227                selections.select_anchor_ranges([position..position])
1228            });
1229
1230            let mut scroll_target_range = None;
1231            if let Some(decorations) = assist.decorations.as_ref() {
1232                scroll_target_range = maybe!({
1233                    let top = editor.row_for_block(decorations.prompt_block_id, cx)?.0 as f64;
1234                    let bottom = editor.row_for_block(decorations.end_block_id, cx)?.0 as f64;
1235                    Some((top, bottom))
1236                });
1237                if scroll_target_range.is_none() {
1238                    log::error!("bug: failed to find blocks for scrolling to inline assist");
1239                }
1240            }
1241            let scroll_target_range = scroll_target_range.unwrap_or_else(|| {
1242                let snapshot = editor.snapshot(window, cx);
1243                let start_row = assist
1244                    .range
1245                    .start
1246                    .to_display_point(&snapshot.display_snapshot)
1247                    .row();
1248                let top = start_row.0 as ScrollOffset;
1249                let bottom = top + 1.0;
1250                (top, bottom)
1251            });
1252            let mut scroll_target_top = scroll_target_range.0;
1253            let mut scroll_target_bottom = scroll_target_range.1;
1254
1255            scroll_target_top -= editor.vertical_scroll_margin() as ScrollOffset;
1256            scroll_target_bottom += editor.vertical_scroll_margin() as ScrollOffset;
1257
1258            let height_in_lines = editor.visible_line_count().unwrap_or(0.);
1259            let scroll_top = editor.scroll_position(cx).y;
1260            let scroll_bottom = scroll_top + height_in_lines;
1261
1262            if scroll_target_top < scroll_top {
1263                editor.set_scroll_position(point(0., scroll_target_top), window, cx);
1264            } else if scroll_target_bottom > scroll_bottom {
1265                if (scroll_target_bottom - scroll_target_top) <= height_in_lines {
1266                    editor.set_scroll_position(
1267                        point(0., scroll_target_bottom - height_in_lines),
1268                        window,
1269                        cx,
1270                    );
1271                } else {
1272                    editor.set_scroll_position(point(0., scroll_target_top), window, cx);
1273                }
1274            }
1275        });
1276    }
1277
1278    fn unlink_assist_group(
1279        &mut self,
1280        assist_group_id: InlineAssistGroupId,
1281        window: &mut Window,
1282        cx: &mut App,
1283    ) -> Vec<InlineAssistId> {
1284        let assist_group = self.assist_groups.get_mut(&assist_group_id).unwrap();
1285        assist_group.linked = false;
1286
1287        for assist_id in &assist_group.assist_ids {
1288            let assist = self.assists.get_mut(assist_id).unwrap();
1289            if let Some(editor_decorations) = assist.decorations.as_ref() {
1290                editor_decorations
1291                    .prompt_editor
1292                    .update(cx, |prompt_editor, cx| prompt_editor.unlink(window, cx));
1293            }
1294        }
1295        assist_group.assist_ids.clone()
1296    }
1297
1298    pub fn start_assist(&mut self, assist_id: InlineAssistId, window: &mut Window, cx: &mut App) {
1299        let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
1300            assist
1301        } else {
1302            return;
1303        };
1304
1305        let assist_group_id = assist.group_id;
1306        if self.assist_groups[&assist_group_id].linked {
1307            for assist_id in self.unlink_assist_group(assist_group_id, window, cx) {
1308                self.start_assist(assist_id, window, cx);
1309            }
1310            return;
1311        }
1312
1313        let Some((user_prompt, mention_set)) = assist.user_prompt(cx).zip(assist.mention_set(cx))
1314        else {
1315            return;
1316        };
1317
1318        self.prompt_history.retain(|prompt| *prompt != user_prompt);
1319        self.prompt_history.push_back(user_prompt.clone());
1320        if self.prompt_history.len() > PROMPT_HISTORY_MAX_LEN {
1321            self.prompt_history.pop_front();
1322        }
1323
1324        let Some(ConfiguredModel { model, .. }) =
1325            LanguageModelRegistry::read_global(cx).inline_assistant_model()
1326        else {
1327            return;
1328        };
1329
1330        let context_task = load_context(&mention_set, cx).shared();
1331        assist
1332            .codegen
1333            .update(cx, |codegen, cx| {
1334                codegen.start(model, user_prompt, context_task, cx)
1335            })
1336            .log_err();
1337    }
1338
1339    pub fn stop_assist(&mut self, assist_id: InlineAssistId, cx: &mut App) {
1340        let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
1341            assist
1342        } else {
1343            return;
1344        };
1345
1346        assist.codegen.update(cx, |codegen, cx| codegen.stop(cx));
1347    }
1348
1349    fn update_editor_highlights(&self, editor: &Entity<Editor>, cx: &mut App) {
1350        let mut gutter_pending_ranges = Vec::new();
1351        let mut gutter_transformed_ranges = Vec::new();
1352        let mut foreground_ranges = Vec::new();
1353        let mut inserted_row_ranges = Vec::new();
1354        let empty_assist_ids = Vec::new();
1355        let assist_ids = self
1356            .assists_by_editor
1357            .get(&editor.downgrade())
1358            .map_or(&empty_assist_ids, |editor_assists| {
1359                &editor_assists.assist_ids
1360            });
1361
1362        for assist_id in assist_ids {
1363            if let Some(assist) = self.assists.get(assist_id) {
1364                let codegen = assist.codegen.read(cx);
1365                let buffer = codegen.buffer(cx).read(cx).read(cx);
1366                foreground_ranges.extend(codegen.last_equal_ranges(cx).iter().cloned());
1367
1368                let pending_range =
1369                    codegen.edit_position(cx).unwrap_or(assist.range.start)..assist.range.end;
1370                if pending_range.end.to_offset(&buffer) > pending_range.start.to_offset(&buffer) {
1371                    gutter_pending_ranges.push(pending_range);
1372                }
1373
1374                if let Some(edit_position) = codegen.edit_position(cx) {
1375                    let edited_range = assist.range.start..edit_position;
1376                    if edited_range.end.to_offset(&buffer) > edited_range.start.to_offset(&buffer) {
1377                        gutter_transformed_ranges.push(edited_range);
1378                    }
1379                }
1380
1381                if assist.decorations.is_some() {
1382                    inserted_row_ranges
1383                        .extend(codegen.diff(cx).inserted_row_ranges.iter().cloned());
1384                }
1385            }
1386        }
1387
1388        let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
1389        merge_ranges(&mut foreground_ranges, &snapshot);
1390        merge_ranges(&mut gutter_pending_ranges, &snapshot);
1391        merge_ranges(&mut gutter_transformed_ranges, &snapshot);
1392        editor.update(cx, |editor, cx| {
1393            enum GutterPendingRange {}
1394            if gutter_pending_ranges.is_empty() {
1395                editor.clear_gutter_highlights::<GutterPendingRange>(cx);
1396            } else {
1397                editor.highlight_gutter::<GutterPendingRange>(
1398                    gutter_pending_ranges,
1399                    |cx| cx.theme().status().info_background,
1400                    cx,
1401                )
1402            }
1403
1404            enum GutterTransformedRange {}
1405            if gutter_transformed_ranges.is_empty() {
1406                editor.clear_gutter_highlights::<GutterTransformedRange>(cx);
1407            } else {
1408                editor.highlight_gutter::<GutterTransformedRange>(
1409                    gutter_transformed_ranges,
1410                    |cx| cx.theme().status().info,
1411                    cx,
1412                )
1413            }
1414
1415            if foreground_ranges.is_empty() {
1416                editor.clear_highlights::<InlineAssist>(cx);
1417            } else {
1418                editor.highlight_text::<InlineAssist>(
1419                    foreground_ranges,
1420                    HighlightStyle {
1421                        fade_out: Some(0.6),
1422                        ..Default::default()
1423                    },
1424                    cx,
1425                );
1426            }
1427
1428            editor.clear_row_highlights::<InlineAssist>();
1429            for row_range in inserted_row_ranges {
1430                editor.highlight_rows::<InlineAssist>(
1431                    row_range,
1432                    cx.theme().status().info_background,
1433                    Default::default(),
1434                    cx,
1435                );
1436            }
1437        });
1438    }
1439
1440    fn update_editor_blocks(
1441        &mut self,
1442        editor: &Entity<Editor>,
1443        assist_id: InlineAssistId,
1444        window: &mut Window,
1445        cx: &mut App,
1446    ) {
1447        let Some(assist) = self.assists.get_mut(&assist_id) else {
1448            return;
1449        };
1450        let Some(decorations) = assist.decorations.as_mut() else {
1451            return;
1452        };
1453
1454        let codegen = assist.codegen.read(cx);
1455        let old_snapshot = codegen.snapshot(cx);
1456        let old_buffer = codegen.old_buffer(cx);
1457        let deleted_row_ranges = codegen.diff(cx).deleted_row_ranges.clone();
1458
1459        editor.update(cx, |editor, cx| {
1460            let old_blocks = mem::take(&mut decorations.removed_line_block_ids);
1461            editor.remove_blocks(old_blocks, None, cx);
1462
1463            let mut new_blocks = Vec::new();
1464            for (new_row, old_row_range) in deleted_row_ranges {
1465                let (_, buffer_start) = old_snapshot
1466                    .point_to_buffer_offset(Point::new(*old_row_range.start(), 0))
1467                    .unwrap();
1468                let (_, buffer_end) = old_snapshot
1469                    .point_to_buffer_offset(Point::new(
1470                        *old_row_range.end(),
1471                        old_snapshot.line_len(MultiBufferRow(*old_row_range.end())),
1472                    ))
1473                    .unwrap();
1474
1475                let deleted_lines_editor = cx.new(|cx| {
1476                    let multi_buffer =
1477                        cx.new(|_| MultiBuffer::without_headers(language::Capability::ReadOnly));
1478                    multi_buffer.update(cx, |multi_buffer, cx| {
1479                        multi_buffer.push_excerpts(
1480                            old_buffer.clone(),
1481                            // todo(lw): buffer_start and buffer_end might come from different snapshots!
1482                            Some(ExcerptRange::new(buffer_start..buffer_end)),
1483                            cx,
1484                        );
1485                    });
1486
1487                    enum DeletedLines {}
1488                    let mut editor = Editor::for_multibuffer(multi_buffer, None, window, cx);
1489                    editor.disable_scrollbars_and_minimap(window, cx);
1490                    editor.set_soft_wrap_mode(language::language_settings::SoftWrap::None, cx);
1491                    editor.set_show_wrap_guides(false, cx);
1492                    editor.set_show_gutter(false, cx);
1493                    editor.set_offset_content(false, cx);
1494                    editor.scroll_manager.set_forbid_vertical_scroll(true);
1495                    editor.set_read_only(true);
1496                    editor.set_show_edit_predictions(Some(false), window, cx);
1497                    editor.highlight_rows::<DeletedLines>(
1498                        Anchor::min()..Anchor::max(),
1499                        cx.theme().status().deleted_background,
1500                        Default::default(),
1501                        cx,
1502                    );
1503                    editor
1504                });
1505
1506                let height =
1507                    deleted_lines_editor.update(cx, |editor, cx| editor.max_point(cx).row().0 + 1);
1508                new_blocks.push(BlockProperties {
1509                    placement: BlockPlacement::Above(new_row),
1510                    height: Some(height),
1511                    style: BlockStyle::Flex,
1512                    render: Arc::new(move |cx| {
1513                        div()
1514                            .block_mouse_except_scroll()
1515                            .bg(cx.theme().status().deleted_background)
1516                            .size_full()
1517                            .h(height as f32 * cx.window.line_height())
1518                            .pl(cx.margins.gutter.full_width())
1519                            .child(deleted_lines_editor.clone())
1520                            .into_any_element()
1521                    }),
1522                    priority: 0,
1523                });
1524            }
1525
1526            decorations.removed_line_block_ids = editor
1527                .insert_blocks(new_blocks, None, cx)
1528                .into_iter()
1529                .collect();
1530        })
1531    }
1532
1533    fn resolve_inline_assist_target(
1534        workspace: &mut Workspace,
1535        agent_panel: Option<Entity<AgentPanel>>,
1536        window: &mut Window,
1537        cx: &mut App,
1538    ) -> Option<InlineAssistTarget> {
1539        if let Some(terminal_panel) = workspace.panel::<TerminalPanel>(cx)
1540            && terminal_panel
1541                .read(cx)
1542                .focus_handle(cx)
1543                .contains_focused(window, cx)
1544            && let Some(terminal_view) = terminal_panel.read(cx).pane().and_then(|pane| {
1545                pane.read(cx)
1546                    .active_item()
1547                    .and_then(|t| t.downcast::<TerminalView>())
1548            })
1549        {
1550            return Some(InlineAssistTarget::Terminal(terminal_view));
1551        }
1552
1553        let text_thread_editor = agent_panel
1554            .and_then(|panel| panel.read(cx).active_text_thread_editor())
1555            .and_then(|editor| {
1556                let editor = &editor.read(cx).editor().clone();
1557                if editor.read(cx).is_focused(window) {
1558                    Some(editor.clone())
1559                } else {
1560                    None
1561                }
1562            });
1563
1564        if let Some(text_thread_editor) = text_thread_editor {
1565            Some(InlineAssistTarget::Editor(text_thread_editor))
1566        } else if let Some(workspace_editor) = workspace
1567            .active_item(cx)
1568            .and_then(|item| item.act_as::<Editor>(cx))
1569        {
1570            Some(InlineAssistTarget::Editor(workspace_editor))
1571        } else {
1572            workspace
1573                .active_item(cx)
1574                .and_then(|item| item.act_as::<TerminalView>(cx))
1575                .map(InlineAssistTarget::Terminal)
1576        }
1577    }
1578}
1579
1580struct EditorInlineAssists {
1581    assist_ids: Vec<InlineAssistId>,
1582    scroll_lock: Option<InlineAssistScrollLock>,
1583    highlight_updates: watch::Sender<()>,
1584    _update_highlights: Task<Result<()>>,
1585    _subscriptions: Vec<gpui::Subscription>,
1586}
1587
1588struct InlineAssistScrollLock {
1589    assist_id: InlineAssistId,
1590    distance_from_top: ScrollOffset,
1591}
1592
1593impl EditorInlineAssists {
1594    fn new(editor: &Entity<Editor>, window: &mut Window, cx: &mut App) -> Self {
1595        let (highlight_updates_tx, mut highlight_updates_rx) = watch::channel(());
1596        Self {
1597            assist_ids: Vec::new(),
1598            scroll_lock: None,
1599            highlight_updates: highlight_updates_tx,
1600            _update_highlights: cx.spawn({
1601                let editor = editor.downgrade();
1602                async move |cx| {
1603                    while let Ok(()) = highlight_updates_rx.changed().await {
1604                        let editor = editor.upgrade().context("editor was dropped")?;
1605                        cx.update_global(|assistant: &mut InlineAssistant, cx| {
1606                            assistant.update_editor_highlights(&editor, cx);
1607                        })?;
1608                    }
1609                    Ok(())
1610                }
1611            }),
1612            _subscriptions: vec![
1613                cx.observe_release_in(editor, window, {
1614                    let editor = editor.downgrade();
1615                    |_, window, cx| {
1616                        InlineAssistant::update_global(cx, |this, cx| {
1617                            this.handle_editor_release(editor, window, cx);
1618                        })
1619                    }
1620                }),
1621                window.observe(editor, cx, move |editor, window, cx| {
1622                    InlineAssistant::update_global(cx, |this, cx| {
1623                        this.handle_editor_change(editor, window, cx)
1624                    })
1625                }),
1626                window.subscribe(editor, cx, move |editor, event, window, cx| {
1627                    InlineAssistant::update_global(cx, |this, cx| {
1628                        this.handle_editor_event(editor, event, window, cx)
1629                    })
1630                }),
1631                editor.update(cx, |editor, cx| {
1632                    let editor_handle = cx.entity().downgrade();
1633                    editor.register_action(move |_: &editor::actions::Newline, window, cx| {
1634                        InlineAssistant::update_global(cx, |this, cx| {
1635                            if let Some(editor) = editor_handle.upgrade() {
1636                                this.handle_editor_newline(editor, window, cx)
1637                            }
1638                        })
1639                    })
1640                }),
1641                editor.update(cx, |editor, cx| {
1642                    let editor_handle = cx.entity().downgrade();
1643                    editor.register_action(move |_: &editor::actions::Cancel, window, cx| {
1644                        InlineAssistant::update_global(cx, |this, cx| {
1645                            if let Some(editor) = editor_handle.upgrade() {
1646                                this.handle_editor_cancel(editor, window, cx)
1647                            }
1648                        })
1649                    })
1650                }),
1651            ],
1652        }
1653    }
1654}
1655
1656struct InlineAssistGroup {
1657    assist_ids: Vec<InlineAssistId>,
1658    linked: bool,
1659    active_assist_id: Option<InlineAssistId>,
1660}
1661
1662impl InlineAssistGroup {
1663    fn new() -> Self {
1664        Self {
1665            assist_ids: Vec::new(),
1666            linked: true,
1667            active_assist_id: None,
1668        }
1669    }
1670}
1671
1672fn build_assist_editor_renderer(editor: &Entity<PromptEditor<BufferCodegen>>) -> RenderBlock {
1673    let editor = editor.clone();
1674
1675    Arc::new(move |cx: &mut BlockContext| {
1676        let editor_margins = editor.read(cx).editor_margins();
1677
1678        *editor_margins.lock() = *cx.margins;
1679        editor.clone().into_any_element()
1680    })
1681}
1682
1683#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
1684struct InlineAssistGroupId(usize);
1685
1686impl InlineAssistGroupId {
1687    fn post_inc(&mut self) -> InlineAssistGroupId {
1688        let id = *self;
1689        self.0 += 1;
1690        id
1691    }
1692}
1693
1694pub struct InlineAssist {
1695    group_id: InlineAssistGroupId,
1696    range: Range<Anchor>,
1697    editor: WeakEntity<Editor>,
1698    decorations: Option<InlineAssistDecorations>,
1699    codegen: Entity<BufferCodegen>,
1700    _subscriptions: Vec<Subscription>,
1701    workspace: WeakEntity<Workspace>,
1702}
1703
1704impl InlineAssist {
1705    fn new(
1706        assist_id: InlineAssistId,
1707        group_id: InlineAssistGroupId,
1708        editor: &Entity<Editor>,
1709        prompt_editor: &Entity<PromptEditor<BufferCodegen>>,
1710        prompt_block_id: CustomBlockId,
1711        tool_description_block_id: CustomBlockId,
1712        end_block_id: CustomBlockId,
1713        range: Range<Anchor>,
1714        codegen: Entity<BufferCodegen>,
1715        workspace: WeakEntity<Workspace>,
1716        window: &mut Window,
1717        cx: &mut App,
1718    ) -> Self {
1719        let prompt_editor_focus_handle = prompt_editor.focus_handle(cx);
1720        InlineAssist {
1721            group_id,
1722            editor: editor.downgrade(),
1723            decorations: Some(InlineAssistDecorations {
1724                prompt_block_id,
1725                prompt_editor: prompt_editor.clone(),
1726                removed_line_block_ids: Default::default(),
1727                model_explanation: Some(tool_description_block_id),
1728                end_block_id,
1729            }),
1730            range,
1731            codegen: codegen.clone(),
1732            workspace,
1733            _subscriptions: vec![
1734                window.on_focus_in(&prompt_editor_focus_handle, cx, move |_, cx| {
1735                    InlineAssistant::update_global(cx, |this, cx| {
1736                        this.handle_prompt_editor_focus_in(assist_id, cx)
1737                    })
1738                }),
1739                window.on_focus_out(&prompt_editor_focus_handle, cx, move |_, _, cx| {
1740                    InlineAssistant::update_global(cx, |this, cx| {
1741                        this.handle_prompt_editor_focus_out(assist_id, cx)
1742                    })
1743                }),
1744                window.subscribe(prompt_editor, cx, |prompt_editor, event, window, cx| {
1745                    InlineAssistant::update_global(cx, |this, cx| {
1746                        this.handle_prompt_editor_event(prompt_editor, event, window, cx)
1747                    })
1748                }),
1749                window.observe(&codegen, cx, {
1750                    let editor = editor.downgrade();
1751                    move |_, window, cx| {
1752                        if let Some(editor) = editor.upgrade() {
1753                            InlineAssistant::update_global(cx, |this, cx| {
1754                                if let Some(editor_assists) =
1755                                    this.assists_by_editor.get_mut(&editor.downgrade())
1756                                {
1757                                    editor_assists.highlight_updates.send(()).ok();
1758                                }
1759
1760                                this.update_editor_blocks(&editor, assist_id, window, cx);
1761                            })
1762                        }
1763                    }
1764                }),
1765                window.subscribe(&codegen, cx, move |codegen, event, window, cx| {
1766                    InlineAssistant::update_global(cx, |this, cx| match event {
1767                        CodegenEvent::Undone => this.finish_assist(assist_id, false, window, cx),
1768                        CodegenEvent::Finished => {
1769                            let assist = if let Some(assist) = this.assists.get(&assist_id) {
1770                                assist
1771                            } else {
1772                                return;
1773                            };
1774
1775                            if let CodegenStatus::Error(error) = codegen.read(cx).status(cx)
1776                                && assist.decorations.is_none()
1777                                && let Some(workspace) = assist.workspace.upgrade()
1778                            {
1779                                #[cfg(any(test, feature = "test-support"))]
1780                                if let Some(sender) = &mut this._inline_assistant_completions {
1781                                    sender
1782                                        .unbounded_send(Err(anyhow::anyhow!(
1783                                            "Inline assistant error: {}",
1784                                            error
1785                                        )))
1786                                        .ok();
1787                                }
1788
1789                                let error = format!("Inline assistant error: {}", error);
1790                                workspace.update(cx, |workspace, cx| {
1791                                    struct InlineAssistantError;
1792
1793                                    let id = NotificationId::composite::<InlineAssistantError>(
1794                                        assist_id.0,
1795                                    );
1796
1797                                    workspace.show_toast(Toast::new(id, error), cx);
1798                                })
1799                            } else {
1800                                #[cfg(any(test, feature = "test-support"))]
1801                                if let Some(sender) = &mut this._inline_assistant_completions {
1802                                    sender.unbounded_send(Ok(assist_id)).ok();
1803                                }
1804                            }
1805
1806                            if assist.decorations.is_none() {
1807                                this.finish_assist(assist_id, false, window, cx);
1808                            }
1809                        }
1810                    })
1811                }),
1812            ],
1813        }
1814    }
1815
1816    fn user_prompt(&self, cx: &App) -> Option<String> {
1817        let decorations = self.decorations.as_ref()?;
1818        Some(decorations.prompt_editor.read(cx).prompt(cx))
1819    }
1820
1821    fn mention_set(&self, cx: &App) -> Option<Entity<MentionSet>> {
1822        let decorations = self.decorations.as_ref()?;
1823        Some(decorations.prompt_editor.read(cx).mention_set().clone())
1824    }
1825}
1826
1827struct InlineAssistDecorations {
1828    prompt_block_id: CustomBlockId,
1829    prompt_editor: Entity<PromptEditor<BufferCodegen>>,
1830    removed_line_block_ids: HashSet<CustomBlockId>,
1831    model_explanation: Option<CustomBlockId>,
1832    end_block_id: CustomBlockId,
1833}
1834
1835struct AssistantCodeActionProvider {
1836    editor: WeakEntity<Editor>,
1837    workspace: WeakEntity<Workspace>,
1838}
1839
1840const ASSISTANT_CODE_ACTION_PROVIDER_ID: &str = "assistant";
1841
1842impl CodeActionProvider for AssistantCodeActionProvider {
1843    fn id(&self) -> Arc<str> {
1844        ASSISTANT_CODE_ACTION_PROVIDER_ID.into()
1845    }
1846
1847    fn code_actions(
1848        &self,
1849        buffer: &Entity<Buffer>,
1850        range: Range<text::Anchor>,
1851        _: &mut Window,
1852        cx: &mut App,
1853    ) -> Task<Result<Vec<CodeAction>>> {
1854        if !AgentSettings::get_global(cx).enabled(cx) {
1855            return Task::ready(Ok(Vec::new()));
1856        }
1857
1858        let snapshot = buffer.read(cx).snapshot();
1859        let mut range = range.to_point(&snapshot);
1860
1861        // Expand the range to line boundaries.
1862        range.start.column = 0;
1863        range.end.column = snapshot.line_len(range.end.row);
1864
1865        let mut has_diagnostics = false;
1866        for diagnostic in snapshot.diagnostics_in_range::<_, Point>(range.clone(), false) {
1867            range.start = cmp::min(range.start, diagnostic.range.start);
1868            range.end = cmp::max(range.end, diagnostic.range.end);
1869            has_diagnostics = true;
1870        }
1871        if has_diagnostics {
1872            let symbols_containing_start = snapshot.symbols_containing(range.start, None);
1873            if let Some(symbol) = symbols_containing_start.last() {
1874                range.start = cmp::min(range.start, symbol.range.start.to_point(&snapshot));
1875                range.end = cmp::max(range.end, symbol.range.end.to_point(&snapshot));
1876            }
1877            let symbols_containing_end = snapshot.symbols_containing(range.end, None);
1878            if let Some(symbol) = symbols_containing_end.last() {
1879                range.start = cmp::min(range.start, symbol.range.start.to_point(&snapshot));
1880                range.end = cmp::max(range.end, symbol.range.end.to_point(&snapshot));
1881            }
1882
1883            Task::ready(Ok(vec![CodeAction {
1884                server_id: language::LanguageServerId(0),
1885                range: snapshot.anchor_before(range.start)..snapshot.anchor_after(range.end),
1886                lsp_action: LspAction::Action(Box::new(lsp::CodeAction {
1887                    title: "Fix with Assistant".into(),
1888                    ..Default::default()
1889                })),
1890                resolved: true,
1891            }]))
1892        } else {
1893            Task::ready(Ok(Vec::new()))
1894        }
1895    }
1896
1897    fn apply_code_action(
1898        &self,
1899        buffer: Entity<Buffer>,
1900        action: CodeAction,
1901        excerpt_id: ExcerptId,
1902        _push_to_history: bool,
1903        window: &mut Window,
1904        cx: &mut App,
1905    ) -> Task<Result<ProjectTransaction>> {
1906        let editor = self.editor.clone();
1907        let workspace = self.workspace.clone();
1908        let prompt_store = PromptStore::global(cx);
1909        window.spawn(cx, async move |cx| {
1910            let workspace = workspace.upgrade().context("workspace was released")?;
1911            let thread_store = cx.update(|_window, cx| {
1912                anyhow::Ok(
1913                    workspace
1914                        .read(cx)
1915                        .panel::<AgentPanel>(cx)
1916                        .context("missing agent panel")?
1917                        .read(cx)
1918                        .thread_store()
1919                        .clone(),
1920                )
1921            })??;
1922            let editor = editor.upgrade().context("editor was released")?;
1923            let range = editor
1924                .update(cx, |editor, cx| {
1925                    editor.buffer().update(cx, |multibuffer, cx| {
1926                        let buffer = buffer.read(cx);
1927                        let multibuffer_snapshot = multibuffer.read(cx);
1928
1929                        let old_context_range =
1930                            multibuffer_snapshot.context_range_for_excerpt(excerpt_id)?;
1931                        let mut new_context_range = old_context_range.clone();
1932                        if action
1933                            .range
1934                            .start
1935                            .cmp(&old_context_range.start, buffer)
1936                            .is_lt()
1937                        {
1938                            new_context_range.start = action.range.start;
1939                        }
1940                        if action.range.end.cmp(&old_context_range.end, buffer).is_gt() {
1941                            new_context_range.end = action.range.end;
1942                        }
1943                        drop(multibuffer_snapshot);
1944
1945                        if new_context_range != old_context_range {
1946                            multibuffer.resize_excerpt(excerpt_id, new_context_range, cx);
1947                        }
1948
1949                        let multibuffer_snapshot = multibuffer.read(cx);
1950                        multibuffer_snapshot.anchor_range_in_excerpt(excerpt_id, action.range)
1951                    })
1952                })?
1953                .context("invalid range")?;
1954
1955            let prompt_store = prompt_store.await.ok();
1956            cx.update_global(|assistant: &mut InlineAssistant, window, cx| {
1957                let assist_id = assistant.suggest_assist(
1958                    &editor,
1959                    range,
1960                    "Fix Diagnostics".into(),
1961                    None,
1962                    true,
1963                    workspace,
1964                    thread_store,
1965                    prompt_store,
1966                    window,
1967                    cx,
1968                );
1969                assistant.start_assist(assist_id, window, cx);
1970            })?;
1971
1972            Ok(ProjectTransaction::default())
1973        })
1974    }
1975}
1976
1977fn merge_ranges(ranges: &mut Vec<Range<Anchor>>, buffer: &MultiBufferSnapshot) {
1978    ranges.sort_unstable_by(|a, b| {
1979        a.start
1980            .cmp(&b.start, buffer)
1981            .then_with(|| b.end.cmp(&a.end, buffer))
1982    });
1983
1984    let mut ix = 0;
1985    while ix + 1 < ranges.len() {
1986        let b = ranges[ix + 1].clone();
1987        let a = &mut ranges[ix];
1988        if a.end.cmp(&b.start, buffer).is_gt() {
1989            if a.end.cmp(&b.end, buffer).is_lt() {
1990                a.end = b.end;
1991            }
1992            ranges.remove(ix + 1);
1993        } else {
1994            ix += 1;
1995        }
1996    }
1997}
1998
1999#[cfg(any(test, feature = "test-support"))]
2000pub mod test {
2001    use std::sync::Arc;
2002
2003    use agent::HistoryStore;
2004    use assistant_text_thread::TextThreadStore;
2005    use client::{Client, UserStore};
2006    use editor::{Editor, MultiBuffer, MultiBufferOffset};
2007    use fs::FakeFs;
2008    use futures::channel::mpsc;
2009    use gpui::{AppContext, TestAppContext, UpdateGlobal as _};
2010    use language::Buffer;
2011    use language_model::LanguageModelRegistry;
2012    use project::Project;
2013    use prompt_store::PromptBuilder;
2014    use smol::stream::StreamExt as _;
2015    use util::test::marked_text_ranges;
2016    use workspace::Workspace;
2017
2018    use crate::InlineAssistant;
2019
2020    pub fn run_inline_assistant_test<SetupF, TestF>(
2021        base_buffer: String,
2022        prompt: String,
2023        setup: SetupF,
2024        test: TestF,
2025        cx: &mut TestAppContext,
2026    ) -> String
2027    where
2028        SetupF: FnOnce(&mut gpui::VisualTestContext),
2029        TestF: FnOnce(&mut gpui::VisualTestContext),
2030    {
2031        let fs = FakeFs::new(cx.executor());
2032        let app_state = cx.update(|cx| workspace::AppState::test(cx));
2033        let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
2034        let http = Arc::new(reqwest_client::ReqwestClient::user_agent("agent tests").unwrap());
2035        let client = cx.update(|cx| {
2036            cx.set_http_client(http);
2037            Client::production(cx)
2038        });
2039        let mut inline_assistant =
2040            InlineAssistant::new(fs.clone(), prompt_builder, client.telemetry().clone());
2041
2042        let (tx, mut completion_rx) = mpsc::unbounded();
2043        inline_assistant.set_completion_receiver(tx);
2044
2045        // Initialize settings and client
2046        cx.update(|cx| {
2047            gpui_tokio::init(cx);
2048            settings::init(cx);
2049            client::init(&client, cx);
2050            workspace::init(app_state.clone(), cx);
2051            let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
2052            language_model::init(client.clone(), cx);
2053            language_models::init(user_store, client.clone(), cx);
2054
2055            cx.set_global(inline_assistant);
2056        });
2057
2058        let project = cx
2059            .executor()
2060            .block_test(async { Project::test(fs.clone(), [], cx).await });
2061
2062        // Create workspace with window
2063        let (workspace, cx) = cx.add_window_view(|window, cx| {
2064            window.activate_window();
2065            Workspace::new(None, project.clone(), app_state.clone(), window, cx)
2066        });
2067
2068        setup(cx);
2069
2070        let (_editor, buffer) = cx.update(|window, cx| {
2071            let buffer = cx.new(|cx| Buffer::local("", cx));
2072            let multibuffer = cx.new(|cx| MultiBuffer::singleton(buffer.clone(), cx));
2073            let editor = cx.new(|cx| Editor::for_multibuffer(multibuffer, None, window, cx));
2074            editor.update(cx, |editor, cx| {
2075                let (unmarked_text, selection_ranges) = marked_text_ranges(&base_buffer, true);
2076                editor.set_text(unmarked_text, window, cx);
2077                editor.change_selections(Default::default(), window, cx, |s| {
2078                    s.select_ranges(
2079                        selection_ranges.into_iter().map(|range| {
2080                            MultiBufferOffset(range.start)..MultiBufferOffset(range.end)
2081                        }),
2082                    )
2083                })
2084            });
2085
2086            let text_thread_store = cx.new(|cx| TextThreadStore::fake(project.clone(), cx));
2087            let history_store = cx.new(|cx| HistoryStore::new(text_thread_store, cx));
2088
2089            // Add editor to workspace
2090            workspace.update(cx, |workspace, cx| {
2091                workspace.add_item_to_active_pane(Box::new(editor.clone()), None, true, window, cx);
2092            });
2093
2094            // Call assist method
2095            InlineAssistant::update_global(cx, |inline_assistant, cx| {
2096                let assist_id = inline_assistant
2097                    .assist(
2098                        &editor,
2099                        workspace.downgrade(),
2100                        project.downgrade(),
2101                        history_store, // thread_store
2102                        None,          // prompt_store
2103                        Some(prompt),
2104                        window,
2105                        cx,
2106                    )
2107                    .unwrap();
2108
2109                inline_assistant.start_assist(assist_id, window, cx);
2110            });
2111
2112            (editor, buffer)
2113        });
2114
2115        cx.run_until_parked();
2116
2117        test(cx);
2118
2119        cx.executor()
2120            .block_test(async { completion_rx.next().await });
2121
2122        buffer.read_with(cx, |buffer, _| buffer.text())
2123    }
2124
2125    #[allow(unused)]
2126    pub fn test_inline_assistant(
2127        base_buffer: &'static str,
2128        llm_output: &'static str,
2129        cx: &mut TestAppContext,
2130    ) -> String {
2131        run_inline_assistant_test(
2132            base_buffer.to_string(),
2133            "Prompt doesn't matter because we're using a fake model".to_string(),
2134            |cx| {
2135                cx.update(|_, cx| LanguageModelRegistry::test(cx));
2136            },
2137            |cx| {
2138                let fake_model = cx.update(|_, cx| {
2139                    LanguageModelRegistry::global(cx)
2140                        .update(cx, |registry, _| registry.fake_model())
2141                });
2142                let fake = fake_model.as_fake();
2143
2144                // let fake = fake_model;
2145                fake.send_last_completion_stream_text_chunk(llm_output.to_string());
2146                fake.end_last_completion_stream();
2147
2148                // Run again to process the model's response
2149                cx.run_until_parked();
2150            },
2151            cx,
2152        )
2153    }
2154}