From b43a03a2354e8acb0e7de2dd2c24cd1a45c15d02 Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Tue, 4 Mar 2025 13:09:04 -0800 Subject: [PATCH 01/14] chat tool response --- crates/goose/src/agents/truncate.rs | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/crates/goose/src/agents/truncate.rs b/crates/goose/src/agents/truncate.rs index 81f754f8c2a2..c43e74ac0d02 100644 --- a/crates/goose/src/agents/truncate.rs +++ b/crates/goose/src/agents/truncate.rs @@ -319,13 +319,16 @@ impl Agent for TruncateAgent { message_tool_response = message_tool_response.with_tool_response( request.id.clone(), Ok(vec![Content::text( - "The following tool call was skipped in Goose chat mode. \ - In chat mode, you cannot run tool calls, instead, you can \ - only provide a detailed plan to the user. Provide an \ - explanation of the proposed tool call as if it were a plan. \ - Only if the user asks, provide a short explanation to the \ - user that they could consider running the tool above on \ - their own or with a different goose mode." + "The tool call was skipped in Goose chat mode. \ + In this mode, tool executions are not allowed. Instead, provide a \ + detailed explanation of what the tool call would do, structured as a \ + plan for the user. Do not apologize for skipping the tool call. \ + **Example Plan:**\n \ + 1. **Identify Task Scope** - Determine the purpose and expected outcome.\n \ + 2. **Outline Execution Steps** - Break down how the tool would be used.\n \ + 3. **Describe Expected Output** - Explain what the user should expect.\n \ + 4. **Suggest Next Steps** - Provide alternatives for the user to proceed.\n \ + If needed, adjust the explanation based on user preferences or questions." )]), ); } From 89ec854d559d06d82c48cb0b5691954a85e3bc4a Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Tue, 4 Mar 2025 13:22:15 -0800 Subject: [PATCH 02/14] plan --- crates/goose/src/agents/truncate.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/crates/goose/src/agents/truncate.rs b/crates/goose/src/agents/truncate.rs index c43e74ac0d02..6bf33833d5d1 100644 --- a/crates/goose/src/agents/truncate.rs +++ b/crates/goose/src/agents/truncate.rs @@ -325,9 +325,7 @@ impl Agent for TruncateAgent { plan for the user. Do not apologize for skipping the tool call. \ **Example Plan:**\n \ 1. **Identify Task Scope** - Determine the purpose and expected outcome.\n \ - 2. **Outline Execution Steps** - Break down how the tool would be used.\n \ - 3. **Describe Expected Output** - Explain what the user should expect.\n \ - 4. **Suggest Next Steps** - Provide alternatives for the user to proceed.\n \ + 2. **Outline Steps** - Break down the steps.\n \ If needed, adjust the explanation based on user preferences or questions." )]), ); From 5c8d8306f9690d88d5c814ed9036f209d1ecf3c0 Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Tue, 4 Mar 2025 13:29:19 -0800 Subject: [PATCH 03/14] no apologies --- crates/goose/src/agents/truncate.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/goose/src/agents/truncate.rs b/crates/goose/src/agents/truncate.rs index 6bf33833d5d1..7dce56d0f6ec 100644 --- a/crates/goose/src/agents/truncate.rs +++ b/crates/goose/src/agents/truncate.rs @@ -322,7 +322,7 @@ impl Agent for TruncateAgent { "The tool call was skipped in Goose chat mode. \ In this mode, tool executions are not allowed. Instead, provide a \ detailed explanation of what the tool call would do, structured as a \ - plan for the user. Do not apologize for skipping the tool call. \ + plan for the user. DO NOT apologize for skipping the tool call. DO NOT say sorry. \ **Example Plan:**\n \ 1. **Identify Task Scope** - Determine the purpose and expected outcome.\n \ 2. **Outline Steps** - Break down the steps.\n \ From 8ea4de26c4f7c47d7df71d08f3a58e5c67642e1d Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Tue, 4 Mar 2025 16:27:49 -0800 Subject: [PATCH 04/14] check in --- Cargo.lock | 26 ++++++++++++ crates/goose-cli/src/session/completion.rs | 34 ++++++++++++++++ crates/goose/Cargo.toml | 3 ++ crates/goose/src/agents/mod.rs | 2 + crates/goose/src/agents/truncate.rs | 46 +++++++++++++++++----- 5 files changed, 102 insertions(+), 9 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5e53102c9c1c..6727102e4da0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -174,6 +174,12 @@ version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7d902e3d592a523def97af8f317b08ce16b7ab854c1985a0c671e6f15cebc236" +[[package]] +name = "arrayref" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb" + [[package]] name = "arrayvec" version = "0.7.6" @@ -887,6 +893,19 @@ version = "2.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6099cdc01846bc367c4e7dd630dc5966dccf36b652fae7a74e17b640411a91b2" +[[package]] +name = "blake3" +version = "1.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "675f87afced0413c9bb02843499dbbd3882a237645883f71a2b59644a6d2f753" +dependencies = [ + "arrayref", + "arrayvec", + "cc", + "cfg-if", + "constant_time_eq", +] + [[package]] name = "block-buffer" version = "0.10.4" @@ -1248,6 +1267,12 @@ dependencies = [ "tiny-keccak", ] +[[package]] +name = "constant_time_eq" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c74b8349d32d297c9134b8c88677813a227df8f779daa29bfc29c183fe3dca6" + [[package]] name = "content_inspector" version = "0.2.4" @@ -2114,6 +2139,7 @@ dependencies = [ "aws-smithy-types", "axum 0.7.9", "base64 0.21.7", + "blake3", "chrono", "criterion", "ctor", diff --git a/crates/goose-cli/src/session/completion.rs b/crates/goose-cli/src/session/completion.rs index 5230a70e552a..d21a2e823ab7 100644 --- a/crates/goose-cli/src/session/completion.rs +++ b/crates/goose-cli/src/session/completion.rs @@ -550,4 +550,38 @@ mod tests { .unwrap(); assert_eq!(candidates.len(), 0); } + + #[test] + fn test_complete_slash_commands_partial() { + let cache = create_test_cache(); + let completer = GooseCompleter::new(cache); + + // Test partial command + let (pos, candidates) = completer.complete_slash_commands("/p").unwrap(); + assert_eq!(pos, 0); + assert!(candidates.iter().any(|c| c.display == "/prompt")); + assert!(candidates.iter().any(|c| c.display == "/prompts")); + } + + #[test] + fn test_complete_prompt_flags_info() { + let cache = create_test_cache(); + let completer = GooseCompleter::new(cache); + + // Test --info flag completion + let (pos, candidates) = completer.complete_prompt_flags("/prompt test_prompt1 --").unwrap(); + assert_eq!(candidates.len(), 1); + assert_eq!(candidates[0].display, "--info"); + } + + #[test] + fn test_complete_argument_keys_required() { + let cache = create_test_cache(); + let completer = GooseCompleter::new(cache); + + // Test required argument completion + let (_, candidates) = completer.complete_argument_keys("/prompt test_prompt1 ").unwrap(); + assert_eq!(candidates.len(), 1); + assert_eq!(candidates[0].display, "required_arg="); + } } diff --git a/crates/goose/Cargo.toml b/crates/goose/Cargo.toml index 0e02205994cd..8702ad8ef0a7 100644 --- a/crates/goose/Cargo.toml +++ b/crates/goose/Cargo.toml @@ -69,6 +69,9 @@ aws-sdk-bedrockruntime = "1.72.0" # For GCP Vertex AI provider auth jsonwebtoken = "9.3.1" +# Added blake3 hashing library as a dependency +blake3 = "1.5" + [target.'cfg(target_os = "windows")'.dependencies] winapi = { version = "0.3", features = ["wincred"] } diff --git a/crates/goose/src/agents/mod.rs b/crates/goose/src/agents/mod.rs index 4aa0659c976b..e970ca620cdc 100644 --- a/crates/goose/src/agents/mod.rs +++ b/crates/goose/src/agents/mod.rs @@ -3,6 +3,7 @@ mod capabilities; pub mod extension; mod factory; mod permission_judge; +mod permission_store; mod reference; mod truncate; @@ -11,3 +12,4 @@ pub use capabilities::Capabilities; pub use extension::ExtensionConfig; pub use factory::{register_agent, AgentFactory}; pub use permission_judge::detect_read_only_tools; +pub use permission_store::ToolPermissionStore; diff --git a/crates/goose/src/agents/truncate.rs b/crates/goose/src/agents/truncate.rs index 7dce56d0f6ec..58515628fb28 100644 --- a/crates/goose/src/agents/truncate.rs +++ b/crates/goose/src/agents/truncate.rs @@ -12,6 +12,7 @@ use super::detect_read_only_tools; use super::Agent; use crate::agents::capabilities::Capabilities; use crate::agents::extension::{ExtensionConfig, ExtensionResult}; +use crate::agents::ToolPermissionStore; use crate::config::Config; use crate::config::ExperimentManager; use crate::message::{Message, ToolRequest}; @@ -265,19 +266,43 @@ impl Agent for TruncateAgent { match mode.as_str() { "approve" => { let mut read_only_tools = Vec::new(); - // Process each tool request sequentially with confirmation - if ExperimentManager::is_enabled("GOOSE_SMART_APPROVE")? { - read_only_tools = detect_read_only_tools(&capabilities, tool_requests.clone()).await; - } + let mut needs_confirmation = Vec::new(); + + // First check permissions for all tools + let store = ToolPermissionStore::load()?; for request in &tool_requests { + if let Ok(tool_call) = request.tool_call.clone() { + if let Some(allowed) = store.check_permission(request) { + if allowed { + let output = capabilities.dispatch_tool_call(tool_call).await; + message_tool_response = message_tool_response.with_tool_response( + request.id.clone(), + output, + ); + } else { + needs_confirmation.push(request.clone()); + } + } else { + needs_confirmation.push(request.clone()); + } + } + } + + // Only check read-only status for tools needing confirmation + if !needs_confirmation.is_empty() && ExperimentManager::is_enabled("GOOSE_SMART_APPROVE")? { + read_only_tools = detect_read_only_tools(&capabilities, needs_confirmation.clone()).await; + } + + // Process remaining tools that need confirmation + for request in &needs_confirmation { if let Ok(tool_call) = request.tool_call.clone() { // Skip confirmation if the tool_call.name is in the read_only_tools list if read_only_tools.contains(&tool_call.name) { let output = capabilities.dispatch_tool_call(tool_call).await; - message_tool_response = message_tool_response.with_tool_response( - request.id.clone(), - output, - ); + message_tool_response = message_tool_response.with_tool_response( + request.id.clone(), + output, + ); } else { let confirmation = Message::user().with_tool_confirmation_request( request.id.clone(), @@ -289,9 +314,12 @@ impl Agent for TruncateAgent { // Wait for confirmation response through the channel let mut rx = self.confirmation_rx.lock().await; - // Loop the recv until we have a matched req_id due to potential duplicate messages. while let Some((req_id, confirmed)) = rx.recv().await { if req_id == request.id { + // Store the user's response + let mut store = ToolPermissionStore::load()?; + store.record_permission(&request, confirmed, None)?; + if confirmed { // User approved - dispatch the tool call let output = capabilities.dispatch_tool_call(tool_call).await; From 389e171447ee80a92443d89a609ac3c7a6e9efda Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Tue, 4 Mar 2025 17:21:28 -0800 Subject: [PATCH 05/14] add permission store --- crates/goose/src/agents/permission_store.rs | 206 ++++++++++++++++++++ 1 file changed, 206 insertions(+) create mode 100644 crates/goose/src/agents/permission_store.rs diff --git a/crates/goose/src/agents/permission_store.rs b/crates/goose/src/agents/permission_store.rs new file mode 100644 index 000000000000..f8197b38de60 --- /dev/null +++ b/crates/goose/src/agents/permission_store.rs @@ -0,0 +1,206 @@ +use std::collections::HashMap; +use serde::{Serialize, Deserialize}; +use anyhow::Result; +use etcetera::{choose_app_strategy, AppStrategy, AppStrategyArgs}; +use std::{ + path::PathBuf, + fs::File, +}; +use crate::message::ToolRequest; +use chrono::Utc; +use std::time::Duration; +use blake3::Hasher; + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct ToolPermissionRecord { + tool_name: String, + allowed: bool, + context_hash: String, // Hash of the tool's arguments/context to differentiate similar calls + timestamp: i64, + expiry: Option, // Optional expiry timestamp +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct ToolPermissionStore { + permissions: HashMap>, + version: u32, // For future schema migrations + #[serde(skip)] // Don't serialize this field + config_dir: PathBuf, +} + +impl ToolPermissionStore { + pub fn new() -> Self { + let config_dir = choose_app_strategy(crate::config::APP_STRATEGY.clone()) + .map(|strategy| strategy.config_dir()) + .unwrap_or_else(|_| PathBuf::from(".config/goose")); + + Self { + permissions: HashMap::new(), + version: 1, + config_dir, + } + } + + pub fn load() -> Result { + let store = Self::new(); + let file_path = store.config_dir.join("tool_permissions.json"); + + if !file_path.exists() { + return Ok(store); + } + + let file = File::open(file_path)?; + let mut permissions: ToolPermissionStore = serde_json::from_reader(file)?; + permissions.config_dir = store.config_dir; + Ok(permissions) + } + + pub fn save(&self) -> anyhow::Result<()> { + let config_dir = self.config_dir.clone(); + std::fs::create_dir_all(&config_dir)?; + + let path = config_dir.join("tool_permissions.json"); + let content = serde_json::to_string_pretty(self)?; + std::fs::write(path, content)?; + Ok(()) + } + + pub fn check_permission(&self, tool_request: &ToolRequest) -> Option { + let context_hash = self.hash_tool_context(tool_request); + let tool_call = tool_request.tool_call.as_ref().unwrap(); + let key = format!("{}:{}", tool_call.name, context_hash); + + self.permissions.get(&key).and_then(|records| { + records + .iter() + .filter(|record| { + record.expiry.map_or(true, |exp| exp > Utc::now().timestamp()) + }) + .last() + .map(|record| record.allowed) + }) + } + + pub fn record_permission( + &mut self, + tool_request: &ToolRequest, + allowed: bool, + expiry_duration: Option + ) -> anyhow::Result<()> { + let context_hash = self.hash_tool_context(tool_request); + let tool_call = tool_request.tool_call.as_ref().unwrap(); + let key = format!("{}:{}", tool_call.name, context_hash); + + let record = ToolPermissionRecord { + tool_name: tool_call.name.clone(), + allowed, + context_hash, + timestamp: Utc::now().timestamp(), + expiry: expiry_duration.map(|d| Utc::now().timestamp() + d.as_secs() as i64), + }; + + self.permissions + .entry(key) + .or_insert_with(Vec::new) + .push(record); + + self.save()?; + Ok(()) + } + + fn hash_tool_context(&self, tool_request: &ToolRequest) -> String { + // Create a hash of the tool's arguments to differentiate similar calls + // This helps identify when the same tool is being used in a different context + let mut hasher = Hasher::new(); + hasher.update(serde_json::to_string(&tool_request.tool_call.as_ref().unwrap().arguments).unwrap_or_default().as_bytes()); + hasher.finalize().to_hex().to_string() + } +} + +#[cfg(test)] +mod tests { + use crate::message::ToolRequest; + use mcp_core::tool::ToolCall; + use crate::agents::permission_store::ToolPermissionStore; + use chrono::Utc; + use std::time::Duration; + + fn create_test_tool_request(name: &str, args: serde_json::Value) -> ToolRequest { + ToolRequest { + id: "test-id".to_string(), + tool_call: Ok(ToolCall { + name: name.to_string(), + arguments: args, + }), + } + } + + #[test] + fn test_permission_store_basic() { + let mut store = ToolPermissionStore::new(); + let tool_request = create_test_tool_request( + "test_tool", + serde_json::json!({"arg1": "value1"}), + ); + + // Initially no permission recorded + assert!(store.check_permission(&tool_request).is_none()); + + // Record a permission + store.record_permission(&tool_request, true, None).unwrap(); + + // Should now find the recorded permission + assert_eq!(store.check_permission(&tool_request), Some(true)); + } + + #[test] + fn test_permission_expiry() { + let mut store = ToolPermissionStore::new(); + let tool_request = create_test_tool_request( + "test_tool", + serde_json::json!({"arg1": "value1"}), + ); + + // Record a permission that expires in 1 second + store.record_permission(&tool_request, true, Some(Duration::from_secs(1))).unwrap(); + + // Should initially be allowed + assert_eq!(store.check_permission(&tool_request), Some(true)); + + // Manually set expiry to the past + if let Some(records) = store.permissions.get_mut(&format!("{}:{}", + tool_request.tool_call.as_ref().unwrap().name, + store.hash_tool_context(&tool_request))) { + if let Some(record) = records.last_mut() { + record.expiry = Some(Utc::now().timestamp() - 2); + } + } + + // Should now be expired (no permission found) + assert!(store.check_permission(&tool_request).is_none()); + } + + #[test] + fn test_different_arguments() { + let mut store = ToolPermissionStore::new(); + + // Create two requests with same tool but different args + let request1 = create_test_tool_request( + "test_tool", + serde_json::json!({"arg": "value1"}), + ); + let request2 = create_test_tool_request( + "test_tool", + serde_json::json!({"arg": "value2"}), + ); + + // Record permission for first request + store.record_permission(&request1, true, None).unwrap(); + + // Should only find permission for first request + assert_eq!(store.check_permission(&request1), Some(true)); + assert!(store.check_permission(&request2).is_none()); + } +} + + From d3e0faccaa8dc97f21ad264841b0a69af8b2cc1b Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Tue, 4 Mar 2025 17:22:06 -0800 Subject: [PATCH 06/14] rm completion tests --- crates/goose-cli/src/session/completion.rs | 35 ---------------------- 1 file changed, 35 deletions(-) diff --git a/crates/goose-cli/src/session/completion.rs b/crates/goose-cli/src/session/completion.rs index d21a2e823ab7..8779c83cf8dc 100644 --- a/crates/goose-cli/src/session/completion.rs +++ b/crates/goose-cli/src/session/completion.rs @@ -550,38 +550,3 @@ mod tests { .unwrap(); assert_eq!(candidates.len(), 0); } - - #[test] - fn test_complete_slash_commands_partial() { - let cache = create_test_cache(); - let completer = GooseCompleter::new(cache); - - // Test partial command - let (pos, candidates) = completer.complete_slash_commands("/p").unwrap(); - assert_eq!(pos, 0); - assert!(candidates.iter().any(|c| c.display == "/prompt")); - assert!(candidates.iter().any(|c| c.display == "/prompts")); - } - - #[test] - fn test_complete_prompt_flags_info() { - let cache = create_test_cache(); - let completer = GooseCompleter::new(cache); - - // Test --info flag completion - let (pos, candidates) = completer.complete_prompt_flags("/prompt test_prompt1 --").unwrap(); - assert_eq!(candidates.len(), 1); - assert_eq!(candidates[0].display, "--info"); - } - - #[test] - fn test_complete_argument_keys_required() { - let cache = create_test_cache(); - let completer = GooseCompleter::new(cache); - - // Test required argument completion - let (_, candidates) = completer.complete_argument_keys("/prompt test_prompt1 ").unwrap(); - assert_eq!(candidates.len(), 1); - assert_eq!(candidates[0].display, "required_arg="); - } -} From fb2a6b476065e0b16cc0b96aae02726b2ad31c56 Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Tue, 4 Mar 2025 17:22:17 -0800 Subject: [PATCH 07/14] rm completion tests --- crates/goose-cli/src/session/completion.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/goose-cli/src/session/completion.rs b/crates/goose-cli/src/session/completion.rs index 8779c83cf8dc..d6a70178be3f 100644 --- a/crates/goose-cli/src/session/completion.rs +++ b/crates/goose-cli/src/session/completion.rs @@ -550,3 +550,4 @@ mod tests { .unwrap(); assert_eq!(candidates.len(), 0); } +} \ No newline at end of file From af71cf9a70e9de1c06dbe577273572e4003d5d5c Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Thu, 6 Mar 2025 17:01:28 -0800 Subject: [PATCH 08/14] store impl --- crates/goose-cli/src/session/completion.rs | 2 +- crates/goose/src/agents/permission_store.rs | 108 ++++++++++---------- crates/goose/src/agents/truncate.rs | 8 +- crates/goose/src/message.rs | 16 +++ 4 files changed, 74 insertions(+), 60 deletions(-) diff --git a/crates/goose-cli/src/session/completion.rs b/crates/goose-cli/src/session/completion.rs index 5c666014530d..9346d499b701 100644 --- a/crates/goose-cli/src/session/completion.rs +++ b/crates/goose-cli/src/session/completion.rs @@ -595,4 +595,4 @@ mod tests { .unwrap(); assert_eq!(candidates.len(), 0); } -} \ No newline at end of file +} diff --git a/crates/goose/src/agents/permission_store.rs b/crates/goose/src/agents/permission_store.rs index f8197b38de60..b4e42ca9f9d5 100644 --- a/crates/goose/src/agents/permission_store.rs +++ b/crates/goose/src/agents/permission_store.rs @@ -1,49 +1,48 @@ -use std::collections::HashMap; -use serde::{Serialize, Deserialize}; -use anyhow::Result; -use etcetera::{choose_app_strategy, AppStrategy, AppStrategyArgs}; -use std::{ - path::PathBuf, - fs::File, -}; use crate::message::ToolRequest; +use anyhow::Result; +use blake3::Hasher; use chrono::Utc; +use etcetera::{choose_app_strategy, AppStrategy}; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; use std::time::Duration; -use blake3::Hasher; +use std::{fs::File, path::PathBuf}; #[derive(Debug, Serialize, Deserialize, Clone)] pub struct ToolPermissionRecord { tool_name: String, allowed: bool, - context_hash: String, // Hash of the tool's arguments/context to differentiate similar calls + context_hash: String, // Hash of the tool's arguments/context to differentiate similar calls + #[serde(skip_serializing_if = "Option::is_none")] // Don't serialize if None + readable_context: Option, // Add this field timestamp: i64, - expiry: Option, // Optional expiry timestamp + expiry: Option, // Optional expiry timestamp } #[derive(Debug, Serialize, Deserialize)] pub struct ToolPermissionStore { permissions: HashMap>, - version: u32, // For future schema migrations - #[serde(skip)] // Don't serialize this field - config_dir: PathBuf, + version: u32, // For future schema migrations + #[serde(skip)] // Don't serialize this field + permissions_dir: PathBuf, } impl ToolPermissionStore { pub fn new() -> Self { - let config_dir = choose_app_strategy(crate::config::APP_STRATEGY.clone()) + let permissions_dir = choose_app_strategy(crate::config::APP_STRATEGY.clone()) .map(|strategy| strategy.config_dir()) .unwrap_or_else(|_| PathBuf::from(".config/goose")); Self { permissions: HashMap::new(), version: 1, - config_dir, + permissions_dir, } } pub fn load() -> Result { let store = Self::new(); - let file_path = store.config_dir.join("tool_permissions.json"); + let file_path = store.permissions_dir.join("tool_permissions.json"); if !file_path.exists() { return Ok(store); @@ -51,15 +50,15 @@ impl ToolPermissionStore { let file = File::open(file_path)?; let mut permissions: ToolPermissionStore = serde_json::from_reader(file)?; - permissions.config_dir = store.config_dir; + permissions.permissions_dir = store.permissions_dir; Ok(permissions) } - + pub fn save(&self) -> anyhow::Result<()> { - let config_dir = self.config_dir.clone(); - std::fs::create_dir_all(&config_dir)?; - - let path = config_dir.join("tool_permissions.json"); + let permissions_dir = self.permissions_dir.clone(); + std::fs::create_dir_all(&permissions_dir)?; + + let path = permissions_dir.join("tool_permissions.json"); let content = serde_json::to_string_pretty(self)?; std::fs::write(path, content)?; Ok(()) @@ -69,12 +68,14 @@ impl ToolPermissionStore { let context_hash = self.hash_tool_context(tool_request); let tool_call = tool_request.tool_call.as_ref().unwrap(); let key = format!("{}:{}", tool_call.name, context_hash); - + self.permissions.get(&key).and_then(|records| { records .iter() .filter(|record| { - record.expiry.map_or(true, |exp| exp > Utc::now().timestamp()) + record + .expiry + .map_or(true, |exp| exp > Utc::now().timestamp()) }) .last() .map(|record| record.allowed) @@ -82,19 +83,20 @@ impl ToolPermissionStore { } pub fn record_permission( - &mut self, - tool_request: &ToolRequest, + &mut self, + tool_request: &ToolRequest, allowed: bool, - expiry_duration: Option + expiry_duration: Option, ) -> anyhow::Result<()> { let context_hash = self.hash_tool_context(tool_request); let tool_call = tool_request.tool_call.as_ref().unwrap(); let key = format!("{}:{}", tool_call.name, context_hash); - + let record = ToolPermissionRecord { tool_name: tool_call.name.clone(), allowed, context_hash, + readable_context: Some(tool_request.to_readable_string()), timestamp: Utc::now().timestamp(), expiry: expiry_duration.map(|d| Utc::now().timestamp() + d.as_secs() as i64), }; @@ -103,7 +105,7 @@ impl ToolPermissionStore { .entry(key) .or_insert_with(Vec::new) .push(record); - + self.save()?; Ok(()) } @@ -112,17 +114,21 @@ impl ToolPermissionStore { // Create a hash of the tool's arguments to differentiate similar calls // This helps identify when the same tool is being used in a different context let mut hasher = Hasher::new(); - hasher.update(serde_json::to_string(&tool_request.tool_call.as_ref().unwrap().arguments).unwrap_or_default().as_bytes()); + hasher.update( + serde_json::to_string(&tool_request.tool_call.as_ref().unwrap().arguments) + .unwrap_or_default() + .as_bytes(), + ); hasher.finalize().to_hex().to_string() } } #[cfg(test)] mod tests { - use crate::message::ToolRequest; - use mcp_core::tool::ToolCall; use crate::agents::permission_store::ToolPermissionStore; + use crate::message::ToolRequest; use chrono::Utc; + use mcp_core::tool::ToolCall; use std::time::Duration; fn create_test_tool_request(name: &str, args: serde_json::Value) -> ToolRequest { @@ -138,10 +144,8 @@ mod tests { #[test] fn test_permission_store_basic() { let mut store = ToolPermissionStore::new(); - let tool_request = create_test_tool_request( - "test_tool", - serde_json::json!({"arg1": "value1"}), - ); + let tool_request = + create_test_tool_request("test_tool", serde_json::json!({"arg1": "value1"})); // Initially no permission recorded assert!(store.check_permission(&tool_request).is_none()); @@ -156,21 +160,23 @@ mod tests { #[test] fn test_permission_expiry() { let mut store = ToolPermissionStore::new(); - let tool_request = create_test_tool_request( - "test_tool", - serde_json::json!({"arg1": "value1"}), - ); + let tool_request = + create_test_tool_request("test_tool", serde_json::json!({"arg1": "value1"})); // Record a permission that expires in 1 second - store.record_permission(&tool_request, true, Some(Duration::from_secs(1))).unwrap(); + store + .record_permission(&tool_request, true, Some(Duration::from_secs(1))) + .unwrap(); // Should initially be allowed assert_eq!(store.check_permission(&tool_request), Some(true)); // Manually set expiry to the past - if let Some(records) = store.permissions.get_mut(&format!("{}:{}", + if let Some(records) = store.permissions.get_mut(&format!( + "{}:{}", tool_request.tool_call.as_ref().unwrap().name, - store.hash_tool_context(&tool_request))) { + store.hash_tool_context(&tool_request) + )) { if let Some(record) = records.last_mut() { record.expiry = Some(Utc::now().timestamp() - 2); } @@ -183,16 +189,10 @@ mod tests { #[test] fn test_different_arguments() { let mut store = ToolPermissionStore::new(); - + // Create two requests with same tool but different args - let request1 = create_test_tool_request( - "test_tool", - serde_json::json!({"arg": "value1"}), - ); - let request2 = create_test_tool_request( - "test_tool", - serde_json::json!({"arg": "value2"}), - ); + let request1 = create_test_tool_request("test_tool", serde_json::json!({"arg": "value1"})); + let request2 = create_test_tool_request("test_tool", serde_json::json!({"arg": "value2"})); // Record permission for first request store.record_permission(&request1, true, None).unwrap(); @@ -202,5 +202,3 @@ mod tests { assert!(store.check_permission(&request2).is_none()); } } - - diff --git a/crates/goose/src/agents/truncate.rs b/crates/goose/src/agents/truncate.rs index cff7dde76e6d..52eefad125cb 100644 --- a/crates/goose/src/agents/truncate.rs +++ b/crates/goose/src/agents/truncate.rs @@ -272,7 +272,7 @@ impl Agent for TruncateAgent { let store = ToolPermissionStore::load()?; for request in &tool_requests { if let Ok(tool_call) = request.tool_call.clone() { - if let Some(allowed) = store.check_permission(request) { + if let Some(allowed) = store.check_permission(*request) { if allowed { let output = capabilities.dispatch_tool_call(tool_call).await; message_tool_response = message_tool_response.with_tool_response( @@ -280,17 +280,17 @@ impl Agent for TruncateAgent { output, ); } else { - needs_confirmation.push(request.clone()); + needs_confirmation.push((*request).clone()); } } else { - needs_confirmation.push(request.clone()); + needs_confirmation.push((*request).clone()); } } } // Only check read-only status for tools needing confirmation if !needs_confirmation.is_empty() && ExperimentManager::is_enabled("GOOSE_SMART_APPROVE")? { - read_only_tools = detect_read_only_tools(&capabilities, needs_confirmation.clone()).await; + read_only_tools = detect_read_only_tools(&capabilities, needs_confirmation.iter().collect()).await; } // Process remaining tools that need confirmation diff --git a/crates/goose/src/message.rs b/crates/goose/src/message.rs index 0e5781e7cca6..c370c174aa8a 100644 --- a/crates/goose/src/message.rs +++ b/crates/goose/src/message.rs @@ -26,6 +26,22 @@ pub struct ToolRequest { pub tool_call: ToolResult, } +impl ToolRequest { + pub fn to_readable_string(&self) -> String { + match &self.tool_call { + Ok(tool_call) => { + format!( + "Tool: {}, Args: {}", + tool_call.name, + serde_json::to_string_pretty(&tool_call.arguments) + .unwrap_or_else(|_| "<>".to_string()) + ) + } + Err(e) => format!("Invalid tool call: {}", e), + } + } +} + #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] #[serde(rename_all = "camelCase")] pub struct ToolResponse { From 0b1e77cd149c6d8e6d6916c17b1c11a0deaea22a Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Thu, 6 Mar 2025 17:16:35 -0800 Subject: [PATCH 09/14] fix dereferencing --- crates/goose/src/agents/permission_store.rs | 17 ++++++++--------- crates/goose/src/agents/truncate.rs | 10 +++++----- 2 files changed, 13 insertions(+), 14 deletions(-) diff --git a/crates/goose/src/agents/permission_store.rs b/crates/goose/src/agents/permission_store.rs index b4e42ca9f9d5..504fb50c07b9 100644 --- a/crates/goose/src/agents/permission_store.rs +++ b/crates/goose/src/agents/permission_store.rs @@ -27,6 +27,12 @@ pub struct ToolPermissionStore { permissions_dir: PathBuf, } +impl Default for ToolPermissionStore { + fn default() -> Self { + Self::new() + } +} + impl ToolPermissionStore { pub fn new() -> Self { let permissions_dir = choose_app_strategy(crate::config::APP_STRATEGY.clone()) @@ -72,11 +78,7 @@ impl ToolPermissionStore { self.permissions.get(&key).and_then(|records| { records .iter() - .filter(|record| { - record - .expiry - .map_or(true, |exp| exp > Utc::now().timestamp()) - }) + .filter(|record| record.expiry.is_none_or(|exp| exp > Utc::now().timestamp())) .last() .map(|record| record.allowed) }) @@ -101,10 +103,7 @@ impl ToolPermissionStore { expiry: expiry_duration.map(|d| Utc::now().timestamp() + d.as_secs() as i64), }; - self.permissions - .entry(key) - .or_insert_with(Vec::new) - .push(record); + self.permissions.entry(key).or_default().push(record); self.save()?; Ok(()) diff --git a/crates/goose/src/agents/truncate.rs b/crates/goose/src/agents/truncate.rs index 52eefad125cb..dc7996015b0c 100644 --- a/crates/goose/src/agents/truncate.rs +++ b/crates/goose/src/agents/truncate.rs @@ -272,7 +272,7 @@ impl Agent for TruncateAgent { let store = ToolPermissionStore::load()?; for request in &tool_requests { if let Ok(tool_call) = request.tool_call.clone() { - if let Some(allowed) = store.check_permission(*request) { + if let Some(allowed) = store.check_permission(request) { if allowed { let output = capabilities.dispatch_tool_call(tool_call).await; message_tool_response = message_tool_response.with_tool_response( @@ -280,17 +280,17 @@ impl Agent for TruncateAgent { output, ); } else { - needs_confirmation.push((*request).clone()); + needs_confirmation.push(request.clone()); } } else { - needs_confirmation.push((*request).clone()); + needs_confirmation.push(request.clone()); } } } // Only check read-only status for tools needing confirmation if !needs_confirmation.is_empty() && ExperimentManager::is_enabled("GOOSE_SMART_APPROVE")? { - read_only_tools = detect_read_only_tools(&capabilities, needs_confirmation.iter().collect()).await; + read_only_tools = detect_read_only_tools(&capabilities, needs_confirmation.iter().map(|r| &**r).collect()).await; } // Process remaining tools that need confirmation @@ -318,7 +318,7 @@ impl Agent for TruncateAgent { if req_id == request.id { // Store the user's response let mut store = ToolPermissionStore::load()?; - store.record_permission(&request, confirmed, None)?; + store.record_permission(request, confirmed, None)?; if confirmed { // User approved - dispatch the tool call From bb266f5acd769d7d142f13da2fd345160634b017 Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Thu, 6 Mar 2025 17:30:32 -0800 Subject: [PATCH 10/14] lint err --- crates/goose/src/agents/truncate.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/crates/goose/src/agents/truncate.rs b/crates/goose/src/agents/truncate.rs index dc7996015b0c..5ce320cea073 100644 --- a/crates/goose/src/agents/truncate.rs +++ b/crates/goose/src/agents/truncate.rs @@ -266,11 +266,11 @@ impl Agent for TruncateAgent { match mode.as_str() { "approve" => { let mut read_only_tools = Vec::new(); - let mut needs_confirmation = Vec::new(); + let mut needs_confirmation = Vec::<&ToolRequest>::new(); // First check permissions for all tools let store = ToolPermissionStore::load()?; - for request in &tool_requests { + for request in tool_requests.iter() { if let Ok(tool_call) = request.tool_call.clone() { if let Some(allowed) = store.check_permission(request) { if allowed { @@ -280,17 +280,17 @@ impl Agent for TruncateAgent { output, ); } else { - needs_confirmation.push(request.clone()); + needs_confirmation.push(request); } } else { - needs_confirmation.push(request.clone()); + needs_confirmation.push(request); } } } // Only check read-only status for tools needing confirmation if !needs_confirmation.is_empty() && ExperimentManager::is_enabled("GOOSE_SMART_APPROVE")? { - read_only_tools = detect_read_only_tools(&capabilities, needs_confirmation.iter().map(|r| &**r).collect()).await; + read_only_tools = detect_read_only_tools(&capabilities, needs_confirmation.clone()).await; } // Process remaining tools that need confirmation From 865386e60d2f1dba76565ff987e8cf91e1a4e43b Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Thu, 6 Mar 2025 17:58:27 -0800 Subject: [PATCH 11/14] add cleanup of expired entries, 30 day expiry --- crates/goose/src/agents/permission_store.rs | 102 +++++++++++++++++++- crates/goose/src/agents/truncate.rs | 5 +- 2 files changed, 100 insertions(+), 7 deletions(-) diff --git a/crates/goose/src/agents/permission_store.rs b/crates/goose/src/agents/permission_store.rs index 504fb50c07b9..fdf8fc866ace 100644 --- a/crates/goose/src/agents/permission_store.rs +++ b/crates/goose/src/agents/permission_store.rs @@ -57,16 +57,26 @@ impl ToolPermissionStore { let file = File::open(file_path)?; let mut permissions: ToolPermissionStore = serde_json::from_reader(file)?; permissions.permissions_dir = store.permissions_dir; + + // Clean up expired entries on load + permissions.cleanup_expired()?; + Ok(permissions) } pub fn save(&self) -> anyhow::Result<()> { - let permissions_dir = self.permissions_dir.clone(); - std::fs::create_dir_all(&permissions_dir)?; + std::fs::create_dir_all(&self.permissions_dir)?; - let path = permissions_dir.join("tool_permissions.json"); + let path = self.permissions_dir.join("tool_permissions.json"); + let temp_path = path.with_extension("tmp"); + + // Write complete content to temporary file let content = serde_json::to_string_pretty(self)?; - std::fs::write(path, content)?; + std::fs::write(&temp_path, &content)?; + + // Atomically rename temp file to target file + std::fs::rename(temp_path, path)?; + Ok(()) } @@ -120,11 +130,27 @@ impl ToolPermissionStore { ); hasher.finalize().to_hex().to_string() } + + pub fn cleanup_expired(&mut self) -> anyhow::Result<()> { + let now = Utc::now().timestamp(); + let mut changed = false; + + self.permissions.retain(|_, records| { + records.retain(|record| record.expiry.map_or(true, |exp| exp > now)); + changed = changed || records.is_empty(); + !records.is_empty() + }); + + if changed { + self.save()?; + } + Ok(()) + } } #[cfg(test)] mod tests { - use crate::agents::permission_store::ToolPermissionStore; + use crate::agents::permission_store::{ToolPermissionRecord, ToolPermissionStore}; use crate::message::ToolRequest; use chrono::Utc; use mcp_core::tool::ToolCall; @@ -200,4 +226,70 @@ mod tests { assert_eq!(store.check_permission(&request1), Some(true)); assert!(store.check_permission(&request2).is_none()); } + + #[test] + fn test_cleanup_expired() { + let mut store = ToolPermissionStore::new(); + let tool_request = + create_test_tool_request("test_tool", serde_json::json!({"arg1": "value1"})); + + // Compute hash and key first + let context_hash = store.hash_tool_context(&tool_request); + let key = format!( + "{}:{}", + tool_request.tool_call.as_ref().unwrap().name, + context_hash + ); + + // Add an expired permission + store + .permissions + .entry(key.clone()) + .or_default() + .push(ToolPermissionRecord { + tool_name: "test_tool".to_string(), + allowed: true, + context_hash: context_hash.clone(), + readable_context: None, + timestamp: Utc::now().timestamp(), + expiry: Some(Utc::now().timestamp() - 1000), + }); + + // Add a valid permission + store + .record_permission(&tool_request, true, Some(Duration::from_secs(3600))) + .unwrap(); + + // Before cleanup - should have 2 records + assert_eq!( + store + .permissions + .get(&format!( + "{}:{}", + tool_request.tool_call.as_ref().unwrap().name, + store.hash_tool_context(&tool_request) + )) + .map(|records| records.len()), + Some(2) + ); + + // Run cleanup + store.cleanup_expired().unwrap(); + + // After cleanup - should have 1 record + assert_eq!( + store + .permissions + .get(&format!( + "{}:{}", + tool_request.tool_call.as_ref().unwrap().name, + store.hash_tool_context(&tool_request) + )) + .map(|records| records.len()), + Some(1) + ); + + // The remaining record should be valid + assert_eq!(store.check_permission(&tool_request), Some(true)); + } } diff --git a/crates/goose/src/agents/truncate.rs b/crates/goose/src/agents/truncate.rs index 5ce320cea073..e4a1f562779e 100644 --- a/crates/goose/src/agents/truncate.rs +++ b/crates/goose/src/agents/truncate.rs @@ -29,6 +29,7 @@ use mcp_core::prompt::Prompt; use mcp_core::protocol::GetPromptResult; use mcp_core::{tool::Tool, Content}; use serde_json::{json, Value}; +use std::time::Duration; const MAX_TRUNCATION_ATTEMPTS: usize = 3; const ESTIMATE_FACTOR_DECAY: f32 = 0.9; @@ -316,9 +317,9 @@ impl Agent for TruncateAgent { let mut rx = self.confirmation_rx.lock().await; while let Some((req_id, confirmed)) = rx.recv().await { if req_id == request.id { - // Store the user's response + // Store the user's response with 30-day expiration let mut store = ToolPermissionStore::load()?; - store.record_permission(request, confirmed, None)?; + store.record_permission(request, confirmed, Some(Duration::from_secs(30 * 24 * 60 * 60)))?; if confirmed { // User approved - dispatch the tool call From 6606c1fdda44082443f6278ec3b2169f42d2039f Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Fri, 7 Mar 2025 10:18:16 -0800 Subject: [PATCH 12/14] rm cleanup --- crates/goose/src/agents/permission_store.rs | 66 --------------------- 1 file changed, 66 deletions(-) diff --git a/crates/goose/src/agents/permission_store.rs b/crates/goose/src/agents/permission_store.rs index fdf8fc866ace..120bdcf3ebce 100644 --- a/crates/goose/src/agents/permission_store.rs +++ b/crates/goose/src/agents/permission_store.rs @@ -226,70 +226,4 @@ mod tests { assert_eq!(store.check_permission(&request1), Some(true)); assert!(store.check_permission(&request2).is_none()); } - - #[test] - fn test_cleanup_expired() { - let mut store = ToolPermissionStore::new(); - let tool_request = - create_test_tool_request("test_tool", serde_json::json!({"arg1": "value1"})); - - // Compute hash and key first - let context_hash = store.hash_tool_context(&tool_request); - let key = format!( - "{}:{}", - tool_request.tool_call.as_ref().unwrap().name, - context_hash - ); - - // Add an expired permission - store - .permissions - .entry(key.clone()) - .or_default() - .push(ToolPermissionRecord { - tool_name: "test_tool".to_string(), - allowed: true, - context_hash: context_hash.clone(), - readable_context: None, - timestamp: Utc::now().timestamp(), - expiry: Some(Utc::now().timestamp() - 1000), - }); - - // Add a valid permission - store - .record_permission(&tool_request, true, Some(Duration::from_secs(3600))) - .unwrap(); - - // Before cleanup - should have 2 records - assert_eq!( - store - .permissions - .get(&format!( - "{}:{}", - tool_request.tool_call.as_ref().unwrap().name, - store.hash_tool_context(&tool_request) - )) - .map(|records| records.len()), - Some(2) - ); - - // Run cleanup - store.cleanup_expired().unwrap(); - - // After cleanup - should have 1 record - assert_eq!( - store - .permissions - .get(&format!( - "{}:{}", - tool_request.tool_call.as_ref().unwrap().name, - store.hash_tool_context(&tool_request) - )) - .map(|records| records.len()), - Some(1) - ); - - // The remaining record should be valid - assert_eq!(store.check_permission(&tool_request), Some(true)); - } } From 1acbadaad11c27736d62fd0d82df727ced5c3a84 Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Fri, 7 Mar 2025 10:24:54 -0800 Subject: [PATCH 13/14] rm tests --- crates/goose/src/agents/permission_store.rs | 80 --------------------- 1 file changed, 80 deletions(-) diff --git a/crates/goose/src/agents/permission_store.rs b/crates/goose/src/agents/permission_store.rs index 120bdcf3ebce..23af2b4bd9ad 100644 --- a/crates/goose/src/agents/permission_store.rs +++ b/crates/goose/src/agents/permission_store.rs @@ -147,83 +147,3 @@ impl ToolPermissionStore { Ok(()) } } - -#[cfg(test)] -mod tests { - use crate::agents::permission_store::{ToolPermissionRecord, ToolPermissionStore}; - use crate::message::ToolRequest; - use chrono::Utc; - use mcp_core::tool::ToolCall; - use std::time::Duration; - - fn create_test_tool_request(name: &str, args: serde_json::Value) -> ToolRequest { - ToolRequest { - id: "test-id".to_string(), - tool_call: Ok(ToolCall { - name: name.to_string(), - arguments: args, - }), - } - } - - #[test] - fn test_permission_store_basic() { - let mut store = ToolPermissionStore::new(); - let tool_request = - create_test_tool_request("test_tool", serde_json::json!({"arg1": "value1"})); - - // Initially no permission recorded - assert!(store.check_permission(&tool_request).is_none()); - - // Record a permission - store.record_permission(&tool_request, true, None).unwrap(); - - // Should now find the recorded permission - assert_eq!(store.check_permission(&tool_request), Some(true)); - } - - #[test] - fn test_permission_expiry() { - let mut store = ToolPermissionStore::new(); - let tool_request = - create_test_tool_request("test_tool", serde_json::json!({"arg1": "value1"})); - - // Record a permission that expires in 1 second - store - .record_permission(&tool_request, true, Some(Duration::from_secs(1))) - .unwrap(); - - // Should initially be allowed - assert_eq!(store.check_permission(&tool_request), Some(true)); - - // Manually set expiry to the past - if let Some(records) = store.permissions.get_mut(&format!( - "{}:{}", - tool_request.tool_call.as_ref().unwrap().name, - store.hash_tool_context(&tool_request) - )) { - if let Some(record) = records.last_mut() { - record.expiry = Some(Utc::now().timestamp() - 2); - } - } - - // Should now be expired (no permission found) - assert!(store.check_permission(&tool_request).is_none()); - } - - #[test] - fn test_different_arguments() { - let mut store = ToolPermissionStore::new(); - - // Create two requests with same tool but different args - let request1 = create_test_tool_request("test_tool", serde_json::json!({"arg": "value1"})); - let request2 = create_test_tool_request("test_tool", serde_json::json!({"arg": "value2"})); - - // Record permission for first request - store.record_permission(&request1, true, None).unwrap(); - - // Should only find permission for first request - assert_eq!(store.check_permission(&request1), Some(true)); - assert!(store.check_permission(&request2).is_none()); - } -} From a271b833fbc460837e3f6f00d9c5e2c02f061708 Mon Sep 17 00:00:00 2001 From: Wendy Tang Date: Fri, 7 Mar 2025 10:40:26 -0800 Subject: [PATCH 14/14] clippy --- crates/goose/src/agents/permission_store.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/goose/src/agents/permission_store.rs b/crates/goose/src/agents/permission_store.rs index 23af2b4bd9ad..d9b1d363068c 100644 --- a/crates/goose/src/agents/permission_store.rs +++ b/crates/goose/src/agents/permission_store.rs @@ -136,7 +136,7 @@ impl ToolPermissionStore { let mut changed = false; self.permissions.retain(|_, records| { - records.retain(|record| record.expiry.map_or(true, |exp| exp > now)); + records.retain(|record| record.expiry.is_none_or(|exp| exp > now)); changed = changed || records.is_empty(); !records.is_empty() });