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 =
 274            ui::text_for_action(&zed_actions::assistant::ToggleFocus, cx)
 275                .map(|keybinding| format!("{keybinding} to chat ― "))
 276                .unwrap_or_default();
 277
 278        format!("{action}… ({assistant_panel_keybinding}↓↑ for history)")
 279    }
 280
 281    pub fn prompt(&self, cx: &AppContext) -> String {
 282        self.editor.read(cx).text(cx)
 283    }
 284
 285    fn toggle_rate_limit_notice(&mut self, _: &ClickEvent, cx: &mut ViewContext<Self>) {
 286        self.show_rate_limit_notice = !self.show_rate_limit_notice;
 287        if self.show_rate_limit_notice {
 288            cx.focus_view(&self.editor);
 289        }
 290        cx.notify();
 291    }
 292
 293    fn handle_prompt_editor_events(
 294        &mut self,
 295        _: View<Editor>,
 296        event: &EditorEvent,
 297        cx: &mut ViewContext<Self>,
 298    ) {
 299        match event {
 300            EditorEvent::Edited { .. } => {
 301                if let Some(workspace) = cx.window_handle().downcast::<Workspace>() {
 302                    workspace
 303                        .update(cx, |workspace, cx| {
 304                            let is_via_ssh = workspace
 305                                .project()
 306                                .update(cx, |project, _| project.is_via_ssh());
 307
 308                            workspace
 309                                .client()
 310                                .telemetry()
 311                                .log_edit_event("inline assist", is_via_ssh);
 312                        })
 313                        .log_err();
 314                }
 315                let prompt = self.editor.read(cx).text(cx);
 316                if self
 317                    .prompt_history_ix
 318                    .map_or(true, |ix| self.prompt_history[ix] != prompt)
 319                {
 320                    self.prompt_history_ix.take();
 321                    self.pending_prompt = prompt;
 322                }
 323
 324                self.edited_since_done = true;
 325                cx.notify();
 326            }
 327            EditorEvent::Blurred => {
 328                if self.show_rate_limit_notice {
 329                    self.show_rate_limit_notice = false;
 330                    cx.notify();
 331                }
 332            }
 333            _ => {}
 334        }
 335    }
 336
 337    fn toggle_context_picker(&mut self, _: &ToggleContextPicker, cx: &mut ViewContext<Self>) {
 338        self.context_picker_menu_handle.toggle(cx);
 339    }
 340
 341    fn toggle_model_selector(&mut self, _: &ToggleModelSelector, cx: &mut ViewContext<Self>) {
 342        self.model_selector_menu_handle.toggle(cx);
 343    }
 344
 345    pub fn remove_all_context(&mut self, _: &RemoveAllContext, cx: &mut ViewContext<Self>) {
 346        self.context_store.update(cx, |store, _cx| store.clear());
 347        cx.notify();
 348    }
 349
 350    fn cancel(&mut self, _: &editor::actions::Cancel, cx: &mut ViewContext<Self>) {
 351        match self.codegen_status(cx) {
 352            CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {
 353                cx.emit(PromptEditorEvent::CancelRequested);
 354            }
 355            CodegenStatus::Pending => {
 356                cx.emit(PromptEditorEvent::StopRequested);
 357            }
 358        }
 359    }
 360
 361    fn confirm(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
 362        match self.codegen_status(cx) {
 363            CodegenStatus::Idle => {
 364                cx.emit(PromptEditorEvent::StartRequested);
 365            }
 366            CodegenStatus::Pending => {
 367                cx.emit(PromptEditorEvent::DismissRequested);
 368            }
 369            CodegenStatus::Done => {
 370                if self.edited_since_done {
 371                    cx.emit(PromptEditorEvent::StartRequested);
 372                } else {
 373                    cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });
 374                }
 375            }
 376            CodegenStatus::Error(_) => {
 377                cx.emit(PromptEditorEvent::StartRequested);
 378            }
 379        }
 380    }
 381
 382    fn move_up(&mut self, _: &MoveUp, cx: &mut ViewContext<Self>) {
 383        if let Some(ix) = self.prompt_history_ix {
 384            if ix > 0 {
 385                self.prompt_history_ix = Some(ix - 1);
 386                let prompt = self.prompt_history[ix - 1].as_str();
 387                self.editor.update(cx, |editor, cx| {
 388                    editor.set_text(prompt, cx);
 389                    editor.move_to_beginning(&Default::default(), cx);
 390                });
 391            }
 392        } else if !self.prompt_history.is_empty() {
 393            self.prompt_history_ix = Some(self.prompt_history.len() - 1);
 394            let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();
 395            self.editor.update(cx, |editor, cx| {
 396                editor.set_text(prompt, cx);
 397                editor.move_to_beginning(&Default::default(), cx);
 398            });
 399        }
 400    }
 401
 402    fn move_down(&mut self, _: &MoveDown, cx: &mut ViewContext<Self>) {
 403        if let Some(ix) = self.prompt_history_ix {
 404            if ix < self.prompt_history.len() - 1 {
 405                self.prompt_history_ix = Some(ix + 1);
 406                let prompt = self.prompt_history[ix + 1].as_str();
 407                self.editor.update(cx, |editor, cx| {
 408                    editor.set_text(prompt, cx);
 409                    editor.move_to_end(&Default::default(), cx)
 410                });
 411            } else {
 412                self.prompt_history_ix = None;
 413                let prompt = self.pending_prompt.as_str();
 414                self.editor.update(cx, |editor, cx| {
 415                    editor.set_text(prompt, cx);
 416                    editor.move_to_end(&Default::default(), cx)
 417                });
 418            }
 419        } else {
 420            cx.focus_view(&self.context_strip);
 421        }
 422    }
 423
 424    fn render_buttons(&self, cx: &mut ViewContext<Self>) -> Vec<AnyElement> {
 425        let mode = match &self.mode {
 426            PromptEditorMode::Buffer { codegen, .. } => {
 427                let codegen = codegen.read(cx);
 428                if codegen.is_insertion {
 429                    GenerationMode::Generate
 430                } else {
 431                    GenerationMode::Transform
 432                }
 433            }
 434            PromptEditorMode::Terminal { .. } => GenerationMode::Generate,
 435        };
 436
 437        let codegen_status = self.codegen_status(cx);
 438
 439        match codegen_status {
 440            CodegenStatus::Idle => {
 441                vec![Button::new("start", mode.start_label())
 442                    .label_size(LabelSize::Small)
 443                    .icon(IconName::Return)
 444                    .icon_size(IconSize::XSmall)
 445                    .icon_color(Color::Muted)
 446                    .on_click(cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StartRequested)))
 447                    .into_any_element()]
 448            }
 449            CodegenStatus::Pending => vec![IconButton::new("stop", IconName::Stop)
 450                .icon_color(Color::Error)
 451                .shape(IconButtonShape::Square)
 452                .tooltip(move |cx| {
 453                    Tooltip::with_meta(
 454                        mode.tooltip_interrupt(),
 455                        Some(&menu::Cancel),
 456                        "Changes won't be discarded",
 457                        cx,
 458                    )
 459                })
 460                .on_click(cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StopRequested)))
 461                .into_any_element()],
 462            CodegenStatus::Done | CodegenStatus::Error(_) => {
 463                let has_error = matches!(codegen_status, CodegenStatus::Error(_));
 464                if has_error || self.edited_since_done {
 465                    vec![IconButton::new("restart", IconName::RotateCw)
 466                        .icon_color(Color::Info)
 467                        .shape(IconButtonShape::Square)
 468                        .tooltip(move |cx| {
 469                            Tooltip::with_meta(
 470                                mode.tooltip_restart(),
 471                                Some(&menu::Confirm),
 472                                "Changes will be discarded",
 473                                cx,
 474                            )
 475                        })
 476                        .on_click(cx.listener(|_, _, cx| {
 477                            cx.emit(PromptEditorEvent::StartRequested);
 478                        }))
 479                        .into_any_element()]
 480                } else {
 481                    let accept = IconButton::new("accept", IconName::Check)
 482                        .icon_color(Color::Info)
 483                        .shape(IconButtonShape::Square)
 484                        .tooltip(move |cx| {
 485                            Tooltip::for_action(mode.tooltip_accept(), &menu::Confirm, cx)
 486                        })
 487                        .on_click(cx.listener(|_, _, cx| {
 488                            cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });
 489                        }))
 490                        .into_any_element();
 491
 492                    match &self.mode {
 493                        PromptEditorMode::Terminal { .. } => vec![
 494                            accept,
 495                            IconButton::new("confirm", IconName::Play)
 496                                .icon_color(Color::Info)
 497                                .shape(IconButtonShape::Square)
 498                                .tooltip(|cx| {
 499                                    Tooltip::for_action(
 500                                        "Execute Generated Command",
 501                                        &menu::SecondaryConfirm,
 502                                        cx,
 503                                    )
 504                                })
 505                                .on_click(cx.listener(|_, _, cx| {
 506                                    cx.emit(PromptEditorEvent::ConfirmRequested { execute: true });
 507                                }))
 508                                .into_any_element(),
 509                        ],
 510                        PromptEditorMode::Buffer { .. } => vec![accept],
 511                    }
 512                }
 513            }
 514        }
 515    }
 516
 517    fn cycle_prev(&mut self, _: &CyclePreviousInlineAssist, cx: &mut ViewContext<Self>) {
 518        match &self.mode {
 519            PromptEditorMode::Buffer { codegen, .. } => {
 520                codegen.update(cx, |codegen, cx| codegen.cycle_prev(cx));
 521            }
 522            PromptEditorMode::Terminal { .. } => {
 523                // no cycle buttons in terminal mode
 524            }
 525        }
 526    }
 527
 528    fn cycle_next(&mut self, _: &CycleNextInlineAssist, cx: &mut ViewContext<Self>) {
 529        match &self.mode {
 530            PromptEditorMode::Buffer { codegen, .. } => {
 531                codegen.update(cx, |codegen, cx| codegen.cycle_next(cx));
 532            }
 533            PromptEditorMode::Terminal { .. } => {
 534                // no cycle buttons in terminal mode
 535            }
 536        }
 537    }
 538
 539    fn render_close_button(&self, cx: &ViewContext<Self>) -> AnyElement {
 540        IconButton::new("cancel", IconName::Close)
 541            .icon_color(Color::Muted)
 542            .shape(IconButtonShape::Square)
 543            .tooltip(|cx| Tooltip::text("Close Assistant", cx))
 544            .on_click(cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)))
 545            .into_any_element()
 546    }
 547
 548    fn render_cycle_controls(&self, codegen: &BufferCodegen, cx: &ViewContext<Self>) -> AnyElement {
 549        let disabled = matches!(codegen.status(cx), CodegenStatus::Idle);
 550
 551        let model_registry = LanguageModelRegistry::read_global(cx);
 552        let default_model = model_registry.active_model();
 553        let alternative_models = model_registry.inline_alternative_models();
 554
 555        let get_model_name = |index: usize| -> String {
 556            let name = |model: &Arc<dyn LanguageModel>| model.name().0.to_string();
 557
 558            match index {
 559                0 => default_model.as_ref().map_or_else(String::new, name),
 560                index if index <= alternative_models.len() => alternative_models
 561                    .get(index - 1)
 562                    .map_or_else(String::new, name),
 563                _ => String::new(),
 564            }
 565        };
 566
 567        let total_models = alternative_models.len() + 1;
 568
 569        if total_models <= 1 {
 570            return div().into_any_element();
 571        }
 572
 573        let current_index = codegen.active_alternative;
 574        let prev_index = (current_index + total_models - 1) % total_models;
 575        let next_index = (current_index + 1) % total_models;
 576
 577        let prev_model_name = get_model_name(prev_index);
 578        let next_model_name = get_model_name(next_index);
 579
 580        h_flex()
 581            .child(
 582                IconButton::new("previous", IconName::ChevronLeft)
 583                    .icon_color(Color::Muted)
 584                    .disabled(disabled || current_index == 0)
 585                    .shape(IconButtonShape::Square)
 586                    .tooltip({
 587                        let focus_handle = self.editor.focus_handle(cx);
 588                        move |cx| {
 589                            cx.new_view(|cx| {
 590                                let mut tooltip = Tooltip::new("Previous Alternative").key_binding(
 591                                    KeyBinding::for_action_in(
 592                                        &CyclePreviousInlineAssist,
 593                                        &focus_handle,
 594                                        cx,
 595                                    ),
 596                                );
 597                                if !disabled && current_index != 0 {
 598                                    tooltip = tooltip.meta(prev_model_name.clone());
 599                                }
 600                                tooltip
 601                            })
 602                            .into()
 603                        }
 604                    })
 605                    .on_click(cx.listener(|this, _, cx| {
 606                        this.cycle_prev(&CyclePreviousInlineAssist, cx);
 607                    })),
 608            )
 609            .child(
 610                Label::new(format!(
 611                    "{}/{}",
 612                    codegen.active_alternative + 1,
 613                    codegen.alternative_count(cx)
 614                ))
 615                .size(LabelSize::Small)
 616                .color(if disabled {
 617                    Color::Disabled
 618                } else {
 619                    Color::Muted
 620                }),
 621            )
 622            .child(
 623                IconButton::new("next", IconName::ChevronRight)
 624                    .icon_color(Color::Muted)
 625                    .disabled(disabled || current_index == total_models - 1)
 626                    .shape(IconButtonShape::Square)
 627                    .tooltip({
 628                        let focus_handle = self.editor.focus_handle(cx);
 629                        move |cx| {
 630                            cx.new_view(|cx| {
 631                                let mut tooltip = Tooltip::new("Next Alternative").key_binding(
 632                                    KeyBinding::for_action_in(
 633                                        &CycleNextInlineAssist,
 634                                        &focus_handle,
 635                                        cx,
 636                                    ),
 637                                );
 638                                if !disabled && current_index != total_models - 1 {
 639                                    tooltip = tooltip.meta(next_model_name.clone());
 640                                }
 641                                tooltip
 642                            })
 643                            .into()
 644                        }
 645                    })
 646                    .on_click(
 647                        cx.listener(|this, _, cx| this.cycle_next(&CycleNextInlineAssist, cx)),
 648                    ),
 649            )
 650            .into_any_element()
 651    }
 652
 653    fn render_rate_limit_notice(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
 654        Popover::new().child(
 655            v_flex()
 656                .occlude()
 657                .p_2()
 658                .child(
 659                    Label::new("Out of Tokens")
 660                        .size(LabelSize::Small)
 661                        .weight(FontWeight::BOLD),
 662                )
 663                .child(Label::new(
 664                    "Try Zed Pro for higher limits, a wider range of models, and more.",
 665                ))
 666                .child(
 667                    h_flex()
 668                        .justify_between()
 669                        .child(CheckboxWithLabel::new(
 670                            "dont-show-again",
 671                            Label::new("Don't show again"),
 672                            if dismissed_rate_limit_notice() {
 673                                ui::ToggleState::Selected
 674                            } else {
 675                                ui::ToggleState::Unselected
 676                            },
 677                            |selection, cx| {
 678                                let is_dismissed = match selection {
 679                                    ui::ToggleState::Unselected => false,
 680                                    ui::ToggleState::Indeterminate => return,
 681                                    ui::ToggleState::Selected => true,
 682                                };
 683
 684                                set_rate_limit_notice_dismissed(is_dismissed, cx)
 685                            },
 686                        ))
 687                        .child(
 688                            h_flex()
 689                                .gap_2()
 690                                .child(
 691                                    Button::new("dismiss", "Dismiss")
 692                                        .style(ButtonStyle::Transparent)
 693                                        .on_click(cx.listener(Self::toggle_rate_limit_notice)),
 694                                )
 695                                .child(Button::new("more-info", "More Info").on_click(
 696                                    |_event, cx| {
 697                                        cx.dispatch_action(Box::new(
 698                                            zed_actions::OpenAccountSettings,
 699                                        ))
 700                                    },
 701                                )),
 702                        ),
 703                ),
 704        )
 705    }
 706
 707    fn render_editor(&mut self, cx: &mut ViewContext<Self>) -> AnyElement {
 708        let font_size = TextSize::Default.rems(cx);
 709        let line_height = font_size.to_pixels(cx.rem_size()) * 1.3;
 710
 711        div()
 712            .key_context("InlineAssistEditor")
 713            .size_full()
 714            .p_2()
 715            .pl_1()
 716            .bg(cx.theme().colors().editor_background)
 717            .child({
 718                let settings = ThemeSettings::get_global(cx);
 719                let text_style = TextStyle {
 720                    color: cx.theme().colors().editor_foreground,
 721                    font_family: settings.buffer_font.family.clone(),
 722                    font_features: settings.buffer_font.features.clone(),
 723                    font_size: font_size.into(),
 724                    line_height: line_height.into(),
 725                    ..Default::default()
 726                };
 727
 728                EditorElement::new(
 729                    &self.editor,
 730                    EditorStyle {
 731                        background: cx.theme().colors().editor_background,
 732                        local_player: cx.theme().players().local(),
 733                        text: text_style,
 734                        ..Default::default()
 735                    },
 736                )
 737            })
 738            .into_any_element()
 739    }
 740
 741    fn handle_context_strip_event(
 742        &mut self,
 743        _context_strip: View<ContextStrip>,
 744        event: &ContextStripEvent,
 745        cx: &mut ViewContext<Self>,
 746    ) {
 747        match event {
 748            ContextStripEvent::PickerDismissed
 749            | ContextStripEvent::BlurredEmpty
 750            | ContextStripEvent::BlurredUp => {
 751                let editor_focus_handle = self.editor.focus_handle(cx);
 752                cx.focus(&editor_focus_handle);
 753            }
 754            ContextStripEvent::BlurredDown => {}
 755        }
 756    }
 757}
 758
 759pub enum PromptEditorMode {
 760    Buffer {
 761        id: InlineAssistId,
 762        codegen: Model<BufferCodegen>,
 763        gutter_dimensions: Arc<Mutex<GutterDimensions>>,
 764    },
 765    Terminal {
 766        id: TerminalInlineAssistId,
 767        codegen: Model<TerminalCodegen>,
 768        height_in_lines: u8,
 769    },
 770}
 771
 772pub enum PromptEditorEvent {
 773    StartRequested,
 774    StopRequested,
 775    ConfirmRequested { execute: bool },
 776    CancelRequested,
 777    DismissRequested,
 778    Resized { height_in_lines: u8 },
 779}
 780
 781#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
 782pub struct InlineAssistId(pub usize);
 783
 784impl InlineAssistId {
 785    pub fn post_inc(&mut self) -> InlineAssistId {
 786        let id = *self;
 787        self.0 += 1;
 788        id
 789    }
 790}
 791
 792impl PromptEditor<BufferCodegen> {
 793    #[allow(clippy::too_many_arguments)]
 794    pub fn new_buffer(
 795        id: InlineAssistId,
 796        gutter_dimensions: Arc<Mutex<GutterDimensions>>,
 797        prompt_history: VecDeque<String>,
 798        prompt_buffer: Model<MultiBuffer>,
 799        codegen: Model<BufferCodegen>,
 800        fs: Arc<dyn Fs>,
 801        context_store: Model<ContextStore>,
 802        workspace: WeakView<Workspace>,
 803        thread_store: Option<WeakModel<ThreadStore>>,
 804        cx: &mut ViewContext<PromptEditor<BufferCodegen>>,
 805    ) -> PromptEditor<BufferCodegen> {
 806        let codegen_subscription = cx.observe(&codegen, Self::handle_codegen_changed);
 807        let mode = PromptEditorMode::Buffer {
 808            id,
 809            codegen,
 810            gutter_dimensions,
 811        };
 812
 813        let prompt_editor = cx.new_view(|cx| {
 814            let mut editor = Editor::new(
 815                EditorMode::AutoHeight {
 816                    max_lines: Self::MAX_LINES as usize,
 817                },
 818                prompt_buffer,
 819                None,
 820                false,
 821                cx,
 822            );
 823            editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
 824            // Since the prompt editors for all inline assistants are linked,
 825            // always show the cursor (even when it isn't focused) because
 826            // typing in one will make what you typed appear in all of them.
 827            editor.set_show_cursor_when_unfocused(true, cx);
 828            editor.set_placeholder_text(Self::placeholder_text(&mode, cx), cx);
 829            editor
 830        });
 831        let context_picker_menu_handle = PopoverMenuHandle::default();
 832        let model_selector_menu_handle = PopoverMenuHandle::default();
 833
 834        let context_strip = cx.new_view(|cx| {
 835            ContextStrip::new(
 836                context_store.clone(),
 837                workspace.clone(),
 838                prompt_editor.downgrade(),
 839                thread_store.clone(),
 840                context_picker_menu_handle.clone(),
 841                SuggestContextKind::Thread,
 842                cx,
 843            )
 844        });
 845
 846        let context_strip_subscription =
 847            cx.subscribe(&context_strip, Self::handle_context_strip_event);
 848
 849        let mut this: PromptEditor<BufferCodegen> = PromptEditor {
 850            editor: prompt_editor.clone(),
 851            context_store,
 852            context_strip,
 853            context_picker_menu_handle,
 854            model_selector: cx.new_view(|cx| {
 855                AssistantModelSelector::new(
 856                    fs,
 857                    model_selector_menu_handle.clone(),
 858                    prompt_editor.focus_handle(cx),
 859                    cx,
 860                )
 861            }),
 862            model_selector_menu_handle,
 863            edited_since_done: false,
 864            prompt_history,
 865            prompt_history_ix: None,
 866            pending_prompt: String::new(),
 867            _codegen_subscription: codegen_subscription,
 868            editor_subscriptions: Vec::new(),
 869            _context_strip_subscription: context_strip_subscription,
 870            show_rate_limit_notice: false,
 871            mode,
 872            _phantom: Default::default(),
 873        };
 874
 875        this.subscribe_to_editor(cx);
 876        this
 877    }
 878
 879    fn handle_codegen_changed(
 880        &mut self,
 881        _: Model<BufferCodegen>,
 882        cx: &mut ViewContext<PromptEditor<BufferCodegen>>,
 883    ) {
 884        match self.codegen_status(cx) {
 885            CodegenStatus::Idle => {
 886                self.editor
 887                    .update(cx, |editor, _| editor.set_read_only(false));
 888            }
 889            CodegenStatus::Pending => {
 890                self.editor
 891                    .update(cx, |editor, _| editor.set_read_only(true));
 892            }
 893            CodegenStatus::Done => {
 894                self.edited_since_done = false;
 895                self.editor
 896                    .update(cx, |editor, _| editor.set_read_only(false));
 897            }
 898            CodegenStatus::Error(error) => {
 899                if cx.has_flag::<ZedPro>()
 900                    && error.error_code() == proto::ErrorCode::RateLimitExceeded
 901                    && !dismissed_rate_limit_notice()
 902                {
 903                    self.show_rate_limit_notice = true;
 904                    cx.notify();
 905                }
 906
 907                self.edited_since_done = false;
 908                self.editor
 909                    .update(cx, |editor, _| editor.set_read_only(false));
 910            }
 911        }
 912    }
 913
 914    pub fn id(&self) -> InlineAssistId {
 915        match &self.mode {
 916            PromptEditorMode::Buffer { id, .. } => *id,
 917            PromptEditorMode::Terminal { .. } => unreachable!(),
 918        }
 919    }
 920
 921    pub fn codegen(&self) -> &Model<BufferCodegen> {
 922        match &self.mode {
 923            PromptEditorMode::Buffer { codegen, .. } => codegen,
 924            PromptEditorMode::Terminal { .. } => unreachable!(),
 925        }
 926    }
 927
 928    pub fn gutter_dimensions(&self) -> &Arc<Mutex<GutterDimensions>> {
 929        match &self.mode {
 930            PromptEditorMode::Buffer {
 931                gutter_dimensions, ..
 932            } => gutter_dimensions,
 933            PromptEditorMode::Terminal { .. } => unreachable!(),
 934        }
 935    }
 936}
 937
 938#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
 939pub struct TerminalInlineAssistId(pub usize);
 940
 941impl TerminalInlineAssistId {
 942    pub fn post_inc(&mut self) -> TerminalInlineAssistId {
 943        let id = *self;
 944        self.0 += 1;
 945        id
 946    }
 947}
 948
 949impl PromptEditor<TerminalCodegen> {
 950    #[allow(clippy::too_many_arguments)]
 951    pub fn new_terminal(
 952        id: TerminalInlineAssistId,
 953        prompt_history: VecDeque<String>,
 954        prompt_buffer: Model<MultiBuffer>,
 955        codegen: Model<TerminalCodegen>,
 956        fs: Arc<dyn Fs>,
 957        context_store: Model<ContextStore>,
 958        workspace: WeakView<Workspace>,
 959        thread_store: Option<WeakModel<ThreadStore>>,
 960        cx: &mut ViewContext<Self>,
 961    ) -> Self {
 962        let codegen_subscription = cx.observe(&codegen, Self::handle_codegen_changed);
 963        let mode = PromptEditorMode::Terminal {
 964            id,
 965            codegen,
 966            height_in_lines: 1,
 967        };
 968
 969        let prompt_editor = cx.new_view(|cx| {
 970            let mut editor = Editor::new(
 971                EditorMode::AutoHeight {
 972                    max_lines: Self::MAX_LINES as usize,
 973                },
 974                prompt_buffer,
 975                None,
 976                false,
 977                cx,
 978            );
 979            editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
 980            editor.set_placeholder_text(Self::placeholder_text(&mode, cx), cx);
 981            editor
 982        });
 983        let context_picker_menu_handle = PopoverMenuHandle::default();
 984        let model_selector_menu_handle = PopoverMenuHandle::default();
 985
 986        let context_strip = cx.new_view(|cx| {
 987            ContextStrip::new(
 988                context_store.clone(),
 989                workspace.clone(),
 990                prompt_editor.downgrade(),
 991                thread_store.clone(),
 992                context_picker_menu_handle.clone(),
 993                SuggestContextKind::Thread,
 994                cx,
 995            )
 996        });
 997
 998        let context_strip_subscription =
 999            cx.subscribe(&context_strip, Self::handle_context_strip_event);
1000
1001        let mut this = Self {
1002            editor: prompt_editor.clone(),
1003            context_store,
1004            context_strip,
1005            context_picker_menu_handle,
1006            model_selector: cx.new_view(|cx| {
1007                AssistantModelSelector::new(
1008                    fs,
1009                    model_selector_menu_handle.clone(),
1010                    prompt_editor.focus_handle(cx),
1011                    cx,
1012                )
1013            }),
1014            model_selector_menu_handle,
1015            edited_since_done: false,
1016            prompt_history,
1017            prompt_history_ix: None,
1018            pending_prompt: String::new(),
1019            _codegen_subscription: codegen_subscription,
1020            editor_subscriptions: Vec::new(),
1021            _context_strip_subscription: context_strip_subscription,
1022            mode,
1023            show_rate_limit_notice: false,
1024            _phantom: Default::default(),
1025        };
1026        this.count_lines(cx);
1027        this.subscribe_to_editor(cx);
1028        this
1029    }
1030
1031    fn count_lines(&mut self, cx: &mut ViewContext<Self>) {
1032        let height_in_lines = cmp::max(
1033            2, // Make the editor at least two lines tall, to account for padding and buttons.
1034            cmp::min(
1035                self.editor
1036                    .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1),
1037                Self::MAX_LINES as u32,
1038            ),
1039        ) as u8;
1040
1041        match &mut self.mode {
1042            PromptEditorMode::Terminal {
1043                height_in_lines: current_height,
1044                ..
1045            } => {
1046                if height_in_lines != *current_height {
1047                    *current_height = height_in_lines;
1048                    cx.emit(PromptEditorEvent::Resized { height_in_lines });
1049                }
1050            }
1051            PromptEditorMode::Buffer { .. } => unreachable!(),
1052        }
1053    }
1054
1055    fn handle_codegen_changed(&mut self, _: Model<TerminalCodegen>, cx: &mut ViewContext<Self>) {
1056        match &self.codegen().read(cx).status {
1057            CodegenStatus::Idle => {
1058                self.editor
1059                    .update(cx, |editor, _| editor.set_read_only(false));
1060            }
1061            CodegenStatus::Pending => {
1062                self.editor
1063                    .update(cx, |editor, _| editor.set_read_only(true));
1064            }
1065            CodegenStatus::Done | CodegenStatus::Error(_) => {
1066                self.edited_since_done = false;
1067                self.editor
1068                    .update(cx, |editor, _| editor.set_read_only(false));
1069            }
1070        }
1071    }
1072
1073    pub fn codegen(&self) -> &Model<TerminalCodegen> {
1074        match &self.mode {
1075            PromptEditorMode::Buffer { .. } => unreachable!(),
1076            PromptEditorMode::Terminal { codegen, .. } => codegen,
1077        }
1078    }
1079
1080    pub fn id(&self) -> TerminalInlineAssistId {
1081        match &self.mode {
1082            PromptEditorMode::Buffer { .. } => unreachable!(),
1083            PromptEditorMode::Terminal { id, .. } => *id,
1084        }
1085    }
1086}
1087
1088const DISMISSED_RATE_LIMIT_NOTICE_KEY: &str = "dismissed-rate-limit-notice";
1089
1090fn dismissed_rate_limit_notice() -> bool {
1091    db::kvp::KEY_VALUE_STORE
1092        .read_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY)
1093        .log_err()
1094        .map_or(false, |s| s.is_some())
1095}
1096
1097fn set_rate_limit_notice_dismissed(is_dismissed: bool, cx: &mut AppContext) {
1098    db::write_and_log(cx, move || async move {
1099        if is_dismissed {
1100            db::kvp::KEY_VALUE_STORE
1101                .write_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY.into(), "1".into())
1102                .await
1103        } else {
1104            db::kvp::KEY_VALUE_STORE
1105                .delete_kvp(DISMISSED_RATE_LIMIT_NOTICE_KEY.into())
1106                .await
1107        }
1108    })
1109}
1110
1111pub enum CodegenStatus {
1112    Idle,
1113    Pending,
1114    Done,
1115    Error(anyhow::Error),
1116}
1117
1118/// This is just CodegenStatus without the anyhow::Error, which causes a lifetime issue for rendering the Cancel button.
1119#[derive(Copy, Clone)]
1120pub enum CancelButtonState {
1121    Idle,
1122    Pending,
1123    Done,
1124    Error,
1125}
1126
1127impl Into<CancelButtonState> for &CodegenStatus {
1128    fn into(self) -> CancelButtonState {
1129        match self {
1130            CodegenStatus::Idle => CancelButtonState::Idle,
1131            CodegenStatus::Pending => CancelButtonState::Pending,
1132            CodegenStatus::Done => CancelButtonState::Done,
1133            CodegenStatus::Error(_) => CancelButtonState::Error,
1134        }
1135    }
1136}
1137
1138#[derive(Copy, Clone)]
1139pub enum GenerationMode {
1140    Generate,
1141    Transform,
1142}
1143
1144impl GenerationMode {
1145    fn start_label(self) -> &'static str {
1146        match self {
1147            GenerationMode::Generate { .. } => "Generate",
1148            GenerationMode::Transform => "Transform",
1149        }
1150    }
1151    fn tooltip_interrupt(self) -> &'static str {
1152        match self {
1153            GenerationMode::Generate { .. } => "Interrupt Generation",
1154            GenerationMode::Transform => "Interrupt Transform",
1155        }
1156    }
1157
1158    fn tooltip_restart(self) -> &'static str {
1159        match self {
1160            GenerationMode::Generate { .. } => "Restart Generation",
1161            GenerationMode::Transform => "Restart Transform",
1162        }
1163    }
1164
1165    fn tooltip_accept(self) -> &'static str {
1166        match self {
1167            GenerationMode::Generate { .. } => "Accept Generation",
1168            GenerationMode::Transform => "Accept Transform",
1169        }
1170    }
1171}