diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 55507e4039..d412517ecf 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::future::Future; use std::pin::Pin; use std::sync::Arc; @@ -35,8 +35,8 @@ use crate::agents::tool_vectordb::generate_table_id; use crate::agents::types::SessionConfig; use crate::agents::types::{FrontendTool, ToolResultReceiver}; use crate::config::{Config, ExtensionConfigManager, PermissionManager}; -use crate::message::{push_message, Message}; -use crate::permission::permission_judge::check_tool_permissions; +use crate::message::{push_message, Message, ToolRequest}; +use crate::permission::permission_judge::{check_tool_permissions, PermissionCheckResult}; use crate::permission::PermissionConfirmation; use crate::providers::base::Provider; use crate::providers::errors::ProviderError; @@ -61,6 +61,26 @@ use crate::conversation_fixer::{debug_conversation_fix, ConversationFixer}; const DEFAULT_MAX_TURNS: u32 = 1000; +/// Context needed for the reply function +pub struct ReplyContext { + pub messages: Vec, + pub tools: Vec, + pub toolshim_tools: Vec, + pub system_prompt: String, + pub goose_mode: String, + pub initial_messages: Vec, + pub config: &'static Config, +} + +/// Result of processing tool requests +pub struct ToolProcessingResult { + pub frontend_requests: Vec, + pub remaining_requests: Vec, + pub filtered_response: Message, + pub readonly_tools: HashSet, + pub regular_tools: HashSet, +} + /// The main goose Agent pub struct Agent { pub(super) provider: Mutex>>, @@ -162,17 +182,6 @@ impl Agent { *tool_monitor = Some(ToolMonitor::new(max_repetitions)); } - pub async fn get_tool_stats(&self) -> Option> { - let tool_monitor = self.tool_monitor.lock().await; - tool_monitor.as_ref().map(|monitor| monitor.get_stats()) - } - - pub async fn reset_tool_monitor(&self) { - if let Some(monitor) = self.tool_monitor.lock().await.as_mut() { - monitor.reset(); - } - } - /// Reset the retry attempts counter to 0 pub async fn reset_retry_attempts(&self) { self.retry_manager.reset_attempts().await; @@ -208,6 +217,119 @@ impl Agent { } } + async fn prepare_reply_context( + &self, + unfixed_messages: &[Message], + session: &Option, + ) -> Result { + let (messages, issues) = ConversationFixer::fix_conversation(Vec::from(unfixed_messages)); + if !issues.is_empty() { + tracing::warn!( + "Conversation issue fixed: {}", + debug_conversation_fix(&messages, unfixed_messages, &issues) + ); + } + let initial_messages = messages.clone(); + let config = Config::global(); + + let (tools, toolshim_tools, system_prompt) = self.prepare_tools_and_prompt().await?; + let goose_mode = Self::determine_goose_mode(session.as_ref(), config); + + Ok(ReplyContext { + messages, + tools, + toolshim_tools, + system_prompt, + goose_mode, + initial_messages, + config, + }) + } + + /// Process tool requests by categorizing them and recording them in the router selector + async fn process_tool_requests( + &self, + response: &Message, + tools: &[rmcp::model::Tool], + ) -> ToolProcessingResult { + let (readonly_tools, regular_tools) = Self::categorize_tools_by_annotation(tools); + + // Categorize tool requests + let (frontend_requests, remaining_requests, filtered_response) = + self.categorize_tool_requests(response).await; + + // Record tool calls in the router selector + let selector = self.router_tool_selector.lock().await.clone(); + if let Some(selector) = selector { + for request in &frontend_requests { + if let Ok(tool_call) = &request.tool_call { + if let Err(e) = selector.record_tool_call(&tool_call.name).await { + error!("Failed to record frontend tool call: {}", e); + } + } + } + for request in &remaining_requests { + if let Ok(tool_call) = &request.tool_call { + if let Err(e) = selector.record_tool_call(&tool_call.name).await { + error!("Failed to record tool call: {}", e); + } + } + } + } + + ToolProcessingResult { + frontend_requests, + remaining_requests, + filtered_response, + readonly_tools, + regular_tools, + } + } + + async fn handle_approved_and_denied_tools( + &self, + permission_check_result: &PermissionCheckResult, + message_tool_response: Arc>, + cancel_token: Option, + ) -> 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()) + .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))) + } + }, + )); + } + } + + // Handle denied tools + for request in &permission_check_result.denied { + let mut response = message_tool_response.lock().await; + *response = response.clone().with_tool_response( + request.id.clone(), + Ok(vec![rmcp::model::Content::text(DECLINED_RESPONSE)]), + ); + } + + Ok(tool_futures) + } + /// Set the scheduler service for this agent pub async fn set_scheduler(&self, scheduler: Arc) { let mut scheduler_service = self.scheduler_service.lock().await; @@ -232,24 +354,6 @@ impl Agent { self.frontend_tools.lock().await.get(name).cloned() } - /// Get all tools from all clients with proper prefixing - pub async fn get_prefixed_tools(&self) -> ExtensionResult> { - let mut tools = self - .extension_manager - .read() - .await - .get_prefixed_tools(None) - .await?; - - // Add frontend tools directly - they don't need prefixing since they're already uniquely named - let frontend_tools = self.frontend_tools.lock().await; - for frontend_tool in frontend_tools.values() { - tools.push(frontend_tool.tool.clone()); - } - - Ok(tools) - } - pub async fn add_final_output_tool(&self, response: Response) { let mut final_output_tool = self.final_output_tool.lock().await; let created_final_output_tool = FinalOutputTool::new(response); @@ -728,22 +832,21 @@ impl Agent { session: Option, cancel_token: Option, ) -> Result>> { - let (mut messages, issues) = - ConversationFixer::fix_conversation(Vec::from(unfixed_messages)); - if !issues.is_empty() { - tracing::warn!( - "Conversation issue fixed: {}", - debug_conversation_fix(&messages, unfixed_messages, &issues) - ); - } - let initial_messages = messages.clone(); + let context = self + .prepare_reply_context(unfixed_messages, &session) + .await?; + let ReplyContext { + mut messages, + mut tools, + mut toolshim_tools, + mut system_prompt, + goose_mode, + initial_messages, + config, + } = context; + let reply_span = tracing::Span::current(); self.reset_retry_attempts().await; - let config = Config::global(); - - let (mut tools, mut toolshim_tools, mut system_prompt) = - self.prepare_tools_and_prompt().await?; - let goose_mode = Self::determine_goose_mode(session.as_ref(), config); if let Some(content) = messages .last() @@ -792,8 +895,7 @@ impl Agent { &messages, &tools, &toolshim_tools, - ) - .await?; + ).await?; let mut added_message = false; let mut messages_to_add = Vec::new(); @@ -836,33 +938,14 @@ impl Agent { } if let Some(response) = response { - let (tools_with_readonly_annotation, tools_without_annotation) = - Self::categorize_tools_by_annotation(&tools); - - // Categorize tool requests - let (frontend_requests, remaining_requests, filtered_response) = - self.categorize_tool_requests(&response).await; - - // Record tool calls in the router selector - let selector = self.router_tool_selector.lock().await.clone(); - if let Some(selector) = selector { - for request in &frontend_requests { - if let Ok(tool_call) = &request.tool_call { - if let Err(e) = selector.record_tool_call(&tool_call.name).await - { - error!("Failed to record frontend tool call: {}", e); - } - } - } - for request in &remaining_requests { - if let Ok(tool_call) = &request.tool_call { - if let Err(e) = selector.record_tool_call(&tool_call.name).await - { - error!("Failed to record tool call: {}", e); - } - } - } - } + let tool_result = self.process_tool_requests(&response, &tools).await; + let ToolProcessingResult { + frontend_requests, + remaining_requests, + filtered_response, + readonly_tools, + regular_tools, + } = tool_result; yield AgentEvent::Message(filtered_response.clone()); tokio::task::yield_now().await; @@ -901,47 +984,17 @@ impl Agent { check_tool_permissions( &remaining_requests, &mode, - tools_with_readonly_annotation.clone(), - tools_without_annotation.clone(), + readonly_tools.clone(), + regular_tools.clone(), &mut permission_manager, self.provider().await?, - ) - .await; + ).await; - 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()) - .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)), - ), - }, - )); - } - } - - for request in &permission_check_result.denied { - let mut response = message_tool_response.lock().await; - *response = response.clone().with_tool_response( - request.id.clone(), - Ok(vec![Content::text(DECLINED_RESPONSE)]), - ); - } + let mut tool_futures = self.handle_approved_and_denied_tools( + &permission_check_result, + message_tool_response.clone(), + cancel_token.clone() + ).await?; let tool_futures_arc = Arc::new(Mutex::new(tool_futures)); diff --git a/crates/goose/src/agents/subagent.rs b/crates/goose/src/agents/subagent.rs index f9cf85d64e..939916e0af 100644 --- a/crates/goose/src/agents/subagent.rs +++ b/crates/goose/src/agents/subagent.rs @@ -91,11 +91,6 @@ impl SubAgent { Ok(subagent) } - /// Get the current status of the subagent - pub async fn get_status(&self) -> SubAgentStatus { - self.status.read().await.clone() - } - /// Update the status of the subagent async fn set_status(&self, status: SubAgentStatus) { // Update the status first, then release the lock @@ -105,26 +100,6 @@ impl SubAgent { } // Write lock is released here! } - /// Get current progress information - pub async fn get_progress(&self) -> SubAgentProgress { - let status = self.get_status().await; - let turn_count = *self.turn_count.lock().await; - - SubAgentProgress { - subagent_id: self.id.clone(), - status: status.clone(), - message: match &status { - SubAgentStatus::Ready => "Ready to process messages".to_string(), - SubAgentStatus::Processing => "Processing request...".to_string(), - SubAgentStatus::Completed(msg) => msg.clone(), - SubAgentStatus::Terminated => "Subagent terminated".to_string(), - }, - turn: turn_count, - max_turns: self.config.max_turns, - timestamp: Utc::now(), - } - } - /// Process a message and generate a response using the subagent's provider #[instrument(skip(self, message))] pub async fn reply_subagent( @@ -283,37 +258,16 @@ impl SubAgent { } /// Add a message to the conversation (for tracking agent responses) - pub async fn add_message(&self, message: Message) { + async fn add_message(&self, message: Message) { let mut conversation = self.conversation.lock().await; conversation.push(message); } /// Get the full conversation history - pub async fn get_conversation(&self) -> Vec { + async fn get_conversation(&self) -> Vec { self.conversation.lock().await.clone() } - /// Check if the subagent has completed its task - pub async fn is_completed(&self) -> bool { - matches!( - self.get_status().await, - SubAgentStatus::Completed(_) | SubAgentStatus::Terminated - ) - } - - /// Terminate the subagent - pub async fn terminate(&self) -> Result<(), anyhow::Error> { - debug!("Terminating subagent {}", self.id); - self.set_status(SubAgentStatus::Terminated).await; - Ok(()) - } - - /// Filter out subagent spawning tools to prevent infinite recursion - fn _filter_subagent_tools(tools: Vec) -> Vec { - // TODO: add this in subagent loop - tools - } - /// Build the system prompt for the subagent using the template async fn build_system_prompt(&self, available_tools: &[Tool]) -> Result { let mut context = HashMap::new();