Inline assistant v2 layout (#22305)

Makes the inline assistant look like @danilo-leal's prototype.

Also fixes the bug with indent guides that @maxdeviant found!

<img width="1059" alt="Screenshot 2024-12-20 at 4 24 56 PM"
src="https://github.com/user-attachments/assets/5e55a3c2-768a-4e9d-bad5-d4ebbe9db9ce"
/>

Release Notes:

- N/A

crates/assistant2/src/assistant.rs

@@ -1,4 +1,5 @@
 mod active_thread;
+mod assistant_model_selector;
 mod assistant_panel;
 mod assistant_settings;
 mod buffer_codegen;

crates/assistant2/src/assistant_model_selector.rs

@@ -0,0 +1,85 @@
+use fs::Fs;
+use gpui::View;
+use language_model::LanguageModelRegistry;
+use language_model_selector::{LanguageModelSelector, LanguageModelSelectorPopoverMenu};
+use settings::update_settings_file;
+use std::sync::Arc;
+use ui::{prelude::*, ButtonLike, PopoverMenuHandle, Tooltip};
+
+use crate::{assistant_settings::AssistantSettings, ToggleModelSelector};
+
+pub struct AssistantModelSelector {
+    selector: View<LanguageModelSelector>,
+    menu_handle: PopoverMenuHandle<LanguageModelSelector>,
+}
+
+impl AssistantModelSelector {
+    pub(crate) fn new(
+        fs: Arc<dyn Fs>,
+        menu_handle: PopoverMenuHandle<LanguageModelSelector>,
+        cx: &mut WindowContext,
+    ) -> Self {
+        Self {
+            selector: cx.new_view(|cx| {
+                let fs = fs.clone();
+                LanguageModelSelector::new(
+                    move |model, cx| {
+                        update_settings_file::<AssistantSettings>(
+                            fs.clone(),
+                            cx,
+                            move |settings, _cx| settings.set_model(model.clone()),
+                        );
+                    },
+                    cx,
+                )
+            }),
+            menu_handle,
+        }
+    }
+}
+
+impl Render for AssistantModelSelector {
+    fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
+        let active_model = LanguageModelRegistry::read_global(cx).active_model();
+        let focus_handle = self.selector.focus_handle(cx).clone();
+
+        LanguageModelSelectorPopoverMenu::new(
+            self.selector.clone(),
+            ButtonLike::new("active-model")
+                .style(ButtonStyle::Subtle)
+                .child(
+                    h_flex()
+                        .w_full()
+                        .gap_0p5()
+                        .child(
+                            div()
+                                .overflow_x_hidden()
+                                .flex_grow()
+                                .whitespace_nowrap()
+                                .child(match active_model {
+                                    Some(model) => h_flex()
+                                        .child(
+                                            Label::new(model.name().0)
+                                                .size(LabelSize::Small)
+                                                .color(Color::Muted),
+                                        )
+                                        .into_any_element(),
+                                    _ => Label::new("No model selected")
+                                        .size(LabelSize::Small)
+                                        .color(Color::Muted)
+                                        .into_any_element(),
+                                }),
+                        )
+                        .child(
+                            Icon::new(IconName::ChevronDown)
+                                .color(Color::Muted)
+                                .size(IconSize::XSmall),
+                        ),
+                )
+                .tooltip(move |cx| {
+                    Tooltip::for_action_in("Change Model", &ToggleModelSelector, &focus_handle, cx)
+                }),
+        )
+        .with_handle(self.menu_handle.clone())
+    }
+}

crates/assistant2/src/inline_prompt_editor.rs

@@ -1,13 +1,12 @@
+use crate::assistant_model_selector::AssistantModelSelector;
 use crate::buffer_codegen::BufferCodegen;
 use crate::context_picker::ContextPicker;
 use crate::context_store::ContextStore;
 use crate::context_strip::ContextStrip;
 use crate::terminal_codegen::TerminalCodegen;
 use crate::thread_store::ThreadStore;
