inline_prompt_editor.rs

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