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

feat: parallel processing in approve mode (#1575)

Browse Source
This commit is contained in:
Wendy Tang
2025-03-07 15:25:00 -08:00
committed by GitHub
parent 00fc3a5de8
commit 3faff03f7c
1 changed files with 38 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/truncate.rs
+38 -22
View File
@@ -28,7 +28,7 @@ use anyhow::{anyhow, Result};
use indoc::indoc; use indoc::indoc;
use mcp_core::prompt::Prompt; use mcp_core::prompt::Prompt;
use mcp_core::protocol::GetPromptResult; use mcp_core::protocol::GetPromptResult;
use mcp_core::{tool::Tool, Content}; use mcp_core::{tool::Tool, Content, ToolError};
use serde_json::{json, Value}; use serde_json::{json, Value};
use std::time::Duration; use std::time::Duration;
@@ -111,6 +111,15 @@ impl TruncateAgent {
&OldestFirstTruncation, &OldestFirstTruncation,
) )
} }
async fn create_tool_future(
capabilities: &Capabilities,
tool_call: mcp_core::tool::ToolCall,
request_id: String,
) -> (String, Result<Vec<Content>, ToolError>) {
let output = capabilities.dispatch_tool_call(tool_call).await;
(request_id, output)
}
} }
#[async_trait] #[async_trait]
@@ -270,6 +279,7 @@ impl Agent for TruncateAgent {
"approve" => { "approve" => {
let mut read_only_tools = Vec::new(); let mut read_only_tools = Vec::new();
let mut needs_confirmation = Vec::<&ToolRequest>::new(); let mut needs_confirmation = Vec::<&ToolRequest>::new();
let mut approved_tools = Vec::new();
// First check permissions for all tools // First check permissions for all tools
let store = ToolPermissionStore::load()?; let store = ToolPermissionStore::load()?;
@@ -277,11 +287,8 @@ impl Agent for TruncateAgent {
if let Ok(tool_call) = request.tool_call.clone() { 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 { if allowed {
let output = capabilities.dispatch_tool_call(tool_call).await; // Instead of executing immediately, collect approved tools
message_tool_response = message_tool_response.with_tool_response( approved_tools.push((request.id.clone(), tool_call));
request.id.clone(),
output,
);
} else { } else {
needs_confirmation.push(request); needs_confirmation.push(request);
} }
@@ -296,16 +303,22 @@ impl Agent for TruncateAgent {
read_only_tools = detect_read_only_tools(&capabilities, needs_confirmation.clone()).await; read_only_tools = detect_read_only_tools(&capabilities, needs_confirmation.clone()).await;
} }
// Process remaining tools that need confirmation // Handle pre-approved and read-only tools in parallel
let mut tool_futures = Vec::new();
// Add pre-approved tools
for (request_id, tool_call) in approved_tools {
let tool_future = Self::create_tool_future(&capabilities, tool_call, request_id.clone());
tool_futures.push(tool_future);
}
// Process read-only tools
for request in &needs_confirmation { for request in &needs_confirmation {
if let Ok(tool_call) = request.tool_call.clone() { if let Ok(tool_call) = request.tool_call.clone() {
// Skip confirmation if the tool_call.name is in the read_only_tools list // Skip confirmation if the tool_call.name is in the read_only_tools list
if read_only_tools.contains(&tool_call.name) { if read_only_tools.contains(&tool_call.name) {
let output = capabilities.dispatch_tool_call(tool_call).await; let tool_future = Self::create_tool_future(&capabilities, tool_call, request.id.clone());
message_tool_response = message_tool_response.with_tool_response( tool_futures.push(tool_future);
request.id.clone(),
output,
);
} else { } else {
let confirmation = Message::user().with_tool_confirmation_request( let confirmation = Message::user().with_tool_confirmation_request(
request.id.clone(), request.id.clone(),
@@ -324,12 +337,9 @@ impl Agent for TruncateAgent {
store.record_permission(request, confirmed, Some(Duration::from_secs(30 * 24 * 60 * 60)))?; store.record_permission(request, confirmed, Some(Duration::from_secs(30 * 24 * 60 * 60)))?;
if confirmed { if confirmed {
// User approved - dispatch the tool call // Add this tool call to the futures collection
let output = capabilities.dispatch_tool_call(tool_call).await; let tool_future = Self::create_tool_future(&capabilities, tool_call, request.id.clone());
message_tool_response = message_tool_response.with_tool_response( tool_futures.push(tool_future);
request.id.clone(),
output,
);
} else { } else {
// User declined - add declined response // User declined - add declined response
message_tool_response = message_tool_response.with_tool_response( message_tool_response = message_tool_response.with_tool_response(
@@ -343,6 +353,14 @@ impl Agent for TruncateAgent {
} }
} }
} }
// Wait for all tool calls to complete
let results = futures::future::join_all(tool_futures).await;
for (request_id, output) in results {
message_tool_response = message_tool_response.with_tool_response(
request_id,
output,
);
}
}, },
"chat" => { "chat" => {
// Skip all tool calls in chat mode // Skip all tool calls in chat mode
@@ -370,10 +388,8 @@ impl Agent for TruncateAgent {
let mut tool_futures = Vec::new(); let mut tool_futures = Vec::new();
for request in &tool_requests { for request in &tool_requests {
if let Ok(tool_call) = request.tool_call.clone() { if let Ok(tool_call) = request.tool_call.clone() {
tool_futures.push(async { let tool_future = Self::create_tool_future(&capabilities, tool_call, request.id.clone());
let output = capabilities.dispatch_tool_call(tool_call).await; tool_futures.push(tool_future);
(request.id.clone(), output)
});
} }
} }
// Wait for all tool calls to complete // Wait for all tool calls to complete