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