From 1abd7dd3c34b6dccb52f2eda30abde34f2ca9bff Mon Sep 17 00:00:00 2001 From: Jack Amadeo Date: Wed, 1 Apr 2026 17:38:40 -0400 Subject: [PATCH] simplify --- crates/goose/src/agents/agent.rs | 62 +------ crates/goose/src/agents/agent_driver.rs | 55 ++---- crates/goose/src/agents/agent_loop.rs | 104 ++++-------- crates/goose/src/agents/tool_execution.rs | 195 +++++++++++++++------- 4 files changed, 172 insertions(+), 244 deletions(-) diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index bf0d76a323..48ff369b39 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -12,7 +12,7 @@ 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, DECLINED_RESPONSE}; +use super::tool_execution::ToolCallResult; use crate::action_required_manager::ActionRequiredManager; use crate::agents::extension::{ExtensionConfig, ExtensionResult, ToolInfo}; use crate::agents::extension_manager::{ @@ -29,7 +29,6 @@ use crate::conversation::message::{ActionRequiredData, Message, MessageContent}; use crate::conversation::{debug_conversation_fix, fix_conversation, Conversation}; use crate::mcp_utils::ToolResult; use crate::permission::permission_inspector::PermissionInspector; -use crate::permission::permission_judge::PermissionCheckResult; use crate::permission::PermissionConfirmation; use crate::providers::base::{PermissionRouting, Provider}; use crate::recipe::{Author, Recipe, Response, Settings}; @@ -370,65 +369,6 @@ impl Agent { }) } - pub(super) async fn handle_approved_and_denied_tools( - &self, - permission_check_result: &PermissionCheckResult, - request_to_response_map: &mut HashMap, - cancel_token: Option, - session: &Session, - ) -> Result> { - let mut tool_futures: Vec<(String, ToolStream)> = Vec::new(); - - // Handle pre-approved and read-only tools - for request in &permission_check_result.approved { - if let Ok(tool_call) = request.tool_call.clone() { - let (req_id, tool_result) = self - .dispatch_tool_call( - tool_call, - request.id.clone(), - cancel_token.clone(), - session, - ) - .await; - - tool_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))) - } - }, - )); - } - } - - Self::handle_denied_tools(permission_check_result, request_to_response_map); - Ok(tool_futures) - } - - pub(super) fn handle_denied_tools( - permission_check_result: &PermissionCheckResult, - request_to_response_map: &mut HashMap, - ) { - for request in &permission_check_result.denied { - if let Some(response) = request_to_response_map.get_mut(&request.id) { - response.add_tool_response_with_metadata( - request.id.clone(), - Ok(CallToolResult::error(vec![rmcp::model::Content::text( - DECLINED_RESPONSE, - )])), - request.metadata.as_ref(), - ); - } - } - } - /// Get a reference count clone to the provider pub async fn provider(&self) -> Result, anyhow::Error> { match &*self.provider.lock().await { diff --git a/crates/goose/src/agents/agent_driver.rs b/crates/goose/src/agents/agent_driver.rs index a66ae274c8..49180d00b8 100644 --- a/crates/goose/src/agents/agent_driver.rs +++ b/crates/goose/src/agents/agent_driver.rs @@ -18,11 +18,9 @@ use crate::agents::types::SessionConfig; use crate::config::GooseMode; use crate::conversation::message::{Message, ToolRequest}; use crate::conversation::Conversation; -use crate::permission::permission_judge::PermissionCheckResult; use crate::providers::base::Provider; use crate::providers::errors::ProviderError; use crate::session::{Session, SessionManager}; -use crate::tool_inspection::InspectionResult; #[async_trait] impl LoopDriver for Agent { @@ -57,59 +55,26 @@ impl LoopDriver for Agent { .await } - async fn inspect_tools( - &self, - session_id: &str, - requests: &[ToolRequest], - messages: &[Message], - goose_mode: GooseMode, - ) -> Result> { - self.tool_inspection_manager - .inspect_tools(session_id, requests, messages, goose_mode) - .await - } - - fn process_permissions( - &self, - requests: &[ToolRequest], - inspection_results: &[InspectionResult], - ) -> Option { - self.tool_inspection_manager - .process_inspection_results_with_permission_inspector(requests, inspection_results) - } - - async fn handle_approved_and_denied_tools( - &self, - permission_check_result: &PermissionCheckResult, - request_to_response_map: &mut HashMap, - cancel_token: Option, - session: &Session, - ) -> Result> { - self.handle_approved_and_denied_tools( - permission_check_result, - request_to_response_map, - cancel_token, - session, - ) - .await - } - - fn handle_approval_tool_requests<'a>( + fn gate_and_dispatch_tools<'a>( &'a self, - tool_requests: &'a [ToolRequest], + session_id: &'a str, + requests: &'a [ToolRequest], + messages: &'a [Message], + goose_mode: GooseMode, tool_futures: &'a mut Vec<(String, ToolStream)>, request_to_response_map: &'a mut HashMap, cancel_token: Option, session: &'a Session, - inspection_results: &'a [InspectionResult], ) -> BoxStream<'a, Result> { - self.handle_approval_tool_requests( - tool_requests, + self.gate_and_dispatch_tools( + session_id, + requests, + messages, + goose_mode, tool_futures, request_to_response_map, cancel_token, session, - inspection_results, ) } diff --git a/crates/goose/src/agents/agent_loop.rs b/crates/goose/src/agents/agent_loop.rs index 5832fd17d1..31923e638e 100644 --- a/crates/goose/src/agents/agent_loop.rs +++ b/crates/goose/src/agents/agent_loop.rs @@ -32,11 +32,9 @@ use crate::conversation::message::{ Message, MessageContent, ProviderMetadata, SystemNotificationType, }; use crate::conversation::Conversation; -use crate::permission::permission_judge::PermissionCheckResult; use crate::providers::base::Provider; use crate::providers::errors::ProviderError; use crate::session::{Session, SessionManager}; -use crate::tool_inspection::InspectionResult; use crate::utils::is_token_cancelled; use rmcp::model::ErrorData; use tokio_util::sync::CancellationToken; @@ -68,41 +66,27 @@ pub trait LoopDriver: Send + Sync { session: &Session, ) -> (String, Result); - /// Run tool inspectors on a batch of requests. - async fn inspect_tools( - &self, - session_id: &str, - requests: &[crate::conversation::message::ToolRequest], - messages: &[Message], - goose_mode: GooseMode, - ) -> Result>; - - /// Process inspection results through the permission inspector. - fn process_permissions( - &self, - requests: &[crate::conversation::message::ToolRequest], - inspection_results: &[InspectionResult], - ) -> Option; - - /// Handle approved and denied tools (dispatch approved, mark denied). - async fn handle_approved_and_denied_tools( - &self, - permission_check_result: &PermissionCheckResult, - request_to_response_map: &mut HashMap, - cancel_token: Option, - session: &Session, - ) -> Result>; - - /// Handle tools that need user approval (interactive). - /// Returns a stream of messages (action-required prompts) and populates tool_futures. - fn handle_approval_tool_requests<'a>( + /// Inspect, permission-check, and dispatch backend tool requests. + /// + /// This consolidates the full tool gate flow: + /// 1. Runs tool inspectors (adversary detection, permission checks, repetition) + /// 2. Partitions into approved / needs_approval / denied + /// 3. Dispatches approved tools → adds to `tool_futures` + /// 4. Marks denied tools with DECLINED_RESPONSE in `request_to_response_map` + /// 5. Streams action-required messages for tools needing user approval; + /// dispatches on confirmation → adds to `tool_futures` + /// + /// Returns a stream of action-required messages to yield to the UI. + fn gate_and_dispatch_tools<'a>( &'a self, - tool_requests: &'a [crate::conversation::message::ToolRequest], + session_id: &'a str, + requests: &'a [crate::conversation::message::ToolRequest], + messages: &'a [Message], + goose_mode: GooseMode, tool_futures: &'a mut Vec<(String, ToolStream)>, request_to_response_map: &'a mut HashMap, cancel_token: Option, session: &'a Session, - inspection_results: &'a [InspectionResult], ) -> BoxStream<'a, Result>; /// Drain pending elicitation messages for a session. @@ -393,6 +377,15 @@ pub fn run_agent_loop<'a>( // Track extension install requests let mut enable_extension_request_ids: Vec = vec![]; + // Track extension install requests + for request in &backend_requests { + if let Ok(tool_call) = &request.tool_call { + if tool_call.name == MANAGE_EXTENSIONS_TOOL_NAME_COMPLETE { + enable_extension_request_ids.push(request.id.clone()); + } + } + } + if goose_mode == GooseMode::Chat { for request in backend_requests.iter() { if let Some(response) = request_to_response_map.get_mut(&request.id) { @@ -404,55 +397,18 @@ pub fn run_agent_loop<'a>( } } } else { - let inspection_results = driver.inspect_tools( + let mut approval_stream = driver.gate_and_dispatch_tools( &session_config.id, &backend_requests, conversation.messages(), goose_mode, - ).await?; - - let permission_check_result = driver.process_permissions( - &backend_requests, - &inspection_results, - ).unwrap_or_else(|| { - let mut result = PermissionCheckResult { - approved: vec![], - needs_approval: vec![], - denied: vec![], - }; - result.needs_approval.extend(backend_requests.iter().cloned()); - result - }); - - for request in &backend_requests { - if let Ok(tool_call) = &request.tool_call { - if tool_call.name == MANAGE_EXTENSIONS_TOOL_NAME_COMPLETE { - enable_extension_request_ids.push(request.id.clone()); - } - } - } - - let mut backend_futures = driver.handle_approved_and_denied_tools( - &permission_check_result, + &mut tool_futures, &mut request_to_response_map, cancel_token.clone(), &session, - ).await?; - tool_futures.append(&mut backend_futures); - - { - let mut tool_approval_stream = driver.handle_approval_tool_requests( - &permission_check_result.needs_approval, - &mut tool_futures, - &mut request_to_response_map, - cancel_token.clone(), - &session, - &inspection_results, - ); - - while let Some(msg) = tool_approval_stream.try_next().await? { - yield AgentEvent::Message(msg); - } + ); + while let Some(msg) = approval_stream.try_next().await? { + yield AgentEvent::Message(msg); } } diff --git a/crates/goose/src/agents/tool_execution.rs b/crates/goose/src/agents/tool_execution.rs index a0fbbd9299..b5d0c11513 100644 --- a/crates/goose/src/agents/tool_execution.rs +++ b/crates/goose/src/agents/tool_execution.rs @@ -56,7 +56,9 @@ impl From> for ToolCallResult { use super::agent::{tool_stream, ToolStream}; use crate::agents::Agent; +use crate::config::GooseMode; use crate::conversation::message::{Message, ToolRequest}; +use crate::permission::permission_judge::PermissionCheckResult; use crate::session::Session; use crate::tool_inspection::get_security_finding_id_from_results; @@ -74,90 +76,155 @@ pub const CHAT_MODE_TOOL_SKIPPED_RESPONSE: &str = "Let the user know the tool ca If needed, adjust the explanation based on user preferences or questions."; impl Agent { - pub(crate) fn handle_approval_tool_requests<'a>( + /// Consolidated tool gate flow: inspect → permission-check → dispatch/deny/approve. + /// + /// Approved tools are dispatched immediately and added to `tool_futures`. + /// Denied tools get a DECLINED_RESPONSE in `request_to_response_map`. + /// Tools needing approval produce action-required messages in the returned stream; + /// once confirmed, they're dispatched and added to `tool_futures`. + pub(crate) fn gate_and_dispatch_tools<'a>( &'a self, - tool_requests: &'a [ToolRequest], + session_id: &'a str, + requests: &'a [ToolRequest], + messages: &'a [Message], + goose_mode: GooseMode, tool_futures: &'a mut Vec<(String, ToolStream)>, request_to_response_map: &'a mut HashMap, - cancellation_token: Option, + cancel_token: Option, session: &'a Session, - inspection_results: &'a [crate::tool_inspection::InspectionResult], ) -> BoxStream<'a, anyhow::Result> { try_stream! { - for request in tool_requests.iter() { - if let Ok(tool_call) = request.tool_call.clone() { - let security_message = inspection_results.iter() - .find(|result| result.tool_request_id == request.id) - .and_then(|result| { - if let crate::tool_inspection::InspectionAction::RequireApproval(Some(message)) = &result.action { - Some(message.clone()) - } else { - None - } - }); + let inspection_results = self + .tool_inspection_manager + .inspect_tools(session_id, requests, messages, goose_mode) + .await?; - let confirmation_rx = self.tool_confirmation_router.register(request.id.clone()).await; + let permission_check_result = self + .tool_inspection_manager + .process_inspection_results_with_permission_inspector(requests, &inspection_results) + .unwrap_or_else(|| { + let mut result = PermissionCheckResult { + approved: vec![], + needs_approval: vec![], + denied: vec![], + }; + result.needs_approval.extend(requests.iter().cloned()); + result + }); - let action_required_msg = Message::assistant() - .with_action_required( + // Dispatch approved tools + for request in &permission_check_result.approved { + if let Ok(tool_call) = request.tool_call.clone() { + let (req_id, tool_result) = self + .dispatch_tool_call( + tool_call, + request.id.clone(), + cancel_token.clone(), + session, + ) + .await; + tool_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))) + } + }, + )); + } + } + + // Mark denied tools + for request in &permission_check_result.denied { + if let Some(response) = request_to_response_map.get_mut(&request.id) { + response.add_tool_response_with_metadata( request.id.clone(), - tool_call.name.to_string().clone(), - tool_call.arguments.clone().unwrap_or_default(), - security_message, - ) - .user_only(); - yield action_required_msg; - - let confirmation = confirmation_rx.await - .map_err(|_| anyhow::anyhow!("Confirmation channel closed for request {}", request.id))?; - - 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" + Ok(CallToolResult::error(vec![Content::text(DECLINED_RESPONSE)])), + request.metadata.as_ref(), ); } + } - 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; + // Handle tools needing approval (yields action-required messages) + for request in permission_check_result.needs_approval.iter() { + if let Ok(tool_call) = request.tool_call.clone() { + let security_message = inspection_results.iter() + .find(|result| result.tool_request_id == request.id) + .and_then(|result| { + if let crate::tool_inspection::InspectionAction::RequireApproval(Some(message)) = &result.action { + Some(message.clone()) + } else { + None + } + }); - tool_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)), - ), - })); + let confirmation_rx = self.tool_confirmation_router.register(request.id.clone()).await; - if confirmation.permission == Permission::AlwaysAllow { - self.tool_inspection_manager - .update_permission_manager(&tool_call.name, PermissionLevel::AlwaysAllow) - .await; - } - } else { - if let Some(response) = request_to_response_map.get_mut(&request.id) { - response.add_tool_response_with_metadata( + let action_required_msg = Message::assistant() + .with_action_required( request.id.clone(), - Ok(CallToolResult::error(vec![Content::text(DECLINED_RESPONSE)])), - request.metadata.as_ref(), + tool_call.name.to_string().clone(), + tool_call.arguments.clone().unwrap_or_default(), + security_message, + ) + .user_only(); + yield action_required_msg; + + let confirmation = confirmation_rx.await + .map_err(|_| anyhow::anyhow!("Confirmation channel closed for request {}", request.id))?; + + 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" ); } - if confirmation.permission == Permission::AlwaysDeny { - self.tool_inspection_manager - .update_permission_manager(&tool_call.name, PermissionLevel::NeverAllow) - .await; + if confirmation.permission == Permission::AllowOnce || confirmation.permission == Permission::AlwaysAllow { + let (req_id, tool_result) = self.dispatch_tool_call(tool_call.clone(), request.id.clone(), cancel_token.clone(), session).await; + + tool_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::AlwaysAllow { + self.tool_inspection_manager + .update_permission_manager(&tool_call.name, PermissionLevel::AlwaysAllow) + .await; + } + } else { + if let Some(response) = request_to_response_map.get_mut(&request.id) { + response.add_tool_response_with_metadata( + request.id.clone(), + Ok(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; + } } } } - } - }.boxed() + }.boxed() } - }