thread_element.rs

  1use anyhow::Result;
  2use editor::{Editor, MultiBuffer};
  3use gpui::{App, Entity, Focusable, SharedString, Subscription, Window, div, prelude::*};
  4use gpui::{FocusHandle, Task};
  5use language::Buffer;
  6use ui::Tooltip;
  7use ui::prelude::*;
  8use zed_actions::agent::Chat;
  9
 10use crate::{AgentThreadEntryContent, Message, MessageChunk, Role, Thread, ThreadEntry};
 11
 12pub struct ThreadElement {
 13    thread: Entity<Thread>,
 14    // todo! use full message editor from agent2
 15    message_editor: Entity<Editor>,
 16    send_task: Option<Task<Result<()>>>,
 17    _subscription: Subscription,
 18}
 19
 20impl ThreadElement {
 21    pub fn new(thread: Entity<Thread>, window: &mut Window, cx: &mut Context<Self>) -> Self {
 22        let message_editor = cx.new(|cx| {
 23            let buffer = cx.new(|cx| Buffer::local("", cx));
 24            let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
 25
 26            let mut editor = Editor::new(
 27                editor::EditorMode::AutoHeight {
 28                    min_lines: 5,
 29                    max_lines: None,
 30                },
 31                buffer,
 32                None,
 33                window,
 34                cx,
 35            );
 36            editor.set_placeholder_text("Send a message", cx);
 37            editor.set_soft_wrap();
 38            editor
 39        });
 40
 41        let subscription = cx.observe(&thread, |_, _, cx| {
 42            cx.notify();
 43        });
 44
 45        Self {
 46            thread,
 47            message_editor,
 48            send_task: None,
 49            _subscription: subscription,
 50        }
 51    }
 52
 53    pub fn title(&self, cx: &App) -> SharedString {
 54        self.thread.read(cx).title()
 55    }
 56
 57    pub fn cancel(&mut self) {
 58        self.send_task.take();
 59    }
 60
 61    fn chat(&mut self, _: &Chat, window: &mut Window, cx: &mut Context<Self>) {
 62        let text = self.message_editor.read(cx).text(cx);
 63        if text.is_empty() {
 64            return;
 65        }
 66
 67        let task = self.thread.update(cx, |thread, cx| {
 68            let message = Message {
 69                role: Role::User,
 70                chunks: vec![MessageChunk::Text { chunk: text.into() }],
 71            };
 72            thread.send(message, cx)
 73        });
 74
 75        self.send_task = Some(cx.spawn(async move |this, cx| {
 76            task.await?;
 77
 78            this.update(cx, |this, _cx| {
 79                this.send_task.take();
 80            })
 81        }));
 82
 83        self.message_editor.update(cx, |editor, cx| {
 84            editor.clear(window, cx);
 85        });
 86    }
 87
 88    fn render_entry(
 89        &self,
 90        entry: &ThreadEntry,
 91        _window: &mut Window,
 92        cx: &Context<Self>,
 93    ) -> AnyElement {
 94        match &entry.content {
 95            AgentThreadEntryContent::Message(message) => {
 96                let message_body = div()
 97                    .children(message.chunks.iter().map(|chunk| match chunk {
 98                        MessageChunk::Text { chunk } => {
 99                            // todo! markdown
100                            Label::new(chunk.clone())
101                        }
102                        _ => todo!(),
103                    }))
104                    .into_any();
105
106                match message.role {
107                    Role::User => div()
108                        .my_1()
109                        .p_2()
110                        .bg(cx.theme().colors().editor_background)
111                        .rounded_lg()
112                        .shadow_md()
113                        .border_1()
114                        .border_color(cx.theme().colors().border)
115                        .child(message_body)
116                        .into_any(),
117                    Role::Assistant => message_body,
118                }
119            }
120            AgentThreadEntryContent::ReadFile { path, content: _ } => {
121                // todo!
122                div()
123                    .child(format!("<Reading file {}>", path.display()))
124                    .into_any()
125            }
126        }
127    }
128}
129
130impl Focusable for ThreadElement {
131    fn focus_handle(&self, cx: &App) -> FocusHandle {
132        self.message_editor.focus_handle(cx)
133    }
134}
135
136impl Render for ThreadElement {
137    fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
138        let text = self.message_editor.read(cx).text(cx);
139        let is_editor_empty = text.is_empty();
140        let focus_handle = self.message_editor.focus_handle(cx);
141
142        v_flex()
143            .key_context("MessageEditor")
144            .on_action(cx.listener(Self::chat))
145            .child(
146                // todo! use gpui::list
147                v_flex().p_2().h_full().gap_1().children(
148                    self.thread
149                        .read(cx)
150                        .entries()
151                        .iter()
152                        .map(|entry| self.render_entry(entry, window, cx)),
153                ),
154            )
155            .when(self.send_task.is_some(), |this| {
156                this.child(
157                    div().p_2().child(
158                        Label::new("Generating...")
159                            .color(Color::Muted)
160                            .size(LabelSize::Small),
161                    ),
162                )
163            })
164            .child(
165                div()
166                    .bg(cx.theme().colors().editor_background)
167                    .border_t_1()
168                    .border_color(cx.theme().colors().border)
169                    .p_2()
170                    .child(self.message_editor.clone()),
171            )
172            .child(
173                h_flex()
174                    .p_2()
175                    .justify_end()
176                    .child(if self.send_task.is_some() {
177                        IconButton::new("stop-generation", IconName::StopFilled)
178                            .icon_color(Color::Error)
179                            .style(ButtonStyle::Tinted(ui::TintColor::Error))
180                            .tooltip(move |window, cx| {
181                                Tooltip::for_action(
182                                    "Stop Generation",
183                                    &editor::actions::Cancel,
184                                    window,
185                                    cx,
186                                )
187                            })
188                            .disabled(is_editor_empty)
189                            .on_click(cx.listener(|this, _event, _, _| this.cancel()))
190                    } else {
191                        IconButton::new("send-message", IconName::Send)
192                            .icon_color(Color::Accent)
193                            .style(ButtonStyle::Filled)
194                            .disabled(is_editor_empty)
195                            .on_click({
196                                let focus_handle = focus_handle.clone();
197                                move |_event, window, cx| {
198                                    focus_handle.dispatch_action(&Chat, window, cx);
199                                }
200                            })
201                            .when(!is_editor_empty, |button| {
202                                button.tooltip(move |window, cx| {
203                                    Tooltip::for_action("Send", &Chat, window, cx)
204                                })
205                            })
206                            .when(is_editor_empty, |button| {
207                                button.tooltip(Tooltip::text("Type a message to submit"))
208                            })
209                    }),
210            )
211    }
212}