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