assistant2: Add headers to chat messages (#11191)

This PR adds headers to the chat messages in the new assistant panel.

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

Release Notes:

- N/A

crates/assistant2/src/assistant2.rs

@@ -33,6 +33,8 @@ use workspace::{
 
 pub use assistant_settings::AssistantSettings;
 
+use crate::ui::{ChatMessageHeader, UserOrAssistant};
+
 const MAX_COMPLETION_CALLS_PER_SUBMISSION: usize = 5;
 
 #[derive(Eq, PartialEq, Copy, Clone, Deserialize)]
@@ -526,7 +528,9 @@ impl AssistantChat {
         match &self.messages[ix] {
             ChatMessage::User(UserMessage { body, .. }) => div()
                 .when(!is_last, |element| element.mb_2())
-                .child(div().p_2().child(Label::new("You").color(Color::Default)))
+                .child(ChatMessageHeader::new(UserOrAssistant::User(
+                    self.user_store.read(cx).current_user(),
+                )))
                 .child(
                     div()
                         .p_2()
@@ -551,11 +555,7 @@ impl AssistantChat {
 
                 div()
                     .when(!is_last, |element| element.mb_2())
-                    .child(
-                        div()
-                            .p_2()
-                            .child(Label::new("Assistant").color(Color::Modified)),
-                    )
+                    .child(ChatMessageHeader::new(UserOrAssistant::Assistant))
                     .child(assistant_body)
                     .child(self.render_error(error.clone(), ix, cx))
                     .children(tool_calls.iter().map(|tool_call| {

crates/assistant2/src/ui.rs

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

crates/assistant2/src/ui/chat_message_header.rs

@@ -0,0 +1,57 @@
+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))
+            }))
+    }
+}