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;
   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                    cx,
 797                )
 798            }),
 799            context_picker_menu_handle,
 800            model_selector: cx.new_view(|cx| {
 801                AssistantModelSelector::new(fs, model_selector_menu_handle.clone(), cx)
 802            }),
 803            model_selector_menu_handle,
 804            edited_since_done: false,
 805            prompt_history,
 806            prompt_history_ix: None,
 807            pending_prompt: String::new(),
 808            _codegen_subscription: codegen_subscription,
 809            editor_subscriptions: Vec::new(),
 810            show_rate_limit_notice: false,
 811            mode,
 812            _phantom: Default::default(),
 813        };
 814
 815        this.subscribe_to_editor(cx);
 816        this
 817    }
 818
 819    fn handle_codegen_changed(
 820        &mut self,
 821        _: Model<BufferCodegen>,
 822        cx: &mut ViewContext<PromptEditor<BufferCodegen>>,
 823    ) {
 824        match self.codegen_status(cx) {
 825            CodegenStatus::Idle => {
 826                self.editor
 827                    .update(cx, |editor, _| editor.set_read_only(false));
 828            }
 829            CodegenStatus::Pending => {
 830                self.editor
 831                    .update(cx, |editor, _| editor.set_read_only(true));
 832            }
 833            CodegenStatus::Done => {
 834                self.edited_since_done = false;
 835                self.editor
 836                    .update(cx, |editor, _| editor.set_read_only(false));
 837            }
 838            CodegenStatus::Error(error) => {
 839                if cx.has_flag::<ZedPro>()
 840                    && error.error_code() == proto::ErrorCode::RateLimitExceeded
 841                    && !dismissed_rate_limit_notice()
 842                {
 843                    self.show_rate_limit_notice = true;
 844                    cx.notify();
 845                }
 846
 847                self.edited_since_done = false;
 848                self.editor
 849                    .update(cx, |editor, _| editor.set_read_only(false));
 850            }
 851        }
 852    }
 853
 854    pub fn id(&self) -> InlineAssistId {
 855        match &self.mode {
 856            PromptEditorMode::Buffer { id, .. } => *id,
 857            PromptEditorMode::Terminal { .. } => unreachable!(),
 858        }
 859    }
 860
 861    pub fn codegen(&self) -> &Model<BufferCodegen> {
 862        match &self.mode {
 863            PromptEditorMode::Buffer { codegen, .. } => codegen,
 864            PromptEditorMode::Terminal { .. } => unreachable!(),
 865        }
 866    }
 867
 868    pub fn gutter_dimensions(&self) -> &Arc<Mutex<GutterDimensions>> {
 869        match &self.mode {
 870            PromptEditorMode::Buffer {
 871                gutter_dimensions, ..
 872            } => gutter_dimensions,
 873            PromptEditorMode::Terminal { .. } => unreachable!(),
 874        }
 875    }
 876}
 877
 878#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
 879pub struct TerminalInlineAssistId(pub usize);
 880
 881impl TerminalInlineAssistId {
 882    pub fn post_inc(&mut self) -> TerminalInlineAssistId {
 883        let id = *self;
 884        self.0 += 1;
 885        id
 886    }
 887}
 888
 889impl PromptEditor<TerminalCodegen> {
 890    #[allow(clippy::too_many_arguments)]
 891    pub fn new_terminal(
 892        id: TerminalInlineAssistId,
 893        prompt_history: VecDeque<String>,
 894        prompt_buffer: Model<MultiBuffer>,
 895        codegen: Model<TerminalCodegen>,
 896        fs: Arc<dyn Fs>,
 897        context_store: Model<ContextStore>,
 898        workspace: WeakView<Workspace>,
 899        thread_store: Option<WeakModel<ThreadStore>>,
 900        cx: &mut ViewContext<Self>,
 901    ) -> Self {
 902        let codegen_subscription = cx.observe(&codegen, Self::handle_codegen_changed);
 903        let mode = PromptEditorMode::Terminal {
 904            id,
 905            codegen,
 906            height_in_lines: 1,
 907        };
 908
 909        let prompt_editor = cx.new_view(|cx| {
 910            let mut editor = Editor::new(
 911                EditorMode::AutoHeight {
 912                    max_lines: Self::MAX_LINES as usize,
 913                },
 914                prompt_buffer,
 915                None,
 916                false,
 917                cx,
 918            );
 919            editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
 920            editor.set_placeholder_text(Self::placeholder_text(&mode, cx), cx);
 921            editor
 922        });
 923        let context_picker_menu_handle = PopoverMenuHandle::default();
 924        let model_selector_menu_handle = PopoverMenuHandle::default();
 925
 926        let mut this = Self {
 927            editor: prompt_editor.clone(),
 928            context_strip: cx.new_view(|cx| {
 929                ContextStrip::new(
 930                    context_store,
 931                    workspace.clone(),
 932                    thread_store.clone(),
 933                    prompt_editor.focus_handle(cx),
 934                    context_picker_menu_handle.clone(),
 935                    cx,
 936                )
 937            }),
 938            context_picker_menu_handle,
 939            model_selector: cx.new_view(|cx| {
 940                AssistantModelSelector::new(fs, model_selector_menu_handle.clone(), cx)
 941            }),
 942            model_selector_menu_handle,
 943            edited_since_done: false,
 944            prompt_history,
 945            prompt_history_ix: None,
 946            pending_prompt: String::new(),
 947            _codegen_subscription: codegen_subscription,
 948            editor_subscriptions: Vec::new(),
 949            mode,
 950            show_rate_limit_notice: false,
 951            _phantom: Default::default(),
 952        };
 953        this.count_lines(cx);
 954        this.subscribe_to_editor(cx);
 955        this
 956    }
 957
 958    fn count_lines(&mut self, cx: &mut ViewContext<Self>) {
 959        let height_in_lines = cmp::max(
 960            2, // Make the editor at least two lines tall, to account for padding and buttons.
 961            cmp::min(
 962                self.editor
 963                    .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1),
 964                Self::MAX_LINES as u32,
 965            ),
 966        ) as u8;
 967
 968        match &mut self.mode {
 969            PromptEditorMode::Terminal {
 970                height_in_lines: current_height,
 971                ..
 972            } => {
 973                if height_in_lines != *current_height {
 974                    *current_height = height_in_lines;
 975                    cx.emit(PromptEditorEvent::Resized { height_in_lines });
 976                }
 977            }
 978            PromptEditorMode::Buffer { .. } => unreachable!(),
 979        }
 980    }
 981
 982    fn handle_codegen_changed(&mut self, _: Model<TerminalCodegen>, cx: &mut ViewContext<Self>) {
 983        match &self.codegen().read(cx).status {
 984            CodegenStatus::Idle => {
 985                self.editor
 986                    .update(cx, |editor, _| editor.set_read_only(false));
 987            }
 988            CodegenStatus::Pending => {
 989                self.editor
 990                    .update(cx, |editor, _| editor.set_read_only(true));
 991            }
 992            CodegenStatus::Done | CodegenStatus::Error(_) => {
 993                self.edited_since_done = false;
 994                self.editor
 995                    .update(cx, |editor, _| editor.set_read_only(false));
 996            }
 997        }
 998    }
 999
