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