-use crate::ToggleContextPicker;
+use crate::{CycleNextInlineAssist, CyclePreviousInlineAssist};
-use crate::{
+use crate::{ToggleContextPicker, ToggleModelSelector};
-    assistant_settings::AssistantSettings, CycleNextInlineAssist, CyclePreviousInlineAssist,
-};
 use client::ErrorExt;
 use collections::VecDeque;
 use editor::{
@@ -22,9 +21,9 @@ use gpui::{
     WeakModel, WeakView, WindowContext,
 };
 use language_model::{LanguageModel, LanguageModelRegistry};
-use language_model_selector::{LanguageModelSelector, LanguageModelSelectorPopoverMenu};
+use language_model_selector::LanguageModelSelector;
 use parking_lot::Mutex;
-use settings::{update_settings_file, Settings};
+use settings::Settings;
 use std::cmp;
 use std::sync::Arc;
 use theme::ThemeSettings;
@@ -39,7 +38,8 @@ pub struct PromptEditor<T> {
     mode: PromptEditorMode,
     context_strip: View<ContextStrip>,
     context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
-    language_model_selector: View<LanguageModelSelector>,
+    model_selector: View<AssistantModelSelector>,
+    model_selector_menu_handle: PopoverMenuHandle<LanguageModelSelector>,
     edited_since_done: bool,
     prompt_history: VecDeque<String>,
     prompt_history_ix: Option<usize>,
@@ -72,23 +72,28 @@ impl<T: 'static> Render for PromptEditor<T> {
 
                 gutter_dimensions.full_width() + (gutter_dimensions.margin / 2.0)
             }
-            PromptEditorMode::Terminal { .. } => Pixels::ZERO,
+            PromptEditorMode::Terminal { .. } => {
+                // Give the equivalent of the same left-padding that we're using on the right
+                Pixels::from(24.0)
+            }
         };
 
         buttons.extend(self.render_buttons(cx));
 
         v_flex()
-            .border_y_1()
-            .border_color(cx.theme().status().info_border)
-            .size_full()
-            .py(cx.line_height() / 2.5)
-            .child(
-                h_flex()
             .key_context("PromptEditor")
             .bg(cx.theme().colors().editor_background)
             .block_mouse_down()
+            .border_y_1()
+            .border_color(cx.theme().status().info_border)
+            .size_full()
+            .pt_1()
+            .pb_2()
+            .child(
+                h_flex()
                     .cursor(CursorStyle::Arrow)
                     .on_action(cx.listener(Self::toggle_context_picker))
+                    .on_action(cx.listener(Self::toggle_model_selector))
                     .on_action(cx.listener(Self::confirm))
                     .on_action(cx.listener(Self::cancel))
                     .on_action(cx.listener(Self::move_up))
@@ -100,27 +105,7 @@ impl<T: 'static> Render for PromptEditor<T> {
                             .w(spacing)
                             .justify_center()
                             .gap_2()
-                            .child(LanguageModelSelectorPopoverMenu::new(
+                            .child(self.render_close_button(cx))
-                                self.language_model_selector.clone(),
-                                IconButton::new("context", IconName::SettingsAlt)
-                                    .shape(IconButtonShape::Square)
-                                    .icon_size(IconSize::Small)
-                                    .icon_color(Color::Muted)
-                                    .tooltip(move |cx| {
-                                        Tooltip::with_meta(
-                                            format!(
-                                                "Using {}",
-                                                LanguageModelRegistry::read_global(cx)
-                                                    .active_model()
-                                                    .map(|model| model.name().0)
-                                                    .unwrap_or_else(|| "No model selected".into()),
-                                            ),
-                                            None,
-                                            "Change Model",
-                                            cx,
-                                        )
-                                    }),
-                            ))
                             .map(|el| {
                                 let CodegenStatus::Error(error) = self.codegen_status(cx) else {
                                     return el;
@@ -172,13 +157,26 @@ impl<T: 'static> Render for PromptEditor<T> {
                                 }
                             }),
                     )
+                    .child(
+                        h_flex()
+                            .w_full()
+                            .justify_between()
                             .child(div().flex_1().child(self.render_editor(cx)))
                             .child(h_flex().gap_2().pr_6().children(buttons)),
+                    ),
             )
             .child(
                 h_flex()
-                    .child(h_flex().w(spacing).justify_center().gap_2())
+                    .child(h_flex().w(spacing).justify_between().gap_2())
-                    .child(self.context_strip.clone()),
+                    .child(
+                        h_flex()
+                            .w_full()
+                            .pl_1()
+                            .pr_6()
+                            .justify_between()
+                            .child(div().pl_1().child(self.context_strip.clone()))
+                            .child(self.model_selector.clone()),
+                    ),
             )
     }
 }
