tool_selector.rs

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