feat: subagent independent extension manager (#3596)

This commit is contained in:
Wendy Tang
2025-07-23 13:08:17 -07:00
committed by GitHub
parent 74140ad0f8
commit a358e264a5
3 changed files with 37 additions and 24 deletions
+1 -3
View File
@@ -344,9 +344,7 @@ impl Agent {
let provider = self.provider().await.ok(); let provider = self.provider().await.ok();
let mcp_tx = self.mcp_tx.lock().await.clone(); let mcp_tx = self.mcp_tx.lock().await.clone();
let task_config = let task_config = TaskConfig::new(provider, mcp_tx);
TaskConfig::new(provider, Some(Arc::clone(&self.extension_manager)), mcp_tx);
subagent_execute_task_tool::run_tasks( subagent_execute_task_tool::run_tasks(
tool_call.arguments.clone(), tool_call.arguments.clone(),
task_config, task_config,
+34 -11
View File
@@ -1,5 +1,7 @@
use crate::{ use crate::{
agents::{Agent, TaskConfig}, agents::extension::ExtensionConfig,
agents::{extension_manager::ExtensionManager, Agent, TaskConfig},
config::ExtensionConfigManager,
message::{Message, MessageContent, ToolRequest}, message::{Message, MessageContent, ToolRequest},
prompt_template::render_global_file, prompt_template::render_global_file,
providers::errors::ProviderError, providers::errors::ProviderError,
@@ -43,6 +45,7 @@ pub struct SubAgent {
pub config: TaskConfig, pub config: TaskConfig,
pub turn_count: Arc<Mutex<usize>>, pub turn_count: Arc<Mutex<usize>>,
pub created_at: DateTime<Utc>, pub created_at: DateTime<Utc>,
pub extension_manager: Arc<RwLock<ExtensionManager>>,
} }
impl SubAgent { impl SubAgent {
@@ -53,6 +56,29 @@ impl SubAgent {
) -> Result<(Arc<Self>, tokio::task::JoinHandle<()>), anyhow::Error> { ) -> Result<(Arc<Self>, tokio::task::JoinHandle<()>), anyhow::Error> {
debug!("Creating new subagent with id: {}", task_config.id); debug!("Creating new subagent with id: {}", task_config.id);
// Create a new extension manager for this subagent
let mut extension_manager = ExtensionManager::new();
// Add extensions based on task_type:
// 1. If executing dynamic task (task_type = 'text_instruction'), default to using all enabled extensions
// 2. (TODO) If executing a sub-recipe task, only use recipe extensions
// Get all enabled extensions from config
let enabled_extensions = ExtensionConfigManager::get_all()
.unwrap_or_default()
.into_iter()
.filter(|ext| ext.enabled)
.map(|ext| ext.config)
.collect::<Vec<ExtensionConfig>>();
// Add enabled extensions to the subagent's extension manager
for extension in enabled_extensions {
if let Err(e) = extension_manager.add_extension(extension).await {
debug!("Failed to add extension to subagent: {}", e);
// Continue with other extensions even if one fails
}
}
let subagent = Arc::new(SubAgent { let subagent = Arc::new(SubAgent {
id: task_config.id.clone(), id: task_config.id.clone(),
conversation: Arc::new(Mutex::new(Vec::new())), conversation: Arc::new(Mutex::new(Vec::new())),
@@ -60,6 +86,7 @@ impl SubAgent {
config: task_config, config: task_config,
turn_count: Arc::new(Mutex::new(0)), turn_count: Arc::new(Mutex::new(0)),
created_at: Utc::now(), created_at: Utc::now(),
extension_manager: Arc::new(RwLock::new(extension_manager)),
}); });
// Send initial MCP notification // Send initial MCP notification
@@ -169,19 +196,13 @@ impl SubAgent {
self.send_mcp_notification("message_processing", &format!("Processing: {}", message)) self.send_mcp_notification("message_processing", &format!("Processing: {}", message))
.await; .await;
// Get provider and extension manager from task config // Get provider from task config
let provider = self let provider = self
.config .config
.provider .provider
.as_ref() .as_ref()
.ok_or_else(|| anyhow!("No provider configured for subagent"))?; .ok_or_else(|| anyhow!("No provider configured for subagent"))?;
let extension_manager = self
.config
.extension_manager
.as_ref()
.ok_or_else(|| anyhow!("No extension manager configured for subagent"))?;
// Check if we've exceeded max turns // Check if we've exceeded max turns
{ {
let turn_count = *self.turn_count.lock().await; let turn_count = *self.turn_count.lock().await;
@@ -220,8 +241,9 @@ impl SubAgent {
// Get the current conversation for context // Get the current conversation for context
let mut messages = self.get_conversation().await; let mut messages = self.get_conversation().await;
// Get tools based on whether we're using a recipe or inheriting from parent // Get tools from the subagent's own extension manager
let tools: Vec<Tool> = extension_manager let tools: Vec<Tool> = self
.extension_manager
.read() .read()
.await .await
.get_prefixed_tools(None) .get_prefixed_tools(None)
@@ -292,7 +314,8 @@ impl SubAgent {
.await; .await;
// Handle platform tools or dispatch to extension manager // Handle platform tools or dispatch to extension manager
let tool_result = match extension_manager let tool_result = match self
.extension_manager
.read() .read()
.await .await
.dispatch_tool_call(tool_call.clone()) .dispatch_tool_call(tool_call.clone())
@@ -1,9 +1,8 @@
use crate::agents::extension_manager::ExtensionManager;
use crate::providers::base::Provider; use crate::providers::base::Provider;
use rmcp::model::JsonRpcMessage; use rmcp::model::JsonRpcMessage;
use std::fmt; use std::fmt;
use std::sync::Arc; use std::sync::Arc;
use tokio::sync::{mpsc, RwLock}; use tokio::sync::mpsc;
use uuid::Uuid; use uuid::Uuid;
/// Configuration for task execution with all necessary dependencies /// Configuration for task execution with all necessary dependencies
@@ -11,7 +10,6 @@ use uuid::Uuid;
pub struct TaskConfig { pub struct TaskConfig {
pub id: String, pub id: String,
pub provider: Option<Arc<dyn Provider>>, pub provider: Option<Arc<dyn Provider>>,
pub extension_manager: Option<Arc<RwLock<ExtensionManager>>>,
pub mcp_tx: mpsc::Sender<JsonRpcMessage>, pub mcp_tx: mpsc::Sender<JsonRpcMessage>,
pub max_turns: Option<usize>, pub max_turns: Option<usize>,
} }
@@ -21,7 +19,6 @@ impl fmt::Debug for TaskConfig {
f.debug_struct("TaskConfig") f.debug_struct("TaskConfig")
.field("id", &self.id) .field("id", &self.id)
.field("provider", &"<dyn Provider>") .field("provider", &"<dyn Provider>")
.field("extension_manager", &"<ExtensionManager>")
.field("max_turns", &self.max_turns) .field("max_turns", &self.max_turns)
.finish() .finish()
} }
@@ -29,15 +26,10 @@ impl fmt::Debug for TaskConfig {
impl TaskConfig { impl TaskConfig {
/// Create a new TaskConfig with all required dependencies /// Create a new TaskConfig with all required dependencies
pub fn new( pub fn new(provider: Option<Arc<dyn Provider>>, mcp_tx: mpsc::Sender<JsonRpcMessage>) -> Self {
provider: Option<Arc<dyn Provider>>,
extension_manager: Option<Arc<RwLock<ExtensionManager>>>,
mcp_tx: mpsc::Sender<JsonRpcMessage>,
) -> Self {
Self { Self {
id: Uuid::new_v4().to_string(), id: Uuid::new_v4().to_string(),
provider, provider,
extension_manager,
mcp_tx, mcp_tx,
max_turns: Some(10), max_turns: Some(10),
} }