1use crate::{
   2    assistant_settings::AssistantSettings, humanize_token_count, prompts::generate_content_prompt,
   3    AssistantPanel, AssistantPanelEvent, CompletionProvider, Hunk, LanguageModelRequest,
   4    LanguageModelRequestMessage, Role, StreamingDiff,
   5};
   6use anyhow::{anyhow, Context as _, Result};
   7use client::telemetry::Telemetry;
   8use collections::{hash_map, HashMap, HashSet, VecDeque};
   9use editor::{
  10    actions::{MoveDown, MoveUp, SelectAll},
  11    display_map::{
  12        BlockContext, BlockDisposition, BlockId, BlockProperties, BlockStyle, RenderBlock,
  13        ToDisplayPoint,
  14    },
  15    Anchor, AnchorRangeExt, Editor, EditorElement, EditorEvent, EditorMode, EditorStyle,
  16    ExcerptRange, GutterDimensions, MultiBuffer, MultiBufferSnapshot, ToOffset, ToPoint,
  17};
  18use fs::Fs;
  19use futures::{channel::mpsc, SinkExt, Stream, StreamExt};
  20use gpui::{
  21    point, AppContext, EventEmitter, FocusHandle, FocusableView, FontStyle, Global, HighlightStyle,
  22    Model, ModelContext, Subscription, Task, TextStyle, UpdateGlobal, View, ViewContext, WeakView,
  23    WhiteSpace, WindowContext,
  24};
  25use language::{Buffer, Point, Selection, TransactionId};
  26use multi_buffer::MultiBufferRow;
  27use parking_lot::Mutex;
  28use rope::Rope;
  29use settings::{update_settings_file, Settings};
  30use similar::TextDiff;
  31use std::{
  32    cmp, mem,
  33    ops::{Range, RangeInclusive},
  34    pin::Pin,
  35    sync::Arc,
  36    task::{self, Poll},
  37    time::{Duration, Instant},
  38};
  39use theme::ThemeSettings;
  40use ui::{prelude::*, ContextMenu, PopoverMenu, Tooltip};
  41use util::RangeExt;
  42use workspace::{notifications::NotificationId, Toast, Workspace};
  43
  44pub fn init(fs: Arc<dyn Fs>, telemetry: Arc<Telemetry>, cx: &mut AppContext) {
  45    cx.set_global(InlineAssistant::new(fs, telemetry));
  46}
  47
  48const PROMPT_HISTORY_MAX_LEN: usize = 20;
  49
  50pub struct InlineAssistant {
  51    next_assist_id: InlineAssistId,
  52    next_assist_group_id: InlineAssistGroupId,
  53    assists: HashMap<InlineAssistId, InlineAssist>,
  54    assists_by_editor: HashMap<WeakView<Editor>, EditorInlineAssists>,
  55    assist_groups: HashMap<InlineAssistGroupId, InlineAssistGroup>,
  56    prompt_history: VecDeque<String>,
  57    telemetry: Option<Arc<Telemetry>>,
  58    fs: Arc<dyn Fs>,
  59}
  60
  61impl Global for InlineAssistant {}
  62
  63impl InlineAssistant {
  64    pub fn new(fs: Arc<dyn Fs>, telemetry: Arc<Telemetry>) -> Self {
  65        Self {
  66            next_assist_id: InlineAssistId::default(),
  67            next_assist_group_id: InlineAssistGroupId::default(),
  68            assists: HashMap::default(),
  69            assists_by_editor: HashMap::default(),
  70            assist_groups: HashMap::default(),
  71            prompt_history: VecDeque::default(),
  72            telemetry: Some(telemetry),
  73            fs,
  74        }
  75    }
  76
  77    pub fn assist(
  78        &mut self,
  79        editor: &View<Editor>,
  80        workspace: Option<WeakView<Workspace>>,
  81        assistant_panel: Option<&View<AssistantPanel>>,
  82        cx: &mut WindowContext,
  83    ) {
  84        let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
  85
  86        let mut selections = Vec::<Selection<Point>>::new();
  87        let mut newest_selection = None;
  88        for mut selection in editor.read(cx).selections.all::<Point>(cx) {
  89            if selection.end > selection.start {
  90                selection.start.column = 0;
  91                // If the selection ends at the start of the line, we don't want to include it.
  92                if selection.end.column == 0 {
  93                    selection.end.row -= 1;
  94                }
  95                selection.end.column = snapshot.line_len(MultiBufferRow(selection.end.row));
  96            }
  97
  98            if let Some(prev_selection) = selections.last_mut() {
  99                if selection.start <= prev_selection.end {
 100                    prev_selection.end = selection.end;
 101                    continue;
 102                }
 103            }
 104
 105            let latest_selection = newest_selection.get_or_insert_with(|| selection.clone());
 106            if selection.id > latest_selection.id {
 107                *latest_selection = selection.clone();
 108            }
 109            selections.push(selection);
 110        }
 111        let newest_selection = newest_selection.unwrap();
 112
 113        let mut codegen_ranges = Vec::new();
 114        for (excerpt_id, buffer, buffer_range) in
 115            snapshot.excerpts_in_ranges(selections.iter().map(|selection| {
 116                snapshot.anchor_before(selection.start)..snapshot.anchor_after(selection.end)
 117            }))
 118        {
 119            let start = Anchor {
 120                buffer_id: Some(buffer.remote_id()),
 121                excerpt_id,
 122                text_anchor: buffer.anchor_before(buffer_range.start),
 123            };
 124            let end = Anchor {
 125                buffer_id: Some(buffer.remote_id()),
 126                excerpt_id,
 127                text_anchor: buffer.anchor_after(buffer_range.end),
 128            };
 129            codegen_ranges.push(start..end);
 130        }
 131
 132        let assist_group_id = self.next_assist_group_id.post_inc();
 133        let prompt_buffer = cx.new_model(|cx| Buffer::local("", cx));
 134        let prompt_buffer = cx.new_model(|cx| MultiBuffer::singleton(prompt_buffer, cx));
 135
 136        let mut assists = Vec::new();
 137        let mut assist_blocks = Vec::new();
 138        let mut assist_to_focus = None;
 139        for range in codegen_ranges {
 140            let assist_id = self.next_assist_id.post_inc();
 141            let codegen = cx.new_model(|cx| {
 142                Codegen::new(
 143                    editor.read(cx).buffer().clone(),
 144                    range.clone(),
 145                    self.telemetry.clone(),
 146                    cx,
 147                )
 148            });
 149
 150            let gutter_dimensions = Arc::new(Mutex::new(GutterDimensions::default()));
 151            let prompt_editor = cx.new_view(|cx| {
 152                PromptEditor::new(
 153                    assist_id,
 154                    gutter_dimensions.clone(),
 155                    self.prompt_history.clone(),
 156                    prompt_buffer.clone(),
 157                    codegen.clone(),
 158                    editor,
 159                    assistant_panel,
 160                    workspace.clone(),
 161                    self.fs.clone(),
 162                    cx,
 163                )
 164            });
 165
 166            if assist_to_focus.is_none() {
 167                let focus_assist = if newest_selection.reversed {
 168                    range.start.to_point(&snapshot) == newest_selection.start
 169                } else {
 170                    range.end.to_point(&snapshot) == newest_selection.end
 171                };
 172                if focus_assist {
 173                    assist_to_focus = Some(assist_id);
 174                }
 175            }
 176
 177            assist_blocks.push(BlockProperties {
 178                style: BlockStyle::Sticky,
 179                position: range.start,
 180                height: prompt_editor.read(cx).height_in_lines,
 181                render: build_assist_editor_renderer(&prompt_editor),
 182                disposition: BlockDisposition::Above,
 183            });
 184            assist_blocks.push(BlockProperties {
 185                style: BlockStyle::Sticky,
 186                position: range.end,
 187                height: 1,
 188                render: Box::new(|cx| {
 189                    v_flex()
 190                        .h_full()
 191                        .w_full()
 192                        .border_t_1()
 193                        .border_color(cx.theme().status().info_border)
 194                        .into_any_element()
 195                }),
 196                disposition: BlockDisposition::Below,
 197            });
 198            assists.push((assist_id, prompt_editor));
 199        }
 200
 201        let assist_block_ids = editor.update(cx, |editor, cx| {
 202            editor.insert_blocks(assist_blocks, None, cx)
 203        });
 204
 205        let editor_assists = self
 206            .assists_by_editor
 207            .entry(editor.downgrade())
 208            .or_insert_with(|| EditorInlineAssists::new(&editor, cx));
 209        let mut assist_group = InlineAssistGroup::new();
 210        for ((assist_id, prompt_editor), block_ids) in
 211            assists.into_iter().zip(assist_block_ids.chunks_exact(2))
 212        {
 213            self.assists.insert(
 214                assist_id,
 215                InlineAssist::new(
 216                    assist_id,
 217                    assist_group_id,
 218                    assistant_panel.is_some(),
 219                    editor,
 220                    &prompt_editor,
 221                    block_ids[0],
 222                    block_ids[1],
 223                    prompt_editor.read(cx).codegen.clone(),
 224                    workspace.clone(),
 225                    cx,
 226                ),
 227            );
 228            assist_group.assist_ids.push(assist_id);
 229            editor_assists.assist_ids.push(assist_id);
 230        }
 231        self.assist_groups.insert(assist_group_id, assist_group);
 232
 233        if let Some(assist_id) = assist_to_focus {
 234            self.focus_assist(assist_id, cx);
 235        }
 236    }
 237
 238    fn handle_prompt_editor_focus_in(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
 239        let assist = &self.assists[&assist_id];
 240        let Some(decorations) = assist.decorations.as_ref() else {
 241            return;
 242        };
 243        let assist_group = self.assist_groups.get_mut(&assist.group_id).unwrap();
 244        let editor_assists = self.assists_by_editor.get_mut(&assist.editor).unwrap();
 245
 246        assist_group.active_assist_id = Some(assist_id);
 247        if assist_group.linked {
 248            for assist_id in &assist_group.assist_ids {
 249                if let Some(decorations) = self.assists[assist_id].decorations.as_ref() {
 250                    decorations.prompt_editor.update(cx, |prompt_editor, cx| {
 251                        prompt_editor.set_show_cursor_when_unfocused(true, cx)
 252                    });
 253                }
 254            }
 255        }
 256
 257        assist
 258            .editor
 259            .update(cx, |editor, cx| {
 260                let scroll_top = editor.scroll_position(cx).y;
 261                let scroll_bottom = scroll_top + editor.visible_line_count().unwrap_or(0.);
 262                let prompt_row = editor
 263                    .row_for_block(decorations.prompt_block_id, cx)
 264                    .unwrap()
 265                    .0 as f32;
 266
 267                if (scroll_top..scroll_bottom).contains(&prompt_row) {
 268                    editor_assists.scroll_lock = Some(InlineAssistScrollLock {
 269                        assist_id,
 270                        distance_from_top: prompt_row - scroll_top,
 271                    });
 272                } else {
 273                    editor_assists.scroll_lock = None;
 274                }
 275            })
 276            .ok();
 277    }
 278
 279    fn handle_prompt_editor_focus_out(
 280        &mut self,
 281        assist_id: InlineAssistId,
 282        cx: &mut WindowContext,
 283    ) {
 284        let assist = &self.assists[&assist_id];
 285        let assist_group = self.assist_groups.get_mut(&assist.group_id).unwrap();
 286        if assist_group.active_assist_id == Some(assist_id) {
 287            assist_group.active_assist_id = None;
 288            if assist_group.linked {
 289                for assist_id in &assist_group.assist_ids {
 290                    if let Some(decorations) = self.assists[assist_id].decorations.as_ref() {
 291                        decorations.prompt_editor.update(cx, |prompt_editor, cx| {
 292                            prompt_editor.set_show_cursor_when_unfocused(false, cx)
 293                        });
 294                    }
 295                }
 296            }
 297        }
 298    }
 299
 300    fn handle_prompt_editor_event(
 301        &mut self,
 302        prompt_editor: View<PromptEditor>,
 303        event: &PromptEditorEvent,
 304        cx: &mut WindowContext,
 305    ) {
 306        let assist_id = prompt_editor.read(cx).id;
 307        match event {
 308            PromptEditorEvent::StartRequested => {
 309                self.start_assist(assist_id, cx);
 310            }
 311            PromptEditorEvent::StopRequested => {
 312                self.stop_assist(assist_id, cx);
 313            }
 314            PromptEditorEvent::ConfirmRequested => {
 315                self.finish_assist(assist_id, false, cx);
 316            }
 317            PromptEditorEvent::CancelRequested => {
 318                self.finish_assist(assist_id, true, cx);
 319            }
 320            PromptEditorEvent::DismissRequested => {
 321                self.dismiss_assist(assist_id, cx);
 322            }
 323            PromptEditorEvent::Resized { height_in_lines } => {
 324                self.resize_assist(assist_id, *height_in_lines, cx);
 325            }
 326        }
 327    }
 328
 329    fn handle_editor_newline(&mut self, editor: View<Editor>, cx: &mut WindowContext) {
 330        let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
 331            return;
 332        };
 333
 334        let editor = editor.read(cx);
 335        if editor.selections.count() == 1 {
 336            let selection = editor.selections.newest::<usize>(cx);
 337            let buffer = editor.buffer().read(cx).snapshot(cx);
 338            for assist_id in &editor_assists.assist_ids {
 339                let assist = &self.assists[assist_id];
 340                let assist_range = assist.codegen.read(cx).range.to_offset(&buffer);
 341                if assist_range.contains(&selection.start) && assist_range.contains(&selection.end)
 342                {
 343                    if matches!(assist.codegen.read(cx).status, CodegenStatus::Pending) {
 344                        self.dismiss_assist(*assist_id, cx);
 345                    } else {
 346                        self.finish_assist(*assist_id, false, cx);
 347                    }
 348
 349                    return;
 350                }
 351            }
 352        }
 353
 354        cx.propagate();
 355    }
 356
 357    fn handle_editor_cancel(&mut self, editor: View<Editor>, cx: &mut WindowContext) {
 358        let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
 359            return;
 360        };
 361
 362        let editor = editor.read(cx);
 363        if editor.selections.count() == 1 {
 364            let selection = editor.selections.newest::<usize>(cx);
 365            let buffer = editor.buffer().read(cx).snapshot(cx);
 366            for assist_id in &editor_assists.assist_ids {
 367                let assist = &self.assists[assist_id];
 368                let assist_range = assist.codegen.read(cx).range.to_offset(&buffer);
 369                if assist.decorations.is_some()
 370                    && assist_range.contains(&selection.start)
 371                    && assist_range.contains(&selection.end)
 372                {
 373                    self.focus_assist(*assist_id, cx);
 374                    return;
 375                }
 376            }
 377        }
 378
 379        cx.propagate();
 380    }
 381
 382    fn handle_editor_change(&mut self, editor: View<Editor>, cx: &mut WindowContext) {
 383        let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
 384            return;
 385        };
 386        let Some(scroll_lock) = editor_assists.scroll_lock.as_ref() else {
 387            return;
 388        };
 389        let assist = &self.assists[&scroll_lock.assist_id];
 390        let Some(decorations) = assist.decorations.as_ref() else {
 391            return;
 392        };
 393
 394        editor.update(cx, |editor, cx| {
 395            let scroll_position = editor.scroll_position(cx);
 396            let target_scroll_top = editor
 397                .row_for_block(decorations.prompt_block_id, cx)
 398                .unwrap()
 399                .0 as f32
 400                - scroll_lock.distance_from_top;
 401            if target_scroll_top != scroll_position.y {
 402                editor.set_scroll_position(point(scroll_position.x, target_scroll_top), cx);
 403            }
 404        });
 405    }
 406
 407    fn handle_editor_event(
 408        &mut self,
 409        editor: View<Editor>,
 410        event: &EditorEvent,
 411        cx: &mut WindowContext,
 412    ) {
 413        let Some(editor_assists) = self.assists_by_editor.get_mut(&editor.downgrade()) else {
 414            return;
 415        };
 416
 417        match event {
 418            EditorEvent::Saved => {
 419                for assist_id in editor_assists.assist_ids.clone() {
 420                    let assist = &self.assists[&assist_id];
 421                    if let CodegenStatus::Done = &assist.codegen.read(cx).status {
 422                        self.finish_assist(assist_id, false, cx)
 423                    }
 424                }
 425            }
 426            EditorEvent::Edited { transaction_id } => {
 427                let buffer = editor.read(cx).buffer().read(cx);
 428                let edited_ranges =
 429                    buffer.edited_ranges_for_transaction::<usize>(*transaction_id, cx);
 430                let snapshot = buffer.snapshot(cx);
 431
 432                for assist_id in editor_assists.assist_ids.clone() {
 433                    let assist = &self.assists[&assist_id];
 434                    if matches!(
 435                        assist.codegen.read(cx).status,
 436                        CodegenStatus::Error(_) | CodegenStatus::Done
 437                    ) {
 438                        let assist_range = assist.codegen.read(cx).range.to_offset(&snapshot);
 439                        if edited_ranges
 440                            .iter()
 441                            .any(|range| range.overlaps(&assist_range))
 442                        {
 443                            self.finish_assist(assist_id, false, cx);
 444                        }
 445                    }
 446                }
 447            }
 448            EditorEvent::ScrollPositionChanged { .. } => {
 449                if let Some(scroll_lock) = editor_assists.scroll_lock.as_ref() {
 450                    let assist = &self.assists[&scroll_lock.assist_id];
 451                    if let Some(decorations) = assist.decorations.as_ref() {
 452                        let distance_from_top = editor.update(cx, |editor, cx| {
 453                            let scroll_top = editor.scroll_position(cx).y;
 454                            let prompt_row = editor
 455                                .row_for_block(decorations.prompt_block_id, cx)
 456                                .unwrap()
 457                                .0 as f32;
 458                            prompt_row - scroll_top
 459                        });
 460
 461                        if distance_from_top != scroll_lock.distance_from_top {
 462                            editor_assists.scroll_lock = None;
 463                        }
 464                    }
 465                }
 466            }
 467            EditorEvent::SelectionsChanged { .. } => {
 468                for assist_id in editor_assists.assist_ids.clone() {
 469                    let assist = &self.assists[&assist_id];
 470                    if let Some(decorations) = assist.decorations.as_ref() {
 471                        if decorations.prompt_editor.focus_handle(cx).is_focused(cx) {
 472                            return;
 473                        }
 474                    }
 475                }
 476
 477                editor_assists.scroll_lock = None;
 478            }
 479            _ => {}
 480        }
 481    }
 482
 483    fn finish_assist(&mut self, assist_id: InlineAssistId, undo: bool, cx: &mut WindowContext) {
 484        if let Some(assist) = self.assists.get(&assist_id) {
 485            let assist_group_id = assist.group_id;
 486            if self.assist_groups[&assist_group_id].linked {
 487                for assist_id in self.unlink_assist_group(assist_group_id, cx) {
 488                    self.finish_assist(assist_id, undo, cx);
 489                }
 490                return;
 491            }
 492        }
 493
 494        self.dismiss_assist(assist_id, cx);
 495
 496        if let Some(assist) = self.assists.remove(&assist_id) {
 497            if let hash_map::Entry::Occupied(mut entry) = self.assist_groups.entry(assist.group_id)
 498            {
 499                entry.get_mut().assist_ids.retain(|id| *id != assist_id);
 500                if entry.get().assist_ids.is_empty() {
 501                    entry.remove();
 502                }
 503            }
 504
 505            if let hash_map::Entry::Occupied(mut entry) =
 506                self.assists_by_editor.entry(assist.editor.clone())
 507            {
 508                entry.get_mut().assist_ids.retain(|id| *id != assist_id);
 509                if entry.get().assist_ids.is_empty() {
 510                    entry.remove();
 511                    if let Some(editor) = assist.editor.upgrade() {
 512                        self.update_editor_highlights(&editor, cx);
 513                    }
 514                } else {
 515                    entry.get().highlight_updates.send(()).ok();
 516                }
 517            }
 518
 519            if undo {
 520                assist.codegen.update(cx, |codegen, cx| codegen.undo(cx));
 521            }
 522        }
 523    }
 524
 525    fn dismiss_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) -> bool {
 526        let Some(assist) = self.assists.get_mut(&assist_id) else {
 527            return false;
 528        };
 529        let Some(editor) = assist.editor.upgrade() else {
 530            return false;
 531        };
 532        let Some(decorations) = assist.decorations.take() else {
 533            return false;
 534        };
 535
 536        editor.update(cx, |editor, cx| {
 537            let mut to_remove = decorations.removed_line_block_ids;
 538            to_remove.insert(decorations.prompt_block_id);
 539            to_remove.insert(decorations.end_block_id);
 540            editor.remove_blocks(to_remove, None, cx);
 541        });
 542
 543        if decorations
 544            .prompt_editor
 545            .focus_handle(cx)
 546            .contains_focused(cx)
 547        {
 548            self.focus_next_assist(assist_id, cx);
 549        }
 550
 551        if let Some(editor_assists) = self.assists_by_editor.get_mut(&editor.downgrade()) {
 552            if editor_assists
 553                .scroll_lock
 554                .as_ref()
 555                .map_or(false, |lock| lock.assist_id == assist_id)
 556            {
 557                editor_assists.scroll_lock = None;
 558            }
 559            editor_assists.highlight_updates.send(()).ok();
 560        }
 561
 562        true
 563    }
 564
 565    fn focus_next_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
 566        let Some(assist) = self.assists.get(&assist_id) else {
 567            return;
 568        };
 569
 570        let assist_group = &self.assist_groups[&assist.group_id];
 571        let assist_ix = assist_group
 572            .assist_ids
 573            .iter()
 574            .position(|id| *id == assist_id)
 575            .unwrap();
 576        let assist_ids = assist_group
 577            .assist_ids
 578            .iter()
 579            .skip(assist_ix + 1)
 580            .chain(assist_group.assist_ids.iter().take(assist_ix));
 581
 582        for assist_id in assist_ids {
 583            let assist = &self.assists[assist_id];
 584            if assist.decorations.is_some() {
 585                self.focus_assist(*assist_id, cx);
 586                return;
 587            }
 588        }
 589
 590        assist.editor.update(cx, |editor, cx| editor.focus(cx)).ok();
 591    }
 592
 593    fn focus_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
 594        let assist = &self.assists[&assist_id];
 595        let Some(editor) = assist.editor.upgrade() else {
 596            return;
 597        };
 598
 599        if let Some(decorations) = assist.decorations.as_ref() {
 600            decorations.prompt_editor.update(cx, |prompt_editor, cx| {
 601                prompt_editor.editor.update(cx, |editor, cx| {
 602                    editor.focus(cx);
 603                    editor.select_all(&SelectAll, cx);
 604                })
 605            });
 606        }
 607
 608        let position = assist.codegen.read(cx).range.start;
 609        editor.update(cx, |editor, cx| {
 610            editor.change_selections(None, cx, |selections| {
 611                selections.select_anchor_ranges([position..position])
 612            });
 613
 614            let mut scroll_target_top;
 615            let mut scroll_target_bottom;
 616            if let Some(decorations) = assist.decorations.as_ref() {
 617                scroll_target_top = editor
 618                    .row_for_block(decorations.prompt_block_id, cx)
 619                    .unwrap()
 620                    .0 as f32;
 621                scroll_target_bottom = editor
 622                    .row_for_block(decorations.end_block_id, cx)
 623                    .unwrap()
 624                    .0 as f32;
 625            } else {
 626                let snapshot = editor.snapshot(cx);
 627                let codegen = assist.codegen.read(cx);
 628                let start_row = codegen
 629                    .range
 630                    .start
 631                    .to_display_point(&snapshot.display_snapshot)
 632                    .row();
 633                scroll_target_top = start_row.0 as f32;
 634                scroll_target_bottom = scroll_target_top + 1.;
 635            }
 636            scroll_target_top -= editor.vertical_scroll_margin() as f32;
 637            scroll_target_bottom += editor.vertical_scroll_margin() as f32;
 638
 639            let height_in_lines = editor.visible_line_count().unwrap_or(0.);
 640            let scroll_top = editor.scroll_position(cx).y;
 641            let scroll_bottom = scroll_top + height_in_lines;
 642
 643            if scroll_target_top < scroll_top {
 644                editor.set_scroll_position(point(0., scroll_target_top), cx);
 645            } else if scroll_target_bottom > scroll_bottom {
 646                if (scroll_target_bottom - scroll_target_top) <= height_in_lines {
 647                    editor
 648                        .set_scroll_position(point(0., scroll_target_bottom - height_in_lines), cx);
 649                } else {
 650                    editor.set_scroll_position(point(0., scroll_target_top), cx);
 651                }
 652            }
 653        });
 654    }
 655
 656    fn resize_assist(
 657        &mut self,
 658        assist_id: InlineAssistId,
 659        height_in_lines: u8,
 660        cx: &mut WindowContext,
 661    ) {
 662        if let Some(assist) = self.assists.get_mut(&assist_id) {
 663            if let Some(editor) = assist.editor.upgrade() {
 664                if let Some(decorations) = assist.decorations.as_ref() {
 665                    let mut new_blocks = HashMap::default();
 666                    new_blocks.insert(
 667                        decorations.prompt_block_id,
 668                        (
 669                            Some(height_in_lines),
 670                            build_assist_editor_renderer(&decorations.prompt_editor),
 671                        ),
 672                    );
 673                    editor.update(cx, |editor, cx| {
 674                        editor
 675                            .display_map
 676                            .update(cx, |map, cx| map.replace_blocks(new_blocks, cx))
 677                    });
 678                }
 679            }
 680        }
 681    }
 682
 683    fn unlink_assist_group(
 684        &mut self,
 685        assist_group_id: InlineAssistGroupId,
 686        cx: &mut WindowContext,
 687    ) -> Vec<InlineAssistId> {
 688        let assist_group = self.assist_groups.get_mut(&assist_group_id).unwrap();
 689        assist_group.linked = false;
 690        for assist_id in &assist_group.assist_ids {
 691            let assist = self.assists.get_mut(assist_id).unwrap();
 692            if let Some(editor_decorations) = assist.decorations.as_ref() {
 693                editor_decorations
 694                    .prompt_editor
 695                    .update(cx, |prompt_editor, cx| prompt_editor.unlink(cx));
 696            }
 697        }
 698        assist_group.assist_ids.clone()
 699    }
 700
 701    fn start_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
 702        let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
 703            assist
 704        } else {
 705            return;
 706        };
 707
 708        let assist_group_id = assist.group_id;
 709        if self.assist_groups[&assist_group_id].linked {
 710            for assist_id in self.unlink_assist_group(assist_group_id, cx) {
 711                self.start_assist(assist_id, cx);
 712            }
 713            return;
 714        }
 715
 716        let Some(user_prompt) = assist
 717            .decorations
 718            .as_ref()
 719            .map(|decorations| decorations.prompt_editor.read(cx).prompt(cx))
 720        else {
 721            return;
 722        };
 723
 724        self.prompt_history.retain(|prompt| *prompt != user_prompt);
 725        self.prompt_history.push_back(user_prompt.clone());
 726        if self.prompt_history.len() > PROMPT_HISTORY_MAX_LEN {
 727            self.prompt_history.pop_front();
 728        }
 729
 730        assist.codegen.update(cx, |codegen, cx| codegen.undo(cx));
 731        let codegen = assist.codegen.clone();
 732        let request = self.request_for_inline_assist(assist_id, cx);
 733
 734        cx.spawn(|mut cx| async move {
 735            let request = request.await?;
 736            codegen.update(&mut cx, |codegen, cx| codegen.start(request, cx))?;
 737            anyhow::Ok(())
 738        })
 739        .detach_and_log_err(cx);
 740    }
 741
 742    fn request_for_inline_assist(
 743        &self,
 744        assist_id: InlineAssistId,
 745        cx: &mut WindowContext,
 746    ) -> Task<Result<LanguageModelRequest>> {
 747        cx.spawn(|mut cx| async move {
 748            let (user_prompt, context_request, project_name, buffer, range, model) = cx
 749                .read_global(|this: &InlineAssistant, cx: &WindowContext| {
 750                    let assist = this.assists.get(&assist_id).context("invalid assist")?;
 751                    let decorations = assist.decorations.as_ref().context("invalid assist")?;
 752                    let editor = assist.editor.upgrade().context("invalid assist")?;
 753                    let user_prompt = decorations.prompt_editor.read(cx).prompt(cx);
 754                    let context_request = if assist.include_context {
 755                        assist.workspace.as_ref().and_then(|workspace| {
 756                            let workspace = workspace.upgrade()?.read(cx);
 757                            let assistant_panel = workspace.panel::<AssistantPanel>(cx)?;
 758                            Some(
 759                                assistant_panel
 760                                    .read(cx)
 761                                    .active_context(cx)?
 762                                    .read(cx)
 763                                    .to_completion_request(cx),
 764                            )
 765                        })
 766                    } else {
 767                        None
 768                    };
 769                    let project_name = assist.workspace.as_ref().and_then(|workspace| {
 770                        let workspace = workspace.upgrade()?;
 771                        Some(
 772                            workspace
 773                                .read(cx)
 774                                .project()
 775                                .read(cx)
 776                                .worktree_root_names(cx)
 777                                .collect::<Vec<&str>>()
 778                                .join("/"),
 779                        )
 780                    });
 781                    let buffer = editor.read(cx).buffer().read(cx).snapshot(cx);
 782                    let range = assist.codegen.read(cx).range.clone();
 783                    let model = CompletionProvider::global(cx).model();
 784                    anyhow::Ok((
 785                        user_prompt,
 786                        context_request,
 787                        project_name,
 788                        buffer,
 789                        range,
 790                        model,
 791                    ))
 792                })??;
 793
 794            let language = buffer.language_at(range.start);
 795            let language_name = if let Some(language) = language.as_ref() {
 796                if Arc::ptr_eq(language, &language::PLAIN_TEXT) {
 797                    None
 798                } else {
 799                    Some(language.name())
 800                }
 801            } else {
 802                None
 803            };
 804
 805            // Higher Temperature increases the randomness of model outputs.
 806            // If Markdown or No Language is Known, increase the randomness for more creative output
 807            // If Code, decrease temperature to get more deterministic outputs
 808            let temperature = if let Some(language) = language_name.clone() {
 809                if language.as_ref() == "Markdown" {
 810                    1.0
 811                } else {
 812                    0.5
 813                }
 814            } else {
 815                1.0
 816            };
 817
 818            let prompt = cx
 819                .background_executor()
 820                .spawn(async move {
 821                    let language_name = language_name.as_deref();
 822                    let start = buffer.point_to_buffer_offset(range.start);
 823                    let end = buffer.point_to_buffer_offset(range.end);
 824                    let (buffer, range) = if let Some((start, end)) = start.zip(end) {
 825                        let (start_buffer, start_buffer_offset) = start;
 826                        let (end_buffer, end_buffer_offset) = end;
 827                        if start_buffer.remote_id() == end_buffer.remote_id() {
 828                            (start_buffer.clone(), start_buffer_offset..end_buffer_offset)
 829                        } else {
 830                            return Err(anyhow!("invalid transformation range"));
 831                        }
 832                    } else {
 833                        return Err(anyhow!("invalid transformation range"));
 834                    };
 835                    generate_content_prompt(user_prompt, language_name, buffer, range, project_name)
 836                })
 837                .await?;
 838
 839            let mut messages = Vec::new();
 840            if let Some(context_request) = context_request {
 841                messages = context_request.messages;
 842            }
 843
 844            messages.push(LanguageModelRequestMessage {
 845                role: Role::User,
 846                content: prompt,
 847            });
 848
 849            Ok(LanguageModelRequest {
 850                model,
 851                messages,
 852                stop: vec!["|END|>".to_string()],
 853                temperature,
 854            })
 855        })
 856    }
 857
 858    fn stop_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
 859        let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
 860            assist
 861        } else {
 862            return;
 863        };
 864
 865        assist.codegen.update(cx, |codegen, cx| codegen.stop(cx));
 866    }
 867
 868    fn update_editor_highlights(&self, editor: &View<Editor>, cx: &mut WindowContext) {
 869        let mut gutter_pending_ranges = Vec::new();
 870        let mut gutter_transformed_ranges = Vec::new();
 871        let mut foreground_ranges = Vec::new();
 872        let mut inserted_row_ranges = Vec::new();
 873        let empty_assist_ids = Vec::new();
 874        let assist_ids = self
 875            .assists_by_editor
 876            .get(&editor.downgrade())
 877            .map_or(&empty_assist_ids, |editor_assists| {
 878                &editor_assists.assist_ids
 879            });
 880
 881        for assist_id in assist_ids {
 882            if let Some(assist) = self.assists.get(assist_id) {
 883                let codegen = assist.codegen.read(cx);
 884                foreground_ranges.extend(codegen.last_equal_ranges().iter().cloned());
 885
 886                if codegen.edit_position != codegen.range.end {
 887                    gutter_pending_ranges.push(codegen.edit_position..codegen.range.end);
 888                }
 889
 890                if codegen.range.start != codegen.edit_position {
 891                    gutter_transformed_ranges.push(codegen.range.start..codegen.edit_position);
 892                }
 893
 894                if assist.decorations.is_some() {
 895                    inserted_row_ranges.extend(codegen.diff.inserted_row_ranges.iter().cloned());
 896                }
 897            }
 898        }
 899
 900        let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
 901        merge_ranges(&mut foreground_ranges, &snapshot);
 902        merge_ranges(&mut gutter_pending_ranges, &snapshot);
 903        merge_ranges(&mut gutter_transformed_ranges, &snapshot);
 904        editor.update(cx, |editor, cx| {
 905            enum GutterPendingRange {}
 906            if gutter_pending_ranges.is_empty() {
 907                editor.clear_gutter_highlights::<GutterPendingRange>(cx);
 908            } else {
 909                editor.highlight_gutter::<GutterPendingRange>(
 910                    &gutter_pending_ranges,
 911                    |cx| cx.theme().status().info_background,
 912                    cx,
 913                )
 914            }
 915
 916            enum GutterTransformedRange {}
 917            if gutter_transformed_ranges.is_empty() {
 918                editor.clear_gutter_highlights::<GutterTransformedRange>(cx);
 919            } else {
 920                editor.highlight_gutter::<GutterTransformedRange>(
 921                    &gutter_transformed_ranges,
 922                    |cx| cx.theme().status().info,
 923                    cx,
 924                )
 925            }
 926
 927            if foreground_ranges.is_empty() {
 928                editor.clear_highlights::<InlineAssist>(cx);
 929            } else {
 930                editor.highlight_text::<InlineAssist>(
 931                    foreground_ranges,
 932                    HighlightStyle {
 933                        fade_out: Some(0.6),
 934                        ..Default::default()
 935                    },
 936                    cx,
 937                );
 938            }
 939
 940            editor.clear_row_highlights::<InlineAssist>();
 941            for row_range in inserted_row_ranges {
 942                editor.highlight_rows::<InlineAssist>(
 943                    row_range,
 944                    Some(cx.theme().status().info_background),
 945                    false,
 946                    cx,
 947                );
 948            }
 949        });
 950    }
 951
 952    fn update_editor_blocks(
 953        &mut self,
 954        editor: &View<Editor>,
 955        assist_id: InlineAssistId,
 956        cx: &mut WindowContext,
 957    ) {
 958        let Some(assist) = self.assists.get_mut(&assist_id) else {
 959            return;
 960        };
 961        let Some(decorations) = assist.decorations.as_mut() else {
 962            return;
 963        };
 964
 965        let codegen = assist.codegen.read(cx);
 966        let old_snapshot = codegen.snapshot.clone();
 967        let old_buffer = codegen.old_buffer.clone();
 968        let deleted_row_ranges = codegen.diff.deleted_row_ranges.clone();
 969
 970        editor.update(cx, |editor, cx| {
 971            let old_blocks = mem::take(&mut decorations.removed_line_block_ids);
 972            editor.remove_blocks(old_blocks, None, cx);
 973
 974            let mut new_blocks = Vec::new();
 975            for (new_row, old_row_range) in deleted_row_ranges {
 976                let (_, buffer_start) = old_snapshot
 977                    .point_to_buffer_offset(Point::new(*old_row_range.start(), 0))
 978                    .unwrap();
 979                let (_, buffer_end) = old_snapshot
 980                    .point_to_buffer_offset(Point::new(
 981                        *old_row_range.end(),
 982                        old_snapshot.line_len(MultiBufferRow(*old_row_range.end())),
 983                    ))
 984                    .unwrap();
 985
 986                let deleted_lines_editor = cx.new_view(|cx| {
 987                    let multi_buffer = cx.new_model(|_| {
 988                        MultiBuffer::without_headers(0, language::Capability::ReadOnly)
 989                    });
 990                    multi_buffer.update(cx, |multi_buffer, cx| {
 991                        multi_buffer.push_excerpts(
 992                            old_buffer.clone(),
 993                            Some(ExcerptRange {
 994                                context: buffer_start..buffer_end,
 995                                primary: None,
 996                            }),
 997                            cx,
 998                        );
 999                    });
1000
1001                    enum DeletedLines {}
1002                    let mut editor = Editor::for_multibuffer(multi_buffer, None, true, cx);
1003                    editor.set_soft_wrap_mode(language::language_settings::SoftWrap::None, cx);
1004                    editor.set_show_wrap_guides(false, cx);
1005                    editor.set_show_gutter(false, cx);
1006                    editor.scroll_manager.set_forbid_vertical_scroll(true);
1007                    editor.set_read_only(true);
1008                    editor.highlight_rows::<DeletedLines>(
1009                        Anchor::min()..=Anchor::max(),
1010                        Some(cx.theme().status().deleted_background),
1011                        false,
1012                        cx,
1013                    );
1014                    editor
1015                });
1016
1017                let height = deleted_lines_editor
1018                    .update(cx, |editor, cx| editor.max_point(cx).row().0 as u8 + 1);
1019                new_blocks.push(BlockProperties {
1020                    position: new_row,
1021                    height,
1022                    style: BlockStyle::Flex,
1023                    render: Box::new(move |cx| {
1024                        div()
1025                            .bg(cx.theme().status().deleted_background)
1026                            .size_full()
1027                            .pl(cx.gutter_dimensions.full_width())
1028                            .child(deleted_lines_editor.clone())
1029                            .into_any_element()
1030                    }),
1031                    disposition: BlockDisposition::Above,
1032                });
1033            }
1034
1035            decorations.removed_line_block_ids = editor
1036                .insert_blocks(new_blocks, None, cx)
1037                .into_iter()
1038                .collect();
1039        })
1040    }
1041}
1042
1043struct EditorInlineAssists {
1044    assist_ids: Vec<InlineAssistId>,
1045    scroll_lock: Option<InlineAssistScrollLock>,
1046    highlight_updates: async_watch::Sender<()>,
1047    _update_highlights: Task<Result<()>>,
1048    _subscriptions: Vec<gpui::Subscription>,
1049}
1050
1051struct InlineAssistScrollLock {
1052    assist_id: InlineAssistId,
1053    distance_from_top: f32,
1054}
1055
1056impl EditorInlineAssists {
1057    #[allow(clippy::too_many_arguments)]
1058    fn new(editor: &View<Editor>, cx: &mut WindowContext) -> Self {
1059        let (highlight_updates_tx, mut highlight_updates_rx) = async_watch::channel(());
1060        Self {
1061            assist_ids: Vec::new(),
1062            scroll_lock: None,
1063            highlight_updates: highlight_updates_tx,
1064            _update_highlights: cx.spawn(|mut cx| {
1065                let editor = editor.downgrade();
1066                async move {
1067                    while let Ok(()) = highlight_updates_rx.changed().await {
1068                        let editor = editor.upgrade().context("editor was dropped")?;
1069                        cx.update_global(|assistant: &mut InlineAssistant, cx| {
1070                            assistant.update_editor_highlights(&editor, cx);
1071                        })?;
1072                    }
1073                    Ok(())
1074                }
1075            }),
1076            _subscriptions: vec![
1077                cx.observe(editor, move |editor, cx| {
1078                    InlineAssistant::update_global(cx, |this, cx| {
1079                        this.handle_editor_change(editor, cx)
1080                    })
1081                }),
1082                cx.subscribe(editor, move |editor, event, cx| {
1083                    InlineAssistant::update_global(cx, |this, cx| {
1084                        this.handle_editor_event(editor, event, cx)
1085                    })
1086                }),
1087                editor.update(cx, |editor, cx| {
1088                    let editor_handle = cx.view().downgrade();
1089                    editor.register_action(
1090                        move |_: &editor::actions::Newline, cx: &mut WindowContext| {
1091                            InlineAssistant::update_global(cx, |this, cx| {
1092                                if let Some(editor) = editor_handle.upgrade() {
1093                                    this.handle_editor_newline(editor, cx)
1094                                }
1095                            })
1096                        },
1097                    )
1098                }),
1099                editor.update(cx, |editor, cx| {
1100                    let editor_handle = cx.view().downgrade();
1101                    editor.register_action(
1102                        move |_: &editor::actions::Cancel, cx: &mut WindowContext| {
1103                            InlineAssistant::update_global(cx, |this, cx| {
1104                                if let Some(editor) = editor_handle.upgrade() {
1105                                    this.handle_editor_cancel(editor, cx)
1106                                }
1107                            })
1108                        },
1109                    )
1110                }),
1111            ],
1112        }
1113    }
1114}
1115
1116struct InlineAssistGroup {
1117    assist_ids: Vec<InlineAssistId>,
1118    linked: bool,
1119    active_assist_id: Option<InlineAssistId>,
1120}
1121
1122impl InlineAssistGroup {
1123    fn new() -> Self {
1124        Self {
1125            assist_ids: Vec::new(),
1126            linked: true,
1127            active_assist_id: None,
1128        }
1129    }
1130}
1131
1132fn build_assist_editor_renderer(editor: &View<PromptEditor>) -> RenderBlock {
1133    let editor = editor.clone();
1134    Box::new(move |cx: &mut BlockContext| {
1135        *editor.read(cx).gutter_dimensions.lock() = *cx.gutter_dimensions;
1136        editor.clone().into_any_element()
1137    })
1138}
1139
1140#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
1141struct InlineAssistId(usize);
1142
1143impl InlineAssistId {
1144    fn post_inc(&mut self) -> InlineAssistId {
1145        let id = *self;
1146        self.0 += 1;
1147        id
1148    }
1149}
1150
1151#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
1152struct InlineAssistGroupId(usize);
1153
1154impl InlineAssistGroupId {
1155    fn post_inc(&mut self) -> InlineAssistGroupId {
1156        let id = *self;
1157        self.0 += 1;
1158        id
1159    }
1160}
1161
1162enum PromptEditorEvent {
1163    StartRequested,
1164    StopRequested,
1165    ConfirmRequested,
1166    CancelRequested,
1167    DismissRequested,
1168    Resized { height_in_lines: u8 },
1169}
1170
1171struct PromptEditor {
1172    id: InlineAssistId,
1173    fs: Arc<dyn Fs>,
1174    height_in_lines: u8,
1175    editor: View<Editor>,
1176    edited_since_done: bool,
1177    gutter_dimensions: Arc<Mutex<GutterDimensions>>,
1178    prompt_history: VecDeque<String>,
1179    prompt_history_ix: Option<usize>,
1180    pending_prompt: String,
1181    codegen: Model<Codegen>,
1182    _codegen_subscription: Subscription,
1183    editor_subscriptions: Vec<Subscription>,
1184    pending_token_count: Task<Result<()>>,
1185    token_count: Option<usize>,
1186    _token_count_subscriptions: Vec<Subscription>,
1187    workspace: Option<WeakView<Workspace>>,
1188}
1189
1190impl EventEmitter<PromptEditorEvent> for PromptEditor {}
1191
1192impl Render for PromptEditor {
1193    fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
1194        let gutter_dimensions = *self.gutter_dimensions.lock();
1195        let fs = self.fs.clone();
1196
1197        let buttons = match &self.codegen.read(cx).status {
1198            CodegenStatus::Idle => {
1199                vec![
1200                    IconButton::new("cancel", IconName::Close)
1201                        .icon_color(Color::Muted)
1202                        .size(ButtonSize::None)
1203                        .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
1204                        .on_click(
1205                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
1206                        ),
1207                    IconButton::new("start", IconName::Sparkle)
1208                        .icon_color(Color::Muted)
1209                        .size(ButtonSize::None)
1210                        .icon_size(IconSize::XSmall)
1211                        .tooltip(|cx| Tooltip::for_action("Transform", &menu::Confirm, cx))
1212                        .on_click(
1213                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StartRequested)),
1214                        ),
1215                ]
1216            }
1217            CodegenStatus::Pending => {
1218                vec![
1219                    IconButton::new("cancel", IconName::Close)
1220                        .icon_color(Color::Muted)
1221                        .size(ButtonSize::None)
1222                        .tooltip(|cx| Tooltip::text("Cancel Assist", cx))
1223                        .on_click(
1224                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
1225                        ),
1226                    IconButton::new("stop", IconName::Stop)
1227                        .icon_color(Color::Error)
1228                        .size(ButtonSize::None)
1229                        .icon_size(IconSize::XSmall)
1230                        .tooltip(|cx| {
1231                            Tooltip::with_meta(
1232                                "Interrupt Transformation",
1233                                Some(&menu::Cancel),
1234                                "Changes won't be discarded",
1235                                cx,
1236                            )
1237                        })
1238                        .on_click(
1239                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StopRequested)),
1240                        ),
1241                ]
1242            }
1243            CodegenStatus::Error(_) | CodegenStatus::Done => {
1244                vec![
1245                    IconButton::new("cancel", IconName::Close)
1246                        .icon_color(Color::Muted)
1247                        .size(ButtonSize::None)
1248                        .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
1249                        .on_click(
1250                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
1251                        ),
1252                    if self.edited_since_done {
1253                        IconButton::new("restart", IconName::RotateCw)
1254                            .icon_color(Color::Info)
1255                            .icon_size(IconSize::XSmall)
1256                            .size(ButtonSize::None)
1257                            .tooltip(|cx| {
1258                                Tooltip::with_meta(
1259                                    "Restart Transformation",
1260                                    Some(&menu::Confirm),
1261                                    "Changes will be discarded",
1262                                    cx,
1263                                )
1264                            })
1265                            .on_click(cx.listener(|_, _, cx| {
1266                                cx.emit(PromptEditorEvent::StartRequested);
1267                            }))
1268                    } else {
1269                        IconButton::new("confirm", IconName::Check)
1270                            .icon_color(Color::Info)
1271                            .size(ButtonSize::None)
1272                            .tooltip(|cx| Tooltip::for_action("Confirm Assist", &menu::Confirm, cx))
1273                            .on_click(cx.listener(|_, _, cx| {
1274                                cx.emit(PromptEditorEvent::ConfirmRequested);
1275                            }))
1276                    },
1277                ]
1278            }
1279        };
1280
1281        h_flex()
1282            .bg(cx.theme().colors().editor_background)
1283            .border_y_1()
1284            .border_color(cx.theme().status().info_border)
1285            .py_1p5()
1286            .h_full()
1287            .w_full()
1288            .on_action(cx.listener(Self::confirm))
1289            .on_action(cx.listener(Self::cancel))
1290            .on_action(cx.listener(Self::move_up))
1291            .on_action(cx.listener(Self::move_down))
1292            .child(
1293                h_flex()
1294                    .w(gutter_dimensions.full_width() + (gutter_dimensions.margin / 2.0))
1295                    .justify_center()
1296                    .gap_2()
1297                    .child(
1298                        PopoverMenu::new("model-switcher")
1299                            .menu(move |cx| {
1300                                ContextMenu::build(cx, |mut menu, cx| {
1301                                    for model in CompletionProvider::global(cx).available_models(cx)
1302                                    {
1303                                        menu = menu.custom_entry(
1304                                            {
1305                                                let model = model.clone();
1306                                                move |_| {
1307                                                    Label::new(model.display_name())
1308                                                        .into_any_element()
1309                                                }
1310                                            },
1311                                            {
1312                                                let fs = fs.clone();
1313                                                let model = model.clone();
1314                                                move |cx| {
1315                                                    let model = model.clone();
1316                                                    update_settings_file::<AssistantSettings>(
1317                                                        fs.clone(),
1318                                                        cx,
1319                                                        move |settings| settings.set_model(model),
1320                                                    );
1321                                                }
1322                                            },
1323                                        );
1324                                    }
1325                                    menu
1326                                })
1327                                .into()
1328                            })
1329                            .trigger(
1330                                IconButton::new("context", IconName::Settings)
1331                                    .size(ButtonSize::None)
1332                                    .icon_size(IconSize::Small)
1333                                    .icon_color(Color::Muted)
1334                                    .tooltip(move |cx| {
1335                                        Tooltip::with_meta(
1336                                            format!(
1337                                                "Using {}",
1338                                                CompletionProvider::global(cx)
1339                                                    .model()
1340                                                    .display_name()
1341                                            ),
1342                                            None,
1343                                            "Click to Change Model",
1344                                            cx,
1345                                        )
1346                                    }),
1347                            )
1348                            .anchor(gpui::AnchorCorner::BottomRight),
1349                    )
1350                    .children(
1351                        if let CodegenStatus::Error(error) = &self.codegen.read(cx).status {
1352                            let error_message = SharedString::from(error.to_string());
1353                            Some(
1354                                div()
1355                                    .id("error")
1356                                    .tooltip(move |cx| Tooltip::text(error_message.clone(), cx))
1357                                    .child(
1358                                        Icon::new(IconName::XCircle)
1359                                            .size(IconSize::Small)
1360                                            .color(Color::Error),
1361                                    ),
1362                            )
1363                        } else {
1364                            None
1365                        },
1366                    ),
1367            )
1368            .child(div().flex_1().child(self.render_prompt_editor(cx)))
1369            .child(
1370                h_flex()
1371                    .gap_2()
1372                    .pr_4()
1373                    .children(self.render_token_count(cx))
1374                    .children(buttons),
1375            )
1376    }
1377}
1378
1379impl FocusableView for PromptEditor {
1380    fn focus_handle(&self, cx: &AppContext) -> FocusHandle {
1381        self.editor.focus_handle(cx)
1382    }
1383}
1384
1385impl PromptEditor {
1386    const MAX_LINES: u8 = 8;
1387
1388    #[allow(clippy::too_many_arguments)]
1389    fn new(
1390        id: InlineAssistId,
1391        gutter_dimensions: Arc<Mutex<GutterDimensions>>,
1392        prompt_history: VecDeque<String>,
1393        prompt_buffer: Model<MultiBuffer>,
1394        codegen: Model<Codegen>,
1395        parent_editor: &View<Editor>,
1396        assistant_panel: Option<&View<AssistantPanel>>,
1397        workspace: Option<WeakView<Workspace>>,
1398        fs: Arc<dyn Fs>,
1399        cx: &mut ViewContext<Self>,
1400    ) -> Self {
1401        let prompt_editor = cx.new_view(|cx| {
1402            let mut editor = Editor::new(
1403                EditorMode::AutoHeight {
1404                    max_lines: Self::MAX_LINES as usize,
1405                },
1406                prompt_buffer,
1407                None,
1408                false,
1409                cx,
1410            );
1411            editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
1412            // Since the prompt editors for all inline assistants are linked,
1413            // always show the cursor (even when it isn't focused) because
1414            // typing in one will make what you typed appear in all of them.
1415            editor.set_show_cursor_when_unfocused(true, cx);
1416            editor.set_placeholder_text("Add a prompt…", cx);
1417            editor
1418        });
1419
1420        let mut token_count_subscriptions = Vec::new();
1421        token_count_subscriptions
1422            .push(cx.subscribe(parent_editor, Self::handle_parent_editor_event));
1423        if let Some(assistant_panel) = assistant_panel {
1424            token_count_subscriptions
1425                .push(cx.subscribe(assistant_panel, Self::handle_assistant_panel_event));
1426        }
1427
1428        let mut this = Self {
1429            id,
1430            height_in_lines: 1,
1431            editor: prompt_editor,
1432            edited_since_done: false,
1433            gutter_dimensions,
1434            prompt_history,
1435            prompt_history_ix: None,
1436            pending_prompt: String::new(),
1437            _codegen_subscription: cx.observe(&codegen, Self::handle_codegen_changed),
1438            editor_subscriptions: Vec::new(),
1439            codegen,
1440            fs,
1441            pending_token_count: Task::ready(Ok(())),
1442            token_count: None,
1443            _token_count_subscriptions: token_count_subscriptions,
1444            workspace,
1445        };
1446        this.count_lines(cx);
1447        this.count_tokens(cx);
1448        this.subscribe_to_editor(cx);
1449        this
1450    }
1451
1452    fn subscribe_to_editor(&mut self, cx: &mut ViewContext<Self>) {
1453        self.editor_subscriptions.clear();
1454        self.editor_subscriptions
1455            .push(cx.observe(&self.editor, Self::handle_prompt_editor_changed));
1456        self.editor_subscriptions
1457            .push(cx.subscribe(&self.editor, Self::handle_prompt_editor_events));
1458    }
1459
1460    fn set_show_cursor_when_unfocused(
1461        &mut self,
1462        show_cursor_when_unfocused: bool,
1463        cx: &mut ViewContext<Self>,
1464    ) {
1465        self.editor.update(cx, |editor, cx| {
1466            editor.set_show_cursor_when_unfocused(show_cursor_when_unfocused, cx)
1467        });
1468    }
1469
1470    fn unlink(&mut self, cx: &mut ViewContext<Self>) {
1471        let prompt = self.prompt(cx);
1472        let focus = self.editor.focus_handle(cx).contains_focused(cx);
1473        self.editor = cx.new_view(|cx| {
1474            let mut editor = Editor::auto_height(Self::MAX_LINES as usize, cx);
1475            editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
1476            editor.set_placeholder_text("Add a prompt…", cx);
1477            editor.set_text(prompt, cx);
1478            if focus {
1479                editor.focus(cx);
1480            }
1481            editor
1482        });
1483        self.subscribe_to_editor(cx);
1484    }
1485
1486    fn prompt(&self, cx: &AppContext) -> String {
1487        self.editor.read(cx).text(cx)
1488    }
1489
1490    fn count_lines(&mut self, cx: &mut ViewContext<Self>) {
1491        let height_in_lines = cmp::max(
1492            2, // Make the editor at least two lines tall, to account for padding and buttons.
1493            cmp::min(
1494                self.editor
1495                    .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1),
1496                Self::MAX_LINES as u32,
1497            ),
1498        ) as u8;
1499
1500        if height_in_lines != self.height_in_lines {
1501            self.height_in_lines = height_in_lines;
1502            cx.emit(PromptEditorEvent::Resized { height_in_lines });
1503        }
1504    }
1505
1506    fn handle_parent_editor_event(
1507        &mut self,
1508        _: View<Editor>,
1509        event: &EditorEvent,
1510        cx: &mut ViewContext<Self>,
1511    ) {
1512        if let EditorEvent::BufferEdited { .. } = event {
1513            self.count_tokens(cx);
1514        }
1515    }
1516
1517    fn handle_assistant_panel_event(
1518        &mut self,
1519        _: View<AssistantPanel>,
1520        event: &AssistantPanelEvent,
1521        cx: &mut ViewContext<Self>,
1522    ) {
1523        let AssistantPanelEvent::ContextEdited { .. } = event;
1524        self.count_tokens(cx);
1525    }
1526
1527    fn count_tokens(&mut self, cx: &mut ViewContext<Self>) {
1528        let assist_id = self.id;
1529        self.pending_token_count = cx.spawn(|this, mut cx| async move {
1530            cx.background_executor().timer(Duration::from_secs(1)).await;
1531            let request = cx
1532                .update_global(|inline_assistant: &mut InlineAssistant, cx| {
1533                    inline_assistant.request_for_inline_assist(assist_id, cx)
1534                })?
1535                .await?;
1536
1537            let token_count = cx
1538                .update(|cx| CompletionProvider::global(cx).count_tokens(request, cx))?
1539                .await?;
1540            this.update(&mut cx, |this, cx| {
1541                this.token_count = Some(token_count);
1542                cx.notify();
1543            })
1544        })
1545    }
1546
1547    fn handle_prompt_editor_changed(&mut self, _: View<Editor>, cx: &mut ViewContext<Self>) {
1548        self.count_lines(cx);
1549    }
1550
1551    fn handle_prompt_editor_events(
1552        &mut self,
1553        _: View<Editor>,
1554        event: &EditorEvent,
1555        cx: &mut ViewContext<Self>,
1556    ) {
1557        match event {
1558            EditorEvent::Edited { .. } => {
1559                let prompt = self.editor.read(cx).text(cx);
1560                if self
1561                    .prompt_history_ix
1562                    .map_or(true, |ix| self.prompt_history[ix] != prompt)
1563                {
1564                    self.prompt_history_ix.take();
1565                    self.pending_prompt = prompt;
1566                }
1567
1568                self.edited_since_done = true;
1569                cx.notify();
1570            }
1571            EditorEvent::BufferEdited => {
1572                self.count_tokens(cx);
1573            }
1574            _ => {}
1575        }
1576    }
1577
1578    fn handle_codegen_changed(&mut self, _: Model<Codegen>, cx: &mut ViewContext<Self>) {
1579        match &self.codegen.read(cx).status {
1580            CodegenStatus::Idle => {
1581                self.editor
1582                    .update(cx, |editor, _| editor.set_read_only(false));
1583            }
1584            CodegenStatus::Pending => {
1585                self.editor
1586                    .update(cx, |editor, _| editor.set_read_only(true));
1587            }
1588            CodegenStatus::Done | CodegenStatus::Error(_) => {
1589                self.edited_since_done = false;
1590                self.editor
1591                    .update(cx, |editor, _| editor.set_read_only(false));
1592            }
1593        }
1594    }
1595
1596    fn cancel(&mut self, _: &editor::actions::Cancel, cx: &mut ViewContext<Self>) {
1597        match &self.codegen.read(cx).status {
1598            CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {
1599                cx.emit(PromptEditorEvent::CancelRequested);
1600            }
1601            CodegenStatus::Pending => {
1602                cx.emit(PromptEditorEvent::StopRequested);
1603            }
1604        }
1605    }
1606
1607    fn confirm(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
1608        match &self.codegen.read(cx).status {
1609            CodegenStatus::Idle => {
1610                cx.emit(PromptEditorEvent::StartRequested);
1611            }
1612            CodegenStatus::Pending => {
1613                cx.emit(PromptEditorEvent::DismissRequested);
1614            }
1615            CodegenStatus::Done | CodegenStatus::Error(_) => {
1616                if self.edited_since_done {
1617                    cx.emit(PromptEditorEvent::StartRequested);
1618                } else {
1619                    cx.emit(PromptEditorEvent::ConfirmRequested);
1620                }
1621            }
1622        }
1623    }
1624
1625    fn move_up(&mut self, _: &MoveUp, cx: &mut ViewContext<Self>) {
1626        if let Some(ix) = self.prompt_history_ix {
1627            if ix > 0 {
1628                self.prompt_history_ix = Some(ix - 1);
1629                let prompt = self.prompt_history[ix - 1].as_str();
1630                self.editor.update(cx, |editor, cx| {
1631                    editor.set_text(prompt, cx);
1632                    editor.move_to_beginning(&Default::default(), cx);
1633                });
1634            }
1635        } else if !self.prompt_history.is_empty() {
1636            self.prompt_history_ix = Some(self.prompt_history.len() - 1);
1637            let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();
1638            self.editor.update(cx, |editor, cx| {
1639                editor.set_text(prompt, cx);
1640                editor.move_to_beginning(&Default::default(), cx);
1641            });
1642        }
1643    }
1644
1645    fn move_down(&mut self, _: &MoveDown, cx: &mut ViewContext<Self>) {
1646        if let Some(ix) = self.prompt_history_ix {
1647            if ix < self.prompt_history.len() - 1 {
1648                self.prompt_history_ix = Some(ix + 1);
1649                let prompt = self.prompt_history[ix + 1].as_str();
1650                self.editor.update(cx, |editor, cx| {
1651                    editor.set_text(prompt, cx);
1652                    editor.move_to_end(&Default::default(), cx)
1653                });
1654            } else {
1655                self.prompt_history_ix = None;
1656                let prompt = self.pending_prompt.as_str();
1657                self.editor.update(cx, |editor, cx| {
1658                    editor.set_text(prompt, cx);
1659                    editor.move_to_end(&Default::default(), cx)
1660                });
1661            }
1662        }
1663    }
1664
1665    fn render_token_count(&self, cx: &mut ViewContext<Self>) -> Option<impl IntoElement> {
1666        let model = CompletionProvider::global(cx).model();
1667        let token_count = self.token_count?;
1668        let max_token_count = model.max_token_count();
1669
1670        let remaining_tokens = max_token_count as isize - token_count as isize;
1671        let token_count_color = if remaining_tokens <= 0 {
1672            Color::Error
1673        } else if token_count as f32 / max_token_count as f32 >= 0.8 {
1674            Color::Warning
1675        } else {
1676            Color::Muted
1677        };
1678
1679        let mut token_count = h_flex()
1680            .id("token_count")
1681            .gap_0p5()
1682            .child(
1683                Label::new(humanize_token_count(token_count))
1684                    .size(LabelSize::Small)
1685                    .color(token_count_color),
1686            )
1687            .child(Label::new("/").size(LabelSize::Small).color(Color::Muted))
1688            .child(
1689                Label::new(humanize_token_count(max_token_count))
1690                    .size(LabelSize::Small)
1691                    .color(Color::Muted),
1692            );
1693        if let Some(workspace) = self.workspace.clone() {
1694            token_count = token_count
1695                .tooltip(|cx| {
1696                    Tooltip::with_meta(
1697                        "Tokens Used by Inline Assistant",
1698                        None,
1699                        "Click to Open Assistant Panel",
1700                        cx,
1701                    )
1702                })
1703                .cursor_pointer()
1704                .on_mouse_down(gpui::MouseButton::Left, |_, cx| cx.stop_propagation())
1705                .on_click(move |_, cx| {
1706                    cx.stop_propagation();
1707                    workspace
1708                        .update(cx, |workspace, cx| {
1709                            workspace.focus_panel::<AssistantPanel>(cx)
1710                        })
1711                        .ok();
1712                });
1713        } else {
1714            token_count = token_count
1715                .cursor_default()
1716                .tooltip(|cx| Tooltip::text("Tokens Used by Inline Assistant", cx));
1717        }
1718
1719        Some(token_count)
1720    }
1721
1722    fn render_prompt_editor(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
1723        let settings = ThemeSettings::get_global(cx);
1724        let text_style = TextStyle {
1725            color: if self.editor.read(cx).read_only(cx) {
1726                cx.theme().colors().text_disabled
1727            } else {
1728                cx.theme().colors().text
1729            },
1730            font_family: settings.ui_font.family.clone(),
1731            font_features: settings.ui_font.features.clone(),
1732            font_size: rems(0.875).into(),
1733            font_weight: settings.ui_font.weight,
1734            font_style: FontStyle::Normal,
1735            line_height: relative(1.3),
1736            background_color: None,
1737            underline: None,
1738            strikethrough: None,
1739            white_space: WhiteSpace::Normal,
1740        };
1741        EditorElement::new(
1742            &self.editor,
1743            EditorStyle {
1744                background: cx.theme().colors().editor_background,
1745                local_player: cx.theme().players().local(),
1746                text: text_style,
1747                ..Default::default()
1748            },
1749        )
1750    }
1751}
1752
1753struct InlineAssist {
1754    group_id: InlineAssistGroupId,
1755    editor: WeakView<Editor>,
1756    decorations: Option<InlineAssistDecorations>,
1757    codegen: Model<Codegen>,
1758    _subscriptions: Vec<Subscription>,
1759    workspace: Option<WeakView<Workspace>>,
1760    include_context: bool,
1761}
1762
1763impl InlineAssist {
1764    #[allow(clippy::too_many_arguments)]
1765    fn new(
1766        assist_id: InlineAssistId,
1767        group_id: InlineAssistGroupId,
1768        include_context: bool,
1769        editor: &View<Editor>,
1770        prompt_editor: &View<PromptEditor>,
1771        prompt_block_id: BlockId,
1772        end_block_id: BlockId,
1773        codegen: Model<Codegen>,
1774        workspace: Option<WeakView<Workspace>>,
1775        cx: &mut WindowContext,
1776    ) -> Self {
1777        let prompt_editor_focus_handle = prompt_editor.focus_handle(cx);
1778        InlineAssist {
1779            group_id,
1780            include_context,
1781            editor: editor.downgrade(),
1782            decorations: Some(InlineAssistDecorations {
1783                prompt_block_id,
1784                prompt_editor: prompt_editor.clone(),
1785                removed_line_block_ids: HashSet::default(),
1786                end_block_id,
1787            }),
1788            codegen: codegen.clone(),
1789            workspace: workspace.clone(),
1790            _subscriptions: vec![
1791                cx.on_focus_in(&prompt_editor_focus_handle, move |cx| {
1792                    InlineAssistant::update_global(cx, |this, cx| {
1793                        this.handle_prompt_editor_focus_in(assist_id, cx)
1794                    })
1795                }),
1796                cx.on_focus_out(&prompt_editor_focus_handle, move |_, cx| {
1797                    InlineAssistant::update_global(cx, |this, cx| {
1798                        this.handle_prompt_editor_focus_out(assist_id, cx)
1799                    })
1800                }),
1801                cx.subscribe(prompt_editor, |prompt_editor, event, cx| {
1802                    InlineAssistant::update_global(cx, |this, cx| {
1803                        this.handle_prompt_editor_event(prompt_editor, event, cx)
1804                    })
1805                }),
1806                cx.observe(&codegen, {
1807                    let editor = editor.downgrade();
1808                    move |_, cx| {
1809                        if let Some(editor) = editor.upgrade() {
1810                            InlineAssistant::update_global(cx, |this, cx| {
1811                                if let Some(editor_assists) =
1812                                    this.assists_by_editor.get(&editor.downgrade())
1813                                {
1814                                    editor_assists.highlight_updates.send(()).ok();
1815                                }
1816
1817                                this.update_editor_blocks(&editor, assist_id, cx);
1818                            })
1819                        }
1820                    }
1821                }),
1822                cx.subscribe(&codegen, move |codegen, event, cx| {
1823                    InlineAssistant::update_global(cx, |this, cx| match event {
1824                        CodegenEvent::Undone => this.finish_assist(assist_id, false, cx),
1825                        CodegenEvent::Finished => {
1826                            let assist = if let Some(assist) = this.assists.get(&assist_id) {
1827                                assist
1828                            } else {
1829                                return;
1830                            };
1831
1832                            if let CodegenStatus::Error(error) = &codegen.read(cx).status {
1833                                if assist.decorations.is_none() {
1834                                    if let Some(workspace) = assist
1835                                        .workspace
1836                                        .as_ref()
1837                                        .and_then(|workspace| workspace.upgrade())
1838                                    {
1839                                        let error = format!("Inline assistant error: {}", error);
1840                                        workspace.update(cx, |workspace, cx| {
1841                                            struct InlineAssistantError;
1842
1843                                            let id =
1844                                                NotificationId::identified::<InlineAssistantError>(
1845                                                    assist_id.0,
1846                                                );
1847
1848                                            workspace.show_toast(Toast::new(id, error), cx);
1849                                        })
1850                                    }
1851                                }
1852                            }
1853
1854                            if assist.decorations.is_none() {
1855                                this.finish_assist(assist_id, false, cx);
1856                            }
1857                        }
1858                    })
1859                }),
1860            ],
1861        }
1862    }
1863}
1864
1865struct InlineAssistDecorations {
1866    prompt_block_id: BlockId,
1867    prompt_editor: View<PromptEditor>,
1868    removed_line_block_ids: HashSet<BlockId>,
1869    end_block_id: BlockId,
1870}
1871
1872#[derive(Debug)]
1873pub enum CodegenEvent {
1874    Finished,
1875    Undone,
1876}
1877
1878pub struct Codegen {
1879    buffer: Model<MultiBuffer>,
1880    old_buffer: Model<Buffer>,
1881    snapshot: MultiBufferSnapshot,
1882    range: Range<Anchor>,
1883    edit_position: Anchor,
1884    last_equal_ranges: Vec<Range<Anchor>>,
1885    transaction_id: Option<TransactionId>,
1886    status: CodegenStatus,
1887    generation: Task<()>,
1888    diff: Diff,
1889    telemetry: Option<Arc<Telemetry>>,
1890    _subscription: gpui::Subscription,
1891}
1892
1893enum CodegenStatus {
1894    Idle,
1895    Pending,
1896    Done,
1897    Error(anyhow::Error),
1898}
1899
1900#[derive(Default)]
1901struct Diff {
1902    task: Option<Task<()>>,
1903    should_update: bool,
1904    deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)>,
1905    inserted_row_ranges: Vec<RangeInclusive<Anchor>>,
1906}
1907
1908impl EventEmitter<CodegenEvent> for Codegen {}
1909
1910impl Codegen {
1911    pub fn new(
1912        buffer: Model<MultiBuffer>,
1913        range: Range<Anchor>,
1914        telemetry: Option<Arc<Telemetry>>,
1915        cx: &mut ModelContext<Self>,
1916    ) -> Self {
1917        let snapshot = buffer.read(cx).snapshot(cx);
1918
1919        let (old_buffer, _, _) = buffer
1920            .read(cx)
1921            .range_to_buffer_ranges(range.clone(), cx)
1922            .pop()
1923            .unwrap();
1924        let old_buffer = cx.new_model(|cx| {
1925            let old_buffer = old_buffer.read(cx);
1926            let text = old_buffer.as_rope().clone();
1927            let line_ending = old_buffer.line_ending();
1928            let language = old_buffer.language().cloned();
1929            let language_registry = old_buffer.language_registry();
1930
1931            let mut buffer = Buffer::local_normalized(text, line_ending, cx);
1932            buffer.set_language(language, cx);
1933            if let Some(language_registry) = language_registry {
1934                buffer.set_language_registry(language_registry)
1935            }
1936            buffer
1937        });
1938
1939        Self {
1940            buffer: buffer.clone(),
1941            old_buffer,
1942            edit_position: range.start,
1943            range,
1944            snapshot,
1945            last_equal_ranges: Default::default(),
1946            transaction_id: Default::default(),
1947            status: CodegenStatus::Idle,
1948            generation: Task::ready(()),
1949            diff: Diff::default(),
1950            telemetry,
1951            _subscription: cx.subscribe(&buffer, Self::handle_buffer_event),
1952        }
1953    }
1954
1955    fn handle_buffer_event(
1956        &mut self,
1957        _buffer: Model<MultiBuffer>,
1958        event: &multi_buffer::Event,
1959        cx: &mut ModelContext<Self>,
1960    ) {
1961        if let multi_buffer::Event::TransactionUndone { transaction_id } = event {
1962            if self.transaction_id == Some(*transaction_id) {
1963                self.transaction_id = None;
1964                self.generation = Task::ready(());
1965                cx.emit(CodegenEvent::Undone);
1966            }
1967        }
1968    }
1969
1970    pub fn last_equal_ranges(&self) -> &[Range<Anchor>] {
1971        &self.last_equal_ranges
1972    }
1973
1974    pub fn start(&mut self, prompt: LanguageModelRequest, cx: &mut ModelContext<Self>) {
1975        let range = self.range.clone();
1976        let snapshot = self.snapshot.clone();
1977        let selected_text = snapshot
1978            .text_for_range(range.start..range.end)
1979            .collect::<Rope>();
1980
1981        let selection_start = range.start.to_point(&snapshot);
1982        let suggested_line_indent = snapshot
1983            .suggested_indents(selection_start.row..selection_start.row + 1, cx)
1984            .into_values()
1985            .next()
1986            .unwrap_or_else(|| snapshot.indent_size_for_line(MultiBufferRow(selection_start.row)));
1987
1988        let model_telemetry_id = prompt.model.telemetry_id();
1989        let response = CompletionProvider::global(cx).complete(prompt, cx);
1990        let telemetry = self.telemetry.clone();
1991        self.edit_position = range.start;
1992        self.diff = Diff::default();
1993        self.status = CodegenStatus::Pending;
1994        self.generation = cx.spawn(|this, mut cx| {
1995            async move {
1996                let response = response.await;
1997                let generate = async {
1998                    let mut edit_start = range.start.to_offset(&snapshot);
1999
2000                    let (mut hunks_tx, mut hunks_rx) = mpsc::channel(1);
2001                    let diff: Task<anyhow::Result<()>> =
2002                        cx.background_executor().spawn(async move {
2003                            let mut response_latency = None;
2004                            let request_start = Instant::now();
2005                            let diff = async {
2006                                let chunks = StripInvalidSpans::new(response.inner.await?);
2007                                futures::pin_mut!(chunks);
2008                                let mut diff = StreamingDiff::new(selected_text.to_string());
2009
2010                                let mut new_text = String::new();
2011                                let mut base_indent = None;
2012                                let mut line_indent = None;
2013                                let mut first_line = true;
2014
2015                                while let Some(chunk) = chunks.next().await {
2016                                    if response_latency.is_none() {
2017                                        response_latency = Some(request_start.elapsed());
2018                                    }
2019                                    let chunk = chunk?;
2020
2021                                    let mut lines = chunk.split('\n').peekable();
2022                                    while let Some(line) = lines.next() {
2023                                        new_text.push_str(line);
2024                                        if line_indent.is_none() {
2025                                            if let Some(non_whitespace_ch_ix) =
2026                                                new_text.find(|ch: char| !ch.is_whitespace())
2027                                            {
2028                                                line_indent = Some(non_whitespace_ch_ix);
2029                                                base_indent = base_indent.or(line_indent);
2030
2031                                                let line_indent = line_indent.unwrap();
2032                                                let base_indent = base_indent.unwrap();
2033                                                let indent_delta =
2034                                                    line_indent as i32 - base_indent as i32;
2035                                                let mut corrected_indent_len = cmp::max(
2036                                                    0,
2037                                                    suggested_line_indent.len as i32 + indent_delta,
2038                                                )
2039                                                    as usize;
2040                                                if first_line {
2041                                                    corrected_indent_len = corrected_indent_len
2042                                                        .saturating_sub(
2043                                                            selection_start.column as usize,
2044                                                        );
2045                                                }
2046
2047                                                let indent_char = suggested_line_indent.char();
2048                                                let mut indent_buffer = [0; 4];
2049                                                let indent_str =
2050                                                    indent_char.encode_utf8(&mut indent_buffer);
2051                                                new_text.replace_range(
2052                                                    ..line_indent,
2053                                                    &indent_str.repeat(corrected_indent_len),
2054                                                );
2055                                            }
2056                                        }
2057
2058                                        if line_indent.is_some() {
2059                                            hunks_tx.send(diff.push_new(&new_text)).await?;
2060                                            new_text.clear();
2061                                        }
2062
2063                                        if lines.peek().is_some() {
2064                                            hunks_tx.send(diff.push_new("\n")).await?;
2065                                            if line_indent.is_none() {
2066                                                // Don't write out the leading indentation in empty lines on the next line
2067                                                // This is the case where the above if statement didn't clear the buffer
2068                                                new_text.clear();
2069                                            }
2070                                            line_indent = None;
2071                                            first_line = false;
2072                                        }
2073                                    }
2074                                }
2075                                hunks_tx.send(diff.push_new(&new_text)).await?;
2076                                hunks_tx.send(diff.finish()).await?;
2077
2078                                anyhow::Ok(())
2079                            };
2080
2081                            let result = diff.await;
2082
2083                            let error_message =
2084                                result.as_ref().err().map(|error| error.to_string());
2085                            if let Some(telemetry) = telemetry {
2086                                telemetry.report_assistant_event(
2087                                    None,
2088                                    telemetry_events::AssistantKind::Inline,
2089                                    model_telemetry_id,
2090                                    response_latency,
2091                                    error_message,
2092                                );
2093                            }
2094
2095                            result?;
2096                            Ok(())
2097                        });
2098
2099                    while let Some(hunks) = hunks_rx.next().await {
2100                        this.update(&mut cx, |this, cx| {
2101                            this.last_equal_ranges.clear();
2102
2103                            let transaction = this.buffer.update(cx, |buffer, cx| {
2104                                // Avoid grouping assistant edits with user edits.
2105                                buffer.finalize_last_transaction(cx);
2106
2107                                buffer.start_transaction(cx);
2108                                buffer.edit(
2109                                    hunks.into_iter().filter_map(|hunk| match hunk {
2110                                        Hunk::Insert { text } => {
2111                                            let edit_start = snapshot.anchor_after(edit_start);
2112                                            Some((edit_start..edit_start, text))
2113                                        }
2114                                        Hunk::Remove { len } => {
2115                                            let edit_end = edit_start + len;
2116                                            let edit_range = snapshot.anchor_after(edit_start)
2117                                                ..snapshot.anchor_before(edit_end);
2118                                            edit_start = edit_end;
2119                                            Some((edit_range, String::new()))
2120                                        }
2121                                        Hunk::Keep { len } => {
2122                                            let edit_end = edit_start + len;
2123                                            let edit_range = snapshot.anchor_after(edit_start)
2124                                                ..snapshot.anchor_before(edit_end);
2125                                            edit_start = edit_end;
2126                                            this.last_equal_ranges.push(edit_range);
2127                                            None
2128                                        }
2129                                    }),
2130                                    None,
2131                                    cx,
2132                                );
2133                                this.edit_position = snapshot.anchor_after(edit_start);
2134
2135                                buffer.end_transaction(cx)
2136                            });
2137
2138                            if let Some(transaction) = transaction {
2139                                if let Some(first_transaction) = this.transaction_id {
2140                                    // Group all assistant edits into the first transaction.
2141                                    this.buffer.update(cx, |buffer, cx| {
2142                                        buffer.merge_transactions(
2143                                            transaction,
2144                                            first_transaction,
2145                                            cx,
2146                                        )
2147                                    });
2148                                } else {
2149                                    this.transaction_id = Some(transaction);
2150                                    this.buffer.update(cx, |buffer, cx| {
2151                                        buffer.finalize_last_transaction(cx)
2152                                    });
2153                                }
2154                            }
2155
2156                            this.update_diff(cx);
2157                            cx.notify();
2158                        })?;
2159                    }
2160
2161                    diff.await?;
2162
2163                    anyhow::Ok(())
2164                };
2165
2166                let result = generate.await;
2167                this.update(&mut cx, |this, cx| {
2168                    this.last_equal_ranges.clear();
2169                    if let Err(error) = result {
2170                        this.status = CodegenStatus::Error(error);
2171                    } else {
2172                        this.status = CodegenStatus::Done;
2173                    }
2174                    cx.emit(CodegenEvent::Finished);
2175                    cx.notify();
2176                })
2177                .ok();
2178            }
2179        });
2180        cx.notify();
2181    }
2182
2183    pub fn stop(&mut self, cx: &mut ModelContext<Self>) {
2184        self.last_equal_ranges.clear();
2185        self.status = CodegenStatus::Done;
2186        self.generation = Task::ready(());
2187        cx.emit(CodegenEvent::Finished);
2188        cx.notify();
2189    }
2190
2191    pub fn undo(&mut self, cx: &mut ModelContext<Self>) {
2192        if let Some(transaction_id) = self.transaction_id.take() {
2193            self.buffer
2194                .update(cx, |buffer, cx| buffer.undo_transaction(transaction_id, cx));
2195        }
2196    }
2197
2198    fn update_diff(&mut self, cx: &mut ModelContext<Self>) {
2199        if self.diff.task.is_some() {
2200            self.diff.should_update = true;
2201        } else {
2202            self.diff.should_update = false;
2203
2204            let old_snapshot = self.snapshot.clone();
2205            let old_range = self.range.to_point(&old_snapshot);
2206            let new_snapshot = self.buffer.read(cx).snapshot(cx);
2207            let new_range = self.range.to_point(&new_snapshot);
2208
2209            self.diff.task = Some(cx.spawn(|this, mut cx| async move {
2210                let (deleted_row_ranges, inserted_row_ranges) = cx
2211                    .background_executor()
2212                    .spawn(async move {
2213                        let old_text = old_snapshot
2214                            .text_for_range(
2215                                Point::new(old_range.start.row, 0)
2216                                    ..Point::new(
2217                                        old_range.end.row,
2218                                        old_snapshot.line_len(MultiBufferRow(old_range.end.row)),
2219                                    ),
2220                            )
2221                            .collect::<String>();
2222                        let new_text = new_snapshot
2223                            .text_for_range(
2224                                Point::new(new_range.start.row, 0)
2225                                    ..Point::new(
2226                                        new_range.end.row,
2227                                        new_snapshot.line_len(MultiBufferRow(new_range.end.row)),
2228                                    ),
2229                            )
2230                            .collect::<String>();
2231
2232                        let mut old_row = old_range.start.row;
2233                        let mut new_row = new_range.start.row;
2234                        let diff = TextDiff::from_lines(old_text.as_str(), new_text.as_str());
2235
2236                        let mut deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)> = Vec::new();
2237                        let mut inserted_row_ranges = Vec::new();
2238                        for change in diff.iter_all_changes() {
2239                            let line_count = change.value().lines().count() as u32;
2240                            match change.tag() {
2241                                similar::ChangeTag::Equal => {
2242                                    old_row += line_count;
2243                                    new_row += line_count;
2244                                }
2245                                similar::ChangeTag::Delete => {
2246                                    let old_end_row = old_row + line_count - 1;
2247                                    let new_row =
2248                                        new_snapshot.anchor_before(Point::new(new_row, 0));
2249
2250                                    if let Some((_, last_deleted_row_range)) =
2251                                        deleted_row_ranges.last_mut()
2252                                    {
2253                                        if *last_deleted_row_range.end() + 1 == old_row {
2254                                            *last_deleted_row_range =
2255                                                *last_deleted_row_range.start()..=old_end_row;
2256                                        } else {
2257                                            deleted_row_ranges
2258                                                .push((new_row, old_row..=old_end_row));
2259                                        }
2260                                    } else {
2261                                        deleted_row_ranges.push((new_row, old_row..=old_end_row));
2262                                    }
2263
2264                                    old_row += line_count;
2265                                }
2266                                similar::ChangeTag::Insert => {
2267                                    let new_end_row = new_row + line_count - 1;
2268                                    let start = new_snapshot.anchor_before(Point::new(new_row, 0));
2269                                    let end = new_snapshot.anchor_before(Point::new(
2270                                        new_end_row,
2271                                        new_snapshot.line_len(MultiBufferRow(new_end_row)),
2272                                    ));
2273                                    inserted_row_ranges.push(start..=end);
2274                                    new_row += line_count;
2275                                }
2276                            }
2277                        }
2278
2279                        (deleted_row_ranges, inserted_row_ranges)
2280                    })
2281                    .await;
2282
2283                this.update(&mut cx, |this, cx| {
2284                    this.diff.deleted_row_ranges = deleted_row_ranges;
2285                    this.diff.inserted_row_ranges = inserted_row_ranges;
2286                    this.diff.task = None;
2287                    if this.diff.should_update {
2288                        this.update_diff(cx);
2289                    }
2290                    cx.notify();
2291                })
2292                .ok();
2293            }));
2294        }
2295    }
2296}
2297
2298struct StripInvalidSpans<T> {
2299    stream: T,
2300    stream_done: bool,
2301    buffer: String,
2302    first_line: bool,
2303    line_end: bool,
2304    starts_with_code_block: bool,
2305}
2306
2307impl<T> StripInvalidSpans<T>
2308where
2309    T: Stream<Item = Result<String>>,
2310{
2311    fn new(stream: T) -> Self {
2312        Self {
2313            stream,
2314            stream_done: false,
2315            buffer: String::new(),
2316            first_line: true,
2317            line_end: false,
2318            starts_with_code_block: false,
2319        }
2320    }
2321}
2322
2323impl<T> Stream for StripInvalidSpans<T>
2324where
2325    T: Stream<Item = Result<String>>,
2326{
2327    type Item = Result<String>;
2328
2329    fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Option<Self::Item>> {
2330        const CODE_BLOCK_DELIMITER: &str = "```";
2331        const CURSOR_SPAN: &str = "<|CURSOR|>";
2332
2333        let this = unsafe { self.get_unchecked_mut() };
2334        loop {
2335            if !this.stream_done {
2336                let mut stream = unsafe { Pin::new_unchecked(&mut this.stream) };
2337                match stream.as_mut().poll_next(cx) {
2338                    Poll::Ready(Some(Ok(chunk))) => {
2339                        this.buffer.push_str(&chunk);
2340                    }
2341                    Poll::Ready(Some(Err(error))) => return Poll::Ready(Some(Err(error))),
2342                    Poll::Ready(None) => {
2343                        this.stream_done = true;
2344                    }
2345                    Poll::Pending => return Poll::Pending,
2346                }
2347            }
2348
2349            let mut chunk = String::new();
2350            let mut consumed = 0;
2351            if !this.buffer.is_empty() {
2352                let mut lines = this.buffer.split('\n').enumerate().peekable();
2353                while let Some((line_ix, line)) = lines.next() {
2354                    if line_ix > 0 {
2355                        this.first_line = false;
2356                    }
2357
2358                    if this.first_line {
2359                        let trimmed_line = line.trim();
2360                        if lines.peek().is_some() {
2361                            if trimmed_line.starts_with(CODE_BLOCK_DELIMITER) {
2362                                consumed += line.len() + 1;
2363                                this.starts_with_code_block = true;
2364                                continue;
2365                            }
2366                        } else if trimmed_line.is_empty()
2367                            || prefixes(CODE_BLOCK_DELIMITER)
2368                                .any(|prefix| trimmed_line.starts_with(prefix))
2369                        {
2370                            break;
2371                        }
2372                    }
2373
2374                    let line_without_cursor = line.replace(CURSOR_SPAN, "");
2375                    if lines.peek().is_some() {
2376                        if this.line_end {
2377                            chunk.push('\n');
2378                        }
2379
2380                        chunk.push_str(&line_without_cursor);
2381                        this.line_end = true;
2382                        consumed += line.len() + 1;
2383                    } else if this.stream_done {
2384                        if !this.starts_with_code_block
2385                            || !line_without_cursor.trim().ends_with(CODE_BLOCK_DELIMITER)
2386                        {
2387                            if this.line_end {
2388                                chunk.push('\n');
2389                            }
2390
2391                            chunk.push_str(&line);
2392                        }
2393
2394                        consumed += line.len();
2395                    } else {
2396                        let trimmed_line = line.trim();
2397                        if trimmed_line.is_empty()
2398                            || prefixes(CURSOR_SPAN).any(|prefix| trimmed_line.ends_with(prefix))
2399                            || prefixes(CODE_BLOCK_DELIMITER)
2400                                .any(|prefix| trimmed_line.ends_with(prefix))
2401                        {
2402                            break;
2403                        } else {
2404                            if this.line_end {
2405                                chunk.push('\n');
2406                                this.line_end = false;
2407                            }
2408
2409                            chunk.push_str(&line_without_cursor);
2410                            consumed += line.len();
2411                        }
2412                    }
2413                }
2414            }
2415
2416            this.buffer = this.buffer.split_off(consumed);
2417            if !chunk.is_empty() {
2418                return Poll::Ready(Some(Ok(chunk)));
2419            } else if this.stream_done {
2420                return Poll::Ready(None);
2421            }
2422        }
2423    }
2424}
2425
2426fn prefixes(text: &str) -> impl Iterator<Item = &str> {
2427    (0..text.len() - 1).map(|ix| &text[..ix + 1])
2428}
2429
2430fn merge_ranges(ranges: &mut Vec<Range<Anchor>>, buffer: &MultiBufferSnapshot) {
2431    ranges.sort_unstable_by(|a, b| {
2432        a.start
2433            .cmp(&b.start, buffer)
2434            .then_with(|| b.end.cmp(&a.end, buffer))
2435    });
2436
2437    let mut ix = 0;
2438    while ix + 1 < ranges.len() {
2439        let b = ranges[ix + 1].clone();
2440        let a = &mut ranges[ix];
2441        if a.end.cmp(&b.start, buffer).is_gt() {
2442            if a.end.cmp(&b.end, buffer).is_lt() {
2443                a.end = b.end;
2444            }
2445            ranges.remove(ix + 1);
2446        } else {
2447            ix += 1;
2448        }
2449    }
2450}
2451
2452#[cfg(test)]
2453mod tests {
2454    use std::sync::Arc;
2455
2456    use crate::FakeCompletionProvider;
2457
2458    use super::*;
2459    use futures::stream::{self};
2460    use gpui::{Context, TestAppContext};
2461    use indoc::indoc;
2462    use language::{
2463        language_settings, tree_sitter_rust, Buffer, Language, LanguageConfig, LanguageMatcher,
2464        Point,
2465    };
2466    use rand::prelude::*;
2467    use serde::Serialize;
2468    use settings::SettingsStore;
2469
2470    #[derive(Serialize)]
2471    pub struct DummyCompletionRequest {
2472        pub name: String,
2473    }
2474
2475    #[gpui::test(iterations = 10)]
2476    async fn test_transform_autoindent(cx: &mut TestAppContext, mut rng: StdRng) {
2477        cx.set_global(cx.update(SettingsStore::test));
2478        let provider = cx.update(|cx| FakeCompletionProvider::setup_test(cx));
2479        cx.update(language_settings::init);
2480
2481        let text = indoc! {"
2482            fn main() {
2483                let x = 0;
2484                for _ in 0..10 {
2485                    x += 1;
2486                }
2487            }
2488        "};
2489        let buffer =
2490            cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
2491        let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
2492        let range = buffer.read_with(cx, |buffer, cx| {
2493            let snapshot = buffer.snapshot(cx);
2494            snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(4, 5))
2495        });
2496        let codegen = cx.new_model(|cx| Codegen::new(buffer.clone(), range, None, cx));
2497
2498        codegen.update(cx, |codegen, cx| {
2499            codegen.start(LanguageModelRequest::default(), cx)
2500        });
2501
2502        cx.background_executor.run_until_parked();
2503
2504        let mut new_text = concat!(
2505            "       let mut x = 0;\n",
2506            "       while x < 10 {\n",
2507            "           x += 1;\n",
2508            "       }",
2509        );
2510        while !new_text.is_empty() {
2511            let max_len = cmp::min(new_text.len(), 10);
2512            let len = rng.gen_range(1..=max_len);
2513            let (chunk, suffix) = new_text.split_at(len);
2514            provider.send_completion(&LanguageModelRequest::default(), chunk.into());
2515            new_text = suffix;
2516            cx.background_executor.run_until_parked();
2517        }
2518        provider.finish_completion(&LanguageModelRequest::default());
2519        cx.background_executor.run_until_parked();
2520
2521        assert_eq!(
2522            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
2523            indoc! {"
2524                fn main() {
2525                    let mut x = 0;
2526                    while x < 10 {
2527                        x += 1;
2528                    }
2529                }
2530            "}
2531        );
2532    }
2533
2534    #[gpui::test(iterations = 10)]
2535    async fn test_autoindent_when_generating_past_indentation(
2536        cx: &mut TestAppContext,
2537        mut rng: StdRng,
2538    ) {
2539        let provider = cx.update(|cx| FakeCompletionProvider::setup_test(cx));
2540        cx.set_global(cx.update(SettingsStore::test));
2541        cx.update(language_settings::init);
2542
2543        let text = indoc! {"
2544            fn main() {
2545                le
2546            }
2547        "};
2548        let buffer =
2549            cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
2550        let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
2551        let range = buffer.read_with(cx, |buffer, cx| {
2552            let snapshot = buffer.snapshot(cx);
2553            snapshot.anchor_before(Point::new(1, 6))..snapshot.anchor_after(Point::new(1, 6))
2554        });
2555        let codegen = cx.new_model(|cx| Codegen::new(buffer.clone(), range, None, cx));
2556
2557        let request = LanguageModelRequest::default();
2558        codegen.update(cx, |codegen, cx| codegen.start(request, cx));
2559
2560        cx.background_executor.run_until_parked();
2561
2562        let mut new_text = concat!(
2563            "t mut x = 0;\n",
2564            "while x < 10 {\n",
2565            "    x += 1;\n",
2566            "}", //
2567        );
2568        while !new_text.is_empty() {
2569            let max_len = cmp::min(new_text.len(), 10);
2570            let len = rng.gen_range(1..=max_len);
2571            let (chunk, suffix) = new_text.split_at(len);
2572            provider.send_completion(&LanguageModelRequest::default(), chunk.into());
2573            new_text = suffix;
2574            cx.background_executor.run_until_parked();
2575        }
2576        provider.finish_completion(&LanguageModelRequest::default());
2577        cx.background_executor.run_until_parked();
2578
2579        assert_eq!(
2580            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
2581            indoc! {"
2582                fn main() {
2583                    let mut x = 0;
2584                    while x < 10 {
2585                        x += 1;
2586                    }
2587                }
2588            "}
2589        );
2590    }
2591
2592    #[gpui::test(iterations = 10)]
2593    async fn test_autoindent_when_generating_before_indentation(
2594        cx: &mut TestAppContext,
2595        mut rng: StdRng,
2596    ) {
2597        let provider = cx.update(|cx| FakeCompletionProvider::setup_test(cx));
2598        cx.set_global(cx.update(SettingsStore::test));
2599        cx.update(language_settings::init);
2600
2601        let text = concat!(
2602            "fn main() {\n",
2603            "  \n",
2604            "}\n" //
2605        );
2606        let buffer =
2607            cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
2608        let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
2609        let range = buffer.read_with(cx, |buffer, cx| {
2610            let snapshot = buffer.snapshot(cx);
2611            snapshot.anchor_before(Point::new(1, 2))..snapshot.anchor_after(Point::new(1, 2))
2612        });
2613        let codegen = cx.new_model(|cx| Codegen::new(buffer.clone(), range, None, cx));
2614
2615        let request = LanguageModelRequest::default();
2616        codegen.update(cx, |codegen, cx| codegen.start(request, cx));
2617
2618        cx.background_executor.run_until_parked();
2619
2620        let mut new_text = concat!(
2621            "let mut x = 0;\n",
2622            "while x < 10 {\n",
2623            "    x += 1;\n",
2624            "}", //
2625        );
2626        while !new_text.is_empty() {
2627            let max_len = cmp::min(new_text.len(), 10);
2628            let len = rng.gen_range(1..=max_len);
2629            let (chunk, suffix) = new_text.split_at(len);
2630            provider.send_completion(&LanguageModelRequest::default(), chunk.into());
2631            new_text = suffix;
2632            cx.background_executor.run_until_parked();
2633        }
2634        provider.finish_completion(&LanguageModelRequest::default());
2635        cx.background_executor.run_until_parked();
2636
2637        assert_eq!(
2638            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
2639            indoc! {"
2640                fn main() {
2641                    let mut x = 0;
2642                    while x < 10 {
2643                        x += 1;
2644                    }
2645                }
2646            "}
2647        );
2648    }
2649
2650    #[gpui::test]
2651    async fn test_strip_invalid_spans_from_codeblock() {
2652        assert_chunks("Lorem ipsum dolor", "Lorem ipsum dolor").await;
2653        assert_chunks("```\nLorem ipsum dolor", "Lorem ipsum dolor").await;
2654        assert_chunks("```\nLorem ipsum dolor\n```", "Lorem ipsum dolor").await;
2655        assert_chunks(
2656            "```html\n```js\nLorem ipsum dolor\n```\n```",
2657            "```js\nLorem ipsum dolor\n```",
2658        )
2659        .await;
2660        assert_chunks("``\nLorem ipsum dolor\n```", "``\nLorem ipsum dolor\n```").await;
2661        assert_chunks("Lorem<|CURSOR|> ipsum", "Lorem ipsum").await;
2662        assert_chunks("Lorem ipsum", "Lorem ipsum").await;
2663        assert_chunks("```\n<|CURSOR|>Lorem ipsum\n```", "Lorem ipsum").await;
2664
2665        async fn assert_chunks(text: &str, expected_text: &str) {
2666            for chunk_size in 1..=text.len() {
2667                let actual_text = StripInvalidSpans::new(chunks(text, chunk_size))
2668                    .map(|chunk| chunk.unwrap())
2669                    .collect::<String>()
2670                    .await;
2671                assert_eq!(
2672                    actual_text, expected_text,
2673                    "failed to strip invalid spans, chunk size: {}",
2674                    chunk_size
2675                );
2676            }
2677        }
2678
2679        fn chunks(text: &str, size: usize) -> impl Stream<Item = Result<String>> {
2680            stream::iter(
2681                text.chars()
2682                    .collect::<Vec<_>>()
2683                    .chunks(size)
2684                    .map(|chunk| Ok(chunk.iter().collect::<String>()))
2685                    .collect::<Vec<_>>(),
2686            )
2687        }
2688    }
2689
2690    fn rust_lang() -> Language {
2691        Language::new(
2692            LanguageConfig {
2693                name: "Rust".into(),
2694                matcher: LanguageMatcher {
2695                    path_suffixes: vec!["rs".to_string()],
2696                    ..Default::default()
2697                },
2698                ..Default::default()
2699            },
2700            Some(tree_sitter_rust::language()),
2701        )
2702        .with_indents_query(
2703            r#"
2704            (call_expression) @indent
2705            (field_expression) @indent
2706            (_ "(" ")" @end) @indent
2707            (_ "{" "}" @end) @indent
2708            "#,
2709        )
2710        .unwrap()
2711    }
2712}