use async_stream::try_stream; use futures::stream::{self, BoxStream}; use futures::{Stream, StreamExt}; use rmcp::model::CallToolResult; use std::collections::HashMap; use std::future::Future; use tokio_util::sync::CancellationToken; use std::path::PathBuf; use crate::config::permission::PermissionLevel; use crate::mcp_utils::ToolResult; use crate::permission::Permission; use rmcp::model::{Content, ServerNotification}; /// Context passed through the tool call dispatch chain. pub struct ToolCallContext { pub session_id: String, pub working_dir: Option, pub tool_call_request_id: Option, } impl ToolCallContext { pub fn new( session_id: String, working_dir: Option, tool_call_request_id: Option, ) -> Self { Self { session_id, working_dir, tool_call_request_id, } } pub fn working_dir_str(&self) -> Option<&str> { self.working_dir.as_ref().and_then(|p| p.to_str()) } } // ToolCallResult combines the result of a tool call with an optional notification stream that // can be used to receive notifications from the tool. pub struct ToolCallResult { pub result: Box> + Send + Unpin>, pub notification_stream: Option + Send + Unpin>>, } impl From> for ToolCallResult { fn from(result: ToolResult) -> Self { Self { result: Box::new(futures::future::ready(result)), notification_stream: None, } } } use super::agent::{tool_stream, ToolStream}; use crate::agents::Agent; use crate::conversation::message::{Message, ToolRequest}; use crate::session::Session; use crate::tool_inspection::get_security_finding_id_from_results; pub const DECLINED_RESPONSE: &str = "The user has declined to run this tool. \ DO NOT attempt to call this tool again. \ If there are no alternative methods to proceed, clearly explain the situation and STOP."; pub const CHAT_MODE_TOOL_SKIPPED_RESPONSE: &str = "Let the user know the tool call was skipped in goose chat mode. \ DO NOT apologize for skipping the tool call. DO NOT say sorry. \ Provide an explanation of what the tool call would do, structured as a \ plan for the user. Again, DO NOT apologize. \ **Example Plan:**\n \ 1. **Identify Task Scope** - Determine the purpose and expected outcome.\n \ 2. **Outline Steps** - Break down the steps.\n \ If needed, adjust the explanation based on user preferences or questions."; impl Agent { pub(crate) fn handle_approval_tool_requests<'a>( &'a self, tool_requests: &'a [ToolRequest], tool_futures: &'a mut Vec<(String, ToolStream)>, request_to_response_map: &'a mut HashMap, cancellation_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 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(), 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::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; 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() } pub(crate) fn handle_frontend_tool_request<'a>( &'a self, tool_request: &'a ToolRequest, message_tool_response: &'a mut Message, ) -> BoxStream<'a, anyhow::Result> { try_stream! { if let Ok(tool_call) = tool_request.tool_call.clone() { if self.is_frontend_tool(&tool_call.name).await { yield Message::assistant().with_frontend_tool_request( tool_request.id.clone(), Ok(tool_call.clone()) ); if let Some((id, result)) = self.tool_result_rx.lock().await.recv().await { message_tool_response.add_tool_response_with_metadata( id, result, tool_request.metadata.as_ref(), ); } } } } .boxed() } }