1use crate::{
   2    prompts::generate_content_prompt, AssistantPanel, CompletionProvider, Hunk,
   3    LanguageModelRequest, LanguageModelRequestMessage, Role, StreamingDiff,
   4};
   5use anyhow::Result;
   6use client::telemetry::Telemetry;
   7use collections::{hash_map, HashMap, HashSet, VecDeque};
   8use editor::{
   9    actions::{MoveDown, MoveUp},
  10    display_map::{
  11        BlockContext, BlockDisposition, BlockId, BlockProperties, BlockStyle, RenderBlock,
  12    },
  13    scroll::{Autoscroll, AutoscrollStrategy},
  14    Anchor, AnchorRangeExt, Editor, EditorElement, EditorEvent, EditorStyle, ExcerptRange,
  15    GutterDimensions, MultiBuffer, MultiBufferSnapshot, ToOffset, ToPoint,
  16};
  17use futures::{channel::mpsc, SinkExt, Stream, StreamExt};
  18use gpui::{
  19    AppContext, EventEmitter, FocusHandle, FocusableView, FontStyle, FontWeight, Global,
  20    HighlightStyle, Model, ModelContext, Subscription, Task, TextStyle, UpdateGlobal, View,
  21    ViewContext, WeakView, WhiteSpace, WindowContext,
  22};
  23use language::{Buffer, Point, TransactionId};
  24use multi_buffer::MultiBufferRow;
  25use parking_lot::Mutex;
  26use rope::Rope;
  27use settings::Settings;
  28use similar::TextDiff;
  29use std::{
  30    cmp, future, mem,
  31    ops::{Range, RangeInclusive},
  32    sync::Arc,
  33    time::Instant,
  34};
  35use theme::ThemeSettings;
  36use ui::{prelude::*, Tooltip};
  37use util::RangeExt;
  38use workspace::{notifications::NotificationId, Toast, Workspace};
  39
  40pub fn init(telemetry: Arc<Telemetry>, cx: &mut AppContext) {
  41    cx.set_global(InlineAssistant::new(telemetry));
  42}
  43
  44const PROMPT_HISTORY_MAX_LEN: usize = 20;
  45
  46pub struct InlineAssistant {
  47    next_assist_id: InlineAssistId,
  48    pending_assists: HashMap<InlineAssistId, PendingInlineAssist>,
  49    pending_assist_ids_by_editor: HashMap<WeakView<Editor>, Vec<InlineAssistId>>,
  50    prompt_history: VecDeque<String>,
  51    telemetry: Option<Arc<Telemetry>>,
  52}
  53
  54impl Global for InlineAssistant {}
  55
  56impl InlineAssistant {
  57    pub fn new(telemetry: Arc<Telemetry>) -> Self {
  58        Self {
  59            next_assist_id: InlineAssistId::default(),
  60            pending_assists: HashMap::default(),
  61            pending_assist_ids_by_editor: HashMap::default(),
  62            prompt_history: VecDeque::default(),
  63            telemetry: Some(telemetry),
  64        }
  65    }
  66
  67    pub fn assist(
  68        &mut self,
  69        editor: &View<Editor>,
  70        workspace: Option<WeakView<Workspace>>,
  71        include_context: bool,
  72        cx: &mut WindowContext,
  73    ) {
  74        let selection = editor.read(cx).selections.newest_anchor().clone();
  75        if selection.start.excerpt_id != selection.end.excerpt_id {
  76            return;
  77        }
  78        let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
  79
  80        // Extend the selection to the start and the end of the line.
  81        let mut point_selection = selection.map(|selection| selection.to_point(&snapshot));
  82        if point_selection.end > point_selection.start {
  83            point_selection.start.column = 0;
  84            // If the selection ends at the start of the line, we don't want to include it.
  85            if point_selection.end.column == 0 {
  86                point_selection.end.row -= 1;
  87            }
  88            point_selection.end.column = snapshot.line_len(MultiBufferRow(point_selection.end.row));
  89        }
  90
  91        let codegen_kind = if point_selection.start == point_selection.end {
  92            CodegenKind::Generate {
  93                position: snapshot.anchor_after(point_selection.start),
  94            }
  95        } else {
  96            CodegenKind::Transform {
  97                range: snapshot.anchor_before(point_selection.start)
  98                    ..snapshot.anchor_after(point_selection.end),
  99            }
 100        };
 101
 102        let assist_id = self.next_assist_id.post_inc();
 103        let codegen = cx.new_model(|cx| {
 104            Codegen::new(
 105                editor.read(cx).buffer().clone(),
 106                codegen_kind,
 107                self.telemetry.clone(),
 108                cx,
 109            )
 110        });
 111
 112        let gutter_dimensions = Arc::new(Mutex::new(GutterDimensions::default()));
 113        let prompt_editor = cx.new_view(|cx| {
 114            InlineAssistEditor::new(
 115                assist_id,
 116                gutter_dimensions.clone(),
 117                self.prompt_history.clone(),
 118                codegen.clone(),
 119                workspace.clone(),
 120                cx,
 121            )
 122        });
 123        let (prompt_block_id, end_block_id) = editor.update(cx, |editor, cx| {
 124            let start_anchor = snapshot.anchor_before(point_selection.start);
 125            let end_anchor = snapshot.anchor_after(point_selection.end);
 126            editor.change_selections(Some(Autoscroll::newest()), cx, |selections| {
 127                selections.select_anchor_ranges([start_anchor..start_anchor])
 128            });
 129            let block_ids = editor.insert_blocks(
 130                [
 131                    BlockProperties {
 132                        style: BlockStyle::Sticky,
 133                        position: start_anchor,
 134                        height: prompt_editor.read(cx).height_in_lines,
 135                        render: build_inline_assist_editor_renderer(
 136                            &prompt_editor,
 137                            gutter_dimensions,
 138                        ),
 139                        disposition: BlockDisposition::Above,
 140                    },
 141                    BlockProperties {
 142                        style: BlockStyle::Sticky,
 143                        position: end_anchor,
 144                        height: 1,
 145                        render: Box::new(|cx| {
 146                            v_flex()
 147                                .h_full()
 148                                .w_full()
 149                                .border_t_1()
 150                                .border_color(cx.theme().status().info_border)
 151                                .into_any_element()
 152                        }),
 153                        disposition: BlockDisposition::Below,
 154                    },
 155                ],
 156                Some(Autoscroll::Strategy(AutoscrollStrategy::Newest)),
 157                cx,
 158            );
 159            (block_ids[0], block_ids[1])
 160        });
 161
 162        self.pending_assists.insert(
 163            assist_id,
 164            PendingInlineAssist {
 165                include_context,
 166                editor: editor.downgrade(),
 167                editor_decorations: Some(PendingInlineAssistDecorations {
 168                    prompt_block_id,
 169                    prompt_editor: prompt_editor.clone(),
 170                    removed_line_block_ids: HashSet::default(),
 171                    end_block_id,
 172                }),
 173                codegen: codegen.clone(),
 174                workspace,
 175                _subscriptions: vec![
 176                    cx.subscribe(&prompt_editor, |inline_assist_editor, event, cx| {
 177                        InlineAssistant::update_global(cx, |this, cx| {
 178                            this.handle_inline_assistant_editor_event(
 179                                inline_assist_editor,
 180                                event,
 181                                cx,
 182                            )
 183                        })
 184                    }),
 185                    editor.update(cx, |editor, _cx| {
 186                        editor.register_action(
 187                            move |_: &editor::actions::Newline, cx: &mut WindowContext| {
 188                                InlineAssistant::update_global(cx, |this, cx| {
 189                                    this.handle_editor_action(assist_id, false, cx)
 190                                })
 191                            },
 192                        )
 193                    }),
 194                    editor.update(cx, |editor, _cx| {
 195                        editor.register_action(
 196                            move |_: &editor::actions::Cancel, cx: &mut WindowContext| {
 197                                InlineAssistant::update_global(cx, |this, cx| {
 198                                    this.handle_editor_action(assist_id, true, cx)
 199                                })
 200                            },
 201                        )
 202                    }),
 203                    cx.subscribe(editor, move |editor, event, cx| {
 204                        InlineAssistant::update_global(cx, |this, cx| {
 205                            this.handle_editor_event(assist_id, editor, event, cx)
 206                        })
 207                    }),
 208                    cx.observe(&codegen, {
 209                        let editor = editor.downgrade();
 210                        move |_, cx| {
 211                            if let Some(editor) = editor.upgrade() {
 212                                InlineAssistant::update_global(cx, |this, cx| {
 213                                    this.update_editor_highlights(&editor, cx);
 214                                    this.update_editor_blocks(&editor, assist_id, cx);
 215                                })
 216                            }
 217                        }
 218                    }),
 219                    cx.subscribe(&codegen, move |codegen, event, cx| {
 220                        InlineAssistant::update_global(cx, |this, cx| match event {
 221                            CodegenEvent::Undone => this.finish_inline_assist(assist_id, false, cx),
 222                            CodegenEvent::Finished => {
 223                                let pending_assist = if let Some(pending_assist) =
 224                                    this.pending_assists.get(&assist_id)
 225                                {
 226                                    pending_assist
 227                                } else {
 228                                    return;
 229                                };
 230
 231                                if let CodegenStatus::Error(error) = &codegen.read(cx).status {
 232                                    if pending_assist.editor_decorations.is_none() {
 233                                        if let Some(workspace) = pending_assist
 234                                            .workspace
 235                                            .as_ref()
 236                                            .and_then(|workspace| workspace.upgrade())
 237                                        {
 238                                            let error =
 239                                                format!("Inline assistant error: {}", error);
 240                                            workspace.update(cx, |workspace, cx| {
 241                                                struct InlineAssistantError;
 242
 243                                                let id = NotificationId::identified::<
 244                                                    InlineAssistantError,
 245                                                >(
 246                                                    assist_id.0
 247                                                );
 248
 249                                                workspace.show_toast(Toast::new(id, error), cx);
 250                                            })
 251                                        }
 252                                    }
 253                                }
 254
 255                                if pending_assist.editor_decorations.is_none() {
 256                                    this.finish_inline_assist(assist_id, false, cx);
 257                                }
 258                            }
 259                        })
 260                    }),
 261                ],
 262            },
 263        );
 264
 265        self.pending_assist_ids_by_editor
 266            .entry(editor.downgrade())
 267            .or_default()
 268            .push(assist_id);
 269        self.update_editor_highlights(editor, cx);
 270    }
 271
 272    fn handle_inline_assistant_editor_event(
 273        &mut self,
 274        inline_assist_editor: View<InlineAssistEditor>,
 275        event: &InlineAssistEditorEvent,
 276        cx: &mut WindowContext,
 277    ) {
 278        let assist_id = inline_assist_editor.read(cx).id;
 279        match event {
 280            InlineAssistEditorEvent::StartRequested => {
 281                self.start_inline_assist(assist_id, cx);
 282            }
 283            InlineAssistEditorEvent::StopRequested => {
 284                self.stop_inline_assist(assist_id, cx);
 285            }
 286            InlineAssistEditorEvent::ConfirmRequested => {
 287                self.finish_inline_assist(assist_id, false, cx);
 288            }
 289            InlineAssistEditorEvent::CancelRequested => {
 290                self.finish_inline_assist(assist_id, true, cx);
 291            }
 292            InlineAssistEditorEvent::DismissRequested => {
 293                self.dismiss_inline_assist(assist_id, cx);
 294            }
 295            InlineAssistEditorEvent::Resized { height_in_lines } => {
 296                self.resize_inline_assist(assist_id, *height_in_lines, cx);
 297            }
 298        }
 299    }
 300
 301    fn handle_editor_action(
 302        &mut self,
 303        assist_id: InlineAssistId,
 304        undo: bool,
 305        cx: &mut WindowContext,
 306    ) {
 307        let Some(assist) = self.pending_assists.get(&assist_id) else {
 308            return;
 309        };
 310        let Some(editor) = assist.editor.upgrade() else {
 311            return;
 312        };
 313
 314        let buffer = editor.read(cx).buffer().read(cx).snapshot(cx);
 315        let assist_range = assist.codegen.read(cx).range().to_offset(&buffer);
 316        let editor = editor.read(cx);
 317        if editor.selections.count() == 1 {
 318            let selection = editor.selections.newest::<usize>(cx);
 319            if assist_range.contains(&selection.start) && assist_range.contains(&selection.end) {
 320                if undo {
 321                    self.finish_inline_assist(assist_id, true, cx);
 322                } else if matches!(assist.codegen.read(cx).status, CodegenStatus::Pending) {
 323                    self.dismiss_inline_assist(assist_id, cx);
 324                } else {
 325                    self.finish_inline_assist(assist_id, false, cx);
 326                }
 327
 328                return;
 329            }
 330        }
 331
 332        cx.propagate();
 333    }
 334
 335    fn handle_editor_event(
 336        &mut self,
 337        assist_id: InlineAssistId,
 338        editor: View<Editor>,
 339        event: &EditorEvent,
 340        cx: &mut WindowContext,
 341    ) {
 342        let Some(assist) = self.pending_assists.get(&assist_id) else {
 343            return;
 344        };
 345
 346        match event {
 347            EditorEvent::SelectionsChanged { local } if *local => {
 348                if let CodegenStatus::Idle = &assist.codegen.read(cx).status {
 349                    self.finish_inline_assist(assist_id, true, cx);
 350                }
 351            }
 352            EditorEvent::Saved => {
 353                if let CodegenStatus::Done = &assist.codegen.read(cx).status {
 354                    self.finish_inline_assist(assist_id, false, cx)
 355                }
 356            }
 357            EditorEvent::Edited { transaction_id }
 358                if matches!(
 359                    assist.codegen.read(cx).status,
 360                    CodegenStatus::Error(_) | CodegenStatus::Done
 361                ) =>
 362            {
 363                let buffer = editor.read(cx).buffer().read(cx);
 364                let edited_ranges =
 365                    buffer.edited_ranges_for_transaction::<usize>(*transaction_id, cx);
 366                let assist_range = assist.codegen.read(cx).range().to_offset(&buffer.read(cx));
 367                if edited_ranges
 368                    .iter()
 369                    .any(|range| range.overlaps(&assist_range))
 370                {
 371                    self.finish_inline_assist(assist_id, false, cx);
 372                }
 373            }
 374            _ => {}
 375        }
 376    }
 377
 378    fn finish_inline_assist(
 379        &mut self,
 380        assist_id: InlineAssistId,
 381        undo: bool,
 382        cx: &mut WindowContext,
 383    ) {
 384        self.dismiss_inline_assist(assist_id, cx);
 385
 386        if let Some(pending_assist) = self.pending_assists.remove(&assist_id) {
 387            if let hash_map::Entry::Occupied(mut entry) = self
 388                .pending_assist_ids_by_editor
 389                .entry(pending_assist.editor.clone())
 390            {
 391                entry.get_mut().retain(|id| *id != assist_id);
 392                if entry.get().is_empty() {
 393                    entry.remove();
 394                }
 395            }
 396
 397            if let Some(editor) = pending_assist.editor.upgrade() {
 398                self.update_editor_highlights(&editor, cx);
 399
 400                if undo {
 401                    pending_assist
 402                        .codegen
 403                        .update(cx, |codegen, cx| codegen.undo(cx));
 404                }
 405            }
 406        }
 407    }
 408
 409    fn dismiss_inline_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) -> bool {
 410        let Some(pending_assist) = self.pending_assists.get_mut(&assist_id) else {
 411            return false;
 412        };
 413        let Some(editor) = pending_assist.editor.upgrade() else {
 414            return false;
 415        };
 416        let Some(decorations) = pending_assist.editor_decorations.take() else {
 417            return false;
 418        };
 419
 420        editor.update(cx, |editor, cx| {
 421            let mut to_remove = decorations.removed_line_block_ids;
 422            to_remove.insert(decorations.prompt_block_id);
 423            to_remove.insert(decorations.end_block_id);
 424            editor.remove_blocks(to_remove, None, cx);
 425            if decorations
 426                .prompt_editor
 427                .focus_handle(cx)
 428                .contains_focused(cx)
 429            {
 430                editor.focus(cx);
 431            }
 432        });
 433
 434        self.update_editor_highlights(&editor, cx);
 435        true
 436    }
 437
 438    fn resize_inline_assist(
 439        &mut self,
 440        assist_id: InlineAssistId,
 441        height_in_lines: u8,
 442        cx: &mut WindowContext,
 443    ) {
 444        if let Some(pending_assist) = self.pending_assists.get_mut(&assist_id) {
 445            if let Some(editor) = pending_assist.editor.upgrade() {
 446                if let Some(decorations) = pending_assist.editor_decorations.as_ref() {
 447                    let gutter_dimensions =
 448                        decorations.prompt_editor.read(cx).gutter_dimensions.clone();
 449                    let mut new_blocks = HashMap::default();
 450                    new_blocks.insert(
 451                        decorations.prompt_block_id,
 452                        (
 453                            Some(height_in_lines),
 454                            build_inline_assist_editor_renderer(
 455                                &decorations.prompt_editor,
 456                                gutter_dimensions,
 457                            ),
 458                        ),
 459                    );
 460                    editor.update(cx, |editor, cx| {
 461                        editor
 462                            .display_map
 463                            .update(cx, |map, cx| map.replace_blocks(new_blocks, cx))
 464                    });
 465                }
 466            }
 467        }
 468    }
 469
 470    fn start_inline_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
 471        let pending_assist = if let Some(pending_assist) = self.pending_assists.get_mut(&assist_id)
 472        {
 473            pending_assist
 474        } else {
 475            return;
 476        };
 477
 478        pending_assist
 479            .codegen
 480            .update(cx, |codegen, cx| codegen.undo(cx));
 481
 482        let Some(user_prompt) = pending_assist
 483            .editor_decorations
 484            .as_ref()
 485            .map(|decorations| decorations.prompt_editor.read(cx).prompt(cx))
 486        else {
 487            return;
 488        };
 489
 490        let context = if pending_assist.include_context {
 491            pending_assist.workspace.as_ref().and_then(|workspace| {
 492                let workspace = workspace.upgrade()?.read(cx);
 493                let assistant_panel = workspace.panel::<AssistantPanel>(cx)?;
 494                assistant_panel.read(cx).active_context(cx)
 495            })
 496        } else {
 497            None
 498        };
 499
 500        let editor = if let Some(editor) = pending_assist.editor.upgrade() {
 501            editor
 502        } else {
 503            return;
 504        };
 505
 506        let project_name = pending_assist.workspace.as_ref().and_then(|workspace| {
 507            let workspace = workspace.upgrade()?;
 508            Some(
 509                workspace
 510                    .read(cx)
 511                    .project()
 512                    .read(cx)
 513                    .worktree_root_names(cx)
 514                    .collect::<Vec<&str>>()
 515                    .join("/"),
 516            )
 517        });
 518
 519        self.prompt_history.retain(|prompt| *prompt != user_prompt);
 520        self.prompt_history.push_back(user_prompt.clone());
 521        if self.prompt_history.len() > PROMPT_HISTORY_MAX_LEN {
 522            self.prompt_history.pop_front();
 523        }
 524
 525        let codegen = pending_assist.codegen.clone();
 526        let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
 527        let range = codegen.read(cx).range();
 528        let start = snapshot.point_to_buffer_offset(range.start);
 529        let end = snapshot.point_to_buffer_offset(range.end);
 530        let (buffer, range) = if let Some((start, end)) = start.zip(end) {
 531            let (start_buffer, start_buffer_offset) = start;
 532            let (end_buffer, end_buffer_offset) = end;
 533            if start_buffer.remote_id() == end_buffer.remote_id() {
 534                (start_buffer.clone(), start_buffer_offset..end_buffer_offset)
 535            } else {
 536                self.finish_inline_assist(assist_id, false, cx);
 537                return;
 538            }
 539        } else {
 540            self.finish_inline_assist(assist_id, false, cx);
 541            return;
 542        };
 543
 544        let language = buffer.language_at(range.start);
 545        let language_name = if let Some(language) = language.as_ref() {
 546            if Arc::ptr_eq(language, &language::PLAIN_TEXT) {
 547                None
 548            } else {
 549                Some(language.name())
 550            }
 551        } else {
 552            None
 553        };
 554
 555        // Higher Temperature increases the randomness of model outputs.
 556        // If Markdown or No Language is Known, increase the randomness for more creative output
 557        // If Code, decrease temperature to get more deterministic outputs
 558        let temperature = if let Some(language) = language_name.clone() {
 559            if language.as_ref() == "Markdown" {
 560                1.0
 561            } else {
 562                0.5
 563            }
 564        } else {
 565            1.0
 566        };
 567
 568        let prompt = cx.background_executor().spawn(async move {
 569            let language_name = language_name.as_deref();
 570            generate_content_prompt(user_prompt, language_name, buffer, range, project_name)
 571        });
 572
 573        let mut messages = Vec::new();
 574        if let Some(context) = context {
 575            let request = context.read(cx).to_completion_request(cx);
 576            messages = request.messages;
 577        }
 578        let model = CompletionProvider::global(cx).model();
 579
 580        cx.spawn(|mut cx| async move {
 581            let prompt = prompt.await?;
 582
 583            messages.push(LanguageModelRequestMessage {
 584                role: Role::User,
 585                content: prompt,
 586            });
 587
 588            let request = LanguageModelRequest {
 589                model,
 590                messages,
 591                stop: vec!["|END|>".to_string()],
 592                temperature,
 593            };
 594
 595            codegen.update(&mut cx, |codegen, cx| codegen.start(request, cx))?;
 596            anyhow::Ok(())
 597        })
 598        .detach_and_log_err(cx);
 599    }
 600
 601    fn stop_inline_assist(&mut self, assist_id: InlineAssistId, cx: &mut WindowContext) {
 602        let pending_assist = if let Some(pending_assist) = self.pending_assists.get_mut(&assist_id)
 603        {
 604            pending_assist
 605        } else {
 606            return;
 607        };
 608
 609        pending_assist
 610            .codegen
 611            .update(cx, |codegen, cx| codegen.stop(cx));
 612    }
 613
 614    fn update_editor_highlights(&self, editor: &View<Editor>, cx: &mut WindowContext) {
 615        let mut gutter_pending_ranges = Vec::new();
 616        let mut gutter_transformed_ranges = Vec::new();
 617        let mut foreground_ranges = Vec::new();
 618        let mut inserted_row_ranges = Vec::new();
 619        let empty_assist_ids = Vec::new();
 620        let assist_ids = self
 621            .pending_assist_ids_by_editor
 622            .get(&editor.downgrade())
 623            .unwrap_or(&empty_assist_ids);
 624
 625        for assist_id in assist_ids {
 626            if let Some(pending_assist) = self.pending_assists.get(assist_id) {
 627                let codegen = pending_assist.codegen.read(cx);
 628                foreground_ranges.extend(codegen.last_equal_ranges().iter().cloned());
 629
 630                if codegen.edit_position != codegen.range().end {
 631                    gutter_pending_ranges.push(codegen.edit_position..codegen.range().end);
 632                }
 633
 634                if codegen.range().start != codegen.edit_position {
 635                    gutter_transformed_ranges.push(codegen.range().start..codegen.edit_position);
 636                }
 637
 638                if pending_assist.editor_decorations.is_some() {
 639                    inserted_row_ranges.extend(codegen.diff.inserted_row_ranges.iter().cloned());
 640                }
 641            }
 642        }
 643
 644        let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
 645        merge_ranges(&mut foreground_ranges, &snapshot);
 646        merge_ranges(&mut gutter_pending_ranges, &snapshot);
 647        merge_ranges(&mut gutter_transformed_ranges, &snapshot);
 648        editor.update(cx, |editor, cx| {
 649            enum GutterPendingRange {}
 650            if gutter_pending_ranges.is_empty() {
 651                editor.clear_gutter_highlights::<GutterPendingRange>(cx);
 652            } else {
 653                editor.highlight_gutter::<GutterPendingRange>(
 654                    &gutter_pending_ranges,
 655                    |cx| cx.theme().status().info_background,
 656                    cx,
 657                )
 658            }
 659
 660            enum GutterTransformedRange {}
 661            if gutter_transformed_ranges.is_empty() {
 662                editor.clear_gutter_highlights::<GutterTransformedRange>(cx);
 663            } else {
 664                editor.highlight_gutter::<GutterTransformedRange>(
 665                    &gutter_transformed_ranges,
 666                    |cx| cx.theme().status().info,
 667                    cx,
 668                )
 669            }
 670
 671            if foreground_ranges.is_empty() {
 672                editor.clear_highlights::<PendingInlineAssist>(cx);
 673            } else {
 674                editor.highlight_text::<PendingInlineAssist>(
 675                    foreground_ranges,
 676                    HighlightStyle {
 677                        fade_out: Some(0.6),
 678                        ..Default::default()
 679                    },
 680                    cx,
 681                );
 682            }
 683
 684            editor.clear_row_highlights::<PendingInlineAssist>();
 685            for row_range in inserted_row_ranges {
 686                editor.highlight_rows::<PendingInlineAssist>(
 687                    row_range,
 688                    Some(cx.theme().status().info_background),
 689                    false,
 690                    cx,
 691                );
 692            }
 693        });
 694    }
 695
 696    fn update_editor_blocks(
 697        &mut self,
 698        editor: &View<Editor>,
 699        assist_id: InlineAssistId,
 700        cx: &mut WindowContext,
 701    ) {
 702        let Some(pending_assist) = self.pending_assists.get_mut(&assist_id) else {
 703            return;
 704        };
 705        let Some(decorations) = pending_assist.editor_decorations.as_mut() else {
 706            return;
 707        };
 708
 709        let codegen = pending_assist.codegen.read(cx);
 710        let old_snapshot = codegen.snapshot.clone();
 711        let old_buffer = codegen.old_buffer.clone();
 712        let deleted_row_ranges = codegen.diff.deleted_row_ranges.clone();
 713
 714        editor.update(cx, |editor, cx| {
 715            let old_blocks = mem::take(&mut decorations.removed_line_block_ids);
 716            editor.remove_blocks(old_blocks, None, cx);
 717
 718            let mut new_blocks = Vec::new();
 719            for (new_row, old_row_range) in deleted_row_ranges {
 720                let (_, buffer_start) = old_snapshot
 721                    .point_to_buffer_offset(Point::new(*old_row_range.start(), 0))
 722                    .unwrap();
 723                let (_, buffer_end) = old_snapshot
 724                    .point_to_buffer_offset(Point::new(
 725                        *old_row_range.end(),
 726                        old_snapshot.line_len(MultiBufferRow(*old_row_range.end())),
 727                    ))
 728                    .unwrap();
 729
 730                let deleted_lines_editor = cx.new_view(|cx| {
 731                    let multi_buffer = cx.new_model(|_| {
 732                        MultiBuffer::without_headers(0, language::Capability::ReadOnly)
 733                    });
 734                    multi_buffer.update(cx, |multi_buffer, cx| {
 735                        multi_buffer.push_excerpts(
 736                            old_buffer.clone(),
 737                            Some(ExcerptRange {
 738                                context: buffer_start..buffer_end,
 739                                primary: None,
 740                            }),
 741                            cx,
 742                        );
 743                    });
 744
 745                    enum DeletedLines {}
 746                    let mut editor = Editor::for_multibuffer(multi_buffer, None, true, cx);
 747                    editor.set_soft_wrap_mode(language::language_settings::SoftWrap::None, cx);
 748                    editor.set_show_wrap_guides(false, cx);
 749                    editor.set_show_gutter(false, cx);
 750                    editor.scroll_manager.set_forbid_vertical_scroll(true);
 751                    editor.set_read_only(true);
 752                    editor.highlight_rows::<DeletedLines>(
 753                        Anchor::min()..=Anchor::max(),
 754                        Some(cx.theme().status().deleted_background),
 755                        false,
 756                        cx,
 757                    );
 758                    editor
 759                });
 760
 761                let height = deleted_lines_editor
 762                    .update(cx, |editor, cx| editor.max_point(cx).row().0 as u8 + 1);
 763                new_blocks.push(BlockProperties {
 764                    position: new_row,
 765                    height,
 766                    style: BlockStyle::Flex,
 767                    render: Box::new(move |cx| {
 768                        div()
 769                            .bg(cx.theme().status().deleted_background)
 770                            .size_full()
 771                            .pl(cx.gutter_dimensions.full_width())
 772                            .child(deleted_lines_editor.clone())
 773                            .into_any_element()
 774                    }),
 775                    disposition: BlockDisposition::Above,
 776                });
 777            }
 778
 779            decorations.removed_line_block_ids = editor
 780                .insert_blocks(new_blocks, None, cx)
 781                .into_iter()
 782                .collect();
 783        })
 784    }
 785}
 786
 787fn build_inline_assist_editor_renderer(
 788    editor: &View<InlineAssistEditor>,
 789    gutter_dimensions: Arc<Mutex<GutterDimensions>>,
 790) -> RenderBlock {
 791    let editor = editor.clone();
 792    Box::new(move |cx: &mut BlockContext| {
 793        *gutter_dimensions.lock() = *cx.gutter_dimensions;
 794        editor.clone().into_any_element()
 795    })
 796}
 797
 798#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
 799struct InlineAssistId(usize);
 800
 801impl InlineAssistId {
 802    fn post_inc(&mut self) -> InlineAssistId {
 803        let id = *self;
 804        self.0 += 1;
 805        id
 806    }
 807}
 808
 809enum InlineAssistEditorEvent {
 810    StartRequested,
 811    StopRequested,
 812    ConfirmRequested,
 813    CancelRequested,
 814    DismissRequested,
 815    Resized { height_in_lines: u8 },
 816}
 817
 818struct InlineAssistEditor {
 819    id: InlineAssistId,
 820    height_in_lines: u8,
 821    prompt_editor: View<Editor>,
 822    edited_since_done: bool,
 823    gutter_dimensions: Arc<Mutex<GutterDimensions>>,
 824    prompt_history: VecDeque<String>,
 825    prompt_history_ix: Option<usize>,
 826    pending_prompt: String,
 827    codegen: Model<Codegen>,
 828    workspace: Option<WeakView<Workspace>>,
 829    _subscriptions: Vec<Subscription>,
 830}
 831
 832impl EventEmitter<InlineAssistEditorEvent> for InlineAssistEditor {}
 833
 834impl Render for InlineAssistEditor {
 835    fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
 836        let gutter_dimensions = *self.gutter_dimensions.lock();
 837
 838        let buttons = match &self.codegen.read(cx).status {
 839            CodegenStatus::Idle => {
 840                vec![
 841                    IconButton::new("start", IconName::Sparkle)
 842                        .icon_color(Color::Muted)
 843                        .size(ButtonSize::None)
 844                        .icon_size(IconSize::XSmall)
 845                        .tooltip(|cx| Tooltip::for_action("Transform", &menu::Confirm, cx))
 846                        .on_click(
 847                            cx.listener(|_, _, cx| {
 848                                cx.emit(InlineAssistEditorEvent::StartRequested)
 849                            }),
 850                        ),
 851                    IconButton::new("cancel", IconName::Close)
 852                        .icon_color(Color::Muted)
 853                        .size(ButtonSize::None)
 854                        .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
 855                        .on_click(cx.listener(|_, _, cx| {
 856                            cx.emit(InlineAssistEditorEvent::CancelRequested)
 857                        })),
 858                ]
 859            }
 860            CodegenStatus::Pending => {
 861                vec![
 862                    IconButton::new("stop", IconName::Stop)
 863                        .icon_color(Color::Error)
 864                        .size(ButtonSize::None)
 865                        .icon_size(IconSize::XSmall)
 866                        .tooltip(|cx| {
 867                            Tooltip::with_meta(
 868                                "Interrupt Transformation",
 869                                Some(&menu::Cancel),
 870                                "Changes won't be discarded",
 871                                cx,
 872                            )
 873                        })
 874                        .on_click(
 875                            cx.listener(|_, _, cx| cx.emit(InlineAssistEditorEvent::StopRequested)),
 876                        ),
 877                    IconButton::new("cancel", IconName::Close)
 878                        .icon_color(Color::Muted)
 879                        .size(ButtonSize::None)
 880                        .tooltip(|cx| Tooltip::text("Cancel Assist", cx))
 881                        .on_click(cx.listener(|_, _, cx| {
 882                            cx.emit(InlineAssistEditorEvent::CancelRequested)
 883                        })),
 884                ]
 885            }
 886            CodegenStatus::Error(_) | CodegenStatus::Done => {
 887                vec![
 888                    if self.edited_since_done {
 889                        IconButton::new("restart", IconName::RotateCw)
 890                            .icon_color(Color::Info)
 891                            .icon_size(IconSize::XSmall)
 892                            .size(ButtonSize::None)
 893                            .tooltip(|cx| {
 894                                Tooltip::with_meta(
 895                                    "Restart Transformation",
 896                                    Some(&menu::Confirm),
 897                                    "Changes will be discarded",
 898                                    cx,
 899                                )
 900                            })
 901                            .on_click(cx.listener(|_, _, cx| {
 902                                cx.emit(InlineAssistEditorEvent::StartRequested);
 903                            }))
 904                    } else {
 905                        IconButton::new("confirm", IconName::Check)
 906                            .icon_color(Color::Info)
 907                            .size(ButtonSize::None)
 908                            .tooltip(|cx| Tooltip::for_action("Confirm Assist", &menu::Confirm, cx))
 909                            .on_click(cx.listener(|_, _, cx| {
 910                                cx.emit(InlineAssistEditorEvent::ConfirmRequested);
 911                            }))
 912                    },
 913                    IconButton::new("cancel", IconName::Close)
 914                        .icon_color(Color::Muted)
 915                        .size(ButtonSize::None)
 916                        .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
 917                        .on_click(cx.listener(|_, _, cx| {
 918                            cx.emit(InlineAssistEditorEvent::CancelRequested)
 919                        })),
 920                ]
 921            }
 922        };
 923
 924        v_flex().h_full().w_full().justify_end().child(
 925            h_flex()
 926                .bg(cx.theme().colors().editor_background)
 927                .border_y_1()
 928                .border_color(cx.theme().status().info_border)
 929                .py_1p5()
 930                .w_full()
 931                .on_action(cx.listener(Self::confirm))
 932                .on_action(cx.listener(Self::cancel))
 933                .on_action(cx.listener(Self::move_up))
 934                .on_action(cx.listener(Self::move_down))
 935                .child(
 936                    h_flex()
 937                        .w(gutter_dimensions.full_width() + (gutter_dimensions.margin / 2.0))
 938                        // .pr(gutter_dimensions.fold_area_width())
 939                        .justify_center()
 940                        .gap_2()
 941                        .children(self.workspace.clone().map(|workspace| {
 942                            IconButton::new("context", IconName::Context)
 943                                .size(ButtonSize::None)
 944                                .icon_size(IconSize::XSmall)
 945                                .icon_color(Color::Muted)
 946                                .on_click({
 947                                    let workspace = workspace.clone();
 948                                    cx.listener(move |_, _, cx| {
 949                                        workspace
 950                                            .update(cx, |workspace, cx| {
 951                                                workspace.focus_panel::<AssistantPanel>(cx);
 952                                            })
 953                                            .ok();
 954                                    })
 955                                })
 956                                .tooltip(move |cx| {
 957                                    let token_count = workspace.upgrade().and_then(|workspace| {
 958                                        let panel =
 959                                            workspace.read(cx).panel::<AssistantPanel>(cx)?;
 960                                        let context = panel.read(cx).active_context(cx)?;
 961                                        context.read(cx).token_count()
 962                                    });
 963                                    if let Some(token_count) = token_count {
 964                                        Tooltip::with_meta(
 965                                            format!(
 966                                                "{} Additional Context Tokens from Assistant",
 967                                                token_count
 968                                            ),
 969                                            Some(&crate::ToggleFocus),
 970                                            "Click to open…",
 971                                            cx,
 972                                        )
 973                                    } else {
 974                                        Tooltip::for_action(
 975                                            "Toggle Assistant Panel",
 976                                            &crate::ToggleFocus,
 977                                            cx,
 978                                        )
 979                                    }
 980                                })
 981                        }))
 982                        .children(
 983                            if let CodegenStatus::Error(error) = &self.codegen.read(cx).status {
 984                                let error_message = SharedString::from(error.to_string());
 985                                Some(
 986                                    div()
 987                                        .id("error")
 988                                        .tooltip(move |cx| Tooltip::text(error_message.clone(), cx))
 989                                        .child(
 990                                            Icon::new(IconName::XCircle)
 991                                                .size(IconSize::Small)
 992                                                .color(Color::Error),
 993                                        ),
 994                                )
 995                            } else {
 996                                None
 997                            },
 998                        ),
 999                )
1000                .child(div().flex_1().child(self.render_prompt_editor(cx)))
1001                .child(h_flex().gap_2().pr_4().children(buttons)),
1002        )
1003    }
1004}
1005
1006impl FocusableView for InlineAssistEditor {
1007    fn focus_handle(&self, cx: &AppContext) -> FocusHandle {
1008        self.prompt_editor.focus_handle(cx)
1009    }
1010}
1011
1012impl InlineAssistEditor {
1013    const MAX_LINES: u8 = 8;
1014
1015    #[allow(clippy::too_many_arguments)]
1016    fn new(
1017        id: InlineAssistId,
1018        gutter_dimensions: Arc<Mutex<GutterDimensions>>,
1019        prompt_history: VecDeque<String>,
1020        codegen: Model<Codegen>,
1021        workspace: Option<WeakView<Workspace>>,
1022        cx: &mut ViewContext<Self>,
1023    ) -> Self {
1024        let prompt_editor = cx.new_view(|cx| {
1025            let mut editor = Editor::auto_height(Self::MAX_LINES as usize, cx);
1026            editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
1027            editor.set_placeholder_text("Add a prompt…", cx);
1028            editor
1029        });
1030        cx.focus_view(&prompt_editor);
1031
1032        let subscriptions = vec![
1033            cx.observe(&codegen, Self::handle_codegen_changed),
1034            cx.observe(&prompt_editor, Self::handle_prompt_editor_changed),
1035            cx.subscribe(&prompt_editor, Self::handle_prompt_editor_events),
1036        ];
1037
1038        let mut this = Self {
1039            id,
1040            height_in_lines: 1,
1041            prompt_editor,
1042            edited_since_done: false,
1043            gutter_dimensions,
1044            prompt_history,
1045            prompt_history_ix: None,
1046            pending_prompt: String::new(),
1047            codegen,
1048            workspace,
1049            _subscriptions: subscriptions,
1050        };
1051        this.count_lines(cx);
1052        this
1053    }
1054
1055    fn prompt(&self, cx: &AppContext) -> String {
1056        self.prompt_editor.read(cx).text(cx)
1057    }
1058
1059    fn count_lines(&mut self, cx: &mut ViewContext<Self>) {
1060        let height_in_lines = cmp::max(
1061            2, // Make the editor at least two lines tall, to account for padding and buttons.
1062            cmp::min(
1063                self.prompt_editor
1064                    .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1),
1065                Self::MAX_LINES as u32,
1066            ),
1067        ) as u8;
1068
1069        if height_in_lines != self.height_in_lines {
1070            self.height_in_lines = height_in_lines;
1071            cx.emit(InlineAssistEditorEvent::Resized { height_in_lines });
1072        }
1073    }
1074
1075    fn handle_prompt_editor_changed(&mut self, _: View<Editor>, cx: &mut ViewContext<Self>) {
1076        self.count_lines(cx);
1077    }
1078
1079    fn handle_prompt_editor_events(
1080        &mut self,
1081        _: View<Editor>,
1082        event: &EditorEvent,
1083        cx: &mut ViewContext<Self>,
1084    ) {
1085        match event {
1086            EditorEvent::Edited { .. } => {
1087                let prompt = self.prompt_editor.read(cx).text(cx);
1088                if self
1089                    .prompt_history_ix
1090                    .map_or(true, |ix| self.prompt_history[ix] != prompt)
1091                {
1092                    self.prompt_history_ix.take();
1093                    self.pending_prompt = prompt;
1094                }
1095
1096                self.edited_since_done = true;
1097                cx.notify();
1098            }
1099            _ => {}
1100        }
1101    }
1102
1103    fn handle_codegen_changed(&mut self, _: Model<Codegen>, cx: &mut ViewContext<Self>) {
1104        match &self.codegen.read(cx).status {
1105            CodegenStatus::Idle => {
1106                self.prompt_editor
1107                    .update(cx, |editor, _| editor.set_read_only(false));
1108            }
1109            CodegenStatus::Pending => {
1110                self.prompt_editor
1111                    .update(cx, |editor, _| editor.set_read_only(true));
1112            }
1113            CodegenStatus::Done | CodegenStatus::Error(_) => {
1114                self.edited_since_done = false;
1115                self.prompt_editor
1116                    .update(cx, |editor, _| editor.set_read_only(false));
1117            }
1118        }
1119    }
1120
1121    fn cancel(&mut self, _: &editor::actions::Cancel, cx: &mut ViewContext<Self>) {
1122        match &self.codegen.read(cx).status {
1123            CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {
1124                cx.emit(InlineAssistEditorEvent::CancelRequested);
1125            }
1126            CodegenStatus::Pending => {
1127                cx.emit(InlineAssistEditorEvent::StopRequested);
1128            }
1129        }
1130    }
1131
1132    fn confirm(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
1133        match &self.codegen.read(cx).status {
1134            CodegenStatus::Idle => {
1135                cx.emit(InlineAssistEditorEvent::StartRequested);
1136            }
1137            CodegenStatus::Pending => {
1138                cx.emit(InlineAssistEditorEvent::DismissRequested);
1139            }
1140            CodegenStatus::Done | CodegenStatus::Error(_) => {
1141                if self.edited_since_done {
1142                    cx.emit(InlineAssistEditorEvent::StartRequested);
1143                } else {
1144                    cx.emit(InlineAssistEditorEvent::ConfirmRequested);
1145                }
1146            }
1147        }
1148    }
1149
1150    fn move_up(&mut self, _: &MoveUp, cx: &mut ViewContext<Self>) {
1151        if let Some(ix) = self.prompt_history_ix {
1152            if ix > 0 {
1153                self.prompt_history_ix = Some(ix - 1);
1154                let prompt = self.prompt_history[ix - 1].as_str();
1155                self.prompt_editor.update(cx, |editor, cx| {
1156                    editor.set_text(prompt, cx);
1157                    editor.move_to_beginning(&Default::default(), cx);
1158                });
1159            }
1160        } else if !self.prompt_history.is_empty() {
1161            self.prompt_history_ix = Some(self.prompt_history.len() - 1);
1162            let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();
1163            self.prompt_editor.update(cx, |editor, cx| {
1164                editor.set_text(prompt, cx);
1165                editor.move_to_beginning(&Default::default(), cx);
1166            });
1167        }
1168    }
1169
1170    fn move_down(&mut self, _: &MoveDown, cx: &mut ViewContext<Self>) {
1171        if let Some(ix) = self.prompt_history_ix {
1172            if ix < self.prompt_history.len() - 1 {
1173                self.prompt_history_ix = Some(ix + 1);
1174                let prompt = self.prompt_history[ix + 1].as_str();
1175                self.prompt_editor.update(cx, |editor, cx| {
1176                    editor.set_text(prompt, cx);
1177                    editor.move_to_end(&Default::default(), cx)
1178                });
1179            } else {
1180                self.prompt_history_ix = None;
1181                let prompt = self.pending_prompt.as_str();
1182                self.prompt_editor.update(cx, |editor, cx| {
1183                    editor.set_text(prompt, cx);
1184                    editor.move_to_end(&Default::default(), cx)
1185                });
1186            }
1187        }
1188    }
1189
1190    fn render_prompt_editor(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
1191        let settings = ThemeSettings::get_global(cx);
1192        let text_style = TextStyle {
1193            color: if self.prompt_editor.read(cx).read_only(cx) {
1194                cx.theme().colors().text_disabled
1195            } else {
1196                cx.theme().colors().text
1197            },
1198            font_family: settings.ui_font.family.clone(),
1199            font_features: settings.ui_font.features.clone(),
1200            font_size: rems(0.875).into(),
1201            font_weight: FontWeight::NORMAL,
1202            font_style: FontStyle::Normal,
1203            line_height: relative(1.3),
1204            background_color: None,
1205            underline: None,
1206            strikethrough: None,
1207            white_space: WhiteSpace::Normal,
1208        };
1209        EditorElement::new(
1210            &self.prompt_editor,
1211            EditorStyle {
1212                background: cx.theme().colors().editor_background,
1213                local_player: cx.theme().players().local(),
1214                text: text_style,
1215                ..Default::default()
1216            },
1217        )
1218    }
1219}
1220
1221struct PendingInlineAssist {
1222    editor: WeakView<Editor>,
1223    editor_decorations: Option<PendingInlineAssistDecorations>,
1224    codegen: Model<Codegen>,
1225    _subscriptions: Vec<Subscription>,
1226    workspace: Option<WeakView<Workspace>>,
1227    include_context: bool,
1228}
1229
1230struct PendingInlineAssistDecorations {
1231    prompt_block_id: BlockId,
1232    prompt_editor: View<InlineAssistEditor>,
1233    removed_line_block_ids: HashSet<BlockId>,
1234    end_block_id: BlockId,
1235}
1236
1237#[derive(Debug)]
1238pub enum CodegenEvent {
1239    Finished,
1240    Undone,
1241}
1242
1243#[derive(Clone)]
1244pub enum CodegenKind {
1245    Transform { range: Range<Anchor> },
1246    Generate { position: Anchor },
1247}
1248
1249impl CodegenKind {
1250    fn range(&self, snapshot: &MultiBufferSnapshot) -> Range<Anchor> {
1251        match self {
1252            CodegenKind::Transform { range } => range.clone(),
1253            CodegenKind::Generate { position } => position.bias_left(snapshot)..*position,
1254        }
1255    }
1256}
1257
1258pub struct Codegen {
1259    buffer: Model<MultiBuffer>,
1260    old_buffer: Model<Buffer>,
1261    snapshot: MultiBufferSnapshot,
1262    kind: CodegenKind,
1263    edit_position: Anchor,
1264    last_equal_ranges: Vec<Range<Anchor>>,
1265    transaction_id: Option<TransactionId>,
1266    status: CodegenStatus,
1267    generation: Task<()>,
1268    diff: Diff,
1269    telemetry: Option<Arc<Telemetry>>,
1270    _subscription: gpui::Subscription,
1271}
1272
1273enum CodegenStatus {
1274    Idle,
1275    Pending,
1276    Done,
1277    Error(anyhow::Error),
1278}
1279
1280#[derive(Default)]
1281struct Diff {
1282    task: Option<Task<()>>,
1283    should_update: bool,
1284    deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)>,
1285    inserted_row_ranges: Vec<RangeInclusive<Anchor>>,
1286}
1287
1288impl EventEmitter<CodegenEvent> for Codegen {}
1289
1290impl Codegen {
1291    pub fn new(
1292        buffer: Model<MultiBuffer>,
1293        kind: CodegenKind,
1294        telemetry: Option<Arc<Telemetry>>,
1295        cx: &mut ModelContext<Self>,
1296    ) -> Self {
1297        let snapshot = buffer.read(cx).snapshot(cx);
1298
1299        let (old_buffer, _, _) = buffer
1300            .read(cx)
1301            .range_to_buffer_ranges(kind.range(&snapshot), cx)
1302            .pop()
1303            .unwrap();
1304        let old_buffer = cx.new_model(|cx| {
1305            let old_buffer = old_buffer.read(cx);
1306            let text = old_buffer.as_rope().clone();
1307            let line_ending = old_buffer.line_ending();
1308            let language = old_buffer.language().cloned();
1309            let language_registry = old_buffer.language_registry();
1310
1311            let mut buffer = Buffer::local_normalized(text, line_ending, cx);
1312            buffer.set_language(language, cx);
1313            if let Some(language_registry) = language_registry {
1314                buffer.set_language_registry(language_registry)
1315            }
1316            buffer
1317        });
1318
1319        Self {
1320            buffer: buffer.clone(),
1321            old_buffer,
1322            edit_position: kind.range(&snapshot).start,
1323            snapshot,
1324            kind,
1325            last_equal_ranges: Default::default(),
1326            transaction_id: Default::default(),
1327            status: CodegenStatus::Idle,
1328            generation: Task::ready(()),
1329            diff: Diff::default(),
1330            telemetry,
1331            _subscription: cx.subscribe(&buffer, Self::handle_buffer_event),
1332        }
1333    }
1334
1335    fn handle_buffer_event(
1336        &mut self,
1337        _buffer: Model<MultiBuffer>,
1338        event: &multi_buffer::Event,
1339        cx: &mut ModelContext<Self>,
1340    ) {
1341        if let multi_buffer::Event::TransactionUndone { transaction_id } = event {
1342            if self.transaction_id == Some(*transaction_id) {
1343                self.transaction_id = None;
1344                self.generation = Task::ready(());
1345                cx.emit(CodegenEvent::Undone);
1346            }
1347        }
1348    }
1349
1350    pub fn range(&self) -> Range<Anchor> {
1351        self.kind.range(&self.snapshot)
1352    }
1353
1354    pub fn last_equal_ranges(&self) -> &[Range<Anchor>] {
1355        &self.last_equal_ranges
1356    }
1357
1358    pub fn start(&mut self, prompt: LanguageModelRequest, cx: &mut ModelContext<Self>) {
1359        let range = self.range();
1360        let snapshot = self.snapshot.clone();
1361        let selected_text = snapshot
1362            .text_for_range(range.start..range.end)
1363            .collect::<Rope>();
1364
1365        let selection_start = range.start.to_point(&snapshot);
1366        let suggested_line_indent = snapshot
1367            .suggested_indents(selection_start.row..selection_start.row + 1, cx)
1368            .into_values()
1369            .next()
1370            .unwrap_or_else(|| snapshot.indent_size_for_line(MultiBufferRow(selection_start.row)));
1371
1372        let model_telemetry_id = prompt.model.telemetry_id();
1373        let response = CompletionProvider::global(cx).complete(prompt);
1374        let telemetry = self.telemetry.clone();
1375        self.edit_position = range.start;
1376        self.diff = Diff::default();
1377        self.status = CodegenStatus::Pending;
1378        self.generation = cx.spawn(|this, mut cx| {
1379            async move {
1380                let generate = async {
1381                    let mut edit_start = range.start.to_offset(&snapshot);
1382
1383                    let (mut hunks_tx, mut hunks_rx) = mpsc::channel(1);
1384                    let diff: Task<anyhow::Result<()>> =
1385                        cx.background_executor().spawn(async move {
1386                            let mut response_latency = None;
1387                            let request_start = Instant::now();
1388                            let diff = async {
1389                                let chunks = strip_invalid_spans_from_codeblock(response.await?);
1390                                futures::pin_mut!(chunks);
1391                                let mut diff = StreamingDiff::new(selected_text.to_string());
1392
1393                                let mut new_text = String::new();
1394                                let mut base_indent = None;
1395                                let mut line_indent = None;
1396                                let mut first_line = true;
1397
1398                                while let Some(chunk) = chunks.next().await {
1399                                    if response_latency.is_none() {
1400                                        response_latency = Some(request_start.elapsed());
1401                                    }
1402                                    let chunk = chunk?;
1403
1404                                    let mut lines = chunk.split('\n').peekable();
1405                                    while let Some(line) = lines.next() {
1406                                        new_text.push_str(line);
1407                                        if line_indent.is_none() {
1408                                            if let Some(non_whitespace_ch_ix) =
1409                                                new_text.find(|ch: char| !ch.is_whitespace())
1410                                            {
1411                                                line_indent = Some(non_whitespace_ch_ix);
1412                                                base_indent = base_indent.or(line_indent);
1413
1414                                                let line_indent = line_indent.unwrap();
1415                                                let base_indent = base_indent.unwrap();
1416                                                let indent_delta =
1417                                                    line_indent as i32 - base_indent as i32;
1418                                                let mut corrected_indent_len = cmp::max(
1419                                                    0,
1420                                                    suggested_line_indent.len as i32 + indent_delta,
1421                                                )
1422                                                    as usize;
1423                                                if first_line {
1424                                                    corrected_indent_len = corrected_indent_len
1425                                                        .saturating_sub(
1426                                                            selection_start.column as usize,
1427                                                        );
1428                                                }
1429
1430                                                let indent_char = suggested_line_indent.char();
1431                                                let mut indent_buffer = [0; 4];
1432                                                let indent_str =
1433                                                    indent_char.encode_utf8(&mut indent_buffer);
1434                                                new_text.replace_range(
1435                                                    ..line_indent,
1436                                                    &indent_str.repeat(corrected_indent_len),
1437                                                );
1438                                            }
1439                                        }
1440
1441                                        if line_indent.is_some() {
1442                                            hunks_tx.send(diff.push_new(&new_text)).await?;
1443                                            new_text.clear();
1444                                        }
1445
1446                                        if lines.peek().is_some() {
1447                                            hunks_tx.send(diff.push_new("\n")).await?;
1448                                            line_indent = None;
1449                                            first_line = false;
1450                                        }
1451                                    }
1452                                }
1453                                hunks_tx.send(diff.push_new(&new_text)).await?;
1454                                hunks_tx.send(diff.finish()).await?;
1455
1456                                anyhow::Ok(())
1457                            };
1458
1459                            let result = diff.await;
1460
1461                            let error_message =
1462                                result.as_ref().err().map(|error| error.to_string());
1463                            if let Some(telemetry) = telemetry {
1464                                telemetry.report_assistant_event(
1465                                    None,
1466                                    telemetry_events::AssistantKind::Inline,
1467                                    model_telemetry_id,
1468                                    response_latency,
1469                                    error_message,
1470                                );
1471                            }
1472
1473                            result?;
1474                            Ok(())
1475                        });
1476
1477                    while let Some(hunks) = hunks_rx.next().await {
1478                        this.update(&mut cx, |this, cx| {
1479                            this.last_equal_ranges.clear();
1480
1481                            let transaction = this.buffer.update(cx, |buffer, cx| {
1482                                // Avoid grouping assistant edits with user edits.
1483                                buffer.finalize_last_transaction(cx);
1484
1485                                buffer.start_transaction(cx);
1486                                buffer.edit(
1487                                    hunks.into_iter().filter_map(|hunk| match hunk {
1488                                        Hunk::Insert { text } => {
1489                                            let edit_start = snapshot.anchor_after(edit_start);
1490                                            Some((edit_start..edit_start, text))
1491                                        }
1492                                        Hunk::Remove { len } => {
1493                                            let edit_end = edit_start + len;
1494                                            let edit_range = snapshot.anchor_after(edit_start)
1495                                                ..snapshot.anchor_before(edit_end);
1496                                            edit_start = edit_end;
1497                                            Some((edit_range, String::new()))
1498                                        }
1499                                        Hunk::Keep { len } => {
1500                                            let edit_end = edit_start + len;
1501                                            let edit_range = snapshot.anchor_after(edit_start)
1502                                                ..snapshot.anchor_before(edit_end);
1503                                            edit_start = edit_end;
1504                                            this.last_equal_ranges.push(edit_range);
1505                                            None
1506                                        }
1507                                    }),
1508                                    None,
1509                                    cx,
1510                                );
1511                                this.edit_position = snapshot.anchor_after(edit_start);
1512
1513                                buffer.end_transaction(cx)
1514                            });
1515
1516                            if let Some(transaction) = transaction {
1517                                if let Some(first_transaction) = this.transaction_id {
1518                                    // Group all assistant edits into the first transaction.
1519                                    this.buffer.update(cx, |buffer, cx| {
1520                                        buffer.merge_transactions(
1521                                            transaction,
1522                                            first_transaction,
1523                                            cx,
1524                                        )
1525                                    });
1526                                } else {
1527                                    this.transaction_id = Some(transaction);
1528                                    this.buffer.update(cx, |buffer, cx| {
1529                                        buffer.finalize_last_transaction(cx)
1530                                    });
1531                                }
1532                            }
1533
1534                            this.update_diff(cx);
1535                            cx.notify();
1536                        })?;
1537                    }
1538
1539                    diff.await?;
1540
1541                    anyhow::Ok(())
1542                };
1543
1544                let result = generate.await;
1545                this.update(&mut cx, |this, cx| {
1546                    this.last_equal_ranges.clear();
1547                    if let Err(error) = result {
1548                        this.status = CodegenStatus::Error(error);
1549                    } else {
1550                        this.status = CodegenStatus::Done;
1551                    }
1552                    cx.emit(CodegenEvent::Finished);
1553                    cx.notify();
1554                })
1555                .ok();
1556            }
1557        });
1558        cx.notify();
1559    }
1560
1561    pub fn stop(&mut self, cx: &mut ModelContext<Self>) {
1562        self.last_equal_ranges.clear();
1563        self.status = CodegenStatus::Done;
1564        self.generation = Task::ready(());
1565        cx.emit(CodegenEvent::Finished);
1566        cx.notify();
1567    }
1568
1569    pub fn undo(&mut self, cx: &mut ModelContext<Self>) {
1570        if let Some(transaction_id) = self.transaction_id.take() {
1571            self.buffer
1572                .update(cx, |buffer, cx| buffer.undo_transaction(transaction_id, cx));
1573        }
1574    }
1575
1576    fn update_diff(&mut self, cx: &mut ModelContext<Self>) {
1577        if self.diff.task.is_some() {
1578            self.diff.should_update = true;
1579        } else {
1580            self.diff.should_update = false;
1581
1582            let old_snapshot = self.snapshot.clone();
1583            let old_range = self.range().to_point(&old_snapshot);
1584            let new_snapshot = self.buffer.read(cx).snapshot(cx);
1585            let new_range = self.range().to_point(&new_snapshot);
1586
1587            self.diff.task = Some(cx.spawn(|this, mut cx| async move {
1588                let (deleted_row_ranges, inserted_row_ranges) = cx
1589                    .background_executor()
1590                    .spawn(async move {
1591                        let old_text = old_snapshot
1592                            .text_for_range(
1593                                Point::new(old_range.start.row, 0)
1594                                    ..Point::new(
1595                                        old_range.end.row,
1596                                        old_snapshot.line_len(MultiBufferRow(old_range.end.row)),
1597                                    ),
1598                            )
1599                            .collect::<String>();
1600                        let new_text = new_snapshot
1601                            .text_for_range(
1602                                Point::new(new_range.start.row, 0)
1603                                    ..Point::new(
1604                                        new_range.end.row,
1605                                        new_snapshot.line_len(MultiBufferRow(new_range.end.row)),
1606                                    ),
1607                            )
1608                            .collect::<String>();
1609
1610                        let mut old_row = old_range.start.row;
1611                        let mut new_row = new_range.start.row;
1612                        let diff = TextDiff::from_lines(old_text.as_str(), new_text.as_str());
1613
1614                        let mut deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)> = Vec::new();
1615                        let mut inserted_row_ranges = Vec::new();
1616                        for change in diff.iter_all_changes() {
1617                            let line_count = change.value().lines().count() as u32;
1618                            match change.tag() {
1619                                similar::ChangeTag::Equal => {
1620                                    old_row += line_count;
1621                                    new_row += line_count;
1622                                }
1623                                similar::ChangeTag::Delete => {
1624                                    let old_end_row = old_row + line_count - 1;
1625                                    let new_row =
1626                                        new_snapshot.anchor_before(Point::new(new_row, 0));
1627
1628                                    if let Some((_, last_deleted_row_range)) =
1629                                        deleted_row_ranges.last_mut()
1630                                    {
1631                                        if *last_deleted_row_range.end() + 1 == old_row {
1632                                            *last_deleted_row_range =
1633                                                *last_deleted_row_range.start()..=old_end_row;
1634                                        } else {
1635                                            deleted_row_ranges
1636                                                .push((new_row, old_row..=old_end_row));
1637                                        }
1638                                    } else {
1639                                        deleted_row_ranges.push((new_row, old_row..=old_end_row));
1640                                    }
1641
1642                                    old_row += line_count;
1643                                }
1644                                similar::ChangeTag::Insert => {
1645                                    let new_end_row = new_row + line_count - 1;
1646                                    let start = new_snapshot.anchor_before(Point::new(new_row, 0));
1647                                    let end = new_snapshot.anchor_before(Point::new(
1648                                        new_end_row,
1649                                        new_snapshot.line_len(MultiBufferRow(new_end_row)),
1650                                    ));
1651                                    inserted_row_ranges.push(start..=end);
1652                                    new_row += line_count;
1653                                }
1654                            }
1655                        }
1656
1657                        (deleted_row_ranges, inserted_row_ranges)
1658                    })
1659                    .await;
1660
1661                this.update(&mut cx, |this, cx| {
1662                    this.diff.deleted_row_ranges = deleted_row_ranges;
1663                    this.diff.inserted_row_ranges = inserted_row_ranges;
1664                    this.diff.task = None;
1665                    if this.diff.should_update {
1666                        this.update_diff(cx);
1667                    }
1668                    cx.notify();
1669                })
1670                .ok();
1671            }));
1672        }
1673    }
1674}
1675
1676fn strip_invalid_spans_from_codeblock(
1677    stream: impl Stream<Item = Result<String>>,
1678) -> impl Stream<Item = Result<String>> {
1679    let mut first_line = true;
1680    let mut buffer = String::new();
1681    let mut starts_with_markdown_codeblock = false;
1682    let mut includes_start_or_end_span = false;
1683    stream.filter_map(move |chunk| {
1684        let chunk = match chunk {
1685            Ok(chunk) => chunk,
1686            Err(err) => return future::ready(Some(Err(err))),
1687        };
1688        buffer.push_str(&chunk);
1689
1690        if buffer.len() > "<|S|".len() && buffer.starts_with("<|S|") {
1691            includes_start_or_end_span = true;
1692
1693            buffer = buffer
1694                .strip_prefix("<|S|>")
1695                .or_else(|| buffer.strip_prefix("<|S|"))
1696                .unwrap_or(&buffer)
1697                .to_string();
1698        } else if buffer.ends_with("|E|>") {
1699            includes_start_or_end_span = true;
1700        } else if buffer.starts_with("<|")
1701            || buffer.starts_with("<|S")
1702            || buffer.starts_with("<|S|")
1703            || buffer.ends_with('|')
1704            || buffer.ends_with("|E")
1705            || buffer.ends_with("|E|")
1706        {
1707            return future::ready(None);
1708        }
1709
1710        if first_line {
1711            if buffer.is_empty() || buffer == "`" || buffer == "``" {
1712                return future::ready(None);
1713            } else if buffer.starts_with("```") {
1714                starts_with_markdown_codeblock = true;
1715                if let Some(newline_ix) = buffer.find('\n') {
1716                    buffer.replace_range(..newline_ix + 1, "");
1717                    first_line = false;
1718                } else {
1719                    return future::ready(None);
1720                }
1721            }
1722        }
1723
1724        let mut text = buffer.to_string();
1725        if starts_with_markdown_codeblock {
1726            text = text
1727                .strip_suffix("\n```\n")
1728                .or_else(|| text.strip_suffix("\n```"))
1729                .or_else(|| text.strip_suffix("\n``"))
1730                .or_else(|| text.strip_suffix("\n`"))
1731                .or_else(|| text.strip_suffix('\n'))
1732                .unwrap_or(&text)
1733                .to_string();
1734        }
1735
1736        if includes_start_or_end_span {
1737            text = text
1738                .strip_suffix("|E|>")
1739                .or_else(|| text.strip_suffix("E|>"))
1740                .or_else(|| text.strip_prefix("|>"))
1741                .or_else(|| text.strip_prefix('>'))
1742                .unwrap_or(&text)
1743                .to_string();
1744        };
1745
1746        if text.contains('\n') {
1747            first_line = false;
1748        }
1749
1750        let remainder = buffer.split_off(text.len());
1751        let result = if buffer.is_empty() {
1752            None
1753        } else {
1754            Some(Ok(buffer.clone()))
1755        };
1756
1757        buffer = remainder;
1758        future::ready(result)
1759    })
1760}
1761
1762fn merge_ranges(ranges: &mut Vec<Range<Anchor>>, buffer: &MultiBufferSnapshot) {
1763    ranges.sort_unstable_by(|a, b| {
1764        a.start
1765            .cmp(&b.start, buffer)
1766            .then_with(|| b.end.cmp(&a.end, buffer))
1767    });
1768
1769    let mut ix = 0;
1770    while ix + 1 < ranges.len() {
1771        let b = ranges[ix + 1].clone();
1772        let a = &mut ranges[ix];
1773        if a.end.cmp(&b.start, buffer).is_gt() {
1774            if a.end.cmp(&b.end, buffer).is_lt() {
1775                a.end = b.end;
1776            }
1777            ranges.remove(ix + 1);
1778        } else {
1779            ix += 1;
1780        }
1781    }
1782}
1783
1784#[cfg(test)]
1785mod tests {
1786    use std::sync::Arc;
1787
1788    use crate::FakeCompletionProvider;
1789
1790    use super::*;
1791    use futures::stream::{self};
1792    use gpui::{Context, TestAppContext};
1793    use indoc::indoc;
1794    use language::{
1795        language_settings, tree_sitter_rust, Buffer, Language, LanguageConfig, LanguageMatcher,
1796        Point,
1797    };
1798    use rand::prelude::*;
1799    use serde::Serialize;
1800    use settings::SettingsStore;
1801
1802    #[derive(Serialize)]
1803    pub struct DummyCompletionRequest {
1804        pub name: String,
1805    }
1806
1807    #[gpui::test(iterations = 10)]
1808    async fn test_transform_autoindent(cx: &mut TestAppContext, mut rng: StdRng) {
1809        let provider = FakeCompletionProvider::default();
1810        cx.set_global(cx.update(SettingsStore::test));
1811        cx.set_global(CompletionProvider::Fake(provider.clone()));
1812        cx.update(language_settings::init);
1813
1814        let text = indoc! {"
1815            fn main() {
1816                let x = 0;
1817                for _ in 0..10 {
1818                    x += 1;
1819                }
1820            }
1821        "};
1822        let buffer =
1823            cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
1824        let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
1825        let range = buffer.read_with(cx, |buffer, cx| {
1826            let snapshot = buffer.snapshot(cx);
1827            snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(4, 5))
1828        });
1829        let codegen = cx.new_model(|cx| {
1830            Codegen::new(buffer.clone(), CodegenKind::Transform { range }, None, cx)
1831        });
1832
1833        let request = LanguageModelRequest::default();
1834        codegen.update(cx, |codegen, cx| codegen.start(request, cx));
1835
1836        let mut new_text = concat!(
1837            "       let mut x = 0;\n",
1838            "       while x < 10 {\n",
1839            "           x += 1;\n",
1840            "       }",
1841        );
1842        while !new_text.is_empty() {
1843            let max_len = cmp::min(new_text.len(), 10);
1844            let len = rng.gen_range(1..=max_len);
1845            let (chunk, suffix) = new_text.split_at(len);
1846            provider.send_completion(chunk.into());
1847            new_text = suffix;
1848            cx.background_executor.run_until_parked();
1849        }
1850        provider.finish_completion();
1851        cx.background_executor.run_until_parked();
1852
1853        assert_eq!(
1854            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1855            indoc! {"
1856                fn main() {
1857                    let mut x = 0;
1858                    while x < 10 {
1859                        x += 1;
1860                    }
1861                }
1862            "}
1863        );
1864    }
1865
1866    #[gpui::test(iterations = 10)]
1867    async fn test_autoindent_when_generating_past_indentation(
1868        cx: &mut TestAppContext,
1869        mut rng: StdRng,
1870    ) {
1871        let provider = FakeCompletionProvider::default();
1872        cx.set_global(CompletionProvider::Fake(provider.clone()));
1873        cx.set_global(cx.update(SettingsStore::test));
1874        cx.update(language_settings::init);
1875
1876        let text = indoc! {"
1877            fn main() {
1878                le
1879            }
1880        "};
1881        let buffer =
1882            cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
1883        let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
1884        let position = buffer.read_with(cx, |buffer, cx| {
1885            let snapshot = buffer.snapshot(cx);
1886            snapshot.anchor_before(Point::new(1, 6))
1887        });
1888        let codegen = cx.new_model(|cx| {
1889            Codegen::new(buffer.clone(), CodegenKind::Generate { position }, None, cx)
1890        });
1891
1892        let request = LanguageModelRequest::default();
1893        codegen.update(cx, |codegen, cx| codegen.start(request, cx));
1894
1895        let mut new_text = concat!(
1896            "t mut x = 0;\n",
1897            "while x < 10 {\n",
1898            "    x += 1;\n",
1899            "}", //
1900        );
1901        while !new_text.is_empty() {
1902            let max_len = cmp::min(new_text.len(), 10);
1903            let len = rng.gen_range(1..=max_len);
1904            let (chunk, suffix) = new_text.split_at(len);
1905            provider.send_completion(chunk.into());
1906            new_text = suffix;
1907            cx.background_executor.run_until_parked();
1908        }
1909        provider.finish_completion();
1910        cx.background_executor.run_until_parked();
1911
1912        assert_eq!(
1913            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1914            indoc! {"
1915                fn main() {
1916                    let mut x = 0;
1917                    while x < 10 {
1918                        x += 1;
1919                    }
1920                }
1921            "}
1922        );
1923    }
1924
1925    #[gpui::test(iterations = 10)]
1926    async fn test_autoindent_when_generating_before_indentation(
1927        cx: &mut TestAppContext,
1928        mut rng: StdRng,
1929    ) {
1930        let provider = FakeCompletionProvider::default();
1931        cx.set_global(CompletionProvider::Fake(provider.clone()));
1932        cx.set_global(cx.update(SettingsStore::test));
1933        cx.update(language_settings::init);
1934
1935        let text = concat!(
1936            "fn main() {\n",
1937            "  \n",
1938            "}\n" //
1939        );
1940        let buffer =
1941            cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
1942        let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
1943        let position = buffer.read_with(cx, |buffer, cx| {
1944            let snapshot = buffer.snapshot(cx);
1945            snapshot.anchor_before(Point::new(1, 2))
1946        });
1947        let codegen = cx.new_model(|cx| {
1948            Codegen::new(buffer.clone(), CodegenKind::Generate { position }, None, cx)
1949        });
1950
1951        let request = LanguageModelRequest::default();
1952        codegen.update(cx, |codegen, cx| codegen.start(request, cx));
1953
1954        let mut new_text = concat!(
1955            "let mut x = 0;\n",
1956            "while x < 10 {\n",
1957            "    x += 1;\n",
1958            "}", //
1959        );
1960        while !new_text.is_empty() {
1961            let max_len = cmp::min(new_text.len(), 10);
1962            let len = rng.gen_range(1..=max_len);
1963            let (chunk, suffix) = new_text.split_at(len);
1964            provider.send_completion(chunk.into());
1965            new_text = suffix;
1966            cx.background_executor.run_until_parked();
1967        }
1968        provider.finish_completion();
1969        cx.background_executor.run_until_parked();
1970
1971        assert_eq!(
1972            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
1973            indoc! {"
1974                fn main() {
1975                    let mut x = 0;
1976                    while x < 10 {
1977                        x += 1;
1978                    }
1979                }
1980            "}
1981        );
1982    }
1983
1984    #[gpui::test]
1985    async fn test_strip_invalid_spans_from_codeblock() {
1986        assert_eq!(
1987            strip_invalid_spans_from_codeblock(chunks("Lorem ipsum dolor", 2))
1988                .map(|chunk| chunk.unwrap())
1989                .collect::<String>()
1990                .await,
1991            "Lorem ipsum dolor"
1992        );
1993        assert_eq!(
1994            strip_invalid_spans_from_codeblock(chunks("```\nLorem ipsum dolor", 2))
1995                .map(|chunk| chunk.unwrap())
1996                .collect::<String>()
1997                .await,
1998            "Lorem ipsum dolor"
1999        );
2000        assert_eq!(
2001            strip_invalid_spans_from_codeblock(chunks("```\nLorem ipsum dolor\n```", 2))
2002                .map(|chunk| chunk.unwrap())
2003                .collect::<String>()
2004                .await,
2005            "Lorem ipsum dolor"
2006        );
2007        assert_eq!(
2008            strip_invalid_spans_from_codeblock(chunks("```\nLorem ipsum dolor\n```\n", 2))
2009                .map(|chunk| chunk.unwrap())
2010                .collect::<String>()
2011                .await,
2012            "Lorem ipsum dolor"
2013        );
2014        assert_eq!(
2015            strip_invalid_spans_from_codeblock(chunks(
2016                "```html\n```js\nLorem ipsum dolor\n```\n```",
2017                2
2018            ))
2019            .map(|chunk| chunk.unwrap())
2020            .collect::<String>()
2021            .await,
2022            "```js\nLorem ipsum dolor\n```"
2023        );
2024        assert_eq!(
2025            strip_invalid_spans_from_codeblock(chunks("``\nLorem ipsum dolor\n```", 2))
2026                .map(|chunk| chunk.unwrap())
2027                .collect::<String>()
2028                .await,
2029            "``\nLorem ipsum dolor\n```"
2030        );
2031        assert_eq!(
2032            strip_invalid_spans_from_codeblock(chunks("<|S|Lorem ipsum|E|>", 2))
2033                .map(|chunk| chunk.unwrap())
2034                .collect::<String>()
2035                .await,
2036            "Lorem ipsum"
2037        );
2038
2039        assert_eq!(
2040            strip_invalid_spans_from_codeblock(chunks("<|S|>Lorem ipsum", 2))
2041                .map(|chunk| chunk.unwrap())
2042                .collect::<String>()
2043                .await,
2044            "Lorem ipsum"
2045        );
2046
2047        assert_eq!(
2048            strip_invalid_spans_from_codeblock(chunks("```\n<|S|>Lorem ipsum\n```", 2))
2049                .map(|chunk| chunk.unwrap())
2050                .collect::<String>()
2051                .await,
2052            "Lorem ipsum"
2053        );
2054        assert_eq!(
2055            strip_invalid_spans_from_codeblock(chunks("```\n<|S|Lorem ipsum|E|>\n```", 2))
2056                .map(|chunk| chunk.unwrap())
2057                .collect::<String>()
2058                .await,
2059            "Lorem ipsum"
2060        );
2061        fn chunks(text: &str, size: usize) -> impl Stream<Item = Result<String>> {
2062            stream::iter(
2063                text.chars()
2064                    .collect::<Vec<_>>()
2065                    .chunks(size)
2066                    .map(|chunk| Ok(chunk.iter().collect::<String>()))
2067                    .collect::<Vec<_>>(),
2068            )
2069        }
2070    }
2071
2072    fn rust_lang() -> Language {
2073        Language::new(
2074            LanguageConfig {
2075                name: "Rust".into(),
2076                matcher: LanguageMatcher {
2077                    path_suffixes: vec!["rs".to_string()],
2078                    ..Default::default()
2079                },
2080                ..Default::default()
2081            },
2082            Some(tree_sitter_rust::language()),
2083        )
2084        .with_indents_query(
2085            r#"
2086            (call_expression) @indent
2087            (field_expression) @indent
2088            (_ "(" ")" @end) @indent
2089            (_ "{" "}" @end) @indent
2090            "#,
2091        )
2092        .unwrap()
2093    }
2094}