inline_prompt_editor.rs

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