@@ -311,6 +309,10 @@ impl<T: 'static> PromptEditor<T> {
         self.context_picker_menu_handle.toggle(cx);
     }
 
+    fn toggle_model_selector(&mut self, _: &ToggleModelSelector, cx: &mut ViewContext<Self>) {
+        self.model_selector_menu_handle.toggle(cx);
+    }
+
     fn cancel(&mut self, _: &editor::actions::Cancel, cx: &mut ViewContext<Self>) {
         match self.codegen_status(cx) {
             CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {
@@ -400,32 +402,14 @@ impl<T: 'static> PromptEditor<T> {
 
         match codegen_status {
             CodegenStatus::Idle => {
-                vec![
+                vec![Button::new("start", mode.start_label())
-                    IconButton::new("cancel", IconName::Close)
-                        .icon_color(Color::Muted)
-                        .shape(IconButtonShape::Square)
-                        .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
-                        .on_click(
-                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
-                        )
-                        .into_any_element(),
-                    Button::new("start", mode.start_label())
                     .icon(IconName::Return)
+                    .label_size(LabelSize::Small)
                     .icon_color(Color::Muted)
-                        .on_click(
+                    .on_click(cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StartRequested)))
-                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StartRequested)),
+                    .into_any_element()]
-                        )
-                        .into_any_element(),
-                ]
             }
-            CodegenStatus::Pending => vec![
+            CodegenStatus::Pending => vec![IconButton::new("stop", IconName::Stop)
-                IconButton::new("cancel", IconName::Close)
-                    .icon_color(Color::Muted)
-                    .shape(IconButtonShape::Square)
-                    .tooltip(|cx| Tooltip::text("Cancel Assist", cx))
-                    .on_click(cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)))
-                    .into_any_element(),
-                IconButton::new("stop", IconName::Stop)
                 .icon_color(Color::Error)
                 .shape(IconButtonShape::Square)
                 .tooltip(move |cx| {
@@ -437,21 +421,11 @@ impl<T: 'static> PromptEditor<T> {
                     )
                 })
                 .on_click(cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StopRequested)))
-                    .into_any_element(),
+                .into_any_element()],
-            ],
             CodegenStatus::Done | CodegenStatus::Error(_) => {
-                let cancel = IconButton::new("cancel", IconName::Close)
-                    .icon_color(Color::Muted)
-                    .shape(IconButtonShape::Square)
-                    .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
-                    .on_click(cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)))
-                    .into_any_element();
-
                 let has_error = matches!(codegen_status, CodegenStatus::Error(_));
                 if has_error || self.edited_since_done {
-                    vec![
+                    vec![IconButton::new("restart", IconName::RotateCw)
-                        cancel,
-                        IconButton::new("restart", IconName::RotateCw)
                         .icon_color(Color::Info)
                         .shape(IconButtonShape::Square)
                         .tooltip(move |cx| {
@@ -465,8 +439,7 @@ impl<T: 'static> PromptEditor<T> {
                         .on_click(cx.listener(|_, _, cx| {
                             cx.emit(PromptEditorEvent::StartRequested);
                         }))
-                            .into_any_element(),
+                        .into_any_element()]
-                    ]
                 } else {
                     let accept = IconButton::new("accept", IconName::Check)
                         .icon_color(Color::Info)
@@ -482,7 +455,6 @@ impl<T: 'static> PromptEditor<T> {
                     match &self.mode {
                         PromptEditorMode::Terminal { .. } => vec![
                             accept,
-                            cancel,
                             IconButton::new("confirm", IconName::Play)
                                 .icon_color(Color::Info)
                                 .shape(IconButtonShape::Square)
@@ -498,7 +470,7 @@ impl<T: 'static> PromptEditor<T> {
                                 }))
                                 .into_any_element(),
                         ],
-                        PromptEditorMode::Buffer { .. } => vec![accept, cancel],
+                        PromptEditorMode::Buffer { .. } => vec![accept],
                     }
                 }
             }
