tool_selector.rs

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