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