@@ -527,6 +499,15 @@ impl<T: 'static> PromptEditor<T> {
         }
     }
 
+    fn render_close_button(&self, cx: &ViewContext<Self>) -> AnyElement {
+        IconButton::new("cancel", IconName::Close)
+            .icon_color(Color::Muted)
+            .shape(IconButtonShape::Square)
+            .tooltip(|cx| Tooltip::text("Close Assistant", cx))
+            .on_click(cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)))
+            .into_any_element()
+    }
+
     fn render_cycle_controls(&self, codegen: &BufferCodegen, cx: &ViewContext<Self>) -> AnyElement {
         let disabled = matches!(codegen.status(cx), CodegenStatus::Idle);
 
@@ -795,6 +776,7 @@ impl PromptEditor<BufferCodegen> {
             editor
         });
         let context_picker_menu_handle = PopoverMenuHandle::default();
+        let model_selector_menu_handle = PopoverMenuHandle::default();
 
         let mut this: PromptEditor<BufferCodegen> = PromptEditor {
             editor: prompt_editor.clone(),
@@ -809,19 +791,10 @@ impl PromptEditor<BufferCodegen> {
                 )
             }),
             context_picker_menu_handle,
-            language_model_selector: cx.new_view(|cx| {
+            model_selector: cx.new_view(|cx| {
-                let fs = fs.clone();
+                AssistantModelSelector::new(fs, model_selector_menu_handle.clone(), cx)
-                LanguageModelSelector::new(
-                    move |model, cx| {
-                        update_settings_file::<AssistantSettings>(
-                            fs.clone(),
-                            cx,
-                            move |settings, _| settings.set_model(model.clone()),
-                        );
-                    },
-                    cx,
-                )
             }),
+            model_selector_menu_handle,
             edited_since_done: false,
             prompt_history,
             prompt_history_ix: None,
@@ -942,6 +915,7 @@ impl PromptEditor<TerminalCodegen> {
             editor
         });
         let context_picker_menu_handle = PopoverMenuHandle::default();
+        let model_selector_menu_handle = PopoverMenuHandle::default();
 
         let mut this = Self {
             editor: prompt_editor.clone(),
@@ -956,19 +930,10 @@ impl PromptEditor<TerminalCodegen> {
                 )
             }),
             context_picker_menu_handle,
