1use std::ops::Range;
2
3use collections::HashMap;
4use editor::display_map::CreaseId;
5use editor::{Addon, AnchorRangeExt, Editor};
6use gpui::{Entity, Subscription};
7use ui::prelude::*;
8
9use crate::{
10 context::{AgentContextHandle, AgentContextKey},
11 context_picker::crease_for_mention,
12 context_store::{ContextStore, ContextStoreEvent},
13};
14
15/// Stored information that can be used to resurrect a context crease when creating an editor for a past message.
16#[derive(Clone, Debug)]
17pub struct MessageCrease {
18 pub range: Range<usize>,
19 pub icon_path: SharedString,
20 pub label: SharedString,
21 /// None for a deserialized message, Some otherwise.
22 pub context: Option<AgentContextHandle>,
23}
24
25#[derive(Default)]
26pub struct ContextCreasesAddon {
27 creases: HashMap<AgentContextKey, Vec<(CreaseId, SharedString)>>,
28 _subscription: Option<Subscription>,
29}
30
31impl Addon for ContextCreasesAddon {
32 fn to_any(&self) -> &dyn std::any::Any {
33 self
34 }
35
36 fn to_any_mut(&mut self) -> Option<&mut dyn std::any::Any> {
37 Some(self)
38 }
39}
40
41impl ContextCreasesAddon {
42 pub fn new() -> Self {
43 Self {
44 creases: HashMap::default(),
45 _subscription: None,
46 }
47 }
48
49 pub fn add_creases(
50 &mut self,
51 context_store: &Entity<ContextStore>,
52 key: AgentContextKey,
53 creases: impl IntoIterator<Item = (CreaseId, SharedString)>,
54 cx: &mut Context<Editor>,
55 ) {
56 self.creases.entry(key).or_default().extend(creases);
57 self._subscription = Some(
58 cx.subscribe(context_store, |editor, _, event, cx| match event {
59 ContextStoreEvent::ContextRemoved(key) => {
60 let Some(this) = editor.addon_mut::<Self>() else {
61 return;
62 };
63 let (crease_ids, replacement_texts): (Vec<_>, Vec<_>) = this
64 .creases
65 .remove(key)
66 .unwrap_or_default()
67 .into_iter()
68 .unzip();
69 let ranges = editor
70 .remove_creases(crease_ids, cx)
71 .into_iter()
72 .map(|(_, range)| range)
73 .collect::<Vec<_>>();
74 editor.unfold_ranges(&ranges, false, false, cx);
75 editor.edit(ranges.into_iter().zip(replacement_texts), cx);
76 cx.notify();
77 }
78 }),
79 )
80 }
81
82 pub fn into_inner(self) -> HashMap<AgentContextKey, Vec<(CreaseId, SharedString)>> {
83 self.creases
84 }
85}
86
87pub fn extract_message_creases(
88 editor: &mut Editor,
89 cx: &mut Context<'_, Editor>,
90) -> Vec<MessageCrease> {
91 let buffer_snapshot = editor.buffer().read(cx).snapshot(cx);
92 let mut contexts_by_crease_id = editor
93 .addon_mut::<ContextCreasesAddon>()
94 .map(std::mem::take)
95 .unwrap_or_default()
96 .into_inner()
97 .into_iter()
98 .flat_map(|(key, creases)| {
99 let context = key.0;
100 creases
101 .into_iter()
102 .map(move |(id, _)| (id, context.clone()))
103 })
104 .collect::<HashMap<_, _>>();
105 // Filter the addon's list of creases based on what the editor reports,
106 // since the addon might have removed creases in it.
107
108 editor.display_map.update(cx, |display_map, cx| {
109 display_map
110 .snapshot(cx)
111 .crease_snapshot
112 .creases()
113 .filter_map(|(id, crease)| {
114 Some((
115 id,
116 (
117 crease.range().to_offset(&buffer_snapshot),
118 crease.metadata()?.clone(),
119 ),
120 ))
121 })
122 .map(|(id, (range, metadata))| {
123 let context = contexts_by_crease_id.remove(&id);
124 MessageCrease {
125 range,
126 context,
127 label: metadata.label,
128 icon_path: metadata.icon_path,
129 }
130 })
131 .collect()
132 })
133}
134
135pub fn insert_message_creases(
136 editor: &mut Editor,
137 message_creases: &[MessageCrease],
138 context_store: &Entity<ContextStore>,
139 window: &mut Window,
140 cx: &mut Context<'_, Editor>,
141) {
142 let buffer_snapshot = editor.buffer().read(cx).snapshot(cx);
143 let creases = message_creases
144 .iter()
145 .map(|crease| {
146 let start = buffer_snapshot.anchor_after(crease.range.start);
147 let end = buffer_snapshot.anchor_before(crease.range.end);
148 crease_for_mention(
149 crease.label.clone(),
150 crease.icon_path.clone(),
151 start..end,
152 cx.weak_entity(),
153 )
154 })
155 .collect::<Vec<_>>();
156 let ids = editor.insert_creases(creases.clone(), cx);
157 editor.fold_creases(creases, false, window, cx);
158 if let Some(addon) = editor.addon_mut::<ContextCreasesAddon>() {
159 for (crease, id) in message_creases.iter().zip(ids) {
160 if let Some(context) = crease.context.as_ref() {
161 let key = AgentContextKey(context.clone());
162 addon.add_creases(context_store, key, vec![(id, crease.label.clone())], cx);
163 }
164 }
165 }
166}