1use crate::{
   2    AssistantPanel, AssistantPanelEvent, CycleNextInlineAssist, CyclePreviousInlineAssist,
   3};
   4use anyhow::{anyhow, Context as _, Result};
   5use assistant_context_editor::{humanize_token_count, RequestType};
   6use assistant_settings::AssistantSettings;
   7use client::{telemetry::Telemetry, ErrorExt};
   8use collections::{hash_map, HashMap, HashSet, VecDeque};
   9use editor::{
  10    actions::{MoveDown, MoveUp, SelectAll},
  11    display_map::{
  12        BlockContext, BlockPlacement, BlockProperties, BlockStyle, CustomBlockId, RenderBlock,
  13        ToDisplayPoint,
  14    },
  15    Anchor, AnchorRangeExt, CodeActionProvider, Editor, EditorElement, EditorEvent, EditorMode,
  16    EditorStyle, ExcerptId, ExcerptRange, GutterDimensions, MultiBuffer, MultiBufferSnapshot,
  17    ToOffset as _, ToPoint,
  18};
  19use feature_flags::{
  20    Assistant2FeatureFlag, FeatureFlagAppExt as _, FeatureFlagViewExt as _, ZedPro,
  21};
  22use fs::Fs;
  23use futures::{
  24    channel::mpsc,
  25    future::{BoxFuture, LocalBoxFuture},
  26    join, SinkExt, Stream, StreamExt,
  27};
  28use gpui::{
  29    anchored, deferred, point, AnyElement, App, ClickEvent, Context, CursorStyle, Entity,
  30    EventEmitter, FocusHandle, Focusable, FontWeight, Global, HighlightStyle, Subscription, Task,
  31    TextStyle, UpdateGlobal, WeakEntity, Window,
  32};
  33use language::{line_diff, Buffer, IndentKind, Point, Selection, TransactionId};
  34use language_model::{
  35    report_assistant_event, LanguageModel, LanguageModelRegistry, LanguageModelRequest,
  36    LanguageModelRequestMessage, LanguageModelTextStream, Role,
  37};
  38use language_model_selector::{InlineLanguageModelSelector, LanguageModelSelector};
  39use multi_buffer::MultiBufferRow;
  40use parking_lot::Mutex;
  41use project::{CodeAction, ProjectTransaction};
  42use prompt_store::PromptBuilder;
  43use rope::Rope;
  44use settings::{update_settings_file, Settings, SettingsStore};
  45use smol::future::FutureExt;
  46use std::{
  47    cmp,
  48    future::{self, Future},
  49    iter, mem,
  50    ops::{Range, RangeInclusive},
  51    pin::Pin,
  52    rc::Rc,
  53    sync::Arc,
  54    task::{self, Poll},
  55    time::{Duration, Instant},
  56};
  57use streaming_diff::{CharOperation, LineDiff, LineOperation, StreamingDiff};
  58use telemetry_events::{AssistantEvent, AssistantKind, AssistantPhase};
  59use terminal_view::terminal_panel::TerminalPanel;
  60use text::{OffsetRangeExt, ToPoint as _};
  61use theme::ThemeSettings;
  62use ui::{
  63    prelude::*, text_for_action, CheckboxWithLabel, IconButtonShape, KeyBinding, Popover, Tooltip,
  64};
  65use util::{RangeExt, ResultExt};
  66use workspace::{notifications::NotificationId, ItemHandle, Toast, Workspace};
  67
  68pub fn init(
  69    fs: Arc<dyn Fs>,
  70    prompt_builder: Arc<PromptBuilder>,
  71    telemetry: Arc<Telemetry>,
  72    cx: &mut App,
  73) {
  74    cx.set_global(InlineAssistant::new(fs, prompt_builder, telemetry));
  75    cx.observe_new(|_, window, cx| {
  76        let Some(window) = window else {
  77            return;
  78        };
  79        let workspace = cx.entity().clone();
  80        InlineAssistant::update_global(cx, |inline_assistant, cx| {
  81            inline_assistant.register_workspace(&workspace, window, cx)
  82        });
  83
  84        cx.observe_flag::<Assistant2FeatureFlag, _>(window, {
  85            |is_assistant2_enabled, _workspace, _window, cx| {
  86                InlineAssistant::update_global(cx, |inline_assistant, _cx| {
  87                    inline_assistant.is_assistant2_enabled = is_assistant2_enabled;
  88                });
  89            }
  90        })
  91        .detach();
  92    })
  93    .detach();
  94}
  95
  96const PROMPT_HISTORY_MAX_LEN: usize = 20;
  97
  98pub struct InlineAssistant {
  99    next_assist_id: InlineAssistId,
 100    next_assist_group_id: InlineAssistGroupId,
 101    assists: HashMap<InlineAssistId, InlineAssist>,
 102    assists_by_editor: HashMap<WeakEntity<Editor>, EditorInlineAssists>,
 103    assist_groups: HashMap<InlineAssistGroupId, InlineAssistGroup>,
 104    confirmed_assists: HashMap<InlineAssistId, Entity<CodegenAlternative>>,
 105    prompt_history: VecDeque<String>,
 106    prompt_builder: Arc<PromptBuilder>,
 107    telemetry: Arc<Telemetry>,
 108    fs: Arc<dyn Fs>,
 109    is_assistant2_enabled: bool,
 110}
 111
 112impl Global for InlineAssistant {}
 113
 114impl InlineAssistant {
 115    pub fn new(
 116        fs: Arc<dyn Fs>,
 117        prompt_builder: Arc<PromptBuilder>,
 118        telemetry: Arc<Telemetry>,
 119    ) -> Self {
 120        Self {
 121            next_assist_id: InlineAssistId::default(),
 122            next_assist_group_id: InlineAssistGroupId::default(),
 123            assists: HashMap::default(),
 124            assists_by_editor: HashMap::default(),
 125            assist_groups: HashMap::default(),
 126            confirmed_assists: HashMap::default(),
 127            prompt_history: VecDeque::default(),
 128            prompt_builder,
 129            telemetry,
 130            fs,
 131            is_assistant2_enabled: false,
 132        }
 133    }
 134
 135    pub fn register_workspace(
 136        &mut self,
 137        workspace: &Entity<Workspace>,
 138        window: &mut Window,
 139        cx: &mut App,
 140    ) {
 141        window
 142            .subscribe(workspace, cx, |workspace, event, window, cx| {
 143                Self::update_global(cx, |this, cx| {
 144                    this.handle_workspace_event(workspace, event, window, cx)
 145                });
 146            })
 147            .detach();
 148
 149        let workspace = workspace.downgrade();
 150        cx.observe_global::<SettingsStore>(move |cx| {
 151            let Some(workspace) = workspace.upgrade() else {
 152                return;
 153            };
 154            let Some(terminal_panel) = workspace.read(cx).panel::<TerminalPanel>(cx) else {
 155                return;
 156            };
 157            let enabled = AssistantSettings::get_global(cx).enabled;
 158            terminal_panel.update(cx, |terminal_panel, cx| {
 159                terminal_panel.set_assistant_enabled(enabled, cx)
 160            });
 161        })
 162        .detach();
 163    }
 164
 165    fn handle_workspace_event(
 166        &mut self,
 167        workspace: Entity<Workspace>,
 168        event: &workspace::Event,
 169        window: &mut Window,
 170        cx: &mut App,
 171    ) {
 172        match event {
 173            workspace::Event::UserSavedItem { item, .. } => {
 174                // When the user manually saves an editor, automatically accepts all finished transformations.
 175                if let Some(editor) = item.upgrade().and_then(|item| item.act_as::<Editor>(cx)) {
 176                    if let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) {
 177                        for assist_id in editor_assists.assist_ids.clone() {
 178                            let assist = &self.assists[&assist_id];
 179                            if let CodegenStatus::Done = assist.codegen.read(cx).status(cx) {
 180                                self.finish_assist(assist_id, false, window, cx)
 181                            }
 182                        }
 183                    }
 184                }
 185            }
 186            workspace::Event::ItemAdded { item } => {
 187                self.register_workspace_item(&workspace, item.as_ref(), window, cx);
 188            }
 189            _ => (),
 190        }
 191    }
 192
 193    fn register_workspace_item(
 194        &mut self,
 195        workspace: &Entity<Workspace>,
 196        item: &dyn ItemHandle,
 197        window: &mut Window,
 198        cx: &mut App,
 199    ) {
 200        let is_assistant2_enabled = self.is_assistant2_enabled;
 201
 202        if let Some(editor) = item.act_as::<Editor>(cx) {
 203            editor.update(cx, |editor, cx| {
 204                if is_assistant2_enabled {
 205                    editor.remove_code_action_provider(
 206                        ASSISTANT_CODE_ACTION_PROVIDER_ID.into(),
 207                        window,
 208                        cx,
 209                    );
 210                } else {
 211                    editor.add_code_action_provider(
 212                        Rc::new(AssistantCodeActionProvider {
 213                            editor: cx.entity().downgrade(),
 214                            workspace: workspace.downgrade(),
 215                        }),
 216                        window,
 217                        cx,
 218                    );
 219                }
 220            });
 221        }
 222    }
 223
 224    pub fn assist(
 225        &mut self,
 226        editor: &Entity<Editor>,
 227        workspace: Option<WeakEntity<Workspace>>,
 228        assistant_panel: Option<&Entity<AssistantPanel>>,
 229        initial_prompt: Option<String>,
 230        window: &mut Window,
 231        cx: &mut App,
 232    ) {
 233        let (snapshot, initial_selections) = editor.update(cx, |editor, cx| {
 234            (
 235                editor.buffer().read(cx).snapshot(cx),
 236                editor.selections.all::<Point>(cx),
 237            )
 238        });
 239
 240        let mut selections = Vec::<Selection<Point>>::new();
 241        let mut newest_selection = None;
 242        for mut selection in initial_selections {
 243            if selection.end > selection.start {
 244                selection.start.column = 0;
 245                // If the selection ends at the start of the line, we don't want to include it.
 246                if selection.end.column == 0 {
 247                    selection.end.row -= 1;
 248                }
 249                selection.end.column = snapshot.line_len(MultiBufferRow(selection.end.row));
 250            }
 251
 252            if let Some(prev_selection) = selections.last_mut() {
 253                if selection.start <= prev_selection.end {
 254                    prev_selection.end = selection.end;
 255                    continue;
 256                }
 257            }
 258
 259            let latest_selection = newest_selection.get_or_insert_with(|| selection.clone());
 260            if selection.id > latest_selection.id {
 261                *latest_selection = selection.clone();
 262            }
 263            selections.push(selection);
 264        }
 265        let newest_selection = newest_selection.unwrap();
 266
 267        let mut codegen_ranges = Vec::new();
 268        for (buffer, buffer_range, excerpt_id) in
 269            snapshot.ranges_to_buffer_ranges(selections.iter().map(|selection| {
 270                snapshot.anchor_before(selection.start)..snapshot.anchor_after(selection.end)
 271            }))
 272        {
 273            let start = buffer.anchor_before(buffer_range.start);
 274            let end = buffer.anchor_after(buffer_range.end);
 275
 276            codegen_ranges.push(Anchor::range_in_buffer(
 277                excerpt_id,
 278                buffer.remote_id(),
 279                start..end,
 280            ));
 281
 282            if let Some(model) = LanguageModelRegistry::read_global(cx).active_model() {
 283                self.telemetry.report_assistant_event(AssistantEvent {
 284                    conversation_id: None,
 285                    kind: AssistantKind::Inline,
 286                    phase: AssistantPhase::Invoked,
 287                    message_id: None,
 288                    model: model.telemetry_id(),
 289                    model_provider: model.provider_id().to_string(),
 290                    response_latency: None,
 291                    error_message: None,
 292                    language_name: buffer.language().map(|language| language.name().to_proto()),
 293                });
 294            }
 295        }
 296
 297        let assist_group_id = self.next_assist_group_id.post_inc();
 298        let prompt_buffer = cx.new(|cx| Buffer::local(initial_prompt.unwrap_or_default(), cx));
 299        let prompt_buffer = cx.new(|cx| MultiBuffer::singleton(prompt_buffer, cx));
 300
 301        let mut assists = Vec::new();
 302        let mut assist_to_focus = None;
 303        for range in codegen_ranges {
 304            let assist_id = self.next_assist_id.post_inc();
 305            let codegen = cx.new(|cx| {
 306                Codegen::new(
 307                    editor.read(cx).buffer().clone(),
 308                    range.clone(),
 309                    None,
 310                    self.telemetry.clone(),
 311                    self.prompt_builder.clone(),
 312                    cx,
 313                )
 314            });
 315
 316            let gutter_dimensions = Arc::new(Mutex::new(GutterDimensions::default()));
 317            let prompt_editor = cx.new(|cx| {
 318                PromptEditor::new(
 319                    assist_id,
 320                    gutter_dimensions.clone(),
 321                    self.prompt_history.clone(),
 322                    prompt_buffer.clone(),
 323                    codegen.clone(),
 324                    editor,
 325                    assistant_panel,
 326                    workspace.clone(),
 327                    self.fs.clone(),
 328                    window,
 329                    cx,
 330                )
 331            });
 332
 333            if assist_to_focus.is_none() {
 334                let focus_assist = if newest_selection.reversed {
 335                    range.start.to_point(&snapshot) == newest_selection.start
 336                } else {
 337                    range.end.to_point(&snapshot) == newest_selection.end
 338                };
 339                if focus_assist {
 340                    assist_to_focus = Some(assist_id);
 341                }
 342            }
 343
 344            let [prompt_block_id, end_block_id] =
 345                self.insert_assist_blocks(editor, &range, &prompt_editor, cx);
 346
 347            assists.push((
 348                assist_id,
 349                range,
 350                prompt_editor,
 351                prompt_block_id,
 352                end_block_id,
 353            ));
 354        }
 355
 356        let editor_assists = self
 357            .assists_by_editor
 358            .entry(editor.downgrade())
 359            .or_insert_with(|| EditorInlineAssists::new(&editor, window, cx));
 360        let mut assist_group = InlineAssistGroup::new();
 361        for (assist_id, range, prompt_editor, prompt_block_id, end_block_id) in assists {
 362            self.assists.insert(
 363                assist_id,
 364                InlineAssist::new(
 365                    assist_id,
 366                    assist_group_id,
 367                    assistant_panel.is_some(),
 368                    editor,
 369                    &prompt_editor,
 370                    prompt_block_id,
 371                    end_block_id,
 372                    range,
 373                    prompt_editor.read(cx).codegen.clone(),
 374                    workspace.clone(),
 375                    window,
 376                    cx,
 377                ),
 378            );
 379            assist_group.assist_ids.push(assist_id);
 380            editor_assists.assist_ids.push(assist_id);
 381        }
 382        self.assist_groups.insert(assist_group_id, assist_group);
 383
 384        if let Some(assist_id) = assist_to_focus {
 385            self.focus_assist(assist_id, window, cx);
 386        }
 387    }
 388
 389    #[allow(clippy::too_many_arguments)]
 390    pub fn suggest_assist(
 391        &mut self,
 392        editor: &Entity<Editor>,
 393        mut range: Range<Anchor>,
 394        initial_prompt: String,
 395        initial_transaction_id: Option<TransactionId>,
 396        focus: bool,
 397        workspace: Option<WeakEntity<Workspace>>,
 398        assistant_panel: Option<&Entity<AssistantPanel>>,
 399        window: &mut Window,
 400        cx: &mut App,
 401    ) -> InlineAssistId {
 402        let assist_group_id = self.next_assist_group_id.post_inc();
 403        let prompt_buffer = cx.new(|cx| Buffer::local(&initial_prompt, cx));
 404        let prompt_buffer = cx.new(|cx| MultiBuffer::singleton(prompt_buffer, cx));
 405
 406        let assist_id = self.next_assist_id.post_inc();
 407
 408        let buffer = editor.read(cx).buffer().clone();
 409        {
 410            let snapshot = buffer.read(cx).read(cx);
 411            range.start = range.start.bias_left(&snapshot);
 412            range.end = range.end.bias_right(&snapshot);
 413        }
 414
 415        let codegen = cx.new(|cx| {
 416            Codegen::new(
 417                editor.read(cx).buffer().clone(),
 418                range.clone(),
 419                initial_transaction_id,
 420                self.telemetry.clone(),
 421                self.prompt_builder.clone(),
 422                cx,
 423            )
 424        });
 425
 426        let gutter_dimensions = Arc::new(Mutex::new(GutterDimensions::default()));
 427        let prompt_editor = cx.new(|cx| {
 428            PromptEditor::new(
 429                assist_id,
 430                gutter_dimensions.clone(),
 431                self.prompt_history.clone(),
 432                prompt_buffer.clone(),
 433                codegen.clone(),
 434                editor,
 435                assistant_panel,
 436                workspace.clone(),
 437                self.fs.clone(),
 438                window,
 439                cx,
 440            )
 441        });
 442
 443        let [prompt_block_id, end_block_id] =
 444            self.insert_assist_blocks(editor, &range, &prompt_editor, cx);
 445
 446        let editor_assists = self
 447            .assists_by_editor
 448            .entry(editor.downgrade())
 449            .or_insert_with(|| EditorInlineAssists::new(&editor, window, cx));
 450
 451        let mut assist_group = InlineAssistGroup::new();
 452        self.assists.insert(
 453            assist_id,
 454            InlineAssist::new(
 455                assist_id,
 456                assist_group_id,
 457                assistant_panel.is_some(),
 458                editor,
 459                &prompt_editor,
 460                prompt_block_id,
 461                end_block_id,
 462                range,
 463                prompt_editor.read(cx).codegen.clone(),
 464                workspace.clone(),
 465                window,
 466                cx,
 467            ),
 468        );
 469        assist_group.assist_ids.push(assist_id);
 470        editor_assists.assist_ids.push(assist_id);
 471        self.assist_groups.insert(assist_group_id, assist_group);
 472
 473        if focus {
 474            self.focus_assist(assist_id, window, cx);
 475        }
 476
 477        assist_id
 478    }
 479
 480    fn insert_assist_blocks(
 481        &self,
 482        editor: &Entity<Editor>,
 483        range: &Range<Anchor>,
 484        prompt_editor: &Entity<PromptEditor>,
 485        cx: &mut App,
 486    ) -> [CustomBlockId; 2] {
 487        let prompt_editor_height = prompt_editor.update(cx, |prompt_editor, cx| {
 488            prompt_editor
 489                .editor
 490                .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1 + 2)
 491        });
 492        let assist_blocks = vec![
 493            BlockProperties {
 494                style: BlockStyle::Sticky,
 495                placement: BlockPlacement::Above(range.start),
 496                height: prompt_editor_height,
 497                render: build_assist_editor_renderer(prompt_editor),
 498                priority: 0,
 499            },
 500            BlockProperties {
 501                style: BlockStyle::Sticky,
 502                placement: BlockPlacement::Below(range.end),
 503                height: 0,
 504                render: Arc::new(|cx| {
 505                    v_flex()
 506                        .h_full()
 507                        .w_full()
 508                        .border_t_1()
 509                        .border_color(cx.theme().status().info_border)
 510                        .into_any_element()
 511                }),
 512                priority: 0,
 513            },
 514        ];
 515
 516        editor.update(cx, |editor, cx| {
 517            let block_ids = editor.insert_blocks(assist_blocks, None, cx);
 518            [block_ids[0], block_ids[1]]
 519        })
 520    }
 521
 522    fn handle_prompt_editor_focus_in(&mut self, assist_id: InlineAssistId, cx: &mut App) {
 523        let assist = &self.assists[&assist_id];
 524        let Some(decorations) = assist.decorations.as_ref() else {
 525            return;
 526        };
 527        let assist_group = self.assist_groups.get_mut(&assist.group_id).unwrap();
 528        let editor_assists = self.assists_by_editor.get_mut(&assist.editor).unwrap();
 529
 530        assist_group.active_assist_id = Some(assist_id);
 531        if assist_group.linked {
 532            for assist_id in &assist_group.assist_ids {
 533                if let Some(decorations) = self.assists[assist_id].decorations.as_ref() {
 534                    decorations.prompt_editor.update(cx, |prompt_editor, cx| {
 535                        prompt_editor.set_show_cursor_when_unfocused(true, cx)
 536                    });
 537                }
 538            }
 539        }
 540
 541        assist
 542            .editor
 543            .update(cx, |editor, cx| {
 544                let scroll_top = editor.scroll_position(cx).y;
 545                let scroll_bottom = scroll_top + editor.visible_line_count().unwrap_or(0.);
 546                let prompt_row = editor
 547                    .row_for_block(decorations.prompt_block_id, cx)
 548                    .unwrap()
 549                    .0 as f32;
 550
 551                if (scroll_top..scroll_bottom).contains(&prompt_row) {
 552                    editor_assists.scroll_lock = Some(InlineAssistScrollLock {
 553                        assist_id,
 554                        distance_from_top: prompt_row - scroll_top,
 555                    });
 556                } else {
 557                    editor_assists.scroll_lock = None;
 558                }
 559            })
 560            .ok();
 561    }
 562
 563    fn handle_prompt_editor_focus_out(&mut self, assist_id: InlineAssistId, cx: &mut App) {
 564        let assist = &self.assists[&assist_id];
 565        let assist_group = self.assist_groups.get_mut(&assist.group_id).unwrap();
 566        if assist_group.active_assist_id == Some(assist_id) {
 567            assist_group.active_assist_id = None;
 568            if assist_group.linked {
 569                for assist_id in &assist_group.assist_ids {
 570                    if let Some(decorations) = self.assists[assist_id].decorations.as_ref() {
 571                        decorations.prompt_editor.update(cx, |prompt_editor, cx| {
 572                            prompt_editor.set_show_cursor_when_unfocused(false, cx)
 573                        });
 574                    }
 575                }
 576            }
 577        }
 578    }
 579
 580    fn handle_prompt_editor_event(
 581        &mut self,
 582        prompt_editor: Entity<PromptEditor>,
 583        event: &PromptEditorEvent,
 584        window: &mut Window,
 585        cx: &mut App,
 586    ) {
 587        let assist_id = prompt_editor.read(cx).id;
 588        match event {
 589            PromptEditorEvent::StartRequested => {
 590                self.start_assist(assist_id, window, cx);
 591            }
 592            PromptEditorEvent::StopRequested => {
 593                self.stop_assist(assist_id, cx);
 594            }
 595            PromptEditorEvent::ConfirmRequested => {
 596                self.finish_assist(assist_id, false, window, cx);
 597            }
 598            PromptEditorEvent::CancelRequested => {
 599                self.finish_assist(assist_id, true, window, cx);
 600            }
 601            PromptEditorEvent::DismissRequested => {
 602                self.dismiss_assist(assist_id, window, cx);
 603            }
 604        }
 605    }
 606
 607    fn handle_editor_newline(&mut self, editor: Entity<Editor>, window: &mut Window, cx: &mut App) {
 608        let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
 609            return;
 610        };
 611
 612        if editor.read(cx).selections.count() == 1 {
 613            let (selection, buffer) = editor.update(cx, |editor, cx| {
 614                (
 615                    editor.selections.newest::<usize>(cx),
 616                    editor.buffer().read(cx).snapshot(cx),
 617                )
 618            });
 619            for assist_id in &editor_assists.assist_ids {
 620                let assist = &self.assists[assist_id];
 621                let assist_range = assist.range.to_offset(&buffer);
 622                if assist_range.contains(&selection.start) && assist_range.contains(&selection.end)
 623                {
 624                    if matches!(assist.codegen.read(cx).status(cx), CodegenStatus::Pending) {
 625                        self.dismiss_assist(*assist_id, window, cx);
 626                    } else {
 627                        self.finish_assist(*assist_id, false, window, cx);
 628                    }
 629
 630                    return;
 631                }
 632            }
 633        }
 634
 635        cx.propagate();
 636    }
 637
 638    fn handle_editor_cancel(&mut self, editor: Entity<Editor>, window: &mut Window, cx: &mut App) {
 639        let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
 640            return;
 641        };
 642
 643        if editor.read(cx).selections.count() == 1 {
 644            let (selection, buffer) = editor.update(cx, |editor, cx| {
 645                (
 646                    editor.selections.newest::<usize>(cx),
 647                    editor.buffer().read(cx).snapshot(cx),
 648                )
 649            });
 650            let mut closest_assist_fallback = None;
 651            for assist_id in &editor_assists.assist_ids {
 652                let assist = &self.assists[assist_id];
 653                let assist_range = assist.range.to_offset(&buffer);
 654                if assist.decorations.is_some() {
 655                    if assist_range.contains(&selection.start)
 656                        && assist_range.contains(&selection.end)
 657                    {
 658                        self.focus_assist(*assist_id, window, cx);
 659                        return;
 660                    } else {
 661                        let distance_from_selection = assist_range
 662                            .start
 663                            .abs_diff(selection.start)
 664                            .min(assist_range.start.abs_diff(selection.end))
 665                            + assist_range
 666                                .end
 667                                .abs_diff(selection.start)
 668                                .min(assist_range.end.abs_diff(selection.end));
 669                        match closest_assist_fallback {
 670                            Some((_, old_distance)) => {
 671                                if distance_from_selection < old_distance {
 672                                    closest_assist_fallback =
 673                                        Some((assist_id, distance_from_selection));
 674                                }
 675                            }
 676                            None => {
 677                                closest_assist_fallback = Some((assist_id, distance_from_selection))
 678                            }
 679                        }
 680                    }
 681                }
 682            }
 683
 684            if let Some((&assist_id, _)) = closest_assist_fallback {
 685                self.focus_assist(assist_id, window, cx);
 686            }
 687        }
 688
 689        cx.propagate();
 690    }
 691
 692    fn handle_editor_release(
 693        &mut self,
 694        editor: WeakEntity<Editor>,
 695        window: &mut Window,
 696        cx: &mut App,
 697    ) {
 698        if let Some(editor_assists) = self.assists_by_editor.get_mut(&editor) {
 699            for assist_id in editor_assists.assist_ids.clone() {
 700                self.finish_assist(assist_id, true, window, cx);
 701            }
 702        }
 703    }
 704
 705    fn handle_editor_change(&mut self, editor: Entity<Editor>, window: &mut Window, cx: &mut App) {
 706        let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) else {
 707            return;
 708        };
 709        let Some(scroll_lock) = editor_assists.scroll_lock.as_ref() else {
 710            return;
 711        };
 712        let assist = &self.assists[&scroll_lock.assist_id];
 713        let Some(decorations) = assist.decorations.as_ref() else {
 714            return;
 715        };
 716
 717        editor.update(cx, |editor, cx| {
 718            let scroll_position = editor.scroll_position(cx);
 719            let target_scroll_top = editor
 720                .row_for_block(decorations.prompt_block_id, cx)
 721                .unwrap()
 722                .0 as f32
 723                - scroll_lock.distance_from_top;
 724            if target_scroll_top != scroll_position.y {
 725                editor.set_scroll_position(point(scroll_position.x, target_scroll_top), window, cx);
 726            }
 727        });
 728    }
 729
 730    fn handle_editor_event(
 731        &mut self,
 732        editor: Entity<Editor>,
 733        event: &EditorEvent,
 734        window: &mut Window,
 735        cx: &mut App,
 736    ) {
 737        let Some(editor_assists) = self.assists_by_editor.get_mut(&editor.downgrade()) else {
 738            return;
 739        };
 740
 741        match event {
 742            EditorEvent::Edited { transaction_id } => {
 743                let buffer = editor.read(cx).buffer().read(cx);
 744                let edited_ranges =
 745                    buffer.edited_ranges_for_transaction::<usize>(*transaction_id, cx);
 746                let snapshot = buffer.snapshot(cx);
 747
 748                for assist_id in editor_assists.assist_ids.clone() {
 749                    let assist = &self.assists[&assist_id];
 750                    if matches!(
 751                        assist.codegen.read(cx).status(cx),
 752                        CodegenStatus::Error(_) | CodegenStatus::Done
 753                    ) {
 754                        let assist_range = assist.range.to_offset(&snapshot);
 755                        if edited_ranges
 756                            .iter()
 757                            .any(|range| range.overlaps(&assist_range))
 758                        {
 759                            self.finish_assist(assist_id, false, window, cx);
 760                        }
 761                    }
 762                }
 763            }
 764            EditorEvent::ScrollPositionChanged { .. } => {
 765                if let Some(scroll_lock) = editor_assists.scroll_lock.as_ref() {
 766                    let assist = &self.assists[&scroll_lock.assist_id];
 767                    if let Some(decorations) = assist.decorations.as_ref() {
 768                        let distance_from_top = editor.update(cx, |editor, cx| {
 769                            let scroll_top = editor.scroll_position(cx).y;
 770                            let prompt_row = editor
 771                                .row_for_block(decorations.prompt_block_id, cx)
 772                                .unwrap()
 773                                .0 as f32;
 774                            prompt_row - scroll_top
 775                        });
 776
 777                        if distance_from_top != scroll_lock.distance_from_top {
 778                            editor_assists.scroll_lock = None;
 779                        }
 780                    }
 781                }
 782            }
 783            EditorEvent::SelectionsChanged { .. } => {
 784                for assist_id in editor_assists.assist_ids.clone() {
 785                    let assist = &self.assists[&assist_id];
 786                    if let Some(decorations) = assist.decorations.as_ref() {
 787                        if decorations
 788                            .prompt_editor
 789                            .focus_handle(cx)
 790                            .is_focused(window)
 791                        {
 792                            return;
 793                        }
 794                    }
 795                }
 796
 797                editor_assists.scroll_lock = None;
 798            }
 799            _ => {}
 800        }
 801    }
 802
 803    pub fn finish_assist(
 804        &mut self,
 805        assist_id: InlineAssistId,
 806        undo: bool,
 807        window: &mut Window,
 808        cx: &mut App,
 809    ) {
 810        if let Some(assist) = self.assists.get(&assist_id) {
 811            let assist_group_id = assist.group_id;
 812            if self.assist_groups[&assist_group_id].linked {
 813                for assist_id in self.unlink_assist_group(assist_group_id, window, cx) {
 814                    self.finish_assist(assist_id, undo, window, cx);
 815                }
 816                return;
 817            }
 818        }
 819
 820        self.dismiss_assist(assist_id, window, cx);
 821
 822        if let Some(assist) = self.assists.remove(&assist_id) {
 823            if let hash_map::Entry::Occupied(mut entry) = self.assist_groups.entry(assist.group_id)
 824            {
 825                entry.get_mut().assist_ids.retain(|id| *id != assist_id);
 826                if entry.get().assist_ids.is_empty() {
 827                    entry.remove();
 828                }
 829            }
 830
 831            if let hash_map::Entry::Occupied(mut entry) =
 832                self.assists_by_editor.entry(assist.editor.clone())
 833            {
 834                entry.get_mut().assist_ids.retain(|id| *id != assist_id);
 835                if entry.get().assist_ids.is_empty() {
 836                    entry.remove();
 837                    if let Some(editor) = assist.editor.upgrade() {
 838                        self.update_editor_highlights(&editor, cx);
 839                    }
 840                } else {
 841                    entry.get().highlight_updates.send(()).ok();
 842                }
 843            }
 844
 845            let active_alternative = assist.codegen.read(cx).active_alternative().clone();
 846            let message_id = active_alternative.read(cx).message_id.clone();
 847
 848            if let Some(model) = LanguageModelRegistry::read_global(cx).active_model() {
 849                let language_name = assist.editor.upgrade().and_then(|editor| {
 850                    let multibuffer = editor.read(cx).buffer().read(cx);
 851                    let multibuffer_snapshot = multibuffer.snapshot(cx);
 852                    let ranges = multibuffer_snapshot.range_to_buffer_ranges(assist.range.clone());
 853                    ranges
 854                        .first()
 855                        .and_then(|(buffer, _, _)| buffer.language())
 856                        .map(|language| language.name())
 857                });
 858                report_assistant_event(
 859                    AssistantEvent {
 860                        conversation_id: None,
 861                        kind: AssistantKind::Inline,
 862                        message_id,
 863                        phase: if undo {
 864                            AssistantPhase::Rejected
 865                        } else {
 866                            AssistantPhase::Accepted
 867                        },
 868                        model: model.telemetry_id(),
 869                        model_provider: model.provider_id().to_string(),
 870                        response_latency: None,
 871                        error_message: None,
 872                        language_name: language_name.map(|name| name.to_proto()),
 873                    },
 874                    Some(self.telemetry.clone()),
 875                    cx.http_client(),
 876                    model.api_key(cx),
 877                    cx.background_executor(),
 878                );
 879            }
 880
 881            if undo {
 882                assist.codegen.update(cx, |codegen, cx| codegen.undo(cx));
 883            } else {
 884                self.confirmed_assists.insert(assist_id, active_alternative);
 885            }
 886        }
 887    }
 888
 889    fn dismiss_assist(
 890        &mut self,
 891        assist_id: InlineAssistId,
 892        window: &mut Window,
 893        cx: &mut App,
 894    ) -> bool {
 895        let Some(assist) = self.assists.get_mut(&assist_id) else {
 896            return false;
 897        };
 898        let Some(editor) = assist.editor.upgrade() else {
 899            return false;
 900        };
 901        let Some(decorations) = assist.decorations.take() else {
 902            return false;
 903        };
 904
 905        editor.update(cx, |editor, cx| {
 906            let mut to_remove = decorations.removed_line_block_ids;
 907            to_remove.insert(decorations.prompt_block_id);
 908            to_remove.insert(decorations.end_block_id);
 909            editor.remove_blocks(to_remove, None, cx);
 910        });
 911
 912        if decorations
 913            .prompt_editor
 914            .focus_handle(cx)
 915            .contains_focused(window, cx)
 916        {
 917            self.focus_next_assist(assist_id, window, cx);
 918        }
 919
 920        if let Some(editor_assists) = self.assists_by_editor.get_mut(&editor.downgrade()) {
 921            if editor_assists
 922                .scroll_lock
 923                .as_ref()
 924                .map_or(false, |lock| lock.assist_id == assist_id)
 925            {
 926                editor_assists.scroll_lock = None;
 927            }
 928            editor_assists.highlight_updates.send(()).ok();
 929        }
 930
 931        true
 932    }
 933
 934    fn focus_next_assist(&mut self, assist_id: InlineAssistId, window: &mut Window, cx: &mut App) {
 935        let Some(assist) = self.assists.get(&assist_id) else {
 936            return;
 937        };
 938
 939        let assist_group = &self.assist_groups[&assist.group_id];
 940        let assist_ix = assist_group
 941            .assist_ids
 942            .iter()
 943            .position(|id| *id == assist_id)
 944            .unwrap();
 945        let assist_ids = assist_group
 946            .assist_ids
 947            .iter()
 948            .skip(assist_ix + 1)
 949            .chain(assist_group.assist_ids.iter().take(assist_ix));
 950
 951        for assist_id in assist_ids {
 952            let assist = &self.assists[assist_id];
 953            if assist.decorations.is_some() {
 954                self.focus_assist(*assist_id, window, cx);
 955                return;
 956            }
 957        }
 958
 959        assist
 960            .editor
 961            .update(cx, |editor, cx| window.focus(&editor.focus_handle(cx)))
 962            .ok();
 963    }
 964
 965    fn focus_assist(&mut self, assist_id: InlineAssistId, window: &mut Window, cx: &mut App) {
 966        let Some(assist) = self.assists.get(&assist_id) else {
 967            return;
 968        };
 969
 970        if let Some(decorations) = assist.decorations.as_ref() {
 971            decorations.prompt_editor.update(cx, |prompt_editor, cx| {
 972                prompt_editor.editor.update(cx, |editor, cx| {
 973                    window.focus(&editor.focus_handle(cx));
 974                    editor.select_all(&SelectAll, window, cx);
 975                })
 976            });
 977        }
 978
 979        self.scroll_to_assist(assist_id, window, cx);
 980    }
 981
 982    pub fn scroll_to_assist(
 983        &mut self,
 984        assist_id: InlineAssistId,
 985        window: &mut Window,
 986        cx: &mut App,
 987    ) {
 988        let Some(assist) = self.assists.get(&assist_id) else {
 989            return;
 990        };
 991        let Some(editor) = assist.editor.upgrade() else {
 992            return;
 993        };
 994
 995        let position = assist.range.start;
 996        editor.update(cx, |editor, cx| {
 997            editor.change_selections(None, window, cx, |selections| {
 998                selections.select_anchor_ranges([position..position])
 999            });
1000
1001            let mut scroll_target_top;
1002            let mut scroll_target_bottom;
1003            if let Some(decorations) = assist.decorations.as_ref() {
1004                scroll_target_top = editor
1005                    .row_for_block(decorations.prompt_block_id, cx)
1006                    .unwrap()
1007                    .0 as f32;
1008                scroll_target_bottom = editor
1009                    .row_for_block(decorations.end_block_id, cx)
1010                    .unwrap()
1011                    .0 as f32;
1012            } else {
1013                let snapshot = editor.snapshot(window, cx);
1014                let start_row = assist
1015                    .range
1016                    .start
1017                    .to_display_point(&snapshot.display_snapshot)
1018                    .row();
1019                scroll_target_top = start_row.0 as f32;
1020                scroll_target_bottom = scroll_target_top + 1.;
1021            }
1022            scroll_target_top -= editor.vertical_scroll_margin() as f32;
1023            scroll_target_bottom += editor.vertical_scroll_margin() as f32;
1024
1025            let height_in_lines = editor.visible_line_count().unwrap_or(0.);
1026            let scroll_top = editor.scroll_position(cx).y;
1027            let scroll_bottom = scroll_top + height_in_lines;
1028
1029            if scroll_target_top < scroll_top {
1030                editor.set_scroll_position(point(0., scroll_target_top), window, cx);
1031            } else if scroll_target_bottom > scroll_bottom {
1032                if (scroll_target_bottom - scroll_target_top) <= height_in_lines {
1033                    editor.set_scroll_position(
1034                        point(0., scroll_target_bottom - height_in_lines),
1035                        window,
1036                        cx,
1037                    );
1038                } else {
1039                    editor.set_scroll_position(point(0., scroll_target_top), window, cx);
1040                }
1041            }
1042        });
1043    }
1044
1045    fn unlink_assist_group(
1046        &mut self,
1047        assist_group_id: InlineAssistGroupId,
1048        window: &mut Window,
1049        cx: &mut App,
1050    ) -> Vec<InlineAssistId> {
1051        let assist_group = self.assist_groups.get_mut(&assist_group_id).unwrap();
1052        assist_group.linked = false;
1053        for assist_id in &assist_group.assist_ids {
1054            let assist = self.assists.get_mut(assist_id).unwrap();
1055            if let Some(editor_decorations) = assist.decorations.as_ref() {
1056                editor_decorations
1057                    .prompt_editor
1058                    .update(cx, |prompt_editor, cx| prompt_editor.unlink(window, cx));
1059            }
1060        }
1061        assist_group.assist_ids.clone()
1062    }
1063
1064    pub fn start_assist(&mut self, assist_id: InlineAssistId, window: &mut Window, cx: &mut App) {
1065        let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
1066            assist
1067        } else {
1068            return;
1069        };
1070
1071        let assist_group_id = assist.group_id;
1072        if self.assist_groups[&assist_group_id].linked {
1073            for assist_id in self.unlink_assist_group(assist_group_id, window, cx) {
1074                self.start_assist(assist_id, window, cx);
1075            }
1076            return;
1077        }
1078
1079        let Some(user_prompt) = assist.user_prompt(cx) else {
1080            return;
1081        };
1082
1083        self.prompt_history.retain(|prompt| *prompt != user_prompt);
1084        self.prompt_history.push_back(user_prompt.clone());
1085        if self.prompt_history.len() > PROMPT_HISTORY_MAX_LEN {
1086            self.prompt_history.pop_front();
1087        }
1088
1089        let assistant_panel_context = assist.assistant_panel_context(cx);
1090
1091        assist
1092            .codegen
1093            .update(cx, |codegen, cx| {
1094                codegen.start(user_prompt, assistant_panel_context, cx)
1095            })
1096            .log_err();
1097    }
1098
1099    pub fn stop_assist(&mut self, assist_id: InlineAssistId, cx: &mut App) {
1100        let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
1101            assist
1102        } else {
1103            return;
1104        };
1105
1106        assist.codegen.update(cx, |codegen, cx| codegen.stop(cx));
1107    }
1108
1109    fn update_editor_highlights(&self, editor: &Entity<Editor>, cx: &mut App) {
1110        let mut gutter_pending_ranges = Vec::new();
1111        let mut gutter_transformed_ranges = Vec::new();
1112        let mut foreground_ranges = Vec::new();
1113        let mut inserted_row_ranges = Vec::new();
1114        let empty_assist_ids = Vec::new();
1115        let assist_ids = self
1116            .assists_by_editor
1117            .get(&editor.downgrade())
1118            .map_or(&empty_assist_ids, |editor_assists| {
1119                &editor_assists.assist_ids
1120            });
1121
1122        for assist_id in assist_ids {
1123            if let Some(assist) = self.assists.get(assist_id) {
1124                let codegen = assist.codegen.read(cx);
1125                let buffer = codegen.buffer(cx).read(cx).read(cx);
1126                foreground_ranges.extend(codegen.last_equal_ranges(cx).iter().cloned());
1127
1128                let pending_range =
1129                    codegen.edit_position(cx).unwrap_or(assist.range.start)..assist.range.end;
1130                if pending_range.end.to_offset(&buffer) > pending_range.start.to_offset(&buffer) {
1131                    gutter_pending_ranges.push(pending_range);
1132                }
1133
1134                if let Some(edit_position) = codegen.edit_position(cx) {
1135                    let edited_range = assist.range.start..edit_position;
1136                    if edited_range.end.to_offset(&buffer) > edited_range.start.to_offset(&buffer) {
1137                        gutter_transformed_ranges.push(edited_range);
1138                    }
1139                }
1140
1141                if assist.decorations.is_some() {
1142                    inserted_row_ranges
1143                        .extend(codegen.diff(cx).inserted_row_ranges.iter().cloned());
1144                }
1145            }
1146        }
1147
1148        let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
1149        merge_ranges(&mut foreground_ranges, &snapshot);
1150        merge_ranges(&mut gutter_pending_ranges, &snapshot);
1151        merge_ranges(&mut gutter_transformed_ranges, &snapshot);
1152        editor.update(cx, |editor, cx| {
1153            enum GutterPendingRange {}
1154            if gutter_pending_ranges.is_empty() {
1155                editor.clear_gutter_highlights::<GutterPendingRange>(cx);
1156            } else {
1157                editor.highlight_gutter::<GutterPendingRange>(
1158                    &gutter_pending_ranges,
1159                    |cx| cx.theme().status().info_background,
1160                    cx,
1161                )
1162            }
1163
1164            enum GutterTransformedRange {}
1165            if gutter_transformed_ranges.is_empty() {
1166                editor.clear_gutter_highlights::<GutterTransformedRange>(cx);
1167            } else {
1168                editor.highlight_gutter::<GutterTransformedRange>(
1169                    &gutter_transformed_ranges,
1170                    |cx| cx.theme().status().info,
1171                    cx,
1172                )
1173            }
1174
1175            if foreground_ranges.is_empty() {
1176                editor.clear_highlights::<InlineAssist>(cx);
1177            } else {
1178                editor.highlight_text::<InlineAssist>(
1179                    foreground_ranges,
1180                    HighlightStyle {
1181                        fade_out: Some(0.6),
1182                        ..Default::default()
1183                    },
1184                    cx,
1185                );
1186            }
1187
1188            editor.clear_row_highlights::<InlineAssist>();
1189            for row_range in inserted_row_ranges {
1190                editor.highlight_rows::<InlineAssist>(
1191                    row_range,
1192                    cx.theme().status().info_background,
1193                    false,
1194                    cx,
1195                );
1196            }
1197        });
1198    }
1199
1200    fn update_editor_blocks(
1201        &mut self,
1202        editor: &Entity<Editor>,
1203        assist_id: InlineAssistId,
1204        window: &mut Window,
1205        cx: &mut App,
1206    ) {
1207        let Some(assist) = self.assists.get_mut(&assist_id) else {
1208            return;
1209        };
1210        let Some(decorations) = assist.decorations.as_mut() else {
1211            return;
1212        };
1213
1214        let codegen = assist.codegen.read(cx);
1215        let old_snapshot = codegen.snapshot(cx);
1216        let old_buffer = codegen.old_buffer(cx);
1217        let deleted_row_ranges = codegen.diff(cx).deleted_row_ranges.clone();
1218
1219        editor.update(cx, |editor, cx| {
1220            let old_blocks = mem::take(&mut decorations.removed_line_block_ids);
1221            editor.remove_blocks(old_blocks, None, cx);
1222
1223            let mut new_blocks = Vec::new();
1224            for (new_row, old_row_range) in deleted_row_ranges {
1225                let (_, buffer_start) = old_snapshot
1226                    .point_to_buffer_offset(Point::new(*old_row_range.start(), 0))
1227                    .unwrap();
1228                let (_, buffer_end) = old_snapshot
1229                    .point_to_buffer_offset(Point::new(
1230                        *old_row_range.end(),
1231                        old_snapshot.line_len(MultiBufferRow(*old_row_range.end())),
1232                    ))
1233                    .unwrap();
1234
1235                let deleted_lines_editor = cx.new(|cx| {
1236                    let multi_buffer =
1237                        cx.new(|_| MultiBuffer::without_headers(language::Capability::ReadOnly));
1238                    multi_buffer.update(cx, |multi_buffer, cx| {
1239                        multi_buffer.push_excerpts(
1240                            old_buffer.clone(),
1241                            Some(ExcerptRange {
1242                                context: buffer_start..buffer_end,
1243                                primary: None,
1244                            }),
1245                            cx,
1246                        );
1247                    });
1248
1249                    enum DeletedLines {}
1250                    let mut editor = Editor::for_multibuffer(multi_buffer, None, true, window, cx);
1251                    editor.set_soft_wrap_mode(language::language_settings::SoftWrap::None, cx);
1252                    editor.set_show_wrap_guides(false, cx);
1253                    editor.set_show_gutter(false, cx);
1254                    editor.scroll_manager.set_forbid_vertical_scroll(true);
1255                    editor.set_show_scrollbars(false, cx);
1256                    editor.set_read_only(true);
1257                    editor.set_show_edit_predictions(Some(false), window, cx);
1258                    editor.highlight_rows::<DeletedLines>(
1259                        Anchor::min()..Anchor::max(),
1260                        cx.theme().status().deleted_background,
1261                        false,
1262                        cx,
1263                    );
1264                    editor
1265                });
1266
1267                let height =
1268                    deleted_lines_editor.update(cx, |editor, cx| editor.max_point(cx).row().0 + 1);
1269                new_blocks.push(BlockProperties {
1270                    placement: BlockPlacement::Above(new_row),
1271                    height,
1272                    style: BlockStyle::Flex,
1273                    render: Arc::new(move |cx| {
1274                        div()
1275                            .block_mouse_down()
1276                            .bg(cx.theme().status().deleted_background)
1277                            .size_full()
1278                            .h(height as f32 * cx.window.line_height())
1279                            .pl(cx.gutter_dimensions.full_width())
1280                            .child(deleted_lines_editor.clone())
1281                            .into_any_element()
1282                    }),
1283                    priority: 0,
1284                });
1285            }
1286
1287            decorations.removed_line_block_ids = editor
1288                .insert_blocks(new_blocks, None, cx)
1289                .into_iter()
1290                .collect();
1291        })
1292    }
1293}
1294
1295struct EditorInlineAssists {
1296    assist_ids: Vec<InlineAssistId>,
1297    scroll_lock: Option<InlineAssistScrollLock>,
1298    highlight_updates: async_watch::Sender<()>,
1299    _update_highlights: Task<Result<()>>,
1300    _subscriptions: Vec<gpui::Subscription>,
1301}
1302
1303struct InlineAssistScrollLock {
1304    assist_id: InlineAssistId,
1305    distance_from_top: f32,
1306}
1307
1308impl EditorInlineAssists {
1309    fn new(editor: &Entity<Editor>, window: &mut Window, cx: &mut App) -> Self {
1310        let (highlight_updates_tx, mut highlight_updates_rx) = async_watch::channel(());
1311        Self {
1312            assist_ids: Vec::new(),
1313            scroll_lock: None,
1314            highlight_updates: highlight_updates_tx,
1315            _update_highlights: cx.spawn(|cx| {
1316                let editor = editor.downgrade();
1317                async move {
1318                    while let Ok(()) = highlight_updates_rx.changed().await {
1319                        let editor = editor.upgrade().context("editor was dropped")?;
1320                        cx.update_global(|assistant: &mut InlineAssistant, cx| {
1321                            assistant.update_editor_highlights(&editor, cx);
1322                        })?;
1323                    }
1324                    Ok(())
1325                }
1326            }),
1327            _subscriptions: vec![
1328                cx.observe_release_in(editor, window, {
1329                    let editor = editor.downgrade();
1330                    |_, window, cx| {
1331                        InlineAssistant::update_global(cx, |this, cx| {
1332                            this.handle_editor_release(editor, window, cx);
1333                        })
1334                    }
1335                }),
1336                window.observe(editor, cx, move |editor, window, cx| {
1337                    InlineAssistant::update_global(cx, |this, cx| {
1338                        this.handle_editor_change(editor, window, cx)
1339                    })
1340                }),
1341                window.subscribe(editor, cx, move |editor, event, window, cx| {
1342                    InlineAssistant::update_global(cx, |this, cx| {
1343                        this.handle_editor_event(editor, event, window, cx)
1344                    })
1345                }),
1346                editor.update(cx, |editor, cx| {
1347                    let editor_handle = cx.entity().downgrade();
1348                    editor.register_action(move |_: &editor::actions::Newline, window, cx| {
1349                        InlineAssistant::update_global(cx, |this, cx| {
1350                            if let Some(editor) = editor_handle.upgrade() {
1351                                this.handle_editor_newline(editor, window, cx)
1352                            }
1353                        })
1354                    })
1355                }),
1356                editor.update(cx, |editor, cx| {
1357                    let editor_handle = cx.entity().downgrade();
1358                    editor.register_action(move |_: &editor::actions::Cancel, window, cx| {
1359                        InlineAssistant::update_global(cx, |this, cx| {
1360                            if let Some(editor) = editor_handle.upgrade() {
1361                                this.handle_editor_cancel(editor, window, cx)
1362                            }
1363                        })
1364                    })
1365                }),
1366            ],
1367        }
1368    }
1369}
1370
1371struct InlineAssistGroup {
1372    assist_ids: Vec<InlineAssistId>,
1373    linked: bool,
1374    active_assist_id: Option<InlineAssistId>,
1375}
1376
1377impl InlineAssistGroup {
1378    fn new() -> Self {
1379        Self {
1380            assist_ids: Vec::new(),
1381            linked: true,
1382            active_assist_id: None,
1383        }
1384    }
1385}
1386
1387fn build_assist_editor_renderer(editor: &Entity<PromptEditor>) -> RenderBlock {
1388    let editor = editor.clone();
1389    Arc::new(move |cx: &mut BlockContext| {
1390        *editor.read(cx).gutter_dimensions.lock() = *cx.gutter_dimensions;
1391        editor.clone().into_any_element()
1392    })
1393}
1394
1395#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
1396pub struct InlineAssistId(usize);
1397
1398impl InlineAssistId {
1399    fn post_inc(&mut self) -> InlineAssistId {
1400        let id = *self;
1401        self.0 += 1;
1402        id
1403    }
1404}
1405
1406#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
1407struct InlineAssistGroupId(usize);
1408
1409impl InlineAssistGroupId {
1410    fn post_inc(&mut self) -> InlineAssistGroupId {
1411        let id = *self;
1412        self.0 += 1;
1413        id
1414    }
1415}
1416
1417enum PromptEditorEvent {
1418    StartRequested,
1419    StopRequested,
1420    ConfirmRequested,
1421    CancelRequested,
1422    DismissRequested,
1423}
1424
1425struct PromptEditor {
1426    id: InlineAssistId,
1427    editor: Entity<Editor>,
1428    language_model_selector: Entity<LanguageModelSelector>,
1429    edited_since_done: bool,
1430    gutter_dimensions: Arc<Mutex<GutterDimensions>>,
1431    prompt_history: VecDeque<String>,
1432    prompt_history_ix: Option<usize>,
1433    pending_prompt: String,
1434    codegen: Entity<Codegen>,
1435    _codegen_subscription: Subscription,
1436    editor_subscriptions: Vec<Subscription>,
1437    pending_token_count: Task<Result<()>>,
1438    token_counts: Option<TokenCounts>,
1439    _token_count_subscriptions: Vec<Subscription>,
1440    workspace: Option<WeakEntity<Workspace>>,
1441    show_rate_limit_notice: bool,
1442}
1443
1444#[derive(Copy, Clone)]
1445pub struct TokenCounts {
1446    total: usize,
1447    assistant_panel: usize,
1448}
1449
1450impl EventEmitter<PromptEditorEvent> for PromptEditor {}
1451
1452impl Render for PromptEditor {
1453    fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
1454        let gutter_dimensions = *self.gutter_dimensions.lock();
1455        let codegen = self.codegen.read(cx);
1456
1457        let mut buttons = Vec::new();
1458        if codegen.alternative_count(cx) > 1 {
1459            buttons.push(self.render_cycle_controls(cx));
1460        }
1461
1462        let status = codegen.status(cx);
1463        buttons.extend(match status {
1464            CodegenStatus::Idle => {
1465                vec![
1466                    IconButton::new("cancel", IconName::Close)
1467                        .icon_color(Color::Muted)
1468                        .shape(IconButtonShape::Square)
1469                        .tooltip(|window, cx| {
1470                            Tooltip::for_action("Cancel Assist", &menu::Cancel, window, cx)
1471                        })
1472                        .on_click(
1473                            cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
1474                        )
1475                        .into_any_element(),
1476                    IconButton::new("start", IconName::SparkleAlt)
1477                        .icon_color(Color::Muted)
1478                        .shape(IconButtonShape::Square)
1479                        .tooltip(|window, cx| {
1480                            Tooltip::for_action("Transform", &menu::Confirm, window, cx)
1481                        })
1482                        .on_click(
1483                            cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::StartRequested)),
1484                        )
1485                        .into_any_element(),
1486                ]
1487            }
1488            CodegenStatus::Pending => {
1489                vec![
1490                    IconButton::new("cancel", IconName::Close)
1491                        .icon_color(Color::Muted)
1492                        .shape(IconButtonShape::Square)
1493                        .tooltip(Tooltip::text("Cancel Assist"))
1494                        .on_click(
1495                            cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
1496                        )
1497                        .into_any_element(),
1498                    IconButton::new("stop", IconName::Stop)
1499                        .icon_color(Color::Error)
1500                        .shape(IconButtonShape::Square)
1501                        .tooltip(|window, cx| {
1502                            Tooltip::with_meta(
1503                                "Interrupt Transformation",
1504                                Some(&menu::Cancel),
1505                                "Changes won't be discarded",
1506                                window,
1507                                cx,
1508                            )
1509                        })
1510                        .on_click(
1511                            cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::StopRequested)),
1512                        )
1513                        .into_any_element(),
1514                ]
1515            }
1516            CodegenStatus::Error(_) | CodegenStatus::Done => {
1517                let must_rerun =
1518                    self.edited_since_done || matches!(status, CodegenStatus::Error(_));
1519                // when accept button isn't visible, then restart maps to confirm
1520                // when accept button is visible, then restart must be mapped to an alternate keyboard shortcut
1521                let restart_key: &dyn gpui::Action = if must_rerun {
1522                    &menu::Confirm
1523                } else {
1524                    &menu::Restart
1525                };
1526                vec![
1527                    IconButton::new("cancel", IconName::Close)
1528                        .icon_color(Color::Muted)
1529                        .shape(IconButtonShape::Square)
1530                        .tooltip(|window, cx| {
1531                            Tooltip::for_action("Cancel Assist", &menu::Cancel, window, cx)
1532                        })
1533                        .on_click(
1534                            cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
1535                        )
1536                        .into_any_element(),
1537                    IconButton::new("restart", IconName::RotateCw)
1538                        .icon_color(Color::Muted)
1539                        .shape(IconButtonShape::Square)
1540                        .tooltip(|window, cx| {
1541                            Tooltip::with_meta(
1542                                "Regenerate Transformation",
1543                                Some(restart_key),
1544                                "Current change will be discarded",
1545                                window,
1546                                cx,
1547                            )
1548                        })
1549                        .on_click(cx.listener(|_, _, _, cx| {
1550                            cx.emit(PromptEditorEvent::StartRequested);
1551                        }))
1552                        .into_any_element(),
1553                    if !must_rerun {
1554                        IconButton::new("confirm", IconName::Check)
1555                            .icon_color(Color::Info)
1556                            .shape(IconButtonShape::Square)
1557                            .tooltip(|window, cx| {
1558                                Tooltip::for_action("Confirm Assist", &menu::Confirm, window, cx)
1559                            })
1560                            .on_click(cx.listener(|_, _, _, cx| {
1561                                cx.emit(PromptEditorEvent::ConfirmRequested);
1562                            }))
1563                            .into_any_element()
1564                    } else {
1565                        div().into_any_element()
1566                    },
1567                ]
1568            }
1569        });
1570
1571        h_flex()
1572            .key_context("PromptEditor")
1573            .bg(cx.theme().colors().editor_background)
1574            .block_mouse_down()
1575            .cursor(CursorStyle::Arrow)
1576            .border_y_1()
1577            .border_color(cx.theme().status().info_border)
1578            .size_full()
1579            .py(window.line_height() / 2.5)
1580            .on_action(cx.listener(Self::confirm))
1581            .on_action(cx.listener(Self::cancel))
1582            .on_action(cx.listener(Self::restart))
1583            .on_action(cx.listener(Self::move_up))
1584            .on_action(cx.listener(Self::move_down))
1585            .capture_action(cx.listener(Self::cycle_prev))
1586            .capture_action(cx.listener(Self::cycle_next))
1587            .child(
1588                h_flex()
1589                    .w(gutter_dimensions.full_width() + (gutter_dimensions.margin / 2.0))
1590                    .justify_center()
1591                    .gap_2()
1592                    .child(
1593                        InlineLanguageModelSelector::new(self.language_model_selector.clone())
1594                            .render(window, cx),
1595                    )
1596                    .map(|el| {
1597                        let CodegenStatus::Error(error) = self.codegen.read(cx).status(cx) else {
1598                            return el;
1599                        };
1600
1601                        let error_message = SharedString::from(error.to_string());
1602                        if error.error_code() == proto::ErrorCode::RateLimitExceeded
1603                            && cx.has_flag::<ZedPro>()
1604                        {
1605                            el.child(
1606                                v_flex()
1607                                    .child(
1608                                        IconButton::new("rate-limit-error", IconName::XCircle)
1609                                            .toggle_state(self.show_rate_limit_notice)
1610                                            .shape(IconButtonShape::Square)
1611                                            .icon_size(IconSize::Small)
1612                                            .on_click(cx.listener(Self::toggle_rate_limit_notice)),
1613                                    )
1614                                    .children(self.show_rate_limit_notice.then(|| {
1615                                        deferred(
1616                                            anchored()
1617                                                .position_mode(gpui::AnchoredPositionMode::Local)
1618                                                .position(point(px(0.), px(24.)))
1619                                                .anchor(gpui::Corner::TopLeft)
1620                                                .child(self.render_rate_limit_notice(cx)),
1621                                        )
1622                                    })),
1623                            )
1624                        } else {
1625                            el.child(
1626                                div()
1627                                    .id("error")
1628                                    .tooltip(Tooltip::text(error_message))
1629                                    .child(
1630                                        Icon::new(IconName::XCircle)
1631                                            .size(IconSize::Small)
1632                                            .color(Color::Error),
1633                                    ),
1634                            )
1635                        }
1636                    }),
1637            )
1638            .child(div().flex_1().child(self.render_prompt_editor(cx)))
1639            .child(
1640                h_flex()
1641                    .gap_2()
1642                    .pr_6()
1643                    .children(self.render_token_count(cx))
1644                    .children(buttons),
1645            )
1646    }
1647}
1648
1649impl Focusable for PromptEditor {
1650    fn focus_handle(&self, cx: &App) -> FocusHandle {
1651        self.editor.focus_handle(cx)
1652    }
1653}
1654
1655impl PromptEditor {
1656    const MAX_LINES: u8 = 8;
1657
1658    #[allow(clippy::too_many_arguments)]
1659    fn new(
1660        id: InlineAssistId,
1661        gutter_dimensions: Arc<Mutex<GutterDimensions>>,
1662        prompt_history: VecDeque<String>,
1663        prompt_buffer: Entity<MultiBuffer>,
1664        codegen: Entity<Codegen>,
1665        parent_editor: &Entity<Editor>,
1666        assistant_panel: Option<&Entity<AssistantPanel>>,
1667        workspace: Option<WeakEntity<Workspace>>,
1668        fs: Arc<dyn Fs>,
1669        window: &mut Window,
1670        cx: &mut Context<Self>,
1671    ) -> Self {
1672        let prompt_editor = cx.new(|cx| {
1673            let mut editor = Editor::new(
1674                EditorMode::AutoHeight {
1675                    max_lines: Self::MAX_LINES as usize,
1676                },
1677                prompt_buffer,
1678                None,
1679                false,
1680                window,
1681                cx,
1682            );
1683            editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
1684            // Since the prompt editors for all inline assistants are linked,
1685            // always show the cursor (even when it isn't focused) because
1686            // typing in one will make what you typed appear in all of them.
1687            editor.set_show_cursor_when_unfocused(true, cx);
1688            editor.set_placeholder_text(Self::placeholder_text(codegen.read(cx), window, cx), cx);
1689            editor
1690        });
1691
1692        let mut token_count_subscriptions = Vec::new();
1693        token_count_subscriptions.push(cx.subscribe_in(
1694            parent_editor,
1695            window,
1696            Self::handle_parent_editor_event,
1697        ));
1698        if let Some(assistant_panel) = assistant_panel {
1699            token_count_subscriptions.push(cx.subscribe_in(
1700                assistant_panel,
1701                window,
1702                Self::handle_assistant_panel_event,
1703            ));
1704        }
1705
1706        let mut this = Self {
1707            id,
1708            editor: prompt_editor,
1709            language_model_selector: cx.new(|cx| {
1710                let fs = fs.clone();
1711                LanguageModelSelector::new(
1712                    move |model, cx| {
1713                        update_settings_file::<AssistantSettings>(
1714                            fs.clone(),
1715                            cx,
1716                            move |settings, _| settings.set_model(model.clone()),
1717                        );
1718                    },
1719                    window,
1720                    cx,
1721                )
1722            }),
1723            edited_since_done: false,
1724            gutter_dimensions,
1725            prompt_history,
1726            prompt_history_ix: None,
1727            pending_prompt: String::new(),
1728            _codegen_subscription: cx.observe(&codegen, Self::handle_codegen_changed),
1729            editor_subscriptions: Vec::new(),
1730            codegen,
1731            pending_token_count: Task::ready(Ok(())),
1732            token_counts: None,
1733            _token_count_subscriptions: token_count_subscriptions,
1734            workspace,
1735            show_rate_limit_notice: false,
1736        };
1737        this.count_tokens(cx);
1738        this.subscribe_to_editor(window, cx);
1739        this
1740    }
1741
1742    fn subscribe_to_editor(&mut self, window: &mut Window, cx: &mut Context<Self>) {
1743        self.editor_subscriptions.clear();
1744        self.editor_subscriptions.push(cx.subscribe_in(
1745            &self.editor,
1746            window,
1747            Self::handle_prompt_editor_events,
1748        ));
1749    }
1750
1751    fn set_show_cursor_when_unfocused(
1752        &mut self,
1753        show_cursor_when_unfocused: bool,
1754        cx: &mut Context<Self>,
1755    ) {
1756        self.editor.update(cx, |editor, cx| {
1757            editor.set_show_cursor_when_unfocused(show_cursor_when_unfocused, cx)
1758        });
1759    }
1760
1761    fn unlink(&mut self, window: &mut Window, cx: &mut Context<Self>) {
1762        let prompt = self.prompt(cx);
1763        let focus = self.editor.focus_handle(cx).contains_focused(window, cx);
1764        self.editor = cx.new(|cx| {
1765            let mut editor = Editor::auto_height(Self::MAX_LINES as usize, window, cx);
1766            editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
1767            editor.set_placeholder_text(
1768                Self::placeholder_text(self.codegen.read(cx), window, cx),
1769                cx,
1770            );
1771            editor.set_placeholder_text("Add a prompt…", cx);
1772            editor.set_text(prompt, window, cx);
1773            if focus {
1774                window.focus(&editor.focus_handle(cx));
1775            }
1776            editor
1777        });
1778        self.subscribe_to_editor(window, cx);
1779    }
1780
1781    fn placeholder_text(codegen: &Codegen, window: &Window, cx: &App) -> String {
1782        let context_keybinding = text_for_action(&zed_actions::assistant::ToggleFocus, window, cx)
1783            .map(|keybinding| format!(" • {keybinding} for context"))
1784            .unwrap_or_default();
1785
1786        let action = if codegen.is_insertion {
1787            "Generate"
1788        } else {
1789            "Transform"
1790        };
1791
1792        format!("{action}…{context_keybinding} • ↓↑ for history")
1793    }
1794
1795    fn prompt(&self, cx: &App) -> String {
1796        self.editor.read(cx).text(cx)
1797    }
1798
1799    fn toggle_rate_limit_notice(
1800        &mut self,
1801        _: &ClickEvent,
1802        window: &mut Window,
1803        cx: &mut Context<Self>,
1804    ) {
1805        self.show_rate_limit_notice = !self.show_rate_limit_notice;
1806        if self.show_rate_limit_notice {
1807            window.focus(&self.editor.focus_handle(cx));
1808        }
1809        cx.notify();
1810    }
1811
1812    fn handle_parent_editor_event(
1813        &mut self,
1814        _: &Entity<Editor>,
1815        event: &EditorEvent,
1816        _: &mut Window,
1817        cx: &mut Context<Self>,
1818    ) {
1819        if let EditorEvent::BufferEdited { .. } = event {
1820            self.count_tokens(cx);
1821        }
1822    }
1823
1824    fn handle_assistant_panel_event(
1825        &mut self,
1826        _: &Entity<AssistantPanel>,
1827        event: &AssistantPanelEvent,
1828        _: &mut Window,
1829        cx: &mut Context<Self>,
1830    ) {
1831        let AssistantPanelEvent::ContextEdited { .. } = event;
1832        self.count_tokens(cx);
1833    }
1834
1835    fn count_tokens(&mut self, cx: &mut Context<Self>) {
1836        let assist_id = self.id;
1837        self.pending_token_count = cx.spawn(|this, mut cx| async move {
1838            cx.background_executor().timer(Duration::from_secs(1)).await;
1839            let token_count = cx
1840                .update_global(|inline_assistant: &mut InlineAssistant, cx| {
1841                    let assist = inline_assistant
1842                        .assists
1843                        .get(&assist_id)
1844                        .context("assist not found")?;
1845                    anyhow::Ok(assist.count_tokens(cx))
1846                })??
1847                .await?;
1848
1849            this.update(&mut cx, |this, cx| {
1850                this.token_counts = Some(token_count);
1851                cx.notify();
1852            })
1853        })
1854    }
1855
1856    fn handle_prompt_editor_events(
1857        &mut self,
1858        _: &Entity<Editor>,
1859        event: &EditorEvent,
1860        window: &mut Window,
1861        cx: &mut Context<Self>,
1862    ) {
1863        match event {
1864            EditorEvent::Edited { .. } => {
1865                if let Some(workspace) = window.root::<Workspace>().flatten() {
1866                    workspace.update(cx, |workspace, cx| {
1867                        let is_via_ssh = workspace
1868                            .project()
1869                            .update(cx, |project, _| project.is_via_ssh());
1870
1871                        workspace
1872                            .client()
1873                            .telemetry()
1874                            .log_edit_event("inline assist", is_via_ssh);
1875                    });
1876                }
1877                let prompt = self.editor.read(cx).text(cx);
1878                if self
1879                    .prompt_history_ix
1880                    .map_or(true, |ix| self.prompt_history[ix] != prompt)
1881                {
1882                    self.prompt_history_ix.take();
1883                    self.pending_prompt = prompt;
1884                }
1885
1886                self.edited_since_done = true;
1887                cx.notify();
1888            }
1889            EditorEvent::BufferEdited => {
1890                self.count_tokens(cx);
1891            }
1892            EditorEvent::Blurred => {
1893                if self.show_rate_limit_notice {
1894                    self.show_rate_limit_notice = false;
1895                    cx.notify();
1896                }
1897            }
1898            _ => {}
1899        }
1900    }
1901
1902    fn handle_codegen_changed(&mut self, _: Entity<Codegen>, cx: &mut Context<Self>) {
1903        match self.codegen.read(cx).status(cx) {
1904            CodegenStatus::Idle => {
1905                self.editor
1906                    .update(cx, |editor, _| editor.set_read_only(false));
1907            }
1908            CodegenStatus::Pending => {
1909                self.editor
1910                    .update(cx, |editor, _| editor.set_read_only(true));
1911            }
1912            CodegenStatus::Done => {
1913                self.edited_since_done = false;
1914                self.editor
1915                    .update(cx, |editor, _| editor.set_read_only(false));
1916            }
1917            CodegenStatus::Error(error) => {
1918                if cx.has_flag::<ZedPro>()
1919                    && error.error_code() == proto::ErrorCode::RateLimitExceeded
1920                    && !dismissed_rate_limit_notice()
1921                {
1922                    self.show_rate_limit_notice = true;
1923                    cx.notify();
1924                }
1925
1926                self.edited_since_done = false;
1927                self.editor
1928                    .update(cx, |editor, _| editor.set_read_only(false));
1929            }
1930        }
1931    }
1932
1933    fn restart(&mut self, _: &menu::Restart, _window: &mut Window, cx: &mut Context<Self>) {
1934        cx.emit(PromptEditorEvent::StartRequested);
1935    }
1936
1937    fn cancel(
1938        &mut self,
1939        _: &editor::actions::Cancel,
1940        _window: &mut Window,
1941        cx: &mut Context<Self>,
1942    ) {
1943        match self.codegen.read(cx).status(cx) {
1944            CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {
1945                cx.emit(PromptEditorEvent::CancelRequested);
1946            }
1947            CodegenStatus::Pending => {
1948                cx.emit(PromptEditorEvent::StopRequested);
1949            }
1950        }
1951    }
1952
1953    fn confirm(&mut self, _: &menu::Confirm, _window: &mut Window, cx: &mut Context<Self>) {
1954        match self.codegen.read(cx).status(cx) {
1955            CodegenStatus::Idle => {
1956                cx.emit(PromptEditorEvent::StartRequested);
1957            }
1958            CodegenStatus::Pending => {
1959                cx.emit(PromptEditorEvent::DismissRequested);
1960            }
1961            CodegenStatus::Done => {
1962                if self.edited_since_done {
1963                    cx.emit(PromptEditorEvent::StartRequested);
1964                } else {
1965                    cx.emit(PromptEditorEvent::ConfirmRequested);
1966                }
1967            }
1968            CodegenStatus::Error(_) => {
1969                cx.emit(PromptEditorEvent::StartRequested);
1970            }
1971        }
1972    }
1973
1974    fn move_up(&mut self, _: &MoveUp, window: &mut Window, cx: &mut Context<Self>) {
1975        if let Some(ix) = self.prompt_history_ix {
1976            if ix > 0 {
1977                self.prompt_history_ix = Some(ix - 1);
1978                let prompt = self.prompt_history[ix - 1].as_str();
1979                self.editor.update(cx, |editor, cx| {
1980                    editor.set_text(prompt, window, cx);
1981                    editor.move_to_beginning(&Default::default(), window, cx);
1982                });
1983            }
1984        } else if !self.prompt_history.is_empty() {
1985            self.prompt_history_ix = Some(self.prompt_history.len() - 1);
1986            let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();
1987            self.editor.update(cx, |editor, cx| {
1988                editor.set_text(prompt, window, cx);
1989                editor.move_to_beginning(&Default::default(), window, cx);
1990            });
1991        }
1992    }
1993
1994    fn move_down(&mut self, _: &MoveDown, window: &mut Window, cx: &mut Context<Self>) {
1995        if let Some(ix) = self.prompt_history_ix {
1996            if ix < self.prompt_history.len() - 1 {
1997                self.prompt_history_ix = Some(ix + 1);
1998                let prompt = self.prompt_history[ix + 1].as_str();
1999                self.editor.update(cx, |editor, cx| {
2000                    editor.set_text(prompt, window, cx);
2001                    editor.move_to_end(&Default::default(), window, cx)
2002                });
2003            } else {
2004                self.prompt_history_ix = None;
2005                let prompt = self.pending_prompt.as_str();
2006                self.editor.update(cx, |editor, cx| {
2007                    editor.set_text(prompt, window, cx);
2008                    editor.move_to_end(&Default::default(), window, cx)
2009                });
2010            }
2011        }
2012    }
2013
2014    fn cycle_prev(
2015        &mut self,
2016        _: &CyclePreviousInlineAssist,
2017        _: &mut Window,
2018        cx: &mut Context<Self>,
2019    ) {
2020        self.codegen
2021            .update(cx, |codegen, cx| codegen.cycle_prev(cx));
2022    }
2023
2024    fn cycle_next(&mut self, _: &CycleNextInlineAssist, _: &mut Window, cx: &mut Context<Self>) {
2025        self.codegen
2026            .update(cx, |codegen, cx| codegen.cycle_next(cx));
2027    }
2028
2029    fn render_cycle_controls(&self, cx: &Context<Self>) -> AnyElement {
2030        let codegen = self.codegen.read(cx);
2031        let disabled = matches!(codegen.status(cx), CodegenStatus::Idle);
2032
2033        let model_registry = LanguageModelRegistry::read_global(cx);
2034        let default_model = model_registry.active_model();
2035        let alternative_models = model_registry.inline_alternative_models();
2036
2037        let get_model_name = |index: usize| -> String {
2038            let name = |model: &Arc<dyn LanguageModel>| model.name().0.to_string();
2039
2040            match index {
2041                0 => default_model.as_ref().map_or_else(String::new, name),
2042                index if index <= alternative_models.len() => alternative_models
2043                    .get(index - 1)
2044                    .map_or_else(String::new, name),
2045                _ => String::new(),
2046            }
2047        };
2048
2049        let total_models = alternative_models.len() + 1;
2050
2051        if total_models <= 1 {
2052            return div().into_any_element();
2053        }
2054
2055        let current_index = codegen.active_alternative;
2056        let prev_index = (current_index + total_models - 1) % total_models;
2057        let next_index = (current_index + 1) % total_models;
2058
2059        let prev_model_name = get_model_name(prev_index);
2060        let next_model_name = get_model_name(next_index);
2061
2062        h_flex()
2063            .child(
2064                IconButton::new("previous", IconName::ChevronLeft)
2065                    .icon_color(Color::Muted)
2066                    .disabled(disabled || current_index == 0)
2067                    .shape(IconButtonShape::Square)
2068                    .tooltip({
2069                        let focus_handle = self.editor.focus_handle(cx);
2070                        move |window, cx| {
2071                            cx.new(|cx| {
2072                                let mut tooltip = Tooltip::new("Previous Alternative").key_binding(
2073                                    KeyBinding::for_action_in(
2074                                        &CyclePreviousInlineAssist,
2075                                        &focus_handle,
2076                                        window,
2077                                        cx,
2078                                    ),
2079                                );
2080                                if !disabled && current_index != 0 {
2081                                    tooltip = tooltip.meta(prev_model_name.clone());
2082                                }
2083                                tooltip
2084                            })
2085                            .into()
2086                        }
2087                    })
2088                    .on_click(cx.listener(|this, _, _, cx| {
2089                        this.codegen
2090                            .update(cx, |codegen, cx| codegen.cycle_prev(cx))
2091                    })),
2092            )
2093            .child(
2094                Label::new(format!(
2095                    "{}/{}",
2096                    codegen.active_alternative + 1,
2097                    codegen.alternative_count(cx)
2098                ))
2099                .size(LabelSize::Small)
2100                .color(if disabled {
2101                    Color::Disabled
2102                } else {
2103                    Color::Muted
2104                }),
2105            )
2106            .child(
2107                IconButton::new("next", IconName::ChevronRight)
2108                    .icon_color(Color::Muted)
2109                    .disabled(disabled || current_index == total_models - 1)
2110                    .shape(IconButtonShape::Square)
2111                    .tooltip({
2112                        let focus_handle = self.editor.focus_handle(cx);
2113                        move |window, cx| {
2114                            cx.new(|cx| {
2115                                let mut tooltip = Tooltip::new("Next Alternative").key_binding(
2116                                    KeyBinding::for_action_in(
2117                                        &CycleNextInlineAssist,
2118                                        &focus_handle,
2119                                        window,
2120                                        cx,
2121                                    ),
2122                                );
2123                                if !disabled && current_index != total_models - 1 {
2124                                    tooltip = tooltip.meta(next_model_name.clone());
2125                                }
2126                                tooltip
2127                            })
2128                            .into()
2129                        }
2130                    })
2131                    .on_click(cx.listener(|this, _, _, cx| {
2132                        this.codegen
2133                            .update(cx, |codegen, cx| codegen.cycle_next(cx))
2134                    })),
2135            )
2136            .into_any_element()
2137    }
2138
2139    fn render_token_count(&self, cx: &mut Context<Self>) -> Option<impl IntoElement> {
2140        let model = LanguageModelRegistry::read_global(cx).active_model()?;
2141        let token_counts = self.token_counts?;
2142        let max_token_count = model.max_token_count();
2143
2144        let remaining_tokens = max_token_count as isize - token_counts.total as isize;
2145        let token_count_color = if remaining_tokens <= 0 {
2146            Color::Error
2147        } else if token_counts.total as f32 / max_token_count as f32 >= 0.8 {
2148            Color::Warning
2149        } else {
2150            Color::Muted
2151        };
2152
2153        let mut token_count = h_flex()
2154            .id("token_count")
2155            .gap_0p5()
2156            .child(
2157                Label::new(humanize_token_count(token_counts.total))
2158                    .size(LabelSize::Small)
2159                    .color(token_count_color),
2160            )
2161            .child(Label::new("/").size(LabelSize::Small).color(Color::Muted))
2162            .child(
2163                Label::new(humanize_token_count(max_token_count))
2164                    .size(LabelSize::Small)
2165                    .color(Color::Muted),
2166            );
2167        if let Some(workspace) = self.workspace.clone() {
2168            token_count = token_count
2169                .tooltip(move |window, cx| {
2170                    Tooltip::with_meta(
2171                        format!(
2172                            "Tokens Used ({} from the Assistant Panel)",
2173                            humanize_token_count(token_counts.assistant_panel)
2174                        ),
2175                        None,
2176                        "Click to open the Assistant Panel",
2177                        window,
2178                        cx,
2179                    )
2180                })
2181                .cursor_pointer()
2182                .on_mouse_down(gpui::MouseButton::Left, |_, _, cx| cx.stop_propagation())
2183                .on_click(move |_, window, cx| {
2184                    cx.stop_propagation();
2185                    workspace
2186                        .update(cx, |workspace, cx| {
2187                            workspace.focus_panel::<AssistantPanel>(window, cx)
2188                        })
2189                        .ok();
2190                });
2191        } else {
2192            token_count = token_count
2193                .cursor_default()
2194                .tooltip(Tooltip::text("Tokens used"));
2195        }
2196
2197        Some(token_count)
2198    }
2199
2200    fn render_prompt_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
2201        let settings = ThemeSettings::get_global(cx);
2202        let text_style = TextStyle {
2203            color: if self.editor.read(cx).read_only(cx) {
2204                cx.theme().colors().text_disabled
2205            } else {
2206                cx.theme().colors().text
2207            },
2208            font_family: settings.buffer_font.family.clone(),
2209            font_fallbacks: settings.buffer_font.fallbacks.clone(),
2210            font_size: settings.buffer_font_size(cx).into(),
2211            font_weight: settings.buffer_font.weight,
2212            line_height: relative(settings.buffer_line_height.value()),
2213            ..Default::default()
2214        };
2215        EditorElement::new(
2216            &self.editor,
2217            EditorStyle {
2218                background: cx.theme().colors().editor_background,
2219                local_player: cx.theme().players().local(),
2220                text: text_style,
2221                ..Default::default()
2222            },
2223        )
2224    }
2225
2226    fn render_rate_limit_notice(&self, cx: &mut Context<Self>) -> impl IntoElement {
2227        Popover::new().child(
2228            v_flex()
2229                .occlude()
2230                .p_2()
2231                .child(
2232                    Label::new("Out of Tokens")
2233                        .size(LabelSize::Small)
2234                        .weight(FontWeight::BOLD),
2235                )
2236                .child(Label::new(
2237                    "Try Zed Pro for higher limits, a wider range of models, and more.",
2238                ))
2239                .child(
2240                    h_flex()
2241                        .justify_between()
2242                        .child(CheckboxWithLabel::new(
2243                            "dont-show-again",
2244                            Label::new("Don't show again"),
2245                            if dismissed_rate_limit_notice() {
2246                                ui::ToggleState::Selected
2247                            } else {
2248                                ui::ToggleState::Unselected
2249                            },
2250                            |selection, _, cx| {
2251                                let is_dismissed = match selection {
2252                                    ui::ToggleState::Unselected => false,
2253                                    ui::ToggleState::Indeterminate => return,
2254                                    ui::ToggleState::Selected => true,
2255                                };
2256
2257                                set_rate_limit_notice_dismissed(is_dismissed, cx)
2258                            },
2259                        ))
2260                        .child(
2261                            h_flex()
2262                                .gap_2()
2263                                .child(
2264                                    Button::new("dismiss", "Dismiss")
2265                                        .style(ButtonStyle::Transparent)
2266                                        .on_click(cx.listener(Self::toggle_rate_limit_notice)),
2267                                )
2268                                .child(Button::new("more-info", "More Info").on_click(
2269                                    |_event, window, cx| {
2270                                        window.dispatch_action(
2271                                            Box::new(zed_actions::OpenAccountSettings),
2272                                            cx,
2273                                        )
2274                                    },
2275                                )),
2276                        ),
2277                ),
2278        )
2279    }
2280}
2281
2282const DISMISSED_RATE_LIMIT_NOTICE_KEY: &str = "dismissed-rate-limit-notice";
2283
2284fn dismissed_rate_limit_notice() -> bool {
2285    db::kvp::KEY_VALUE_STORE
2286        .read_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY)
2287        .log_err()
2288        .map_or(false, |s| s.is_some())
2289}
2290
2291fn set_rate_limit_notice_dismissed(is_dismissed: bool, cx: &mut App) {
2292    db::write_and_log(cx, move || async move {
2293        if is_dismissed {
2294            db::kvp::KEY_VALUE_STORE
2295                .write_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY.into(), "1".into())
2296                .await
2297        } else {
2298            db::kvp::KEY_VALUE_STORE
2299                .delete_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY.into())
2300                .await
2301        }
2302    })
2303}
2304
2305struct InlineAssist {
2306    group_id: InlineAssistGroupId,
2307    range: Range<Anchor>,
2308    editor: WeakEntity<Editor>,
2309    decorations: Option<InlineAssistDecorations>,
2310    codegen: Entity<Codegen>,
2311    _subscriptions: Vec<Subscription>,
2312    workspace: Option<WeakEntity<Workspace>>,
2313    include_context: bool,
2314}
2315
2316impl InlineAssist {
2317    #[allow(clippy::too_many_arguments)]
2318    fn new(
2319        assist_id: InlineAssistId,
2320        group_id: InlineAssistGroupId,
2321        include_context: bool,
2322        editor: &Entity<Editor>,
2323        prompt_editor: &Entity<PromptEditor>,
2324        prompt_block_id: CustomBlockId,
2325        end_block_id: CustomBlockId,
2326        range: Range<Anchor>,
2327        codegen: Entity<Codegen>,
2328        workspace: Option<WeakEntity<Workspace>>,
2329        window: &mut Window,
2330        cx: &mut App,
2331    ) -> Self {
2332        let prompt_editor_focus_handle = prompt_editor.focus_handle(cx);
2333        InlineAssist {
2334            group_id,
2335            include_context,
2336            editor: editor.downgrade(),
2337            decorations: Some(InlineAssistDecorations {
2338                prompt_block_id,
2339                prompt_editor: prompt_editor.clone(),
2340                removed_line_block_ids: HashSet::default(),
2341                end_block_id,
2342            }),
2343            range,
2344            codegen: codegen.clone(),
2345            workspace: workspace.clone(),
2346            _subscriptions: vec![
2347                window.on_focus_in(&prompt_editor_focus_handle, cx, move |_, cx| {
2348                    InlineAssistant::update_global(cx, |this, cx| {
2349                        this.handle_prompt_editor_focus_in(assist_id, cx)
2350                    })
2351                }),
2352                window.on_focus_out(&prompt_editor_focus_handle, cx, move |_, _, cx| {
2353                    InlineAssistant::update_global(cx, |this, cx| {
2354                        this.handle_prompt_editor_focus_out(assist_id, cx)
2355                    })
2356                }),
2357                window.subscribe(
2358                    prompt_editor,
2359                    cx,
2360                    move |prompt_editor, event, window, cx| {
2361                        InlineAssistant::update_global(cx, |this, cx| {
2362                            this.handle_prompt_editor_event(prompt_editor, event, window, cx)
2363                        })
2364                    },
2365                ),
2366                window.observe(&codegen, cx, {
2367                    let editor = editor.downgrade();
2368                    move |_, window, cx| {
2369                        if let Some(editor) = editor.upgrade() {
2370                            InlineAssistant::update_global(cx, |this, cx| {
2371                                if let Some(editor_assists) =
2372                                    this.assists_by_editor.get(&editor.downgrade())
2373                                {
2374                                    editor_assists.highlight_updates.send(()).ok();
2375                                }
2376
2377                                this.update_editor_blocks(&editor, assist_id, window, cx);
2378                            })
2379                        }
2380                    }
2381                }),
2382                window.subscribe(&codegen, cx, move |codegen, event, window, cx| {
2383                    InlineAssistant::update_global(cx, |this, cx| match event {
2384                        CodegenEvent::Undone => this.finish_assist(assist_id, false, window, cx),
2385                        CodegenEvent::Finished => {
2386                            let assist = if let Some(assist) = this.assists.get(&assist_id) {
2387                                assist
2388                            } else {
2389                                return;
2390                            };
2391
2392                            if let CodegenStatus::Error(error) = codegen.read(cx).status(cx) {
2393                                if assist.decorations.is_none() {
2394                                    if let Some(workspace) = assist
2395                                        .workspace
2396                                        .as_ref()
2397                                        .and_then(|workspace| workspace.upgrade())
2398                                    {
2399                                        let error = format!("Inline assistant error: {}", error);
2400                                        workspace.update(cx, |workspace, cx| {
2401                                            struct InlineAssistantError;
2402
2403                                            let id =
2404                                                NotificationId::composite::<InlineAssistantError>(
2405                                                    assist_id.0,
2406                                                );
2407
2408                                            workspace.show_toast(Toast::new(id, error), cx);
2409                                        })
2410                                    }
2411                                }
2412                            }
2413
2414                            if assist.decorations.is_none() {
2415                                this.finish_assist(assist_id, false, window, cx);
2416                            }
2417                        }
2418                    })
2419                }),
2420            ],
2421        }
2422    }
2423
2424    fn user_prompt(&self, cx: &App) -> Option<String> {
2425        let decorations = self.decorations.as_ref()?;
2426        Some(decorations.prompt_editor.read(cx).prompt(cx))
2427    }
2428
2429    fn assistant_panel_context(&self, cx: &mut App) -> Option<LanguageModelRequest> {
2430        if self.include_context {
2431            let workspace = self.workspace.as_ref()?;
2432            let workspace = workspace.upgrade()?.read(cx);
2433            let assistant_panel = workspace.panel::<AssistantPanel>(cx)?;
2434            Some(
2435                assistant_panel
2436                    .read(cx)
2437                    .active_context(cx)?
2438                    .read(cx)
2439                    .to_completion_request(RequestType::Chat, cx),
2440            )
2441        } else {
2442            None
2443        }
2444    }
2445
2446    pub fn count_tokens(&self, cx: &mut App) -> BoxFuture<'static, Result<TokenCounts>> {
2447        let Some(user_prompt) = self.user_prompt(cx) else {
2448            return future::ready(Err(anyhow!("no user prompt"))).boxed();
2449        };
2450        let assistant_panel_context = self.assistant_panel_context(cx);
2451        self.codegen
2452            .read(cx)
2453            .count_tokens(user_prompt, assistant_panel_context, cx)
2454    }
2455}
2456
2457struct InlineAssistDecorations {
2458    prompt_block_id: CustomBlockId,
2459    prompt_editor: Entity<PromptEditor>,
2460    removed_line_block_ids: HashSet<CustomBlockId>,
2461    end_block_id: CustomBlockId,
2462}
2463
2464#[derive(Copy, Clone, Debug)]
2465pub enum CodegenEvent {
2466    Finished,
2467    Undone,
2468}
2469
2470pub struct Codegen {
2471    alternatives: Vec<Entity<CodegenAlternative>>,
2472    active_alternative: usize,
2473    seen_alternatives: HashSet<usize>,
2474    subscriptions: Vec<Subscription>,
2475    buffer: Entity<MultiBuffer>,
2476    range: Range<Anchor>,
2477    initial_transaction_id: Option<TransactionId>,
2478    telemetry: Arc<Telemetry>,
2479    builder: Arc<PromptBuilder>,
2480    is_insertion: bool,
2481}
2482
2483impl Codegen {
2484    pub fn new(
2485        buffer: Entity<MultiBuffer>,
2486        range: Range<Anchor>,
2487        initial_transaction_id: Option<TransactionId>,
2488        telemetry: Arc<Telemetry>,
2489        builder: Arc<PromptBuilder>,
2490        cx: &mut Context<Self>,
2491    ) -> Self {
2492        let codegen = cx.new(|cx| {
2493            CodegenAlternative::new(
2494                buffer.clone(),
2495                range.clone(),
2496                false,
2497                Some(telemetry.clone()),
2498                builder.clone(),
2499                cx,
2500            )
2501        });
2502        let mut this = Self {
2503            is_insertion: range.to_offset(&buffer.read(cx).snapshot(cx)).is_empty(),
2504            alternatives: vec![codegen],
2505            active_alternative: 0,
2506            seen_alternatives: HashSet::default(),
2507            subscriptions: Vec::new(),
2508            buffer,
2509            range,
2510            initial_transaction_id,
2511            telemetry,
2512            builder,
2513        };
2514        this.activate(0, cx);
2515        this
2516    }
2517
2518    fn subscribe_to_alternative(&mut self, cx: &mut Context<Self>) {
2519        let codegen = self.active_alternative().clone();
2520        self.subscriptions.clear();
2521        self.subscriptions
2522            .push(cx.observe(&codegen, |_, _, cx| cx.notify()));
2523        self.subscriptions
2524            .push(cx.subscribe(&codegen, |_, _, event, cx| cx.emit(*event)));
2525    }
2526
2527    fn active_alternative(&self) -> &Entity<CodegenAlternative> {
2528        &self.alternatives[self.active_alternative]
2529    }
2530
2531    fn status<'a>(&self, cx: &'a App) -> &'a CodegenStatus {
2532        &self.active_alternative().read(cx).status
2533    }
2534
2535    fn alternative_count(&self, cx: &App) -> usize {
2536        LanguageModelRegistry::read_global(cx)
2537            .inline_alternative_models()
2538            .len()
2539            + 1
2540    }
2541
2542    pub fn cycle_prev(&mut self, cx: &mut Context<Self>) {
2543        let next_active_ix = if self.active_alternative == 0 {
2544            self.alternatives.len() - 1
2545        } else {
2546            self.active_alternative - 1
2547        };
2548        self.activate(next_active_ix, cx);
2549    }
2550
2551    pub fn cycle_next(&mut self, cx: &mut Context<Self>) {
2552        let next_active_ix = (self.active_alternative + 1) % self.alternatives.len();
2553        self.activate(next_active_ix, cx);
2554    }
2555
2556    fn activate(&mut self, index: usize, cx: &mut Context<Self>) {
2557        self.active_alternative()
2558            .update(cx, |codegen, cx| codegen.set_active(false, cx));
2559        self.seen_alternatives.insert(index);
2560        self.active_alternative = index;
2561        self.active_alternative()
2562            .update(cx, |codegen, cx| codegen.set_active(true, cx));
2563        self.subscribe_to_alternative(cx);
2564        cx.notify();
2565    }
2566
2567    pub fn start(
2568        &mut self,
2569        user_prompt: String,
2570        assistant_panel_context: Option<LanguageModelRequest>,
2571        cx: &mut Context<Self>,
2572    ) -> Result<()> {
2573        let alternative_models = LanguageModelRegistry::read_global(cx)
2574            .inline_alternative_models()
2575            .to_vec();
2576
2577        self.active_alternative()
2578            .update(cx, |alternative, cx| alternative.undo(cx));
2579        self.activate(0, cx);
2580        self.alternatives.truncate(1);
2581
2582        for _ in 0..alternative_models.len() {
2583            self.alternatives.push(cx.new(|cx| {
2584                CodegenAlternative::new(
2585                    self.buffer.clone(),
2586                    self.range.clone(),
2587                    false,
2588                    Some(self.telemetry.clone()),
2589                    self.builder.clone(),
2590                    cx,
2591                )
2592            }));
2593        }
2594
2595        let primary_model = LanguageModelRegistry::read_global(cx)
2596            .active_model()
2597            .context("no active model")?;
2598
2599        for (model, alternative) in iter::once(primary_model)
2600            .chain(alternative_models)
2601            .zip(&self.alternatives)
2602        {
2603            alternative.update(cx, |alternative, cx| {
2604                alternative.start(
2605                    user_prompt.clone(),
2606                    assistant_panel_context.clone(),
2607                    model.clone(),
2608                    cx,
2609                )
2610            })?;
2611        }
2612
2613        Ok(())
2614    }
2615
2616    pub fn stop(&mut self, cx: &mut Context<Self>) {
2617        for codegen in &self.alternatives {
2618            codegen.update(cx, |codegen, cx| codegen.stop(cx));
2619        }
2620    }
2621
2622    pub fn undo(&mut self, cx: &mut Context<Self>) {
2623        self.active_alternative()
2624            .update(cx, |codegen, cx| codegen.undo(cx));
2625
2626        self.buffer.update(cx, |buffer, cx| {
2627            if let Some(transaction_id) = self.initial_transaction_id.take() {
2628                buffer.undo_transaction(transaction_id, cx);
2629                buffer.refresh_preview(cx);
2630            }
2631        });
2632    }
2633
2634    pub fn count_tokens(
2635        &self,
2636        user_prompt: String,
2637        assistant_panel_context: Option<LanguageModelRequest>,
2638        cx: &App,
2639    ) -> BoxFuture<'static, Result<TokenCounts>> {
2640        self.active_alternative()
2641            .read(cx)
2642            .count_tokens(user_prompt, assistant_panel_context, cx)
2643    }
2644
2645    pub fn buffer(&self, cx: &App) -> Entity<MultiBuffer> {
2646        self.active_alternative().read(cx).buffer.clone()
2647    }
2648
2649    pub fn old_buffer(&self, cx: &App) -> Entity<Buffer> {
2650        self.active_alternative().read(cx).old_buffer.clone()
2651    }
2652
2653    pub fn snapshot(&self, cx: &App) -> MultiBufferSnapshot {
2654        self.active_alternative().read(cx).snapshot.clone()
2655    }
2656
2657    pub fn edit_position(&self, cx: &App) -> Option<Anchor> {
2658        self.active_alternative().read(cx).edit_position
2659    }
2660
2661    fn diff<'a>(&self, cx: &'a App) -> &'a Diff {
2662        &self.active_alternative().read(cx).diff
2663    }
2664
2665    pub fn last_equal_ranges<'a>(&self, cx: &'a App) -> &'a [Range<Anchor>] {
2666        self.active_alternative().read(cx).last_equal_ranges()
2667    }
2668}
2669
2670impl EventEmitter<CodegenEvent> for Codegen {}
2671
2672pub struct CodegenAlternative {
2673    buffer: Entity<MultiBuffer>,
2674    old_buffer: Entity<Buffer>,
2675    snapshot: MultiBufferSnapshot,
2676    edit_position: Option<Anchor>,
2677    range: Range<Anchor>,
2678    last_equal_ranges: Vec<Range<Anchor>>,
2679    transformation_transaction_id: Option<TransactionId>,
2680    status: CodegenStatus,
2681    generation: Task<()>,
2682    diff: Diff,
2683    telemetry: Option<Arc<Telemetry>>,
2684    _subscription: gpui::Subscription,
2685    builder: Arc<PromptBuilder>,
2686    active: bool,
2687    edits: Vec<(Range<Anchor>, String)>,
2688    line_operations: Vec<LineOperation>,
2689    request: Option<LanguageModelRequest>,
2690    elapsed_time: Option<f64>,
2691    completion: Option<String>,
2692    message_id: Option<String>,
2693}
2694
2695enum CodegenStatus {
2696    Idle,
2697    Pending,
2698    Done,
2699    Error(anyhow::Error),
2700}
2701
2702#[derive(Default)]
2703struct Diff {
2704    deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)>,
2705    inserted_row_ranges: Vec<Range<Anchor>>,
2706}
2707
2708impl Diff {
2709    fn is_empty(&self) -> bool {
2710        self.deleted_row_ranges.is_empty() && self.inserted_row_ranges.is_empty()
2711    }
2712}
2713
2714impl EventEmitter<CodegenEvent> for CodegenAlternative {}
2715
2716impl CodegenAlternative {
2717    pub fn new(
2718        multi_buffer: Entity<MultiBuffer>,
2719        range: Range<Anchor>,
2720        active: bool,
2721        telemetry: Option<Arc<Telemetry>>,
2722        builder: Arc<PromptBuilder>,
2723        cx: &mut Context<Self>,
2724    ) -> Self {
2725        let snapshot = multi_buffer.read(cx).snapshot(cx);
2726
2727        let (buffer, _, _) = snapshot
2728            .range_to_buffer_ranges(range.clone())
2729            .pop()
2730            .unwrap();
2731        let old_buffer = cx.new(|cx| {
2732            let text = buffer.as_rope().clone();
2733            let line_ending = buffer.line_ending();
2734            let language = buffer.language().cloned();
2735            let language_registry = multi_buffer
2736                .read(cx)
2737                .buffer(buffer.remote_id())
2738                .unwrap()
2739                .read(cx)
2740                .language_registry();
2741
2742            let mut buffer = Buffer::local_normalized(text, line_ending, cx);
2743            buffer.set_language(language, cx);
2744            if let Some(language_registry) = language_registry {
2745                buffer.set_language_registry(language_registry)
2746            }
2747            buffer
2748        });
2749
2750        Self {
2751            buffer: multi_buffer.clone(),
2752            old_buffer,
2753            edit_position: None,
2754            message_id: None,
2755            snapshot,
2756            last_equal_ranges: Default::default(),
2757            transformation_transaction_id: None,
2758            status: CodegenStatus::Idle,
2759            generation: Task::ready(()),
2760            diff: Diff::default(),
2761            telemetry,
2762            _subscription: cx.subscribe(&multi_buffer, Self::handle_buffer_event),
2763            builder,
2764            active,
2765            edits: Vec::new(),
2766            line_operations: Vec::new(),
2767            range,
2768            request: None,
2769            elapsed_time: None,
2770            completion: None,
2771        }
2772    }
2773
2774    fn set_active(&mut self, active: bool, cx: &mut Context<Self>) {
2775        if active != self.active {
2776            self.active = active;
2777
2778            if self.active {
2779                let edits = self.edits.clone();
2780                self.apply_edits(edits, cx);
2781                if matches!(self.status, CodegenStatus::Pending) {
2782                    let line_operations = self.line_operations.clone();
2783                    self.reapply_line_based_diff(line_operations, cx);
2784                } else {
2785                    self.reapply_batch_diff(cx).detach();
2786                }
2787            } else if let Some(transaction_id) = self.transformation_transaction_id.take() {
2788                self.buffer.update(cx, |buffer, cx| {
2789                    buffer.undo_transaction(transaction_id, cx);
2790                    buffer.forget_transaction(transaction_id, cx);
2791                });
2792            }
2793        }
2794    }
2795
2796    fn handle_buffer_event(
2797        &mut self,
2798        _buffer: Entity<MultiBuffer>,
2799        event: &multi_buffer::Event,
2800        cx: &mut Context<Self>,
2801    ) {
2802        if let multi_buffer::Event::TransactionUndone { transaction_id } = event {
2803            if self.transformation_transaction_id == Some(*transaction_id) {
2804                self.transformation_transaction_id = None;
2805                self.generation = Task::ready(());
2806                cx.emit(CodegenEvent::Undone);
2807            }
2808        }
2809    }
2810
2811    pub fn last_equal_ranges(&self) -> &[Range<Anchor>] {
2812        &self.last_equal_ranges
2813    }
2814
2815    pub fn count_tokens(
2816        &self,
2817        user_prompt: String,
2818        assistant_panel_context: Option<LanguageModelRequest>,
2819        cx: &App,
2820    ) -> BoxFuture<'static, Result<TokenCounts>> {
2821        if let Some(model) = LanguageModelRegistry::read_global(cx).active_model() {
2822            let request = self.build_request(user_prompt, assistant_panel_context.clone(), cx);
2823            match request {
2824                Ok(request) => {
2825                    let total_count = model.count_tokens(request.clone(), cx);
2826                    let assistant_panel_count = assistant_panel_context
2827                        .map(|context| model.count_tokens(context, cx))
2828                        .unwrap_or_else(|| future::ready(Ok(0)).boxed());
2829
2830                    async move {
2831                        Ok(TokenCounts {
2832                            total: total_count.await?,
2833                            assistant_panel: assistant_panel_count.await?,
2834                        })
2835                    }
2836                    .boxed()
2837                }
2838                Err(error) => futures::future::ready(Err(error)).boxed(),
2839            }
2840        } else {
2841            future::ready(Err(anyhow!("no active model"))).boxed()
2842        }
2843    }
2844
2845    pub fn start(
2846        &mut self,
2847        user_prompt: String,
2848        assistant_panel_context: Option<LanguageModelRequest>,
2849        model: Arc<dyn LanguageModel>,
2850        cx: &mut Context<Self>,
2851    ) -> Result<()> {
2852        if let Some(transformation_transaction_id) = self.transformation_transaction_id.take() {
2853            self.buffer.update(cx, |buffer, cx| {
2854                buffer.undo_transaction(transformation_transaction_id, cx);
2855            });
2856        }
2857
2858        self.edit_position = Some(self.range.start.bias_right(&self.snapshot));
2859
2860        let api_key = model.api_key(cx);
2861        let telemetry_id = model.telemetry_id();
2862        let provider_id = model.provider_id();
2863        let stream: LocalBoxFuture<Result<LanguageModelTextStream>> =
2864            if user_prompt.trim().to_lowercase() == "delete" {
2865                async { Ok(LanguageModelTextStream::default()) }.boxed_local()
2866            } else {
2867                let request = self.build_request(user_prompt, assistant_panel_context, cx)?;
2868                self.request = Some(request.clone());
2869
2870                cx.spawn(|_, cx| async move { model.stream_completion_text(request, &cx).await })
2871                    .boxed_local()
2872            };
2873        self.handle_stream(telemetry_id, provider_id.to_string(), api_key, stream, cx);
2874        Ok(())
2875    }
2876
2877    fn build_request(
2878        &self,
2879        user_prompt: String,
2880        assistant_panel_context: Option<LanguageModelRequest>,
2881        cx: &App,
2882    ) -> Result<LanguageModelRequest> {
2883        let buffer = self.buffer.read(cx).snapshot(cx);
2884        let language = buffer.language_at(self.range.start);
2885        let language_name = if let Some(language) = language.as_ref() {
2886            if Arc::ptr_eq(language, &language::PLAIN_TEXT) {
2887                None
2888            } else {
2889                Some(language.name())
2890            }
2891        } else {
2892            None
2893        };
2894
2895        let language_name = language_name.as_ref();
2896        let start = buffer.point_to_buffer_offset(self.range.start);
2897        let end = buffer.point_to_buffer_offset(self.range.end);
2898        let (buffer, range) = if let Some((start, end)) = start.zip(end) {
2899            let (start_buffer, start_buffer_offset) = start;
2900            let (end_buffer, end_buffer_offset) = end;
2901            if start_buffer.remote_id() == end_buffer.remote_id() {
2902                (start_buffer.clone(), start_buffer_offset..end_buffer_offset)
2903            } else {
2904                return Err(anyhow::anyhow!("invalid transformation range"));
2905            }
2906        } else {
2907            return Err(anyhow::anyhow!("invalid transformation range"));
2908        };
2909
2910        let prompt = self
2911            .builder
2912            .generate_inline_transformation_prompt(user_prompt, language_name, buffer, range)
2913            .map_err(|e| anyhow::anyhow!("Failed to generate content prompt: {}", e))?;
2914
2915        let mut messages = Vec::new();
2916        if let Some(context_request) = assistant_panel_context {
2917            messages = context_request.messages;
2918        }
2919
2920        messages.push(LanguageModelRequestMessage {
2921            role: Role::User,
2922            content: vec![prompt.into()],
2923            cache: false,
2924        });
2925
2926        Ok(LanguageModelRequest {
2927            messages,
2928            tools: Vec::new(),
2929            stop: Vec::new(),
2930            temperature: None,
2931        })
2932    }
2933
2934    pub fn handle_stream(
2935        &mut self,
2936        model_telemetry_id: String,
2937        model_provider_id: String,
2938        model_api_key: Option<String>,
2939        stream: impl 'static + Future<Output = Result<LanguageModelTextStream>>,
2940        cx: &mut Context<Self>,
2941    ) {
2942        let start_time = Instant::now();
2943        let snapshot = self.snapshot.clone();
2944        let selected_text = snapshot
2945            .text_for_range(self.range.start..self.range.end)
2946            .collect::<Rope>();
2947
2948        let selection_start = self.range.start.to_point(&snapshot);
2949
2950        // Start with the indentation of the first line in the selection
2951        let mut suggested_line_indent = snapshot
2952            .suggested_indents(selection_start.row..=selection_start.row, cx)
2953            .into_values()
2954            .next()
2955            .unwrap_or_else(|| snapshot.indent_size_for_line(MultiBufferRow(selection_start.row)));
2956
2957        // If the first line in the selection does not have indentation, check the following lines
2958        if suggested_line_indent.len == 0 && suggested_line_indent.kind == IndentKind::Space {
2959            for row in selection_start.row..=self.range.end.to_point(&snapshot).row {
2960                let line_indent = snapshot.indent_size_for_line(MultiBufferRow(row));
2961                // Prefer tabs if a line in the selection uses tabs as indentation
2962                if line_indent.kind == IndentKind::Tab {
2963                    suggested_line_indent.kind = IndentKind::Tab;
2964                    break;
2965                }
2966            }
2967        }
2968
2969        let http_client = cx.http_client().clone();
2970        let telemetry = self.telemetry.clone();
2971        let language_name = {
2972            let multibuffer = self.buffer.read(cx);
2973            let snapshot = multibuffer.snapshot(cx);
2974            let ranges = snapshot.range_to_buffer_ranges(self.range.clone());
2975            ranges
2976                .first()
2977                .and_then(|(buffer, _, _)| buffer.language())
2978                .map(|language| language.name())
2979        };
2980
2981        self.diff = Diff::default();
2982        self.status = CodegenStatus::Pending;
2983        let mut edit_start = self.range.start.to_offset(&snapshot);
2984        let completion = Arc::new(Mutex::new(String::new()));
2985        let completion_clone = completion.clone();
2986
2987        self.generation = cx.spawn(|codegen, mut cx| {
2988            async move {
2989                let stream = stream.await;
2990                let message_id = stream
2991                    .as_ref()
2992                    .ok()
2993                    .and_then(|stream| stream.message_id.clone());
2994                let generate = async {
2995                    let (mut diff_tx, mut diff_rx) = mpsc::channel(1);
2996                    let executor = cx.background_executor().clone();
2997                    let message_id = message_id.clone();
2998                    let line_based_stream_diff: Task<anyhow::Result<()>> =
2999                        cx.background_spawn(async move {
3000                            let mut response_latency = None;
3001                            let request_start = Instant::now();
3002                            let diff = async {
3003                                let chunks = StripInvalidSpans::new(stream?.stream);
3004                                futures::pin_mut!(chunks);
3005                                let mut diff = StreamingDiff::new(selected_text.to_string());
3006                                let mut line_diff = LineDiff::default();
3007
3008                                let mut new_text = String::new();
3009                                let mut base_indent = None;
3010                                let mut line_indent = None;
3011                                let mut first_line = true;
3012
3013                                while let Some(chunk) = chunks.next().await {
3014                                    if response_latency.is_none() {
3015                                        response_latency = Some(request_start.elapsed());
3016                                    }
3017                                    let chunk = chunk?;
3018                                    completion_clone.lock().push_str(&chunk);
3019
3020                                    let mut lines = chunk.split('\n').peekable();
3021                                    while let Some(line) = lines.next() {
3022                                        new_text.push_str(line);
3023                                        if line_indent.is_none() {
3024                                            if let Some(non_whitespace_ch_ix) =
3025                                                new_text.find(|ch: char| !ch.is_whitespace())
3026                                            {
3027                                                line_indent = Some(non_whitespace_ch_ix);
3028                                                base_indent = base_indent.or(line_indent);
3029
3030                                                let line_indent = line_indent.unwrap();
3031                                                let base_indent = base_indent.unwrap();
3032                                                let indent_delta =
3033                                                    line_indent as i32 - base_indent as i32;
3034                                                let mut corrected_indent_len = cmp::max(
3035                                                    0,
3036                                                    suggested_line_indent.len as i32 + indent_delta,
3037                                                )
3038                                                    as usize;
3039                                                if first_line {
3040                                                    corrected_indent_len = corrected_indent_len
3041                                                        .saturating_sub(
3042                                                            selection_start.column as usize,
3043                                                        );
3044                                                }
3045
3046                                                let indent_char = suggested_line_indent.char();
3047                                                let mut indent_buffer = [0; 4];
3048                                                let indent_str =
3049                                                    indent_char.encode_utf8(&mut indent_buffer);
3050                                                new_text.replace_range(
3051                                                    ..line_indent,
3052                                                    &indent_str.repeat(corrected_indent_len),
3053                                                );
3054                                            }
3055                                        }
3056
3057                                        if line_indent.is_some() {
3058                                            let char_ops = diff.push_new(&new_text);
3059                                            line_diff
3060                                                .push_char_operations(&char_ops, &selected_text);
3061                                            diff_tx
3062                                                .send((char_ops, line_diff.line_operations()))
3063                                                .await?;
3064                                            new_text.clear();
3065                                        }
3066
3067                                        if lines.peek().is_some() {
3068                                            let char_ops = diff.push_new("\n");
3069                                            line_diff
3070                                                .push_char_operations(&char_ops, &selected_text);
3071                                            diff_tx
3072                                                .send((char_ops, line_diff.line_operations()))
3073                                                .await?;
3074                                            if line_indent.is_none() {
3075                                                // Don't write out the leading indentation in empty lines on the next line
3076                                                // This is the case where the above if statement didn't clear the buffer
3077                                                new_text.clear();
3078                                            }
3079                                            line_indent = None;
3080                                            first_line = false;
3081                                        }
3082                                    }
3083                                }
3084
3085                                let mut char_ops = diff.push_new(&new_text);
3086                                char_ops.extend(diff.finish());
3087                                line_diff.push_char_operations(&char_ops, &selected_text);
3088                                line_diff.finish(&selected_text);
3089                                diff_tx
3090                                    .send((char_ops, line_diff.line_operations()))
3091                                    .await?;
3092
3093                                anyhow::Ok(())
3094                            };
3095
3096                            let result = diff.await;
3097
3098                            let error_message =
3099                                result.as_ref().err().map(|error| error.to_string());
3100                            report_assistant_event(
3101                                AssistantEvent {
3102                                    conversation_id: None,
3103                                    message_id,
3104                                    kind: AssistantKind::Inline,
3105                                    phase: AssistantPhase::Response,
3106                                    model: model_telemetry_id,
3107                                    model_provider: model_provider_id.to_string(),
3108                                    response_latency,
3109                                    error_message,
3110                                    language_name: language_name.map(|name| name.to_proto()),
3111                                },
3112                                telemetry,
3113                                http_client,
3114                                model_api_key,
3115                                &executor,
3116                            );
3117
3118                            result?;
3119                            Ok(())
3120                        });
3121
3122                    while let Some((char_ops, line_ops)) = diff_rx.next().await {
3123                        codegen.update(&mut cx, |codegen, cx| {
3124                            codegen.last_equal_ranges.clear();
3125
3126                            let edits = char_ops
3127                                .into_iter()
3128                                .filter_map(|operation| match operation {
3129                                    CharOperation::Insert { text } => {
3130                                        let edit_start = snapshot.anchor_after(edit_start);
3131                                        Some((edit_start..edit_start, text))
3132                                    }
3133                                    CharOperation::Delete { bytes } => {
3134                                        let edit_end = edit_start + bytes;
3135                                        let edit_range = snapshot.anchor_after(edit_start)
3136                                            ..snapshot.anchor_before(edit_end);
3137                                        edit_start = edit_end;
3138                                        Some((edit_range, String::new()))
3139                                    }
3140                                    CharOperation::Keep { bytes } => {
3141                                        let edit_end = edit_start + bytes;
3142                                        let edit_range = snapshot.anchor_after(edit_start)
3143                                            ..snapshot.anchor_before(edit_end);
3144                                        edit_start = edit_end;
3145                                        codegen.last_equal_ranges.push(edit_range);
3146                                        None
3147                                    }
3148                                })
3149                                .collect::<Vec<_>>();
3150
3151                            if codegen.active {
3152                                codegen.apply_edits(edits.iter().cloned(), cx);
3153                                codegen.reapply_line_based_diff(line_ops.iter().cloned(), cx);
3154                            }
3155                            codegen.edits.extend(edits);
3156                            codegen.line_operations = line_ops;
3157                            codegen.edit_position = Some(snapshot.anchor_after(edit_start));
3158
3159                            cx.notify();
3160                        })?;
3161                    }
3162
3163                    // Streaming stopped and we have the new text in the buffer, and a line-based diff applied for the whole new buffer.
3164                    // That diff is not what a regular diff is and might look unexpected, ergo apply a regular diff.
3165                    // It's fine to apply even if the rest of the line diffing fails, as no more hunks are coming through `diff_rx`.
3166                    let batch_diff_task =
3167                        codegen.update(&mut cx, |codegen, cx| codegen.reapply_batch_diff(cx))?;
3168                    let (line_based_stream_diff, ()) =
3169                        join!(line_based_stream_diff, batch_diff_task);
3170                    line_based_stream_diff?;
3171
3172                    anyhow::Ok(())
3173                };
3174
3175                let result = generate.await;
3176                let elapsed_time = start_time.elapsed().as_secs_f64();
3177
3178                codegen
3179                    .update(&mut cx, |this, cx| {
3180                        this.message_id = message_id;
3181                        this.last_equal_ranges.clear();
3182                        if let Err(error) = result {
3183                            this.status = CodegenStatus::Error(error);
3184                        } else {
3185                            this.status = CodegenStatus::Done;
3186                        }
3187                        this.elapsed_time = Some(elapsed_time);
3188                        this.completion = Some(completion.lock().clone());
3189                        cx.emit(CodegenEvent::Finished);
3190                        cx.notify();
3191                    })
3192                    .ok();
3193            }
3194        });
3195        cx.notify();
3196    }
3197
3198    pub fn stop(&mut self, cx: &mut Context<Self>) {
3199        self.last_equal_ranges.clear();
3200        if self.diff.is_empty() {
3201            self.status = CodegenStatus::Idle;
3202        } else {
3203            self.status = CodegenStatus::Done;
3204        }
3205        self.generation = Task::ready(());
3206        cx.emit(CodegenEvent::Finished);
3207        cx.notify();
3208    }
3209
3210    pub fn undo(&mut self, cx: &mut Context<Self>) {
3211        self.buffer.update(cx, |buffer, cx| {
3212            if let Some(transaction_id) = self.transformation_transaction_id.take() {
3213                buffer.undo_transaction(transaction_id, cx);
3214                buffer.refresh_preview(cx);
3215            }
3216        });
3217    }
3218
3219    fn apply_edits(
3220        &mut self,
3221        edits: impl IntoIterator<Item = (Range<Anchor>, String)>,
3222        cx: &mut Context<CodegenAlternative>,
3223    ) {
3224        let transaction = self.buffer.update(cx, |buffer, cx| {
3225            // Avoid grouping assistant edits with user edits.
3226            buffer.finalize_last_transaction(cx);
3227            buffer.start_transaction(cx);
3228            buffer.edit(edits, None, cx);
3229            buffer.end_transaction(cx)
3230        });
3231
3232        if let Some(transaction) = transaction {
3233            if let Some(first_transaction) = self.transformation_transaction_id {
3234                // Group all assistant edits into the first transaction.
3235                self.buffer.update(cx, |buffer, cx| {
3236                    buffer.merge_transactions(transaction, first_transaction, cx)
3237                });
3238            } else {
3239                self.transformation_transaction_id = Some(transaction);
3240                self.buffer
3241                    .update(cx, |buffer, cx| buffer.finalize_last_transaction(cx));
3242            }
3243        }
3244    }
3245
3246    fn reapply_line_based_diff(
3247        &mut self,
3248        line_operations: impl IntoIterator<Item = LineOperation>,
3249        cx: &mut Context<Self>,
3250    ) {
3251        let old_snapshot = self.snapshot.clone();
3252        let old_range = self.range.to_point(&old_snapshot);
3253        let new_snapshot = self.buffer.read(cx).snapshot(cx);
3254        let new_range = self.range.to_point(&new_snapshot);
3255
3256        let mut old_row = old_range.start.row;
3257        let mut new_row = new_range.start.row;
3258
3259        self.diff.deleted_row_ranges.clear();
3260        self.diff.inserted_row_ranges.clear();
3261        for operation in line_operations {
3262            match operation {
3263                LineOperation::Keep { lines } => {
3264                    old_row += lines;
3265                    new_row += lines;
3266                }
3267                LineOperation::Delete { lines } => {
3268                    let old_end_row = old_row + lines - 1;
3269                    let new_row = new_snapshot.anchor_before(Point::new(new_row, 0));
3270
3271                    if let Some((_, last_deleted_row_range)) =
3272                        self.diff.deleted_row_ranges.last_mut()
3273                    {
3274                        if *last_deleted_row_range.end() + 1 == old_row {
3275                            *last_deleted_row_range = *last_deleted_row_range.start()..=old_end_row;
3276                        } else {
3277                            self.diff
3278                                .deleted_row_ranges
3279                                .push((new_row, old_row..=old_end_row));
3280                        }
3281                    } else {
3282                        self.diff
3283                            .deleted_row_ranges
3284                            .push((new_row, old_row..=old_end_row));
3285                    }
3286
3287                    old_row += lines;
3288                }
3289                LineOperation::Insert { lines } => {
3290                    let new_end_row = new_row + lines - 1;
3291                    let start = new_snapshot.anchor_before(Point::new(new_row, 0));
3292                    let end = new_snapshot.anchor_before(Point::new(
3293                        new_end_row,
3294                        new_snapshot.line_len(MultiBufferRow(new_end_row)),
3295                    ));
3296                    self.diff.inserted_row_ranges.push(start..end);
3297                    new_row += lines;
3298                }
3299            }
3300
3301            cx.notify();
3302        }
3303    }
3304
3305    fn reapply_batch_diff(&mut self, cx: &mut Context<Self>) -> Task<()> {
3306        let old_snapshot = self.snapshot.clone();
3307        let old_range = self.range.to_point(&old_snapshot);
3308        let new_snapshot = self.buffer.read(cx).snapshot(cx);
3309        let new_range = self.range.to_point(&new_snapshot);
3310
3311        cx.spawn(|codegen, mut cx| async move {
3312            let (deleted_row_ranges, inserted_row_ranges) = cx
3313                .background_spawn(async move {
3314                    let old_text = old_snapshot
3315                        .text_for_range(
3316                            Point::new(old_range.start.row, 0)
3317                                ..Point::new(
3318                                    old_range.end.row,
3319                                    old_snapshot.line_len(MultiBufferRow(old_range.end.row)),
3320                                ),
3321                        )
3322                        .collect::<String>();
3323                    let new_text = new_snapshot
3324                        .text_for_range(
3325                            Point::new(new_range.start.row, 0)
3326                                ..Point::new(
3327                                    new_range.end.row,
3328                                    new_snapshot.line_len(MultiBufferRow(new_range.end.row)),
3329                                ),
3330                        )
3331                        .collect::<String>();
3332
3333                    let old_start_row = old_range.start.row;
3334                    let new_start_row = new_range.start.row;
3335                    let mut deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)> = Vec::new();
3336                    let mut inserted_row_ranges = Vec::new();
3337                    for (old_rows, new_rows) in line_diff(&old_text, &new_text) {
3338                        let old_rows = old_start_row + old_rows.start..old_start_row + old_rows.end;
3339                        let new_rows = new_start_row + new_rows.start..new_start_row + new_rows.end;
3340                        if !old_rows.is_empty() {
3341                            deleted_row_ranges.push((
3342                                new_snapshot.anchor_before(Point::new(new_rows.start, 0)),
3343                                old_rows.start..=old_rows.end - 1,
3344                            ));
3345                        }
3346                        if !new_rows.is_empty() {
3347                            let start = new_snapshot.anchor_before(Point::new(new_rows.start, 0));
3348                            let new_end_row = new_rows.end - 1;
3349                            let end = new_snapshot.anchor_before(Point::new(
3350                                new_end_row,
3351                                new_snapshot.line_len(MultiBufferRow(new_end_row)),
3352                            ));
3353                            inserted_row_ranges.push(start..end);
3354                        }
3355                    }
3356                    (deleted_row_ranges, inserted_row_ranges)
3357                })
3358                .await;
3359
3360            codegen
3361                .update(&mut cx, |codegen, cx| {
3362                    codegen.diff.deleted_row_ranges = deleted_row_ranges;
3363                    codegen.diff.inserted_row_ranges = inserted_row_ranges;
3364                    cx.notify();
3365                })
3366                .ok();
3367        })
3368    }
3369}
3370
3371struct StripInvalidSpans<T> {
3372    stream: T,
3373    stream_done: bool,
3374    buffer: String,
3375    first_line: bool,
3376    line_end: bool,
3377    starts_with_code_block: bool,
3378}
3379
3380impl<T> StripInvalidSpans<T>
3381where
3382    T: Stream<Item = Result<String>>,
3383{
3384    fn new(stream: T) -> Self {
3385        Self {
3386            stream,
3387            stream_done: false,
3388            buffer: String::new(),
3389            first_line: true,
3390            line_end: false,
3391            starts_with_code_block: false,
3392        }
3393    }
3394}
3395
3396impl<T> Stream for StripInvalidSpans<T>
3397where
3398    T: Stream<Item = Result<String>>,
3399{
3400    type Item = Result<String>;
3401
3402    fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Option<Self::Item>> {
3403        const CODE_BLOCK_DELIMITER: &str = "```";
3404        const CURSOR_SPAN: &str = "<|CURSOR|>";
3405
3406        let this = unsafe { self.get_unchecked_mut() };
3407        loop {
3408            if !this.stream_done {
3409                let mut stream = unsafe { Pin::new_unchecked(&mut this.stream) };
3410                match stream.as_mut().poll_next(cx) {
3411                    Poll::Ready(Some(Ok(chunk))) => {
3412                        this.buffer.push_str(&chunk);
3413                    }
3414                    Poll::Ready(Some(Err(error))) => return Poll::Ready(Some(Err(error))),
3415                    Poll::Ready(None) => {
3416                        this.stream_done = true;
3417                    }
3418                    Poll::Pending => return Poll::Pending,
3419                }
3420            }
3421
3422            let mut chunk = String::new();
3423            let mut consumed = 0;
3424            if !this.buffer.is_empty() {
3425                let mut lines = this.buffer.split('\n').enumerate().peekable();
3426                while let Some((line_ix, line)) = lines.next() {
3427                    if line_ix > 0 {
3428                        this.first_line = false;
3429                    }
3430
3431                    if this.first_line {
3432                        let trimmed_line = line.trim();
3433                        if lines.peek().is_some() {
3434                            if trimmed_line.starts_with(CODE_BLOCK_DELIMITER) {
3435                                consumed += line.len() + 1;
3436                                this.starts_with_code_block = true;
3437                                continue;
3438                            }
3439                        } else if trimmed_line.is_empty()
3440                            || prefixes(CODE_BLOCK_DELIMITER)
3441                                .any(|prefix| trimmed_line.starts_with(prefix))
3442                        {
3443                            break;
3444                        }
3445                    }
3446
3447                    let line_without_cursor = line.replace(CURSOR_SPAN, "");
3448                    if lines.peek().is_some() {
3449                        if this.line_end {
3450                            chunk.push('\n');
3451                        }
3452
3453                        chunk.push_str(&line_without_cursor);
3454                        this.line_end = true;
3455                        consumed += line.len() + 1;
3456                    } else if this.stream_done {
3457                        if !this.starts_with_code_block
3458                            || !line_without_cursor.trim().ends_with(CODE_BLOCK_DELIMITER)
3459                        {
3460                            if this.line_end {
3461                                chunk.push('\n');
3462                            }
3463
3464                            chunk.push_str(&line);
3465                        }
3466
3467                        consumed += line.len();
3468                    } else {
3469                        let trimmed_line = line.trim();
3470                        if trimmed_line.is_empty()
3471                            || prefixes(CURSOR_SPAN).any(|prefix| trimmed_line.ends_with(prefix))
3472                            || prefixes(CODE_BLOCK_DELIMITER)
3473                                .any(|prefix| trimmed_line.ends_with(prefix))
3474                        {
3475                            break;
3476                        } else {
3477                            if this.line_end {
3478                                chunk.push('\n');
3479                                this.line_end = false;
3480                            }
3481
3482                            chunk.push_str(&line_without_cursor);
3483                            consumed += line.len();
3484                        }
3485                    }
3486                }
3487            }
3488
3489            this.buffer = this.buffer.split_off(consumed);
3490            if !chunk.is_empty() {
3491                return Poll::Ready(Some(Ok(chunk)));
3492            } else if this.stream_done {
3493                return Poll::Ready(None);
3494            }
3495        }
3496    }
3497}
3498
3499struct AssistantCodeActionProvider {
3500    editor: WeakEntity<Editor>,
3501    workspace: WeakEntity<Workspace>,
3502}
3503
3504const ASSISTANT_CODE_ACTION_PROVIDER_ID: &str = "assistant";
3505
3506impl CodeActionProvider for AssistantCodeActionProvider {
3507    fn id(&self) -> Arc<str> {
3508        ASSISTANT_CODE_ACTION_PROVIDER_ID.into()
3509    }
3510
3511    fn code_actions(
3512        &self,
3513        buffer: &Entity<Buffer>,
3514        range: Range<text::Anchor>,
3515        _: &mut Window,
3516        cx: &mut App,
3517    ) -> Task<Result<Vec<CodeAction>>> {
3518        if !AssistantSettings::get_global(cx).enabled {
3519            return Task::ready(Ok(Vec::new()));
3520        }
3521
3522        let snapshot = buffer.read(cx).snapshot();
3523        let mut range = range.to_point(&snapshot);
3524
3525        // Expand the range to line boundaries.
3526        range.start.column = 0;
3527        range.end.column = snapshot.line_len(range.end.row);
3528
3529        let mut has_diagnostics = false;
3530        for diagnostic in snapshot.diagnostics_in_range::<_, Point>(range.clone(), false) {
3531            range.start = cmp::min(range.start, diagnostic.range.start);
3532            range.end = cmp::max(range.end, diagnostic.range.end);
3533            has_diagnostics = true;
3534        }
3535        if has_diagnostics {
3536            if let Some(symbols_containing_start) = snapshot.symbols_containing(range.start, None) {
3537                if let Some(symbol) = symbols_containing_start.last() {
3538                    range.start = cmp::min(range.start, symbol.range.start.to_point(&snapshot));
3539                    range.end = cmp::max(range.end, symbol.range.end.to_point(&snapshot));
3540                }
3541            }
3542
3543            if let Some(symbols_containing_end) = snapshot.symbols_containing(range.end, None) {
3544                if let Some(symbol) = symbols_containing_end.last() {
3545                    range.start = cmp::min(range.start, symbol.range.start.to_point(&snapshot));
3546                    range.end = cmp::max(range.end, symbol.range.end.to_point(&snapshot));
3547                }
3548            }
3549
3550            Task::ready(Ok(vec![CodeAction {
3551                server_id: language::LanguageServerId(0),
3552                range: snapshot.anchor_before(range.start)..snapshot.anchor_after(range.end),
3553                lsp_action: lsp::CodeAction {
3554                    title: "Fix with Assistant".into(),
3555                    ..Default::default()
3556                },
3557            }]))
3558        } else {
3559            Task::ready(Ok(Vec::new()))
3560        }
3561    }
3562
3563    fn apply_code_action(
3564        &self,
3565        buffer: Entity<Buffer>,
3566        action: CodeAction,
3567        excerpt_id: ExcerptId,
3568        _push_to_history: bool,
3569        window: &mut Window,
3570        cx: &mut App,
3571    ) -> Task<Result<ProjectTransaction>> {
3572        let editor = self.editor.clone();
3573        let workspace = self.workspace.clone();
3574        window.spawn(cx, |mut cx| async move {
3575            let editor = editor.upgrade().context("editor was released")?;
3576            let range = editor
3577                .update(&mut cx, |editor, cx| {
3578                    editor.buffer().update(cx, |multibuffer, cx| {
3579                        let buffer = buffer.read(cx);
3580                        let multibuffer_snapshot = multibuffer.read(cx);
3581
3582                        let old_context_range =
3583                            multibuffer_snapshot.context_range_for_excerpt(excerpt_id)?;
3584                        let mut new_context_range = old_context_range.clone();
3585                        if action
3586                            .range
3587                            .start
3588                            .cmp(&old_context_range.start, buffer)
3589                            .is_lt()
3590                        {
3591                            new_context_range.start = action.range.start;
3592                        }
3593                        if action.range.end.cmp(&old_context_range.end, buffer).is_gt() {
3594                            new_context_range.end = action.range.end;
3595                        }
3596                        drop(multibuffer_snapshot);
3597
3598                        if new_context_range != old_context_range {
3599                            multibuffer.resize_excerpt(excerpt_id, new_context_range, cx);
3600                        }
3601
3602                        let multibuffer_snapshot = multibuffer.read(cx);
3603                        Some(
3604                            multibuffer_snapshot
3605                                .anchor_in_excerpt(excerpt_id, action.range.start)?
3606                                ..multibuffer_snapshot
3607                                    .anchor_in_excerpt(excerpt_id, action.range.end)?,
3608                        )
3609                    })
3610                })?
3611                .context("invalid range")?;
3612            let assistant_panel = workspace.update(&mut cx, |workspace, cx| {
3613                workspace
3614                    .panel::<AssistantPanel>(cx)
3615                    .context("assistant panel was released")
3616            })??;
3617
3618            cx.update_global(|assistant: &mut InlineAssistant, window, cx| {
3619                let assist_id = assistant.suggest_assist(
3620                    &editor,
3621                    range,
3622                    "Fix Diagnostics".into(),
3623                    None,
3624                    true,
3625                    Some(workspace),
3626                    Some(&assistant_panel),
3627                    window,
3628                    cx,
3629                );
3630                assistant.start_assist(assist_id, window, cx);
3631            })?;
3632
3633            Ok(ProjectTransaction::default())
3634        })
3635    }
3636}
3637
3638fn prefixes(text: &str) -> impl Iterator<Item = &str> {
3639    (0..text.len() - 1).map(|ix| &text[..ix + 1])
3640}
3641
3642fn merge_ranges(ranges: &mut Vec<Range<Anchor>>, buffer: &MultiBufferSnapshot) {
3643    ranges.sort_unstable_by(|a, b| {
3644        a.start
3645            .cmp(&b.start, buffer)
3646            .then_with(|| b.end.cmp(&a.end, buffer))
3647    });
3648
3649    let mut ix = 0;
3650    while ix + 1 < ranges.len() {
3651        let b = ranges[ix + 1].clone();
3652        let a = &mut ranges[ix];
3653        if a.end.cmp(&b.start, buffer).is_gt() {
3654            if a.end.cmp(&b.end, buffer).is_lt() {
3655                a.end = b.end;
3656            }
3657            ranges.remove(ix + 1);
3658        } else {
3659            ix += 1;
3660        }
3661    }
3662}
3663
3664#[cfg(test)]
3665mod tests {
3666    use super::*;
3667    use futures::stream::{self};
3668    use gpui::TestAppContext;
3669    use indoc::indoc;
3670    use language::{
3671        language_settings, tree_sitter_rust, Buffer, Language, LanguageConfig, LanguageMatcher,
3672        Point,
3673    };
3674    use language_model::LanguageModelRegistry;
3675    use rand::prelude::*;
3676    use serde::Serialize;
3677    use settings::SettingsStore;
3678    use std::{future, sync::Arc};
3679
3680    #[derive(Serialize)]
3681    pub struct DummyCompletionRequest {
3682        pub name: String,
3683    }
3684
3685    #[gpui::test(iterations = 10)]
3686    async fn test_transform_autoindent(cx: &mut TestAppContext, mut rng: StdRng) {
3687        cx.set_global(cx.update(SettingsStore::test));
3688        cx.update(language_model::LanguageModelRegistry::test);
3689        cx.update(language_settings::init);
3690
3691        let text = indoc! {"
3692            fn main() {
3693                let x = 0;
3694                for _ in 0..10 {
3695                    x += 1;
3696                }
3697            }
3698        "};
3699        let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
3700        let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
3701        let range = buffer.read_with(cx, |buffer, cx| {
3702            let snapshot = buffer.snapshot(cx);
3703            snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(4, 5))
3704        });
3705        let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
3706        let codegen = cx.new(|cx| {
3707            CodegenAlternative::new(
3708                buffer.clone(),
3709                range.clone(),
3710                true,
3711                None,
3712                prompt_builder,
3713                cx,
3714            )
3715        });
3716
3717        let chunks_tx = simulate_response_stream(codegen.clone(), cx);
3718
3719        let mut new_text = concat!(
3720            "       let mut x = 0;\n",
3721            "       while x < 10 {\n",
3722            "           x += 1;\n",
3723            "       }",
3724        );
3725        while !new_text.is_empty() {
3726            let max_len = cmp::min(new_text.len(), 10);
3727            let len = rng.gen_range(1..=max_len);
3728            let (chunk, suffix) = new_text.split_at(len);
3729            chunks_tx.unbounded_send(chunk.to_string()).unwrap();
3730            new_text = suffix;
3731            cx.background_executor.run_until_parked();
3732        }
3733        drop(chunks_tx);
3734        cx.background_executor.run_until_parked();
3735
3736        assert_eq!(
3737            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
3738            indoc! {"
3739                fn main() {
3740                    let mut x = 0;
3741                    while x < 10 {
3742                        x += 1;
3743                    }
3744                }
3745            "}
3746        );
3747    }
3748
3749    #[gpui::test(iterations = 10)]
3750    async fn test_autoindent_when_generating_past_indentation(
3751        cx: &mut TestAppContext,
3752        mut rng: StdRng,
3753    ) {
3754        cx.set_global(cx.update(SettingsStore::test));
3755        cx.update(language_settings::init);
3756
3757        let text = indoc! {"
3758            fn main() {
3759                le
3760            }
3761        "};
3762        let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
3763        let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
3764        let range = buffer.read_with(cx, |buffer, cx| {
3765            let snapshot = buffer.snapshot(cx);
3766            snapshot.anchor_before(Point::new(1, 6))..snapshot.anchor_after(Point::new(1, 6))
3767        });
3768        let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
3769        let codegen = cx.new(|cx| {
3770            CodegenAlternative::new(
3771                buffer.clone(),
3772                range.clone(),
3773                true,
3774                None,
3775                prompt_builder,
3776                cx,
3777            )
3778        });
3779
3780        let chunks_tx = simulate_response_stream(codegen.clone(), cx);
3781
3782        cx.background_executor.run_until_parked();
3783
3784        let mut new_text = concat!(
3785            "t mut x = 0;\n",
3786            "while x < 10 {\n",
3787            "    x += 1;\n",
3788            "}", //
3789        );
3790        while !new_text.is_empty() {
3791            let max_len = cmp::min(new_text.len(), 10);
3792            let len = rng.gen_range(1..=max_len);
3793            let (chunk, suffix) = new_text.split_at(len);
3794            chunks_tx.unbounded_send(chunk.to_string()).unwrap();
3795            new_text = suffix;
3796            cx.background_executor.run_until_parked();
3797        }
3798        drop(chunks_tx);
3799        cx.background_executor.run_until_parked();
3800
3801        assert_eq!(
3802            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
3803            indoc! {"
3804                fn main() {
3805                    let mut x = 0;
3806                    while x < 10 {
3807                        x += 1;
3808                    }
3809                }
3810            "}
3811        );
3812    }
3813
3814    #[gpui::test(iterations = 10)]
3815    async fn test_autoindent_when_generating_before_indentation(
3816        cx: &mut TestAppContext,
3817        mut rng: StdRng,
3818    ) {
3819        cx.update(LanguageModelRegistry::test);
3820        cx.set_global(cx.update(SettingsStore::test));
3821        cx.update(language_settings::init);
3822
3823        let text = concat!(
3824            "fn main() {\n",
3825            "  \n",
3826            "}\n" //
3827        );
3828        let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
3829        let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
3830        let range = buffer.read_with(cx, |buffer, cx| {
3831            let snapshot = buffer.snapshot(cx);
3832            snapshot.anchor_before(Point::new(1, 2))..snapshot.anchor_after(Point::new(1, 2))
3833        });
3834        let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
3835        let codegen = cx.new(|cx| {
3836            CodegenAlternative::new(
3837                buffer.clone(),
3838                range.clone(),
3839                true,
3840                None,
3841                prompt_builder,
3842                cx,
3843            )
3844        });
3845
3846        let chunks_tx = simulate_response_stream(codegen.clone(), cx);
3847
3848        cx.background_executor.run_until_parked();
3849
3850        let mut new_text = concat!(
3851            "let mut x = 0;\n",
3852            "while x < 10 {\n",
3853            "    x += 1;\n",
3854            "}", //
3855        );
3856        while !new_text.is_empty() {
3857            let max_len = cmp::min(new_text.len(), 10);
3858            let len = rng.gen_range(1..=max_len);
3859            let (chunk, suffix) = new_text.split_at(len);
3860            chunks_tx.unbounded_send(chunk.to_string()).unwrap();
3861            new_text = suffix;
3862            cx.background_executor.run_until_parked();
3863        }
3864        drop(chunks_tx);
3865        cx.background_executor.run_until_parked();
3866
3867        assert_eq!(
3868            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
3869            indoc! {"
3870                fn main() {
3871                    let mut x = 0;
3872                    while x < 10 {
3873                        x += 1;
3874                    }
3875                }
3876            "}
3877        );
3878    }
3879
3880    #[gpui::test(iterations = 10)]
3881    async fn test_autoindent_respects_tabs_in_selection(cx: &mut TestAppContext) {
3882        cx.update(LanguageModelRegistry::test);
3883        cx.set_global(cx.update(SettingsStore::test));
3884        cx.update(language_settings::init);
3885
3886        let text = indoc! {"
3887            func main() {
3888            \tx := 0
3889            \tfor i := 0; i < 10; i++ {
3890            \t\tx++
3891            \t}
3892            }
3893        "};
3894        let buffer = cx.new(|cx| Buffer::local(text, cx));
3895        let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
3896        let range = buffer.read_with(cx, |buffer, cx| {
3897            let snapshot = buffer.snapshot(cx);
3898            snapshot.anchor_before(Point::new(0, 0))..snapshot.anchor_after(Point::new(4, 2))
3899        });
3900        let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
3901        let codegen = cx.new(|cx| {
3902            CodegenAlternative::new(
3903                buffer.clone(),
3904                range.clone(),
3905                true,
3906                None,
3907                prompt_builder,
3908                cx,
3909            )
3910        });
3911
3912        let chunks_tx = simulate_response_stream(codegen.clone(), cx);
3913        let new_text = concat!(
3914            "func main() {\n",
3915            "\tx := 0\n",
3916            "\tfor x < 10 {\n",
3917            "\t\tx++\n",
3918            "\t}", //
3919        );
3920        chunks_tx.unbounded_send(new_text.to_string()).unwrap();
3921        drop(chunks_tx);
3922        cx.background_executor.run_until_parked();
3923
3924        assert_eq!(
3925            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
3926            indoc! {"
3927                func main() {
3928                \tx := 0
3929                \tfor x < 10 {
3930                \t\tx++
3931                \t}
3932                }
3933            "}
3934        );
3935    }
3936
3937    #[gpui::test]
3938    async fn test_inactive_codegen_alternative(cx: &mut TestAppContext) {
3939        cx.update(LanguageModelRegistry::test);
3940        cx.set_global(cx.update(SettingsStore::test));
3941        cx.update(language_settings::init);
3942
3943        let text = indoc! {"
3944            fn main() {
3945                let x = 0;
3946            }
3947        "};
3948        let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
3949        let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
3950        let range = buffer.read_with(cx, |buffer, cx| {
3951            let snapshot = buffer.snapshot(cx);
3952            snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(1, 14))
3953        });
3954        let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
3955        let codegen = cx.new(|cx| {
3956            CodegenAlternative::new(
3957                buffer.clone(),
3958                range.clone(),
3959                false,
3960                None,
3961                prompt_builder,
3962                cx,
3963            )
3964        });
3965
3966        let chunks_tx = simulate_response_stream(codegen.clone(), cx);
3967        chunks_tx
3968            .unbounded_send("let mut x = 0;\nx += 1;".to_string())
3969            .unwrap();
3970        drop(chunks_tx);
3971        cx.run_until_parked();
3972
3973        // The codegen is inactive, so the buffer doesn't get modified.
3974        assert_eq!(
3975            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
3976            text
3977        );
3978
3979        // Activating the codegen applies the changes.
3980        codegen.update(cx, |codegen, cx| codegen.set_active(true, cx));
3981        assert_eq!(
3982            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
3983            indoc! {"
3984                fn main() {
3985                    let mut x = 0;
3986                    x += 1;
3987                }
3988            "}
3989        );
3990
3991        // Deactivating the codegen undoes the changes.
3992        codegen.update(cx, |codegen, cx| codegen.set_active(false, cx));
3993        cx.run_until_parked();
3994        assert_eq!(
3995            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
3996            text
3997        );
3998    }
3999
4000    #[gpui::test]
4001    async fn test_strip_invalid_spans_from_codeblock() {
4002        assert_chunks("Lorem ipsum dolor", "Lorem ipsum dolor").await;
4003        assert_chunks("```\nLorem ipsum dolor", "Lorem ipsum dolor").await;
4004        assert_chunks("```\nLorem ipsum dolor\n```", "Lorem ipsum dolor").await;
4005        assert_chunks(
4006            "```html\n```js\nLorem ipsum dolor\n```\n```",
4007            "```js\nLorem ipsum dolor\n```",
4008        )
4009        .await;
4010        assert_chunks("``\nLorem ipsum dolor\n```", "``\nLorem ipsum dolor\n```").await;
4011        assert_chunks("Lorem<|CURSOR|> ipsum", "Lorem ipsum").await;
4012        assert_chunks("Lorem ipsum", "Lorem ipsum").await;
4013        assert_chunks("```\n<|CURSOR|>Lorem ipsum\n```", "Lorem ipsum").await;
4014
4015        async fn assert_chunks(text: &str, expected_text: &str) {
4016            for chunk_size in 1..=text.len() {
4017                let actual_text = StripInvalidSpans::new(chunks(text, chunk_size))
4018                    .map(|chunk| chunk.unwrap())
4019                    .collect::<String>()
4020                    .await;
4021                assert_eq!(
4022                    actual_text, expected_text,
4023                    "failed to strip invalid spans, chunk size: {}",
4024                    chunk_size
4025                );
4026            }
4027        }
4028
4029        fn chunks(text: &str, size: usize) -> impl Stream<Item = Result<String>> {
4030            stream::iter(
4031                text.chars()
4032                    .collect::<Vec<_>>()
4033                    .chunks(size)
4034                    .map(|chunk| Ok(chunk.iter().collect::<String>()))
4035                    .collect::<Vec<_>>(),
4036            )
4037        }
4038    }
4039
4040    fn simulate_response_stream(
4041        codegen: Entity<CodegenAlternative>,
4042        cx: &mut TestAppContext,
4043    ) -> mpsc::UnboundedSender<String> {
4044        let (chunks_tx, chunks_rx) = mpsc::unbounded();
4045        codegen.update(cx, |codegen, cx| {
4046            codegen.handle_stream(
4047                String::new(),
4048                String::new(),
4049                None,
4050                future::ready(Ok(LanguageModelTextStream {
4051                    message_id: None,
4052                    stream: chunks_rx.map(Ok).boxed(),
4053                })),
4054                cx,
4055            );
4056        });
4057        chunks_tx
4058    }
4059
4060    fn rust_lang() -> Language {
4061        Language::new(
4062            LanguageConfig {
4063                name: "Rust".into(),
4064                matcher: LanguageMatcher {
4065                    path_suffixes: vec!["rs".to_string()],
4066                    ..Default::default()
4067                },
4068                ..Default::default()
4069            },
4070            Some(tree_sitter_rust::LANGUAGE.into()),
4071        )
4072        .with_indents_query(
4073            r#"
4074            (call_expression) @indent
4075            (field_expression) @indent
4076            (_ "(" ")" @end) @indent
4077            (_ "{" "}" @end) @indent
4078            "#,
4079        )
4080        .unwrap()
4081    }
4082}