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, ClipboardEntry, Context, CursorStyle, Entity, EventEmitter, FocusHandle,
  15    Focusable, 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::{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…", window, 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        let images = cx
 276            .read_from_clipboard()
 277            .map(|item| {
 278                item.into_entries()
 279                    .filter_map(|entry| {
 280                        if let ClipboardEntry::Image(image) = entry {
 281                            Some(image)
 282                        } else {
 283                            None
 284                        }
 285                    })
 286                    .collect::<Vec<_>>()
 287            })
 288            .unwrap_or_default();
 289
 290        if images.is_empty() {
 291            return;
 292        }
 293        cx.stop_propagation();
 294
 295        self.context_store.update(cx, |store, cx| {
 296            for image in images {
 297                store.add_image_instance(Arc::new(image), cx);
 298            }
 299        });
 300    }
 301
 302    fn handle_prompt_editor_events(
 303        &mut self,
 304        _: &Entity<Editor>,
 305        event: &EditorEvent,
 306        window: &mut Window,
 307        cx: &mut Context<Self>,
 308    ) {
 309        match event {
 310            EditorEvent::Edited { .. } => {
 311                if let Some(workspace) = window.root::<Workspace>().flatten() {
 312                    workspace.update(cx, |workspace, cx| {
 313                        let is_via_ssh = workspace.project().read(cx).is_via_remote_server();
 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                    .is_none_or(|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(cx));
 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            CodegenStatus::Done => {
 385                if self.edited_since_done {
 386                    cx.emit(PromptEditorEvent::StartRequested);
 387                } else {
 388                    cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });
 389                }
 390            }
 391            CodegenStatus::Error(_) => {
 392                cx.emit(PromptEditorEvent::StartRequested);
 393            }
 394        }
 395    }
 396
 397    fn move_up(&mut self, _: &MoveUp, window: &mut Window, cx: &mut Context<Self>) {
 398        if let Some(ix) = self.prompt_history_ix {
 399            if ix > 0 {
 400                self.prompt_history_ix = Some(ix - 1);
 401                let prompt = self.prompt_history[ix - 1].as_str();
 402                self.editor.update(cx, |editor, cx| {
 403                    editor.set_text(prompt, window, cx);
 404                    editor.move_to_beginning(&Default::default(), window, cx);
 405                });
 406            }
 407        } else if !self.prompt_history.is_empty() {
 408            self.prompt_history_ix = Some(self.prompt_history.len() - 1);
 409            let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();
 410            self.editor.update(cx, |editor, cx| {
 411                editor.set_text(prompt, window, cx);
 412                editor.move_to_beginning(&Default::default(), window, cx);
 413            });
 414        }
 415    }
 416
 417    fn move_down(&mut self, _: &MoveDown, window: &mut Window, cx: &mut Context<Self>) {
 418        if let Some(ix) = self.prompt_history_ix {
 419            if ix < self.prompt_history.len() - 1 {
 420                self.prompt_history_ix = Some(ix + 1);
 421                let prompt = self.prompt_history[ix + 1].as_str();
 422                self.editor.update(cx, |editor, cx| {
 423                    editor.set_text(prompt, window, cx);
 424                    editor.move_to_end(&Default::default(), window, cx)
 425                });
 426            } else {
 427                self.prompt_history_ix = None;
 428                let prompt = self.pending_prompt.as_str();
 429                self.editor.update(cx, |editor, cx| {
 430                    editor.set_text(prompt, window, cx);
 431                    editor.move_to_end(&Default::default(), window, cx)
 432                });
 433            }
 434        } else if self.context_strip.read(cx).has_context_items(cx) {
 435            self.context_strip.focus_handle(cx).focus(window);
 436        }
 437    }
 438
 439    fn render_buttons(&self, _window: &mut Window, cx: &mut Context<Self>) -> Vec<AnyElement> {
 440        let mode = match &self.mode {
 441            PromptEditorMode::Buffer { codegen, .. } => {
 442                let codegen = codegen.read(cx);
 443                if codegen.is_insertion {
 444                    GenerationMode::Generate
 445                } else {
 446                    GenerationMode::Transform
 447                }
 448            }
 449            PromptEditorMode::Terminal { .. } => GenerationMode::Generate,
 450        };
 451
 452        let codegen_status = self.codegen_status(cx);
 453
 454        match codegen_status {
 455            CodegenStatus::Idle => {
 456                vec![
 457                    Button::new("start", mode.start_label())
 458                        .label_size(LabelSize::Small)
 459                        .icon(IconName::Return)
 460                        .icon_size(IconSize::XSmall)
 461                        .icon_color(Color::Muted)
 462                        .on_click(
 463                            cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::StartRequested)),
 464                        )
 465                        .into_any_element(),
 466                ]
 467            }
 468            CodegenStatus::Pending => vec![
 469                IconButton::new("stop", IconName::Stop)
 470                    .icon_color(Color::Error)
 471                    .tooltip(move |window, cx| {
 472                        Tooltip::with_meta(
 473                            mode.tooltip_interrupt(),
 474                            Some(&menu::Cancel),
 475                            "Changes won't be discarded",
 476                            window,
 477                            cx,
 478                        )
 479                    })
 480                    .on_click(cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::StopRequested)))
 481                    .into_any_element(),
 482            ],
 483            CodegenStatus::Done | CodegenStatus::Error(_) => {
 484                let has_error = matches!(codegen_status, CodegenStatus::Error(_));
 485                if has_error || self.edited_since_done {
 486                    vec![
 487                        IconButton::new("restart", IconName::RotateCw)
 488                            .icon_color(Color::Info)
 489                            .tooltip(move |window, cx| {
 490                                Tooltip::with_meta(
 491                                    mode.tooltip_restart(),
 492                                    Some(&menu::Confirm),
 493                                    "Changes will be discarded",
 494                                    window,
 495                                    cx,
 496                                )
 497                            })
 498                            .on_click(cx.listener(|_, _, _, cx| {
 499                                cx.emit(PromptEditorEvent::StartRequested);
 500                            }))
 501                            .into_any_element(),
 502                    ]
 503                } else {
 504                    let accept = IconButton::new("accept", IconName::Check)
 505                        .icon_color(Color::Info)
 506                        .tooltip(move |window, cx| {
 507                            Tooltip::for_action(mode.tooltip_accept(), &menu::Confirm, window, cx)
 508                        })
 509                        .on_click(cx.listener(|_, _, _, cx| {
 510                            cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });
 511                        }))
 512                        .into_any_element();
 513
 514                    match &self.mode {
 515                        PromptEditorMode::Terminal { .. } => vec![
 516                            accept,
 517                            IconButton::new("confirm", IconName::PlayFilled)
 518                                .icon_color(Color::Info)
 519                                .tooltip(|window, cx| {
 520                                    Tooltip::for_action(
 521                                        "Execute Generated Command",
 522                                        &menu::SecondaryConfirm,
 523                                        window,
 524                                        cx,
 525                                    )
 526                                })
 527                                .on_click(cx.listener(|_, _, _, cx| {
 528                                    cx.emit(PromptEditorEvent::ConfirmRequested { execute: true });
 529                                }))
 530                                .into_any_element(),
 531                        ],
 532                        PromptEditorMode::Buffer { .. } => vec![accept],
 533                    }
 534                }
 535            }
 536        }
 537    }
 538
 539    fn cycle_prev(
 540        &mut self,
 541        _: &CyclePreviousInlineAssist,
 542        _: &mut Window,
 543        cx: &mut Context<Self>,
 544    ) {
 545        match &self.mode {
 546            PromptEditorMode::Buffer { codegen, .. } => {
 547                codegen.update(cx, |codegen, cx| codegen.cycle_prev(cx));
 548            }
 549            PromptEditorMode::Terminal { .. } => {
 550                // no cycle buttons in terminal mode
 551            }
 552        }
 553    }
 554
 555    fn cycle_next(&mut self, _: &CycleNextInlineAssist, _: &mut Window, cx: &mut Context<Self>) {
 556        match &self.mode {
 557            PromptEditorMode::Buffer { codegen, .. } => {
 558                codegen.update(cx, |codegen, cx| codegen.cycle_next(cx));
 559            }
 560            PromptEditorMode::Terminal { .. } => {
 561                // no cycle buttons in terminal mode
 562            }
 563        }
 564    }
 565
 566    fn render_close_button(&self, cx: &mut Context<Self>) -> AnyElement {
 567        IconButton::new("cancel", IconName::Close)
 568            .icon_color(Color::Muted)
 569            .tooltip(Tooltip::text("Close Assistant"))
 570            .on_click(cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::CancelRequested)))
 571            .into_any_element()
 572    }
 573
 574    fn render_cycle_controls(&self, codegen: &BufferCodegen, cx: &Context<Self>) -> AnyElement {
 575        let disabled = matches!(codegen.status(cx), CodegenStatus::Idle);
 576
 577        let model_registry = LanguageModelRegistry::read_global(cx);
 578        let default_model = model_registry.default_model().map(|default| default.model);
 579        let alternative_models = model_registry.inline_alternative_models();
 580
 581        let get_model_name = |index: usize| -> String {
 582            let name = |model: &Arc<dyn LanguageModel>| model.name().0.to_string();
 583
 584            match index {
 585                0 => default_model.as_ref().map_or_else(String::new, name),
 586                index if index <= alternative_models.len() => alternative_models
 587                    .get(index - 1)
 588                    .map_or_else(String::new, name),
 589                _ => String::new(),
 590            }
 591        };
 592
 593        let total_models = alternative_models.len() + 1;
 594
 595        if total_models <= 1 {
 596            return div().into_any_element();
 597        }
 598
 599        let current_index = codegen.active_alternative;
 600        let prev_index = (current_index + total_models - 1) % total_models;
 601        let next_index = (current_index + 1) % total_models;
 602
 603        let prev_model_name = get_model_name(prev_index);
 604        let next_model_name = get_model_name(next_index);
 605
 606        h_flex()
 607            .child(
 608                IconButton::new("previous", IconName::ChevronLeft)
 609                    .icon_color(Color::Muted)
 610                    .disabled(disabled || current_index == 0)
 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                    .tooltip({
 653                        let focus_handle = self.editor.focus_handle(cx);
 654                        move |window, cx| {
 655                            cx.new(|cx| {
 656                                let mut tooltip = Tooltip::new("Next Alternative").key_binding(
 657                                    KeyBinding::for_action_in(
 658                                        &CycleNextInlineAssist,
 659                                        &focus_handle,
 660                                        window,
 661                                        cx,
 662                                    ),
 663                                );
 664                                if !disabled && current_index != total_models - 1 {
 665                                    tooltip = tooltip.meta(next_model_name.clone());
 666                                }
 667                                tooltip
 668                            })
 669                            .into()
 670                        }
 671                    })
 672                    .on_click(cx.listener(|this, _, window, cx| {
 673                        this.cycle_next(&CycleNextInlineAssist, window, cx)
 674                    })),
 675            )
 676            .into_any_element()
 677    }
 678
 679    fn render_editor(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> AnyElement {
 680        let colors = cx.theme().colors();
 681
 682        div()
 683            .key_context("InlineAssistEditor")
 684            .size_full()
 685            .p_2()
 686            .pl_1()
 687            .bg(colors.editor_background)
 688            .child({
 689                let settings = ThemeSettings::get_global(cx);
 690                let font_size = settings.buffer_font_size(cx);
 691                let line_height = font_size * 1.2;
 692
 693                let text_style = TextStyle {
 694                    color: colors.editor_foreground,
 695                    font_family: settings.buffer_font.family.clone(),
 696                    font_features: settings.buffer_font.features.clone(),
 697                    font_size: font_size.into(),
 698                    line_height: line_height.into(),
 699                    ..Default::default()
 700                };
 701
 702                EditorElement::new(
 703                    &self.editor,
 704                    EditorStyle {
 705                        background: colors.editor_background,
 706                        local_player: cx.theme().players().local(),
 707                        text: text_style,
 708                        ..Default::default()
 709                    },
 710                )
 711            })
 712            .into_any_element()
 713    }
 714
 715    fn handle_context_strip_event(
 716        &mut self,
 717        _context_strip: &Entity<ContextStrip>,
 718        event: &ContextStripEvent,
 719        window: &mut Window,
 720        cx: &mut Context<Self>,
 721    ) {
 722        match event {
 723            ContextStripEvent::PickerDismissed
 724            | ContextStripEvent::BlurredEmpty
 725            | ContextStripEvent::BlurredUp => self.editor.focus_handle(cx).focus(window),
 726            ContextStripEvent::BlurredDown => {}
 727        }
 728    }
 729}
 730
 731pub enum PromptEditorMode {
 732    Buffer {
 733        id: InlineAssistId,
 734        codegen: Entity<BufferCodegen>,
 735        editor_margins: Arc<Mutex<EditorMargins>>,
 736    },
 737    Terminal {
 738        id: TerminalInlineAssistId,
 739        codegen: Entity<TerminalCodegen>,
 740        height_in_lines: u8,
 741    },
 742}
 743
 744pub enum PromptEditorEvent {
 745    StartRequested,
 746    StopRequested,
 747    ConfirmRequested { execute: bool },
 748    CancelRequested,
 749    Resized { height_in_lines: u8 },
 750}
 751
 752#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
 753pub struct InlineAssistId(pub usize);
 754
 755impl InlineAssistId {
 756    pub fn post_inc(&mut self) -> InlineAssistId {
 757        let id = *self;
 758        self.0 += 1;
 759        id
 760    }
 761}
 762
 763impl PromptEditor<BufferCodegen> {
 764    pub fn new_buffer(
 765        id: InlineAssistId,
 766        editor_margins: Arc<Mutex<EditorMargins>>,
 767        prompt_history: VecDeque<String>,
 768        prompt_buffer: Entity<MultiBuffer>,
 769        codegen: Entity<BufferCodegen>,
 770        fs: Arc<dyn Fs>,
 771        context_store: Entity<ContextStore>,
 772        workspace: WeakEntity<Workspace>,
 773        thread_store: Option<WeakEntity<ThreadStore>>,
 774        text_thread_store: Option<WeakEntity<TextThreadStore>>,
 775        window: &mut Window,
 776        cx: &mut Context<PromptEditor<BufferCodegen>>,
 777    ) -> PromptEditor<BufferCodegen> {
 778        let codegen_subscription = cx.observe(&codegen, Self::handle_codegen_changed);
 779        let codegen_buffer = codegen.read(cx).buffer(cx).read(cx).as_singleton();
 780        let mode = PromptEditorMode::Buffer {
 781            id,
 782            codegen,
 783            editor_margins,
 784        };
 785
 786        let prompt_editor = cx.new(|cx| {
 787            let mut editor = Editor::new(
 788                EditorMode::AutoHeight {
 789                    min_lines: 1,
 790                    max_lines: Some(Self::MAX_LINES as usize),
 791                },
 792                prompt_buffer,
 793                None,
 794                window,
 795                cx,
 796            );
 797            editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
 798            // Since the prompt editors for all inline assistants are linked,
 799            // always show the cursor (even when it isn't focused) because
 800            // typing in one will make what you typed appear in all of them.
 801            editor.set_show_cursor_when_unfocused(true, cx);
 802            editor.set_placeholder_text(&Self::placeholder_text(&mode, window, cx), window, cx);
 803            editor.register_addon(ContextCreasesAddon::new());
 804            editor.set_context_menu_options(ContextMenuOptions {
 805                min_entries_visible: 12,
 806                max_entries_visible: 12,
 807                placement: None,
 808            });
 809
 810            editor
 811        });
 812
 813        let prompt_editor_entity = prompt_editor.downgrade();
 814        prompt_editor.update(cx, |editor, _| {
 815            editor.set_completion_provider(Some(Rc::new(ContextPickerCompletionProvider::new(
 816                workspace.clone(),
 817                context_store.downgrade(),
 818                thread_store.clone(),
 819                text_thread_store.clone(),
 820                prompt_editor_entity,
 821                codegen_buffer.as_ref().map(Entity::downgrade),
 822            ))));
 823        });
 824
 825        let context_picker_menu_handle = PopoverMenuHandle::default();
 826        let model_selector_menu_handle = PopoverMenuHandle::default();
 827
 828        let context_strip = cx.new(|cx| {
 829            ContextStrip::new(
 830                context_store.clone(),
 831                workspace.clone(),
 832                thread_store.clone(),
 833                text_thread_store.clone(),
 834                context_picker_menu_handle.clone(),
 835                SuggestContextKind::Thread,
 836                ModelUsageContext::InlineAssistant,
 837                window,
 838                cx,
 839            )
 840        });
 841
 842        let context_strip_subscription =
 843            cx.subscribe_in(&context_strip, window, Self::handle_context_strip_event);
 844
 845        let mut this: PromptEditor<BufferCodegen> = PromptEditor {
 846            editor: prompt_editor.clone(),
 847            context_store,
 848            context_strip,
 849            context_picker_menu_handle,
 850            model_selector: cx.new(|cx| {
 851                AgentModelSelector::new(
 852                    fs,
 853                    model_selector_menu_handle,
 854                    prompt_editor.focus_handle(cx),
 855                    ModelUsageContext::InlineAssistant,
 856                    window,
 857                    cx,
 858                )
 859            }),
 860            edited_since_done: false,
 861            prompt_history,
 862            prompt_history_ix: None,
 863            pending_prompt: String::new(),
 864            _codegen_subscription: codegen_subscription,
 865            editor_subscriptions: Vec::new(),
 866            _context_strip_subscription: context_strip_subscription,
 867            show_rate_limit_notice: false,
 868            mode,
 869            _phantom: Default::default(),
 870        };
 871
 872        this.subscribe_to_editor(window, cx);
 873        this
 874    }
 875
 876    fn handle_codegen_changed(
 877        &mut self,
 878        _: Entity<BufferCodegen>,
 879        cx: &mut Context<PromptEditor<BufferCodegen>>,
 880    ) {
 881        match self.codegen_status(cx) {
 882            CodegenStatus::Idle => {
 883                self.editor
 884                    .update(cx, |editor, _| editor.set_read_only(false));
 885            }
 886            CodegenStatus::Pending => {
 887                self.editor
 888                    .update(cx, |editor, _| editor.set_read_only(true));
 889            }
 890            CodegenStatus::Done => {
 891                self.edited_since_done = false;
 892                self.editor
 893                    .update(cx, |editor, _| editor.set_read_only(false));
 894            }
 895            CodegenStatus::Error(_error) => {
 896                self.edited_since_done = false;
 897                self.editor
 898                    .update(cx, |editor, _| editor.set_read_only(false));
 899            }
 900        }
 901    }
 902
 903    pub fn id(&self) -> InlineAssistId {
 904        match &self.mode {
 905            PromptEditorMode::Buffer { id, .. } => *id,
 906            PromptEditorMode::Terminal { .. } => unreachable!(),
 907        }
 908    }
 909
 910    pub fn codegen(&self) -> &Entity<BufferCodegen> {
 911        match &self.mode {
 912            PromptEditorMode::Buffer { codegen, .. } => codegen,
 913            PromptEditorMode::Terminal { .. } => unreachable!(),
 914        }
 915    }
 916
 917    pub fn editor_margins(&self) -> &Arc<Mutex<EditorMargins>> {
 918        match &self.mode {
 919            PromptEditorMode::Buffer { editor_margins, .. } => editor_margins,
 920            PromptEditorMode::Terminal { .. } => unreachable!(),
 921        }
 922    }
 923}
 924
 925#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
 926pub struct TerminalInlineAssistId(pub usize);
 927
 928impl TerminalInlineAssistId {
 929    pub fn post_inc(&mut self) -> TerminalInlineAssistId {
 930        let id = *self;
 931        self.0 += 1;
 932        id
 933    }
 934}
 935
 936impl PromptEditor<TerminalCodegen> {
 937    pub fn new_terminal(
 938        id: TerminalInlineAssistId,
 939        prompt_history: VecDeque<String>,
 940        prompt_buffer: Entity<MultiBuffer>,
 941        codegen: Entity<TerminalCodegen>,
 942        fs: Arc<dyn Fs>,
 943        context_store: Entity<ContextStore>,
 944        workspace: WeakEntity<Workspace>,
 945        thread_store: Option<WeakEntity<ThreadStore>>,
 946        text_thread_store: Option<WeakEntity<TextThreadStore>>,
 947        window: &mut Window,
 948        cx: &mut Context<Self>,
 949    ) -> Self {
 950        let codegen_subscription = cx.observe(&codegen, Self::handle_codegen_changed);
 951        let mode = PromptEditorMode::Terminal {
 952            id,
 953            codegen,
 954            height_in_lines: 1,
 955        };
 956
 957        let prompt_editor = cx.new(|cx| {
 958            let mut editor = Editor::new(
 959                EditorMode::AutoHeight {
 960                    min_lines: 1,
 961                    max_lines: Some(Self::MAX_LINES as usize),
 962                },
 963                prompt_buffer,
 964                None,
 965                window,
 966                cx,
 967            );
 968            editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
 969            editor.set_placeholder_text(&Self::placeholder_text(&mode, window, cx), window, cx);
 970            editor.set_context_menu_options(ContextMenuOptions {
 971                min_entries_visible: 12,
 972                max_entries_visible: 12,
 973                placement: None,
 974            });
 975            editor
 976        });
 977
 978        let prompt_editor_entity = prompt_editor.downgrade();
 979        prompt_editor.update(cx, |editor, _| {
 980            editor.set_completion_provider(Some(Rc::new(ContextPickerCompletionProvider::new(
 981                workspace.clone(),
 982                context_store.downgrade(),
 983                thread_store.clone(),
 984                text_thread_store.clone(),
 985                prompt_editor_entity,
 986                None,
 987            ))));
 988        });
 989
 990        let context_picker_menu_handle = PopoverMenuHandle::default();
 991        let model_selector_menu_handle = PopoverMenuHandle::default();
 992
 993        let context_strip = cx.new(|cx| {
 994            ContextStrip::new(
 995                context_store.clone(),
 996                workspace.clone(),
 997                thread_store.clone(),
 998                text_thread_store.clone(),
 999                context_picker_menu_handle.clone(),
1000                SuggestContextKind::Thread,
1001                ModelUsageContext::InlineAssistant,
1002                window,
1003                cx,
1004            )
1005        });
1006
1007        let context_strip_subscription =
1008            cx.subscribe_in(&context_strip, window, Self::handle_context_strip_event);
1009
1010        let mut this = Self {
1011            editor: prompt_editor.clone(),
1012            context_store,
1013            context_strip,
1014            context_picker_menu_handle,
1015            model_selector: cx.new(|cx| {
1016                AgentModelSelector::new(
1017                    fs,
1018                    model_selector_menu_handle.clone(),
1019                    prompt_editor.focus_handle(cx),
1020                    ModelUsageContext::InlineAssistant,
1021                    window,
1022                    cx,
1023                )
1024            }),
1025            edited_since_done: false,
1026            prompt_history,
1027            prompt_history_ix: None,
1028            pending_prompt: String::new(),
1029            _codegen_subscription: codegen_subscription,
1030            editor_subscriptions: Vec::new(),
1031            _context_strip_subscription: context_strip_subscription,
1032            mode,
1033            show_rate_limit_notice: false,
1034            _phantom: Default::default(),
1035        };
1036        this.count_lines(cx);
1037        this.subscribe_to_editor(window, cx);
1038        this
1039    }
1040
1041    fn count_lines(&mut self, cx: &mut Context<Self>) {
1042        let height_in_lines = cmp::max(
1043            2, // Make the editor at least two lines tall, to account for padding and buttons.
1044            cmp::min(
1045                self.editor
1046                    .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1),
1047                Self::MAX_LINES as u32,
1048            ),
1049        ) as u8;
1050
1051        match &mut self.mode {
1052            PromptEditorMode::Terminal {
1053                height_in_lines: current_height,
1054                ..
1055            } => {
1056                if height_in_lines != *current_height {
1057                    *current_height = height_in_lines;
1058                    cx.emit(PromptEditorEvent::Resized { height_in_lines });
1059                }
1060            }
1061            PromptEditorMode::Buffer { .. } => unreachable!(),
1062        }
1063    }
1064
1065    fn handle_codegen_changed(&mut self, _: Entity<TerminalCodegen>, cx: &mut Context<Self>) {
1066        match &self.codegen().read(cx).status {
1067            CodegenStatus::Idle => {
1068                self.editor
1069                    .update(cx, |editor, _| editor.set_read_only(false));
1070            }
1071            CodegenStatus::Pending => {
1072                self.editor
1073                    .update(cx, |editor, _| editor.set_read_only(true));
1074            }
1075            CodegenStatus::Done | CodegenStatus::Error(_) => {
1076                self.edited_since_done = false;
1077                self.editor
1078                    .update(cx, |editor, _| editor.set_read_only(false));
1079            }
1080        }
1081    }
1082
1083    pub fn codegen(&self) -> &Entity<TerminalCodegen> {
1084        match &self.mode {
1085            PromptEditorMode::Buffer { .. } => unreachable!(),
1086            PromptEditorMode::Terminal { codegen, .. } => codegen,
1087        }
1088    }
1089
1090    pub fn id(&self) -> TerminalInlineAssistId {
1091        match &self.mode {
1092            PromptEditorMode::Buffer { .. } => unreachable!(),
1093            PromptEditorMode::Terminal { id, .. } => *id,
1094        }
1095    }
1096}
1097
1098pub enum CodegenStatus {
1099    Idle,
1100    Pending,
1101    Done,
1102    Error(anyhow::Error),
1103}
1104
1105/// This is just CodegenStatus without the anyhow::Error, which causes a lifetime issue for rendering the Cancel button.
1106#[derive(Copy, Clone)]
1107pub enum CancelButtonState {
1108    Idle,
1109    Pending,
1110    Done,
1111    Error,
1112}
1113
1114impl Into<CancelButtonState> for &CodegenStatus {
1115    fn into(self) -> CancelButtonState {
1116        match self {
1117            CodegenStatus::Idle => CancelButtonState::Idle,
1118            CodegenStatus::Pending => CancelButtonState::Pending,
1119            CodegenStatus::Done => CancelButtonState::Done,
1120            CodegenStatus::Error(_) => CancelButtonState::Error,
1121        }
1122    }
1123}
1124
1125#[derive(Copy, Clone)]
1126pub enum GenerationMode {
1127    Generate,
1128    Transform,
1129}
1130
1131impl GenerationMode {
1132    fn start_label(self) -> &'static str {
1133        match self {
1134            GenerationMode::Generate => "Generate",
1135            GenerationMode::Transform => "Transform",
1136        }
1137    }
1138    fn tooltip_interrupt(self) -> &'static str {
1139        match self {
1140            GenerationMode::Generate => "Interrupt Generation",
1141            GenerationMode::Transform => "Interrupt Transform",
1142        }
1143    }
1144
1145    fn tooltip_restart(self) -> &'static str {
1146        match self {
1147            GenerationMode::Generate => "Restart Generation",
1148            GenerationMode::Transform => "Restart Transform",
1149        }
1150    }
1151
1152    fn tooltip_accept(self) -> &'static str {
1153        match self {
1154            GenerationMode::Generate => "Accept Generation",
1155            GenerationMode::Transform => "Accept Transform",
1156        }
1157    }
1158}