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(cx);
  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, cx)
 275                .map(|keybinding| format!("{keybinding} to chat ― "))
 276                .unwrap_or_default();
 277
 278        format!("{action}… ({assistant_panel_keybinding}↓↑ for history)")
 279    }
 280
 281    pub fn prompt(&self, cx: &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.root::<Workspace>().flatten() {
 308                    workspace.update(cx, |workspace, cx| {
 309                        let is_via_ssh = workspace
 310                            .project()
 311                            .update(cx, |project, _| project.is_via_ssh());
 312
 313                        workspace
 314                            .client()
 315                            .telemetry()
 316                            .log_edit_event("inline assist", is_via_ssh);
 317                    });
 318                }
 319                let prompt = self.editor.read(cx).text(cx);
 320                if self
 321                    .prompt_history_ix
 322                    .map_or(true, |ix| self.prompt_history[ix] != prompt)
 323                {
 324                    self.prompt_history_ix.take();
 325                    self.pending_prompt = prompt;
 326                }
 327
 328                self.edited_since_done = true;
 329                cx.notify();
 330            }
 331            EditorEvent::Blurred => {
 332                if self.show_rate_limit_notice {
 333                    self.show_rate_limit_notice = false;
 334                    cx.notify();
 335                }
 336            }
 337            _ => {}
 338        }
 339    }
 340
 341    fn toggle_context_picker(
 342        &mut self,
 343        _: &ToggleContextPicker,
 344        window: &mut Window,
 345        cx: &mut Context<Self>,
 346    ) {
 347        self.context_picker_menu_handle.toggle(window, cx);
 348    }
 349
 350    fn toggle_model_selector(
 351        &mut self,
 352        _: &ToggleModelSelector,
 353        window: &mut Window,
 354        cx: &mut Context<Self>,
 355    ) {
 356        self.model_selector_menu_handle.toggle(window, cx);
 357    }
 358
 359    pub fn remove_all_context(
 360        &mut self,
 361        _: &RemoveAllContext,
 362        _window: &mut Window,
 363        cx: &mut Context<Self>,
 364    ) {
 365        self.context_store.update(cx, |store, _cx| store.clear());
 366        cx.notify();
 367    }
 368
 369    fn cancel(
 370        &mut self,
 371        _: &editor::actions::Cancel,
 372        _window: &mut Window,
 373        cx: &mut Context<Self>,
 374    ) {
 375        match self.codegen_status(cx) {
 376            CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {
 377                cx.emit(PromptEditorEvent::CancelRequested);
 378            }
 379            CodegenStatus::Pending => {
 380                cx.emit(PromptEditorEvent::StopRequested);
 381            }
 382        }
 383    }
 384
 385    fn confirm(&mut self, _: &menu::Confirm, _window: &mut Window, cx: &mut Context<Self>) {
 386        match self.codegen_status(cx) {
 387            CodegenStatus::Idle => {
 388                cx.emit(PromptEditorEvent::StartRequested);
 389            }
 390            CodegenStatus::Pending => {
 391                cx.emit(PromptEditorEvent::DismissRequested);
 392            }
 393            CodegenStatus::Done => {
 394                if self.edited_since_done {
 395                    cx.emit(PromptEditorEvent::StartRequested);
 396                } else {
 397                    cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });
 398                }
 399            }
 400            CodegenStatus::Error(_) => {
 401                cx.emit(PromptEditorEvent::StartRequested);
 402            }
 403        }
 404    }
 405
 406    fn move_up(&mut self, _: &MoveUp, window: &mut Window, cx: &mut Context<Self>) {
 407        if let Some(ix) = self.prompt_history_ix {
 408            if ix > 0 {
 409                self.prompt_history_ix = Some(ix - 1);
 410                let prompt = self.prompt_history[ix - 1].as_str();
 411                self.editor.update(cx, |editor, cx| {
 412                    editor.set_text(prompt, window, cx);
 413                    editor.move_to_beginning(&Default::default(), window, cx);
 414                });
 415            }
 416        } else if !self.prompt_history.is_empty() {
 417            self.prompt_history_ix = Some(self.prompt_history.len() - 1);
 418            let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();
 419            self.editor.update(cx, |editor, cx| {
 420                editor.set_text(prompt, window, cx);
 421                editor.move_to_beginning(&Default::default(), window, cx);
 422            });
 423        }
 424    }
 425
 426    fn move_down(&mut self, _: &MoveDown, window: &mut Window, cx: &mut Context<Self>) {
 427        if let Some(ix) = self.prompt_history_ix {
 428            if ix < self.prompt_history.len() - 1 {
 429                self.prompt_history_ix = Some(ix + 1);
 430                let prompt = self.prompt_history[ix + 1].as_str();
 431                self.editor.update(cx, |editor, cx| {
 432                    editor.set_text(prompt, window, cx);
 433                    editor.move_to_end(&Default::default(), window, cx)
 434                });
 435            } else {
 436                self.prompt_history_ix = None;
 437                let prompt = self.pending_prompt.as_str();
 438                self.editor.update(cx, |editor, cx| {
 439                    editor.set_text(prompt, window, cx);
 440                    editor.move_to_end(&Default::default(), window, cx)
 441                });
 442            }
 443        } else {
 444            self.context_strip.focus_handle(cx).focus(window);
 445        }
 446    }
 447
 448    fn render_buttons(&self, _window: &mut Window, cx: &mut Context<Self>) -> Vec<AnyElement> {
 449        let mode = match &self.mode {
 450            PromptEditorMode::Buffer { codegen, .. } => {
 451                let codegen = codegen.read(cx);
 452                if codegen.is_insertion {
 453                    GenerationMode::Generate
 454                } else {
 455                    GenerationMode::Transform
 456                }
 457            }
 458            PromptEditorMode::Terminal { .. } => GenerationMode::Generate,
 459        };
 460
 461        let codegen_status = self.codegen_status(cx);
 462
 463        match codegen_status {
 464            CodegenStatus::Idle => {
 465                vec![Button::new("start", mode.start_label())
 466                    .label_size(LabelSize::Small)
 467                    .icon(IconName::Return)
 468                    .icon_size(IconSize::XSmall)
 469                    .icon_color(Color::Muted)
 470                    .on_click(cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::StartRequested)))
 471                    .into_any_element()]
 472            }
 473            CodegenStatus::Pending => vec![IconButton::new("stop", IconName::Stop)
 474                .icon_color(Color::Error)
 475                .shape(IconButtonShape::Square)
 476                .tooltip(move |window, cx| {
 477                    Tooltip::with_meta(
 478                        mode.tooltip_interrupt(),
 479                        Some(&menu::Cancel),
 480                        "Changes won't be discarded",
 481                        window,
 482                        cx,
 483                    )
 484                })
 485                .on_click(cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::StopRequested)))
 486                .into_any_element()],
 487            CodegenStatus::Done | CodegenStatus::Error(_) => {
 488                let has_error = matches!(codegen_status, CodegenStatus::Error(_));
 489                if has_error || self.edited_since_done {
 490                    vec![IconButton::new("restart", IconName::RotateCw)
 491                        .icon_color(Color::Info)
 492                        .shape(IconButtonShape::Square)
 493                        .tooltip(move |window, cx| {
 494                            Tooltip::with_meta(
 495                                mode.tooltip_restart(),
 496                                Some(&menu::Confirm),
 497                                "Changes will be discarded",
 498                                window,
 499                                cx,
 500                            )
 501                        })
 502                        .on_click(cx.listener(|_, _, _, cx| {
 503                            cx.emit(PromptEditorEvent::StartRequested);
 504                        }))
 505                        .into_any_element()]
 506                } else {
 507                    let accept = IconButton::new("accept", IconName::Check)
 508                        .icon_color(Color::Info)
 509                        .shape(IconButtonShape::Square)
 510                        .tooltip(move |window, cx| {
 511                            Tooltip::for_action(mode.tooltip_accept(), &menu::Confirm, window, cx)
 512                        })
 513                        .on_click(cx.listener(|_, _, _, cx| {
 514                            cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });
 515                        }))
 516                        .into_any_element();
 517
 518                    match &self.mode {
 519                        PromptEditorMode::Terminal { .. } => vec![
 520                            accept,
 521                            IconButton::new("confirm", IconName::Play)
 522                                .icon_color(Color::Info)
 523                                .shape(IconButtonShape::Square)
 524                                .tooltip(|window, cx| {
 525                                    Tooltip::for_action(
 526                                        "Execute Generated Command",
 527                                        &menu::SecondaryConfirm,
 528                                        window,
 529                                        cx,
 530                                    )
 531                                })
 532                                .on_click(cx.listener(|_, _, _, cx| {
 533                                    cx.emit(PromptEditorEvent::ConfirmRequested { execute: true });
 534                                }))
 535                                .into_any_element(),
 536                        ],
 537                        PromptEditorMode::Buffer { .. } => vec![accept],
 538                    }
 539                }
 540            }
 541        }
 542    }
 543
 544    fn cycle_prev(
 545        &mut self,
 546        _: &CyclePreviousInlineAssist,
 547        _: &mut Window,
 548        cx: &mut Context<Self>,
 549    ) {
 550        match &self.mode {
 551            PromptEditorMode::Buffer { codegen, .. } => {
 552                codegen.update(cx, |codegen, cx| codegen.cycle_prev(cx));
 553            }
 554            PromptEditorMode::Terminal { .. } => {
 555                // no cycle buttons in terminal mode
 556            }
 557        }
 558    }
 559
 560    fn cycle_next(&mut self, _: &CycleNextInlineAssist, _: &mut Window, cx: &mut Context<Self>) {
 561        match &self.mode {
 562            PromptEditorMode::Buffer { codegen, .. } => {
 563                codegen.update(cx, |codegen, cx| codegen.cycle_next(cx));
 564            }
 565            PromptEditorMode::Terminal { .. } => {
 566                // no cycle buttons in terminal mode
 567            }
 568        }
 569    }
 570
 571    fn render_close_button(&self, cx: &mut Context<Self>) -> AnyElement {
 572        IconButton::new("cancel", IconName::Close)
 573            .icon_color(Color::Muted)
 574            .shape(IconButtonShape::Square)
 575            .tooltip(Tooltip::text("Close Assistant"))
 576            .on_click(cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::CancelRequested)))
 577            .into_any_element()
 578    }
 579
 580    fn render_cycle_controls(&self, codegen: &BufferCodegen, cx: &Context<Self>) -> AnyElement {
 581        let disabled = matches!(codegen.status(cx), CodegenStatus::Idle);
 582
 583        let model_registry = LanguageModelRegistry::read_global(cx);
 584        let default_model = model_registry.active_model();
 585        let alternative_models = model_registry.inline_alternative_models();
 586
 587        let get_model_name = |index: usize| -> String {
 588            let name = |model: &Arc<dyn LanguageModel>| model.name().0.to_string();
 589
 590            match index {
 591                0 => default_model.as_ref().map_or_else(String::new, name),
 592                index if index <= alternative_models.len() => alternative_models
 593                    .get(index - 1)
 594                    .map_or_else(String::new, name),
 595                _ => String::new(),
 596            }
 597        };
 598
 599        let total_models = alternative_models.len() + 1;
 600
 601        if total_models <= 1 {
 602            return div().into_any_element();
 603        }
 604
 605        let current_index = codegen.active_alternative;
 606        let prev_index = (current_index + total_models - 1) % total_models;
 607        let next_index = (current_index + 1) % total_models;
 608
 609        let prev_model_name = get_model_name(prev_index);
 610        let next_model_name = get_model_name(next_index);
 611
 612        h_flex()
 613            .child(
 614                IconButton::new("previous", IconName::ChevronLeft)
 615                    .icon_color(Color::Muted)
 616                    .disabled(disabled || current_index == 0)
 617                    .shape(IconButtonShape::Square)
 618                    .tooltip({
 619                        let focus_handle = self.editor.focus_handle(cx);
 620                        move |window, cx| {
 621                            cx.new(|cx| {
 622                                let mut tooltip = Tooltip::new("Previous Alternative").key_binding(
 623                                    KeyBinding::for_action_in(
 624                                        &CyclePreviousInlineAssist,
 625                                        &focus_handle,
 626                                        window,
 627                                        cx,
 628                                    ),
 629                                );
 630                                if !disabled && current_index != 0 {
 631                                    tooltip = tooltip.meta(prev_model_name.clone());
 632                                }
 633                                tooltip
 634                            })
 635                            .into()
 636                        }
 637                    })
 638                    .on_click(cx.listener(|this, _, window, cx| {
 639                        this.cycle_prev(&CyclePreviousInlineAssist, window, cx);
 640                    })),
 641            )
 642            .child(
 643                Label::new(format!(
 644                    "{}/{}",
 645                    codegen.active_alternative + 1,
 646                    codegen.alternative_count(cx)
 647                ))
 648                .size(LabelSize::Small)
 649                .color(if disabled {
 650                    Color::Disabled
 651                } else {
 652                    Color::Muted
 653                }),
 654            )
 655            .child(
 656                IconButton::new("next", IconName::ChevronRight)
 657                    .icon_color(Color::Muted)
 658                    .disabled(disabled || current_index == total_models - 1)
 659                    .shape(IconButtonShape::Square)
 660                    .tooltip({
 661                        let focus_handle = self.editor.focus_handle(cx);
 662                        move |window, cx| {
 663                            cx.new(|cx| {
 664                                let mut tooltip = Tooltip::new("Next Alternative").key_binding(
 665                                    KeyBinding::for_action_in(
 666                                        &CycleNextInlineAssist,
 667                                        &focus_handle,
 668                                        window,
 669                                        cx,
 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}