inline_prompt_editor.rs

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