1use std::sync::Arc;
2
3use collections::HashMap;
4use editor::{Editor, MultiBuffer};
5use gpui::{
6 list, AbsoluteLength, AnyElement, App, ClickEvent, DefiniteLength, EdgesRefinement, Empty,
7 Entity, Focusable, Length, ListAlignment, ListOffset, ListState, StyleRefinement, Subscription,
8 Task, TextStyleRefinement, UnderlineStyle,
9};
10use language::{Buffer, LanguageRegistry};
11use language_model::{LanguageModelRegistry, LanguageModelToolUseId, Role};
12use markdown::{Markdown, MarkdownStyle};
13use settings::Settings as _;
14use theme::ThemeSettings;
15use ui::{prelude::*, Disclosure, KeyBinding};
16use util::ResultExt as _;
17
18use crate::thread::{MessageId, RequestKind, Thread, ThreadError, ThreadEvent};
19use crate::thread_store::ThreadStore;
20use crate::tool_use::{ToolUse, ToolUseStatus};
21use crate::ui::ContextPill;
22
23pub struct ActiveThread {
24 language_registry: Arc<LanguageRegistry>,
25 thread_store: Entity<ThreadStore>,
26 thread: Entity<Thread>,
27 save_thread_task: Option<Task<()>>,
28 messages: Vec<MessageId>,
29 list_state: ListState,
30 rendered_messages_by_id: HashMap<MessageId, Entity<Markdown>>,
31 editing_message: Option<(MessageId, EditMessageState)>,
32 expanded_tool_uses: HashMap<LanguageModelToolUseId, bool>,
33 last_error: Option<ThreadError>,
34 _subscriptions: Vec<Subscription>,
35}
36
37struct EditMessageState {
38 editor: Entity<Editor>,
39}
40
41impl ActiveThread {
42 pub fn new(
43 thread: Entity<Thread>,
44 thread_store: Entity<ThreadStore>,
45 language_registry: Arc<LanguageRegistry>,
46 window: &mut Window,
47 cx: &mut Context<Self>,
48 ) -> Self {
49 let subscriptions = vec![
50 cx.observe(&thread, |_, _, cx| cx.notify()),
51 cx.subscribe_in(&thread, window, Self::handle_thread_event),
52 ];
53
54 let mut this = Self {
55 language_registry,
56 thread_store,
57 thread: thread.clone(),
58 save_thread_task: None,
59 messages: Vec::new(),
60 rendered_messages_by_id: HashMap::default(),
61 expanded_tool_uses: HashMap::default(),
62 list_state: ListState::new(0, ListAlignment::Bottom, px(1024.), {
63 let this = cx.entity().downgrade();
64 move |ix, window: &mut Window, cx: &mut App| {
65 this.update(cx, |this, cx| this.render_message(ix, window, cx))
66 .unwrap()
67 }
68 }),
69 editing_message: None,
70 last_error: None,
71 _subscriptions: subscriptions,
72 };
73
74 for message in thread.read(cx).messages().cloned().collect::<Vec<_>>() {
75 this.push_message(&message.id, message.text.clone(), window, cx);
76 }
77
78 this
79 }
80
81 pub fn thread(&self) -> &Entity<Thread> {
82 &self.thread
83 }
84
85 pub fn is_empty(&self) -> bool {
86 self.messages.is_empty()
87 }
88
89 pub fn summary(&self, cx: &App) -> Option<SharedString> {
90 self.thread.read(cx).summary()
91 }
92
93 pub fn summary_or_default(&self, cx: &App) -> SharedString {
94 self.thread.read(cx).summary_or_default()
95 }
96
97 pub fn cancel_last_completion(&mut self, cx: &mut App) -> bool {
98 self.last_error.take();
99 self.thread
100 .update(cx, |thread, _cx| thread.cancel_last_completion())
101 }
102
103 pub fn last_error(&self) -> Option<ThreadError> {
104 self.last_error.clone()
105 }
106
107 pub fn clear_last_error(&mut self) {
108 self.last_error.take();
109 }
110
111 fn push_message(
112 &mut self,
113 id: &MessageId,
114 text: String,
115 window: &mut Window,
116 cx: &mut Context<Self>,
117 ) {
118 let old_len = self.messages.len();
119 self.messages.push(*id);
120 self.list_state.splice(old_len..old_len, 1);
121
122 let markdown = self.render_markdown(text.into(), window, cx);
123 self.rendered_messages_by_id.insert(*id, markdown);
124 self.list_state.scroll_to(ListOffset {
125 item_ix: old_len,
126 offset_in_item: Pixels(0.0),
127 });
128 }
129
130 fn edited_message(
131 &mut self,
132 id: &MessageId,
133 text: String,
134 window: &mut Window,
135 cx: &mut Context<Self>,
136 ) {
137 let Some(index) = self.messages.iter().position(|message_id| message_id == id) else {
138 return;
139 };
140 self.list_state.splice(index..index + 1, 1);
141 let markdown = self.render_markdown(text.into(), window, cx);
142 self.rendered_messages_by_id.insert(*id, markdown);
143 }
144
145 fn deleted_message(&mut self, id: &MessageId) {
146 let Some(index) = self.messages.iter().position(|message_id| message_id == id) else {
147 return;
148 };
149 self.messages.remove(index);
150 self.list_state.splice(index..index + 1, 0);
151 self.rendered_messages_by_id.remove(id);
152 }
153
154 fn render_markdown(
155 &self,
156 text: SharedString,
157 window: &Window,
158 cx: &mut Context<Self>,
159 ) -> Entity<Markdown> {
160 let theme_settings = ThemeSettings::get_global(cx);
161 let colors = cx.theme().colors();
162 let ui_font_size = TextSize::Default.rems(cx);
163 let buffer_font_size = TextSize::Small.rems(cx);
164 let mut text_style = window.text_style();
165
166 text_style.refine(&TextStyleRefinement {
167 font_family: Some(theme_settings.ui_font.family.clone()),
168 font_fallbacks: theme_settings.ui_font.fallbacks.clone(),
169 font_features: Some(theme_settings.ui_font.features.clone()),
170 font_size: Some(ui_font_size.into()),
171 color: Some(cx.theme().colors().text),
172 ..Default::default()
173 });
174
175 let markdown_style = MarkdownStyle {
176 base_text_style: text_style,
177 syntax: cx.theme().syntax().clone(),
178 selection_background_color: cx.theme().players().local().selection,
179 code_block_overflow_x_scroll: true,
180 table_overflow_x_scroll: true,
181 code_block: StyleRefinement {
182 margin: EdgesRefinement {
183 top: Some(Length::Definite(rems(0.).into())),
184 left: Some(Length::Definite(rems(0.).into())),
185 right: Some(Length::Definite(rems(0.).into())),
186 bottom: Some(Length::Definite(rems(0.5).into())),
187 },
188 padding: EdgesRefinement {
189 top: Some(DefiniteLength::Absolute(AbsoluteLength::Pixels(Pixels(8.)))),
190 left: Some(DefiniteLength::Absolute(AbsoluteLength::Pixels(Pixels(8.)))),
191 right: Some(DefiniteLength::Absolute(AbsoluteLength::Pixels(Pixels(8.)))),
192 bottom: Some(DefiniteLength::Absolute(AbsoluteLength::Pixels(Pixels(8.)))),
193 },
194 background: Some(colors.editor_background.into()),
195 border_color: Some(colors.border_variant),
196 border_widths: EdgesRefinement {
197 top: Some(AbsoluteLength::Pixels(Pixels(1.))),
198 left: Some(AbsoluteLength::Pixels(Pixels(1.))),
199 right: Some(AbsoluteLength::Pixels(Pixels(1.))),
200 bottom: Some(AbsoluteLength::Pixels(Pixels(1.))),
201 },
202 text: Some(TextStyleRefinement {
203 font_family: Some(theme_settings.buffer_font.family.clone()),
204 font_fallbacks: theme_settings.buffer_font.fallbacks.clone(),
205 font_features: Some(theme_settings.buffer_font.features.clone()),
206 font_size: Some(buffer_font_size.into()),
207 ..Default::default()
208 }),
209 ..Default::default()
210 },
211 inline_code: TextStyleRefinement {
212 font_family: Some(theme_settings.buffer_font.family.clone()),
213 font_fallbacks: theme_settings.buffer_font.fallbacks.clone(),
214 font_features: Some(theme_settings.buffer_font.features.clone()),
215 font_size: Some(buffer_font_size.into()),
216 background_color: Some(colors.editor_foreground.opacity(0.1)),
217 ..Default::default()
218 },
219 link: TextStyleRefinement {
220 background_color: Some(colors.editor_foreground.opacity(0.025)),
221 underline: Some(UnderlineStyle {
222 color: Some(colors.text_accent.opacity(0.5)),
223 thickness: px(1.),
224 ..Default::default()
225 }),
226 ..Default::default()
227 },
228 ..Default::default()
229 };
230
231 cx.new(|cx| {
232 Markdown::new(
233 text,
234 markdown_style,
235 Some(self.language_registry.clone()),
236 None,
237 cx,
238 )
239 })
240 }
241
242 fn handle_thread_event(
243 &mut self,
244 _thread: &Entity<Thread>,
245 event: &ThreadEvent,
246 window: &mut Window,
247 cx: &mut Context<Self>,
248 ) {
249 match event {
250 ThreadEvent::ShowError(error) => {
251 self.last_error = Some(error.clone());
252 }
253 ThreadEvent::StreamedCompletion | ThreadEvent::SummaryChanged => {
254 self.save_thread(cx);
255 }
256 ThreadEvent::StreamedAssistantText(message_id, text) => {
257 if let Some(markdown) = self.rendered_messages_by_id.get_mut(&message_id) {
258 markdown.update(cx, |markdown, cx| {
259 markdown.append(text, cx);
260 });
261 }
262 }
263 ThreadEvent::MessageAdded(message_id) => {
264 if let Some(message_text) = self
265 .thread
266 .read(cx)
267 .message(*message_id)
268 .map(|message| message.text.clone())
269 {
270 self.push_message(message_id, message_text, window, cx);
271 }
272
273 self.save_thread(cx);
274 cx.notify();
275 }
276 ThreadEvent::MessageEdited(message_id) => {
277 if let Some(message_text) = self
278 .thread
279 .read(cx)
280 .message(*message_id)
281 .map(|message| message.text.clone())
282 {
283 self.edited_message(message_id, message_text, window, cx);
284 }
285
286 self.save_thread(cx);
287 cx.notify();
288 }
289 ThreadEvent::MessageDeleted(message_id) => {
290 self.deleted_message(message_id);
291 self.save_thread(cx);
292 cx.notify();
293 }
294 ThreadEvent::UsePendingTools => {
295 self.thread.update(cx, |thread, cx| {
296 thread.use_pending_tools(cx);
297 });
298 }
299 ThreadEvent::ToolFinished { .. } => {
300 if self.thread.read(cx).all_tools_finished() {
301 let model_registry = LanguageModelRegistry::read_global(cx);
302 if let Some(model) = model_registry.active_model() {
303 self.thread.update(cx, |thread, cx| {
304 thread.send_tool_results_to_model(model, cx);
305 });
306 }
307 }
308 }
309 ThreadEvent::ScriptFinished => {
310 let model_registry = LanguageModelRegistry::read_global(cx);
311 if let Some(model) = model_registry.active_model() {
312 self.thread.update(cx, |thread, cx| {
313 thread.send_to_model(model, RequestKind::Chat, false, cx);
314 });
315 }
316 }
317 }
318 }
319
320 /// Spawns a task to save the active thread.
321 ///
322 /// Only one task to save the thread will be in flight at a time.
323 fn save_thread(&mut self, cx: &mut Context<Self>) {
324 let thread = self.thread.clone();
325 self.save_thread_task = Some(cx.spawn(|this, mut cx| async move {
326 let task = this
327 .update(&mut cx, |this, cx| {
328 this.thread_store
329 .update(cx, |thread_store, cx| thread_store.save_thread(&thread, cx))
330 })
331 .ok();
332
333 if let Some(task) = task {
334 task.await.log_err();
335 }
336 }));
337 }
338
339 fn start_editing_message(
340 &mut self,
341 message_id: MessageId,
342 message_text: String,
343 window: &mut Window,
344 cx: &mut Context<Self>,
345 ) {
346 let buffer = cx.new(|cx| {
347 MultiBuffer::singleton(cx.new(|cx| Buffer::local(message_text.clone(), cx)), cx)
348 });
349 let editor = cx.new(|cx| {
350 let mut editor = Editor::new(
351 editor::EditorMode::AutoHeight { max_lines: 8 },
352 buffer,
353 None,
354 false,
355 window,
356 cx,
357 );
358 editor.focus_handle(cx).focus(window);
359 editor.move_to_end(&editor::actions::MoveToEnd, window, cx);
360 editor
361 });
362 self.editing_message = Some((
363 message_id,
364 EditMessageState {
365 editor: editor.clone(),
366 },
367 ));
368 cx.notify();
369 }
370
371 fn cancel_editing_message(&mut self, _: &menu::Cancel, _: &mut Window, cx: &mut Context<Self>) {
372 self.editing_message.take();
373 cx.notify();
374 }
375
376 fn confirm_editing_message(
377 &mut self,
378 _: &menu::Confirm,
379 _: &mut Window,
380 cx: &mut Context<Self>,
381 ) {
382 let Some((message_id, state)) = self.editing_message.take() else {
383 return;
384 };
385 let edited_text = state.editor.read(cx).text(cx);
386 self.thread.update(cx, |thread, cx| {
387 thread.edit_message(message_id, Role::User, edited_text, cx);
388 for message_id in self.messages_after(message_id) {
389 thread.delete_message(*message_id, cx);
390 }
391 });
392
393 let provider = LanguageModelRegistry::read_global(cx).active_provider();
394 if provider
395 .as_ref()
396 .map_or(false, |provider| provider.must_accept_terms(cx))
397 {
398 cx.notify();
399 return;
400 }
401 let model_registry = LanguageModelRegistry::read_global(cx);
402 let Some(model) = model_registry.active_model() else {
403 return;
404 };
405
406 self.thread.update(cx, |thread, cx| {
407 thread.send_to_model(model, RequestKind::Chat, false, cx)
408 });
409 cx.notify();
410 }
411
412 fn last_user_message(&self, cx: &Context<Self>) -> Option<MessageId> {
413 self.messages
414 .iter()
415 .rev()
416 .find(|message_id| {
417 self.thread
418 .read(cx)
419 .message(**message_id)
420 .map_or(false, |message| message.role == Role::User)
421 })
422 .cloned()
423 }
424
425 fn messages_after(&self, message_id: MessageId) -> &[MessageId] {
426 self.messages
427 .iter()
428 .position(|id| *id == message_id)
429 .map(|index| &self.messages[index + 1..])
430 .unwrap_or(&[])
431 }
432
433 fn handle_cancel_click(&mut self, _: &ClickEvent, window: &mut Window, cx: &mut Context<Self>) {
434 self.cancel_editing_message(&menu::Cancel, window, cx);
435 }
436
437 fn handle_regenerate_click(
438 &mut self,
439 _: &ClickEvent,
440 window: &mut Window,
441 cx: &mut Context<Self>,
442 ) {
443 self.confirm_editing_message(&menu::Confirm, window, cx);
444 }
445
446 fn render_message(&self, ix: usize, window: &mut Window, cx: &mut Context<Self>) -> AnyElement {
447 let message_id = self.messages[ix];
448 let Some(message) = self.thread.read(cx).message(message_id) else {
449 return Empty.into_any();
450 };
451
452 let Some(markdown) = self.rendered_messages_by_id.get(&message_id) else {
453 return Empty.into_any();
454 };
455
456 let thread = self.thread.read(cx);
457
458 let context = thread.context_for_message(message_id);
459 let tool_uses = thread.tool_uses_for_message(message_id);
460 let scripting_tool_uses = thread.scripting_tool_uses_for_message(message_id);
461
462 // Don't render user messages that are just there for returning tool results.
463 if message.role == Role::User
464 && (thread.message_has_tool_results(message_id)
465 || thread.message_has_scripting_tool_results(message_id))
466 {
467 return Empty.into_any();
468 }
469
470 let allow_editing_message =
471 message.role == Role::User && self.last_user_message(cx) == Some(message_id);
472
473 let edit_message_editor = self
474 .editing_message
475 .as_ref()
476 .filter(|(id, _)| *id == message_id)
477 .map(|(_, state)| state.editor.clone());
478
479 let colors = cx.theme().colors();
480
481 let message_content = v_flex()
482 .child(
483 if let Some(edit_message_editor) = edit_message_editor.clone() {
484 div()
485 .key_context("EditMessageEditor")
486 .on_action(cx.listener(Self::cancel_editing_message))
487 .on_action(cx.listener(Self::confirm_editing_message))
488 .p_2p5()
489 .child(edit_message_editor)
490 } else {
491 div().p_2p5().text_ui(cx).child(markdown.clone())
492 },
493 )
494 .when_some(context, |parent, context| {
495 if !context.is_empty() {
496 parent.child(
497 h_flex().flex_wrap().gap_1().px_1p5().pb_1p5().children(
498 context
499 .into_iter()
500 .map(|context| ContextPill::added(context, false, false, None)),
501 ),
502 )
503 } else {
504 parent
505 }
506 });
507
508 let styled_message = match message.role {
509 Role::User => v_flex()
510 .id(("message-container", ix))
511 .pt_2p5()
512 .px_2p5()
513 .child(
514 v_flex()
515 .bg(colors.editor_background)
516 .rounded_lg()
517 .border_1()
518 .border_color(colors.border)
519 .shadow_sm()
520 .child(
521 h_flex()
522 .py_1()
523 .pl_2()
524 .pr_1()
525 .bg(colors.editor_foreground.opacity(0.05))
526 .border_b_1()
527 .border_color(colors.border)
528 .justify_between()
529 .rounded_t(px(6.))
530 .child(
531 h_flex()
532 .gap_1p5()
533 .child(
534 Icon::new(IconName::PersonCircle)
535 .size(IconSize::XSmall)
536 .color(Color::Muted),
537 )
538 .child(
539 Label::new("You")
540 .size(LabelSize::Small)
541 .color(Color::Muted),
542 ),
543 )
544 .when_some(
545 edit_message_editor.clone(),
546 |this, edit_message_editor| {
547 let focus_handle = edit_message_editor.focus_handle(cx);
548 this.child(
549 h_flex()
550 .gap_1()
551 .child(
552 Button::new("cancel-edit-message", "Cancel")
553 .label_size(LabelSize::Small)
554 .key_binding(
555 KeyBinding::for_action_in(
556 &menu::Cancel,
557 &focus_handle,
558 window,
559 cx,
560 )
561 .map(|kb| kb.size(rems_from_px(12.))),
562 )
563 .on_click(
564 cx.listener(Self::handle_cancel_click),
565 ),
566 )
567 .child(
568 Button::new(
569 "confirm-edit-message",
570 "Regenerate",
571 )
572 .label_size(LabelSize::Small)
573 .key_binding(
574 KeyBinding::for_action_in(
575 &menu::Confirm,
576 &focus_handle,
577 window,
578 cx,
579 )
580 .map(|kb| kb.size(rems_from_px(12.))),
581 )
582 .on_click(
583 cx.listener(Self::handle_regenerate_click),
584 ),
585 ),
586 )
587 },
588 )
589 .when(
590 edit_message_editor.is_none() && allow_editing_message,
591 |this| {
592 this.child(
593 Button::new("edit-message", "Edit")
594 .label_size(LabelSize::Small)
595 .on_click(cx.listener({
596 let message_text = message.text.clone();
597 move |this, _, window, cx| {
598 this.start_editing_message(
599 message_id,
600 message_text.clone(),
601 window,
602 cx,
603 );
604 }
605 })),
606 )
607 },
608 ),
609 )
610 .child(message_content),
611 ),
612 Role::Assistant => div()
613 .id(("message-container", ix))
614 .child(message_content)
615 .map(|parent| {
616 if tool_uses.is_empty() && scripting_tool_uses.is_empty() {
617 return parent;
618 }
619
620 parent.child(
621 v_flex()
622 .children(
623 tool_uses
624 .into_iter()
625 .map(|tool_use| self.render_tool_use(tool_use, cx)),
626 )
627 .children(
628 scripting_tool_uses
629 .into_iter()
630 .map(|tool_use| self.render_scripting_tool_use(tool_use, cx)),
631 ),
632 )
633 }),
634 Role::System => div().id(("message-container", ix)).py_1().px_2().child(
635 v_flex()
636 .bg(colors.editor_background)
637 .rounded_sm()
638 .child(message_content),
639 ),
640 };
641
642 styled_message.into_any()
643 }
644
645 fn render_tool_use(&self, tool_use: ToolUse, cx: &mut Context<Self>) -> impl IntoElement {
646 let is_open = self
647 .expanded_tool_uses
648 .get(&tool_use.id)
649 .copied()
650 .unwrap_or_default();
651
652 div().px_2p5().child(
653 v_flex()
654 .gap_1()
655 .rounded_lg()
656 .border_1()
657 .border_color(cx.theme().colors().border)
658 .child(
659 h_flex()
660 .justify_between()
661 .py_0p5()
662 .pl_1()
663 .pr_2()
664 .bg(cx.theme().colors().editor_foreground.opacity(0.02))
665 .when(is_open, |element| element.border_b_1().rounded_t(px(6.)))
666 .when(!is_open, |element| element.rounded_md())
667 .border_color(cx.theme().colors().border)
668 .child(
669 h_flex()
670 .gap_1()
671 .child(Disclosure::new("tool-use-disclosure", is_open).on_click(
672 cx.listener({
673 let tool_use_id = tool_use.id.clone();
674 move |this, _event, _window, _cx| {
675 let is_open = this
676 .expanded_tool_uses
677 .entry(tool_use_id.clone())
678 .or_insert(false);
679
680 *is_open = !*is_open;
681 }
682 }),
683 ))
684 .child(Label::new(tool_use.name)),
685 )
686 .child(
687 Label::new(match tool_use.status {
688 ToolUseStatus::Pending => "Pending",
689 ToolUseStatus::Running => "Running",
690 ToolUseStatus::Finished(_) => "Finished",
691 ToolUseStatus::Error(_) => "Error",
692 })
693 .size(LabelSize::XSmall)
694 .buffer_font(cx),
695 ),
696 )
697 .map(|parent| {
698 if !is_open {
699 return parent;
700 }
701
702 parent.child(
703 v_flex()
704 .child(
705 v_flex()
706 .gap_0p5()
707 .py_1()
708 .px_2p5()
709 .border_b_1()
710 .border_color(cx.theme().colors().border)
711 .child(Label::new("Input:"))
712 .child(Label::new(
713 serde_json::to_string_pretty(&tool_use.input)
714 .unwrap_or_default(),
715 )),
716 )
717 .map(|parent| match tool_use.status {
718 ToolUseStatus::Finished(output) => parent.child(
719 v_flex()
720 .gap_0p5()
721 .py_1()
722 .px_2p5()
723 .child(Label::new("Result:"))
724 .child(Label::new(output)),
725 ),
726 ToolUseStatus::Error(err) => parent.child(
727 v_flex()
728 .gap_0p5()
729 .py_1()
730 .px_2p5()
731 .child(Label::new("Error:"))
732 .child(Label::new(err)),
733 ),
734 ToolUseStatus::Pending | ToolUseStatus::Running => parent,
735 }),
736 )
737 }),
738 )
739 }
740
741 fn render_scripting_tool_use(
742 &self,
743 tool_use: ToolUse,
744 cx: &mut Context<Self>,
745 ) -> impl IntoElement {
746 // TODO: Add custom rendering for scripting tool uses.
747 self.render_tool_use(tool_use, cx)
748 }
749}
750
751impl Render for ActiveThread {
752 fn render(&mut self, _window: &mut Window, _cx: &mut Context<Self>) -> impl IntoElement {
753 v_flex()
754 .size_full()
755 .child(list(self.list_state.clone()).flex_grow())
756 }
757}