inline_prompt_editor.rs

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