inline_prompt_editor.rs

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