tenant.openagents/omega

1use std::{collections::BTreeMap, sync::Arc};
2
3use agent::ContextServerRegistry;
4use agent_settings::{AgentProfileId, AgentProfileSettings};
5use fs::Fs;
6use gpui::{App, Context, DismissEvent, Entity, EventEmitter, Focusable, Task, WeakEntity, Window};
7use picker::{Picker, PickerDelegate};
8use settings::{AgentProfileContent, ContextServerPresetContent, update_settings_file};
9use ui::{ListItem, ListItemSpacing, prelude::*};
10use util::ResultExt as _;
11
12pub struct ToolPicker {
13    picker: Entity<Picker<ToolPickerDelegate>>,
14}
15
16#[derive(Clone, Copy, Debug, PartialEq)]
17enum ToolPickerMode {
18    BuiltinTools,
19    McpTools,
20}
21
22impl ToolPicker {
23    pub fn builtin_tools(
24        delegate: ToolPickerDelegate,
25        window: &mut Window,
26        cx: &mut Context<Self>,
27    ) -> Self {
28        let picker = cx.new(|cx| Picker::uniform_list(delegate, window, cx).embedded());
29        Self { picker }
30    }
31
32    pub fn mcp_tools(
33        delegate: ToolPickerDelegate,
34        window: &mut Window,
35        cx: &mut Context<Self>,
36    ) -> Self {
37        let picker = cx.new(|cx| Picker::list(delegate, window, cx).embedded());
38        Self { picker }
39    }
40}
41
42impl EventEmitter<DismissEvent> for ToolPicker {}
43
44impl Focusable for ToolPicker {
45    fn focus_handle(&self, cx: &App) -> gpui::FocusHandle {
46        self.picker.focus_handle(cx)
47    }
48}
49
50impl Render for ToolPicker {
51    fn render(&mut self, _window: &mut Window, _cx: &mut Context<Self>) -> impl IntoElement {
52        v_flex().child(self.picker.clone())
53    }
54}
55
56#[derive(Debug, Clone)]
57pub enum PickerItem {
58    Tool {
59        server_id: Option<Arc<str>>,
60        name: Arc<str>,
61    },
62    ContextServer {
63        server_id: Arc<str>,
64    },
65}
66
67pub struct ToolPickerDelegate {
68    tool_picker: WeakEntity<ToolPicker>,
69    fs: Arc<dyn Fs>,
70    items: Arc<Vec<PickerItem>>,
71    profile_id: AgentProfileId,
72    profile_settings: AgentProfileSettings,
73    filtered_items: Vec<PickerItem>,
74    selected_index: usize,
75    mode: ToolPickerMode,
76}
77
78impl ToolPickerDelegate {
79    pub fn builtin_tools(
80        tool_names: Vec<Arc<str>>,
81        fs: Arc<dyn Fs>,
82        profile_id: AgentProfileId,
83        profile_settings: AgentProfileSettings,
84        cx: &mut Context<ToolPicker>,
85    ) -> Self {
86        Self::new(
87            Arc::new(
88                tool_names
89                    .into_iter()
90                    .map(|name| PickerItem::Tool {
91                        name,
92                        server_id: None,
93                    })
94                    .collect(),
95            ),
96            ToolPickerMode::BuiltinTools,
97            fs,
98            profile_id,
99            profile_settings,
100            cx,
101        )
102    }
103
104    pub fn mcp_tools(
105        registry: &Entity<ContextServerRegistry>,
106        fs: Arc<dyn Fs>,
107        profile_id: AgentProfileId,
108        profile_settings: AgentProfileSettings,
109        cx: &mut Context<ToolPicker>,
110    ) -> Self {
111        let mut items = Vec::new();
112
113        for (id, tools) in registry.read(cx).servers() {
114            let server_id = id.clone().0;
115            items.push(PickerItem::ContextServer {
116                server_id: server_id.clone(),
117            });
118            items.extend(tools.keys().map(|tool_name| PickerItem::Tool {
119                name: tool_name.clone().into(),
120                server_id: Some(server_id.clone()),
121            }));
122        }
123
124        Self::new(
125            Arc::new(items),
126            ToolPickerMode::McpTools,
127            fs,
128            profile_id,
129            profile_settings,
130            cx,
131        )
132    }
133
134    fn new(
135        items: Arc<Vec<PickerItem>>,
136        mode: ToolPickerMode,
137        fs: Arc<dyn Fs>,
138        profile_id: AgentProfileId,
139        profile_settings: AgentProfileSettings,
140        cx: &mut Context<ToolPicker>,
141    ) -> Self {
142        Self {
143            tool_picker: cx.entity().downgrade(),
144            mode,
145            fs,
146            items,
147            profile_id,
148            profile_settings,
149            filtered_items: Vec::new(),
150            selected_index: 0,
151        }
152    }
153}
154
155impl PickerDelegate for ToolPickerDelegate {
156    type ListItem = AnyElement;
157
158    fn name() -> &'static str {
159        "tool picker"
160    }
161
162    fn match_count(&self) -> usize {
163        self.filtered_items.len()
164    }
165
166    fn selected_index(&self) -> usize {
167        self.selected_index
168    }
169
170    fn set_selected_index(
171        &mut self,
172        ix: usize,
173        _window: &mut Window,
174        _cx: &mut Context<Picker<Self>>,
175    ) {
176        self.selected_index = ix;
177    }
178
179    fn can_select(&self, ix: usize, _window: &mut Window, _cx: &mut Context<Picker<Self>>) -> bool {
180        matches!(self.filtered_items.get(ix), Some(PickerItem::Tool { .. }))
181    }
182
183    fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
184        match self.mode {
185            ToolPickerMode::BuiltinTools => "Search built-in tools…",
186            ToolPickerMode::McpTools => "Search MCP tools…",
187        }
188        .into()
189    }
190
191    fn update_matches(
192        &mut self,
193        query: String,
194        window: &mut Window,
195        cx: &mut Context<Picker<Self>>,
196    ) -> Task<()> {
197        let all_items = self.items.clone();
198
199        cx.spawn_in(window, async move |this, cx| {
200            let filtered_items = cx
201                .background_spawn(async move {
202                    let mut tools_by_provider: BTreeMap<Option<Arc<str>>, Vec<Arc<str>>> =
203                        BTreeMap::default();
204
205                    for item in all_items.iter() {
206                        if let PickerItem::Tool { server_id, name } = item.clone()
207                            && name.contains(&query)
208                        {
209                            tools_by_provider.entry(server_id).or_default().push(name);
210                        }
211                    }
212
213                    let mut items = Vec::new();
214
215                    for (server_id, names) in tools_by_provider {
216                        if let Some(server_id) = server_id.clone() {
217                            items.push(PickerItem::ContextServer { server_id });
218                        }
219                        for name in names {
220                            items.push(PickerItem::Tool {
221                                server_id: server_id.clone(),
222                                name,
223                            });
224                        }
225                    }
226
227                    items
228                })
229                .await;
230
231            this.update(cx, |this, _cx| {
232                this.delegate.filtered_items = filtered_items;
233                this.delegate.selected_index = this
234                    .delegate
235                    .selected_index
236                    .min(this.delegate.filtered_items.len().saturating_sub(1));
237            })
238            .log_err();
239        })
240    }
241
242    fn confirm(&mut self, _secondary: bool, window: &mut Window, cx: &mut Context<Picker<Self>>) {
243        if self.filtered_items.is_empty() {
244            self.dismissed(window, cx);
245            return;
246        }
247
248        let Some(item) = self.filtered_items.get(self.selected_index) else {
249            return;
250        };
251
252        let PickerItem::Tool {
253            name: tool_name,
254            server_id,
255        } = item
256        else {
257            return;
258        };
259
260        let is_currently_enabled = if let Some(server_id) = server_id.clone() {
261            let preset = self
262                .profile_settings
263                .context_servers
264                .entry(server_id)
265                .or_default();
266            let is_enabled = *preset.tools.entry(tool_name.clone()).or_default();
267            *preset.tools.entry(tool_name.clone()).or_default() = !is_enabled;
268            is_enabled
269        } else {
270            let is_enabled = *self
271                .profile_settings
272                .tools
273                .entry(tool_name.clone())
274                .or_default();
275            *self
276                .profile_settings
277                .tools
278                .entry(tool_name.clone())
279                .or_default() = !is_enabled;
280            is_enabled
281        };
282
283        update_settings_file(self.fs.clone(), cx, {
284            let profile_id = self.profile_id.clone();
285            let default_profile = self.profile_settings.clone();
286            let server_id = server_id.clone();
287            let tool_name = tool_name.clone();
288            move |settings, _cx| {
289                let profiles = settings
290                    .agent
291                    .get_or_insert_default()
292                    .profiles
293                    .get_or_insert_default();
294                let profile = profiles
295                    .entry(profile_id.0)
296                    .or_insert_with(|| AgentProfileContent {
297                        name: default_profile.name.into(),
298                        tools: default_profile.tools,
299                        enable_all_context_servers: Some(
300                            default_profile.enable_all_context_servers,
301                        ),
302                        context_servers: default_profile
303                            .context_servers
304                            .into_iter()
305                            .map(|(server_id, preset)| {
306                                (
307                                    server_id,
308                                    ContextServerPresetContent {
309                                        tools: preset.tools,
310                                    },
311                                )
312                            })
313                            .collect(),
314                        default_model: default_profile.default_model.clone(),
315                    });
316
317                if let Some(server_id) = server_id {
318                    let preset = profile.context_servers.entry(server_id).or_default();
319                    *preset.tools.entry(tool_name).or_default() = !is_currently_enabled;
320                } else {
321                    *profile.tools.entry(tool_name).or_default() = !is_currently_enabled;
322                }
323            }
324        });
325    }
326
327    fn dismissed(&mut self, _window: &mut Window, cx: &mut Context<Picker<Self>>) {
328        self.tool_picker
329            .update(cx, |_this, cx| cx.emit(DismissEvent))
330            .log_err();
331    }
332
333    fn render_match(
334        &self,
335        ix: usize,
336        selected: bool,
337        _window: &mut Window,
338        cx: &mut Context<Picker<Self>>,
339    ) -> Option<Self::ListItem> {
340        let item = &self.filtered_items.get(ix)?;
341        match item {
342            PickerItem::ContextServer { server_id, .. } => Some(
343                div()
344                    .px_2()
345                    .pb_1()
346                    .when(ix > 1, |this| {
347                        this.mt_1()
348                            .pt_2()
349                            .border_t_1()
350                            .border_color(cx.theme().colors().border_variant)
351                    })
352                    .child(
353                        Label::new(server_id)
354                            .size(LabelSize::XSmall)
355                            .color(Color::Muted),
356                    )
357                    .into_any_element(),
358            ),
359            PickerItem::Tool { name, server_id } => {
360                let is_enabled = if let Some(server_id) = server_id {
361                    self.profile_settings
362                        .context_servers
363                        .get(server_id.as_ref())
364                        .and_then(|preset| preset.tools.get(name))
365                        .copied()
366                        .unwrap_or(self.profile_settings.enable_all_context_servers)
367                } else {
368                    self.profile_settings
369                        .tools
370                        .get(name)
371                        .copied()
372                        .unwrap_or(false)
373                };
374
375                Some(
376                    ListItem::new(ix)
377                        .inset(true)
378                        .spacing(ListItemSpacing::Sparse)
379                        .toggle_state(selected)
380                        .child(Label::new(name.clone()))
381                        .end_slot::<Icon>(is_enabled.then(|| {
382                            Icon::new(IconName::Check)
383                                .size(IconSize::Small)
384                                .color(Color::Success)
385                        }))
386                        .into_any_element(),
387                )
388            }
389        }
390    }
391}
392