From f812ca12ffb4d3978c8f3f2f0f40d5bea70fc7e3 Mon Sep 17 00:00:00 2001
From: Yingjie He <yingjiehe@squareup.com>
Date: Tue, 8 Apr 2025 13:25:48 -0700
Subject: [PATCH] feat: remove permission when deleting extension (#2089)

---
 crates/goose-cli/src/commands/configure.rs |  6 +++-
 crates/goose/src/config/permission.rs      | 41 ++++++++++++++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/crates/goose-cli/src/commands/configure.rs b/crates/goose-cli/src/commands/configure.rs
index ce88ab01..ae1940cf 100644
--- a/crates/goose-cli/src/commands/configure.rs
+++ b/crates/goose-cli/src/commands/configure.rs
@@ -2,7 +2,9 @@ use cliclack::spinner;
 use console::style;
 use goose::agents::{extension::Envs, ExtensionConfig};
 use goose::config::extensions::name_to_key;
-use goose::config::{Config, ConfigError, ExperimentManager, ExtensionEntry, ExtensionManager};
+use goose::config::{
+    Config, ConfigError, ExperimentManager, ExtensionEntry, ExtensionManager, PermissionManager,
+};
 use goose::message::Message;
 use goose::providers::{create, providers};
 use mcp_core::tool::ToolAnnotations;
@@ -738,6 +740,8 @@ pub fn remove_extension_dialog() -> Result<(), Box<dyn Error>> {
 
     for name in selected {
         ExtensionManager::remove(&name_to_key(name))?;
+        let mut permission_manager = PermissionManager::default();
+        permission_manager.remove_extension(&name_to_key(name));
         cliclack::outro(format!("Removed {} extension", style(name).green()))?;
     }
 
diff --git a/crates/goose/src/config/permission.rs b/crates/goose/src/config/permission.rs
index 5b2b4cc6..ba259326 100644
--- a/crates/goose/src/config/permission.rs
+++ b/crates/goose/src/config/permission.rs
@@ -179,6 +179,25 @@ impl PermissionManager {
             .expect("Failed to serialize permission config");
         fs::write(&self.config_path, yaml_content).expect("Failed to write to permission.yaml");
     }
+
+    /// Removes all entries where the principal name starts with the given extension name.
+    pub fn remove_extension(&mut self, extension_name: &str) {
+        for permission_config in self.permission_map.values_mut() {
+            permission_config
+                .always_allow
+                .retain(|p| !p.starts_with(extension_name));
+            permission_config
+                .ask_before
+                .retain(|p| !p.starts_with(extension_name));
+            permission_config
+                .never_allow
+                .retain(|p| !p.starts_with(extension_name));
+        }
+
+        let yaml_content = serde_yaml::to_string(&self.permission_map)
+            .expect("Failed to serialize permission config");
+        fs::write(&self.config_path, yaml_content).expect("Failed to write to permission.yaml");
+    }
 }
 
 #[cfg(test)]
@@ -273,4 +292,26 @@ mod tests {
         assert!(!config.ask_before.contains(&"tool7".to_string()));
         assert!(config.never_allow.contains(&"tool7".to_string()));
     }
+
+    #[test]
+    fn test_remove_extension() {
+        let mut manager = create_test_permission_manager();
+        manager.update_user_permission("prefix__tool1", PermissionLevel::AlwaysAllow);
+        manager.update_user_permission("nonprefix__tool2", PermissionLevel::AlwaysAllow);
+        manager.update_user_permission("prefix__tool3", PermissionLevel::AskBefore);
+
+        // Remove entries starting with "prefix"
+        manager.remove_extension("prefix");
+
+        let config = manager.permission_map.get(USER_PERMISSION).unwrap();
+
+        // Verify entries with "prefix" are removed
+        assert!(!config.always_allow.contains(&"prefix__tool1".to_string()));
+        assert!(!config.ask_before.contains(&"prefix__tool3".to_string()));
+
+        // Verify other entries remain
+        assert!(config
+            .always_allow
+            .contains(&"nonprefix__tool2".to_string()));
+    }
 }