inline_prompt_editor.rs

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