use std::collections::HashMap; use std::future::Future; use std::pin::Pin; use std::sync::Arc; use anyhow::{anyhow, Context, Result}; use futures::stream::BoxStream; use futures::{stream, FutureExt, Stream, StreamExt, TryStreamExt}; use uuid::Uuid; use super::container::Container; use super::final_output_tool::FinalOutputTool; use super::platform_tools; 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}; use crate::agents::extension_manager::{get_parameter_names, ExtensionManager}; use crate::agents::extension_manager_extension::MANAGE_EXTENSIONS_TOOL_NAME_COMPLETE; use crate::agents::final_output_tool::{FINAL_OUTPUT_CONTINUATION_MESSAGE, FINAL_OUTPUT_TOOL_NAME}; use crate::agents::platform_tools::PLATFORM_MANAGE_SCHEDULE_TOOL_NAME; use crate::agents::prompt_manager::PromptManager; use crate::agents::retry::{RetryManager, RetryResult}; use crate::agents::subagent_task_config::TaskConfig; use crate::agents::subagent_tool::{ create_subagent_tool, handle_subagent_tool, SUBAGENT_TOOL_NAME, }; use crate::agents::types::{FrontendTool, SessionConfig, SharedProvider, ToolResultReceiver}; use crate::config::permission::PermissionManager; use crate::config::{get_enabled_extensions, Config, GooseMode}; use crate::context_mgmt::{ check_if_compaction_needed, compact_messages, DEFAULT_COMPACTION_THRESHOLD, }; use crate::conversation::message::{ ActionRequiredData, Message, MessageContent, ProviderMetadata, SystemNotificationType, ToolRequest, }; use crate::conversation::tool_result_serde::call_tool_result; 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::Provider; use crate::providers::errors::ProviderError; use crate::recipe::{Author, Recipe, Response, Settings, SubRecipe}; use crate::scheduler_trait::SchedulerTrait; use crate::security::security_inspector::SecurityInspector; use crate::session::extension_data::{EnabledExtensionsState, ExtensionState}; use crate::session::{Session, SessionManager, SessionType}; use crate::tool_inspection::ToolInspectionManager; use crate::tool_monitor::RepetitionInspector; use crate::utils::is_token_cancelled; use regex::Regex; use rmcp::model::{ CallToolRequestParams, CallToolResult, Content, ErrorCode, ErrorData, GetPromptResult, Prompt, ServerNotification, Tool, }; use serde_json::Value; use tokio::sync::{mpsc, Mutex}; use tokio_util::sync::CancellationToken; use tracing::{debug, error, info, instrument, warn}; const DEFAULT_MAX_TURNS: u32 = 1000; const COMPACTION_THINKING_TEXT: &str = "goose is compacting the conversation..."; /// Context needed for the reply function pub struct ReplyContext { pub conversation: Conversation, pub tools: Vec, pub toolshim_tools: Vec, pub system_prompt: String, pub goose_mode: GooseMode, pub tool_call_cut_off: usize, pub initial_messages: Vec, } pub struct ToolCategorizeResult { pub frontend_requests: Vec, pub remaining_requests: Vec, pub filtered_response: Message, } #[derive(Debug, Clone, serde::Serialize, utoipa::ToSchema)] pub struct ExtensionLoadResult { pub name: String, pub success: bool, #[serde(skip_serializing_if = "Option::is_none")] pub error: Option, } #[derive(Clone)] pub struct AgentConfig { pub session_manager: Arc, pub permission_manager: Arc, pub scheduler_service: Option>, pub goose_mode: GooseMode, pub disable_session_naming: bool, } impl AgentConfig { pub fn new( session_manager: Arc, permission_manager: Arc, scheduler_service: Option>, goose_mode: GooseMode, disable_session_naming: bool, ) -> Self { Self { session_manager, permission_manager, scheduler_service, goose_mode, disable_session_naming, } } } /// The main goose Agent pub struct Agent { pub(super) provider: SharedProvider, pub config: AgentConfig, pub extension_manager: Arc, pub(super) sub_recipes: Mutex>, pub(super) final_output_tool: Arc>>, 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(super) tool_result_tx: mpsc::Sender<(String, ToolResult)>, pub(super) tool_result_rx: ToolResultReceiver, pub(super) retry_manager: RetryManager, pub(super) tool_inspection_manager: ToolInspectionManager, container: Mutex>, } #[derive(Clone, Debug)] pub enum AgentEvent { Message(Message), McpNotification((String, ServerNotification)), ModelChange { model: String, mode: String }, HistoryReplaced(Conversation), } impl Default for Agent { fn default() -> Self { Self::new() } } pub enum ToolStreamItem { Message(ServerNotification), Result(T), } pub type ToolStream = Pin>> + Send>>; // tool_stream combines a stream of ServerNotifications with a future representing the // final result of the tool call. MCP notifications are not request-scoped, but // this lets us capture all notifications emitted during the tool call for // simpler consumption pub fn tool_stream(rx: S, done: F) -> ToolStream where S: Stream + Send + Unpin + 'static, F: Future> + Send + 'static, { Box::pin(async_stream::stream! { tokio::pin!(done); let mut rx = rx; loop { tokio::select! { Some(msg) = rx.next() => { yield ToolStreamItem::Message(msg); } r = &mut done => { yield ToolStreamItem::Result(r); break; } } } }) } impl Agent { pub fn new() -> Self { Self::with_config(AgentConfig::new( Arc::new(SessionManager::instance()), PermissionManager::instance(), None, Config::global().get_goose_mode().unwrap_or(GooseMode::Auto), Config::global() .get_goose_disable_session_naming() .unwrap_or(false), )) } 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)); let session_manager = Arc::clone(&config.session_manager); let permission_manager = Arc::clone(&config.permission_manager); Self { provider: provider.clone(), config, extension_manager: Arc::new(ExtensionManager::new(provider.clone(), session_manager)), sub_recipes: Mutex::new(HashMap::new()), final_output_tool: Arc::new(Mutex::new(None)), 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_result_tx: tool_tx, tool_result_rx: Arc::new(Mutex::new(tool_rx)), retry_manager: RetryManager::new(), tool_inspection_manager: Self::create_tool_inspection_manager(permission_manager), container: Mutex::new(None), } } /// Create a tool inspection manager with default inspectors fn create_tool_inspection_manager( permission_manager: Arc, ) -> ToolInspectionManager { let mut tool_inspection_manager = ToolInspectionManager::new(); // Add security inspector (highest priority - runs first) tool_inspection_manager.add_inspector(Box::new(SecurityInspector::new())); // Add permission inspector (medium-high priority) tool_inspection_manager.add_inspector(Box::new(PermissionInspector::new( std::collections::HashSet::new(), // readonly tools - will be populated from extension manager std::collections::HashSet::new(), // regular tools - will be populated from extension manager permission_manager, ))); // Add repetition inspector (lower priority - basic repetition checking) tool_inspection_manager.add_inspector(Box::new(RepetitionInspector::new(None))); tool_inspection_manager } /// Reset the retry attempts counter to 0 pub async fn reset_retry_attempts(&self) { self.retry_manager.reset_attempts().await; } /// Increment the retry attempts counter and return the new value pub async fn increment_retry_attempts(&self) -> u32 { self.retry_manager.increment_attempts().await } /// Get the current retry attempts count pub async fn get_retry_attempts(&self) -> u32 { self.retry_manager.get_attempts().await } async fn handle_retry_logic( &self, messages: &mut Conversation, session_config: &SessionConfig, initial_messages: &[Message], ) -> Result { let result = self .retry_manager .handle_retry_logic( messages, session_config, initial_messages, &self.final_output_tool, ) .await?; match result { RetryResult::Retried => Ok(true), RetryResult::Skipped | RetryResult::MaxAttemptsReached | RetryResult::SuccessChecksPassed => Ok(false), } } async fn drain_elicitation_messages(&self, session_id: &str) -> Vec { let mut messages = Vec::new(); let manager = self.config.session_manager.clone(); let mut elicitation_rx = ActionRequiredManager::global().request_rx.lock().await; while let Ok(mut elicitation_message) = elicitation_rx.try_recv() { if elicitation_message.id.is_none() { elicitation_message = elicitation_message.with_generated_id(); } if let Err(e) = manager.add_message(session_id, &elicitation_message).await { warn!("Failed to save elicitation message to session: {}", e); } messages.push(elicitation_message); } messages } async fn prepare_reply_context( &self, session_id: &str, unfixed_conversation: Conversation, working_dir: &std::path::Path, ) -> Result { let unfixed_messages = unfixed_conversation.messages().clone(); let (conversation, issues) = fix_conversation(unfixed_conversation.clone()); if !issues.is_empty() { debug!( "Conversation issue fixed: {}", debug_conversation_fix( unfixed_messages.as_slice(), conversation.messages(), &issues ) ); } let initial_messages = conversation.messages().clone(); let (tools, toolshim_tools, system_prompt) = self .prepare_tools_and_prompt(session_id, working_dir) .await?; Ok(ReplyContext { conversation, tools, toolshim_tools, system_prompt, goose_mode: self.config.goose_mode, tool_call_cut_off: Config::global() .get_param::("GOOSE_TOOL_CALL_CUTOFF") .unwrap_or(10), initial_messages, }) } async fn categorize_tools( &self, response: &Message, tools: &[rmcp::model::Tool], ) -> ToolCategorizeResult { // Categorize tool requests let (frontend_requests, remaining_requests, filtered_response) = self.categorize_tool_requests(response, tools).await; ToolCategorizeResult { frontend_requests, remaining_requests, filtered_response, } } async fn handle_approved_and_denied_tools( &self, permission_check_result: &PermissionCheckResult, request_to_response_map: &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).await; Ok(tool_futures) } async fn handle_denied_tools( permission_check_result: &PermissionCheckResult, request_to_response_map: &HashMap>>, ) { for request in &permission_check_result.denied { 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(CallToolResult { content: vec![rmcp::model::Content::text(DECLINED_RESPONSE)], structured_content: None, is_error: Some(true), meta: None, }), 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 { Some(provider) => Ok(Arc::clone(provider)), None => Err(anyhow!("Provider not set")), } } /// When set, all stdio extensions will be started via `docker exec` in the specified container. pub async fn set_container(&self, container: Option) { *self.container.lock().await = container.clone(); } pub async fn container(&self) -> Option { self.container.lock().await.clone() } /// Check if a tool is a frontend tool pub async fn is_frontend_tool(&self, name: &str) -> bool { self.frontend_tools.lock().await.contains_key(name) } /// Get a reference to a frontend tool pub async fn get_frontend_tool(&self, name: &str) -> Option { self.frontend_tools.lock().await.get(name).cloned() } 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); let final_output_system_prompt = created_final_output_tool.system_prompt(); *final_output_tool = Some(created_final_output_tool); self.extend_system_prompt(final_output_system_prompt).await; } pub async fn add_sub_recipes(&self, sub_recipes_to_add: Vec) { let mut sub_recipes = self.sub_recipes.lock().await; for sr in sub_recipes_to_add { sub_recipes.insert(sr.name.clone(), sr); } } pub async fn apply_recipe_components( &self, sub_recipes: Option>, response: Option, include_final_output: bool, ) { if let Some(sub_recipes) = sub_recipes { self.add_sub_recipes(sub_recipes).await; } if include_final_output { if let Some(response) = response { self.add_final_output_tool(response).await; } } } /// Dispatch a single tool call to the appropriate client #[instrument(skip(self, tool_call, request_id), fields(input, output))] pub async fn dispatch_tool_call( &self, tool_call: CallToolRequestParams, request_id: String, cancellation_token: Option, session: &Session, ) -> (String, Result) { // Prevent subagents from creating other subagents if session.session_type == SessionType::SubAgent && tool_call.name == SUBAGENT_TOOL_NAME { return ( request_id, Err(ErrorData::new( ErrorCode::INVALID_REQUEST, "Subagents cannot create other subagents".to_string(), None, )), ); } if tool_call.name == PLATFORM_MANAGE_SCHEDULE_TOOL_NAME { let arguments = tool_call .arguments .map(Value::Object) .unwrap_or(Value::Object(serde_json::Map::new())); let result = self .handle_schedule_management(arguments, request_id.clone()) .await; let wrapped_result = result.map(|content| CallToolResult { content, structured_content: None, is_error: Some(false), meta: None, }); return (request_id, Ok(ToolCallResult::from(wrapped_result))); } if tool_call.name == FINAL_OUTPUT_TOOL_NAME { return if let Some(final_output_tool) = self.final_output_tool.lock().await.as_mut() { let result = final_output_tool.execute_tool_call(tool_call.clone()).await; (request_id, Ok(result)) } else { ( request_id, Err(ErrorData::new( ErrorCode::INTERNAL_ERROR, "Final output tool not defined".to_string(), None, )), ) }; } debug!("WAITING_TOOL_START: {}", tool_call.name); let result: ToolCallResult = if tool_call.name == SUBAGENT_TOOL_NAME { let provider = match self.provider().await { Ok(p) => p, Err(_) => { return ( request_id, Err(ErrorData::new( ErrorCode::INTERNAL_ERROR, "Provider is required".to_string(), None, )), ); } }; let extensions = self.get_extension_configs().await; let max_turns_from_recipe = session .recipe .as_ref() .and_then(|r| r.settings.as_ref()) .and_then(|s| s.max_turns); let task_config = TaskConfig::new(provider, &session.id, &session.working_dir, extensions) .with_max_turns(max_turns_from_recipe); let sub_recipes = self.sub_recipes.lock().await.clone(); let arguments = tool_call .arguments .clone() .map(Value::Object) .unwrap_or(Value::Object(serde_json::Map::new())); handle_subagent_tool( &self.config, arguments, task_config, sub_recipes, session.working_dir.clone(), cancellation_token, ) } else if self.is_frontend_tool(&tool_call.name).await { // For frontend tools, return an error indicating we need frontend execution ToolCallResult::from(Err(ErrorData::new( ErrorCode::INTERNAL_ERROR, "Frontend tool execution required".to_string(), None, ))) } else { // Clone the result to ensure no references to extension_manager are returned let result = self .extension_manager .dispatch_tool_call( &session.id, tool_call.clone(), Some(session.working_dir.as_path()), cancellation_token.unwrap_or_default(), ) .await; result.unwrap_or_else(|e| { crate::posthog::emit_error( "tool_execution_failed", &format!("{}: {}", tool_call.name, e), ); // Try to downcast to ErrorData to avoid double wrapping let error_data = e.downcast::().unwrap_or_else(|e| { ErrorData::new(ErrorCode::INTERNAL_ERROR, e.to_string(), None) }); ToolCallResult::from(Err(error_data)) }) }; debug!("WAITING_TOOL_END: {}", tool_call.name); ( request_id, Ok(ToolCallResult { notification_stream: result.notification_stream, result: Box::new( result .result .map(super::large_response_handler::process_tool_response), ), }), ) } /// Save current extension state to session metadata /// Should be called after any extension add/remove operation pub async fn save_extension_state(&self, session: &SessionConfig) -> Result<()> { let extension_configs = self.extension_manager.get_extension_configs().await; let extensions_state = EnabledExtensionsState::new(extension_configs); let session_manager = self.config.session_manager.clone(); let mut session_data = session_manager.get_session(&session.id, false).await?; if let Err(e) = extensions_state.to_extension_data(&mut session_data.extension_data) { warn!("Failed to serialize extension state: {}", e); return Err(anyhow!("Extension state serialization failed: {}", e)); } session_manager .update(&session.id) .extension_data(session_data.extension_data) .apply() .await?; Ok(()) } /// Save current extension state to session by session_id pub async fn persist_extension_state(&self, session_id: &str) -> Result<()> { let extension_configs = self.extension_manager.get_extension_configs().await; let extensions_state = EnabledExtensionsState::new(extension_configs); let session_manager = self.config.session_manager.clone(); let session = session_manager.get_session(session_id, false).await?; let mut extension_data = session.extension_data.clone(); extensions_state .to_extension_data(&mut extension_data) .map_err(|e| anyhow!("Failed to serialize extension state: {}", e))?; session_manager .update(session_id) .extension_data(extension_data) .apply() .await?; Ok(()) } /// Load extensions from session into the agent /// Skips extensions that are already loaded /// Uses the session's working_dir for extension initialization pub async fn load_extensions_from_session( self: &Arc, session: &Session, ) -> Vec { let session_extensions = EnabledExtensionsState::from_extension_data(&session.extension_data); let enabled_configs = match session_extensions { Some(state) => state.extensions, None => { tracing::warn!( "No extensions found in session {}. This is unexpected.", session.id ); return vec![]; } }; let session_id = session.id.clone(); let extension_futures = enabled_configs .into_iter() .map(|config| { let config_clone = config.clone(); let agent_ref = self.clone(); let session_id_clone = session_id.clone(); async move { let name = config_clone.name().to_string(); if agent_ref .extension_manager .is_extension_enabled(&name) .await { tracing::debug!("Extension {} already loaded, skipping", name); return ExtensionLoadResult { name, success: true, error: None, }; } match agent_ref .add_extension(config_clone, &session_id_clone) .await { Ok(_) => ExtensionLoadResult { name, success: true, error: None, }, Err(e) => { let error_msg = e.to_string(); warn!("Failed to load extension {}: {}", name, error_msg); ExtensionLoadResult { name, success: false, error: Some(error_msg), } } } } }) .collect::>(); futures::future::join_all(extension_futures).await } pub async fn add_extension( &self, extension: ExtensionConfig, session_id: &str, ) -> ExtensionResult<()> { let session = self .config .session_manager .get_session(session_id, false) .await .map_err(|e| { crate::agents::extension::ExtensionError::SetupError(format!( "Failed to get session '{}': {}", session_id, e )) })?; let working_dir = Some(session.working_dir); match &extension { ExtensionConfig::Frontend { tools, instructions, .. } => { // For frontend tools, just store them in the frontend_tools map let mut frontend_tools = self.frontend_tools.lock().await; for tool in tools { let frontend_tool = FrontendTool { name: tool.name.to_string(), tool: tool.clone(), }; frontend_tools.insert(tool.name.to_string(), frontend_tool); } // Store instructions if provided, using "frontend" as the key let mut frontend_instructions = self.frontend_instructions.lock().await; if let Some(instructions) = instructions { *frontend_instructions = Some(instructions.clone()); } else { // Default frontend instructions if none provided *frontend_instructions = Some( "The following tools are provided directly by the frontend and will be executed by the frontend when called.".to_string(), ); } } _ => { let container = self.container.lock().await; self.extension_manager .add_extension(extension.clone(), working_dir, container.as_ref()) .await?; } } // Persist extension state after successful add self.persist_extension_state(session_id) .await .map_err(|e| { error!("Failed to persist extension state: {}", e); crate::agents::extension::ExtensionError::SetupError(format!( "Failed to persist extension state: {}", e )) })?; Ok(()) } pub async fn subagents_enabled(&self, session_id: &str) -> bool { if self.config.goose_mode != GooseMode::Auto { return false; } let context = self.extension_manager.get_context(); if matches!( context .session_manager .get_session(session_id, false) .await .ok() .map(|session| session.session_type), Some(SessionType::SubAgent) ) { return false; } !self .extension_manager .list_extensions() .await .map(|ext| ext.is_empty()) .unwrap_or(true) } pub async fn list_tools(&self, session_id: &str, extension_name: Option) -> Vec { let mut prefixed_tools = self .extension_manager .get_prefixed_tools(session_id, extension_name.clone()) .await .unwrap_or_default(); let subagents_enabled = self.subagents_enabled(session_id).await; if (extension_name.is_none() || extension_name.as_deref() == Some("platform")) && self.config.scheduler_service.is_some() { prefixed_tools.push(platform_tools::manage_schedule_tool()); } if extension_name.is_none() { if let Some(final_output_tool) = self.final_output_tool.lock().await.as_ref() { prefixed_tools.push(final_output_tool.tool()); } if subagents_enabled { let sub_recipes = self.sub_recipes.lock().await; let sub_recipes_vec: Vec<_> = sub_recipes.values().cloned().collect(); prefixed_tools.push(create_subagent_tool(&sub_recipes_vec)); } } prefixed_tools } pub async fn remove_extension(&self, name: &str, session_id: &str) -> Result<()> { self.extension_manager.remove_extension(name).await?; // Persist extension state after successful removal self.persist_extension_state(session_id) .await .map_err(|e| { error!("Failed to persist extension state: {}", e); anyhow!("Failed to persist extension state: {}", e) })?; Ok(()) } pub async fn list_extensions(&self) -> Vec { self.extension_manager .list_extensions() .await .expect("Failed to list extensions") } pub async fn get_extension_configs(&self) -> Vec { self.extension_manager.get_extension_configs().await } /// Handle a confirmation response for a tool request pub async fn handle_confirmation( &self, request_id: String, confirmation: PermissionConfirmation, ) { if let Err(e) = self.confirmation_tx.send((request_id, confirmation)).await { error!("Failed to send confirmation: {}", e); } } #[instrument(skip(self, user_message, session_config), fields(user_message))] pub async fn reply( &self, user_message: Message, session_config: SessionConfig, cancel_token: Option, ) -> Result>> { let session_manager = self.config.session_manager.clone(); for content in &user_message.content { if let MessageContent::ActionRequired(action_required) = content { if let ActionRequiredData::ElicitationResponse { id, user_data } = &action_required.data { if let Err(e) = ActionRequiredManager::global() .submit_response(id.clone(), user_data.clone()) .await { let error_text = format!("Failed to submit elicitation response: {}", e); error!(error_text); return Ok(Box::pin(stream::once(async { Ok(AgentEvent::Message( Message::assistant().with_text(error_text), )) }))); } session_manager .add_message(&session_config.id, &user_message) .await?; return Ok(Box::pin(futures::stream::empty())); } } } let message_text = user_message.as_concat_text(); // Track custom slash command usage (don't track command name for privacy) if message_text.trim().starts_with('/') { let command = message_text.split_whitespace().next(); if let Some(cmd) = command { if crate::slash_commands::get_recipe_for_command(cmd).is_some() { crate::posthog::emit_custom_slash_command_used(); } } } let command_result = self .execute_command(&message_text, &session_config.id) .await; match command_result { Err(e) => { let error_message = Message::assistant() .with_text(e.to_string()) .with_visibility(true, false); return Ok(Box::pin(stream::once(async move { Ok(AgentEvent::Message(error_message)) }))); } Ok(Some(response)) if response.role == rmcp::model::Role::Assistant => { session_manager .add_message( &session_config.id, &user_message.clone().with_visibility(true, false), ) .await?; session_manager .add_message( &session_config.id, &response.clone().with_visibility(true, false), ) .await?; // Check if this was a command that modifies conversation history let modifies_history = crate::agents::execute_commands::COMPACT_TRIGGERS .contains(&message_text.trim()) || message_text.trim() == "/clear"; return Ok(Box::pin(async_stream::try_stream! { yield AgentEvent::Message(user_message); yield AgentEvent::Message(response); // After commands that modify history, notify UI that history was replaced if modifies_history { let updated_session = session_manager.get_session(&session_config.id, true) .await .map_err(|e| anyhow!("Failed to fetch updated session: {}", e))?; let updated_conversation = updated_session .conversation .ok_or_else(|| anyhow!("Session has no conversation after history modification"))?; yield AgentEvent::HistoryReplaced(updated_conversation); } })); } Ok(Some(resolved_message)) => { session_manager .add_message( &session_config.id, &user_message.clone().with_visibility(true, false), ) .await?; session_manager .add_message( &session_config.id, &resolved_message.clone().with_visibility(false, true), ) .await?; } Ok(None) => { session_manager .add_message(&session_config.id, &user_message) .await?; } } let session = session_manager .get_session(&session_config.id, true) .await?; let conversation = session .conversation .clone() .ok_or_else(|| anyhow::anyhow!("Session {} has no conversation", session_config.id))?; let needs_auto_compact = check_if_compaction_needed( self.provider().await?.as_ref(), &conversation, None, &session, ) .await?; let conversation_to_compact = conversation.clone(); Ok(Box::pin(async_stream::try_stream! { let final_conversation = if !needs_auto_compact { conversation } else { let config = Config::global(); let threshold = config .get_param::("GOOSE_AUTO_COMPACT_THRESHOLD") .unwrap_or(DEFAULT_COMPACTION_THRESHOLD); let threshold_percentage = (threshold * 100.0) as u32; let inline_msg = format!( "Exceeded auto-compact threshold of {}%. Performing auto-compaction...", threshold_percentage ); yield AgentEvent::Message( Message::assistant().with_system_notification( SystemNotificationType::InlineMessage, inline_msg, ) ); yield AgentEvent::Message( Message::assistant().with_system_notification( SystemNotificationType::ThinkingMessage, COMPACTION_THINKING_TEXT, ) ); match compact_messages( self.provider().await?.as_ref(), &session_config.id, &conversation_to_compact, false, ) .await { Ok((compacted_conversation, summarization_usage)) => { session_manager.replace_conversation(&session_config.id, &compacted_conversation).await?; self.update_session_metrics(&session_config.id, session_config.schedule_id.clone(), &summarization_usage, true).await?; yield AgentEvent::HistoryReplaced(compacted_conversation.clone()); yield AgentEvent::Message( Message::assistant().with_system_notification( SystemNotificationType::InlineMessage, "Compaction complete", ) ); compacted_conversation } Err(e) => { yield AgentEvent::Message( Message::assistant().with_text( format!("Ran into this error trying to compact: {e}.\n\nPlease try again or create a new session") ) ); return; } } }; let mut reply_stream = self.reply_internal(final_conversation, session_config, session, cancel_token).await?; while let Some(event) = reply_stream.next().await { yield event?; } })) } async fn reply_internal( &self, conversation: Conversation, session_config: SessionConfig, session: Session, cancel_token: Option, ) -> Result>> { let context = self .prepare_reply_context(&session.id, conversation, session.working_dir.as_path()) .await?; let ReplyContext { mut conversation, mut tools, mut toolshim_tools, mut system_prompt, tool_call_cut_off, goose_mode, initial_messages, } = context; let reply_span = tracing::Span::current(); self.reset_retry_attempts().await; let provider = self.provider().await?; let session_manager = self.config.session_manager.clone(); let session_id = session_config.id.clone(); if !self.config.disable_session_naming { let manager_for_spawn = session_manager.clone(); tokio::spawn(async move { if let Err(e) = manager_for_spawn .maybe_update_name(&session_id, provider) .await { warn!("Failed to generate session description: {}", e); } }); } let working_dir = session.working_dir.clone(); Ok(Box::pin(async_stream::try_stream! { let _ = reply_span.enter(); let mut turns_taken = 0u32; let max_turns = session_config.max_turns.unwrap_or(DEFAULT_MAX_TURNS); let mut compaction_attempts = 0; loop { if is_token_cancelled(&cancel_token) { break; } if let Some(final_output_tool) = self.final_output_tool.lock().await.as_ref() { if final_output_tool.final_output.is_some() { let final_event = AgentEvent::Message( Message::assistant().with_text(final_output_tool.final_output.clone().unwrap()) ); yield final_event; break; } } turns_taken += 1; if turns_taken > max_turns { yield AgentEvent::Message( Message::assistant().with_text( "I've reached the maximum number of actions I can do without user input. Would you like me to continue?" ) ); break; } let tool_pair_summarization_task = crate::context_mgmt::maybe_summarize_tool_pair( self.provider().await?, session_config.id.clone(), conversation.clone(), tool_call_cut_off, ); let conversation_with_moim = super::moim::inject_moim( &session_config.id, conversation.clone(), &self.extension_manager, &working_dir, ).await; let mut stream = Self::stream_response_from_provider( self.provider().await?, &session_config.id, &system_prompt, conversation_with_moim.messages(), &tools, &toolshim_tools, ).await?; let mut no_tools_called = true; let mut messages_to_add = Conversation::default(); let mut tools_updated = false; let mut did_recovery_compact_this_iteration = false; while let Some(next) = stream.next().await { if is_token_cancelled(&cancel_token) { break; } match next { Ok((response, usage)) => { compaction_attempts = 0; // Emit model change event if provider is lead-worker let provider = self.provider().await?; if let Some(lead_worker) = provider.as_lead_worker() { if let Some(ref usage) = usage { let active_model = usage.model.clone(); let (lead_model, worker_model) = lead_worker.get_model_info(); let mode = if active_model == lead_model { "lead" } else if active_model == worker_model { "worker" } else { "unknown" }; yield AgentEvent::ModelChange { model: active_model, mode: mode.to_string(), }; } } if let Some(ref usage) = usage { self.update_session_metrics(&session_config.id, session_config.schedule_id.clone(), usage, false).await?; } if let Some(response) = response { let ToolCategorizeResult { frontend_requests, remaining_requests, filtered_response, } = self.categorize_tools(&response, &tools).await; yield AgentEvent::Message(filtered_response.clone()); tokio::task::yield_now().await; let num_tool_requests = frontend_requests.len() + remaining_requests.len(); if num_tool_requests == 0 { messages_to_add.push(response.clone()); continue; } let tool_response_messages: Vec>> = (0..num_tool_requests) .map(|_| Arc::new(Mutex::new(Message::user().with_generated_id()))) .collect(); let mut request_to_response_map = HashMap::new(); let mut request_metadata: HashMap> = HashMap::new(); for (idx, request) in frontend_requests.iter().chain(remaining_requests.iter()).enumerate() { request_to_response_map.insert(request.id.clone(), tool_response_messages[idx].clone()); request_metadata.insert(request.id.clone(), request.metadata.clone()); } for (idx, request) in frontend_requests.iter().enumerate() { let mut frontend_tool_stream = self.handle_frontend_tool_request( request, tool_response_messages[idx].clone(), ); while let Some(msg) = frontend_tool_stream.try_next().await? { yield AgentEvent::Message(msg); } } if goose_mode == GooseMode::Chat { // Skip all remaining tool calls in chat mode for request in remaining_requests.iter() { 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(CallToolResult { content: vec![Content::text(CHAT_MODE_TOOL_SKIPPED_RESPONSE)], structured_content: None, is_error: Some(false), meta: None, }), request.metadata.as_ref(), ); } } } else { // Run all tool inspectors let inspection_results = self.tool_inspection_manager .inspect_tools( &remaining_requests, conversation.messages(), goose_mode, ) .await?; let permission_check_result = self.tool_inspection_manager .process_inspection_results_with_permission_inspector( &remaining_requests, &inspection_results, ) .unwrap_or_else(|| { let mut result = PermissionCheckResult { approved: vec![], needs_approval: vec![], denied: vec![], }; result.needs_approval.extend(remaining_requests.iter().cloned()); result }); // Track extension requests let mut enable_extension_request_ids = vec![]; for request in &remaining_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 tool_futures = self.handle_approved_and_denied_tools( &permission_check_result, &request_to_response_map, cancel_token.clone(), &session, ).await?; let tool_futures_arc = Arc::new(Mutex::new(tool_futures)); let mut tool_approval_stream = self.handle_approval_tool_requests( &permission_check_result.needs_approval, tool_futures_arc.clone(), &request_to_response_map, cancel_token.clone(), &session, &inspection_results, ); while let Some(msg) = tool_approval_stream.try_next().await? { yield AgentEvent::Message(msg); } tool_futures = { let mut futures_lock = tool_futures_arc.lock().await; futures_lock.drain(..).collect::>() }; let with_id = tool_futures .into_iter() .map(|(request_id, stream)| { stream.map(move |item| (request_id.clone(), item)) }) .collect::>(); let mut combined = stream::select_all(with_id); let mut all_install_successful = true; loop { if is_token_cancelled(&cancel_token) { break; } for msg in self.drain_elicitation_messages(&session_config.id).await { yield AgentEvent::Message(msg); } tokio::select! { biased; tool_item = combined.next() => { match tool_item { Some((request_id, item)) => { match item { ToolStreamItem::Result(output) => { let output = call_tool_result::validate(output); if let Ok(ref call_result) = output { if let Some(ref meta) = call_result.meta { if let Some(notification_data) = meta.0.get("platform_notification") { if let Some(method) = notification_data.get("method").and_then(|v| v.as_str()) { let params = notification_data.get("params").cloned(); let custom_notification = rmcp::model::CustomNotification::new( method.to_string(), params, ); let server_notification = rmcp::model::ServerNotification::CustomNotification(custom_notification); yield AgentEvent::McpNotification((request_id.clone(), server_notification)); } } } } if enable_extension_request_ids.contains(&request_id) && output.is_err() { all_install_successful = false; } if let Some(response_msg) = request_to_response_map.get(&request_id) { let metadata = request_metadata.get(&request_id).and_then(|m| m.as_ref()); let mut response = response_msg.lock().await; *response = response.clone().with_tool_response_with_metadata(request_id, output, metadata); } } ToolStreamItem::Message(msg) => { yield AgentEvent::McpNotification((request_id, msg)); } } } None => break, } } _ = tokio::time::sleep(std::time::Duration::from_millis(100)) => { // Continue loop to drain elicitation messages } } } // check for remaining elicitation messages after all tools complete for msg in self.drain_elicitation_messages(&session_config.id).await { yield AgentEvent::Message(msg); } if all_install_successful && !enable_extension_request_ids.is_empty() { if let Err(e) = self.save_extension_state(&session_config).await { warn!("Failed to save extension state after runtime changes: {}", e); } tools_updated = true; } } // Preserve thinking content from the original response // Gemini (and other thinking models) require thinking to be echoed back let thinking_content: Vec = response.content.iter() .filter(|c| matches!(c, MessageContent::Thinking(_))) .cloned() .collect(); if !thinking_content.is_empty() { let thinking_msg = Message::new( response.role.clone(), response.created, thinking_content, ).with_id(format!("msg_{}", Uuid::new_v4())); messages_to_add.push(thinking_msg); } for (idx, request) in frontend_requests.iter().chain(remaining_requests.iter()).enumerate() { if request.tool_call.is_ok() { let request_msg = Message::assistant() .with_id(format!("msg_{}", Uuid::new_v4())) .with_tool_request_with_metadata( request.id.clone(), request.tool_call.clone(), request.metadata.as_ref(), request.tool_meta.clone(), ); messages_to_add.push(request_msg); let final_response = tool_response_messages[idx] .lock().await.clone(); yield AgentEvent::Message(final_response.clone()); messages_to_add.push(final_response); } } no_tools_called = false; } } Err(ref provider_err @ ProviderError::ContextLengthExceeded(_)) => { crate::posthog::emit_error(provider_err.telemetry_type(), &provider_err.to_string()); compaction_attempts += 1; if compaction_attempts >= 2 { error!("Context limit exceeded after compaction - prompt too large"); yield AgentEvent::Message( Message::assistant().with_system_notification( SystemNotificationType::InlineMessage, "Unable to continue: Context limit still exceeded after compaction. Try using a shorter message, a model with a larger context window, or start a new session." ) ); break; } yield AgentEvent::Message( Message::assistant().with_system_notification( SystemNotificationType::InlineMessage, "Context limit reached. Compacting to continue conversation...", ) ); yield AgentEvent::Message( Message::assistant().with_system_notification( SystemNotificationType::ThinkingMessage, COMPACTION_THINKING_TEXT, ) ); match compact_messages( self.provider().await?.as_ref(), &session_config.id, &conversation, false, ) .await { Ok((compacted_conversation, usage)) => { session_manager.replace_conversation(&session_config.id, &compacted_conversation).await?; self.update_session_metrics(&session_config.id, session_config.schedule_id.clone(), &usage, true).await?; conversation = compacted_conversation; did_recovery_compact_this_iteration = true; yield AgentEvent::HistoryReplaced(conversation.clone()); break; } Err(e) => { crate::posthog::emit_error("compaction_failed", &e.to_string()); error!("Compaction failed: {}", e); break; } } } Err(ref provider_err) => { crate::posthog::emit_error(provider_err.telemetry_type(), &provider_err.to_string()); error!("Error: {}", provider_err); yield AgentEvent::Message( Message::assistant().with_text( format!("Ran into this error: {provider_err}.\n\nPlease retry if you think this is a transient or recoverable error.") ) ); break; } } } if tools_updated { (tools, toolshim_tools, system_prompt) = self.prepare_tools_and_prompt(&session_config.id, &session.working_dir).await?; } let mut exit_chat = false; if no_tools_called { if let Some(final_output_tool) = self.final_output_tool.lock().await.as_ref() { if final_output_tool.final_output.is_none() { warn!("Final output tool has not been called yet. Continuing agent loop."); let message = Message::user().with_text(FINAL_OUTPUT_CONTINUATION_MESSAGE); messages_to_add.push(message.clone()); yield AgentEvent::Message(message); } else { let message = Message::assistant().with_text(final_output_tool.final_output.clone().unwrap()); messages_to_add.push(message.clone()); yield AgentEvent::Message(message); exit_chat = true; } } else if did_recovery_compact_this_iteration { // Avoid setting exit_chat; continue from last user message in the conversation } else { match self.handle_retry_logic(&mut conversation, &session_config, &initial_messages).await { Ok(should_retry) => { if should_retry { info!("Retry logic triggered, restarting agent loop"); } else { exit_chat = true; } } Err(e) => { error!("Retry logic failed: {}", e); yield AgentEvent::Message( Message::assistant().with_text( format!("Retry logic encountered an error: {}", e) ) ); exit_chat = true; } } } } if let Ok(Some((summary_msg, tool_id))) = tool_pair_summarization_task.await { let mut updated_messages = conversation.messages().clone(); let matching: Vec<&mut Message> = updated_messages .iter_mut() .filter(|msg| { msg.id.is_some() && msg.content.iter().any(|c| match c { MessageContent::ToolRequest(req) => req.id == tool_id, MessageContent::ToolResponse(resp) => resp.id == tool_id, _ => false, }) }) .collect(); if matching.len() == 2 { for msg in matching { let id = msg.id.as_ref().unwrap(); msg.metadata = msg.metadata.with_agent_invisible(); SessionManager::update_message_metadata(&session_config.id, id, |metadata| { metadata.with_agent_invisible() }).await?; } conversation = Conversation::new_unvalidated(updated_messages); messages_to_add.push(summary_msg); } else { warn!("Expected a tool request/reply pair, but found {} matching messages", matching.len()); } } for msg in &messages_to_add { session_manager.add_message(&session_config.id, msg).await?; } conversation.extend(messages_to_add); if exit_chat { break; } tokio::task::yield_now().await; } })) } pub async fn extend_system_prompt(&self, instruction: String) { let mut prompt_manager = self.prompt_manager.lock().await; prompt_manager.add_system_prompt_extra(instruction); } pub async fn update_provider( &self, provider: Arc, session_id: &str, ) -> Result<()> { let provider_name = provider.get_name().to_string(); let model_config = provider.get_model_config(); let mut current_provider = self.provider.lock().await; *current_provider = Some(provider); self.config .session_manager .clone() .update(session_id) .provider_name(&provider_name) .model_config(model_config) .apply() .await .context("Failed to persist provider config to session") } /// Restore the provider from session data or fall back to global config /// This is used when resuming a session to restore the provider state pub async fn restore_provider_from_session(&self, session: &Session) -> Result<()> { let config = Config::global(); let provider_name = session .provider_name .clone() .or_else(|| config.get_goose_provider().ok()) .ok_or_else(|| anyhow!("Could not configure agent: missing provider"))?; let model_config = match session.model_config.clone() { Some(saved_config) => saved_config, None => { let model_name = config .get_goose_model() .map_err(|_| anyhow!("Could not configure agent: missing model"))?; crate::model::ModelConfig::new(&model_name) .map_err(|e| anyhow!("Could not configure agent: invalid model {}", e))? } }; let provider = crate::providers::create(&provider_name, model_config) .await .map_err(|e| anyhow!("Could not create provider: {}", e))?; self.update_provider(provider, &session.id).await } /// Override the system prompt with a custom template pub async fn override_system_prompt(&self, template: String) { let mut prompt_manager = self.prompt_manager.lock().await; prompt_manager.set_system_prompt_override(template); } pub async fn list_extension_prompts(&self, session_id: &str) -> HashMap> { self.extension_manager .list_prompts(session_id, CancellationToken::default()) .await .expect("Failed to list prompts") } pub async fn get_prompt( &self, session_id: &str, name: &str, arguments: Value, ) -> Result { // First find which extension has this prompt let prompts = self .extension_manager .list_prompts(session_id, CancellationToken::default()) .await .map_err(|e| anyhow!("Failed to list prompts: {}", e))?; if let Some(extension) = prompts .iter() .find(|(_, prompt_list)| prompt_list.iter().any(|p| p.name == name)) .map(|(extension, _)| extension) { return self .extension_manager .get_prompt( session_id, extension, name, arguments, CancellationToken::default(), ) .await .map_err(|e| anyhow!("Failed to get prompt: {}", e)); } Err(anyhow!("Prompt '{}' not found", name)) } pub async fn get_plan_prompt(&self, session_id: &str) -> Result { let tools = self .extension_manager .get_prefixed_tools(session_id, None) .await?; let tools_info = tools .into_iter() .map(|tool| { ToolInfo::new( &tool.name, tool.description .as_ref() .map(|d| d.as_ref()) .unwrap_or_default(), get_parameter_names(&tool), None, ) }) .collect(); let plan_prompt = self.extension_manager.get_planning_prompt(tools_info).await; Ok(plan_prompt) } pub async fn handle_tool_result(&self, id: String, result: ToolResult) { if let Err(e) = self.tool_result_tx.send((id, result)).await { error!("Failed to send tool result: {}", e); } } pub async fn create_recipe( &self, session_id: &str, mut messages: Conversation, ) -> Result { tracing::info!("Starting recipe creation with {} messages", messages.len()); let extensions_info = self.extension_manager.get_extensions_info().await; tracing::debug!("Retrieved {} extensions info", extensions_info.len()); let (extension_count, tool_count) = self .extension_manager .get_extension_and_tool_counts(session_id) .await; // Get model name from provider let provider = self.provider().await.map_err(|e| { tracing::error!("Failed to get provider for recipe creation: {}", e); e })?; let model_config = provider.get_model_config(); let model_name = &model_config.model_name; tracing::debug!("Using model: {}", model_name); let prompt_manager = self.prompt_manager.lock().await; let system_prompt = prompt_manager .builder() .with_extensions(extensions_info.into_iter()) .with_frontend_instructions(self.frontend_instructions.lock().await.clone()) .with_extension_and_tool_counts(extension_count, tool_count) .build(); let recipe_prompt = prompt_manager.get_recipe_prompt().await; let tools = self .extension_manager .get_prefixed_tools(session_id, None) .await .map_err(|e| { tracing::error!("Failed to get tools for recipe creation: {}", e); e })?; messages.push(Message::user().with_text(recipe_prompt)); let (messages, issues) = fix_conversation(messages); if !issues.is_empty() { issues .iter() .for_each(|issue| tracing::warn!(recipe.conversation.issue = issue)); } tracing::debug!( "Added recipe prompt to messages, total messages: {}", messages.len() ); tracing::info!("Calling provider to generate recipe content"); let (result, _usage) = self .provider .lock() .await .as_ref() .ok_or_else(|| { let error = anyhow!("Provider not available during recipe creation"); tracing::error!("{}", error); error })? .complete(session_id, &system_prompt, messages.messages(), &tools) .await .map_err(|e| { tracing::error!("Provider completion failed during recipe creation: {}", e); e })?; let content = result.as_concat_text(); tracing::debug!( "Provider returned content with {} characters", content.len() ); // the response may be contained in ```json ```, strip that before parsing json let re = Regex::new(r"(?s)```[^\n]*\n(.*?)\n```").unwrap(); let clean_content = re .captures(&content) .and_then(|caps| caps.get(1).map(|m| m.as_str())) .unwrap_or(&content) .trim() .to_string(); let (instructions, activities) = if let Ok(json_content) = serde_json::from_str::(&clean_content) { let instructions = json_content .get("instructions") .ok_or_else(|| anyhow!("Missing 'instructions' in json response"))? .as_str() .ok_or_else(|| anyhow!("instructions' is not a string"))? .to_string(); let activities = json_content .get("activities") .ok_or_else(|| anyhow!("Missing 'activities' in json response"))? .as_array() .ok_or_else(|| anyhow!("'activities' is not an array'"))? .iter() .map(|act| { act.as_str() .map(|s| s.to_string()) .ok_or(anyhow!("'activities' array element is not a string")) }) .collect::>()?; (instructions, activities) } else { tracing::warn!("Failed to parse JSON, falling back to string parsing"); // If we can't get valid JSON, try string parsing // Use split_once to get the content after "Instructions:". let after_instructions = content .split_once("instructions:") .map(|(_, rest)| rest) .unwrap_or(&content); // Split once more to separate instructions from activities. let (instructions_part, activities_text) = after_instructions .split_once("activities:") .unwrap_or((after_instructions, "")); let instructions = instructions_part .trim_end_matches(|c: char| c.is_whitespace() || c == '#') .trim() .to_string(); let activities_text = activities_text.trim(); // Regex to remove bullet markers or numbers with an optional dot. let bullet_re = Regex::new(r"^[•\-*\d]+\.?\s*").expect("Invalid regex"); // Process each line in the activities section. let activities: Vec = activities_text .lines() .map(|line| bullet_re.replace(line, "").to_string()) .map(|s| s.trim().to_string()) .filter(|line| !line.is_empty()) .collect(); (instructions, activities) }; let extension_configs = get_enabled_extensions(); let author = Author { contact: std::env::var("USER") .or_else(|_| std::env::var("USERNAME")) .ok(), metadata: None, }; // Ideally we'd get the name of the provider we are using from the provider itself, // but it doesn't know and the plumbing looks complicated. let config = Config::global(); let provider_name: String = config .get_goose_provider() .expect("No provider configured. Run 'goose configure' first"); let settings = Settings { goose_provider: Some(provider_name.clone()), goose_model: Some(model_name.clone()), temperature: Some(model_config.temperature.unwrap_or(0.0)), max_turns: None, }; tracing::debug!( "Building recipe with {} activities and {} extensions", activities.len(), extension_configs.len() ); let (title, description) = if let Ok(json_content) = serde_json::from_str::(&clean_content) { let title = json_content .get("title") .and_then(|t| t.as_str()) .unwrap_or("Custom recipe from chat") .to_string(); let description = json_content .get("description") .and_then(|d| d.as_str()) .unwrap_or("a custom recipe instance from this chat session") .to_string(); (title, description) } else { ( "Custom recipe from chat".to_string(), "a custom recipe instance from this chat session".to_string(), ) }; let recipe = Recipe::builder() .title(title) .description(description) .instructions(instructions) .activities(activities) .extensions(extension_configs) .settings(settings) .author(author) .build() .map_err(|e| { tracing::error!("Failed to build recipe: {}", e); anyhow!("Recipe build failed: {}", e) })?; tracing::info!("Recipe creation completed successfully"); Ok(recipe) } } #[cfg(test)] mod tests { use super::*; use crate::recipe::Response; #[tokio::test] async fn test_add_final_output_tool() -> Result<()> { let agent = Agent::new(); let response = Response { json_schema: Some(serde_json::json!({ "type": "object", "properties": { "result": {"type": "string"} } })), }; agent.add_final_output_tool(response).await; let tools = agent.list_tools("test-session-id", None).await; let final_output_tool = tools .iter() .find(|tool| tool.name == FINAL_OUTPUT_TOOL_NAME); assert!( final_output_tool.is_some(), "Final output tool should be present after adding" ); let prompt_manager = agent.prompt_manager.lock().await; let system_prompt = prompt_manager.builder().build(); let final_output_tool_ref = agent.final_output_tool.lock().await; let final_output_tool_system_prompt = final_output_tool_ref.as_ref().unwrap().system_prompt(); assert!(system_prompt.contains(&final_output_tool_system_prompt)); Ok(()) } #[tokio::test] async fn test_tool_inspection_manager_has_all_inspectors() -> Result<()> { let agent = Agent::new(); // Verify that the tool inspection manager has all expected inspectors let inspector_names = agent.tool_inspection_manager.inspector_names(); assert!( inspector_names.contains(&"repetition"), "Tool inspection manager should contain repetition inspector" ); assert!( inspector_names.contains(&"permission"), "Tool inspection manager should contain permission inspector" ); assert!( inspector_names.contains(&"security"), "Tool inspection manager should contain security inspector" ); Ok(()) } }