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

feat: use the same permission flow for enable extensions (#2302)

Browse Source
This commit is contained in:
Yingjie He
2025-04-23 08:54:04 -07:00
committed by GitHub
parent 08682507d9
commit cc755100f0
16 changed files with 77 additions and 420 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+37 -38
View File
@@ -22,8 +22,8 @@ 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::platform_tools::{
PLATFORM_LIST_RESOURCES_TOOL_NAME, PLATFORM_READ_RESOURCE_TOOL_NAME,
PLATFORM_SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME,
PLATFORM_ENABLE_EXTENSION_TOOL_NAME, PLATFORM_LIST_RESOURCES_TOOL_NAME,
PLATFORM_READ_RESOURCE_TOOL_NAME, PLATFORM_SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME,
};
use crate::agents::prompt_manager::PromptManager;
use crate::agents::types::SessionConfig;
@@ -33,9 +33,7 @@ use mcp_core::{
};
use super::platform_tools;
use super::tool_execution::{
ExtensionInstallResult, ToolFuture, CHAT_MODE_TOOL_SKIPPED_RESPONSE, DECLINED_RESPONSE,
};
use super::tool_execution::{ToolFuture, CHAT_MODE_TOOL_SKIPPED_RESPONSE, DECLINED_RESPONSE};
const MAX_TRUNCATION_ATTEMPTS: usize = 3;
const ESTIMATE_FACTOR_DECAY: f32 = 0.9;
@@ -114,6 +112,16 @@ impl Agent {
tool_call: mcp_core::tool::ToolCall,
request_id: String,
) -> (String, Result<Vec<Content>, ToolError>) {
if tool_call.name == PLATFORM_ENABLE_EXTENSION_TOOL_NAME {
let extension_name = tool_call
.arguments
.get("extension_name")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
return self.enable_extension(extension_name, request_id).await;
}
let extension_manager = self.extension_manager.lock().await;
let result = if tool_call.name == PLATFORM_READ_RESOURCE_TOOL_NAME {
// Check if the tool is read_resource and handle it separately
@@ -418,26 +426,26 @@ impl Agent {
// What remains is handling the remaining tool requests (enable extension,
// regular tool calls) in goose_mode == ["auto", "approve" or "smart_approve"]
let mut permission_manager = PermissionManager::default();
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;
let (permission_check_result, enable_extension_request_ids) = check_tool_permissions(
&remaining_requests,
&mode,
tools_with_readonly_annotation.clone(),
tools_without_annotation.clone(),
&mut permission_manager,
self.provider(),
).await;
// Handle pre-approved and read-only tools in parallel
let mut tool_futures: Vec<ToolFuture> = Vec::new();
let mut install_results: Vec<ExtensionInstallResult> = Vec::new();
let install_results_arc = Arc::new(Mutex::new(install_results));
// Skip the confirmation for approved tools
for request in &permission_check_result.approved {
if let Ok(tool_call) = request.tool_call.clone() {
let tool_future = self.dispatch_tool_call(tool_call, request.id.clone());
tool_futures.push(Box::pin(tool_future));
tool_futures.push(Box::pin(tool_future));
}
}
for request in &permission_check_result.denied {
let mut response = message_tool_response.lock().await;
*response = response.clone().with_tool_response(
@@ -446,51 +454,42 @@ impl Agent {
);
}
// we need interior mutability in handle_approval_tool_requests
// We need interior mutability in handle_approval_tool_requests
let tool_futures_arc = Arc::new(Mutex::new(tool_futures));
// Process tools requiring approval (enable extension, regular tool calls)
let mut tool_approval_stream = self.handle_approval_tool_requests(
&permission_check_result.needs_approval,
install_results_arc.clone(),
tool_futures_arc.clone(),
&mut permission_manager,
message_tool_response.clone()
message_tool_response.clone(),
);
// we have a stream of tool_approval_requests to handle
// execution is yeield back to this reply loop, and is of the same Message
// type, so we can yield the Message back up to be handled and grab and
// We have a stream of tool_approval_requests to handle
// Execution is yielded back to this reply loop, and is of the same Message
// type, so we can yield the Message back up to be handled and grab any
// confirmations or denials
while let Some(msg) = tool_approval_stream.try_next().await? {
yield msg;
}
tool_futures = {
// Lock the mutex asynchronously.
// Lock the mutex asynchronously
let mut futures_lock = tool_futures_arc.lock().await;
// Drain the vector and collect into a new Vec.
// Drain the vector and collect into a new Vec
futures_lock.drain(..).collect::<Vec<_>>()
};
install_results = {
// Lock the mutex asynchronously.
let mut results_lock = install_results_arc.lock().await;
// Drain the vector and collect into a new Vec.
results_lock.drain(..).collect::<Vec<_>>()
};
// Wait for all tool calls to complete
let results = futures::future::join_all(tool_futures).await;
let mut all_install_successful = true;
// Check if any install results had errors before processing them
let all_install_successful = !install_results.iter().any(|(_, result)| result.is_err());
for (request_id, output) in results.into_iter().chain(install_results.into_iter()) {
for (request_id, output) in results.into_iter() {
if enable_extension_request_ids.contains(&request_id) && output.is_err(){
all_install_successful = false;
}
let mut response = message_tool_response.lock().await;
*response = response.clone().with_tool_response(
request_id,
output
);
*response = response.clone().with_tool_response(request_id, output);
}
// Update system prompt and tools if installations were successful
crates/goose/src/agents/tool_execution.rs
+11 -24
View File
@@ -10,8 +10,6 @@ use tokio::sync::Mutex;
use crate::config::permission::PermissionLevel;
use crate::config::PermissionManager;
use crate::message::{Message, ToolRequest};
use crate::permission::permission_confirmation::PrincipalType;
use crate::permission::permission_judge::get_confirmation_message;
use crate::permission::Permission;
use mcp_core::{Content, ToolError};
@@ -19,9 +17,6 @@ use mcp_core::{Content, ToolError};
pub(crate) type ToolFuture<'a> =
Pin<Box<dyn Future<Output = (String, Result<Vec<Content>, ToolError>)> + Send + 'a>>;
pub(crate) type ToolFuturesVec<'a> = Arc<Mutex<Vec<ToolFuture<'a>>>>;
// Type alias for extension installation results
pub(crate) type ExtensionInstallResult = (String, Result<Vec<Content>, ToolError>);
pub(crate) type ExtensionInstallResults = Arc<Mutex<Vec<ExtensionInstallResult>>>;
use crate::agents::Agent;
@@ -42,7 +37,6 @@ impl Agent {
pub(crate) fn handle_approval_tool_requests<'a>(
&'a self,
tool_requests: &'a [ToolRequest],
install_results: ExtensionInstallResults,
tool_futures: ToolFuturesVec<'a>,
permission_manager: &'a mut PermissionManager,
message_tool_response: Arc<Mutex<Message>>,
@@ -50,31 +44,24 @@ impl Agent {
try_stream! {
for request in tool_requests {
if let Ok(tool_call) = request.tool_call.clone() {
let (principal_type, confirmation) = get_confirmation_message(&request.id.clone(), tool_call.clone());
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 mut rx = self.confirmation_rx.lock().await;
while let Some((req_id, confirmation)) = rx.recv().await {
if req_id == request.id {
if confirmation.permission == Permission::AllowOnce || confirmation.permission == Permission::AlwaysAllow {
if principal_type == PrincipalType::Extension {
let extension_name = tool_call.arguments.get("extension_name")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let tool_future = self.dispatch_tool_call(tool_call.clone(), request.id.clone());
let mut futures = tool_futures.lock().await;
futures.push(Box::pin(tool_future));
let mut results = install_results.lock().await;
let install_result = self.enable_extension(extension_name, request.id.clone()).await;
results.push(install_result);
} else {
// Add this tool call to the futures collection
let tool_future = self.dispatch_tool_call(tool_call.clone(), request.id.clone());
let mut futures = tool_futures.lock().await;
futures.push(Box::pin(tool_future));
if confirmation.permission == Permission::AlwaysAllow {
permission_manager.update_user_permission(&tool_call.name, PermissionLevel::AlwaysAllow);
}
if confirmation.permission == Permission::AlwaysAllow {
permission_manager.update_user_permission(&tool_call.name, PermissionLevel::AlwaysAllow);
}
} else {
// User declined - add declined response