inline_prompt_editor.rs

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