mirror of
https://github.com/aaif-goose/goose.git
synced 2026-07-03 14:10:03 +02:00
simplify
This commit is contained in:
Jack Amadeo
@@ -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<String, Message>,
|
||||
cancel_token: Option<tokio_util::sync::CancellationToken>,
|
||||
session: &Session,
|
||||
) -> Result<Vec<(String, ToolStream)>> {
|
||||
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<String, Message>,
|
||||
) {
|
||||
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<Arc<dyn Provider>, anyhow::Error> {
|
||||
match &*self.provider.lock().await {
|
||||
|
||||
@@ -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<Vec<InspectionResult>> {
|
||||
self.tool_inspection_manager
|
||||
.inspect_tools(session_id, requests, messages, goose_mode)
|
||||
.await
|
||||
}
|
||||
|
||||
fn process_permissions(
|
||||
&self,
|
||||
requests: &[ToolRequest],
|
||||
inspection_results: &[InspectionResult],
|
||||
) -> Option<PermissionCheckResult> {
|
||||
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<String, Message>,
|
||||
cancel_token: Option<CancellationToken>,
|
||||
session: &Session,
|
||||
) -> Result<Vec<(String, ToolStream)>> {
|
||||
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<String, Message>,
|
||||
cancel_token: Option<CancellationToken>,
|
||||
session: &'a Session,
|
||||
inspection_results: &'a [InspectionResult],
|
||||
) -> BoxStream<'a, Result<Message>> {
|
||||
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,
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
@@ -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<ToolCallResult, ErrorData>);
|
||||
|
||||
/// 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<Vec<InspectionResult>>;
|
||||
|
||||
/// Process inspection results through the permission inspector.
|
||||
fn process_permissions(
|
||||
&self,
|
||||
requests: &[crate::conversation::message::ToolRequest],
|
||||
inspection_results: &[InspectionResult],
|
||||
) -> Option<PermissionCheckResult>;
|
||||
|
||||
/// 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<String, Message>,
|
||||
cancel_token: Option<CancellationToken>,
|
||||
session: &Session,
|
||||
) -> Result<Vec<(String, ToolStream)>>;
|
||||
|
||||
/// 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<String, Message>,
|
||||
cancel_token: Option<CancellationToken>,
|
||||
session: &'a Session,
|
||||
inspection_results: &'a [InspectionResult],
|
||||
) -> BoxStream<'a, Result<Message>>;
|
||||
|
||||
/// 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<String> = 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);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -56,7 +56,9 @@ impl From<ToolResult<rmcp::model::CallToolResult>> 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<String, Message>,
|
||||
cancellation_token: Option<CancellationToken>,
|
||||
cancel_token: Option<CancellationToken>,
|
||||
session: &'a Session,
|
||||
inspection_results: &'a [crate::tool_inspection::InspectionResult],
|
||||
) -> BoxStream<'a, anyhow::Result<Message>> {
|
||||
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()
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user