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