1000    pub fn codegen(&self) -> &Model<TerminalCodegen> {
1001        match &self.mode {
1002            PromptEditorMode::Buffer { .. } => unreachable!(),
1003            PromptEditorMode::Terminal { codegen, .. } => codegen,
1004        }
1005    }
1006
1007    pub fn id(&self) -> TerminalInlineAssistId {
1008        match &self.mode {
1009            PromptEditorMode::Buffer { .. } => unreachable!(),
1010            PromptEditorMode::Terminal { id, .. } => *id,
1011        }
1012    }
1013}
1014
1015const DISMISSED_RATE_LIMIT_NOTICE_KEY: &str = "dismissed-rate-limit-notice";
1016
1017fn dismissed_rate_limit_notice() -> bool {
1018    db::kvp::KEY_VALUE_STORE
1019        .read_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY)
1020        .log_err()
1021        .map_or(false, |s| s.is_some())
1022}
1023
1024fn set_rate_limit_notice_dismissed(is_dismissed: bool, cx: &mut AppContext) {
1025    db::write_and_log(cx, move || async move {
1026        if is_dismissed {
1027            db::kvp::KEY_VALUE_STORE
1028                .write_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY.into(), "1".into())
1029                .await
1030        } else {
1031            db::kvp::KEY_VALUE_STORE
1032                .delete_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY.into())
1033                .await
1034        }
1035    })
1036}
1037
1038pub enum CodegenStatus {
1039    Idle,
1040    Pending,
1041    Done,
1042    Error(anyhow::Error),
1043}
1044
1045/// This is just CodegenStatus without the anyhow::Error, which causes a lifetime issue for rendering the Cancel button.
1046#[derive(Copy, Clone)]
1047pub enum CancelButtonState {
1048    Idle,
1049    Pending,
1050    Done,
1051    Error,
1052}
1053
1054impl Into<CancelButtonState> for &CodegenStatus {
1055    fn into(self) -> CancelButtonState {
1056        match self {
1057            CodegenStatus::Idle => CancelButtonState::Idle,
1058            CodegenStatus::Pending => CancelButtonState::Pending,
1059            CodegenStatus::Done => CancelButtonState::Done,
1060            CodegenStatus::Error(_) => CancelButtonState::Error,
1061        }
1062    }
1063}
1064
1065#[derive(Copy, Clone)]
1066pub enum GenerationMode {
1067    Generate,
1068    Transform,
1069}
1070
1071impl GenerationMode {
1072    fn start_label(self) -> &'static str {
1073        match self {
1074            GenerationMode::Generate { .. } => "Generate",
1075            GenerationMode::Transform => "Transform",
1076        }
1077    }
1078    fn tooltip_interrupt(self) -> &'static str {
1079        match self {
1080            GenerationMode::Generate { .. } => "Interrupt Generation",
1081            GenerationMode::Transform => "Interrupt Transform",
1082        }
1083    }
1084
1085    fn tooltip_restart(self) -> &'static str {
1086        match self {
1087            GenerationMode::Generate { .. } => "Restart Generation",
1088            GenerationMode::Transform => "Restart Transform",
1089        }
1090    }
1091
1092    fn tooltip_accept(self) -> &'static str {
1093        match self {
1094            GenerationMode::Generate { .. } => "Accept Generation",
1095            GenerationMode::Transform => "Accept Transform",
1096        }
1097    }
1098}