assistant2: Allow profiles to manage context server tools (#27452)

Marshall Bowers created 9 months ago

This PR updates the agent profiles with support for managing context
server tools.

Release Notes:

- N/A

Change summary

crates/assistant2/src/tool_selector.rs              | 16 +++++++++++
crates/assistant_settings/src/agent_profile.rs      |  2 -
crates/assistant_settings/src/assistant_settings.rs | 20 ++++++++++++++
3 files changed, 34 insertions(+), 4 deletions(-)

Detailed changes

crates/assistant2/src/tool_selector.rs

@@ -50,7 +50,8 @@ impl ToolSelector {
                 menu = menu.toggleable_entry(profile.name.clone(), false, icon_position, None, {
                     let tools = tool_set.clone();
                     move |_window, cx| {
-                        tools.disable_source(ToolSource::Native, cx);
+                        tools.disable_all_tools(cx);
+
                         tools.enable(
                             ToolSource::Native,
                             &profile
@@ -59,6 +60,19 @@ impl ToolSelector {
                                 .filter_map(|(tool, enabled)| enabled.then(|| tool.clone()))
                                 .collect::<Vec<_>>(),
                         );
+
+                        for (context_server_id, preset) in &profile.context_servers {
+                            tools.enable(
+                                ToolSource::ContextServer {
+                                    id: context_server_id.clone().into(),
+                                },
+                                &preset
+                                    .tools
+                                    .iter()
+                                    .filter_map(|(tool, enabled)| enabled.then(|| tool.clone()))
+                                    .collect::<Vec<_>>(),
+                            )
+                        }
                     }
                 });
             }

crates/assistant_settings/src/agent_profile.rs

@@ -9,12 +9,10 @@ pub struct AgentProfile {
     /// The name of the profile.
     pub name: SharedString,
     pub tools: IndexMap<Arc<str>, bool>,
-    #[allow(dead_code)]
     pub context_servers: IndexMap<Arc<str>, ContextServerPreset>,
 }
 
 #[derive(Debug, Clone)]
 pub struct ContextServerPreset {
-    #[allow(dead_code)]
     pub tools: IndexMap<Arc<str>, bool>,
 }

crates/assistant_settings/src/assistant_settings.rs

@@ -412,6 +412,13 @@ impl Default for LanguageModelSelection {
 pub struct AgentProfileContent {
     pub name: Arc<str>,
     pub tools: IndexMap<Arc<str>, bool>,
+    #[serde(default)]
+    pub context_servers: IndexMap<Arc<str>, ContextServerPresetContent>,
+}
+
+#[derive(Debug, PartialEq, Clone, Serialize, Deserialize, JsonSchema)]
+pub struct ContextServerPresetContent {
+    pub tools: IndexMap<Arc<str>, bool>,
 }
 
 #[derive(Clone, Serialize, Deserialize, JsonSchema, Debug)]
@@ -522,7 +529,18 @@ impl Settings for AssistantSettings {
                             AgentProfile {
                                 name: profile.name.into(),
                                 tools: profile.tools,
-                                context_servers: IndexMap::default(),
+                                context_servers: profile
+                                    .context_servers
+                                    .into_iter()
+                                    .map(|(context_server_id, preset)| {
+                                        (
+                                            context_server_id,
+                                            ContextServerPreset {
+                                                tools: preset.tools.clone(),
+                                            },
+                                        )
+                                    })
+                                    .collect(),
                             },
                         )
                     }));