Show custom header for assistant messages

crates/ai/src/assistant.rs

@@ -1,13 +1,15 @@
 use crate::{OpenAIRequest, OpenAIResponseStreamEvent, RequestMessage, Role};
 use anyhow::{anyhow, Result};
-use editor::{Editor, MultiBuffer};
+use collections::HashMap;
+use editor::{Editor, ExcerptId, ExcerptRange, MultiBuffer};
 use futures::{io::BufReader, AsyncBufReadExt, AsyncReadExt, Stream, StreamExt};
 use gpui::{
     actions, elements::*, executor::Background, Action, AppContext, AsyncAppContext, Entity,
-    ModelHandle, Subscription, Task, View, ViewContext, ViewHandle, WeakViewHandle, WindowContext,
+    ModelContext, ModelHandle, Subscription, Task, View, ViewContext, ViewHandle, WeakViewHandle,
+    WindowContext,
 };
 use isahc::{http::StatusCode, Request, RequestExt};
-use language::{language_settings::SoftWrap, Anchor, Buffer, Language, LanguageRegistry};
+use language::{language_settings::SoftWrap, Buffer, Language, LanguageRegistry};
 use std::{io, sync::Arc};
 use util::{post_inc, ResultExt, TryFutureExt};
 use workspace::{
@@ -19,8 +21,8 @@ use workspace::{
 actions!(assistant, [NewContext, Assist, CancelLastAssist]);
 
 pub fn init(cx: &mut AppContext) {
-    cx.add_action(Assistant::assist);
-    cx.capture_action(Assistant::cancel_last_assist);
+    cx.add_action(AssistantEditor::assist);
+    cx.capture_action(AssistantEditor::cancel_last_assist);
 }
 
 pub enum AssistantPanelEvent {
@@ -188,7 +190,7 @@ impl Panel for AssistantPanel {
                     .await?;
                 workspace.update(&mut cx, |workspace, cx| {
                     let editor = Box::new(cx.add_view(|cx| {
-                        Assistant::new(markdown, workspace.app_state().languages.clone(), cx)
+                        AssistantEditor::new(markdown, workspace.app_state().languages.clone(), cx)
                     }));
                     Pane::add_item(workspace, &pane, editor, true, focus, None, cx);
                 })?;
@@ -230,38 +232,31 @@ impl Panel for AssistantPanel {
 }
 
 struct Assistant {
+    buffer: ModelHandle<MultiBuffer>,
     messages: Vec<Message>,
-    editor: ViewHandle<Editor>,
+    messages_by_id: HashMap<ExcerptId, Message>,
     completion_count: usize,
     pending_completions: Vec<PendingCompletion>,
     markdown: Arc<Language>,
     language_registry: Arc<LanguageRegistry>,
 }
 
-struct PendingCompletion {
-    id: usize,
-    _task: Task<Option<()>>,
+impl Entity for Assistant {
+    type Event = ();
 }
 
 impl Assistant {
     fn new(
         markdown: Arc<Language>,
         language_registry: Arc<LanguageRegistry>,
-        cx: &mut ViewContext<Self>,
+        cx: &mut ModelContext<Self>,
     ) -> Self {
-        let editor = cx.add_view(|cx| {
-            let multibuffer = cx.add_model(|_| MultiBuffer::new(0));
-            let mut editor = Editor::for_multibuffer(multibuffer, None, cx);
-            editor.set_soft_wrap_mode(SoftWrap::EditorWidth, cx);
-            editor.set_show_gutter(false, cx);
-            editor
-        });
-
         let mut this = Self {
+            buffer: cx.add_model(|_| MultiBuffer::new(0)),
             messages: Default::default(),
-            editor,
-            completion_count: 0,
-            pending_completions: Vec::new(),
+            messages_by_id: Default::default(),
+            completion_count: Default::default(),
+            pending_completions: Default::default(),
             markdown,
             language_registry,
         };
@@ -269,7 +264,7 @@ impl Assistant {
         this
     }
 
-    fn assist(&mut self, _: &Assist, cx: &mut ViewContext<Self>) {
+    fn assist(&mut self, cx: &mut ModelContext<Self>) {
         let messages = self
             .messages
             .iter()
@@ -285,8 +280,8 @@ impl Assistant {
         };
 
         if let Some(api_key) = std::env::var("OPENAI_API_KEY").log_err() {
-            let stream = stream_completion(api_key, cx.background_executor().clone(), request);
-            let response_buffer = self.push_message(Role::Assistant, cx);
+            let stream = stream_completion(api_key, cx.background().clone(), request);
+            let response = self.push_message(Role::Assistant, cx);
             self.push_message(Role::User, cx);
             let task = cx.spawn(|this, mut cx| {
                 async move {
@@ -295,7 +290,7 @@ impl Assistant {
                     while let Some(message) = messages.next().await {
                         let mut message = message?;
                         if let Some(choice) = message.choices.pop() {
-                            response_buffer.update(&mut cx, |content, cx| {
+                            response.content.update(&mut cx, |content, cx| {
                                 let text: Arc<str> = choice.delta.content?.into();
                                 content.edit([(content.len()..content.len(), text)], None, cx);
                                 Some(())
@@ -306,8 +301,7 @@ impl Assistant {
                     this.update(&mut cx, |this, _| {
                         this.pending_completions
                             .retain(|completion| completion.id != this.completion_count);
-                    })
-                    .ok();
+                    });
 
                     anyhow::Ok(())
                 }
@@ -321,45 +315,123 @@ impl Assistant {
         }
     }
 
-    fn cancel_last_assist(&mut self, _: &editor::Cancel, cx: &mut ViewContext<Self>) {
-        if self.pending_completions.pop().is_none() {
-            cx.propagate_action();
-        }
+    fn cancel_last_assist(&mut self) -> bool {
+        self.pending_completions.pop().is_some()
     }
 
-    fn push_message(&mut self, role: Role, cx: &mut ViewContext<Self>) -> ModelHandle<Buffer> {
+    fn push_message(&mut self, role: Role, cx: &mut ModelContext<Self>) -> Message {
         let content = cx.add_model(|cx| {
             let mut buffer = Buffer::new(0, "", cx);
             buffer.set_language(Some(self.markdown.clone()), cx);
             buffer.set_language_registry(self.language_registry.clone());
             buffer
         });
+        let excerpt_id = self.buffer.update(cx, |buffer, cx| {
+            buffer
+                .push_excerpts(
+                    content.clone(),
+                    vec![ExcerptRange {
+                        context: 0..0,
+                        primary: None,
+                    }],
+                    cx,
+                )
+                .pop()
+                .unwrap()
+        });
+
         let message = Message {
             role,
             content: content.clone(),
         };
-        self.messages.push(message);
-
-        self.editor.update(cx, |editor, cx| {
-            editor.buffer().update(cx, |buffer, cx| {
-                buffer.push_excerpts_with_context_lines(
-                    content.clone(),
-                    vec![Anchor::MIN..Anchor::MAX],
-                    0,
-                    cx,
-                )
-            });
-        });
-
-        content
+        self.messages.push(message.clone());
+        self.messages_by_id.insert(excerpt_id, message.clone());
+        message
     }
 }
 
-impl Entity for Assistant {
+struct PendingCompletion {
+    id: usize,
+    _task: Task<Option<()>>,
+}
+
+struct AssistantEditor {
+    assistant: ModelHandle<Assistant>,
+    editor: ViewHandle<Editor>,
+}
+
+impl AssistantEditor {
+    fn new(
+        markdown: Arc<Language>,
+        language_registry: Arc<LanguageRegistry>,
+        cx: &mut ViewContext<Self>,
+    ) -> Self {
+        let assistant = cx.add_model(|cx| Assistant::new(markdown, language_registry, cx));
+        let editor = cx.add_view(|cx| {
+            let mut editor = Editor::for_multibuffer(assistant.read(cx).buffer.clone(), None, cx);
+            editor.set_soft_wrap_mode(SoftWrap::EditorWidth, cx);
+            editor.set_show_gutter(false, cx);
+            editor.set_render_excerpt_header(
+                {
+                    let assistant = assistant.clone();
+                    move |editor, params: editor::RenderExcerptHeaderParams, cx| {
+                        let style = &theme::current(cx).assistant;
+                        if let Some(message) = assistant.read(cx).messages_by_id.get(&params.id) {
+                            let sender = match message.role {
+                                Role::User => Label::new("You", style.user_sender.text.clone())
+                                    .contained()
+                                    .with_style(style.user_sender.container),
+                                Role::Assistant => {
+                                    Label::new("Assistant", style.assistant_sender.text.clone())
+                                        .contained()
+                                        .with_style(style.assistant_sender.container)
+                                }
+                                Role::System => {
+                                    Label::new("System", style.assistant_sender.text.clone())
+                                        .contained()
+                                        .with_style(style.assistant_sender.container)
+                                }
+                            };
+
+                            Flex::row()
+                                .with_child(sender)
+                                .aligned()
+                                .left()
+                                .contained()
+                                .with_style(style.header)
+                                .into_any()
+                        } else {
+                            Empty::new().into_any()
+                        }
+                    }
+                },
+                cx,
+            );
+            editor
+        });
+        Self { assistant, editor }
+    }
+
+    fn assist(&mut self, _: &Assist, cx: &mut ViewContext<Self>) {
+        self.assistant
+            .update(cx, |assistant, cx| assistant.assist(cx));
+    }
+
+    fn cancel_last_assist(&mut self, _: &editor::Cancel, cx: &mut ViewContext<Self>) {
+        if !self
+            .assistant
+            .update(cx, |assistant, _| assistant.cancel_last_assist())
+        {
+            cx.propagate_action();
+        }
+    }
+}
+
+impl Entity for AssistantEditor {
     type Event = ();
 }
 
-impl View for Assistant {
+impl View for AssistantEditor {
     fn ui_name() -> &'static str {
         "ContextEditor"
     }
@@ -374,7 +446,7 @@ impl View for Assistant {
     }
 }
 
-impl Item for Assistant {
+impl Item for AssistantEditor {
     fn tab_content<V: View>(
         &self,
         _: Option<usize>,
@@ -385,6 +457,7 @@ impl Item for Assistant {
     }
 }
 
+#[derive(Clone)]
 struct Message {
     role: Role,
     content: ModelHandle<Buffer>,

crates/editor/src/editor.rs

@@ -46,7 +46,8 @@ use gpui::{
     platform::{CursorStyle, MouseButton},
     serde_json::{self, json},
     AnyElement, AnyViewHandle, AppContext, AsyncAppContext, ClipboardItem, Element, Entity,
-    ModelHandle, Subscription, Task, View, ViewContext, ViewHandle, WeakViewHandle, WindowContext,
+    LayoutContext, ModelHandle, Subscription, Task, View, ViewContext, ViewHandle, WeakViewHandle,
+    WindowContext,
 };
 use highlight_matching_bracket::refresh_matching_bracket_highlights;
 use hover_popover::{hide_hover, HoverState};
@@ -498,6 +499,7 @@ pub struct Editor {
     mode: EditorMode,
     show_gutter: bool,
     placeholder_text: Option<Arc<str>>,
+    render_excerpt_header: Option<element::RenderExcerptHeader>,
     highlighted_rows: Option<Range<u32>>,
     #[allow(clippy::type_complexity)]
     background_highlights: BTreeMap<TypeId, (fn(&Theme) -> Color, Vec<Range<Anchor>>)>,
@@ -1301,6 +1303,7 @@ impl Editor {
             mode,
             show_gutter: mode == EditorMode::Full,
             placeholder_text: None,
+            render_excerpt_header: None,
             highlighted_rows: None,
             background_highlights: Default::default(),
             nav_history: None,
@@ -6663,6 +6666,20 @@ impl Editor {
         cx.notify();
     }
 
+    pub fn set_render_excerpt_header(
+        &mut self,
+        render_excerpt_header: impl 'static
+            + Fn(
+                &mut Editor,
+                RenderExcerptHeaderParams,
+                &mut LayoutContext<Editor>,
+            ) -> AnyElement<Editor>,
+        cx: &mut ViewContext<Self>,
+    ) {
+        self.render_excerpt_header = Some(Arc::new(render_excerpt_header));
+        cx.notify();
+    }
+
     pub fn reveal_in_finder(&mut self, _: &RevealInFinder, cx: &mut ViewContext<Self>) {
         if let Some(buffer) = self.buffer().read(cx).as_singleton() {
             if let Some(file) = buffer.read(cx).file().and_then(|f| f.as_local()) {
@@ -7308,8 +7325,12 @@ impl View for Editor {
             });
         }
 
+        let mut editor = EditorElement::new(style.clone());
+        if let Some(render_excerpt_header) = self.render_excerpt_header.clone() {
+            editor = editor.with_render_excerpt_header(render_excerpt_header);
+        }
         Stack::new()
-            .with_child(EditorElement::new(style.clone()))
+            .with_child(editor)
             .with_child(ChildView::new(&self.mouse_context_menu, cx))
             .into_any()
     }

crates/editor/src/element.rs

@@ -91,18 +91,41 @@ impl SelectionLayout {
     }
 }
 
-#[derive(Clone)]
+pub struct RenderExcerptHeaderParams<'a> {
+    pub id: crate::ExcerptId,
+    pub buffer: &'a language::BufferSnapshot,
+    pub range: &'a crate::ExcerptRange<text::Anchor>,
+    pub starts_new_buffer: bool,
+    pub gutter_padding: f32,
+    pub editor_style: &'a EditorStyle,
+}
+
+pub type RenderExcerptHeader = Arc<
+    dyn Fn(
+        &mut Editor,
+        RenderExcerptHeaderParams,
+        &mut LayoutContext<Editor>,
+    ) -> AnyElement<Editor>,
+>;
+
 pub struct EditorElement {
     style: Arc<EditorStyle>,
+    render_excerpt_header: RenderExcerptHeader,
 }
 
 impl EditorElement {
     pub fn new(style: EditorStyle) -> Self {
         Self {
             style: Arc::new(style),
+            render_excerpt_header: Arc::new(render_excerpt_header),
         }
     }
 
+    pub fn with_render_excerpt_header(mut self, render: RenderExcerptHeader) -> Self {
+        self.render_excerpt_header = render;
+        self
+    }
+
     fn attach_mouse_handlers(
         scene: &mut SceneBuilder,
         position_map: &Arc<PositionMap>,
@@ -1465,11 +1488,9 @@ impl EditorElement {
         line_height: f32,
         style: &EditorStyle,
         line_layouts: &[LineWithInvisibles],
-        include_root: bool,
         editor: &mut Editor,
         cx: &mut LayoutContext<Editor>,
     ) -> (f32, Vec<BlockLayout>) {
-        let tooltip_style = theme::current(cx).tooltip.clone();
         let scroll_x = snapshot.scroll_anchor.offset.x();
         let (fixed_blocks, non_fixed_blocks) = snapshot
             .blocks_in_range(rows.clone())
@@ -1510,112 +1531,18 @@ impl EditorElement {
                     range,
                     starts_new_buffer,
                     ..
-                } => {
-                    let id = *id;
-                    let jump_icon = project::File::from_dyn(buffer.file()).map(|file| {
-                        let jump_path = ProjectPath {
-                            worktree_id: file.worktree_id(cx),
-                            path: file.path.clone(),
-                        };
-                        let jump_anchor = range
-                            .primary
-                            .as_ref()
-                            .map_or(range.context.start, |primary| primary.start);
-                        let jump_position = language::ToPoint::to_point(&jump_anchor, buffer);
-
-                        enum JumpIcon {}
-                        MouseEventHandler::<JumpIcon, _>::new(id.into(), cx, |state, _| {
-                            let style = style.jump_icon.style_for(state, false);
-                            Svg::new("icons/arrow_up_right_8.svg")
-                                .with_color(style.color)
-                                .constrained()
-                                .with_width(style.icon_width)
-                                .aligned()
-                                .contained()
-                                .with_style(style.container)
-                                .constrained()
-                                .with_width(style.button_width)
-                                .with_height(style.button_width)
-                        })
-                        .with_cursor_style(CursorStyle::PointingHand)
-                        .on_click(MouseButton::Left, move |_, editor, cx| {
-                            if let Some(workspace) = editor
-                                .workspace
-                                .as_ref()
-                                .and_then(|(workspace, _)| workspace.upgrade(cx))
-                            {
-                                workspace.update(cx, |workspace, cx| {
-                                    Editor::jump(
-                                        workspace,
-                                        jump_path.clone(),
-                                        jump_position,
-                                        jump_anchor,
-                                        cx,
-                                    );
-                                });
-                            }
-                        })
-                        .with_tooltip::<JumpIcon>(
-                            id.into(),
-                            "Jump to Buffer".to_string(),
-                            Some(Box::new(crate::OpenExcerpts)),
-                            tooltip_style.clone(),
-                            cx,
-                        )
-                        .aligned()
-                        .flex_float()
-                    });
-
-                    if *starts_new_buffer {
-                        let style = &self.style.diagnostic_path_header;
-                        let font_size =
-                            (style.text_scale_factor * self.style.text.font_size).round();
-
-                        let path = buffer.resolve_file_path(cx, include_root);
-                        let mut filename = None;
-                        let mut parent_path = None;
-                        // Can't use .and_then() because `.file_name()` and `.parent()` return references :(
-                        if let Some(path) = path {
-                            filename = path.file_name().map(|f| f.to_string_lossy().to_string());
-                            parent_path =
-                                path.parent().map(|p| p.to_string_lossy().to_string() + "/");
-                        }
-
-                        Flex::row()
-                            .with_child(
-                                Label::new(
-                                    filename.unwrap_or_else(|| "untitled".to_string()),
-                                    style.filename.text.clone().with_font_size(font_size),
-                                )
-                                .contained()
-                                .with_style(style.filename.container)
-                                .aligned(),
-                            )
-                            .with_children(parent_path.map(|path| {
-                                Label::new(path, style.path.text.clone().with_font_size(font_size))
-                                    .contained()
-                                    .with_style(style.path.container)
-                                    .aligned()
-                            }))
-                            .with_children(jump_icon)
-                            .contained()
-                            .with_style(style.container)
-                            .with_padding_left(gutter_padding)
-                            .with_padding_right(gutter_padding)
-                            .expanded()
-                            .into_any_named("path header block")
-                    } else {
-                        let text_style = self.style.text.clone();
-                        Flex::row()
-                            .with_child(Label::new("⋯", text_style))
-                            .with_children(jump_icon)
-                            .contained()
-                            .with_padding_left(gutter_padding)
-                            .with_padding_right(gutter_padding)
-                            .expanded()
-                            .into_any_named("collapsed context")
-                    }
-                }
+                } => (self.render_excerpt_header)(
+                    editor,
+                    RenderExcerptHeaderParams {
+                        id: *id,
+                        buffer,
+                        range,
+                        starts_new_buffer: *starts_new_buffer,
+                        gutter_padding,
+                        editor_style: style,
+                    },
+                    cx,
+                ),
             };
 
             element.layout(
@@ -2080,12 +2007,6 @@ impl Element<Editor> for EditorElement {
             ShowScrollbar::Never => false,
         };
 
-        let include_root = editor
-            .project
-            .as_ref()
-            .map(|project| project.read(cx).visible_worktrees(cx).count() > 1)
-            .unwrap_or_default();
-
         let fold_ranges: Vec<(BufferRow, Range<DisplayPoint>, Color)> = fold_ranges
             .into_iter()
             .map(|(id, fold)| {
@@ -2144,7 +2065,6 @@ impl Element<Editor> for EditorElement {
             line_height,
             &style,
             &line_layouts,
-            include_root,
             editor,
             cx,
         );
@@ -2759,6 +2679,121 @@ impl HighlightedRange {
     }
 }
 
+fn render_excerpt_header(
+    editor: &mut Editor,
+    RenderExcerptHeaderParams {
+        id,
+        buffer,
+        range,
+        starts_new_buffer,
+        gutter_padding,
+        editor_style,
+    }: RenderExcerptHeaderParams,
+    cx: &mut LayoutContext<Editor>,
+) -> AnyElement<Editor> {
+    let tooltip_style = theme::current(cx).tooltip.clone();
+    let include_root = editor
+        .project
+        .as_ref()
+        .map(|project| project.read(cx).visible_worktrees(cx).count() > 1)
+        .unwrap_or_default();
+    let jump_icon = project::File::from_dyn(buffer.file()).map(|file| {
+        let jump_path = ProjectPath {
+            worktree_id: file.worktree_id(cx),
+            path: file.path.clone(),
+        };
+        let jump_anchor = range
+            .primary
+            .as_ref()
+            .map_or(range.context.start, |primary| primary.start);
+        let jump_position = language::ToPoint::to_point(&jump_anchor, buffer);
+
+        enum JumpIcon {}
+        MouseEventHandler::<JumpIcon, _>::new(id.into(), cx, |state, _| {
+            let style = editor_style.jump_icon.style_for(state, false);
+            Svg::new("icons/arrow_up_right_8.svg")
+                .with_color(style.color)
+                .constrained()
+                .with_width(style.icon_width)
+                .aligned()
+                .contained()
+                .with_style(style.container)
+                .constrained()
+                .with_width(style.button_width)
+                .with_height(style.button_width)
+        })
+        .with_cursor_style(CursorStyle::PointingHand)
+        .on_click(MouseButton::Left, move |_, editor, cx| {
+            if let Some(workspace) = editor
+                .workspace
+                .as_ref()
+                .and_then(|(workspace, _)| workspace.upgrade(cx))
+            {
+                workspace.update(cx, |workspace, cx| {
+                    Editor::jump(workspace, jump_path.clone(), jump_position, jump_anchor, cx);
+                });
+            }
+        })
+        .with_tooltip::<JumpIcon>(
+            id.into(),
+            "Jump to Buffer".to_string(),
+            Some(Box::new(crate::OpenExcerpts)),
+            tooltip_style.clone(),
+            cx,
+        )
+        .aligned()
+        .flex_float()
+    });
+
+    if starts_new_buffer {
+        let style = &editor_style.diagnostic_path_header;
+        let font_size = (style.text_scale_factor * editor_style.text.font_size).round();
+
+        let path = buffer.resolve_file_path(cx, include_root);
+        let mut filename = None;
+        let mut parent_path = None;
+        // Can't use .and_then() because `.file_name()` and `.parent()` return references :(
+        if let Some(path) = path {
+            filename = path.file_name().map(|f| f.to_string_lossy().to_string());
+            parent_path = path.parent().map(|p| p.to_string_lossy().to_string() + "/");
+        }
+
+        Flex::row()
+            .with_child(
+                Label::new(
+                    filename.unwrap_or_else(|| "untitled".to_string()),
+                    style.filename.text.clone().with_font_size(font_size),
+                )
+                .contained()
+                .with_style(style.filename.container)
+                .aligned(),
+            )
+            .with_children(parent_path.map(|path| {
+                Label::new(path, style.path.text.clone().with_font_size(font_size))
+                    .contained()
+                    .with_style(style.path.container)
+                    .aligned()
+            }))
+            .with_children(jump_icon)
+            .contained()
+            .with_style(style.container)
+            .with_padding_left(gutter_padding)
+            .with_padding_right(gutter_padding)
+            .expanded()
+            .into_any_named("path header block")
+    } else {
+        let text_style = editor_style.text.clone();
+        Flex::row()
+            .with_child(Label::new("⋯", text_style))
+            .with_children(jump_icon)
+            .contained()
+            .with_padding_left(gutter_padding)
+            .with_padding_right(gutter_padding)
+            .expanded()
+            .into_any_named("collapsed context")
+    }
+}
+
 fn position_to_display_point(
     position: Vector2F,
     text_bounds: RectF,

crates/theme/src/theme.rs

@@ -971,6 +971,9 @@ pub struct TerminalStyle {
 #[derive(Clone, Deserialize, Default)]
 pub struct AssistantStyle {
     pub container: ContainerStyle,
+    pub header: ContainerStyle,
+    pub user_sender: ContainedText,
+    pub assistant_sender: ContainedText,
 }
 
 #[derive(Clone, Deserialize, Default)]

styles/src/styleTree/assistant.ts

@@ -1,13 +1,23 @@
 import { ColorScheme } from "../themes/common/colorScheme"
+import { text, border } from "./components"
 import editor from "./editor"
 
 export default function assistant(colorScheme: ColorScheme) {
+    const layer = colorScheme.highest;
     return {
       container: {
         background: editor(colorScheme).background,
-        padding: {
-          left: 10,
-        }
+        padding: { left: 12 }
+      },
+      header: {
+        border: border(layer, "default", { bottom: true, top: true }),
+        margin: { bottom: 6, top: 6 }
+      },
+      user_sender: {
+        ...text(layer, "sans", "default", { size: "sm", weight: "bold" }),
+      },
+      assistant_sender: {
+        ...text(layer, "sans", "accent", { size: "sm", weight: "bold" }),
       }
     }
 }