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}