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_all_tools(cx);
 54
 55                        tools.enable(
 56                            ToolSource::Native,
 57                            &profile
 58                                .tools
 59                                .iter()
 60                                .filter_map(|(tool, enabled)| enabled.then(|| tool.clone()))
 61                                .collect::<Vec<_>>(),
 62                        );
 63
 64                        for (context_server_id, preset) in &profile.context_servers {
 65                            tools.enable(
 66                                ToolSource::ContextServer {
 67                                    id: context_server_id.clone().into(),
 68                                },
 69                                &preset
 70                                    .tools
 71                                    .iter()
 72                                    .filter_map(|(tool, enabled)| enabled.then(|| tool.clone()))
 73                                    .collect::<Vec<_>>(),
 74                            )
 75                        }
 76                    }
 77                });
 78            }
 79
 80            menu = menu.separator();
 81
 82            let tools_by_source = tool_set.tools_by_source(cx);
 83
 84            let all_tools_enabled = tool_set.are_all_tools_enabled();
 85            menu = menu.toggleable_entry("All Tools", all_tools_enabled, icon_position, None, {
 86                let tools = tool_set.clone();
 87                move |_window, cx| {
 88                    if all_tools_enabled {
 89                        tools.disable_all_tools(cx);
 90                    } else {
 91                        tools.enable_all_tools();
 92                    }
 93                }
 94            });
 95
 96            for (source, tools) in tools_by_source {
 97                let mut tools = tools
 98                    .into_iter()
 99                    .map(|tool| {
100                        let source = tool.source();
101                        let name = tool.name().into();
102                        let is_enabled = tool_set.is_enabled(&source, &name);
103
104                        (source, name, is_enabled)
105                    })
106                    .collect::<Vec<_>>();
107
108                if ToolSource::Native == source {
109                    tools.sort_by(|(_, name_a, _), (_, name_b, _)| name_a.cmp(name_b));
110                }
111
112                menu = match &source {
113                    ToolSource::Native => menu.separator().header("Zed Tools"),
114                    ToolSource::ContextServer { id } => {
115                        let all_tools_from_source_enabled =
116                            tool_set.are_all_tools_from_source_enabled(&source);
117
118                        menu.separator().header(id).toggleable_entry(
119                            "All Tools",
120                            all_tools_from_source_enabled,
121                            icon_position,
122                            None,
123                            {
124                                let tools = tool_set.clone();
125                                let source = source.clone();
126                                move |_window, cx| {
127                                    if all_tools_from_source_enabled {
128                                        tools.disable_source(source.clone(), cx);
129                                    } else {
130                                        tools.enable_source(&source);
131                                    }
132                                }
133                            },
134                        )
135                    }
136                };
137
138                for (source, name, is_enabled) in tools {
139                    menu = menu.toggleable_entry(name.clone(), is_enabled, icon_position, None, {
140                        let tools = tool_set.clone();
141                        move |_window, _cx| {
142                            if is_enabled {
143                                tools.disable(source.clone(), &[name.clone()]);
144                            } else {
145                                tools.enable(source.clone(), &[name.clone()]);
146                            }
147                        }
148                    });
149                }
150            }
151
152            menu
153        })
154    }
155}
156
157impl Render for ToolSelector {
158    fn render(&mut self, _window: &mut Window, cx: &mut Context<'_, Self>) -> impl IntoElement {
159        let this = cx.entity().clone();
160        PopoverMenu::new("tool-selector")
161            .menu(move |window, cx| {
162                Some(this.update(cx, |this, cx| this.build_context_menu(window, cx)))
163            })
164            .trigger_with_tooltip(
165                IconButton::new("tool-selector-button", IconName::SettingsAlt)
166                    .icon_size(IconSize::Small)
167                    .icon_color(Color::Muted),
168                Tooltip::text("Customize Tools"),
169            )
170            .anchor(gpui::Corner::BottomLeft)
171    }
172}