diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index e925fb44..2c62b0b7 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -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>, pub(super) frontend_instructions: Mutex>, pub(super) prompt_manager: Mutex, - pub(super) confirmation_tx: mpsc::Sender<(String, PermissionConfirmation)>, - pub(super) confirmation_rx: Mutex>, + pub tool_confirmation_router: ToolConfirmationRouter, pub(super) tool_result_tx: mpsc::Sender<(String, ToolResult)>, 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); - // 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] diff --git a/crates/goose/src/agents/mod.rs b/crates/goose/src/agents/mod.rs index b3bc386a..eba1f019 100644 --- a/crates/goose/src/agents/mod.rs +++ b/crates/goose/src/agents/mod.rs @@ -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; diff --git a/crates/goose/src/agents/tool_confirmation_router.rs b/crates/goose/src/agents/tool_confirmation_router.rs new file mode 100644 index 00000000..2148e505 --- /dev/null +++ b/crates/goose/src/agents/tool_confirmation_router.rs @@ -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>>, +} + +impl ToolConfirmationRouter { + pub fn new() -> Self { + Self { + pending: Mutex::new(HashMap::new()), + } + } + + pub async fn register(&self, request_id: String) -> oneshot::Receiver { + 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); + } +} diff --git a/crates/goose/src/agents/tool_execution.rs b/crates/goose/src/agents/tool_execution.rs index 2811d81c..38c9e4b9 100644 --- a/crates/goose/src/agents/tool_execution.rs +++ b/crates/goose/src/agents/tool_execution.rs @@ -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; } } }