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

fix: tool confirmation handling for multiple requests (#7856)

Browse Source
This commit is contained in:
Lifei Zhou
2026-03-16 09:46:34 +11:00
committed by GitHub
parent 831cb9bb82
commit 291bbd8330
4 changed files with 221 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+26 -18
View File
@@ -12,6 +12,7 @@ use uuid::Uuid;
use super::container::Container;
use super::final_output_tool::FinalOutputTool;
use super::platform_tools;
use super::tool_confirmation_router::ToolConfirmationRouter;
use super::tool_execution::{ToolCallResult, CHAT_MODE_TOOL_SKIPPED_RESPONSE, DECLINED_RESPONSE};
use crate::action_required_manager::ActionRequiredManager;
use crate::agents::extension::{ExtensionConfig, ExtensionResult, ToolInfo};
@@ -141,8 +142,7 @@ pub struct Agent {
pub(super) frontend_tools: Mutex<HashMap<String, FrontendTool>>,
pub(super) frontend_instructions: Mutex<Option<String>>,
pub(super) prompt_manager: Mutex<PromptManager>,
pub(super) confirmation_tx: mpsc::Sender<(String, PermissionConfirmation)>,
pub(super) confirmation_rx: Mutex<mpsc::Receiver<(String, PermissionConfirmation)>>,
pub tool_confirmation_router: ToolConfirmationRouter,
pub(super) tool_result_tx: mpsc::Sender<(String, ToolResult<CallToolResult>)>,
pub(super) tool_result_rx: ToolResultReceiver,
@@ -215,8 +215,6 @@ impl Agent {
}
pub fn with_config(config: AgentConfig) -> Self {
// Create channels with buffer size 32 (adjust if needed)
let (confirm_tx, confirm_rx) = mpsc::channel(32);
let (tool_tx, tool_rx) = mpsc::channel(32);
let provider = Arc::new(Mutex::new(None));
@@ -240,8 +238,7 @@ impl Agent {
frontend_tools: Mutex::new(HashMap::new()),
frontend_instructions: Mutex::new(None),
prompt_manager: Mutex::new(PromptManager::new()),
confirmation_tx: confirm_tx,
confirmation_rx: Mutex::new(confirm_rx),
tool_confirmation_router: ToolConfirmationRouter::new(),
tool_result_tx: tool_tx,
tool_result_rx: Arc::new(Mutex::new(tool_rx)),
retry_manager: RetryManager::new(),
@@ -861,8 +858,12 @@ impl Agent {
return;
}
}
if let Err(e) = self.confirmation_tx.send((request_id, confirmation)).await {
error!("Failed to send confirmation: {}", e);
if !self
.tool_confirmation_router
.deliver(request_id, confirmation)
.await
{
error!("Failed to deliver confirmation");
}
}
@@ -2136,7 +2137,7 @@ mod tests {
*agent.provider.lock().await =
Some(provider.clone() as Arc<dyn crate::providers::base::Provider>);
// Known request_id → provider handles it, confirmation_tx NOT called
// Known request_id → provider handles it, confirmation_router NOT called
agent
.handle_confirmation(
"known".to_string(),
@@ -2148,7 +2149,12 @@ mod tests {
.await;
assert_eq!(provider.handled.lock().await.len(), 1);
// Unknown request_id → provider returns false, falls through to confirmation_tx
// Unknown request_id → provider returns false, falls through to confirmation_router
// Register first so deliver() has somewhere to send
let rx = agent
.tool_confirmation_router
.register("unknown".to_string())
.await;
agent
.handle_confirmation(
"unknown".to_string(),
@@ -2159,17 +2165,20 @@ mod tests {
)
.await;
assert_eq!(provider.handled.lock().await.len(), 2);
// Verify the fallthrough went to confirmation_rx
let mut rx = agent.confirmation_rx.lock().await;
let (id, conf) = rx.recv().await.unwrap();
assert_eq!(id, "unknown");
// Verify the fallthrough went to confirmation_router
let conf = rx.await.unwrap();
assert_eq!(conf.permission, crate::permission::Permission::DenyOnce);
}
#[tokio::test]
async fn test_handle_confirmation_noop_provider() {
let agent = Agent::new();
// No provider set → Noop routing, goes straight to confirmation_tx
// No provider set → Noop routing, goes straight to confirmation_router
// Register first so deliver() has somewhere to send
let rx = agent
.tool_confirmation_router
.register("any".to_string())
.await;
agent
.handle_confirmation(
"any".to_string(),
@@ -2180,9 +2189,8 @@ mod tests {
)
.await;
let mut rx = agent.confirmation_rx.lock().await;
let (id, _) = rx.recv().await.unwrap();
assert_eq!(id, "any");
let conf = rx.await.unwrap();
assert_eq!(conf.permission, crate::permission::Permission::AllowOnce);
}
#[tokio::test]
crates/goose/src/agents/mod.rs
+1
View File
@@ -18,6 +18,7 @@ mod schedule_tool;
pub mod subagent_execution_tool;
pub(crate) mod subagent_handler;
pub(crate) mod subagent_task_config;
mod tool_confirmation_router;
mod tool_execution;
pub mod types;
pub mod validate_extensions;
crates/goose/src/agents/tool_confirmation_router.rs
+144
View File
@@ -0,0 +1,144 @@
use std::collections::HashMap;
use tokio::sync::{oneshot, Mutex};
use tracing::warn;
use crate::permission::PermissionConfirmation;
pub struct ToolConfirmationRouter {
pending: Mutex<HashMap<String, oneshot::Sender<PermissionConfirmation>>>,
}
impl ToolConfirmationRouter {
pub fn new() -> Self {
Self {
pending: Mutex::new(HashMap::new()),
}
}
pub async fn register(&self, request_id: String) -> oneshot::Receiver<PermissionConfirmation> {
let (tx, rx) = oneshot::channel();
let mut pending = self.pending.lock().await;
pending.retain(|_, sender| !sender.is_closed());
pending.insert(request_id, tx);
rx
}
pub async fn deliver(&self, request_id: String, confirmation: PermissionConfirmation) -> bool {
if let Some(tx) = self.pending.lock().await.remove(&request_id) {
if tx.send(confirmation).is_err() {
warn!(
request_id = %request_id,
"Confirmation receiver was dropped (task cancelled)"
);
false
} else {
true
}
} else {
warn!(
request_id = %request_id,
"No task waiting for confirmation"
);
false
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::permission::permission_confirmation::PrincipalType;
use crate::permission::Permission;
fn test_confirmation() -> PermissionConfirmation {
PermissionConfirmation {
principal_type: PrincipalType::Tool,
permission: Permission::AllowOnce,
}
}
#[tokio::test]
async fn test_register_then_deliver() {
let router = ToolConfirmationRouter::new();
let rx = router.register("req_1".to_string()).await;
assert!(
router
.deliver("req_1".to_string(), test_confirmation())
.await
);
let confirmation = rx.await.unwrap();
assert_eq!(confirmation.permission, Permission::AllowOnce);
}
#[tokio::test]
async fn test_deliver_unknown_request() {
let router = ToolConfirmationRouter::new();
assert!(
!router
.deliver("unknown".to_string(), test_confirmation())
.await
);
}
#[tokio::test]
async fn test_cancelled_receiver() {
let router = ToolConfirmationRouter::new();
let rx = router.register("req_1".to_string()).await;
drop(rx); // simulate task cancellation
assert!(
!router
.deliver("req_1".to_string(), test_confirmation())
.await
);
}
#[tokio::test]
async fn test_stale_entries_pruned_on_register() {
let router = ToolConfirmationRouter::new();
let rx = router.register("req_1".to_string()).await;
drop(rx); // simulate task cancellation — entry is now stale
assert_eq!(router.pending.lock().await.len(), 1);
let _rx2 = router.register("req_2".to_string()).await;
assert_eq!(router.pending.lock().await.len(), 1); // only req_2 remains
assert!(router.pending.lock().await.contains_key("req_2"));
}
#[tokio::test]
async fn test_concurrent_requests_out_of_order() {
use std::sync::Arc;
let router = Arc::new(ToolConfirmationRouter::new());
// Register two requests
let rx1 = router.register("req_1".to_string()).await;
let rx2 = router.register("req_2".to_string()).await;
// Deliver in reverse order
assert!(
router
.deliver(
"req_2".to_string(),
PermissionConfirmation {
principal_type: PrincipalType::Tool,
permission: Permission::DenyOnce,
}
)
.await
);
assert_eq!(router.pending.lock().await.len(), 1);
assert!(
router
.deliver("req_1".to_string(), test_confirmation())
.await
);
assert_eq!(router.pending.lock().await.len(), 0);
let c1 = rx1.await.unwrap();
assert_eq!(c1.permission, Permission::AllowOnce);
let c2 = rx2.await.unwrap();
assert_eq!(c2.permission, Permission::DenyOnce);
}
}
crates/goose/src/agents/tool_execution.rs
+50 -51
View File
@@ -99,7 +99,9 @@ impl Agent {
}
});
let confirmation = Message::assistant()
let confirmation_rx = self.tool_confirmation_router.register(request.id.clone()).await;
let action_required_msg = Message::assistant()
.with_action_required(
request.id.clone(),
tool_call.name.to_string().clone(),
@@ -107,61 +109,58 @@ impl Agent {
security_message,
)
.user_only();
yield confirmation;
yield action_required_msg;
let mut rx = self.confirmation_rx.lock().await;
while let Some((req_id, confirmation)) = rx.recv().await {
if req_id == request.id {
// Log user decision if this was a security alert
if let Some(finding_id) = get_security_finding_id_from_results(&request.id, inspection_results) {
tracing::info!(
monotonic_counter.goose.prompt_injection_user_decisions = 1,
decision = ?confirmation.permission,
finding_id = %finding_id,
tool_request_id = %request.id,
"Prompt injection detection: user decision on command injection finding"
);
}
let confirmation = confirmation_rx.await
.map_err(|_| anyhow::anyhow!("Confirmation channel closed for request {}", request.id))?;
if confirmation.permission == Permission::AllowOnce || confirmation.permission == Permission::AlwaysAllow {
let (req_id, tool_result) = self.dispatch_tool_call(tool_call.clone(), request.id.clone(), cancellation_token.clone(), session).await;
let mut futures = tool_futures.lock().await;
// Log user decision if this was a security alert
if let Some(finding_id) = get_security_finding_id_from_results(&request.id, inspection_results) {
tracing::info!(
monotonic_counter.goose.prompt_injection_user_decisions = 1,
decision = ?confirmation.permission,
finding_id = %finding_id,
tool_request_id = %request.id,
"Prompt injection detection: user decision on command injection finding"
);
}
futures.push((req_id, match tool_result {
Ok(result) => tool_stream(
result.notification_stream.unwrap_or_else(|| Box::new(stream::empty())),
result.result,
),
Err(e) => tool_stream(
Box::new(stream::empty()),
futures::future::ready(Err(e)),
),
}));
if confirmation.permission == Permission::AllowOnce || confirmation.permission == Permission::AlwaysAllow {
let (req_id, tool_result) = self.dispatch_tool_call(tool_call.clone(), request.id.clone(), cancellation_token.clone(), session).await;
let mut futures = tool_futures.lock().await;
// Update the shared permission manager when user selects "Always Allow"
if confirmation.permission == Permission::AlwaysAllow {
self.tool_inspection_manager
.update_permission_manager(&tool_call.name, PermissionLevel::AlwaysAllow)
.await;
}
} else {
// User declined - update the specific response message for this request
if let Some(response_msg) = request_to_response_map.get(&request.id) {
let mut response = response_msg.lock().await;
*response = response.clone().with_tool_response_with_metadata(
request.id.clone(),
Ok(rmcp::model::CallToolResult::error(vec![Content::text(DECLINED_RESPONSE)])),
request.metadata.as_ref(),
);
}
futures.push((req_id, match tool_result {
Ok(result) => tool_stream(
result.notification_stream.unwrap_or_else(|| Box::new(stream::empty())),
result.result,
),
Err(e) => tool_stream(
Box::new(stream::empty()),
futures::future::ready(Err(e)),
),
}));
if confirmation.permission == Permission::AlwaysDeny {
self.tool_inspection_manager
.update_permission_manager(&tool_call.name, PermissionLevel::NeverAllow)
.await;
}
}
break; // Exit the loop once the matching `req_id` is found
// Update the shared permission manager when user selects "Always Allow"
if confirmation.permission == Permission::AlwaysAllow {
self.tool_inspection_manager
.update_permission_manager(&tool_call.name, PermissionLevel::AlwaysAllow)
.await;
}
} else {
// User declined - update the specific response message for this request
if let Some(response_msg) = request_to_response_map.get(&request.id) {
let mut response = response_msg.lock().await;
*response = response.clone().with_tool_response_with_metadata(
request.id.clone(),
Ok(rmcp::model::CallToolResult::error(vec![Content::text(DECLINED_RESPONSE)])),
request.metadata.as_ref(),
);
}
if confirmation.permission == Permission::AlwaysDeny {
self.tool_inspection_manager
.update_permission_manager(&tool_call.name, PermissionLevel::NeverAllow)
.await;
}
}
}