message_editor.rs

  1use crate::acp::completion_provider::ContextPickerCompletionProvider;
  2use crate::acp::completion_provider::MentionImage;
  3use crate::acp::completion_provider::MentionSet;
  4use acp_thread::MentionUri;
  5use agent::TextThreadStore;
  6use agent::ThreadStore;
  7use agent_client_protocol as acp;
  8use anyhow::Result;
  9use collections::HashSet;
 10use editor::ExcerptId;
 11use editor::actions::Paste;
 12use editor::display_map::CreaseId;
 13use editor::{
 14    AnchorRangeExt, ContextMenuOptions, ContextMenuPlacement, Editor, EditorElement, EditorMode,
 15    EditorStyle, MultiBuffer,
 16};
 17use futures::FutureExt as _;
 18use gpui::ClipboardEntry;
 19use gpui::Image;
 20use gpui::ImageFormat;
 21use gpui::{
 22    AppContext, Context, Entity, EventEmitter, FocusHandle, Focusable, Task, TextStyle, WeakEntity,
 23};
 24use language::Buffer;
 25use language::Language;
 26use language_model::LanguageModelImage;
 27use parking_lot::Mutex;
 28use project::{CompletionIntent, Project};
 29use settings::Settings;
 30use std::fmt::Write;
 31use std::path::Path;
 32use std::rc::Rc;
 33use std::sync::Arc;
 34use theme::ThemeSettings;
 35use ui::IconName;
 36use ui::SharedString;
 37use ui::{
 38    ActiveTheme, App, InteractiveElement, IntoElement, ParentElement, Render, Styled, TextSize,
 39    Window, div,
 40};
 41use util::ResultExt;
 42use workspace::Workspace;
 43use workspace::notifications::NotifyResultExt as _;
 44use zed_actions::agent::Chat;
 45
 46use super::completion_provider::Mention;
 47
 48pub struct MessageEditor {
 49    editor: Entity<Editor>,
 50    project: Entity<Project>,
 51    thread_store: Entity<ThreadStore>,
 52    text_thread_store: Entity<TextThreadStore>,
 53    mention_set: Arc<Mutex<MentionSet>>,
 54}
 55
 56pub enum MessageEditorEvent {
 57    Send,
 58    Cancel,
 59}
 60
 61impl EventEmitter<MessageEditorEvent> for MessageEditor {}
 62
 63impl MessageEditor {
 64    pub fn new(
 65        workspace: WeakEntity<Workspace>,
 66        project: Entity<Project>,
 67        thread_store: Entity<ThreadStore>,
 68        text_thread_store: Entity<TextThreadStore>,
 69        mode: EditorMode,
 70        window: &mut Window,
 71        cx: &mut Context<Self>,
 72    ) -> Self {
 73        let language = Language::new(
 74            language::LanguageConfig {
 75                completion_query_characters: HashSet::from_iter(['.', '-', '_', '@']),
 76                ..Default::default()
 77            },
 78            None,
 79        );
 80
 81        let mention_set = Arc::new(Mutex::new(MentionSet::default()));
 82        let editor = cx.new(|cx| {
 83            let buffer = cx.new(|cx| Buffer::local("", cx).with_language(Arc::new(language), cx));
 84            let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
 85
 86            let mut editor = Editor::new(mode, buffer, None, window, cx);
 87            editor.set_placeholder_text("Message the agent - @ to include files", cx);
 88            editor.set_show_indent_guides(false, cx);
 89            editor.set_soft_wrap();
 90            editor.set_use_modal_editing(true);
 91            editor.set_completion_provider(Some(Rc::new(ContextPickerCompletionProvider::new(
 92                mention_set.clone(),
 93                workspace,
 94                thread_store.downgrade(),
 95                text_thread_store.downgrade(),
 96                cx.weak_entity(),
 97            ))));
 98            editor.set_context_menu_options(ContextMenuOptions {
 99                min_entries_visible: 12,
100                max_entries_visible: 12,
101                placement: Some(ContextMenuPlacement::Above),
102            });
103            editor
104        });
105
106        Self {
107            editor,
108            project,
109            mention_set,
110            thread_store,
111            text_thread_store,
112        }
113    }
114
115    pub fn is_empty(&self, cx: &App) -> bool {
116        self.editor.read(cx).is_empty(cx)
117    }
118
119    pub fn contents(
120        &self,
121        window: &mut Window,
122        cx: &mut Context<Self>,
123    ) -> Task<Result<Vec<acp::ContentBlock>>> {
124        let contents = self.mention_set.lock().contents(
125            self.project.clone(),
126            self.thread_store.clone(),
127            self.text_thread_store.clone(),
128            window,
129            cx,
130        );
131        let editor = self.editor.clone();
132
133        cx.spawn(async move |_, cx| {
134            let contents = contents.await?;
135
136            editor.update(cx, |editor, cx| {
137                let mut ix = 0;
138                let mut chunks: Vec<acp::ContentBlock> = Vec::new();
139                let text = editor.text(cx);
140                editor.display_map.update(cx, |map, cx| {
141                    let snapshot = map.snapshot(cx);
142                    for (crease_id, crease) in snapshot.crease_snapshot.creases() {
143                        // Skip creases that have been edited out of the message buffer.
144                        if !crease.range().start.is_valid(&snapshot.buffer_snapshot) {
145                            continue;
146                        }
147
148                        let Some(mention) = contents.get(&crease_id) else {
149                            continue;
150                        };
151
152                        let crease_range = crease.range().to_offset(&snapshot.buffer_snapshot);
153                        if crease_range.start > ix {
154                            chunks.push(text[ix..crease_range.start].into());
155                        }
156                        let chunk = match mention {
157                            Mention::Text { uri, content } => {
158                                acp::ContentBlock::Resource(acp::EmbeddedResource {
159                                    annotations: None,
160                                    resource: acp::EmbeddedResourceResource::TextResourceContents(
161                                        acp::TextResourceContents {
162                                            mime_type: None,
163                                            text: content.clone(),
164                                            uri: uri.to_uri().to_string(),
165                                        },
166                                    ),
167                                })
168                            }
169                            Mention::Image(mention_image) => {
170                                acp::ContentBlock::Image(acp::ImageContent {
171                                    annotations: None,
172                                    data: mention_image.data.to_string(),
173                                    mime_type: mention_image.format.mime_type().into(),
174                                    uri: mention_image
175                                        .abs_path
176                                        .as_ref()
177                                        .map(|path| format!("file://{}", path.display())),
178                                })
179                            }
180                        };
181                        chunks.push(chunk);
182                        ix = crease_range.end;
183                    }
184
185                    if ix < text.len() {
186                        let last_chunk = text[ix..].trim_end();
187                        if !last_chunk.is_empty() {
188                            chunks.push(last_chunk.into());
189                        }
190                    }
191                });
192
193                chunks
194            })
195        })
196    }
197
198    pub fn clear(&mut self, window: &mut Window, cx: &mut Context<Self>) {
199        self.editor.update(cx, |editor, cx| {
200            editor.clear(window, cx);
201            editor.remove_creases(self.mention_set.lock().drain(), cx)
202        });
203    }
204
205    fn chat(&mut self, _: &Chat, _: &mut Window, cx: &mut Context<Self>) {
206        cx.emit(MessageEditorEvent::Send)
207    }
208
209    fn cancel(&mut self, _: &menu::Cancel, _: &mut Window, cx: &mut Context<Self>) {
210        cx.emit(MessageEditorEvent::Cancel)
211    }
212
213    fn paste(&mut self, _: &Paste, window: &mut Window, cx: &mut Context<Self>) {
214        let images = cx
215            .read_from_clipboard()
216            .map(|item| {
217                item.into_entries()
218                    .filter_map(|entry| {
219                        if let ClipboardEntry::Image(image) = entry {
220                            Some(image)
221                        } else {
222                            None
223                        }
224                    })
225                    .collect::<Vec<_>>()
226            })
227            .unwrap_or_default();
228
229        if images.is_empty() {
230            return;
231        }
232        cx.stop_propagation();
233
234        let replacement_text = "image";
235        for image in images {
236            let (excerpt_id, anchor) = self.editor.update(cx, |message_editor, cx| {
237                let snapshot = message_editor.snapshot(window, cx);
238                let (excerpt_id, _, snapshot) = snapshot.buffer_snapshot.as_singleton().unwrap();
239
240                let anchor = snapshot.anchor_before(snapshot.len());
241                message_editor.edit(
242                    [(
243                        multi_buffer::Anchor::max()..multi_buffer::Anchor::max(),
244                        format!("{replacement_text} "),
245                    )],
246                    cx,
247                );
248                (*excerpt_id, anchor)
249            });
250
251            self.insert_image(
252                excerpt_id,
253                anchor,
254                replacement_text.len(),
255                Arc::new(image),
256                None,
257                window,
258                cx,
259            );
260        }
261    }
262
263    pub fn insert_dragged_files(
264        &self,
265        paths: Vec<project::ProjectPath>,
266        window: &mut Window,
267        cx: &mut Context<Self>,
268    ) {
269        let buffer = self.editor.read(cx).buffer().clone();
270        let Some((&excerpt_id, _, _)) = buffer.read(cx).snapshot(cx).as_singleton() else {
271            return;
272        };
273        let Some(buffer) = buffer.read(cx).as_singleton() else {
274            return;
275        };
276        for path in paths {
277            let Some(entry) = self.project.read(cx).entry_for_path(&path, cx) else {
278                continue;
279            };
280            let Some(abs_path) = self.project.read(cx).absolute_path(&path, cx) else {
281                continue;
282            };
283
284            let anchor = buffer.update(cx, |buffer, _cx| buffer.anchor_before(buffer.len()));
285            let path_prefix = abs_path
286                .file_name()
287                .unwrap_or(path.path.as_os_str())
288                .display()
289                .to_string();
290            let Some(completion) = ContextPickerCompletionProvider::completion_for_path(
291                path,
292                &path_prefix,
293                false,
294                entry.is_dir(),
295                excerpt_id,
296                anchor..anchor,
297                self.editor.clone(),
298                self.mention_set.clone(),
299                self.project.clone(),
300                cx,
301            ) else {
302                continue;
303            };
304
305            self.editor.update(cx, |message_editor, cx| {
306                message_editor.edit(
307                    [(
308                        multi_buffer::Anchor::max()..multi_buffer::Anchor::max(),
309                        completion.new_text,
310                    )],
311                    cx,
312                );
313            });
314            if let Some(confirm) = completion.confirm.clone() {
315                confirm(CompletionIntent::Complete, window, cx);
316            }
317        }
318    }
319
320    fn insert_image(
321        &mut self,
322        excerpt_id: ExcerptId,
323        crease_start: text::Anchor,
324        content_len: usize,
325        image: Arc<Image>,
326        abs_path: Option<Arc<Path>>,
327        window: &mut Window,
328        cx: &mut Context<Self>,
329    ) {
330        let Some(crease_id) = insert_crease_for_image(
331            excerpt_id,
332            crease_start,
333            content_len,
334            self.editor.clone(),
335            window,
336            cx,
337        ) else {
338            return;
339        };
340        self.editor.update(cx, |_editor, cx| {
341            let format = image.format;
342            let convert = LanguageModelImage::from_image(image, cx);
343
344            let task = cx
345                .spawn_in(window, async move |editor, cx| {
346                    if let Some(image) = convert.await {
347                        Ok(MentionImage {
348                            abs_path,
349                            data: image.source,
350                            format,
351                        })
352                    } else {
353                        editor
354                            .update(cx, |editor, cx| {
355                                let snapshot = editor.buffer().read(cx).snapshot(cx);
356                                let Some(anchor) =
357                                    snapshot.anchor_in_excerpt(excerpt_id, crease_start)
358                                else {
359                                    return;
360                                };
361                                editor.display_map.update(cx, |display_map, cx| {
362                                    display_map.unfold_intersecting(vec![anchor..anchor], true, cx);
363                                });
364                                editor.remove_creases([crease_id], cx);
365                            })
366                            .ok();
367                        Err("Failed to convert image".to_string())
368                    }
369                })
370                .shared();
371
372            cx.spawn_in(window, {
373                let task = task.clone();
374                async move |_, cx| task.clone().await.notify_async_err(cx)
375            })
376            .detach();
377
378            self.mention_set.lock().insert_image(crease_id, task);
379        });
380    }
381
382    pub fn set_mode(&mut self, mode: EditorMode, cx: &mut Context<Self>) {
383        self.editor.update(cx, |editor, cx| {
384            editor.set_mode(mode);
385            cx.notify()
386        });
387    }
388
389    pub fn set_message(
390        &mut self,
391        message: Vec<acp::ContentBlock>,
392        window: &mut Window,
393        cx: &mut Context<Self>,
394    ) {
395        let mut text = String::new();
396        let mut mentions = Vec::new();
397        let mut images = Vec::new();
398
399        for chunk in message {
400            match chunk {
401                acp::ContentBlock::Text(text_content) => {
402                    text.push_str(&text_content.text);
403                }
404                acp::ContentBlock::Resource(acp::EmbeddedResource {
405                    resource: acp::EmbeddedResourceResource::TextResourceContents(resource),
406                    ..
407                }) => {
408                    if let Some(mention_uri) = MentionUri::parse(&resource.uri).log_err() {
409                        let start = text.len();
410                        write!(&mut text, "{}", mention_uri.as_link()).ok();
411                        let end = text.len();
412                        mentions.push((start..end, mention_uri));
413                    }
414                }
415                acp::ContentBlock::Image(content) => {
416                    let start = text.len();
417                    text.push_str("image");
418                    let end = text.len();
419                    images.push((start..end, content));
420                }
421                acp::ContentBlock::Audio(_)
422                | acp::ContentBlock::Resource(_)
423                | acp::ContentBlock::ResourceLink(_) => {}
424            }
425        }
426
427        let snapshot = self.editor.update(cx, |editor, cx| {
428            editor.set_text(text, window, cx);
429            editor.buffer().read(cx).snapshot(cx)
430        });
431
432        self.mention_set.lock().clear();
433        for (range, mention_uri) in mentions {
434            let anchor = snapshot.anchor_before(range.start);
435            let crease_id = crate::context_picker::insert_crease_for_mention(
436                anchor.excerpt_id,
437                anchor.text_anchor,
438                range.end - range.start,
439                mention_uri.name().into(),
440                mention_uri.icon_path(cx),
441                self.editor.clone(),
442                window,
443                cx,
444            );
445
446            if let Some(crease_id) = crease_id {
447                self.mention_set.lock().insert_uri(crease_id, mention_uri);
448            }
449        }
450        for (range, content) in images {
451            let Some(format) = ImageFormat::from_mime_type(&content.mime_type) else {
452                continue;
453            };
454            let anchor = snapshot.anchor_before(range.start);
455            let abs_path = content
456                .uri
457                .as_ref()
458                .and_then(|uri| uri.strip_prefix("file://").map(|s| Path::new(s).into()));
459
460            let name = content
461                .uri
462                .as_ref()
463                .and_then(|uri| {
464                    uri.strip_prefix("file://")
465                        .and_then(|path| Path::new(path).file_name())
466                })
467                .map(|name| name.to_string_lossy().to_string())
468                .unwrap_or("Image".to_owned());
469            let crease_id = crate::context_picker::insert_crease_for_mention(
470                anchor.excerpt_id,
471                anchor.text_anchor,
472                range.end - range.start,
473                name.into(),
474                IconName::Image.path().into(),
475                self.editor.clone(),
476                window,
477                cx,
478            );
479            let data: SharedString = content.data.to_string().into();
480
481            if let Some(crease_id) = crease_id {
482                self.mention_set.lock().insert_image(
483                    crease_id,
484                    Task::ready(Ok(MentionImage {
485                        abs_path,
486                        data,
487                        format,
488                    }))
489                    .shared(),
490                );
491            }
492        }
493        cx.notify();
494    }
495
496    #[cfg(test)]
497    pub fn set_text(&mut self, text: &str, window: &mut Window, cx: &mut Context<Self>) {
498        self.editor.update(cx, |editor, cx| {
499            editor.set_text(text, window, cx);
500        });
501    }
502}
503
504impl Focusable for MessageEditor {
505    fn focus_handle(&self, cx: &App) -> FocusHandle {
506        self.editor.focus_handle(cx)
507    }
508}
509
510impl Render for MessageEditor {
511    fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
512        div()
513            .key_context("MessageEditor")
514            .on_action(cx.listener(Self::chat))
515            .on_action(cx.listener(Self::cancel))
516            .capture_action(cx.listener(Self::paste))
517            .flex_1()
518            .child({
519                let settings = ThemeSettings::get_global(cx);
520                let font_size = TextSize::Small
521                    .rems(cx)
522                    .to_pixels(settings.agent_font_size(cx));
523                let line_height = settings.buffer_line_height.value() * font_size;
524
525                let text_style = TextStyle {
526                    color: cx.theme().colors().text,
527                    font_family: settings.buffer_font.family.clone(),
528                    font_fallbacks: settings.buffer_font.fallbacks.clone(),
529                    font_features: settings.buffer_font.features.clone(),
530                    font_size: font_size.into(),
531                    line_height: line_height.into(),
532                    ..Default::default()
533                };
534
535                EditorElement::new(
536                    &self.editor,
537                    EditorStyle {
538                        background: cx.theme().colors().editor_background,
539                        local_player: cx.theme().players().local(),
540                        text: text_style,
541                        syntax: cx.theme().syntax().clone(),
542                        ..Default::default()
543                    },
544                )
545            })
546    }
547}
548
549pub(crate) fn insert_crease_for_image(
550    excerpt_id: ExcerptId,
551    anchor: text::Anchor,
552    content_len: usize,
553    editor: Entity<Editor>,
554    window: &mut Window,
555    cx: &mut App,
556) -> Option<CreaseId> {
557    crate::context_picker::insert_crease_for_mention(
558        excerpt_id,
559        anchor,
560        content_len,
561        "Image".into(),
562        IconName::Image.path().into(),
563        editor,
564        window,
565        cx,
566    )
567}
568
569#[cfg(test)]
570mod tests {
571    use std::path::Path;
572
573    use agent::{TextThreadStore, ThreadStore};
574    use agent_client_protocol as acp;
575    use editor::EditorMode;
576    use fs::FakeFs;
577    use gpui::{AppContext, TestAppContext};
578    use lsp::{CompletionContext, CompletionTriggerKind};
579    use project::{CompletionIntent, Project};
580    use serde_json::json;
581    use util::path;
582    use workspace::Workspace;
583
584    use crate::acp::{message_editor::MessageEditor, thread_view::tests::init_test};
585
586    #[gpui::test]
587    async fn test_at_mention_removal(cx: &mut TestAppContext) {
588        init_test(cx);
589
590        let fs = FakeFs::new(cx.executor());
591        fs.insert_tree("/project", json!({"file": ""})).await;
592        let project = Project::test(fs, [Path::new(path!("/project"))], cx).await;
593
594        let (workspace, cx) =
595            cx.add_window_view(|window, cx| Workspace::test_new(project.clone(), window, cx));
596
597        let thread_store = cx.new(|cx| ThreadStore::fake(project.clone(), cx));
598        let text_thread_store = cx.new(|cx| TextThreadStore::fake(project.clone(), cx));
599
600        let message_editor = cx.update(|window, cx| {
601            cx.new(|cx| {
602                MessageEditor::new(
603                    workspace.downgrade(),
604                    project.clone(),
605                    thread_store.clone(),
606                    text_thread_store.clone(),
607                    EditorMode::AutoHeight {
608                        min_lines: 1,
609                        max_lines: None,
610                    },
611                    window,
612                    cx,
613                )
614            })
615        });
616        let editor = message_editor.update(cx, |message_editor, _| message_editor.editor.clone());
617
618        cx.run_until_parked();
619
620        let excerpt_id = editor.update(cx, |editor, cx| {
621            editor
622                .buffer()
623                .read(cx)
624                .excerpt_ids()
625                .into_iter()
626                .next()
627                .unwrap()
628        });
629        let completions = editor.update_in(cx, |editor, window, cx| {
630            editor.set_text("Hello @file ", window, cx);
631            let buffer = editor.buffer().read(cx).as_singleton().unwrap();
632            let completion_provider = editor.completion_provider().unwrap();
633            completion_provider.completions(
634                excerpt_id,
635                &buffer,
636                text::Anchor::MAX,
637                CompletionContext {
638                    trigger_kind: CompletionTriggerKind::TRIGGER_CHARACTER,
639                    trigger_character: Some("@".into()),
640                },
641                window,
642                cx,
643            )
644        });
645        let [_, completion]: [_; 2] = completions
646            .await
647            .unwrap()
648            .into_iter()
649            .flat_map(|response| response.completions)
650            .collect::<Vec<_>>()
651            .try_into()
652            .unwrap();
653
654        editor.update_in(cx, |editor, window, cx| {
655            let snapshot = editor.buffer().read(cx).snapshot(cx);
656            let start = snapshot
657                .anchor_in_excerpt(excerpt_id, completion.replace_range.start)
658                .unwrap();
659            let end = snapshot
660                .anchor_in_excerpt(excerpt_id, completion.replace_range.end)
661                .unwrap();
662            editor.edit([(start..end, completion.new_text)], cx);
663            (completion.confirm.unwrap())(CompletionIntent::Complete, window, cx);
664        });
665
666        cx.run_until_parked();
667
668        // Backspace over the inserted crease (and the following space).
669        editor.update_in(cx, |editor, window, cx| {
670            editor.backspace(&Default::default(), window, cx);
671            editor.backspace(&Default::default(), window, cx);
672        });
673
674        let content = message_editor
675            .update_in(cx, |message_editor, window, cx| {
676                message_editor.contents(window, cx)
677            })
678            .await
679            .unwrap();
680
681        // We don't send a resource link for the deleted crease.
682        pretty_assertions::assert_matches!(content.as_slice(), [acp::ContentBlock::Text { .. }]);
683    }
684}