-            language_model_selector: cx.new_view(|cx| {
+            model_selector: cx.new_view(|cx| {
-                let fs = fs.clone();
+                AssistantModelSelector::new(fs, model_selector_menu_handle.clone(), cx)
-                LanguageModelSelector::new(
-                    move |model, cx| {
-                        update_settings_file::<AssistantSettings>(
-                            fs.clone(),
-                            cx,
-                            move |settings, _| settings.set_model(model.clone()),
-                        );
-                    },
-                    cx,
-                )
             }),
+            model_selector_menu_handle,
             edited_since_done: false,
             prompt_history,
             prompt_history_ix: None,

crates/assistant2/src/message_editor.rs

@@ -7,17 +7,17 @@ use gpui::{
     WeakView,
 };
 use language_model::{LanguageModelRegistry, LanguageModelRequestTool};
-use language_model_selector::{LanguageModelSelector, LanguageModelSelectorPopoverMenu};
+use language_model_selector::LanguageModelSelector;
 use rope::Point;
-use settings::{update_settings_file, Settings};
+use settings::Settings;
 use theme::ThemeSettings;
 use ui::{
     prelude::*, ButtonLike, CheckboxWithLabel, ElevationIndex, KeyBinding, PopoverMenu,
-    PopoverMenuHandle, Tooltip,
+    PopoverMenuHandle,
 };
 use workspace::Workspace;
 
-use crate::assistant_settings::AssistantSettings;
+use crate::assistant_model_selector::AssistantModelSelector;
 use crate::context_picker::{ConfirmBehavior, ContextPicker};
 use crate::context_store::ContextStore;
 use crate::context_strip::ContextStrip;
@@ -33,8 +33,8 @@ pub struct MessageEditor {
     context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
     inline_context_picker: View<ContextPicker>,
     inline_context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
-    language_model_selector: View<LanguageModelSelector>,
+    model_selector: View<AssistantModelSelector>,
-    language_model_selector_menu_handle: PopoverMenuHandle<LanguageModelSelector>,
+    model_selector_menu_handle: PopoverMenuHandle<LanguageModelSelector>,
     use_tools: bool,
     _subscriptions: Vec<Subscription>,
 }
@@ -50,6 +50,7 @@ impl MessageEditor {
         let context_store = cx.new_model(|_cx| ContextStore::new());
         let context_picker_menu_handle = PopoverMenuHandle::default();
         let inline_context_picker_menu_handle = PopoverMenuHandle::default();
+        let model_selector_menu_handle = PopoverMenuHandle::default();
 
         let editor = cx.new_view(|cx| {
             let mut editor = Editor::auto_height(10, cx);
@@ -92,27 +93,17 @@ impl MessageEditor {
             context_picker_menu_handle,
             inline_context_picker,
             inline_context_picker_menu_handle,
-            language_model_selector: cx.new_view(|cx| {
+            model_selector: cx.new_view(|cx| {
-                let fs = fs.clone();
+                AssistantModelSelector::new(fs, model_selector_menu_handle.clone(), cx)
-                LanguageModelSelector::new(
-                    move |model, cx| {
-                        update_settings_file::<AssistantSettings>(
-                            fs.clone(),
-                            cx,
-                            move |settings, _cx| settings.set_model(model.clone()),
-                        );
-                    },
-                    cx,
-                )
             }),
-            language_model_selector_menu_handle: PopoverMenuHandle::default(),
+            model_selector_menu_handle,
             use_tools: false,
             _subscriptions: subscriptions,
         }
     }
 
     fn toggle_model_selector(&mut self, _: &ToggleModelSelector, cx: &mut ViewContext<Self>) {
-        self.language_model_selector_menu_handle.toggle(cx);
+        self.model_selector_menu_handle.toggle(cx)
     }
 
     fn toggle_context_picker(&mut self, _: &ToggleContextPicker, cx: &mut ViewContext<Self>) {
@@ -203,50 +194,6 @@ impl MessageEditor {
         let editor_focus_handle = self.editor.focus_handle(cx);
         cx.focus(&editor_focus_handle);
     }
-
-    fn render_language_model_selector(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
-        let active_model = LanguageModelRegistry::read_global(cx).active_model();
-        let focus_handle = self.language_model_selector.focus_handle(cx).clone();
-
-        LanguageModelSelectorPopoverMenu::new(
-            self.language_model_selector.clone(),
-            ButtonLike::new("active-model")
-                .style(ButtonStyle::Subtle)
-                .child(
-                    h_flex()
-                        .w_full()
-                        .gap_0p5()
-                        .child(
-                            div()
-                                .overflow_x_hidden()
-                                .flex_grow()
-                                .whitespace_nowrap()
-                                .child(match active_model {
-                                    Some(model) => h_flex()
-                                        .child(
-                                            Label::new(model.name().0)
-                                                .size(LabelSize::Small)
-                                                .color(Color::Muted),
-                                        )
-                                        .into_any_element(),
-                                    _ => Label::new("No model selected")
-                                        .size(LabelSize::Small)
-                                        .color(Color::Muted)
-                                        .into_any_element(),
-                                }),
-                        )
-                        .child(
-                            Icon::new(IconName::ChevronDown)
-                                .color(Color::Muted)
-                                .size(IconSize::XSmall),
-                        ),
-                )
-                .tooltip(move |cx| {
-                    Tooltip::for_action_in("Change Model", &ToggleModelSelector, &focus_handle, cx)
-                }),
-        )
-        .with_handle(self.language_model_selector_menu_handle.clone())
-    }
 }
 
 impl FocusableView for MessageEditor {
@@ -321,10 +268,7 @@ impl Render for MessageEditor {
                         }),
                     ))
                     .child(
-                        h_flex()
+                        h_flex().gap_1().child(self.model_selector.clone()).child(
-                            .gap_1()
-                            .child(self.render_language_model_selector(cx))
-                            .child(
                             ButtonLike::new("chat")
                                 .style(ButtonStyle::Filled)
                                 .layer(ElevationIndex::ModalSurface)