tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

feat: use tool level permission control (#2081)

Browse Source
This commit is contained in:
Yingjie He
2025-04-09 13:45:15 -07:00
committed by GitHub
parent 513d5c8f5a
commit 9610406065
6 changed files with 115 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/session/mod.rs
+6 -9
View File
@@ -628,17 +628,15 @@ impl Session {
output::hide_thinking();
// Format the confirmation prompt
let prompt = "Goose would like to call the above tool, do you approve?".to_string();
let prompt = "Goose would like to call the above tool, do you allow?".to_string();
// Get confirmation from user
let confirmed = cliclack::confirm(prompt).initial_value(true).interact()?;
let permission = if confirmed {
Permission::AllowOnce
} else {
Permission::DenyOnce
};
let permission = cliclack::select(prompt)
.item(Permission::AllowOnce, "Allow", "Allow the tool call once")
.item(Permission::AlwaysAllow, "Always Allow", "Always allow the tool call")
.item(Permission::DenyOnce, "Deny", "Deny the tool call")
.interact()?;
self.agent.handle_confirmation(confirmation.id.clone(), PermissionConfirmation {
principal_name: "tool_name_placeholder".to_string(),
principal_type: PrincipalType::Tool,
permission,
},).await;
@@ -656,7 +654,6 @@ impl Session {
Permission::DenyOnce
};
self.agent.handle_confirmation(enable_extension_request.id.clone(), PermissionConfirmation {
principal_name: "extension_name_placeholder".to_string(),
principal_type: PrincipalType::Extension,
permission,
},).await;
crates/goose-server/src/routes/reply.rs
-1
View File
@@ -398,7 +398,6 @@ async fn confirm_handler(
.handle_confirmation(
request.id.clone(),
PermissionConfirmation {
principal_name: "tool_name_placeholder".to_string(),
principal_type: PrincipalType::Tool,
permission,
},
crates/goose/src/agents/agent.rs
+76 -91
View File
@@ -11,11 +11,11 @@ use tracing::{debug, error, instrument, warn};
use crate::agents::extension::{ExtensionConfig, ExtensionResult, ToolInfo};
use crate::agents::extension_manager::{get_parameter_names, ExtensionManager};
use crate::agents::types::ToolResultReceiver;
use crate::config::{Config, ExtensionConfigManager};
use crate::config::permission::PermissionLevel;
use crate::config::{Config, ExtensionConfigManager, PermissionManager};
use crate::message::{Message, MessageContent, ToolRequest};
use crate::permission::{
detect_read_only_tools, Permission, PermissionConfirmation, ToolPermissionStore,
};
use crate::permission::permission_judge::check_tool_permissions;
use crate::permission::{Permission, PermissionConfirmation};
use crate::providers::base::Provider;
use crate::providers::errors::ProviderError;
use crate::providers::toolshim::{
@@ -331,16 +331,18 @@ impl Agent {
tools.push(platform_tools::search_available_extensions_tool());
tools.push(platform_tools::enable_extension_tool());
let (tools_with_readonly_annotation, tools_without_annotation): (Vec<String>, Vec<String>) =
tools.iter().fold((vec![], vec![]), |mut acc, tool| {
let (tools_with_readonly_annotation, tools_without_annotation): (
HashSet<String>,
HashSet<String>,
) = tools
.iter()
.fold((HashSet::new(), HashSet::new()), |mut acc, tool| {
match &tool.annotations {
Some(annotations) => {
if annotations.read_only_hint {
acc.0.push(tool.name.clone());
}
Some(annotations) if annotations.read_only_hint => {
acc.0.insert(tool.name.clone());
}
None => {
acc.1.push(tool.name.clone());
_ => {
acc.1.insert(tool.name.clone());
}
}
acc
@@ -481,53 +483,22 @@ impl Agent {
// If there are install extension requests, always require confirmation
// or if goose_mode is approve or smart_approve, check permissions for all tools
if !enable_extension_requests.is_empty() || mode.as_str() == "approve" || mode.as_str() == "smart_approve" {
let mut needs_confirmation = Vec::<&ToolRequest>::new();
let mut approved_tools = Vec::new();
let mut llm_detect_candidates = Vec::<&ToolRequest>::new();
let mut detected_read_only_tools = Vec::new();
let mut permission_manager = PermissionManager::default();
// Skip the platform tools
remaining_requests.retain(|req| {
if let Ok(tool_call) = &req.tool_call {
!tool_call.name.starts_with("platform__")
} else {
true // If there's an error (Err), don't skip the request
}
});
let permission_check_result = check_tool_permissions(remaining_requests,
&mode,
tools_with_readonly_annotation.clone(),
tools_without_annotation.clone(),
&mut permission_manager,
self.provider()).await;
// If approve mode or smart approve mode, check permissions for all tools
if mode.as_str() == "approve" || mode.as_str() == "smart_approve" {
let store = ToolPermissionStore::load()?;
for request in &non_enable_extension_requests {
if let Ok(tool_call) = request.tool_call.clone() {
// Regular permission checking for other tools
if tools_with_readonly_annotation.contains(&tool_call.name) {
approved_tools.push((request.id.clone(), tool_call));
} else if let Some(allowed) = store.check_permission(request) {
if allowed {
// Instead of executing immediately, collect approved tools
approved_tools.push((request.id.clone(), tool_call));
} else {
// If the tool doesn't have any annotation, we can use llm-as-a-judge to check permission.
if tools_without_annotation.contains(&tool_call.name) {
llm_detect_candidates.push(request);
}
needs_confirmation.push(request);
}
} else {
if tools_without_annotation.contains(&tool_call.name) {
llm_detect_candidates.push(request);
}
needs_confirmation.push(request);
}
}
}
}
// Only check read-only status for tools needing confirmation
if !llm_detect_candidates.is_empty() && mode == "smart_approve" {
detected_read_only_tools = detect_read_only_tools(self.provider(), llm_detect_candidates.clone()).await;
// Remove install extensions from read-only tools
if !enable_extension_requests.is_empty() {
detected_read_only_tools.retain(|tool_name| {
!enable_extension_requests.iter().any(|req| {
req.tool_call.as_ref()
.map(|call| call.name == *tool_name)
.unwrap_or(false)
})
});
}
}
// Handle pre-approved and read-only tools in parallel
let mut tool_futures = Vec::new();
@@ -562,44 +533,58 @@ impl Agent {
}
}
// Process read-only tools
for request in &needs_confirmation {
// Skip the confirmation for approved tools
for request in &permission_check_result.approved {
if let Ok(tool_call) = request.tool_call.clone() {
let is_frontend_tool = self.is_frontend_tool(&tool_call.name);
// Skip confirmation if the tool_call.name is in the read_only_tools list
if detected_read_only_tools.contains(&tool_call.name) {
let tool_future = Self::create_tool_future(&extension_manager, tool_call, is_frontend_tool, request.id.clone());
tool_futures.push(tool_future);
} else {
let confirmation = Message::user().with_tool_confirmation_request(
request.id.clone(),
tool_call.name.clone(),
tool_call.arguments.clone(),
Some("Goose would like to call the above tool. Allow? (y/n):".to_string()),
);
yield confirmation;
let tool_future = Self::create_tool_future(&extension_manager, tool_call, is_frontend_tool, request.id.clone());
tool_futures.push(tool_future);
}
}
// Wait for confirmation response through the channel
let mut rx = self.confirmation_rx.lock().await;
while let Some((req_id, tool_confirmation)) = rx.recv().await {
if req_id == request.id {
let confirmed = tool_confirmation.permission == Permission::AllowOnce || tool_confirmation.permission == Permission::AlwaysAllow;
if confirmed {
// Add this tool call to the futures collection
let tool_future = Self::create_tool_future(&extension_manager, tool_call, is_frontend_tool, request.id.clone());
tool_futures.push(tool_future);
} else {
// User declined - add declined response
message_tool_response = message_tool_response.with_tool_response(
request.id.clone(),
Ok(vec![Content::text(
"The user has declined to run this tool. \
DO NOT attempt to call this tool again. \
If there are no alternative methods to proceed, clearly explain the situation and STOP.")]),
);
let denied_content_text = Content::text(
"The user has declined to run this tool. \
DO NOT attempt to call this tool again. \
If there are no alternative methods to proceed, clearly explain the situation and STOP.");
for request in &permission_check_result.denied {
message_tool_response = message_tool_response.with_tool_response(
request.id.clone(),
Ok(vec![denied_content_text.clone()]),
);
}
// Process read-only tools
for request in &permission_check_result.needs_approval {
if let Ok(tool_call) = request.tool_call.clone() {
let is_frontend_tool = self.is_frontend_tool(&tool_call.name);
let confirmation = Message::user().with_tool_confirmation_request(
request.id.clone(),
tool_call.name.clone(),
tool_call.arguments.clone(),
Some("Goose would like to call the above tool. Allow? (y/n):".to_string()),
);
yield confirmation;
// Wait for confirmation response through the channel
let mut rx = self.confirmation_rx.lock().await;
while let Some((req_id, tool_confirmation)) = rx.recv().await {
if req_id == request.id {
let confirmed = tool_confirmation.permission == Permission::AllowOnce || tool_confirmation.permission == Permission::AlwaysAllow;
if confirmed {
// Add this tool call to the futures collection
let tool_future = Self::create_tool_future(&extension_manager, tool_call.clone(), is_frontend_tool, request.id.clone());
tool_futures.push(tool_future);
if tool_confirmation.permission == Permission::AlwaysAllow {
permission_manager.update_user_permission(&tool_call.name, PermissionLevel::AlwaysAllow);
}
break; // Exit the loop once the matching `req_id` is found
} else {
// User declined - add declined response
message_tool_response = message_tool_response.with_tool_response(
request.id.clone(),
Ok(vec![denied_content_text.clone()]),
);
}
break; // Exit the loop once the matching `req_id` is found
}
}
}
crates/goose/src/config/permission.rs
-9
View File
@@ -1,6 +1,5 @@
use super::APP_STRATEGY;
use etcetera::{choose_app_strategy, AppStrategy};
use once_cell::sync::OnceCell;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs;
@@ -31,9 +30,6 @@ pub struct PermissionManager {
permission_map: HashMap<String, PermissionConfig>, // Mapping of permission names to configurations
}
// Global singleton for the PermissionManager
static GLOBAL_PERMISSION_MANAGER: OnceCell<PermissionManager> = OnceCell::new();
// Constants representing specific permission categories
const USER_PERMISSION: &str = "user";
const SMART_APPROVE_PERMISSION: &str = "smart_approve";
@@ -68,11 +64,6 @@ impl Default for PermissionManager {
}
impl PermissionManager {
/// Returns the global instance of the PermissionManager, initializing it if necessary.
pub fn global() -> &'static PermissionManager {
GLOBAL_PERMISSION_MANAGER.get_or_init(PermissionManager::default)
}
/// Creates a new `PermissionManager` with a specified config path.
pub fn new<P: AsRef<Path>>(config_path: P) -> Self {
let config_path = config_path.as_ref().to_path_buf();
crates/goose/src/permission/permission_confirmation.rs
-1
View File
@@ -15,7 +15,6 @@ pub enum PrincipalType {
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct PermissionConfirmation {
pub principal_name: String,
pub principal_type: PrincipalType,
pub permission: Permission,
}
crates/goose/src/permission/permission_judge.rs
+33 -18
View File
@@ -33,11 +33,12 @@ fn create_read_only_tool() -> Tool {
- `INSERT`, `UPDATE`, or `DELETE` in SQL.
- Writing or appending to a file.
- Modifying system configurations.
- Sending messages to Slack channel.
How to analyze tool requests:
- Inspect each tool request to identify its purpose based on its name and arguments.
- Categorize the operation as read-only if it does not involve any state or data modification.
- Return a list of tool names that are strictly read-only.
- Return a list of tool names that are strictly read-only. If you cannot make the decision, then it is not read-only.
Use this analysis to generate the list of tools performing read-only operations from the provided tool requests.
"#}
@@ -67,6 +68,16 @@ fn create_read_only_tool() -> Tool {
/// Builds the message to be sent to the LLM for detecting read-only operations.
fn create_check_messages(tool_requests: Vec<&ToolRequest>) -> Vec<Message> {
let tool_names: Vec<String> = tool_requests
.iter()
.filter_map(|req| {
if let Ok(tool_call) = &req.tool_call {
Some(tool_call.name.clone())
} else {
None // Skip requests with errors in tool_call
}
})
.collect();
let mut check_messages = vec![];
check_messages.push(Message {
role: mcp_core::Role::User,
@@ -79,7 +90,7 @@ fn create_check_messages(tool_requests: Vec<&ToolRequest>) -> Vec<Message> {
\n- Examples include file reading, SELECT queries in SQL, and directory listing. \
\n- Write operations include INSERT, UPDATE, DELETE, and file writing. \
\n\nPlease provide a list of tool names that qualify as read-only:",
tool_requests,
tool_names.join(", "),
),
annotations: None,
})],
@@ -147,7 +158,7 @@ pub struct PermissionCheckResult {
}
pub async fn check_tool_permissions(
remaining_requests: Vec<ToolRequest>,
remaining_requests: Vec<&&ToolRequest>,
mode: &str,
tools_with_readonly_annotation: HashSet<String>,
tools_without_annotation: HashSet<String>,
@@ -159,14 +170,14 @@ pub async fn check_tool_permissions(
let mut denied = vec![];
let mut llm_detect_candidates = vec![];
for request in remaining_requests {
for &&request in &remaining_requests {
if let Ok(tool_call) = request.tool_call.clone() {
// 1. Check user-defined permission
if let Some(level) = permission_manager.get_user_permission(&tool_call.name) {
match level {
PermissionLevel::AlwaysAllow => approved.push(request),
PermissionLevel::AskBefore => needs_approval.push(request),
PermissionLevel::NeverAllow => denied.push(request),
PermissionLevel::AlwaysAllow => approved.push(request.clone()),
PermissionLevel::AskBefore => needs_approval.push(request.clone()),
PermissionLevel::NeverAllow => denied.push(request.clone()),
}
continue;
}
@@ -174,30 +185,30 @@ pub async fn check_tool_permissions(
// 2. Fallback based on mode
match mode {
"manual_approve" => {
needs_approval.push(request);
needs_approval.push(request.clone());
}
"smart_approve" => {
if let Some(level) =
permission_manager.get_smart_approve_permission(&tool_call.name)
{
match level {
PermissionLevel::AlwaysAllow => approved.push(request),
PermissionLevel::AskBefore => needs_approval.push(request),
PermissionLevel::NeverAllow => denied.push(request),
PermissionLevel::AlwaysAllow => approved.push(request.clone()),
PermissionLevel::AskBefore => needs_approval.push(request.clone()),
PermissionLevel::NeverAllow => denied.push(request.clone()),
}
continue;
}
if tools_with_readonly_annotation.contains(&tool_call.name) {
approved.push(request);
approved.push(request.clone());
} else if tools_without_annotation.contains(&tool_call.name) {
llm_detect_candidates.push(request);
llm_detect_candidates.push(request.clone());
} else {
needs_approval.push(request);
needs_approval.push(request.clone());
}
}
_ => {
needs_approval.push(request);
needs_approval.push(request.clone());
}
}
}
@@ -210,13 +221,13 @@ pub async fn check_tool_permissions(
for request in llm_detect_candidates {
if let Ok(tool_call) = request.tool_call.clone() {
if detected_readonly_tools.contains(&tool_call.name) {
approved.push(request);
approved.push(request.clone());
permission_manager.update_smart_approve_permission(
&tool_call.name,
PermissionLevel::AlwaysAllow,
);
} else {
needs_approval.push(request);
needs_approval.push(request.clone());
permission_manager.update_smart_approve_permission(
&tool_call.name,
PermissionLevel::AskBefore,
@@ -401,7 +412,11 @@ mod tests {
}),
};
let remaining_requests = vec![tool_request_1, tool_request_2];
// Store ToolRequests in a Vec
let tool_requests = vec![&tool_request_1, &tool_request_2];
// Create a Vec of references to ToolRequests
let remaining_requests: Vec<&&ToolRequest> = tool_requests.iter().collect();
// Call the function under test
let result = check_tool_permissions(