inline_prompt_editor.rs

   1use crate::assistant_model_selector::AssistantModelSelector;
   2use crate::buffer_codegen::BufferCodegen;
   3use crate::context_picker::ContextPicker;
   4use crate::context_store::ContextStore;
   5use crate::context_strip::{ContextStrip, ContextStripEvent, SuggestContextKind};
   6use crate::terminal_codegen::TerminalCodegen;
   7use crate::thread_store::ThreadStore;
   8use crate::{CycleNextInlineAssist, CyclePreviousInlineAssist};
   9use crate::{RemoveAllContext, ToggleContextPicker, ToggleModelSelector};
  10use client::ErrorExt;
  11use collections::VecDeque;
  12use editor::{
  13    actions::{MoveDown, MoveUp},
  14    Editor, EditorElement, EditorEvent, EditorMode, EditorStyle, GutterDimensions, MultiBuffer,
  15};
  16use feature_flags::{FeatureFlagAppExt as _, ZedPro};
  17use fs::Fs;
  18use gpui::{
  19    anchored, deferred, point, AnyElement, AppContext, ClickEvent, CursorStyle, EventEmitter,
  20    FocusHandle, FocusableView, FontWeight, Model, Subscription, TextStyle, View, ViewContext,
  21    WeakModel, WeakView, WindowContext,
  22};
  23use language_model::{LanguageModel, LanguageModelRegistry};
  24use language_model_selector::LanguageModelSelector;
  25use parking_lot::Mutex;
  26use settings::Settings;
  27use std::cmp;
  28use std::sync::Arc;
  29use theme::ThemeSettings;
  30use ui::utils::WithRemSize;
  31use ui::{
  32    prelude::*, CheckboxWithLabel, IconButtonShape, KeyBinding, Popover, PopoverMenuHandle, Tooltip,
  33};
  34use util::ResultExt;
  35use workspace::Workspace;
  36
  37pub struct PromptEditor<T> {
  38    pub editor: View<Editor>,
  39    mode: PromptEditorMode,
  40    context_store: Model<ContextStore>,
  41    context_strip: View<ContextStrip>,
  42    context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
  43    model_selector: View<AssistantModelSelector>,
  44    model_selector_menu_handle: PopoverMenuHandle<LanguageModelSelector>,
  45    edited_since_done: bool,
  46    prompt_history: VecDeque<String>,
  47    prompt_history_ix: Option<usize>,
  48    pending_prompt: String,
  49    _codegen_subscription: Subscription,
  50    editor_subscriptions: Vec<Subscription>,
  51    _context_strip_subscription: Subscription,
  52    show_rate_limit_notice: bool,
  53    _phantom: std::marker::PhantomData<T>,
  54}
  55
  56impl<T: 'static> EventEmitter<PromptEditorEvent> for PromptEditor<T> {}
  57
  58impl<T: 'static> Render for PromptEditor<T> {
  59    fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
  60        let ui_font_size = ThemeSettings::get_global(cx).ui_font_size;
  61        let mut buttons = Vec::new();
  62
  63        let left_gutter_width = match &self.mode {
  64            PromptEditorMode::Buffer {
  65                id: _,
  66                codegen,
  67                gutter_dimensions,
  68            } => {
  69                let codegen = codegen.read(cx);
  70
  71                if codegen.alternative_count(cx) > 1 {
  72                    buttons.push(self.render_cycle_controls(&codegen, cx));
  73                }
  74
  75                let gutter_dimensions = gutter_dimensions.lock();
  76
  77                gutter_dimensions.full_width() + (gutter_dimensions.margin / 2.0)
  78            }
  79            PromptEditorMode::Terminal { .. } => {
  80                // Give the equivalent of the same left-padding that we're using on the right
  81                Pixels::from(40.0)
  82            }
  83        };
  84
  85        let bottom_padding = match &self.mode {
  86            PromptEditorMode::Buffer { .. } => Pixels::from(0.),
  87            PromptEditorMode::Terminal { .. } => Pixels::from(8.0),
  88        };
  89
  90        buttons.extend(self.render_buttons(cx));
  91
  92        v_flex()
  93            .key_context("PromptEditor")
  94            .bg(cx.theme().colors().editor_background)
  95            .block_mouse_down()
  96            .gap_0p5()
  97            .border_y_1()
  98            .border_color(cx.theme().status().info_border)
  99            .size_full()
 100            .pt_0p5()
 101            .pb(bottom_padding)
 102            .pr_6()
 103            .child(
 104                h_flex()
 105                    .items_start()
 106                    .cursor(CursorStyle::Arrow)
 107                    .on_action(cx.listener(Self::toggle_context_picker))
 108                    .on_action(cx.listener(Self::toggle_model_selector))
 109                    .on_action(cx.listener(Self::confirm))
 110                    .on_action(cx.listener(Self::cancel))
 111                    .on_action(cx.listener(Self::move_up))
 112                    .on_action(cx.listener(Self::move_down))
 113                    .on_action(cx.listener(Self::remove_all_context))
 114                    .capture_action(cx.listener(Self::cycle_prev))
 115                    .capture_action(cx.listener(Self::cycle_next))
 116                    .child(
 117                        WithRemSize::new(ui_font_size)
 118                            .flex()
 119                            .flex_row()
 120                            .flex_shrink_0()
 121                            .items_center()
 122                            .h_full()
 123                            .w(left_gutter_width)
 124                            .justify_center()
 125                            .gap_2()
 126                            .child(self.render_close_button(cx))
 127                            .map(|el| {
 128                                let CodegenStatus::Error(error) = self.codegen_status(cx) else {
 129                                    return el;
 130                                };
 131
 132                                let error_message = SharedString::from(error.to_string());
 133                                if error.error_code() == proto::ErrorCode::RateLimitExceeded
 134                                    && cx.has_flag::<ZedPro>()
 135                                {
 136                                    el.child(
 137                                        v_flex()
 138                                            .child(
 139                                                IconButton::new(
 140                                                    "rate-limit-error",
 141                                                    IconName::XCircle,
 142                                                )
 143                                                .toggle_state(self.show_rate_limit_notice)
 144                                                .shape(IconButtonShape::Square)
 145                                                .icon_size(IconSize::Small)
 146                                                .on_click(
 147                                                    cx.listener(Self::toggle_rate_limit_notice),
 148                                                ),
 149                                            )
 150                                            .children(self.show_rate_limit_notice.then(|| {
 151                                                deferred(
 152                                                    anchored()
 153                                                        .position_mode(
 154                                                            gpui::AnchoredPositionMode::Local,
 155                                                        )
 156                                                        .position(point(px(0.), px(24.)))
 157                                                        .anchor(gpui::Corner::TopLeft)
 158                                                        .child(self.render_rate_limit_notice(cx)),
 159                                                )
 160                                            })),
 161                                    )
 162                                } else {
 163                                    el.child(
 164                                        div()
 165                                            .id("error")
 166                                            .tooltip(move |cx| {
 167                                                Tooltip::text(error_message.clone(), cx)
 168                                            })
 169                                            .child(
 170                                                Icon::new(IconName::XCircle)
 171                                                    .size(IconSize::Small)
 172                                                    .color(Color::Error),
 173                                            ),
 174                                    )
 175                                }
 176                            }),
 177                    )
 178                    .child(
 179                        h_flex()
 180                            .w_full()
 181                            .justify_between()
 182                            .child(div().flex_1().child(self.render_editor(cx)))
 183                            .child(
 184                                WithRemSize::new(ui_font_size)
 185                                    .flex()
 186                                    .flex_row()
 187                                    .items_center()
 188                                    .gap_1()
 189                                    .children(buttons),
 190                            ),
 191                    ),
 192            )
 193            .child(
 194                WithRemSize::new(ui_font_size)
 195                    .flex()
 196                    .flex_row()
 197                    .items_center()
 198                    .child(h_flex().flex_shrink_0().w(left_gutter_width))
 199                    .child(
 200                        h_flex()
 201                            .w_full()
 202                            .pl_1()
 203                            .items_start()
 204                            .justify_between()
 205                            .child(self.context_strip.clone())
 206                            .child(self.model_selector.clone()),
 207                    ),
 208            )
 209    }
 210}
 211
 212impl<T: 'static> FocusableView for PromptEditor<T> {
 213    fn focus_handle(&self, cx: &AppContext) -> FocusHandle {
 214        self.editor.focus_handle(cx)
 215    }
 216}
 217
 218impl<T: 'static> PromptEditor<T> {
 219    const MAX_LINES: u8 = 8;
 220
 221    fn codegen_status<'a>(&'a self, cx: &'a AppContext) -> &'a CodegenStatus {
 222        match &self.mode {
 223            PromptEditorMode::Buffer { codegen, .. } => codegen.read(cx).status(cx),
 224            PromptEditorMode::Terminal { codegen, .. } => &codegen.read(cx).status,
 225        }
 226    }
 227
 228    fn subscribe_to_editor(&mut self, cx: &mut ViewContext<Self>) {
 229        self.editor_subscriptions.clear();
 230        self.editor_subscriptions
 231            .push(cx.subscribe(&self.editor, Self::handle_prompt_editor_events));
 232    }
 233
 234    pub fn set_show_cursor_when_unfocused(
 235        &mut self,
 236        show_cursor_when_unfocused: bool,
 237        cx: &mut ViewContext<Self>,
 238    ) {
 239        self.editor.update(cx, |editor, cx| {
 240            editor.set_show_cursor_when_unfocused(show_cursor_when_unfocused, cx)
 241        });
 242    }
 243
 244    pub fn unlink(&mut self, cx: &mut ViewContext<Self>) {
 245        let prompt = self.prompt(cx);
 246        let focus = self.editor.focus_handle(cx).contains_focused(cx);
 247        self.editor = cx.new_view(|cx| {
 248            let mut editor = Editor::auto_height(Self::MAX_LINES as usize, cx);
 249            editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
 250            editor.set_placeholder_text(Self::placeholder_text(&self.mode, cx), cx);
 251            editor.set_placeholder_text("Add a prompt…", cx);
 252            editor.set_text(prompt, cx);
 253            if focus {
 254                editor.focus(cx);
 255            }
 256            editor
 257        });
 258        self.subscribe_to_editor(cx);
 259    }
 260
 261    pub fn placeholder_text(mode: &PromptEditorMode, cx: &WindowContext) -> String {
 262        let action = match mode {
 263            PromptEditorMode::Buffer { codegen, .. } => {
 264                if codegen.read(cx).is_insertion {
 265                    "Generate"
 266                } else {
 267                    "Transform"
 268                }
 269            }
 270            PromptEditorMode::Terminal { .. } => "Generate",
 271        };
 272
 273        let assistant_panel_keybinding = ui::text_for_action(&crate::ToggleFocus, cx)
 274            .map(|keybinding| format!("{keybinding} to chat ― "))
 275            .unwrap_or_default();
 276
 277        format!("{action}… ({assistant_panel_keybinding}↓↑ for history)")
 278    }
 279
 280    pub fn prompt(&self, cx: &AppContext) -> String {
 281        self.editor.read(cx).text(cx)
 282    }
 283
 284    fn toggle_rate_limit_notice(&mut self, _: &ClickEvent, cx: &mut ViewContext<Self>) {
 285        self.show_rate_limit_notice = !self.show_rate_limit_notice;
 286        if self.show_rate_limit_notice {
 287            cx.focus_view(&self.editor);
 288        }
 289        cx.notify();
 290    }
 291
 292    fn handle_prompt_editor_events(
 293        &mut self,
 294        _: View<Editor>,
 295        event: &EditorEvent,
 296        cx: &mut ViewContext<Self>,
 297    ) {
 298        match event {
 299            EditorEvent::Edited { .. } => {
 300                if let Some(workspace) = cx.window_handle().downcast::<Workspace>() {
 301                    workspace
 302                        .update(cx, |workspace, cx| {
 303                            let is_via_ssh = workspace
 304                                .project()
 305                                .update(cx, |project, _| project.is_via_ssh());
 306
 307                            workspace
 308                                .client()
 309                                .telemetry()
 310                                .log_edit_event("inline assist", is_via_ssh);
 311                        })
 312                        .log_err();
 313                }
 314                let prompt = self.editor.read(cx).text(cx);
 315                if self
 316                    .prompt_history_ix
 317                    .map_or(true, |ix| self.prompt_history[ix] != prompt)
 318                {
 319                    self.prompt_history_ix.take();
 320                    self.pending_prompt = prompt;
 321                }
 322
 323                self.edited_since_done = true;
 324                cx.notify();
 325            }
 326            EditorEvent::Blurred => {
 327                if self.show_rate_limit_notice {
 328                    self.show_rate_limit_notice = false;
 329                    cx.notify();
 330                }
 331            }
 332            _ => {}
 333        }
 334    }
 335
 336    fn toggle_context_picker(&mut self, _: &ToggleContextPicker, cx: &mut ViewContext<Self>) {
 337        self.context_picker_menu_handle.toggle(cx);
 338    }
 339
 340    fn toggle_model_selector(&mut self, _: &ToggleModelSelector, cx: &mut ViewContext<Self>) {
 341        self.model_selector_menu_handle.toggle(cx);
 342    }
 343
 344    pub fn remove_all_context(&mut self, _: &RemoveAllContext, cx: &mut ViewContext<Self>) {
 345        self.context_store.update(cx, |store, _cx| store.clear());
 346        cx.notify();
 347    }
 348
 349    fn cancel(&mut self, _: &editor::actions::Cancel, cx: &mut ViewContext<Self>) {
 350        match self.codegen_status(cx) {
 351            CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {
 352                cx.emit(PromptEditorEvent::CancelRequested);
 353            }
 354            CodegenStatus::Pending => {
 355                cx.emit(PromptEditorEvent::StopRequested);
 356            }
 357        }
 358    }
 359
 360    fn confirm(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
 361        match self.codegen_status(cx) {
 362            CodegenStatus::Idle => {
 363                cx.emit(PromptEditorEvent::StartRequested);
 364            }
 365            CodegenStatus::Pending => {
 366                cx.emit(PromptEditorEvent::DismissRequested);
 367            }
 368            CodegenStatus::Done => {
 369                if self.edited_since_done {
 370                    cx.emit(PromptEditorEvent::StartRequested);
 371                } else {
 372                    cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });
 373                }
 374            }
 375            CodegenStatus::Error(_) => {
 376                cx.emit(PromptEditorEvent::StartRequested);
 377            }
 378        }
 379    }
 380
 381    fn move_up(&mut self, _: &MoveUp, cx: &mut ViewContext<Self>) {
 382        if let Some(ix) = self.prompt_history_ix {
 383            if ix > 0 {
 384                self.prompt_history_ix = Some(ix - 1);
 385                let prompt = self.prompt_history[ix - 1].as_str();
 386                self.editor.update(cx, |editor, cx| {
 387                    editor.set_text(prompt, cx);
 388                    editor.move_to_beginning(&Default::default(), cx);
 389                });
 390            }
 391        } else if !self.prompt_history.is_empty() {
 392            self.prompt_history_ix = Some(self.prompt_history.len() - 1);
 393            let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();
 394            self.editor.update(cx, |editor, cx| {
 395                editor.set_text(prompt, cx);
 396                editor.move_to_beginning(&Default::default(), cx);
 397            });
 398        }
 399    }
 400
 401    fn move_down(&mut self, _: &MoveDown, cx: &mut ViewContext<Self>) {
 402        if let Some(ix) = self.prompt_history_ix {
 403            if ix < self.prompt_history.len() - 1 {
 404                self.prompt_history_ix = Some(ix + 1);
 405                let prompt = self.prompt_history[ix + 1].as_str();
 406                self.editor.update(cx, |editor, cx| {
 407                    editor.set_text(prompt, cx);
 408                    editor.move_to_end(&Default::default(), cx)
 409                });
 410            } else {
 411                self.prompt_history_ix = None;
 412                let prompt = self.pending_prompt.as_str();
 413                self.editor.update(cx, |editor, cx| {
 414                    editor.set_text(prompt, cx);
 415                    editor.move_to_end(&Default::default(), cx)
 416                });
 417            }
 418        } else {
 419            cx.focus_view(&self.context_strip);
 420        }
 421    }
 422
 423    fn render_buttons(&self, cx: &mut ViewContext<Self>) -> Vec<AnyElement> {
 424        let mode = match &self.mode {
 425            PromptEditorMode::Buffer { codegen, .. } => {
 426                let codegen = codegen.read(cx);
 427                if codegen.is_insertion {
 428                    GenerationMode::Generate
 429                } else {
 430                    GenerationMode::Transform
 431                }
 432            }
 433            PromptEditorMode::Terminal { .. } => GenerationMode::Generate,
 434        };
 435
 436        let codegen_status = self.codegen_status(cx);
 437
 438        match codegen_status {
 439            CodegenStatus::Idle => {
 440                vec![Button::new("start", mode.start_label())
 441                    .label_size(LabelSize::Small)
 442                    .icon(IconName::Return)
 443                    .icon_size(IconSize::XSmall)
 444                    .icon_color(Color::Muted)
 445                    .on_click(cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StartRequested)))
 446                    .into_any_element()]
 447            }
 448            CodegenStatus::Pending => vec![IconButton::new("stop", IconName::Stop)
 449                .icon_color(Color::Error)
 450                .shape(IconButtonShape::Square)
 451                .tooltip(move |cx| {
 452                    Tooltip::with_meta(
 453                        mode.tooltip_interrupt(),
 454                        Some(&menu::Cancel),
 455                        "Changes won't be discarded",
 456                        cx,
 457                    )
 458                })
 459                .on_click(cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StopRequested)))
 460                .into_any_element()],
 461            CodegenStatus::Done | CodegenStatus::Error(_) => {
 462                let has_error = matches!(codegen_status, CodegenStatus::Error(_));
 463                if has_error || self.edited_since_done {
 464                    vec![IconButton::new("restart", IconName::RotateCw)
 465                        .icon_color(Color::Info)
 466                        .shape(IconButtonShape::Square)
 467                        .tooltip(move |cx| {
 468                            Tooltip::with_meta(
 469                                mode.tooltip_restart(),
 470                                Some(&menu::Confirm),
 471                                "Changes will be discarded",
 472                                cx,
 473                            )
 474                        })
 475                        .on_click(cx.listener(|_, _, cx| {
 476                            cx.emit(PromptEditorEvent::StartRequested);
 477                        }))
 478                        .into_any_element()]
 479                } else {
 480                    let accept = IconButton::new("accept", IconName::Check)
 481                        .icon_color(Color::Info)
 482                        .shape(IconButtonShape::Square)
 483                        .tooltip(move |cx| {
 484                            Tooltip::for_action(mode.tooltip_accept(), &menu::Confirm, cx)
 485                        })
 486                        .on_click(cx.listener(|_, _, cx| {
 487                            cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });
 488                        }))
 489                        .into_any_element();
 490
 491                    match &self.mode {
 492                        PromptEditorMode::Terminal { .. } => vec![
 493                            accept,
 494                            IconButton::new("confirm", IconName::Play)
 495                                .icon_color(Color::Info)
 496                                .shape(IconButtonShape::Square)
 497                                .tooltip(|cx| {
 498                                    Tooltip::for_action(
 499                                        "Execute Generated Command",
 500                                        &menu::SecondaryConfirm,
 501                                        cx,
 502                                    )
 503                                })
 504                                .on_click(cx.listener(|_, _, cx| {
 505                                    cx.emit(PromptEditorEvent::ConfirmRequested { execute: true });
 506                                }))
 507                                .into_any_element(),
 508                        ],
 509                        PromptEditorMode::Buffer { .. } => vec![accept],
 510                    }
 511                }
 512            }
 513        }
 514    }
 515
 516    fn cycle_prev(&mut self, _: &CyclePreviousInlineAssist, cx: &mut ViewContext<Self>) {
 517        match &self.mode {
 518            PromptEditorMode::Buffer { codegen, .. } => {
 519                codegen.update(cx, |codegen, cx| codegen.cycle_prev(cx));
 520            }
 521            PromptEditorMode::Terminal { .. } => {
 522                // no cycle buttons in terminal mode
 523            }
 524        }
 525    }
 526
 527    fn cycle_next(&mut self, _: &CycleNextInlineAssist, cx: &mut ViewContext<Self>) {
 528        match &self.mode {
 529            PromptEditorMode::Buffer { codegen, .. } => {
 530                codegen.update(cx, |codegen, cx| codegen.cycle_next(cx));
 531            }
 532            PromptEditorMode::Terminal { .. } => {
 533                // no cycle buttons in terminal mode
 534            }
 535        }
 536    }
 537
 538    fn render_close_button(&self, cx: &ViewContext<Self>) -> AnyElement {
 539        IconButton::new("cancel", IconName::Close)
 540            .icon_color(Color::Muted)
 541            .shape(IconButtonShape::Square)
 542            .tooltip(|cx| Tooltip::text("Close Assistant", cx))
 543            .on_click(cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)))
 544            .into_any_element()
 545    }
 546
 547    fn render_cycle_controls(&self, codegen: &BufferCodegen, cx: &ViewContext<Self>) -> AnyElement {
 548        let disabled = matches!(codegen.status(cx), CodegenStatus::Idle);
 549
 550        let model_registry = LanguageModelRegistry::read_global(cx);
 551        let default_model = model_registry.active_model();
 552        let alternative_models = model_registry.inline_alternative_models();
 553
 554        let get_model_name = |index: usize| -> String {
 555            let name = |model: &Arc<dyn LanguageModel>| model.name().0.to_string();
 556
 557            match index {
 558                0 => default_model.as_ref().map_or_else(String::new, name),
 559                index if index <= alternative_models.len() => alternative_models
 560                    .get(index - 1)
 561                    .map_or_else(String::new, name),
 562                _ => String::new(),
 563            }
 564        };
 565
 566        let total_models = alternative_models.len() + 1;
 567
 568        if total_models <= 1 {
 569            return div().into_any_element();
 570        }
 571
 572        let current_index = codegen.active_alternative;
 573        let prev_index = (current_index + total_models - 1) % total_models;
 574        let next_index = (current_index + 1) % total_models;
 575
 576        let prev_model_name = get_model_name(prev_index);
 577        let next_model_name = get_model_name(next_index);
 578
 579        h_flex()
 580            .child(
 581                IconButton::new("previous", IconName::ChevronLeft)
 582                    .icon_color(Color::Muted)
 583                    .disabled(disabled || current_index == 0)
 584                    .shape(IconButtonShape::Square)
 585                    .tooltip({
 586                        let focus_handle = self.editor.focus_handle(cx);
 587                        move |cx| {
 588                            cx.new_view(|cx| {
 589                                let mut tooltip = Tooltip::new("Previous Alternative").key_binding(
 590                                    KeyBinding::for_action_in(
 591                                        &CyclePreviousInlineAssist,
 592                                        &focus_handle,
 593                                        cx,
 594                                    ),
 595                                );
 596                                if !disabled && current_index != 0 {
 597                                    tooltip = tooltip.meta(prev_model_name.clone());
 598                                }
 599                                tooltip
 600                            })
 601                            .into()
 602                        }
 603                    })
 604                    .on_click(cx.listener(|this, _, cx| {
 605                        this.cycle_prev(&CyclePreviousInlineAssist, cx);
 606                    })),
 607            )
 608            .child(
 609                Label::new(format!(
 610                    "{}/{}",
 611                    codegen.active_alternative + 1,
 612                    codegen.alternative_count(cx)
 613                ))
 614                .size(LabelSize::Small)
 615                .color(if disabled {
 616                    Color::Disabled
 617                } else {
 618                    Color::Muted
 619                }),
 620            )
 621            .child(
 622                IconButton::new("next", IconName::ChevronRight)
 623                    .icon_color(Color::Muted)
 624                    .disabled(disabled || current_index == total_models - 1)
 625                    .shape(IconButtonShape::Square)
 626                    .tooltip({
 627                        let focus_handle = self.editor.focus_handle(cx);
 628                        move |cx| {
 629                            cx.new_view(|cx| {
 630                                let mut tooltip = Tooltip::new("Next Alternative").key_binding(
 631                                    KeyBinding::for_action_in(
 632                                        &CycleNextInlineAssist,
 633                                        &focus_handle,
 634                                        cx,
 635                                    ),
 636                                );
 637                                if !disabled && current_index != total_models - 1 {
 638                                    tooltip = tooltip.meta(next_model_name.clone());
 639                                }
 640                                tooltip
 641                            })
 642                            .into()
 643                        }
 644                    })
 645                    .on_click(
 646                        cx.listener(|this, _, cx| this.cycle_next(&CycleNextInlineAssist, cx)),
 647                    ),
 648            )
 649            .into_any_element()
 650    }
 651
 652    fn render_rate_limit_notice(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
 653        Popover::new().child(
 654            v_flex()
 655                .occlude()
 656                .p_2()
 657                .child(
 658                    Label::new("Out of Tokens")
 659                        .size(LabelSize::Small)
 660                        .weight(FontWeight::BOLD),
 661                )
 662                .child(Label::new(
 663                    "Try Zed Pro for higher limits, a wider range of models, and more.",
 664                ))
 665                .child(
 666                    h_flex()
 667                        .justify_between()
 668                        .child(CheckboxWithLabel::new(
 669                            "dont-show-again",
 670                            Label::new("Don't show again"),
 671                            if dismissed_rate_limit_notice() {
 672                                ui::ToggleState::Selected
 673                            } else {
 674                                ui::ToggleState::Unselected
 675                            },
 676                            |selection, cx| {
 677                                let is_dismissed = match selection {
 678                                    ui::ToggleState::Unselected => false,
 679                                    ui::ToggleState::Indeterminate => return,
 680                                    ui::ToggleState::Selected => true,
 681                                };
 682
 683                                set_rate_limit_notice_dismissed(is_dismissed, cx)
 684                            },
 685                        ))
 686                        .child(
 687                            h_flex()
 688                                .gap_2()
 689                                .child(
 690                                    Button::new("dismiss", "Dismiss")
 691                                        .style(ButtonStyle::Transparent)
 692                                        .on_click(cx.listener(Self::toggle_rate_limit_notice)),
 693                                )
 694                                .child(Button::new("more-info", "More Info").on_click(
 695                                    |_event, cx| {
 696                                        cx.dispatch_action(Box::new(
 697                                            zed_actions::OpenAccountSettings,
 698                                        ))
 699                                    },
 700                                )),
 701                        ),
 702                ),
 703        )
 704    }
 705
 706    fn render_editor(&mut self, cx: &mut ViewContext<Self>) -> AnyElement {
 707        let font_size = TextSize::Default.rems(cx);
 708        let line_height = font_size.to_pixels(cx.rem_size()) * 1.3;
 709
 710        div()
 711            .key_context("InlineAssistEditor")
 712            .size_full()
 713            .p_2()
 714            .pl_1()
 715            .bg(cx.theme().colors().editor_background)
 716            .child({
 717                let settings = ThemeSettings::get_global(cx);
 718                let text_style = TextStyle {
 719                    color: cx.theme().colors().editor_foreground,
 720                    font_family: settings.buffer_font.family.clone(),
 721                    font_features: settings.buffer_font.features.clone(),
 722                    font_size: font_size.into(),
 723                    line_height: line_height.into(),
 724                    ..Default::default()
 725                };
 726
 727                EditorElement::new(
 728                    &self.editor,
 729                    EditorStyle {
 730                        background: cx.theme().colors().editor_background,
 731                        local_player: cx.theme().players().local(),
 732                        text: text_style,
 733                        ..Default::default()
 734                    },
 735                )
 736            })
 737            .into_any_element()
 738    }
 739
 740    fn handle_context_strip_event(
 741        &mut self,
 742        _context_strip: View<ContextStrip>,
 743        event: &ContextStripEvent,
 744        cx: &mut ViewContext<Self>,
 745    ) {
 746        match event {
 747            ContextStripEvent::PickerDismissed
 748            | ContextStripEvent::BlurredEmpty
 749            | ContextStripEvent::BlurredUp => {
 750                let editor_focus_handle = self.editor.focus_handle(cx);
 751                cx.focus(&editor_focus_handle);
 752            }
 753            ContextStripEvent::BlurredDown => {}
 754        }
 755    }
 756}
 757
 758pub enum PromptEditorMode {
 759    Buffer {
 760        id: InlineAssistId,
 761        codegen: Model<BufferCodegen>,
 762        gutter_dimensions: Arc<Mutex<GutterDimensions>>,
 763    },
 764    Terminal {
 765        id: TerminalInlineAssistId,
 766        codegen: Model<TerminalCodegen>,
 767        height_in_lines: u8,
 768    },
 769}
 770
 771pub enum PromptEditorEvent {
 772    StartRequested,
 773    StopRequested,
 774    ConfirmRequested { execute: bool },
 775    CancelRequested,
 776    DismissRequested,
 777    Resized { height_in_lines: u8 },
 778}
 779
 780#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
 781pub struct InlineAssistId(pub usize);
 782
 783impl InlineAssistId {
 784    pub fn post_inc(&mut self) -> InlineAssistId {
 785        let id = *self;
 786        self.0 += 1;
 787        id
 788    }
 789}
 790
 791impl PromptEditor<BufferCodegen> {
 792    #[allow(clippy::too_many_arguments)]
 793    pub fn new_buffer(
 794        id: InlineAssistId,
 795        gutter_dimensions: Arc<Mutex<GutterDimensions>>,
 796        prompt_history: VecDeque<String>,
 797        prompt_buffer: Model<MultiBuffer>,
 798        codegen: Model<BufferCodegen>,
 799        fs: Arc<dyn Fs>,
 800        context_store: Model<ContextStore>,
 801        workspace: WeakView<Workspace>,
 802        thread_store: Option<WeakModel<ThreadStore>>,
 803        cx: &mut ViewContext<PromptEditor<BufferCodegen>>,
 804    ) -> PromptEditor<BufferCodegen> {
 805        let codegen_subscription = cx.observe(&codegen, Self::handle_codegen_changed);
 806        let mode = PromptEditorMode::Buffer {
 807            id,
 808            codegen,
 809            gutter_dimensions,
 810        };
 811
 812        let prompt_editor = cx.new_view(|cx| {
 813            let mut editor = Editor::new(
 814                EditorMode::AutoHeight {
 815                    max_lines: Self::MAX_LINES as usize,
 816                },
 817                prompt_buffer,
 818                None,
 819                false,
 820                cx,
 821            );
 822            editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
 823            // Since the prompt editors for all inline assistants are linked,
 824            // always show the cursor (even when it isn't focused) because
 825            // typing in one will make what you typed appear in all of them.
 826            editor.set_show_cursor_when_unfocused(true, cx);
 827            editor.set_placeholder_text(Self::placeholder_text(&mode, cx), cx);
 828            editor
 829        });
 830        let context_picker_menu_handle = PopoverMenuHandle::default();
 831        let model_selector_menu_handle = PopoverMenuHandle::default();
 832
 833        let context_strip = cx.new_view(|cx| {
 834            ContextStrip::new(
 835                context_store.clone(),
 836                workspace.clone(),
 837                thread_store.clone(),
 838                context_picker_menu_handle.clone(),
 839                SuggestContextKind::Thread,
 840                cx,
 841            )
 842        });
 843
 844        let context_strip_subscription =
 845            cx.subscribe(&context_strip, Self::handle_context_strip_event);
 846
 847        let mut this: PromptEditor<BufferCodegen> = PromptEditor {
 848            editor: prompt_editor.clone(),
 849            context_store,
 850            context_strip,
 851            context_picker_menu_handle,
 852            model_selector: cx.new_view(|cx| {
 853                AssistantModelSelector::new(
 854                    fs,
 855                    model_selector_menu_handle.clone(),
 856                    prompt_editor.focus_handle(cx),
 857                    cx,
 858                )
 859            }),
 860            model_selector_menu_handle,
 861            edited_since_done: false,
 862            prompt_history,
 863            prompt_history_ix: None,
 864            pending_prompt: String::new(),
 865            _codegen_subscription: codegen_subscription,
 866            editor_subscriptions: Vec::new(),
 867            _context_strip_subscription: context_strip_subscription,
 868            show_rate_limit_notice: false,
 869            mode,
 870            _phantom: Default::default(),
 871        };
 872
 873        this.subscribe_to_editor(cx);
 874        this
 875    }
 876
 877    fn handle_codegen_changed(
 878        &mut self,
 879        _: Model<BufferCodegen>,
 880        cx: &mut ViewContext<PromptEditor<BufferCodegen>>,
 881    ) {
 882        match self.codegen_status(cx) {
 883            CodegenStatus::Idle => {
 884                self.editor
 885                    .update(cx, |editor, _| editor.set_read_only(false));
 886            }
 887            CodegenStatus::Pending => {
 888                self.editor
 889                    .update(cx, |editor, _| editor.set_read_only(true));
 890            }
 891            CodegenStatus::Done => {
 892                self.edited_since_done = false;
 893                self.editor
 894                    .update(cx, |editor, _| editor.set_read_only(false));
 895            }
 896            CodegenStatus::Error(error) => {
 897                if cx.has_flag::<ZedPro>()
 898                    && error.error_code() == proto::ErrorCode::RateLimitExceeded
 899                    && !dismissed_rate_limit_notice()
 900                {
 901                    self.show_rate_limit_notice = true;
 902                    cx.notify();
 903                }
 904
 905                self.edited_since_done = false;
 906                self.editor
 907                    .update(cx, |editor, _| editor.set_read_only(false));
 908            }
 909        }
 910    }
 911
 912    pub fn id(&self) -> InlineAssistId {
 913        match &self.mode {
 914            PromptEditorMode::Buffer { id, .. } => *id,
 915            PromptEditorMode::Terminal { .. } => unreachable!(),
 916        }
 917    }
 918
 919    pub fn codegen(&self) -> &Model<BufferCodegen> {
 920        match &self.mode {
 921            PromptEditorMode::Buffer { codegen, .. } => codegen,
 922            PromptEditorMode::Terminal { .. } => unreachable!(),
 923        }
 924    }
 925
 926    pub fn gutter_dimensions(&self) -> &Arc<Mutex<GutterDimensions>> {
 927        match &self.mode {
 928            PromptEditorMode::Buffer {
 929                gutter_dimensions, ..
 930            } => gutter_dimensions,
 931            PromptEditorMode::Terminal { .. } => unreachable!(),
 932        }
 933    }
 934}
 935
 936#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
 937pub struct TerminalInlineAssistId(pub usize);
 938
 939impl TerminalInlineAssistId {
 940    pub fn post_inc(&mut self) -> TerminalInlineAssistId {
 941        let id = *self;
 942        self.0 += 1;
 943        id
 944    }
 945}
 946
 947impl PromptEditor<TerminalCodegen> {
 948    #[allow(clippy::too_many_arguments)]
 949    pub fn new_terminal(
 950        id: TerminalInlineAssistId,
 951        prompt_history: VecDeque<String>,
 952        prompt_buffer: Model<MultiBuffer>,
 953        codegen: Model<TerminalCodegen>,
 954        fs: Arc<dyn Fs>,
 955        context_store: Model<ContextStore>,
 956        workspace: WeakView<Workspace>,
 957        thread_store: Option<WeakModel<ThreadStore>>,
 958        cx: &mut ViewContext<Self>,
 959    ) -> Self {
 960        let codegen_subscription = cx.observe(&codegen, Self::handle_codegen_changed);
 961        let mode = PromptEditorMode::Terminal {
 962            id,
 963            codegen,
 964            height_in_lines: 1,
 965        };
 966
 967        let prompt_editor = cx.new_view(|cx| {
 968            let mut editor = Editor::new(
 969                EditorMode::AutoHeight {
 970                    max_lines: Self::MAX_LINES as usize,
 971                },
 972                prompt_buffer,
 973                None,
 974                false,
 975                cx,
 976            );
 977            editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
 978            editor.set_placeholder_text(Self::placeholder_text(&mode, cx), cx);
 979            editor
 980        });
 981        let context_picker_menu_handle = PopoverMenuHandle::default();
 982        let model_selector_menu_handle = PopoverMenuHandle::default();
 983
 984        let context_strip = cx.new_view(|cx| {
 985            ContextStrip::new(
 986                context_store.clone(),
 987                workspace.clone(),
 988                thread_store.clone(),
 989                context_picker_menu_handle.clone(),
 990                SuggestContextKind::Thread,
 991                cx,
 992            )
 993        });
 994
 995        let context_strip_subscription =
 996            cx.subscribe(&context_strip, Self::handle_context_strip_event);
 997
 998        let mut this = Self {
 999            editor: prompt_editor.clone(),
1000            context_store,
1001            context_strip,
1002            context_picker_menu_handle,
1003            model_selector: cx.new_view(|cx| {
1004                AssistantModelSelector::new(
1005                    fs,
1006                    model_selector_menu_handle.clone(),
1007                    prompt_editor.focus_handle(cx),
1008                    cx,
1009                )
1010            }),
1011            model_selector_menu_handle,
1012            edited_since_done: false,
1013            prompt_history,
1014            prompt_history_ix: None,
1015            pending_prompt: String::new(),
1016            _codegen_subscription: codegen_subscription,
1017            editor_subscriptions: Vec::new(),
1018            _context_strip_subscription: context_strip_subscription,
1019            mode,
1020            show_rate_limit_notice: false,
1021            _phantom: Default::default(),
1022        };
1023        this.count_lines(cx);
1024        this.subscribe_to_editor(cx);
1025        this
1026    }
1027
1028    fn count_lines(&mut self, cx: &mut ViewContext<Self>) {
1029        let height_in_lines = cmp::max(
1030            2, // Make the editor at least two lines tall, to account for padding and buttons.
1031            cmp::min(
1032                self.editor
1033                    .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1),
1034                Self::MAX_LINES as u32,
1035            ),
1036        ) as u8;
1037
1038        match &mut self.mode {
1039            PromptEditorMode::Terminal {
1040                height_in_lines: current_height,
1041                ..
1042            } => {
1043                if height_in_lines != *current_height {
1044                    *current_height = height_in_lines;
1045                    cx.emit(PromptEditorEvent::Resized { height_in_lines });
1046                }
1047            }
1048            PromptEditorMode::Buffer { .. } => unreachable!(),
1049        }
1050    }
1051
1052    fn handle_codegen_changed(&mut self, _: Model<TerminalCodegen>, cx: &mut ViewContext<Self>) {
1053        match &self.codegen().read(cx).status {
1054            CodegenStatus::Idle => {
1055                self.editor
1056                    .update(cx, |editor, _| editor.set_read_only(false));
1057            }
1058            CodegenStatus::Pending => {
1059                self.editor
1060                    .update(cx, |editor, _| editor.set_read_only(true));
1061            }
1062            CodegenStatus::Done | CodegenStatus::Error(_) => {
1063                self.edited_since_done = false;
1064                self.editor
1065                    .update(cx, |editor, _| editor.set_read_only(false));
1066            }
1067        }
1068    }
1069
1070    pub fn codegen(&self) -> &Model<TerminalCodegen> {
1071        match &self.mode {
1072            PromptEditorMode::Buffer { .. } => unreachable!(),
1073            PromptEditorMode::Terminal { codegen, .. } => codegen,
1074        }
1075    }
1076
1077    pub fn id(&self) -> TerminalInlineAssistId {
1078        match &self.mode {
1079            PromptEditorMode::Buffer { .. } => unreachable!(),
1080            PromptEditorMode::Terminal { id, .. } => *id,
1081        }
1082    }
1083}
1084
1085const DISMISSED_RATE_LIMIT_NOTICE_KEY: &str = "dismissed-rate-limit-notice";
1086
1087fn dismissed_rate_limit_notice() -> bool {
1088    db::kvp::KEY_VALUE_STORE
1089        .read_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY)
1090        .log_err()
1091        .map_or(false, |s| s.is_some())
1092}
1093
1094fn set_rate_limit_notice_dismissed(is_dismissed: bool, cx: &mut AppContext) {
1095    db::write_and_log(cx, move || async move {
1096        if is_dismissed {
1097            db::kvp::KEY_VALUE_STORE
1098                .write_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY.into(), "1".into())
1099                .await
1100        } else {
1101            db::kvp::KEY_VALUE_STORE
1102                .delete_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY.into())
1103                .await
1104        }
1105    })
1106}
1107
1108pub enum CodegenStatus {
1109    Idle,
1110    Pending,
1111    Done,
1112    Error(anyhow::Error),
1113}
1114
1115/// This is just CodegenStatus without the anyhow::Error, which causes a lifetime issue for rendering the Cancel button.
1116#[derive(Copy, Clone)]
1117pub enum CancelButtonState {
1118    Idle,
1119    Pending,
1120    Done,
1121    Error,
1122}
1123
1124impl Into<CancelButtonState> for &CodegenStatus {
1125    fn into(self) -> CancelButtonState {
1126        match self {
1127            CodegenStatus::Idle => CancelButtonState::Idle,
1128            CodegenStatus::Pending => CancelButtonState::Pending,
1129            CodegenStatus::Done => CancelButtonState::Done,
1130            CodegenStatus::Error(_) => CancelButtonState::Error,
1131        }
1132    }
1133}
1134
1135#[derive(Copy, Clone)]
1136pub enum GenerationMode {
1137    Generate,
1138    Transform,
1139}
1140
1141impl GenerationMode {
1142    fn start_label(self) -> &'static str {
1143        match self {
1144            GenerationMode::Generate { .. } => "Generate",
1145            GenerationMode::Transform => "Transform",
1146        }
1147    }
1148    fn tooltip_interrupt(self) -> &'static str {
1149        match self {
1150            GenerationMode::Generate { .. } => "Interrupt Generation",
1151            GenerationMode::Transform => "Interrupt Transform",
1152        }
1153    }
1154
1155    fn tooltip_restart(self) -> &'static str {
1156        match self {
1157            GenerationMode::Generate { .. } => "Restart Generation",
1158            GenerationMode::Transform => "Restart Transform",
1159        }
1160    }
1161
1162    fn tooltip_accept(self) -> &'static str {
1163        match self {
1164            GenerationMode::Generate { .. } => "Accept Generation",
1165            GenerationMode::Transform => "Accept Transform",
1166        }
1167    }
1168}