tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

subagents (#2797)

Browse Source 
Co-authored-by: Lily Delalande <ldelalande@squareup.com>
This commit is contained in:
Wendy Tang
2025-06-25 16:18:48 -07:00
committed by GitHub
parent c3acddc360
commit d9d7eb0697
17 changed files with 1709 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+187 -68
View File
@@ -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) = &notif.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) = &notif.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