tool_selector.rs

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