From ffdde20410352a6ad390ee7b61040a7d00c77b3b Mon Sep 17 00:00:00 2001 From: Mic Neale Date: Mon, 23 Dec 2024 16:41:20 +1100 Subject: [PATCH 1/4] first pass at layering in trust checks --- crates/goose/src/developer.rs | 31 ++++++ crates/goose/src/lib.rs | 1 + crates/goose/src/trust_risk.rs | 174 +++++++++++++++++++++++++++++++++ 3 files changed, 206 insertions(+) create mode 100644 crates/goose/src/trust_risk.rs diff --git a/crates/goose/src/developer.rs b/crates/goose/src/developer.rs index 3e0cdca56..61a6881be 100644 --- a/crates/goose/src/developer.rs +++ b/crates/goose/src/developer.rs @@ -16,6 +16,7 @@ use xcap::{Monitor, Window}; use crate::errors::{AgentError, AgentResult}; use crate::systems::System; +use crate::trust_risk::{TrustLevel, TrustManager}; use mcp_core::{Content, Resource, Role, Tool, ToolCall}; pub struct DeveloperSystem { @@ -24,6 +25,7 @@ pub struct DeveloperSystem { active_resources: Mutex>, // Use URI string as key instead of PathBuf file_history: Mutex>>, instructions: String, + trust_manager: TrustManager, } impl Default for DeveloperSystem { @@ -33,6 +35,17 @@ impl Default for DeveloperSystem { } impl DeveloperSystem { + /// Changes the trust level for operations + pub fn set_trust_level(&self, level: TrustLevel) { + self.trust_manager.set_level(level); + } + + /// Gets the current trust level + pub fn get_trust_level(&self) -> TrustLevel { + self.trust_manager.get_level() + } + + // Reads a resource from a URI and returns its content. // The resource must already exist in active_resources. pub async fn read_resource(&self, uri: &str) -> AgentResult { @@ -121,6 +134,7 @@ impl DeveloperSystem { } pub fn new() -> Self { + let trust_manager = TrustManager::new(); let list_windows_tool = Tool::new( "list_windows", indoc! {r#" @@ -279,6 +293,7 @@ impl DeveloperSystem { }, file_history: Mutex::new(HashMap::new()), instructions, + trust_manager, } } @@ -305,6 +320,13 @@ impl DeveloperSystem { .ok_or(AgentError::InvalidParameters( "The command string is required".into(), ))?; + + // Only block destructive operations at level 0 + if self.trust_manager.is_destructive_command(command) && self.trust_manager.get_level() == TrustLevel::NoDestructive { + return Err(AgentError::ExecutionError( + "Do not run any destructive or editing commands. Can run commands but not if they make file changes or system changes.".into() + )); + } // Disallow commands that should use other tools if command.trim_start().starts_with("cat") { @@ -370,6 +392,13 @@ impl DeveloperSystem { .and_then(|v| v.as_str()) .ok_or_else(|| AgentError::InvalidParameters("Missing 'path' parameter".into()))?; + // Only block destructive operations at level 0 + if (command == "write" || command == "str_replace") && self.trust_manager.get_level() == TrustLevel::NoDestructive { + return Err(AgentError::ExecutionError( + "File modifications are not allowed at NoDestructive trust level (0)".into() + )); + } + let path = self.resolve_path(path_str)?; match command { @@ -903,6 +932,8 @@ mod tests { #[tokio::test] async fn test_text_editor_write_and_view_file() { let system = get_system().await; + // Set trust level to AllowAll for this test + system.set_trust_level(TrustLevel::AllowAll); let temp_dir = tempfile::tempdir().unwrap(); let file_path = temp_dir.path().join("test.txt"); diff --git a/crates/goose/src/lib.rs b/crates/goose/src/lib.rs index 394e6450c..e3224d1da 100644 --- a/crates/goose/src/lib.rs +++ b/crates/goose/src/lib.rs @@ -9,3 +9,4 @@ pub mod prompt_template; pub mod providers; pub mod systems; pub mod token_counter; +pub mod trust_risk; diff --git a/crates/goose/src/trust_risk.rs b/crates/goose/src/trust_risk.rs new file mode 100644 index 000000000..34d2b65a3 --- /dev/null +++ b/crates/goose/src/trust_risk.rs @@ -0,0 +1,174 @@ +use std::sync::RwLock; + +/// Represents different trust levels for performing operations +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum TrustLevel { + /// No destructive actions allowed (0) + NoDestructive, + /// Confirm with user before destructive actions (1) + ConfirmDestructive, + /// Allow all actions without confirmation (2) + AllowAll, +} + +impl Default for TrustLevel { + fn default() -> Self { + TrustLevel::ConfirmDestructive // Default to middle ground - confirm destructive actions + } +} + +impl From for TrustLevel { + fn from(value: u8) -> Self { + match value { + 0 => TrustLevel::NoDestructive, + 1 => TrustLevel::ConfirmDestructive, + _ => TrustLevel::AllowAll, + } + } +} + +impl From for u8 { + fn from(level: TrustLevel) -> Self { + match level { + TrustLevel::NoDestructive => 0, + TrustLevel::ConfirmDestructive => 1, + TrustLevel::AllowAll => 2, + } + } +} + +/// Manages trust level state and provides methods to check and modify trust settings +pub struct TrustManager { + level: RwLock, +} + +impl Default for TrustManager { + fn default() -> Self { + Self::new() + } +} + +impl TrustManager { + /// Creates a new TrustManager with default trust level + pub fn new() -> Self { + Self { + level: RwLock::new(TrustLevel::default()), + } + } + + /// Creates a new TrustManager with a specific trust level + pub fn with_level(level: TrustLevel) -> Self { + Self { + level: RwLock::new(level), + } + } + + /// Gets the current trust level + pub fn get_level(&self) -> TrustLevel { + *self.level.read().unwrap() + } + + /// Sets a new trust level + pub fn set_level(&self, new_level: TrustLevel) { + *self.level.write().unwrap() = new_level; + } + + /// Checks if a destructive action is allowed + /// Returns true if the action should proceed, false if it should be blocked + pub fn can_perform_destructive(&self) -> bool { + match self.get_level() { + TrustLevel::NoDestructive => false, + TrustLevel::AllowAll => true, + TrustLevel::ConfirmDestructive => { + // In a real implementation, this would interact with the user + // For now, we'll just return false to be safe + false + } + } + } + + /// Checks if a command is potentially destructive + /// This is a basic implementation that could be expanded + pub fn is_destructive_command(&self, command: &str) -> bool { + let command = command.trim().to_lowercase(); + + // List of destructive command prefixes + let destructive_prefixes = [ + "rm", "del", "remove", + "write", "overwrite", + "delete", + "drop", + ]; + + destructive_prefixes.iter().any(|prefix| command.starts_with(prefix)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_trust_level_conversion() { + assert_eq!(TrustLevel::from(0), TrustLevel::NoDestructive); + assert_eq!(TrustLevel::from(1), TrustLevel::ConfirmDestructive); + assert_eq!(TrustLevel::from(2), TrustLevel::AllowAll); + assert_eq!(TrustLevel::from(255), TrustLevel::AllowAll); // Any value > 1 is AllowAll + + assert_eq!(u8::from(TrustLevel::NoDestructive), 0); + assert_eq!(u8::from(TrustLevel::ConfirmDestructive), 1); + assert_eq!(u8::from(TrustLevel::AllowAll), 2); + } + + #[test] + fn test_trust_manager_defaults() { + let manager = TrustManager::new(); + assert_eq!(manager.get_level(), TrustLevel::ConfirmDestructive); + } + + #[test] + fn test_trust_level_changes() { + let manager = TrustManager::new(); + + manager.set_level(TrustLevel::NoDestructive); + assert_eq!(manager.get_level(), TrustLevel::NoDestructive); + + manager.set_level(TrustLevel::AllowAll); + assert_eq!(manager.get_level(), TrustLevel::AllowAll); + } + + #[test] + fn test_destructive_command_detection() { + let manager = TrustManager::new(); + + // Test destructive commands + assert!(manager.is_destructive_command("rm -rf /")); + assert!(manager.is_destructive_command("delete file.txt")); + assert!(manager.is_destructive_command("remove old_data")); + assert!(manager.is_destructive_command("write new content")); + + // Test non-destructive commands + assert!(!manager.is_destructive_command("ls")); + assert!(!manager.is_destructive_command("cd /")); + assert!(!manager.is_destructive_command("echo hello")); + assert!(!manager.is_destructive_command("pwd")); + } + + #[test] + fn test_destructive_action_permissions() { + let manager = TrustManager::new(); + + // Test NoDestructive level + manager.set_level(TrustLevel::NoDestructive); + assert!(!manager.can_perform_destructive()); + + // Test AllowAll level + manager.set_level(TrustLevel::AllowAll); + assert!(manager.can_perform_destructive()); + + // Test ConfirmDestructive level + manager.set_level(TrustLevel::ConfirmDestructive); + // Currently returns false as user interaction is not implemented + assert!(!manager.can_perform_destructive()); + } +} \ No newline at end of file From d78991d72c78984bf96491b694768371d8843d25 Mon Sep 17 00:00:00 2001 From: Mic Neale Date: Thu, 26 Dec 2024 12:52:58 +1100 Subject: [PATCH 2/4] fmt --- crates/goose/src/developer.rs | 13 ++++++++----- crates/goose/src/trust_risk.rs | 29 +++++++++++++++++------------ 2 files changed, 25 insertions(+), 17 deletions(-) diff --git a/crates/goose/src/developer.rs b/crates/goose/src/developer.rs index 61a6881be..1c958296b 100644 --- a/crates/goose/src/developer.rs +++ b/crates/goose/src/developer.rs @@ -45,7 +45,6 @@ impl DeveloperSystem { self.trust_manager.get_level() } - // Reads a resource from a URI and returns its content. // The resource must already exist in active_resources. pub async fn read_resource(&self, uri: &str) -> AgentResult { @@ -320,9 +319,11 @@ impl DeveloperSystem { .ok_or(AgentError::InvalidParameters( "The command string is required".into(), ))?; - + // Only block destructive operations at level 0 - if self.trust_manager.is_destructive_command(command) && self.trust_manager.get_level() == TrustLevel::NoDestructive { + if self.trust_manager.is_destructive_command(command) + && self.trust_manager.get_level() == TrustLevel::NoDestructive + { return Err(AgentError::ExecutionError( "Do not run any destructive or editing commands. Can run commands but not if they make file changes or system changes.".into() )); @@ -393,9 +394,11 @@ impl DeveloperSystem { .ok_or_else(|| AgentError::InvalidParameters("Missing 'path' parameter".into()))?; // Only block destructive operations at level 0 - if (command == "write" || command == "str_replace") && self.trust_manager.get_level() == TrustLevel::NoDestructive { + if (command == "write" || command == "str_replace") + && self.trust_manager.get_level() == TrustLevel::NoDestructive + { return Err(AgentError::ExecutionError( - "File modifications are not allowed at NoDestructive trust level (0)".into() + "File modifications are not allowed at NoDestructive trust level (0)".into(), )); } diff --git a/crates/goose/src/trust_risk.rs b/crates/goose/src/trust_risk.rs index 34d2b65a3..cd1aeeac5 100644 --- a/crates/goose/src/trust_risk.rs +++ b/crates/goose/src/trust_risk.rs @@ -91,16 +91,21 @@ impl TrustManager { /// This is a basic implementation that could be expanded pub fn is_destructive_command(&self, command: &str) -> bool { let command = command.trim().to_lowercase(); - + // List of destructive command prefixes let destructive_prefixes = [ - "rm", "del", "remove", - "write", "overwrite", + "rm", + "del", + "remove", + "write", + "overwrite", "delete", "drop", ]; - destructive_prefixes.iter().any(|prefix| command.starts_with(prefix)) + destructive_prefixes + .iter() + .any(|prefix| command.starts_with(prefix)) } } @@ -129,10 +134,10 @@ mod tests { #[test] fn test_trust_level_changes() { let manager = TrustManager::new(); - + manager.set_level(TrustLevel::NoDestructive); assert_eq!(manager.get_level(), TrustLevel::NoDestructive); - + manager.set_level(TrustLevel::AllowAll); assert_eq!(manager.get_level(), TrustLevel::AllowAll); } @@ -140,13 +145,13 @@ mod tests { #[test] fn test_destructive_command_detection() { let manager = TrustManager::new(); - + // Test destructive commands assert!(manager.is_destructive_command("rm -rf /")); assert!(manager.is_destructive_command("delete file.txt")); assert!(manager.is_destructive_command("remove old_data")); assert!(manager.is_destructive_command("write new content")); - + // Test non-destructive commands assert!(!manager.is_destructive_command("ls")); assert!(!manager.is_destructive_command("cd /")); @@ -157,18 +162,18 @@ mod tests { #[test] fn test_destructive_action_permissions() { let manager = TrustManager::new(); - + // Test NoDestructive level manager.set_level(TrustLevel::NoDestructive); assert!(!manager.can_perform_destructive()); - + // Test AllowAll level manager.set_level(TrustLevel::AllowAll); assert!(manager.can_perform_destructive()); - + // Test ConfirmDestructive level manager.set_level(TrustLevel::ConfirmDestructive); // Currently returns false as user interaction is not implemented assert!(!manager.can_perform_destructive()); } -} \ No newline at end of file +} From e9912a1a7e89fa0e0b2acb0240a3f891ba5ead6e Mon Sep 17 00:00:00 2001 From: Mic Neale Date: Mon, 30 Dec 2024 18:01:27 +1100 Subject: [PATCH 3/4] show how to set the trust level --- crates/goose-server/src/state.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/crates/goose-server/src/state.rs b/crates/goose-server/src/state.rs index fa430c316..1afc27b11 100644 --- a/crates/goose-server/src/state.rs +++ b/crates/goose-server/src/state.rs @@ -1,5 +1,6 @@ use anyhow::Result; use goose::providers::configs::GroqProviderConfig; +use goose::trust_risk::TrustLevel; use goose::{ agent::Agent, developer::DeveloperSystem, @@ -21,7 +22,9 @@ impl AppState { pub fn new(provider_config: ProviderConfig, secret_key: String) -> Result { let provider = factory::get_provider(provider_config.clone())?; let mut agent = Agent::new(provider); - agent.add_system(Box::new(DeveloperSystem::new())); + let developer_system = DeveloperSystem::new(); + //developer_system.set_trust_level(TrustLevel::NoDestructive); + agent.add_system(Box::new(developer_system)); // Add memory system only if GOOSE_SERVER__MEMORY is set to "true" if let Ok(memory_enabled) = env::var("GOOSE_SERVER__MEMORY") { From cece465bc0af1e36d4f8b22cc926aae77bef1b8d Mon Sep 17 00:00:00 2001 From: Mic Neale Date: Mon, 30 Dec 2024 18:13:29 +1100 Subject: [PATCH 4/4] not needed import --- crates/goose-server/src/state.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/goose-server/src/state.rs b/crates/goose-server/src/state.rs index 1afc27b11..ba49ffc28 100644 --- a/crates/goose-server/src/state.rs +++ b/crates/goose-server/src/state.rs @@ -1,6 +1,6 @@ use anyhow::Result; use goose::providers::configs::GroqProviderConfig; -use goose::trust_risk::TrustLevel; +//use goose::trust_risk::TrustLevel; use goose::{ agent::Agent, developer::DeveloperSystem,