Start displaying messages in new thread element

Agus Zubiaga and Smit Barmase created 5 months ago

Co-authored-by: Smit Barmase <[e-mail]>

Change summary

crates/agent2/src/acp.rs            | 16 +++++++
crates/agent2/src/agent2.rs         | 22 ++++++++++
crates/agent2/src/thread_element.rs | 63 +++++++++++++++++++++++++++---
3 files changed, 93 insertions(+), 8 deletions(-)

Detailed changes

crates/agent2/src/acp.rs

@@ -82,8 +82,22 @@ impl acp::Client for AcpClientDelegate {
 
     async fn stream_message_chunk(
         &self,
-        chunk: acp::StreamMessageChunkParams,
+        params: acp::StreamMessageChunkParams,
     ) -> Result<acp::StreamMessageChunkResponse> {
+        let cx = &mut self.cx.clone();
+
+        cx.update(|cx| {
+            self.update_thread(&params.thread_id.into(), cx, |thread, cx| {
+                let acp::MessageChunk::Text { chunk } = &params.chunk;
+                thread.push_assistant_chunk(
+                    MessageChunk::Text {
+                        chunk: chunk.into(),
+                    },
+                    cx,
+                )
+            });
+        })?;
+
         Ok(acp::StreamMessageChunkResponse)
     }
 

crates/agent2/src/agent2.rs

@@ -219,6 +219,28 @@ impl Thread {
         cx.notify();
     }
 
+    pub fn push_assistant_chunk(&mut self, chunk: MessageChunk, cx: &mut Context<Self>) {
+        if let Some(last_entry) = self.entries.last_mut() {
+            if let AgentThreadEntryContent::Message(Message {
+                ref mut chunks,
+                role: Role::Assistant,
+            }) = last_entry.content
+            {
+                chunks.push(chunk);
+                return;
+            }
+        }
+
+        self.entries.push(ThreadEntry {
+            id: self.next_entry_id.post_inc(),
+            content: AgentThreadEntryContent::Message(Message {
+                role: Role::Assistant,
+                chunks: vec![chunk],
+            }),
+        });
+        cx.notify();
+    }
+
     pub fn send(&mut self, message: Message, cx: &mut Context<Self>) -> Task<Result<()>> {
         let agent = self.agent.clone();
         let id = self.id.clone();

crates/agent2/src/thread_element.rs

@@ -1,21 +1,20 @@
-use std::sync::Arc;
-
 use anyhow::Result;
 use editor::{Editor, MultiBuffer};
-use gpui::{App, Entity, Focusable, SharedString, Window, div, prelude::*};
+use gpui::{App, Entity, Focusable, SharedString, Subscription, Window, div, prelude::*};
 use gpui::{FocusHandle, Task};
 use language::Buffer;
 use ui::Tooltip;
 use ui::prelude::*;
 use zed_actions::agent::Chat;
 
-use crate::{Message, MessageChunk, Role, Thread};
+use crate::{AgentThreadEntryContent, Message, MessageChunk, Role, Thread, ThreadEntry};
 
 pub struct ThreadElement {
     thread: Entity<Thread>,
     // todo! use full message editor from agent2
     message_editor: Entity<Editor>,
     send_task: Option<Task<Result<()>>>,
+    _subscription: Subscription,
 }
 
 impl ThreadElement {
@@ -39,10 +38,15 @@ impl ThreadElement {
             editor
         });
 
+        let subscription = cx.observe(&thread, |_, _, cx| {
+            cx.notify();
+        });
+
         Self {
             thread,
             message_editor,
             send_task: None,
+            _subscription: subscription,
         }
     }
 
@@ -60,18 +64,48 @@ impl ThreadElement {
             return;
         }
 
-        self.send_task = Some(self.thread.update(cx, |thread, cx| {
+        let task = self.thread.update(cx, |thread, cx| {
             let message = Message {
                 role: Role::User,
                 chunks: vec![MessageChunk::Text { chunk: text.into() }],
             };
             thread.send(message, cx)
+        });
+
+        self.send_task = Some(cx.spawn(async move |this, cx| {
+            task.await?;
+
+            this.update(cx, |this, _cx| {
+                this.send_task.take();
+            })
         }));
 
         self.message_editor.update(cx, |editor, cx| {
             editor.clear(window, cx);
         });
     }
+
+    fn render_entry(
+        &self,
+        entry: &ThreadEntry,
+        _window: &mut Window,
+        _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::ReadFile { path, content: _ } => {
+                // todo!
+                div()
+                    .child(format!("<Reading file {}>", path.display()))
+                    .into_any()
+            }
+        }
+    }
 }
 
 impl Focusable for ThreadElement {
@@ -81,7 +115,7 @@ impl Focusable for ThreadElement {
 }
 
 impl Render for ThreadElement {
-    fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
+    fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
         let text = self.message_editor.read(cx).text(cx);
         let is_editor_empty = text.is_empty();
         let focus_handle = self.message_editor.focus_handle(cx);
@@ -89,7 +123,22 @@ impl Render for ThreadElement {
         v_flex()
             .key_context("MessageEditor")
             .on_action(cx.listener(Self::chat))
-            .child(div().h_full())
+            .child(
+                v_flex().h_full().gap_1().children(
+                    self.thread
+                        .read(cx)
+                        .entries()
+                        .iter()
+                        .map(|entry| self.render_entry(entry, window, cx)),
+                ),
+            )
+            .when(self.send_task.is_some(), |this| {
+                this.child(
+                    Label::new("Generating...")
+                        .color(Color::Muted)
+                        .size(LabelSize::Small),
+                )
+            })
             .child(
                 div()
                     .bg(cx.theme().colors().editor_background)