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