1use std::sync::Arc;
2
3use assistant_settings::{
4 AgentProfile, AgentProfileContent, AssistantSettings, AssistantSettingsContent,
5 ContextServerPresetContent, VersionedAssistantSettingsContent,
6};
7use assistant_tool::{ToolSource, ToolWorkingSet};
8use fs::Fs;
9use fuzzy::{StringMatch, StringMatchCandidate, match_strings};
10use gpui::{App, Context, DismissEvent, Entity, EventEmitter, Focusable, Task, WeakEntity, Window};
11use picker::{Picker, PickerDelegate};
12use settings::{Settings as _, update_settings_file};
13use ui::{HighlightedLabel, ListItem, ListItemSpacing, prelude::*};
14use util::ResultExt as _;
15
16use crate::ThreadStore;
17
18pub struct ToolPicker {
19 picker: Entity<Picker<ToolPickerDelegate>>,
20}
21
22impl ToolPicker {
23 pub fn new(delegate: ToolPickerDelegate, window: &mut Window, cx: &mut Context<Self>) -> Self {
24 let picker = cx.new(|cx| Picker::uniform_list(delegate, window, cx).modal(false));
25 Self { picker }
26 }
27}
28
29impl EventEmitter<DismissEvent> for ToolPicker {}
30
31impl Focusable for ToolPicker {
32 fn focus_handle(&self, cx: &App) -> gpui::FocusHandle {
33 self.picker.focus_handle(cx)
34 }
35}
36
37impl Render for ToolPicker {
38 fn render(&mut self, _window: &mut Window, _cx: &mut Context<Self>) -> impl IntoElement {
39 v_flex().w(rems(34.)).child(self.picker.clone())
40 }
41}
42
43#[derive(Debug, Clone)]
44pub struct ToolEntry {
45 pub name: Arc<str>,
46 pub source: ToolSource,
47}
48
49pub struct ToolPickerDelegate {
50 tool_picker: WeakEntity<ToolPicker>,
51 thread_store: WeakEntity<ThreadStore>,
52 fs: Arc<dyn Fs>,
53 tools: Vec<ToolEntry>,
54 profile_id: Arc<str>,
55 profile: AgentProfile,
56 matches: Vec<StringMatch>,
57 selected_index: usize,
58}
59
60impl ToolPickerDelegate {
61 pub fn new(
62 fs: Arc<dyn Fs>,
63 tool_set: Arc<ToolWorkingSet>,
64 thread_store: WeakEntity<ThreadStore>,
65 profile_id: Arc<str>,
66 profile: AgentProfile,
67 cx: &mut Context<ToolPicker>,
68 ) -> Self {
69 let mut tool_entries = Vec::new();
70
71 for (source, tools) in tool_set.tools_by_source(cx) {
72 tool_entries.extend(tools.into_iter().map(|tool| ToolEntry {
73 name: tool.name().into(),
74 source: source.clone(),
75 }));
76 }
77
78 Self {
79 tool_picker: cx.entity().downgrade(),
80 thread_store,
81 fs,
82 tools: tool_entries,
83 profile_id,
84 profile,
85 matches: Vec::new(),
86 selected_index: 0,
87 }
88 }
89}
90
91impl PickerDelegate for ToolPickerDelegate {
92 type ListItem = ListItem;
93
94 fn match_count(&self) -> usize {
95 self.matches.len()
96 }
97
98 fn selected_index(&self) -> usize {
99 self.selected_index
100 }
101
102 fn set_selected_index(
103 &mut self,
104 ix: usize,
105 _window: &mut Window,
106 _cx: &mut Context<Picker<Self>>,
107 ) {
108 self.selected_index = ix;
109 }
110
111 fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
112 "Search tools…".into()
113 }
114
115 fn update_matches(
116 &mut self,
117 query: String,
118 window: &mut Window,
119 cx: &mut Context<Picker<Self>>,
120 ) -> Task<()> {
121 let background = cx.background_executor().clone();
122 let candidates = self
123 .tools
124 .iter()
125 .enumerate()
126 .map(|(id, profile)| StringMatchCandidate::new(id, profile.name.as_ref()))
127 .collect::<Vec<_>>();
128
129 cx.spawn_in(window, async move |this, cx| {
130 let matches = if query.is_empty() {
131 candidates
132 .into_iter()
133 .enumerate()
134 .map(|(index, candidate)| StringMatch {
135 candidate_id: index,
136 string: candidate.string,
137 positions: Vec::new(),
138 score: 0.,
139 })
140 .collect()
141 } else {
142 match_strings(
143 &candidates,
144 &query,
145 false,
146 100,
147 &Default::default(),
148 background,
149 )
150 .await
151 };
152
153 this.update(cx, |this, _cx| {
154 this.delegate.matches = matches;
155 this.delegate.selected_index = this
156 .delegate
157 .selected_index
158 .min(this.delegate.matches.len().saturating_sub(1));
159 })
160 .log_err();
161 })
162 }
163
164 fn confirm(&mut self, _secondary: bool, window: &mut Window, cx: &mut Context<Picker<Self>>) {
165 if self.matches.is_empty() {
166 self.dismissed(window, cx);
167 return;
168 }
169
170 let candidate_id = self.matches[self.selected_index].candidate_id;
171 let tool = &self.tools[candidate_id];
172
173 let is_enabled = match &tool.source {
174 ToolSource::Native => {
175 let is_enabled = self.profile.tools.entry(tool.name.clone()).or_default();
176 *is_enabled = !*is_enabled;
177 *is_enabled
178 }
179 ToolSource::ContextServer { id } => {
180 let preset = self
181 .profile
182 .context_servers
183 .entry(id.clone().into())
184 .or_default();
185 let is_enabled = preset.tools.entry(tool.name.clone()).or_default();
186 *is_enabled = !*is_enabled;
187 *is_enabled
188 }
189 };
190
191 let active_profile_id = &AssistantSettings::get_global(cx).default_profile;
192 if active_profile_id == &self.profile_id {
193 self.thread_store
194 .update(cx, |this, _cx| {
195 this.load_profile(&self.profile);
196 })
197 .log_err();
198 }
199
200 update_settings_file::<AssistantSettings>(self.fs.clone(), cx, {
201 let profile_id = self.profile_id.clone();
202 let default_profile = self.profile.clone();
203 let tool = tool.clone();
204 move |settings, _cx| match settings {
205 AssistantSettingsContent::Versioned(VersionedAssistantSettingsContent::V2(
206 settings,
207 )) => {
208 let profiles = settings.profiles.get_or_insert_default();
209 let profile =
210 profiles
211 .entry(profile_id)
212 .or_insert_with(|| AgentProfileContent {
213 name: default_profile.name.into(),
214 tools: default_profile.tools,
215 context_servers: default_profile
216 .context_servers
217 .into_iter()
218 .map(|(server_id, preset)| {
219 (
220 server_id,
221 ContextServerPresetContent {
222 tools: preset.tools,
223 },
224 )
225 })
226 .collect(),
227 });
228
229 match tool.source {
230 ToolSource::Native => {
231 *profile.tools.entry(tool.name).or_default() = is_enabled;
232 }
233 ToolSource::ContextServer { id } => {
234 let preset = profile
235 .context_servers
236 .entry(id.clone().into())
237 .or_default();
238 *preset.tools.entry(tool.name.clone()).or_default() = is_enabled;
239 }
240 }
241 }
242 _ => {}
243 }
244 });
245 }
246
247 fn dismissed(&mut self, _window: &mut Window, cx: &mut Context<Picker<Self>>) {
248 self.tool_picker
249 .update(cx, |_this, cx| cx.emit(DismissEvent))
250 .log_err();
251 }
252
253 fn render_match(
254 &self,
255 ix: usize,
256 selected: bool,
257 _window: &mut Window,
258 _cx: &mut Context<Picker<Self>>,
259 ) -> Option<Self::ListItem> {
260 let tool_match = &self.matches[ix];
261 let tool = &self.tools[tool_match.candidate_id];
262
263 let is_enabled = match &tool.source {
264 ToolSource::Native => self.profile.tools.get(&tool.name).copied().unwrap_or(false),
265 ToolSource::ContextServer { id } => self
266 .profile
267 .context_servers
268 .get(id.as_ref())
269 .and_then(|preset| preset.tools.get(&tool.name))
270 .copied()
271 .unwrap_or(false),
272 };
273
274 Some(
275 ListItem::new(ix)
276 .inset(true)
277 .spacing(ListItemSpacing::Sparse)
278 .toggle_state(selected)
279 .child(
280 h_flex()
281 .gap_2()
282 .child(HighlightedLabel::new(
283 tool_match.string.clone(),
284 tool_match.positions.clone(),
285 ))
286 .map(|parent| match &tool.source {
287 ToolSource::Native => parent,
288 ToolSource::ContextServer { id } => parent
289 .child(Label::new(id).size(LabelSize::XSmall).color(Color::Muted)),
290 }),
291 )
292 .end_slot::<Icon>(is_enabled.then(|| {
293 Icon::new(IconName::Check)
294 .size(IconSize::Small)
295 .color(Color::Success)
296 })),
297 )
298 }
299}