message_editor.rs

  1use crate::acp::completion_provider::ContextPickerCompletionProvider;
  2use crate::acp::completion_provider::MentionSet;
  3use acp_thread::MentionUri;
  4use agent::TextThreadStore;
  5use agent::ThreadStore;
  6use agent_client_protocol as acp;
  7use anyhow::Result;
  8use collections::HashSet;
  9use editor::{
 10    AnchorRangeExt, ContextMenuOptions, ContextMenuPlacement, Editor, EditorElement, EditorMode,
 11    EditorStyle, MultiBuffer,
 12};
 13use gpui::{
 14    AppContext, Context, Entity, EventEmitter, FocusHandle, Focusable, Task, TextStyle, WeakEntity,
 15};
 16use language::Buffer;
 17use language::Language;
 18use parking_lot::Mutex;
 19use project::{CompletionIntent, Project};
 20use settings::Settings;
 21use std::fmt::Write;
 22use std::rc::Rc;
 23use std::sync::Arc;
 24use theme::ThemeSettings;
 25use ui::{
 26    ActiveTheme, App, InteractiveElement, IntoElement, ParentElement, Render, Styled, TextSize,
 27    Window, div,
 28};
 29use util::ResultExt;
 30use workspace::Workspace;
 31use zed_actions::agent::Chat;
 32
 33pub struct MessageEditor {
 34    editor: Entity<Editor>,
 35    project: Entity<Project>,
 36    thread_store: Entity<ThreadStore>,
 37    text_thread_store: Entity<TextThreadStore>,
 38    mention_set: Arc<Mutex<MentionSet>>,
 39}
 40
 41pub enum MessageEditorEvent {
 42    Send,
 43    Cancel,
 44}
 45
 46impl EventEmitter<MessageEditorEvent> for MessageEditor {}
 47
 48impl MessageEditor {
 49    pub fn new(
 50        workspace: WeakEntity<Workspace>,
 51        project: Entity<Project>,
 52        thread_store: Entity<ThreadStore>,
 53        text_thread_store: Entity<TextThreadStore>,
 54        mode: EditorMode,
 55        window: &mut Window,
 56        cx: &mut Context<Self>,
 57    ) -> Self {
 58        let language = Language::new(
 59            language::LanguageConfig {
 60                completion_query_characters: HashSet::from_iter(['.', '-', '_', '@']),
 61                ..Default::default()
 62            },
 63            None,
 64        );
 65
 66        let mention_set = Arc::new(Mutex::new(MentionSet::default()));
 67        let editor = cx.new(|cx| {
 68            let buffer = cx.new(|cx| Buffer::local("", cx).with_language(Arc::new(language), cx));
 69            let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
 70
 71            let mut editor = Editor::new(mode, buffer, None, window, cx);
 72            editor.set_placeholder_text("Message the agent - @ to include files", cx);
 73            editor.set_show_indent_guides(false, cx);
 74            editor.set_soft_wrap();
 75            editor.set_use_modal_editing(true);
 76            editor.set_completion_provider(Some(Rc::new(ContextPickerCompletionProvider::new(
 77                mention_set.clone(),
 78                workspace,
 79                thread_store.downgrade(),
 80                text_thread_store.downgrade(),
 81                cx.weak_entity(),
 82            ))));
 83            editor.set_context_menu_options(ContextMenuOptions {
 84                min_entries_visible: 12,
 85                max_entries_visible: 12,
 86                placement: Some(ContextMenuPlacement::Above),
 87            });
 88            editor
 89        });
 90
 91        Self {
 92            editor,
 93            project,
 94            mention_set,
 95            thread_store,
 96            text_thread_store,
 97        }
 98    }
 99
100    pub fn is_empty(&self, cx: &App) -> bool {
101        self.editor.read(cx).is_empty(cx)
102    }
103
104    pub fn contents(
105        &self,
106        window: &mut Window,
107        cx: &mut Context<Self>,
108    ) -> Task<Result<Vec<acp::ContentBlock>>> {
109        let contents = self.mention_set.lock().contents(
110            self.project.clone(),
111            self.thread_store.clone(),
112            self.text_thread_store.clone(),
113            window,
114            cx,
115        );
116        let editor = self.editor.clone();
117
118        cx.spawn(async move |_, cx| {
119            let contents = contents.await?;
120
121            editor.update(cx, |editor, cx| {
122                let mut ix = 0;
123                let mut chunks: Vec<acp::ContentBlock> = Vec::new();
124                let text = editor.text(cx);
125                editor.display_map.update(cx, |map, cx| {
126                    let snapshot = map.snapshot(cx);
127                    for (crease_id, crease) in snapshot.crease_snapshot.creases() {
128                        // Skip creases that have been edited out of the message buffer.
129                        if !crease.range().start.is_valid(&snapshot.buffer_snapshot) {
130                            continue;
131                        }
132
133                        if let Some(mention) = contents.get(&crease_id) {
134                            let crease_range = crease.range().to_offset(&snapshot.buffer_snapshot);
135                            if crease_range.start > ix {
136                                chunks.push(text[ix..crease_range.start].into());
137                            }
138                            chunks.push(acp::ContentBlock::Resource(acp::EmbeddedResource {
139                                annotations: None,
140                                resource: acp::EmbeddedResourceResource::TextResourceContents(
141                                    acp::TextResourceContents {
142                                        mime_type: None,
143                                        text: mention.content.clone(),
144                                        uri: mention.uri.to_uri().to_string(),
145                                    },
146                                ),
147                            }));
148                            ix = crease_range.end;
149                        }
150                    }
151
152                    if ix < text.len() {
153                        let last_chunk = text[ix..].trim_end();
154                        if !last_chunk.is_empty() {
155                            chunks.push(last_chunk.into());
156                        }
157                    }
158                });
159
160                chunks
161            })
162        })
163    }
164
165    pub fn clear(&mut self, window: &mut Window, cx: &mut Context<Self>) {
166        self.editor.update(cx, |editor, cx| {
167            editor.clear(window, cx);
168            editor.remove_creases(self.mention_set.lock().drain(), cx)
169        });
170    }
171
172    fn chat(&mut self, _: &Chat, _: &mut Window, cx: &mut Context<Self>) {
173        cx.emit(MessageEditorEvent::Send)
174    }
175
176    fn cancel(&mut self, _: &menu::Cancel, _: &mut Window, cx: &mut Context<Self>) {
177        cx.emit(MessageEditorEvent::Cancel)
178    }
179
180    pub fn insert_dragged_files(
181        &self,
182        paths: Vec<project::ProjectPath>,
183        window: &mut Window,
184        cx: &mut Context<Self>,
185    ) {
186        let buffer = self.editor.read(cx).buffer().clone();
187        let Some((&excerpt_id, _, _)) = buffer.read(cx).snapshot(cx).as_singleton() else {
188            return;
189        };
190        let Some(buffer) = buffer.read(cx).as_singleton() else {
191            return;
192        };
193        for path in paths {
194            let Some(entry) = self.project.read(cx).entry_for_path(&path, cx) else {
195                continue;
196            };
197            let Some(abs_path) = self.project.read(cx).absolute_path(&path, cx) else {
198                continue;
199            };
200
201            let anchor = buffer.update(cx, |buffer, _cx| buffer.anchor_before(buffer.len()));
202            let path_prefix = abs_path
203                .file_name()
204                .unwrap_or(path.path.as_os_str())
205                .display()
206                .to_string();
207            let Some(completion) = ContextPickerCompletionProvider::completion_for_path(
208                path,
209                &path_prefix,
210                false,
211                entry.is_dir(),
212                excerpt_id,
213                anchor..anchor,
214                self.editor.clone(),
215                self.mention_set.clone(),
216                self.project.clone(),
217                cx,
218            ) else {
219                continue;
220            };
221
222            self.editor.update(cx, |message_editor, cx| {
223                message_editor.edit(
224                    [(
225                        multi_buffer::Anchor::max()..multi_buffer::Anchor::max(),
226                        completion.new_text,
227                    )],
228                    cx,
229                );
230            });
231            if let Some(confirm) = completion.confirm.clone() {
232                confirm(CompletionIntent::Complete, window, cx);
233            }
234        }
235    }
236
237    pub fn set_mode(&mut self, mode: EditorMode, cx: &mut Context<Self>) {
238        self.editor.update(cx, |editor, cx| {
239            editor.set_mode(mode);
240            cx.notify()
241        });
242    }
243
244    pub fn set_message(
245        &mut self,
246        message: &[acp::ContentBlock],
247        window: &mut Window,
248        cx: &mut Context<Self>,
249    ) {
250        let mut text = String::new();
251        let mut mentions = Vec::new();
252
253        for chunk in message {
254            match chunk {
255                acp::ContentBlock::Text(text_content) => {
256                    text.push_str(&text_content.text);
257                }
258                acp::ContentBlock::Resource(acp::EmbeddedResource {
259                    resource: acp::EmbeddedResourceResource::TextResourceContents(resource),
260                    ..
261                }) => {
262                    if let Some(mention_uri) = MentionUri::parse(&resource.uri).log_err() {
263                        let start = text.len();
264                        write!(&mut text, "{}", mention_uri.as_link()).ok();
265                        let end = text.len();
266                        mentions.push((start..end, mention_uri));
267                    }
268                }
269                acp::ContentBlock::Image(_)
270                | acp::ContentBlock::Audio(_)
271                | acp::ContentBlock::Resource(_)
272                | acp::ContentBlock::ResourceLink(_) => {}
273            }
274        }
275
276        let snapshot = self.editor.update(cx, |editor, cx| {
277            editor.set_text(text, window, cx);
278            editor.buffer().read(cx).snapshot(cx)
279        });
280
281        self.mention_set.lock().clear();
282        for (range, mention_uri) in mentions {
283            let anchor = snapshot.anchor_before(range.start);
284            let crease_id = crate::context_picker::insert_crease_for_mention(
285                anchor.excerpt_id,
286                anchor.text_anchor,
287                range.end - range.start,
288                mention_uri.name().into(),
289                mention_uri.icon_path(cx),
290                self.editor.clone(),
291                window,
292                cx,
293            );
294
295            if let Some(crease_id) = crease_id {
296                self.mention_set.lock().insert(crease_id, mention_uri);
297            }
298        }
299        cx.notify();
300    }
301
302    #[cfg(test)]
303    pub fn set_text(&mut self, text: &str, window: &mut Window, cx: &mut Context<Self>) {
304        self.editor.update(cx, |editor, cx| {
305            editor.set_text(text, window, cx);
306        });
307    }
308}
309
310impl Focusable for MessageEditor {
311    fn focus_handle(&self, cx: &App) -> FocusHandle {
312        self.editor.focus_handle(cx)
313    }
314}
315
316impl Render for MessageEditor {
317    fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
318        div()
319            .key_context("MessageEditor")
320            .on_action(cx.listener(Self::chat))
321            .on_action(cx.listener(Self::cancel))
322            .flex_1()
323            .child({
324                let settings = ThemeSettings::get_global(cx);
325                let font_size = TextSize::Small
326                    .rems(cx)
327                    .to_pixels(settings.agent_font_size(cx));
328                let line_height = settings.buffer_line_height.value() * font_size;
329
330                let text_style = TextStyle {
331                    color: cx.theme().colors().text,
332                    font_family: settings.buffer_font.family.clone(),
333                    font_fallbacks: settings.buffer_font.fallbacks.clone(),
334                    font_features: settings.buffer_font.features.clone(),
335                    font_size: font_size.into(),
336                    line_height: line_height.into(),
337                    ..Default::default()
338                };
339
340                EditorElement::new(
341                    &self.editor,
342                    EditorStyle {
343                        background: cx.theme().colors().editor_background,
344                        local_player: cx.theme().players().local(),
345                        text: text_style,
346                        syntax: cx.theme().syntax().clone(),
347                        ..Default::default()
348                    },
349                )
350            })
351    }
352}
353
354#[cfg(test)]
355mod tests {
356    use std::path::Path;
357
358    use agent::{TextThreadStore, ThreadStore};
359    use agent_client_protocol as acp;
360    use editor::EditorMode;
361    use fs::FakeFs;
362    use gpui::{AppContext, TestAppContext};
363    use lsp::{CompletionContext, CompletionTriggerKind};
364    use project::{CompletionIntent, Project};
365    use serde_json::json;
366    use util::path;
367    use workspace::Workspace;
368
369    use crate::acp::{message_editor::MessageEditor, thread_view::tests::init_test};
370
371    #[gpui::test]
372    async fn test_at_mention_removal(cx: &mut TestAppContext) {
373        init_test(cx);
374
375        let fs = FakeFs::new(cx.executor());
376        fs.insert_tree("/project", json!({"file": ""})).await;
377        let project = Project::test(fs, [Path::new(path!("/project"))], cx).await;
378
379        let (workspace, cx) =
380            cx.add_window_view(|window, cx| Workspace::test_new(project.clone(), window, cx));
381
382        let thread_store = cx.new(|cx| ThreadStore::fake(project.clone(), cx));
383        let text_thread_store = cx.new(|cx| TextThreadStore::fake(project.clone(), cx));
384
385        let message_editor = cx.update(|window, cx| {
386            cx.new(|cx| {
387                MessageEditor::new(
388                    workspace.downgrade(),
389                    project.clone(),
390                    thread_store.clone(),
391                    text_thread_store.clone(),
392                    EditorMode::AutoHeight {
393                        min_lines: 1,
394                        max_lines: None,
395                    },
396                    window,
397                    cx,
398                )
399            })
400        });
401        let editor = message_editor.update(cx, |message_editor, _| message_editor.editor.clone());
402
403        cx.run_until_parked();
404
405        let excerpt_id = editor.update(cx, |editor, cx| {
406            editor
407                .buffer()
408                .read(cx)
409                .excerpt_ids()
410                .into_iter()
411                .next()
412                .unwrap()
413        });
414        let completions = editor.update_in(cx, |editor, window, cx| {
415            editor.set_text("Hello @file ", window, cx);
416            let buffer = editor.buffer().read(cx).as_singleton().unwrap();
417            let completion_provider = editor.completion_provider().unwrap();
418            completion_provider.completions(
419                excerpt_id,
420                &buffer,
421                text::Anchor::MAX,
422                CompletionContext {
423                    trigger_kind: CompletionTriggerKind::TRIGGER_CHARACTER,
424                    trigger_character: Some("@".into()),
425                },
426                window,
427                cx,
428            )
429        });
430        let [_, completion]: [_; 2] = completions
431            .await
432            .unwrap()
433            .into_iter()
434            .flat_map(|response| response.completions)
435            .collect::<Vec<_>>()
436            .try_into()
437            .unwrap();
438
439        editor.update_in(cx, |editor, window, cx| {
440            let snapshot = editor.buffer().read(cx).snapshot(cx);
441            let start = snapshot
442                .anchor_in_excerpt(excerpt_id, completion.replace_range.start)
443                .unwrap();
444            let end = snapshot
445                .anchor_in_excerpt(excerpt_id, completion.replace_range.end)
446                .unwrap();
447            editor.edit([(start..end, completion.new_text)], cx);
448            (completion.confirm.unwrap())(CompletionIntent::Complete, window, cx);
449        });
450
451        cx.run_until_parked();
452
453        // Backspace over the inserted crease (and the following space).
454        editor.update_in(cx, |editor, window, cx| {
455            editor.backspace(&Default::default(), window, cx);
456            editor.backspace(&Default::default(), window, cx);
457        });
458
459        let content = message_editor
460            .update_in(cx, |message_editor, window, cx| {
461                message_editor.contents(window, cx)
462            })
463            .await
464            .unwrap();
465
466        // We don't send a resource link for the deleted crease.
467        pretty_assertions::assert_matches!(content.as_slice(), [acp::ContentBlock::Text { .. }]);
468    }
469}