subagents (#2797)
Co-authored-by: Lily Delalande <ldelalande@squareup.com>
This commit is contained in:
@@ -22,7 +22,7 @@ use crate::scheduler_trait::SchedulerTrait;
|
||||
use crate::tool_monitor::{ToolCall, ToolMonitor};
|
||||
use regex::Regex;
|
||||
use serde_json::Value;
|
||||
use tokio::sync::{mpsc, Mutex};
|
||||
use tokio::sync::{mpsc, Mutex, RwLock};
|
||||
use tracing::{debug, error, instrument};
|
||||
|
||||
use crate::agents::extension::{ExtensionConfig, ExtensionError, ExtensionResult, ToolInfo};
|
||||
@@ -45,14 +45,18 @@ use mcp_core::{
|
||||
prompt::Prompt, protocol::GetPromptResult, tool::Tool, Content, ToolError, ToolResult,
|
||||
};
|
||||
|
||||
use crate::agents::subagent_tools::SUBAGENT_RUN_TASK_TOOL_NAME;
|
||||
|
||||
use super::platform_tools;
|
||||
use super::router_tools;
|
||||
use super::subagent_manager::SubAgentManager;
|
||||
use super::subagent_tools;
|
||||
use super::tool_execution::{ToolCallResult, CHAT_MODE_TOOL_SKIPPED_RESPONSE, DECLINED_RESPONSE};
|
||||
|
||||
/// The main goose Agent
|
||||
pub struct Agent {
|
||||
pub(super) provider: Mutex<Option<Arc<dyn Provider>>>,
|
||||
pub(super) extension_manager: Mutex<ExtensionManager>,
|
||||
pub(super) extension_manager: RwLock<ExtensionManager>,
|
||||
pub(super) sub_recipe_manager: Mutex<SubRecipeManager>,
|
||||
pub(super) frontend_tools: Mutex<HashMap<String, FrontendTool>>,
|
||||
pub(super) frontend_instructions: Mutex<Option<String>>,
|
||||
@@ -64,6 +68,8 @@ pub struct Agent {
|
||||
pub(super) tool_monitor: Mutex<Option<ToolMonitor>>,
|
||||
pub(super) router_tool_selector: Mutex<Option<Arc<Box<dyn RouterToolSelector>>>>,
|
||||
pub(super) scheduler_service: Mutex<Option<Arc<dyn SchedulerTrait>>>,
|
||||
pub(super) subagent_manager: Mutex<Option<SubAgentManager>>,
|
||||
pub(super) mcp_notification_rx: Arc<Mutex<mpsc::Receiver<JsonRpcMessage>>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
@@ -73,52 +79,6 @@ pub enum AgentEvent {
|
||||
ModelChange { model: String, mode: String },
|
||||
}
|
||||
|
||||
impl Agent {
|
||||
pub fn new() -> 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);
|
||||
|
||||
Self {
|
||||
provider: Mutex::new(None),
|
||||
extension_manager: Mutex::new(ExtensionManager::new()),
|
||||
sub_recipe_manager: Mutex::new(SubRecipeManager::new()),
|
||||
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)),
|
||||
tool_monitor: Mutex::new(None),
|
||||
router_tool_selector: Mutex::new(None),
|
||||
scheduler_service: Mutex::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn configure_tool_monitor(&self, max_repetitions: Option<u32>) {
|
||||
let mut tool_monitor = self.tool_monitor.lock().await;
|
||||
*tool_monitor = Some(ToolMonitor::new(max_repetitions));
|
||||
}
|
||||
|
||||
pub async fn get_tool_stats(&self) -> Option<HashMap<String, u32>> {
|
||||
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();
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the scheduler service for this agent
|
||||
pub async fn set_scheduler(&self, scheduler: Arc<dyn SchedulerTrait>) {
|
||||
let mut scheduler_service = self.scheduler_service.lock().await;
|
||||
*scheduler_service = Some(scheduler);
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for Agent {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
@@ -160,6 +120,55 @@ where
|
||||
}
|
||||
|
||||
impl Agent {
|
||||
pub fn new() -> 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);
|
||||
// Add MCP notification channel
|
||||
let (mcp_tx, mcp_rx) = mpsc::channel(100);
|
||||
|
||||
Self {
|
||||
provider: Mutex::new(None),
|
||||
extension_manager: RwLock::new(ExtensionManager::new()),
|
||||
sub_recipe_manager: Mutex::new(SubRecipeManager::new()),
|
||||
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)),
|
||||
tool_monitor: Mutex::new(None),
|
||||
router_tool_selector: Mutex::new(None),
|
||||
scheduler_service: Mutex::new(None),
|
||||
// Initialize with MCP notification support
|
||||
subagent_manager: Mutex::new(Some(SubAgentManager::new(mcp_tx))),
|
||||
mcp_notification_rx: Arc::new(Mutex::new(mcp_rx)),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn configure_tool_monitor(&self, max_repetitions: Option<u32>) {
|
||||
let mut tool_monitor = self.tool_monitor.lock().await;
|
||||
*tool_monitor = Some(ToolMonitor::new(max_repetitions));
|
||||
}
|
||||
|
||||
pub async fn get_tool_stats(&self) -> Option<HashMap<String, u32>> {
|
||||
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();
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the scheduler service for this agent
|
||||
pub async fn set_scheduler(&self, scheduler: Arc<dyn SchedulerTrait>) {
|
||||
let mut scheduler_service = self.scheduler_service.lock().await;
|
||||
*scheduler_service = Some(scheduler);
|
||||
}
|
||||
|
||||
/// Get a reference count clone to the provider
|
||||
pub async fn provider(&self) -> Result<Arc<dyn Provider>, anyhow::Error> {
|
||||
match &*self.provider.lock().await {
|
||||
@@ -182,7 +191,7 @@ impl Agent {
|
||||
pub async fn get_prefixed_tools(&self) -> ExtensionResult<Vec<Tool>> {
|
||||
let mut tools = self
|
||||
.extension_manager
|
||||
.lock()
|
||||
.read()
|
||||
.await
|
||||
.get_prefixed_tools(None)
|
||||
.await?;
|
||||
@@ -249,7 +258,7 @@ impl Agent {
|
||||
return (request_id, Ok(ToolCallResult::from(result)));
|
||||
}
|
||||
|
||||
let extension_manager = self.extension_manager.lock().await;
|
||||
let extension_manager = self.extension_manager.read().await;
|
||||
let sub_recipe_manager = self.sub_recipe_manager.lock().await;
|
||||
|
||||
let result: ToolCallResult = if sub_recipe_manager.is_sub_recipe_tool(&tool_call.name) {
|
||||
@@ -271,6 +280,11 @@ impl Agent {
|
||||
)
|
||||
} else if tool_call.name == PLATFORM_SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME {
|
||||
ToolCallResult::from(extension_manager.search_available_extensions().await)
|
||||
} else if tool_call.name == SUBAGENT_RUN_TASK_TOOL_NAME {
|
||||
ToolCallResult::from(
|
||||
self.handle_run_subagent_task(tool_call.arguments.clone())
|
||||
.await,
|
||||
)
|
||||
} else if self.is_frontend_tool(&tool_call.name).await {
|
||||
// For frontend tools, return an error indicating we need frontend execution
|
||||
ToolCallResult::from(Err(ToolError::ExecutionError(
|
||||
@@ -333,13 +347,13 @@ impl Agent {
|
||||
extension_name: String,
|
||||
request_id: String,
|
||||
) -> (String, Result<Vec<Content>, ToolError>) {
|
||||
let mut extension_manager = self.extension_manager.lock().await;
|
||||
let mut extension_manager = self.extension_manager.write().await;
|
||||
|
||||
let selector = self.router_tool_selector.lock().await.clone();
|
||||
if ToolRouterIndexManager::is_tool_router_enabled(&selector) {
|
||||
if let Some(selector) = selector {
|
||||
let selector_action = if action == "disable" { "remove" } else { "add" };
|
||||
let extension_manager = self.extension_manager.lock().await;
|
||||
let extension_manager = self.extension_manager.read().await;
|
||||
let selector = Arc::new(selector);
|
||||
if let Err(e) = ToolRouterIndexManager::update_extension_tools(
|
||||
&selector,
|
||||
@@ -407,6 +421,33 @@ impl Agent {
|
||||
})
|
||||
.map_err(|e| ToolError::ExecutionError(e.to_string()));
|
||||
|
||||
// Update vector index if operation was successful and vector routing is enabled
|
||||
if result.is_ok() {
|
||||
let selector = self.router_tool_selector.lock().await.clone();
|
||||
if ToolRouterIndexManager::is_tool_router_enabled(&selector) {
|
||||
if let Some(selector) = selector {
|
||||
let vector_action = if action == "disable" { "remove" } else { "add" };
|
||||
let extension_manager = self.extension_manager.read().await;
|
||||
let selector = Arc::new(selector);
|
||||
if let Err(e) = ToolRouterIndexManager::update_extension_tools(
|
||||
&selector,
|
||||
&extension_manager,
|
||||
&extension_name,
|
||||
vector_action,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return (
|
||||
request_id,
|
||||
Err(ToolError::ExecutionError(format!(
|
||||
"Failed to update vector index: {}",
|
||||
e
|
||||
))),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
(request_id, result)
|
||||
}
|
||||
|
||||
@@ -439,7 +480,7 @@ impl Agent {
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
let mut extension_manager = self.extension_manager.lock().await;
|
||||
let mut extension_manager = self.extension_manager.write().await;
|
||||
extension_manager.add_extension(extension.clone()).await?;
|
||||
}
|
||||
}
|
||||
@@ -448,7 +489,7 @@ impl Agent {
|
||||
let selector = self.router_tool_selector.lock().await.clone();
|
||||
if ToolRouterIndexManager::is_tool_router_enabled(&selector) {
|
||||
if let Some(selector) = selector {
|
||||
let extension_manager = self.extension_manager.lock().await;
|
||||
let extension_manager = self.extension_manager.read().await;
|
||||
let selector = Arc::new(selector);
|
||||
if let Err(e) = ToolRouterIndexManager::update_extension_tools(
|
||||
&selector,
|
||||
@@ -471,7 +512,7 @@ impl Agent {
|
||||
}
|
||||
|
||||
pub async fn list_tools(&self, extension_name: Option<String>) -> Vec<Tool> {
|
||||
let extension_manager = self.extension_manager.lock().await;
|
||||
let extension_manager = self.extension_manager.read().await;
|
||||
let mut prefixed_tools = extension_manager
|
||||
.get_prefixed_tools(extension_name.clone())
|
||||
.await
|
||||
@@ -485,6 +526,12 @@ impl Agent {
|
||||
platform_tools::manage_schedule_tool(),
|
||||
]);
|
||||
|
||||
// Add subagent tool (only if ALPHA_FEATURES is enabled)
|
||||
let config = Config::global();
|
||||
if config.get_param::<bool>("ALPHA_FEATURES").unwrap_or(false) {
|
||||
prefixed_tools.push(subagent_tools::run_task_subagent_tool());
|
||||
}
|
||||
|
||||
// Add resource tools if supported
|
||||
if extension_manager.supports_resources() {
|
||||
prefixed_tools.extend([
|
||||
@@ -521,7 +568,7 @@ impl Agent {
|
||||
let selector = self.router_tool_selector.lock().await.clone();
|
||||
if let Some(selector) = selector {
|
||||
if let Ok(recent_calls) = selector.get_recent_tool_calls(20).await {
|
||||
let extension_manager = self.extension_manager.lock().await;
|
||||
let extension_manager = self.extension_manager.read().await;
|
||||
// Add recent tool calls to the list, avoiding duplicates
|
||||
for tool_name in recent_calls {
|
||||
// Find the tool in the extension manager's tools
|
||||
@@ -541,11 +588,14 @@ impl Agent {
|
||||
}
|
||||
|
||||
pub async fn remove_extension(&self, name: &str) -> Result<()> {
|
||||
let mut extension_manager = self.extension_manager.write().await;
|
||||
extension_manager.remove_extension(name).await?;
|
||||
|
||||
// If vector tool selection is enabled, remove tools from the index
|
||||
let selector = self.router_tool_selector.lock().await.clone();
|
||||
if ToolRouterIndexManager::is_tool_router_enabled(&selector) {
|
||||
if let Some(selector) = selector {
|
||||
let extension_manager = self.extension_manager.lock().await;
|
||||
let extension_manager = self.extension_manager.read().await;
|
||||
ToolRouterIndexManager::update_extension_tools(
|
||||
&selector,
|
||||
&extension_manager,
|
||||
@@ -556,14 +606,11 @@ impl Agent {
|
||||
}
|
||||
}
|
||||
|
||||
let mut extension_manager = self.extension_manager.lock().await;
|
||||
extension_manager.remove_extension(name).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn list_extensions(&self) -> Vec<String> {
|
||||
let extension_manager = self.extension_manager.lock().await;
|
||||
let extension_manager = self.extension_manager.read().await;
|
||||
extension_manager
|
||||
.list_extensions()
|
||||
.await
|
||||
@@ -631,6 +678,28 @@ impl Agent {
|
||||
Ok(Box::pin(async_stream::try_stream! {
|
||||
let _ = reply_span.enter();
|
||||
loop {
|
||||
// Check for MCP notifications from subagents
|
||||
let mcp_notifications = self.get_mcp_notifications().await;
|
||||
for notification in mcp_notifications {
|
||||
// Extract subagent info from the notification data
|
||||
if let JsonRpcMessage::Notification(ref notif) = notification {
|
||||
if let Some(params) = ¬if.params {
|
||||
if let Some(data) = params.get("data") {
|
||||
if let (Some(subagent_id), Some(_message)) = (
|
||||
data.get("subagent_id").and_then(|v| v.as_str()),
|
||||
data.get("message").and_then(|v| v.as_str())
|
||||
) {
|
||||
// Emit as McpNotification event
|
||||
yield AgentEvent::McpNotification((
|
||||
subagent_id.to_string(),
|
||||
notification.clone()
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
match Self::generate_response_from_provider(
|
||||
self.provider().await?,
|
||||
&system_prompt,
|
||||
@@ -832,6 +901,31 @@ impl Agent {
|
||||
|
||||
messages.push(response);
|
||||
messages.push(final_message_tool_resp);
|
||||
|
||||
// Check for MCP notifications from subagents again before next iteration
|
||||
// Note: These are already handled as McpNotification events above,
|
||||
// so we don't need to convert them to assistant messages here.
|
||||
// This was causing duplicate plain-text notifications.
|
||||
// let mcp_notifications = self.get_mcp_notifications().await;
|
||||
// for notification in mcp_notifications {
|
||||
// // Extract subagent info from the notification data for assistant messages
|
||||
// if let JsonRpcMessage::Notification(ref notif) = notification {
|
||||
// if let Some(params) = ¬if.params {
|
||||
// if let Some(data) = params.get("data") {
|
||||
// if let (Some(subagent_id), Some(message)) = (
|
||||
// data.get("subagent_id").and_then(|v| v.as_str()),
|
||||
// data.get("message").and_then(|v| v.as_str())
|
||||
// ) {
|
||||
// yield AgentEvent::Message(
|
||||
// Message::assistant().with_text(
|
||||
// format!("Subagent {}: {}", subagent_id, message)
|
||||
// )
|
||||
// );
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
},
|
||||
Err(ProviderError::ContextLengthExceeded(_)) => {
|
||||
// At this point, the last message should be a user message
|
||||
@@ -862,9 +956,32 @@ impl Agent {
|
||||
prompt_manager.add_system_prompt_extra(instruction);
|
||||
}
|
||||
|
||||
/// Update the provider used by this agent
|
||||
/// Get MCP notifications from subagents
|
||||
pub async fn get_mcp_notifications(&self) -> Vec<JsonRpcMessage> {
|
||||
let mut notifications = Vec::new();
|
||||
let mut rx = self.mcp_notification_rx.lock().await;
|
||||
|
||||
while let Ok(notification) = rx.try_recv() {
|
||||
notifications.push(notification);
|
||||
}
|
||||
|
||||
notifications
|
||||
}
|
||||
|
||||
/// Update the provider
|
||||
pub async fn update_provider(&self, provider: Arc<dyn Provider>) -> Result<()> {
|
||||
*self.provider.lock().await = Some(provider.clone());
|
||||
let mut current_provider = self.provider.lock().await;
|
||||
*current_provider = Some(provider.clone());
|
||||
|
||||
// Initialize subagent manager with MCP notification support
|
||||
// Need to recreate the MCP channel since we're replacing the manager
|
||||
let (mcp_tx, mcp_rx) = mpsc::channel(100);
|
||||
{
|
||||
let mut rx_guard = self.mcp_notification_rx.lock().await;
|
||||
*rx_guard = mcp_rx;
|
||||
}
|
||||
*self.subagent_manager.lock().await = Some(SubAgentManager::new(mcp_tx));
|
||||
|
||||
self.update_router_tool_selector(Some(provider), None)
|
||||
.await?;
|
||||
Ok(())
|
||||
@@ -876,7 +993,7 @@ impl Agent {
|
||||
reindex_all: Option<bool>,
|
||||
) -> Result<()> {
|
||||
let config = Config::global();
|
||||
let extension_manager = self.extension_manager.lock().await;
|
||||
let _extension_manager = self.extension_manager.read().await;
|
||||
let provider = match provider {
|
||||
Some(p) => p,
|
||||
None => self.provider().await?,
|
||||
@@ -910,6 +1027,7 @@ impl Agent {
|
||||
};
|
||||
|
||||
// First index platform tools
|
||||
let extension_manager = self.extension_manager.read().await;
|
||||
ToolRouterIndexManager::index_platform_tools(&selector, &extension_manager).await?;
|
||||
|
||||
if reindex_all.unwrap_or(false) {
|
||||
@@ -934,6 +1052,7 @@ impl Agent {
|
||||
|
||||
// Update the selector
|
||||
*self.router_tool_selector.lock().await = Some(selector.clone());
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -944,7 +1063,7 @@ impl Agent {
|
||||
}
|
||||
|
||||
pub async fn list_extension_prompts(&self) -> HashMap<String, Vec<Prompt>> {
|
||||
let extension_manager = self.extension_manager.lock().await;
|
||||
let extension_manager = self.extension_manager.read().await;
|
||||
extension_manager
|
||||
.list_prompts()
|
||||
.await
|
||||
@@ -952,7 +1071,7 @@ impl Agent {
|
||||
}
|
||||
|
||||
pub async fn get_prompt(&self, name: &str, arguments: Value) -> Result<GetPromptResult> {
|
||||
let extension_manager = self.extension_manager.lock().await;
|
||||
let extension_manager = self.extension_manager.read().await;
|
||||
|
||||
// First find which extension has this prompt
|
||||
let prompts = extension_manager
|
||||
@@ -975,7 +1094,7 @@ impl Agent {
|
||||
}
|
||||
|
||||
pub async fn get_plan_prompt(&self) -> anyhow::Result<String> {
|
||||
let extension_manager = self.extension_manager.lock().await;
|
||||
let extension_manager = self.extension_manager.read().await;
|
||||
let tools = extension_manager.get_prefixed_tools(None).await?;
|
||||
let tools_info = tools
|
||||
.into_iter()
|
||||
@@ -1001,7 +1120,7 @@ impl Agent {
|
||||
}
|
||||
|
||||
pub async fn create_recipe(&self, mut messages: Vec<Message>) -> Result<Recipe> {
|
||||
let extension_manager = self.extension_manager.lock().await;
|
||||
let extension_manager = self.extension_manager.read().await;
|
||||
let extensions_info = extension_manager.get_extensions_info().await;
|
||||
|
||||
// Get model name from provider
|
||||
|
||||
@@ -11,6 +11,11 @@ mod router_tool_selector;
|
||||
mod router_tools;
|
||||
mod schedule_tool;
|
||||
pub mod sub_recipe_manager;
|
||||
pub mod subagent;
|
||||
pub mod subagent_handler;
|
||||
pub mod subagent_manager;
|
||||
pub mod subagent_tools;
|
||||
pub mod subagent_types;
|
||||
mod tool_execution;
|
||||
mod tool_router_index_manager;
|
||||
pub(crate) mod tool_vectordb;
|
||||
@@ -20,4 +25,7 @@ pub use agent::{Agent, AgentEvent};
|
||||
pub use extension::ExtensionConfig;
|
||||
pub use extension_manager::ExtensionManager;
|
||||
pub use prompt_manager::PromptManager;
|
||||
pub use subagent::{SubAgent, SubAgentConfig, SubAgentProgress, SubAgentStatus};
|
||||
pub use subagent_manager::SubAgentManager;
|
||||
pub use subagent_types::SpawnSubAgentArgs;
|
||||
pub use types::{FrontendTool, SessionConfig};
|
||||
|
||||
@@ -52,7 +52,7 @@ impl Agent {
|
||||
}
|
||||
|
||||
// Prepare system prompt
|
||||
let extension_manager = self.extension_manager.lock().await;
|
||||
let extension_manager = self.extension_manager.read().await;
|
||||
let extensions_info = extension_manager.get_extensions_info().await;
|
||||
|
||||
// Get model name from provider
|
||||
|
||||
@@ -0,0 +1,757 @@
|
||||
use crate::{
|
||||
agents::{extension_manager::ExtensionManager, Agent},
|
||||
message::{Message, MessageContent, ToolRequest},
|
||||
prompt_template::render_global_file,
|
||||
providers::base::Provider,
|
||||
providers::errors::ProviderError,
|
||||
recipe::Recipe,
|
||||
};
|
||||
use anyhow::anyhow;
|
||||
use chrono::{DateTime, Utc};
|
||||
use mcp_core::protocol::{JsonRpcMessage, JsonRpcNotification};
|
||||
use mcp_core::{handler::ToolError, role::Role, tool::Tool};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{self, json};
|
||||
use std::{collections::HashMap, sync::Arc};
|
||||
use tokio::sync::{mpsc, Mutex, RwLock};
|
||||
use tracing::{debug, error, instrument};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::agents::platform_tools::{
|
||||
self, PLATFORM_LIST_RESOURCES_TOOL_NAME, PLATFORM_READ_RESOURCE_TOOL_NAME,
|
||||
PLATFORM_SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME,
|
||||
};
|
||||
use crate::agents::subagent_tools::SUBAGENT_RUN_TASK_TOOL_NAME;
|
||||
|
||||
/// Status of a subagent
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub enum SubAgentStatus {
|
||||
Ready, // Ready to process messages
|
||||
Processing, // Currently working on a task
|
||||
Completed(String), // Task completed (with optional message for success/error)
|
||||
Terminated, // Manually terminated
|
||||
}
|
||||
|
||||
/// Configuration for a subagent
|
||||
#[derive(Debug)]
|
||||
pub struct SubAgentConfig {
|
||||
pub id: String,
|
||||
pub recipe: Option<Recipe>,
|
||||
pub instructions: Option<String>,
|
||||
pub max_turns: Option<usize>,
|
||||
pub timeout_seconds: Option<u64>,
|
||||
}
|
||||
|
||||
impl SubAgentConfig {
|
||||
pub fn new_with_recipe(recipe: Recipe) -> Self {
|
||||
Self {
|
||||
id: Uuid::new_v4().to_string(),
|
||||
recipe: Some(recipe),
|
||||
instructions: None,
|
||||
max_turns: None,
|
||||
timeout_seconds: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn new_with_instructions(instructions: String) -> Self {
|
||||
Self {
|
||||
id: Uuid::new_v4().to_string(),
|
||||
recipe: None,
|
||||
instructions: Some(instructions),
|
||||
max_turns: None,
|
||||
timeout_seconds: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_max_turns(mut self, max_turns: usize) -> Self {
|
||||
self.max_turns = Some(max_turns);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_timeout(mut self, timeout_seconds: u64) -> Self {
|
||||
self.timeout_seconds = Some(timeout_seconds);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Progress information for a subagent
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SubAgentProgress {
|
||||
pub subagent_id: String,
|
||||
pub status: SubAgentStatus,
|
||||
pub message: String,
|
||||
pub turn: usize,
|
||||
pub max_turns: Option<usize>,
|
||||
pub timestamp: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// A specialized agent that can handle specific tasks independently
|
||||
pub struct SubAgent {
|
||||
pub id: String,
|
||||
pub conversation: Arc<Mutex<Vec<Message>>>,
|
||||
pub status: Arc<RwLock<SubAgentStatus>>,
|
||||
pub config: SubAgentConfig,
|
||||
pub turn_count: Arc<Mutex<usize>>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub recipe_extensions: Arc<Mutex<Vec<String>>>,
|
||||
pub missing_extensions: Arc<Mutex<Vec<String>>>, // Track extensions that weren't enabled
|
||||
pub mcp_notification_tx: mpsc::Sender<JsonRpcMessage>, // For MCP notifications
|
||||
}
|
||||
|
||||
impl SubAgent {
|
||||
/// Create a new subagent with the given configuration and provider
|
||||
#[instrument(skip(config, _provider, extension_manager, mcp_notification_tx))]
|
||||
pub async fn new(
|
||||
config: SubAgentConfig,
|
||||
_provider: Arc<dyn Provider>,
|
||||
extension_manager: Arc<tokio::sync::RwLockReadGuard<'_, ExtensionManager>>,
|
||||
mcp_notification_tx: mpsc::Sender<JsonRpcMessage>,
|
||||
) -> Result<(Arc<Self>, tokio::task::JoinHandle<()>), anyhow::Error> {
|
||||
debug!("Creating new subagent with id: {}", config.id);
|
||||
|
||||
let mut missing_extensions = Vec::new();
|
||||
let mut recipe_extensions = Vec::new();
|
||||
|
||||
// Check if extensions from recipe exist in the extension manager
|
||||
if let Some(recipe) = &config.recipe {
|
||||
if let Some(extensions) = &recipe.extensions {
|
||||
for extension in extensions {
|
||||
let extension_name = extension.name();
|
||||
let existing_extensions = extension_manager.list_extensions().await?;
|
||||
|
||||
if !existing_extensions.contains(&extension_name) {
|
||||
missing_extensions.push(extension_name);
|
||||
} else {
|
||||
recipe_extensions.push(extension_name);
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// If no recipe, inherit all extensions from the parent agent
|
||||
let existing_extensions = extension_manager.list_extensions().await?;
|
||||
recipe_extensions = existing_extensions;
|
||||
}
|
||||
|
||||
let subagent = Arc::new(SubAgent {
|
||||
id: config.id.clone(),
|
||||
conversation: Arc::new(Mutex::new(Vec::new())),
|
||||
status: Arc::new(RwLock::new(SubAgentStatus::Ready)),
|
||||
config,
|
||||
turn_count: Arc::new(Mutex::new(0)),
|
||||
created_at: Utc::now(),
|
||||
recipe_extensions: Arc::new(Mutex::new(recipe_extensions)),
|
||||
missing_extensions: Arc::new(Mutex::new(missing_extensions)),
|
||||
mcp_notification_tx,
|
||||
});
|
||||
|
||||
// Send initial MCP notification
|
||||
let subagent_clone = Arc::clone(&subagent);
|
||||
subagent_clone
|
||||
.send_mcp_notification("subagent_created", "Subagent created and ready")
|
||||
.await;
|
||||
|
||||
// Create a background task handle (for future use with streaming/monitoring)
|
||||
let subagent_clone = Arc::clone(&subagent);
|
||||
let handle = tokio::spawn(async move {
|
||||
// This could be used for background monitoring, cleanup, etc.
|
||||
debug!("Subagent {} background task started", subagent_clone.id);
|
||||
});
|
||||
|
||||
debug!("Subagent {} created successfully", subagent.id);
|
||||
Ok((subagent, handle))
|
||||
}
|
||||
|
||||
/// 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
|
||||
{
|
||||
let mut current_status = self.status.write().await;
|
||||
*current_status = status.clone();
|
||||
} // Write lock is released here!
|
||||
|
||||
// Send MCP notifications based on status
|
||||
match &status {
|
||||
SubAgentStatus::Processing => {
|
||||
self.send_mcp_notification("status_changed", "Processing request")
|
||||
.await;
|
||||
}
|
||||
SubAgentStatus::Completed(msg) => {
|
||||
self.send_mcp_notification("completed", &format!("Completed: {}", msg))
|
||||
.await;
|
||||
}
|
||||
SubAgentStatus::Terminated => {
|
||||
self.send_mcp_notification("terminated", "Subagent terminated")
|
||||
.await;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
/// Send an MCP notification about the subagent's activity
|
||||
pub async fn send_mcp_notification(&self, notification_type: &str, message: &str) {
|
||||
let notification = JsonRpcMessage::Notification(JsonRpcNotification {
|
||||
jsonrpc: "2.0".to_string(),
|
||||
method: "notifications/message".to_string(),
|
||||
params: Some(json!({
|
||||
"level": "info",
|
||||
"logger": format!("subagent_{}", self.id),
|
||||
"data": {
|
||||
"subagent_id": self.id,
|
||||
"type": notification_type,
|
||||
"message": message,
|
||||
"timestamp": Utc::now().to_rfc3339()
|
||||
}
|
||||
})),
|
||||
});
|
||||
|
||||
if let Err(e) = self.mcp_notification_tx.send(notification).await {
|
||||
error!(
|
||||
"Failed to send MCP notification from subagent {}: {}",
|
||||
self.id, e
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// 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, provider, extension_manager))]
|
||||
pub async fn reply_subagent(
|
||||
&self,
|
||||
message: String,
|
||||
provider: Arc<dyn Provider>,
|
||||
extension_manager: Arc<tokio::sync::RwLockReadGuard<'_, ExtensionManager>>,
|
||||
) -> Result<Message, anyhow::Error> {
|
||||
debug!("Processing message for subagent {}", self.id);
|
||||
self.send_mcp_notification("message_processing", &format!("Processing: {}", message))
|
||||
.await;
|
||||
|
||||
// Check if we've exceeded max turns
|
||||
{
|
||||
let turn_count = *self.turn_count.lock().await;
|
||||
if let Some(max_turns) = self.config.max_turns {
|
||||
if turn_count >= max_turns {
|
||||
self.set_status(SubAgentStatus::Completed(
|
||||
"Maximum turns exceeded".to_string(),
|
||||
))
|
||||
.await;
|
||||
return Err(anyhow!("Maximum turns ({}) exceeded", max_turns));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Set status to processing
|
||||
self.set_status(SubAgentStatus::Processing).await;
|
||||
|
||||
// Add user message to conversation
|
||||
let user_message = Message::user().with_text(message.clone());
|
||||
{
|
||||
let mut conversation = self.conversation.lock().await;
|
||||
conversation.push(user_message.clone());
|
||||
}
|
||||
|
||||
// Increment turn count
|
||||
{
|
||||
let mut turn_count = self.turn_count.lock().await;
|
||||
*turn_count += 1;
|
||||
self.send_mcp_notification(
|
||||
"turn_progress",
|
||||
&format!("Turn {}/{}", turn_count, self.config.max_turns.unwrap_or(0)),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
// Get the current conversation for context
|
||||
let mut messages = self.get_conversation().await;
|
||||
|
||||
// Get tools based on whether we're using a recipe or inheriting from parent
|
||||
let tools: Vec<Tool> = if self.config.recipe.is_some() {
|
||||
// Recipe mode: only get tools from the recipe's extensions
|
||||
let recipe_extensions = self.recipe_extensions.lock().await;
|
||||
let mut recipe_tools = Vec::new();
|
||||
|
||||
debug!(
|
||||
"Subagent {} operating in recipe mode with {} extensions",
|
||||
self.id,
|
||||
recipe_extensions.len()
|
||||
);
|
||||
|
||||
for extension_name in recipe_extensions.iter() {
|
||||
match extension_manager
|
||||
.get_prefixed_tools(Some(extension_name.clone()))
|
||||
.await
|
||||
{
|
||||
Ok(mut ext_tools) => {
|
||||
debug!(
|
||||
"Added {} tools from extension {}",
|
||||
ext_tools.len(),
|
||||
extension_name
|
||||
);
|
||||
recipe_tools.append(&mut ext_tools);
|
||||
}
|
||||
Err(e) => {
|
||||
debug!(
|
||||
"Failed to get tools for extension {}: {}",
|
||||
extension_name, e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
debug!(
|
||||
"Subagent {} has {} total recipe tools before filtering",
|
||||
self.id,
|
||||
recipe_tools.len()
|
||||
);
|
||||
// Filter out subagent tools from recipe tools
|
||||
let mut filtered_tools = Self::filter_subagent_tools(recipe_tools);
|
||||
|
||||
// Add platform tools (except subagent tools)
|
||||
Self::add_platform_tools(&mut filtered_tools, &extension_manager).await;
|
||||
|
||||
debug!(
|
||||
"Subagent {} has {} tools after filtering and adding platform tools",
|
||||
self.id,
|
||||
filtered_tools.len()
|
||||
);
|
||||
filtered_tools
|
||||
} else {
|
||||
// No recipe: inherit all tools from parent (but filter out subagent tools)
|
||||
debug!(
|
||||
"Subagent {} operating in inheritance mode, using all parent tools",
|
||||
self.id
|
||||
);
|
||||
let parent_tools = extension_manager.get_prefixed_tools(None).await?;
|
||||
debug!(
|
||||
"Subagent {} has {} parent tools before filtering",
|
||||
self.id,
|
||||
parent_tools.len()
|
||||
);
|
||||
let mut filtered_tools = Self::filter_subagent_tools(parent_tools);
|
||||
|
||||
// Add platform tools (except subagent tools)
|
||||
Self::add_platform_tools(&mut filtered_tools, &extension_manager).await;
|
||||
|
||||
debug!(
|
||||
"Subagent {} has {} tools after filtering and adding platform tools",
|
||||
self.id,
|
||||
filtered_tools.len()
|
||||
);
|
||||
filtered_tools
|
||||
};
|
||||
|
||||
let toolshim_tools: Vec<Tool> = vec![];
|
||||
|
||||
// Build system prompt using the template
|
||||
let system_prompt = self.build_system_prompt(&tools).await?;
|
||||
|
||||
// Generate response from provider
|
||||
loop {
|
||||
match Agent::generate_response_from_provider(
|
||||
Arc::clone(&provider),
|
||||
&system_prompt,
|
||||
&messages,
|
||||
&tools,
|
||||
&toolshim_tools,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok((response, _usage)) => {
|
||||
// Process any tool calls in the response
|
||||
let tool_requests: Vec<ToolRequest> = response
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|content| {
|
||||
if let MessageContent::ToolRequest(req) = content {
|
||||
Some(req.clone())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
// If there are no tool requests, we're done
|
||||
if tool_requests.is_empty() {
|
||||
self.add_message(response.clone()).await;
|
||||
|
||||
// Send notification about response
|
||||
self.send_mcp_notification(
|
||||
"response_generated",
|
||||
&format!("Responded: {}", response.as_concat_text()),
|
||||
)
|
||||
.await;
|
||||
|
||||
// Add delay before completion to ensure all processing finishes
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
|
||||
|
||||
// Set status back to ready and return the final response
|
||||
self.set_status(SubAgentStatus::Completed("Completed!".to_string()))
|
||||
.await;
|
||||
break Ok(response);
|
||||
}
|
||||
|
||||
// Add the assistant message with tool calls to the conversation
|
||||
messages.push(response.clone());
|
||||
|
||||
// Process each tool request and create user response messages
|
||||
for request in &tool_requests {
|
||||
if let Ok(tool_call) = &request.tool_call {
|
||||
// Send notification about tool usage
|
||||
self.send_mcp_notification(
|
||||
"tool_usage",
|
||||
&format!("Using tool: {}", tool_call.name),
|
||||
)
|
||||
.await;
|
||||
|
||||
// Handle platform tools or dispatch to extension manager
|
||||
let tool_result = if self.is_platform_tool(&tool_call.name) {
|
||||
self.handle_platform_tool_call(
|
||||
tool_call.clone(),
|
||||
&extension_manager,
|
||||
)
|
||||
.await
|
||||
} else {
|
||||
match extension_manager
|
||||
.dispatch_tool_call(tool_call.clone())
|
||||
.await
|
||||
{
|
||||
Ok(result) => result.result.await,
|
||||
Err(e) => Err(ToolError::ExecutionError(e.to_string())),
|
||||
}
|
||||
};
|
||||
|
||||
match tool_result {
|
||||
Ok(result) => {
|
||||
// Create a user message with the tool response
|
||||
let tool_response_message = Message::user()
|
||||
.with_tool_response(request.id.clone(), Ok(result.clone()));
|
||||
messages.push(tool_response_message);
|
||||
|
||||
// Send notification about tool completion
|
||||
self.send_mcp_notification(
|
||||
"tool_completed",
|
||||
&format!("Tool {} completed successfully", tool_call.name),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
Err(e) => {
|
||||
// Create a user message with the tool error
|
||||
let tool_error_message = Message::user().with_tool_response(
|
||||
request.id.clone(),
|
||||
Err(ToolError::ExecutionError(e.to_string())),
|
||||
);
|
||||
messages.push(tool_error_message);
|
||||
|
||||
// Send notification about tool error
|
||||
self.send_mcp_notification(
|
||||
"tool_error",
|
||||
&format!("Tool {} error: {}", tool_call.name, e),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Continue the loop to get the next response from the provider
|
||||
}
|
||||
Err(ProviderError::ContextLengthExceeded(_)) => {
|
||||
self.set_status(SubAgentStatus::Completed(
|
||||
"Context length exceeded".to_string(),
|
||||
))
|
||||
.await;
|
||||
break Ok(Message::assistant().with_context_length_exceeded(
|
||||
"The context length of the model has been exceeded. Please start a new session and try again.",
|
||||
));
|
||||
}
|
||||
Err(ProviderError::RateLimitExceeded(_)) => {
|
||||
self.set_status(SubAgentStatus::Completed("Rate limit exceeded".to_string()))
|
||||
.await;
|
||||
break Ok(Message::assistant()
|
||||
.with_text("Rate limit exceeded. Please try again later."));
|
||||
}
|
||||
Err(e) => {
|
||||
self.set_status(SubAgentStatus::Completed(format!("Error: {}", e)))
|
||||
.await;
|
||||
error!("Error: {}", e);
|
||||
break Ok(Message::assistant().with_text(format!("Ran into this error: {e}.\n\nPlease retry if you think this is a transient or recoverable error.")));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Add a message to the conversation (for tracking agent responses)
|
||||
pub 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<Message> {
|
||||
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(())
|
||||
}
|
||||
|
||||
/// Get formatted conversation for display
|
||||
pub async fn get_formatted_conversation(&self) -> String {
|
||||
let conversation = self.conversation.lock().await;
|
||||
|
||||
let mut formatted = format!("=== Subagent {} Conversation ===\n", self.id);
|
||||
|
||||
if let Some(recipe) = &self.config.recipe {
|
||||
formatted.push_str(&format!("Recipe: {}\n", recipe.title));
|
||||
} else if let Some(instructions) = &self.config.instructions {
|
||||
formatted.push_str(&format!("Instructions: {}\n", instructions));
|
||||
} else {
|
||||
formatted.push_str("Mode: Ad-hoc subagent\n");
|
||||
}
|
||||
|
||||
formatted.push_str(&format!(
|
||||
"Created: {}\n",
|
||||
self.created_at.format("%Y-%m-%d %H:%M:%S UTC")
|
||||
));
|
||||
|
||||
let progress = self.get_progress().await;
|
||||
|
||||
formatted.push_str(&format!("Status: {:?}\n", progress.status));
|
||||
formatted.push_str(&format!("Turn: {}", progress.turn));
|
||||
if let Some(max_turns) = progress.max_turns {
|
||||
formatted.push_str(&format!("/{}", max_turns));
|
||||
}
|
||||
formatted.push_str("\n\n");
|
||||
|
||||
for (i, message) in conversation.iter().enumerate() {
|
||||
formatted.push_str(&format!(
|
||||
"{}. {}: {}\n",
|
||||
i + 1,
|
||||
match message.role {
|
||||
Role::User => "User",
|
||||
Role::Assistant => "Assistant",
|
||||
},
|
||||
message.as_concat_text()
|
||||
));
|
||||
}
|
||||
|
||||
formatted.push_str("=== End Conversation ===\n");
|
||||
|
||||
formatted
|
||||
}
|
||||
|
||||
/// Get the list of extensions that weren't enabled
|
||||
pub async fn get_missing_extensions(&self) -> Vec<String> {
|
||||
self.missing_extensions.lock().await.clone()
|
||||
}
|
||||
|
||||
/// Filter out subagent spawning tools to prevent infinite recursion
|
||||
fn filter_subagent_tools(tools: Vec<Tool>) -> Vec<Tool> {
|
||||
let original_count = tools.len();
|
||||
let filtered_tools: Vec<Tool> = tools
|
||||
.into_iter()
|
||||
.filter(|tool| {
|
||||
let should_keep = tool.name != SUBAGENT_RUN_TASK_TOOL_NAME;
|
||||
if !should_keep {
|
||||
debug!("Filtering out subagent tool: {}", tool.name);
|
||||
}
|
||||
should_keep
|
||||
})
|
||||
.collect();
|
||||
|
||||
let filtered_count = filtered_tools.len();
|
||||
if filtered_count < original_count {
|
||||
debug!(
|
||||
"Filtered {} subagent tool(s) from {} total tools",
|
||||
original_count - filtered_count,
|
||||
original_count
|
||||
);
|
||||
}
|
||||
|
||||
filtered_tools
|
||||
}
|
||||
|
||||
/// Add platform tools to the subagent's tool list (excluding dangerous tools)
|
||||
async fn add_platform_tools(tools: &mut Vec<Tool>, extension_manager: &ExtensionManager) {
|
||||
debug!("Adding safe platform tools to subagent");
|
||||
|
||||
// Add safe platform tools - subagents can search for extensions but can't manage them or schedules
|
||||
tools.push(platform_tools::search_available_extensions_tool());
|
||||
debug!("Added search_available_extensions tool");
|
||||
|
||||
// Add resource tools if supported - these are generally safe for subagents
|
||||
if extension_manager.supports_resources() {
|
||||
tools.extend([
|
||||
platform_tools::read_resource_tool(),
|
||||
platform_tools::list_resources_tool(),
|
||||
]);
|
||||
debug!("Added 2 resource platform tools");
|
||||
}
|
||||
|
||||
// Note: We explicitly do NOT add these tools for security reasons:
|
||||
// - manage_extensions (could interfere with parent agent's extensions)
|
||||
// - manage_schedule (could interfere with parent agent's scheduling)
|
||||
// - subagent spawning tools (prevent recursion)
|
||||
debug!("Platform tools added successfully (dangerous tools excluded)");
|
||||
}
|
||||
|
||||
/// Check if a tool name is a platform tool that subagents can use
|
||||
fn is_platform_tool(&self, tool_name: &str) -> bool {
|
||||
matches!(
|
||||
tool_name,
|
||||
PLATFORM_SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME
|
||||
| PLATFORM_READ_RESOURCE_TOOL_NAME
|
||||
| PLATFORM_LIST_RESOURCES_TOOL_NAME
|
||||
)
|
||||
}
|
||||
|
||||
/// Handle platform tool calls that are safe for subagents
|
||||
async fn handle_platform_tool_call(
|
||||
&self,
|
||||
tool_call: mcp_core::tool::ToolCall,
|
||||
extension_manager: &ExtensionManager,
|
||||
) -> Result<Vec<mcp_core::Content>, ToolError> {
|
||||
debug!("Handling platform tool: {}", tool_call.name);
|
||||
|
||||
match tool_call.name.as_str() {
|
||||
PLATFORM_SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME => extension_manager
|
||||
.search_available_extensions()
|
||||
.await
|
||||
.map_err(|e| ToolError::ExecutionError(e.to_string())),
|
||||
PLATFORM_READ_RESOURCE_TOOL_NAME => extension_manager
|
||||
.read_resource(tool_call.arguments)
|
||||
.await
|
||||
.map_err(|e| ToolError::ExecutionError(e.to_string())),
|
||||
PLATFORM_LIST_RESOURCES_TOOL_NAME => extension_manager
|
||||
.list_resources(tool_call.arguments)
|
||||
.await
|
||||
.map_err(|e| ToolError::ExecutionError(e.to_string())),
|
||||
_ => Err(ToolError::ExecutionError(format!(
|
||||
"Platform tool '{}' is not available to subagents for security reasons",
|
||||
tool_call.name
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the system prompt for the subagent using the template
|
||||
async fn build_system_prompt(&self, available_tools: &[Tool]) -> Result<String, anyhow::Error> {
|
||||
let mut context = HashMap::new();
|
||||
|
||||
// Add basic context
|
||||
context.insert(
|
||||
"current_date_time",
|
||||
serde_json::Value::String(Utc::now().format("%Y-%m-%d %H:%M:%S UTC").to_string()),
|
||||
);
|
||||
context.insert("subagent_id", serde_json::Value::String(self.id.clone()));
|
||||
|
||||
// Add recipe information if available
|
||||
if let Some(recipe) = &self.config.recipe {
|
||||
context.insert(
|
||||
"recipe_title",
|
||||
serde_json::Value::String(recipe.title.clone()),
|
||||
);
|
||||
}
|
||||
|
||||
// Add max turns if configured
|
||||
if let Some(max_turns) = self.config.max_turns {
|
||||
context.insert(
|
||||
"max_turns",
|
||||
serde_json::Value::Number(serde_json::Number::from(max_turns)),
|
||||
);
|
||||
}
|
||||
|
||||
// Add task instructions
|
||||
let instructions = if let Some(recipe) = &self.config.recipe {
|
||||
recipe.instructions.as_deref().unwrap_or("")
|
||||
} else {
|
||||
self.config.instructions.as_deref().unwrap_or("")
|
||||
};
|
||||
context.insert(
|
||||
"task_instructions",
|
||||
serde_json::Value::String(instructions.to_string()),
|
||||
);
|
||||
|
||||
// Add available extensions (only if we have a recipe and extensions)
|
||||
if self.config.recipe.is_some() {
|
||||
let extensions: Vec<String> = self.recipe_extensions.lock().await.clone();
|
||||
if !extensions.is_empty() {
|
||||
context.insert(
|
||||
"extensions",
|
||||
serde_json::Value::Array(
|
||||
extensions
|
||||
.into_iter()
|
||||
.map(serde_json::Value::String)
|
||||
.collect(),
|
||||
),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Add available tools with descriptions for better context
|
||||
let tools_with_descriptions: Vec<String> = available_tools
|
||||
.iter()
|
||||
.map(|t| {
|
||||
if t.description.is_empty() {
|
||||
t.name.clone()
|
||||
} else {
|
||||
format!("{}: {}", t.name, t.description)
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
context.insert(
|
||||
"available_tools",
|
||||
serde_json::Value::String(if tools_with_descriptions.is_empty() {
|
||||
"None".to_string()
|
||||
} else {
|
||||
tools_with_descriptions.join(", ")
|
||||
}),
|
||||
);
|
||||
|
||||
// Add tool count for context
|
||||
context.insert(
|
||||
"tool_count",
|
||||
serde_json::Value::Number(serde_json::Number::from(available_tools.len())),
|
||||
);
|
||||
|
||||
// Render the subagent system prompt template
|
||||
let system_prompt = render_global_file("subagent_system.md", &context)
|
||||
.map_err(|e| anyhow!("Failed to render subagent system prompt: {}", e))?;
|
||||
|
||||
Ok(system_prompt)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
use anyhow::Result;
|
||||
use mcp_core::{Content, ToolError};
|
||||
use serde_json::Value;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::agents::subagent_types::SpawnSubAgentArgs;
|
||||
use crate::agents::Agent;
|
||||
|
||||
impl Agent {
|
||||
/// Handle running a complete subagent task (replaces the individual spawn/send/check tools)
|
||||
pub async fn handle_run_subagent_task(
|
||||
&self,
|
||||
arguments: Value,
|
||||
) -> Result<Vec<Content>, ToolError> {
|
||||
let subagent_manager = self.subagent_manager.lock().await;
|
||||
let manager = subagent_manager.as_ref().ok_or_else(|| {
|
||||
ToolError::ExecutionError("Subagent manager not initialized".to_string())
|
||||
})?;
|
||||
|
||||
// Parse arguments - using "task" as the main message parameter
|
||||
let message = arguments
|
||||
.get("task")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::ExecutionError("Missing task parameter".to_string()))?
|
||||
.to_string();
|
||||
|
||||
// Either recipe_name or instructions must be provided
|
||||
let recipe_name = arguments
|
||||
.get("recipe_name")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
let instructions = arguments
|
||||
.get("instructions")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
|
||||
let mut args = if let Some(recipe_name) = recipe_name {
|
||||
SpawnSubAgentArgs::new_with_recipe(recipe_name, message.clone())
|
||||
} else if let Some(instructions) = instructions {
|
||||
SpawnSubAgentArgs::new_with_instructions(instructions, message.clone())
|
||||
} else {
|
||||
return Err(ToolError::ExecutionError(
|
||||
"Either recipe_name or instructions parameter must be provided".to_string(),
|
||||
));
|
||||
};
|
||||
|
||||
// Set max_turns with default of 10
|
||||
let max_turns = arguments
|
||||
.get("max_turns")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(10) as usize;
|
||||
args = args.with_max_turns(max_turns);
|
||||
|
||||
if let Some(timeout) = arguments.get("timeout_seconds").and_then(|v| v.as_u64()) {
|
||||
args = args.with_timeout(timeout);
|
||||
}
|
||||
|
||||
// Get the provider from the parent agent
|
||||
let provider = self
|
||||
.provider()
|
||||
.await
|
||||
.map_err(|e| ToolError::ExecutionError(format!("Failed to get provider: {}", e)))?;
|
||||
|
||||
// Get the extension manager from the parent agent
|
||||
let extension_manager = Arc::new(self.extension_manager.read().await);
|
||||
|
||||
// Run the complete subagent task
|
||||
match manager
|
||||
.run_complete_subagent_task(args, provider, extension_manager)
|
||||
.await
|
||||
{
|
||||
Ok(result) => Ok(vec![Content::text(result)]),
|
||||
Err(e) => Err(ToolError::ExecutionError(format!(
|
||||
"Failed to run subagent task: {}",
|
||||
e
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,404 @@
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use mcp_core::protocol::JsonRpcMessage;
|
||||
use tokio::sync::{mpsc, Mutex, RwLock};
|
||||
use tracing::{debug, error, instrument, warn};
|
||||
|
||||
use crate::agents::extension_manager::ExtensionManager;
|
||||
use crate::agents::subagent::{SubAgent, SubAgentConfig, SubAgentProgress, SubAgentStatus};
|
||||
use crate::agents::subagent_types::SpawnSubAgentArgs;
|
||||
use crate::providers::base::Provider;
|
||||
use crate::recipe::Recipe;
|
||||
|
||||
/// Manages the lifecycle of subagents
|
||||
pub struct SubAgentManager {
|
||||
subagents: Arc<RwLock<HashMap<String, Arc<SubAgent>>>>,
|
||||
handles: Arc<Mutex<HashMap<String, tokio::task::JoinHandle<()>>>>,
|
||||
mcp_notification_tx: mpsc::Sender<JsonRpcMessage>,
|
||||
}
|
||||
|
||||
impl SubAgentManager {
|
||||
/// Create a new subagent manager
|
||||
pub fn new(mcp_notification_tx: mpsc::Sender<JsonRpcMessage>) -> Self {
|
||||
Self {
|
||||
subagents: Arc::new(RwLock::new(HashMap::new())),
|
||||
handles: Arc::new(Mutex::new(HashMap::new())),
|
||||
mcp_notification_tx,
|
||||
}
|
||||
}
|
||||
|
||||
/// Spawn a new interactive subagent
|
||||
#[instrument(skip(self, args, provider, extension_manager))]
|
||||
pub async fn spawn_interactive_subagent(
|
||||
&self,
|
||||
args: SpawnSubAgentArgs,
|
||||
provider: Arc<dyn Provider>,
|
||||
extension_manager: Arc<tokio::sync::RwLockReadGuard<'_, ExtensionManager>>,
|
||||
) -> Result<String> {
|
||||
debug!("Spawning interactive subagent");
|
||||
|
||||
// Create subagent config based on whether we have a recipe or instructions
|
||||
let mut config = if let Some(recipe_name) = args.recipe_name {
|
||||
debug!("Using recipe: {}", recipe_name);
|
||||
// Load the recipe
|
||||
let recipe = self.load_recipe(&recipe_name).await?;
|
||||
SubAgentConfig::new_with_recipe(recipe)
|
||||
} else if let Some(instructions) = args.instructions {
|
||||
debug!("Using direct instructions");
|
||||
SubAgentConfig::new_with_instructions(instructions)
|
||||
} else {
|
||||
return Err(anyhow!(
|
||||
"Either recipe_name or instructions must be provided"
|
||||
));
|
||||
};
|
||||
|
||||
if let Some(max_turns) = args.max_turns {
|
||||
config = config.with_max_turns(max_turns);
|
||||
}
|
||||
if let Some(timeout) = args.timeout_seconds {
|
||||
config = config.with_timeout(timeout);
|
||||
}
|
||||
|
||||
// Create the subagent with the parent agent's provider
|
||||
let (subagent, handle) = SubAgent::new(
|
||||
config,
|
||||
Arc::clone(&provider),
|
||||
Arc::clone(&extension_manager),
|
||||
self.mcp_notification_tx.clone(),
|
||||
)
|
||||
.await?;
|
||||
let subagent_id = subagent.id.clone();
|
||||
|
||||
// Store the subagent and its handle
|
||||
{
|
||||
let mut subagents = self.subagents.write().await;
|
||||
subagents.insert(subagent_id.clone(), Arc::clone(&subagent));
|
||||
}
|
||||
{
|
||||
let mut handles = self.handles.lock().await;
|
||||
handles.insert(subagent_id.clone(), handle);
|
||||
}
|
||||
|
||||
// Return immediately - no initial message processing
|
||||
Ok(subagent_id)
|
||||
}
|
||||
|
||||
/// Get a subagent by ID
|
||||
pub async fn get_subagent(&self, id: &str) -> Option<Arc<SubAgent>> {
|
||||
let subagents = self.subagents.read().await;
|
||||
subagents.get(id).cloned()
|
||||
}
|
||||
|
||||
/// List all active subagent IDs
|
||||
pub async fn list_subagents(&self) -> Vec<String> {
|
||||
let subagents = self.subagents.read().await;
|
||||
subagents.keys().cloned().collect()
|
||||
}
|
||||
|
||||
/// Get status of all subagents
|
||||
pub async fn get_subagent_status(&self) -> HashMap<String, SubAgentStatus> {
|
||||
let subagents = self.subagents.read().await;
|
||||
let mut status_map = HashMap::new();
|
||||
|
||||
for (id, subagent) in subagents.iter() {
|
||||
status_map.insert(id.clone(), subagent.get_status().await);
|
||||
}
|
||||
|
||||
status_map
|
||||
}
|
||||
|
||||
/// Get progress of all subagents
|
||||
pub async fn get_subagent_progress(&self) -> HashMap<String, SubAgentProgress> {
|
||||
let subagents = self.subagents.read().await;
|
||||
let mut progress_map = HashMap::new();
|
||||
|
||||
for (id, subagent) in subagents.iter() {
|
||||
progress_map.insert(id.clone(), subagent.get_progress().await);
|
||||
}
|
||||
|
||||
progress_map
|
||||
}
|
||||
|
||||
/// Send a message to a specific subagent
|
||||
#[instrument(skip(self, message, provider, extension_manager))]
|
||||
pub async fn send_message_to_subagent(
|
||||
&self,
|
||||
subagent_id: &str,
|
||||
message: String,
|
||||
provider: Arc<dyn Provider>,
|
||||
extension_manager: Arc<tokio::sync::RwLockReadGuard<'_, ExtensionManager>>,
|
||||
) -> Result<String> {
|
||||
let subagent = self
|
||||
.get_subagent(subagent_id)
|
||||
.await
|
||||
.ok_or_else(|| anyhow!("Subagent {} not found", subagent_id))?;
|
||||
|
||||
// Process the message and get a reply
|
||||
match subagent
|
||||
.reply_subagent(message, provider, extension_manager)
|
||||
.await
|
||||
{
|
||||
Ok(response) => Ok(format!(
|
||||
"Message sent to subagent {}. Response:\n{}",
|
||||
subagent_id,
|
||||
response.as_concat_text()
|
||||
)),
|
||||
Err(e) => Err(anyhow!("Failed to process message in subagent: {}", e)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Terminate a specific subagent
|
||||
#[instrument(skip(self))]
|
||||
pub async fn terminate_subagent(&self, id: &str) -> Result<()> {
|
||||
debug!("Terminating subagent {}", id);
|
||||
|
||||
// Get and terminate the subagent
|
||||
let subagent = {
|
||||
let mut subagents = self.subagents.write().await;
|
||||
subagents.remove(id)
|
||||
};
|
||||
|
||||
if let Some(subagent) = subagent {
|
||||
subagent.terminate().await?;
|
||||
} else {
|
||||
warn!("Attempted to terminate non-existent subagent {}", id);
|
||||
return Err(anyhow!("Subagent {} not found", id));
|
||||
}
|
||||
|
||||
// Clean up the background handle
|
||||
let handle = {
|
||||
let mut handles = self.handles.lock().await;
|
||||
handles.remove(id)
|
||||
};
|
||||
|
||||
if let Some(handle) = handle {
|
||||
handle.abort();
|
||||
}
|
||||
|
||||
debug!("Subagent {} terminated successfully", id);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Terminate all subagents
|
||||
#[instrument(skip(self))]
|
||||
pub async fn terminate_all_subagents(&self) -> Result<()> {
|
||||
debug!("Terminating all subagents");
|
||||
|
||||
let subagent_ids: Vec<String> = {
|
||||
let subagents = self.subagents.read().await;
|
||||
subagents.keys().cloned().collect()
|
||||
};
|
||||
|
||||
for id in subagent_ids {
|
||||
if let Err(e) = self.terminate_subagent(&id).await {
|
||||
error!("Failed to terminate subagent {}: {}", id, e);
|
||||
}
|
||||
}
|
||||
|
||||
debug!("All subagents terminated");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get formatted conversation from a subagent
|
||||
pub async fn get_subagent_conversation(&self, id: &str) -> Result<String> {
|
||||
let subagent = self
|
||||
.get_subagent(id)
|
||||
.await
|
||||
.ok_or_else(|| anyhow!("Subagent {} not found", id))?;
|
||||
|
||||
Ok(subagent.get_formatted_conversation().await)
|
||||
}
|
||||
|
||||
/// Clean up completed or failed subagents
|
||||
pub async fn cleanup_completed_subagents(&self) -> Result<usize> {
|
||||
let mut completed_ids = Vec::new();
|
||||
|
||||
// Find completed subagents
|
||||
{
|
||||
let subagents = self.subagents.read().await;
|
||||
for (id, subagent) in subagents.iter() {
|
||||
if subagent.is_completed().await {
|
||||
completed_ids.push(id.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Remove completed subagents
|
||||
let count = completed_ids.len();
|
||||
for id in completed_ids {
|
||||
if let Err(e) = self.terminate_subagent(&id).await {
|
||||
error!("Failed to cleanup completed subagent {}: {}", id, e);
|
||||
}
|
||||
}
|
||||
|
||||
debug!("Cleaned up {} completed subagents", count);
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
/// Load a recipe from file
|
||||
async fn load_recipe(&self, recipe_name: &str) -> Result<Recipe> {
|
||||
// Try to load from current directory first
|
||||
let recipe_path = if recipe_name.ends_with(".yaml") || recipe_name.ends_with(".yml") {
|
||||
recipe_name.to_string()
|
||||
} else {
|
||||
format!("{}.yaml", recipe_name)
|
||||
};
|
||||
|
||||
if Path::new(&recipe_path).exists() {
|
||||
let content = tokio::fs::read_to_string(&recipe_path).await?;
|
||||
let recipe: Recipe = serde_yaml::from_str(&content)?;
|
||||
return Ok(recipe);
|
||||
}
|
||||
|
||||
// Try some common recipe locations
|
||||
let common_paths = [
|
||||
format!("recipes/{}", recipe_path),
|
||||
format!("./recipes/{}", recipe_path),
|
||||
format!("../recipes/{}", recipe_path),
|
||||
];
|
||||
|
||||
for path in &common_paths {
|
||||
if Path::new(path).exists() {
|
||||
let content = tokio::fs::read_to_string(path).await?;
|
||||
let recipe: Recipe = serde_yaml::from_str(&content)?;
|
||||
return Ok(recipe);
|
||||
}
|
||||
}
|
||||
|
||||
Err(anyhow!(
|
||||
"Recipe file '{}' not found in current directory or common recipe locations",
|
||||
recipe_name
|
||||
))
|
||||
}
|
||||
|
||||
/// Get count of active subagents
|
||||
pub async fn get_active_count(&self) -> usize {
|
||||
let subagents = self.subagents.read().await;
|
||||
subagents.len()
|
||||
}
|
||||
|
||||
/// Check if a subagent exists
|
||||
pub async fn has_subagent(&self, id: &str) -> bool {
|
||||
let subagents = self.subagents.read().await;
|
||||
subagents.contains_key(id)
|
||||
}
|
||||
|
||||
/// Run a complete subagent task (spawn, execute, cleanup)
|
||||
#[instrument(skip(self, args, provider, extension_manager))]
|
||||
pub async fn run_complete_subagent_task(
|
||||
&self,
|
||||
args: SpawnSubAgentArgs,
|
||||
provider: Arc<dyn Provider>,
|
||||
extension_manager: Arc<tokio::sync::RwLockReadGuard<'_, ExtensionManager>>,
|
||||
) -> Result<String> {
|
||||
debug!("Running complete subagent task");
|
||||
|
||||
// Create subagent config based on whether we have a recipe or instructions
|
||||
let mut config = if let Some(recipe_name) = args.recipe_name {
|
||||
debug!("Using recipe: {}", recipe_name);
|
||||
// Load the recipe
|
||||
let recipe = self.load_recipe(&recipe_name).await?;
|
||||
SubAgentConfig::new_with_recipe(recipe)
|
||||
} else if let Some(instructions) = args.instructions {
|
||||
debug!("Using direct instructions");
|
||||
SubAgentConfig::new_with_instructions(instructions)
|
||||
} else {
|
||||
return Err(anyhow!(
|
||||
"Either recipe_name or instructions must be provided"
|
||||
));
|
||||
};
|
||||
|
||||
// Set default max_turns if not provided
|
||||
let max_turns = args.max_turns.unwrap_or(10);
|
||||
config = config.with_max_turns(max_turns);
|
||||
|
||||
if let Some(timeout) = args.timeout_seconds {
|
||||
config = config.with_timeout(timeout);
|
||||
}
|
||||
|
||||
// Create the subagent with the parent agent's provider
|
||||
let (subagent, handle) = SubAgent::new(
|
||||
config,
|
||||
Arc::clone(&provider),
|
||||
Arc::clone(&extension_manager),
|
||||
self.mcp_notification_tx.clone(),
|
||||
)
|
||||
.await?;
|
||||
let subagent_id = subagent.id.clone();
|
||||
|
||||
// Store the subagent and its handle temporarily
|
||||
{
|
||||
let mut subagents = self.subagents.write().await;
|
||||
subagents.insert(subagent_id.clone(), Arc::clone(&subagent));
|
||||
}
|
||||
{
|
||||
let mut handles = self.handles.lock().await;
|
||||
handles.insert(subagent_id.clone(), handle);
|
||||
}
|
||||
|
||||
// Run the complete conversation
|
||||
let mut conversation_result = String::new();
|
||||
let turn_count = 0;
|
||||
let current_message = args.message.clone();
|
||||
|
||||
// For now, we just complete after one turn since we don't have a mechanism
|
||||
// for the subagent to continue autonomously without user input
|
||||
// In a future iteration, we could add logic for the subagent to continue
|
||||
// working on multi-step tasks with proper turn management
|
||||
match subagent
|
||||
.reply_subagent(
|
||||
current_message,
|
||||
Arc::clone(&provider),
|
||||
Arc::clone(&extension_manager),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(response) => {
|
||||
let response_text = response.as_concat_text();
|
||||
conversation_result.push_str(&format!(
|
||||
"\n--- Turn {} ---\n{}",
|
||||
turn_count + 1,
|
||||
response_text
|
||||
));
|
||||
conversation_result.push_str(&format!(
|
||||
"\n[Task completed after {} turns]",
|
||||
turn_count + 1
|
||||
));
|
||||
}
|
||||
Err(e) => {
|
||||
conversation_result
|
||||
.push_str(&format!("\n[Error after {} turns: {}]", turn_count, e));
|
||||
}
|
||||
}
|
||||
|
||||
// Clean up the subagent
|
||||
if let Err(e) = self.terminate_subagent(&subagent_id).await {
|
||||
debug!("Failed to cleanup subagent {}: {}", subagent_id, e);
|
||||
}
|
||||
|
||||
// Return the complete conversation result
|
||||
Ok(format!("Subagent task completed:\n{}", conversation_result))
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for SubAgentManager {
|
||||
fn default() -> Self {
|
||||
// Create a dummy channel for default implementation
|
||||
// In practice, this should not be used - SubAgentManager should be created
|
||||
// with a proper MCP notification sender
|
||||
let (tx, _rx) = mpsc::channel(1);
|
||||
Self::new(tx)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for SubAgentManager {
|
||||
fn drop(&mut self) {
|
||||
// Note: In a real implementation, you might want to spawn a task to clean up
|
||||
// subagents gracefully, but for now we'll rely on the Drop implementations
|
||||
// of the individual components
|
||||
debug!("SubAgentManager dropped");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
use indoc::indoc;
|
||||
use mcp_core::tool::{Tool, ToolAnnotations};
|
||||
use serde_json::json;
|
||||
|
||||
pub const SUBAGENT_RUN_TASK_TOOL_NAME: &str = "subagent__run_task";
|
||||
|
||||
pub fn run_task_subagent_tool() -> Tool {
|
||||
Tool::new(
|
||||
SUBAGENT_RUN_TASK_TOOL_NAME.to_string(),
|
||||
indoc! {r#"
|
||||
Spawn a specialized subagent to handle a specific task completely and automatically.
|
||||
|
||||
This tool creates a subagent, processes your task through a complete conversation,
|
||||
and returns the final result. The subagent is automatically cleaned up after completion.
|
||||
|
||||
You can configure the subagent in two ways:
|
||||
1. Using a recipe file that defines instructions, extensions, and behavior
|
||||
2. Providing direct instructions for ad-hoc tasks
|
||||
|
||||
The subagent will work autonomously until the task is complete, it reaches max_turns,
|
||||
or it encounters an error. You'll get the final result without needing to manage
|
||||
the subagent lifecycle manually.
|
||||
|
||||
Examples:
|
||||
- "Convert these unittest files to pytest format: file1.py, file2.py"
|
||||
- "Research the latest developments in AI and provide a comprehensive summary"
|
||||
- "Review this code for security vulnerabilities and suggest fixes"
|
||||
- "Refactor this legacy code to use modern Python patterns"
|
||||
"#}
|
||||
.to_string(),
|
||||
json!({
|
||||
"type": "object",
|
||||
"required": ["task"],
|
||||
"properties": {
|
||||
"recipe_name": {
|
||||
"type": "string",
|
||||
"description": "Name of the recipe file to configure the subagent (e.g., 'research_assistant_recipe.yaml'). Either this or 'instructions' must be provided."
|
||||
},
|
||||
"instructions": {
|
||||
"type": "string",
|
||||
"description": "Direct instructions for the subagent's task. Either this or 'recipe_name' must be provided. Example: 'You are a code refactoring assistant. Help convert unittest tests to pytest format.'"
|
||||
},
|
||||
"task": {
|
||||
"type": "string",
|
||||
"description": "The task description or initial message for the subagent to work on"
|
||||
},
|
||||
"max_turns": {
|
||||
"type": "integer",
|
||||
"description": "Maximum number of conversation turns before auto-completion (default: 10)",
|
||||
"minimum": 1,
|
||||
"default": 10
|
||||
},
|
||||
"timeout_seconds": {
|
||||
"type": "integer",
|
||||
"description": "Optional timeout for the entire task in seconds",
|
||||
"minimum": 1
|
||||
}
|
||||
}
|
||||
}),
|
||||
Some(ToolAnnotations {
|
||||
title: Some("Run subagent task".to_string()),
|
||||
read_only_hint: false,
|
||||
destructive_hint: false,
|
||||
idempotent_hint: false,
|
||||
open_world_hint: false,
|
||||
}),
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SpawnSubAgentArgs {
|
||||
pub recipe_name: Option<String>,
|
||||
pub instructions: Option<String>,
|
||||
pub message: String,
|
||||
pub max_turns: Option<usize>,
|
||||
pub timeout_seconds: Option<u64>,
|
||||
}
|
||||
|
||||
impl SpawnSubAgentArgs {
|
||||
pub fn new_with_recipe(recipe_name: String, message: String) -> Self {
|
||||
Self {
|
||||
recipe_name: Some(recipe_name),
|
||||
instructions: None,
|
||||
message,
|
||||
max_turns: None,
|
||||
timeout_seconds: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn new_with_instructions(instructions: String, message: String) -> Self {
|
||||
Self {
|
||||
recipe_name: None,
|
||||
instructions: Some(instructions),
|
||||
message,
|
||||
max_turns: None,
|
||||
timeout_seconds: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_max_turns(mut self, max_turns: usize) -> Self {
|
||||
self.max_turns = Some(max_turns);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_timeout(mut self, timeout_seconds: u64) -> Self {
|
||||
self.timeout_seconds = Some(timeout_seconds);
|
||||
self
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user