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