From f383a7626fa7cb1141b943cdb28f24ed316b8404 Mon Sep 17 00:00:00 2001
From: Agus Zubiaga <agus@zed.dev>
Date: Thu, 26 Jun 2025 14:16:30 -0300
Subject: [PATCH] Improve user message

Co-authored-by: Smit Barmase <heysmitbarmase@gmail.com>
---
 crates/agent2/src/agent2.rs         |  1 +
 crates/agent2/src/thread_element.rs | 36 ++++++++++++++++++++++-------
 2 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/crates/agent2/src/agent2.rs b/crates/agent2/src/agent2.rs
index d901e416a289555ceb2da8dabb43851f1bad3efc..28031f1e055d41a35275b2d5e86e9f05f8c1c1b1 100644
--- a/crates/agent2/src/agent2.rs
+++ b/crates/agent2/src/agent2.rs
@@ -226,6 +226,7 @@ impl Thread {
                 role: Role::Assistant,
             }) = last_entry.content
             {
+                // todo! merge with last chunk if same type
                 chunks.push(chunk);
                 return;
             }
diff --git a/crates/agent2/src/thread_element.rs b/crates/agent2/src/thread_element.rs
index 09c8d357188482aef1c787802ba480514045ab73..71b8bde8fd258cf6c471c9caaf5e0e2f6ff53f51 100644
--- a/crates/agent2/src/thread_element.rs
+++ b/crates/agent2/src/thread_element.rs
@@ -89,15 +89,34 @@ impl ThreadElement {
         &self,
         entry: &ThreadEntry,
         _window: &mut Window,
-        _cx: &Context<Self>,
+        cx: &Context<Self>,
     ) -> AnyElement {
         match &entry.content {
-            AgentThreadEntryContent::Message(message) => div()
-                .children(message.chunks.iter().map(|chunk| match chunk {
-                    MessageChunk::Text { chunk } => div().child(chunk.clone()),
-                    _ => todo!(),
-                }))
-                .into_any(),
+            AgentThreadEntryContent::Message(message) => {
+                let message_body = div()
+                    .children(message.chunks.iter().map(|chunk| match chunk {
+                        MessageChunk::Text { chunk } => {
+                            // todo! markdown
+                            Label::new(chunk.clone())
+                        }
+                        _ => todo!(),
+                    }))
+                    .into_any();
+
+                match message.role {
+                    Role::User => div()
+                        .my_1()
+                        .p_2()
+                        .bg(cx.theme().colors().editor_background)
+                        .rounded_lg()
+                        .shadow_md()
+                        .border_1()
+                        .border_color(cx.theme().colors().border)
+                        .child(message_body)
+                        .into_any(),
+                    Role::Assistant => message_body,
+                }
+            }
             AgentThreadEntryContent::ReadFile { path, content: _ } => {
                 // todo!
                 div()
@@ -121,6 +140,7 @@ impl Render for ThreadElement {
         let focus_handle = self.message_editor.focus_handle(cx);
 
         v_flex()
+            .p_2()
             .key_context("MessageEditor")
             .on_action(cx.listener(Self::chat))
             .child(
@@ -133,7 +153,7 @@ impl Render for ThreadElement {
                 ),
             )
             .when(self.send_task.is_some(), |this| {
-                this.child(
+                this.my_1().child(
                     Label::new("Generating...")
                         .color(Color::Muted)
                         .size(LabelSize::Small),