tool_selector.rs

  1use std::sync::Arc;
  2
  3use assistant_settings::{AgentProfile, AssistantSettings};
  4use assistant_tool::{ToolSource, ToolWorkingSet};
  5use gpui::{Entity, Subscription};
  6use indexmap::IndexMap;
  7use settings::{Settings as _, SettingsStore};
  8use ui::{prelude::*, ContextMenu, PopoverMenu, Tooltip};
  9
 10pub struct ToolSelector {
 11    profiles: IndexMap<Arc<str>, AgentProfile>,
 12    tools: Arc<ToolWorkingSet>,
 13    _subscriptions: Vec<Subscription>,
 14}
 15
 16impl ToolSelector {
 17    pub fn new(tools: Arc<ToolWorkingSet>, cx: &mut Context<Self>) -> Self {
 18        let settings_subscription = cx.observe_global::<SettingsStore>(move |this, cx| {
 19            this.refresh_profiles(cx);
 20        });
 21
 22        let mut this = Self {
 23            profiles: IndexMap::default(),
 24            tools,
 25            _subscriptions: vec![settings_subscription],
 26        };
 27        this.refresh_profiles(cx);
 28
 29        this
 30    }
 31
 32    fn refresh_profiles(&mut self, cx: &mut Context<Self>) {
 33        let settings = AssistantSettings::get_global(cx);
 34
 35        self.profiles = settings.profiles.clone();
 36    }
 37
 38    fn build_context_menu(
 39        &self,
 40        window: &mut Window,
 41        cx: &mut Context<Self>,
 42    ) -> Entity<ContextMenu> {
 43        let profiles = self.profiles.clone();
 44        let tool_set = self.tools.clone();
 45        ContextMenu::build_persistent(window, cx, move |mut menu, _window, cx| {
 46            let icon_position = IconPosition::End;
 47
 48            menu = menu.header("Profiles");
 49            for (_id, profile) in profiles.clone() {
 50                menu = menu.toggleable_entry(profile.name.clone(), false, icon_position, None, {
 51                    let tools = tool_set.clone();
 52                    move |_window, cx| {
 53                        tools.disable_source(ToolSource::Native, cx);
 54                        tools.enable(
 55                            ToolSource::Native,
 56                            &profile
 57                                .tools
 58                                .iter()
 59                                .filter_map(|(tool, enabled)| enabled.then(|| tool.clone()))
 60                                .collect::<Vec<_>>(),
 61                        );
 62                    }
 63                });
 64            }
 65
 66            menu = menu.separator();
 67
 68            let tools_by_source = tool_set.tools_by_source(cx);
 69
 70            let all_tools_enabled = tool_set.are_all_tools_enabled();
 71            menu = menu.toggleable_entry("All Tools", all_tools_enabled, icon_position, None, {
 72                let tools = tool_set.clone();
 73                move |_window, cx| {
 74                    if all_tools_enabled {
 75                        tools.disable_all_tools(cx);
 76                    } else {
 77                        tools.enable_all_tools();
 78                    }
 79                }
 80            });
 81
 82            for (source, tools) in tools_by_source {
 83                let mut tools = tools
 84                    .into_iter()
 85                    .map(|tool| {
 86                        let source = tool.source();
 87                        let name = tool.name().into();
 88                        let is_enabled = tool_set.is_enabled(&source, &name);
 89
 90                        (source, name, is_enabled)
 91                    })
 92                    .collect::<Vec<_>>();
 93
 94                if ToolSource::Native == source {
 95                    tools.sort_by(|(_, name_a, _), (_, name_b, _)| name_a.cmp(name_b));
 96                }
 97
 98                menu = match &source {
 99                    ToolSource::Native => menu.separator().header("Zed Tools"),
100                    ToolSource::ContextServer { id } => {
101                        let all_tools_from_source_enabled =
102                            tool_set.are_all_tools_from_source_enabled(&source);
103
104                        menu.separator().header(id).toggleable_entry(
105                            "All Tools",
106                            all_tools_from_source_enabled,
107                            icon_position,
108                            None,
109                            {
110                                let tools = tool_set.clone();
111                                let source = source.clone();
112                                move |_window, cx| {
113                                    if all_tools_from_source_enabled {
114                                        tools.disable_source(source.clone(), cx);
115                                    } else {
116                                        tools.enable_source(&source);
117                                    }
118                                }
119                            },
120                        )
121                    }
122                };
123
124                for (source, name, is_enabled) in tools {
125                    menu = menu.toggleable_entry(name.clone(), is_enabled, icon_position, None, {
126                        let tools = tool_set.clone();
127                        move |_window, _cx| {
128                            if is_enabled {
129                                tools.disable(source.clone(), &[name.clone()]);
130                            } else {
131                                tools.enable(source.clone(), &[name.clone()]);
132                            }
133                        }
134                    });
135                }
136            }
137
138            menu
139        })
140    }
141}
142
143impl Render for ToolSelector {
144    fn render(&mut self, _window: &mut Window, cx: &mut Context<'_, Self>) -> impl IntoElement {
145        let this = cx.entity().clone();
146        PopoverMenu::new("tool-selector")
147            .menu(move |window, cx| {
148                Some(this.update(cx, |this, cx| this.build_context_menu(window, cx)))
149            })
150            .trigger_with_tooltip(
151                IconButton::new("tool-selector-button", IconName::SettingsAlt)
152                    .icon_size(IconSize::Small)
153                    .icon_color(Color::Muted),
154                Tooltip::text("Customize Tools"),
155            )
156            .anchor(gpui::Corner::BottomLeft)
157    }
158}