Store profile per thread (#31907)

This allows storing the profile per thread, as well as moving the logic of which tools are enabled or not to the profile itself. This makes it much easier to switch between profiles, means there is less global state being changed on every profile change. Release Notes: - agent panel: allow saving the profile per thread --------- Co-authored-by: Ben Kunkle <ben.kunkle@gmail.com>
2025-06-06 14:05:27 +02:00 · 2025-06-06 14:05:27 +02:00 · 709523bf36
commit 709523bf36
parent 7afee64119
21 changed files with 556 additions and 369 deletions
--- a/crates/agent/src/agent_profile.rs
+++ b/crates/agent/src/agent_profile.rs
@ -0,0 +1,334 @@
+use std::sync::Arc;
+
+use agent_settings::{AgentProfileId, AgentProfileSettings, AgentSettings};
+use assistant_tool::{Tool, ToolSource, ToolWorkingSet};
+use collections::IndexMap;
+use convert_case::{Case, Casing};
+use fs::Fs;
+use gpui::{App, Entity};
+use settings::{Settings, update_settings_file};
+use ui::SharedString;
+use util::ResultExt;
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct AgentProfile {
+    id: AgentProfileId,
+    tool_set: Entity<ToolWorkingSet>,
+}
+
+pub type AvailableProfiles = IndexMap<AgentProfileId, SharedString>;
+
+impl AgentProfile {
+    pub fn new(id: AgentProfileId, tool_set: Entity<ToolWorkingSet>) -> Self {
+        Self { id, tool_set }
+    }
+
+    /// Saves a new profile to the settings.
+    pub fn create(
+        name: String,
+        base_profile_id: Option<AgentProfileId>,
+        fs: Arc<dyn Fs>,
+        cx: &App,
+    ) -> AgentProfileId {
+        let id = AgentProfileId(name.to_case(Case::Kebab).into());
+
+        let base_profile =
+            base_profile_id.and_then(|id| AgentSettings::get_global(cx).profiles.get(&id).cloned());
+
+        let profile_settings = AgentProfileSettings {
+            name: name.into(),
+            tools: base_profile
+                .as_ref()
+                .map(|profile| profile.tools.clone())
+                .unwrap_or_default(),
+            enable_all_context_servers: base_profile
+                .as_ref()
+                .map(|profile| profile.enable_all_context_servers)
+                .unwrap_or_default(),
+            context_servers: base_profile
+                .map(|profile| profile.context_servers)
+                .unwrap_or_default(),
+        };
+
+        update_settings_file::<AgentSettings>(fs, cx, {
+            let id = id.clone();
+            move |settings, _cx| {
+                settings.create_profile(id, profile_settings).log_err();
+            }
+        });
+
+        id
+    }
+
+    /// Returns a map of AgentProfileIds to their names
+    pub fn available_profiles(cx: &App) -> AvailableProfiles {
+        let mut profiles = AvailableProfiles::default();
+        for (id, profile) in AgentSettings::get_global(cx).profiles.iter() {
+            profiles.insert(id.clone(), profile.name.clone());
+        }
+        profiles
+    }
+
+    pub fn id(&self) -> &AgentProfileId {
+        &self.id
+    }
+
+    pub fn enabled_tools(&self, cx: &App) -> Vec<Arc<dyn Tool>> {
+        let Some(settings) = AgentSettings::get_global(cx).profiles.get(&self.id) else {
+            return Vec::new();
+        };
+
+        self.tool_set
+            .read(cx)
+            .tools(cx)
+            .into_iter()
+            .filter(|tool| Self::is_enabled(settings, tool.source(), tool.name()))
+            .collect()
+    }
+
+    fn is_enabled(settings: &AgentProfileSettings, source: ToolSource, name: String) -> bool {
+        match source {
+            ToolSource::Native => *settings.tools.get(name.as_str()).unwrap_or(&false),
+            ToolSource::ContextServer { id } => {
+                if settings.enable_all_context_servers {
+                    return true;
+                }
+
+                let Some(preset) = settings.context_servers.get(id.as_ref()) else {
+                    return false;
+                };
+                *preset.tools.get(name.as_str()).unwrap_or(&false)
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use agent_settings::ContextServerPreset;
+    use assistant_tool::ToolRegistry;
+    use collections::IndexMap;
+    use gpui::{AppContext, TestAppContext};
+    use http_client::FakeHttpClient;
+    use project::Project;
+    use settings::{Settings, SettingsStore};
+    use ui::SharedString;
+
+    use super::*;
+
+    #[gpui::test]
+    async fn test_enabled_built_in_tools_for_profile(cx: &mut TestAppContext) {
+        init_test_settings(cx);
+
+        let id = AgentProfileId::default();
+        let profile_settings = cx.read(|cx| {
+            AgentSettings::get_global(cx)
+                .profiles
+                .get(&id)
+                .unwrap()
+                .clone()
+        });
+        let tool_set = default_tool_set(cx);
+
+        let profile = AgentProfile::new(id.clone(), tool_set);
+
+        let mut enabled_tools = cx
+            .read(|cx| profile.enabled_tools(cx))
+            .into_iter()
+            .map(|tool| tool.name())
+            .collect::<Vec<_>>();
+        enabled_tools.sort();
+
+        let mut expected_tools = profile_settings
+            .tools
+            .into_iter()
+            .filter_map(|(tool, enabled)| enabled.then_some(tool.to_string()))
+            // Provider dependent
+            .filter(|tool| tool != "web_search")
+            .collect::<Vec<_>>();
+        // Plus all registered MCP tools
+        expected_tools.extend(["enabled_mcp_tool".into(), "disabled_mcp_tool".into()]);
+        expected_tools.sort();
+
+        assert_eq!(enabled_tools, expected_tools);
+    }
+
+    #[gpui::test]
+    async fn test_custom_mcp_settings(cx: &mut TestAppContext) {
+        init_test_settings(cx);
+
+        let id = AgentProfileId("custom_mcp".into());
+        let profile_settings = cx.read(|cx| {
+            AgentSettings::get_global(cx)
+                .profiles
+                .get(&id)
+                .unwrap()
+                .clone()
+        });
+        let tool_set = default_tool_set(cx);
+
+        let profile = AgentProfile::new(id.clone(), tool_set);
+
+        let mut enabled_tools = cx
+            .read(|cx| profile.enabled_tools(cx))
+            .into_iter()
+            .map(|tool| tool.name())
+            .collect::<Vec<_>>();
+        enabled_tools.sort();
+
+        let mut expected_tools = profile_settings.context_servers["mcp"]
+            .tools
+            .iter()
+            .filter_map(|(key, enabled)| enabled.then(|| key.to_string()))
+            .collect::<Vec<_>>();
+        expected_tools.sort();
+
+        assert_eq!(enabled_tools, expected_tools);
+    }
+
+    #[gpui::test]
+    async fn test_only_built_in(cx: &mut TestAppContext) {
+        init_test_settings(cx);
+
+        let id = AgentProfileId("write_minus_mcp".into());
+        let profile_settings = cx.read(|cx| {
+            AgentSettings::get_global(cx)
+                .profiles
+                .get(&id)
+                .unwrap()
+                .clone()
+        });
+        let tool_set = default_tool_set(cx);
+
+        let profile = AgentProfile::new(id.clone(), tool_set);
+
+        let mut enabled_tools = cx
+            .read(|cx| profile.enabled_tools(cx))
+            .into_iter()
+            .map(|tool| tool.name())
+            .collect::<Vec<_>>();
+        enabled_tools.sort();
+
+        let mut expected_tools = profile_settings
+            .tools
+            .into_iter()
+            .filter_map(|(tool, enabled)| enabled.then_some(tool.to_string()))
+            // Provider dependent
+            .filter(|tool| tool != "web_search")
+            .collect::<Vec<_>>();
+        expected_tools.sort();
+
+        assert_eq!(enabled_tools, expected_tools);
+    }
+
+    fn init_test_settings(cx: &mut TestAppContext) {
+        cx.update(|cx| {
+            let settings_store = SettingsStore::test(cx);
+            cx.set_global(settings_store);
+            Project::init_settings(cx);
+            AgentSettings::register(cx);
+            language_model::init_settings(cx);
+            ToolRegistry::default_global(cx);
+            assistant_tools::init(FakeHttpClient::with_404_response(), cx);
+        });
+
+        cx.update(|cx| {
+            let mut agent_settings = AgentSettings::get_global(cx).clone();
+            agent_settings.profiles.insert(
+                AgentProfileId("write_minus_mcp".into()),
+                AgentProfileSettings {
+                    name: "write_minus_mcp".into(),
+                    enable_all_context_servers: false,
+                    ..agent_settings.profiles[&AgentProfileId::default()].clone()
+                },
+            );
+            agent_settings.profiles.insert(
+                AgentProfileId("custom_mcp".into()),
+                AgentProfileSettings {
+                    name: "mcp".into(),
+                    tools: IndexMap::default(),
+                    enable_all_context_servers: false,
+                    context_servers: IndexMap::from_iter([("mcp".into(), context_server_preset())]),
+                },
+            );
+            AgentSettings::override_global(agent_settings, cx);
+        })
+    }
+
+    fn context_server_preset() -> ContextServerPreset {
+        ContextServerPreset {
+            tools: IndexMap::from_iter([
+                ("enabled_mcp_tool".into(), true),
+                ("disabled_mcp_tool".into(), false),
+            ]),
+        }
+    }
+
+    fn default_tool_set(cx: &mut TestAppContext) -> Entity<ToolWorkingSet> {
+        cx.new(|_| {
+            let mut tool_set = ToolWorkingSet::default();
+            tool_set.insert(Arc::new(FakeTool::new("enabled_mcp_tool", "mcp")));
+            tool_set.insert(Arc::new(FakeTool::new("disabled_mcp_tool", "mcp")));
+            tool_set
+        })
+    }
+
+    struct FakeTool {
+        name: String,
+        source: SharedString,
+    }
+
+    impl FakeTool {
+        fn new(name: impl Into<String>, source: impl Into<SharedString>) -> Self {
+            Self {
+                name: name.into(),
+                source: source.into(),
+            }
+        }
+    }
+
+    impl Tool for FakeTool {
+        fn name(&self) -> String {
+            self.name.clone()
+        }
+
+        fn source(&self) -> ToolSource {
+            ToolSource::ContextServer {
+                id: self.source.clone(),
+            }
+        }
+
+        fn description(&self) -> String {
+            unimplemented!()
+        }
+
+        fn icon(&self) -> ui::IconName {
+            unimplemented!()
+        }
+
+        fn needs_confirmation(&self, _input: &serde_json::Value, _cx: &App) -> bool {
+            unimplemented!()
+        }
+
+        fn ui_text(&self, _input: &serde_json::Value) -> String {
+            unimplemented!()
+        }
+
+        fn run(
+            self: Arc<Self>,
+            _input: serde_json::Value,
+            _request: Arc<language_model::LanguageModelRequest>,
+            _project: Entity<Project>,
+            _action_log: Entity<assistant_tool::ActionLog>,
+            _model: Arc<dyn language_model::LanguageModel>,
+            _window: Option<gpui::AnyWindowHandle>,
+            _cx: &mut App,
+        ) -> assistant_tool::ToolResult {
+            unimplemented!()
+        }
+
+        fn may_perform_edits(&self) -> bool {
+            unimplemented!()
+        }
+    }
+}