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