assistant2: Use ChatMessage component to render chat messages (#11193)

This PR updates the new assistant panel to use the `ChatMessage`
component to render its chat messages.

This also lays the foundation for collapsing the messages, though that
has yet to be wired up.

Adapted from the work on the `assistant-chat-ui` branch.

Release Notes:

- N/A

crates/assistant2/src/assistant2.rs

@@ -22,7 +22,6 @@ use semantic_index::{CloudEmbeddingProvider, SemanticIndex};
 use serde::Deserialize;
 use settings::Settings;
 use std::sync::Arc;
-use theme::ThemeSettings;
 use tools::ProjectIndexTool;
 use ui::Composer;
 use util::{paths::EMBEDDINGS_DIR, ResultExt};
@@ -33,7 +32,7 @@ use workspace::{
 
 pub use assistant_settings::AssistantSettings;
 
-use crate::ui::{ChatMessageHeader, UserOrAssistant};
+use crate::ui::UserOrAssistant;
 
 const MAX_COMPLETION_CALLS_PER_SUBMISSION: usize = 5;
 
@@ -526,19 +525,15 @@ impl AssistantChat {
         let is_last = ix == self.messages.len() - 1;
 
         match &self.messages[ix] {
-            ChatMessage::User(UserMessage { body, .. }) => div()
+            ChatMessage::User(UserMessage { id, body }) => div()
                 .when(!is_last, |element| element.mb_2())
-                .child(ChatMessageHeader::new(UserOrAssistant::User(
+                .child(crate::ui::ChatMessage::new(
-                    self.user_store.read(cx).current_user(),
+                    *id,
-                )))
+                    UserOrAssistant::User(self.user_store.read(cx).current_user()),
-                .child(
+                    body.clone().into_any_element(),
-                    div()
+                    false,
-                        .p_2()
+                    Box::new(|_, _| {}),
-                        .text_color(cx.theme().colors().editor_foreground)
+                ))
-                        .font(ThemeSettings::get_global(cx).buffer_font.clone())
-                        .bg(cx.theme().colors().editor_background)
-                        .child(body.clone()),
-                )
                 .into_any(),
             ChatMessage::Assistant(AssistantMessage {
                 id,
@@ -555,8 +550,14 @@ impl AssistantChat {
 
                 div()
                     .when(!is_last, |element| element.mb_2())
-                    .child(ChatMessageHeader::new(UserOrAssistant::Assistant))
+                    .child(crate::ui::ChatMessage::new(
-                    .child(assistant_body)
+                        *id,
+                        UserOrAssistant::Assistant,
+                        assistant_body.into_any_element(),
+                        false,
+                        Box::new(|_, _| {}),
+                    ))
+                    // TODO: Should the errors and tool calls get passed into `ChatMessage`?
                     .child(self.render_error(error.clone(), ix, cx))
                     .children(tool_calls.iter().map(|tool_call| {
                         let result = &tool_call.result;

crates/assistant2/src/ui.rs

@@ -1,5 +1,5 @@
-mod chat_message_header;
+mod chat_message;
 mod composer;
 
-pub use chat_message_header::*;
+pub use chat_message::*;
 pub use composer::*;

crates/assistant2/src/ui/chat_message.rs

@@ -0,0 +1,131 @@
+use std::sync::Arc;
+
+use client::User;
+use gpui::AnyElement;
+use ui::{prelude::*, Avatar};
+
+use crate::MessageId;
+
+pub enum UserOrAssistant {
+    User(Option<Arc<User>>),
+    Assistant,
+}
+
+#[derive(IntoElement)]
+pub struct ChatMessage {
+    id: MessageId,
+    player: UserOrAssistant,
+    message: AnyElement,
+    collapsed: bool,
+    on_collapse: Box<dyn Fn(bool, &mut WindowContext) + 'static>,
+}
+
+impl ChatMessage {
+    pub fn new(
+        id: MessageId,
+        player: UserOrAssistant,
+        message: AnyElement,
+        collapsed: bool,
+        on_collapse: Box<dyn Fn(bool, &mut WindowContext) + 'static>,
+    ) -> Self {
+        Self {
+            id,
+            player,
+            message,
+            collapsed,
+            on_collapse,
+        }
+    }
+}
+
+impl RenderOnce for ChatMessage {
+    fn render(self, cx: &mut WindowContext) -> impl IntoElement {
+        // TODO: This should be top padding + 1.5x line height
+        // Set the message height to cut off at exactly 1.5 lines when collapsed
+        let collapsed_height = rems(2.875);
+
+        let collapse_handle_id = SharedString::from(format!("{}_collapse_handle", self.id.0));
+        let collapse_handle = h_flex()
+            .id(collapse_handle_id.clone())
+            .group(collapse_handle_id.clone())
+            .flex_none()
+            .justify_center()
+            .w_1()
+            .mx_2()
+            .h_full()
+            .on_click(move |_event, cx| (self.on_collapse)(!self.collapsed, cx))
+            .child(
+                div()
+                    .w_px()
+                    .h_full()
+                    .rounded_lg()
+                    .overflow_hidden()
+                    .bg(cx.theme().colors().element_background)
+                    .group_hover(collapse_handle_id, |this| {
+                        this.bg(cx.theme().colors().element_hover)
+                    }),
+            );
+        let content = div()
+            .overflow_hidden()
+            .w_full()
+            .p_4()
+            .rounded_lg()
+            .when(self.collapsed, |this| this.h(collapsed_height))
+            .bg(cx.theme().colors().surface_background)
+            .child(self.message);
+
+        v_flex()
+            .gap_1()
+            .child(ChatMessageHeader::new(self.player))
+            .child(h_flex().gap_3().child(collapse_handle).child(content))
+    }
+}
+
+#[derive(IntoElement)]
+struct ChatMessageHeader {
+    player: UserOrAssistant,
+    contexts: Vec<()>,
+}
+
+impl ChatMessageHeader {
+    fn new(player: UserOrAssistant) -> Self {
+        Self {
+            player,
+            contexts: Vec::new(),
+        }
+    }
+}
+
+impl RenderOnce for ChatMessageHeader {
+    fn render(self, _cx: &mut WindowContext) -> impl IntoElement {
+        let (username, avatar_uri) = match self.player {
+            UserOrAssistant::Assistant => (
+                "Assistant".into(),
+                Some("https://zed.dev/assistant_avatar.png".into()),
+            ),
+            UserOrAssistant::User(Some(user)) => {
+                (user.github_login.clone(), Some(user.avatar_uri.clone()))
+            }
+            UserOrAssistant::User(None) => ("You".into(), None),
+        };
+
+        h_flex()
+            .justify_between()
+            .child(
+                h_flex()
+                    .gap_3()
+                    .map(|this| {
+                        let avatar_size = rems(20.0 / 16.0);
+                        if let Some(avatar_uri) = avatar_uri {
+                            this.child(Avatar::new(avatar_uri).size(avatar_size))
+                        } else {
+                            this.child(div().size(avatar_size))
+                        }
+                    })
+                    .child(Label::new(username).color(Color::Default)),
+            )
+            .child(div().when(!self.contexts.is_empty(), |this| {
+                this.child(Label::new(self.contexts.len().to_string()).color(Color::Muted))
+            }))
+    }
+}

crates/assistant2/src/ui/chat_message_header.rs

@@ -1,57 +0,0 @@
-use client::User;
-use std::sync::Arc;
-use ui::{prelude::*, Avatar};
-
-pub enum UserOrAssistant {
-    User(Option<Arc<User>>),
-    Assistant,
-}
-
-#[derive(IntoElement)]
-pub struct ChatMessageHeader {
-    player: UserOrAssistant,
-    contexts: Vec<()>,
-}
-
-impl ChatMessageHeader {
-    pub fn new(player: UserOrAssistant) -> Self {
-        Self {
-            player,
-            contexts: Vec::new(),
-        }
-    }
-}
-
-impl RenderOnce for ChatMessageHeader {
-    fn render(self, _cx: &mut WindowContext) -> impl IntoElement {
-        let (username, avatar_uri) = match self.player {
-            UserOrAssistant::Assistant => (
-                "Assistant".into(),
-                Some("https://zed.dev/assistant_avatar.png".into()),
-            ),
-            UserOrAssistant::User(Some(user)) => {
-                (user.github_login.clone(), Some(user.avatar_uri.clone()))
-            }
-            UserOrAssistant::User(None) => ("You".into(), None),
-        };
-
-        h_flex()
-            .justify_between()
-            .child(
-                h_flex()
-                    .gap_3()
-                    .map(|this| {
-                        let avatar_size = rems(20.0 / 16.0);
-                        if let Some(avatar_uri) = avatar_uri {
-                            this.child(Avatar::new(avatar_uri).size(avatar_size))
-                        } else {
-                            this.child(div().size(avatar_size))
-                        }
-                    })
-                    .child(Label::new(username).color(Color::Default)),
-            )
-            .child(div().when(!self.contexts.is_empty(), |this| {
-                this.child(Label::new(self.contexts.len().to_string()).color(Color::Muted))
-            }))
-    }
-}