feat: consolidate subagent execution for dynamic tasks (#3444)

Co-authored-by: Lifei Zhou <lifei@squareup.com>
This commit is contained in:
Wendy Tang
2025-07-17 18:48:00 -07:00
committed by GitHub
parent 9f65455cc3
commit 4771ffbb5b
36 changed files with 751 additions and 1062 deletions
+33 -69
View File
@@ -1,79 +1,43 @@
use crate::agents::subagent::SubAgent;
use crate::agents::subagent_task_config::TaskConfig;
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;
/// Standalone function to run a complete subagent task
pub async fn run_complete_subagent_task(
task_arguments: Value,
task_config: TaskConfig,
) -> Result<Vec<Content>, ToolError> {
// Parse arguments - using "task" as the main message parameter
let text_instruction = task_arguments
.get("text_instruction")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::ExecutionError("Missing text_instruction parameter".to_string()))?
.to_string();
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())
})?;
// Create the subagent with the parent agent's provider
let (subagent, handle) = SubAgent::new(task_config.clone())
.await
.map_err(|e| ToolError::ExecutionError(format!("Failed to create subagent: {}", e)))?;
// 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);
// Execute the subagent task
let result = match subagent.reply_subagent(text_instruction, task_config).await {
Ok(response) => {
let response_text = response.as_concat_text();
Ok(vec![Content::text(response_text)])
}
Err(e) => Err(ToolError::ExecutionError(format!(
"Subagent execution failed: {}",
e
))),
};
// 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
))),
}
// Clean up the subagent handle
if let Err(e) = handle.await {
tracing::debug!("Subagent handle cleanup error: {}", e);
}
// Return the result
result
}