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

Revert "Revert "Rewrite extension management tools"" (#5273)

Browse Source 
Signed-off-by: Angela Ning <aning@squareup.com>
This commit is contained in:
Angela Ning
2025-10-21 11:34:17 -04:00
committed by GitHub
parent 3546ddb8c2
commit 09752aa7d0
19 changed files with 933 additions and 357 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/agent.rs
+16 -180
View File
@@ -10,12 +10,9 @@ use uuid::Uuid;
use crate::agents::extension::{ExtensionConfig, ExtensionError, ExtensionResult, ToolInfo};
use crate::agents::extension_manager::{get_parameter_names, ExtensionManager};
use crate::agents::extension_manager_extension::MANAGE_EXTENSIONS_TOOL_NAME_COMPLETE;
use crate::agents::final_output_tool::{FINAL_OUTPUT_CONTINUATION_MESSAGE, FINAL_OUTPUT_TOOL_NAME};
use crate::agents::platform_tools::{
PLATFORM_LIST_RESOURCES_TOOL_NAME, PLATFORM_MANAGE_EXTENSIONS_TOOL_NAME,
PLATFORM_MANAGE_SCHEDULE_TOOL_NAME, PLATFORM_READ_RESOURCE_TOOL_NAME,
PLATFORM_SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME,
};
use crate::agents::platform_tools::PLATFORM_MANAGE_SCHEDULE_TOOL_NAME;
use crate::agents::prompt_manager::PromptManager;
use crate::agents::recipe_tools::dynamic_task_tools::{
create_dynamic_task, create_dynamic_task_tool, DYNAMIC_TASK_TOOL_NAME_PREFIX,
@@ -32,7 +29,7 @@ use crate::agents::tool_route_manager::ToolRouteManager;
use crate::agents::tool_router_index_manager::ToolRouterIndexManager;
use crate::agents::types::SessionConfig;
use crate::agents::types::{FrontendTool, ToolResultReceiver};
use crate::config::{get_enabled_extensions, get_extension_by_name, Config};
use crate::config::{get_enabled_extensions, Config};
use crate::context_mgmt::DEFAULT_COMPACTION_THRESHOLD;
use crate::conversation::{debug_conversation_fix, fix_conversation, Conversation};
use crate::mcp_utils::ToolResult;
@@ -88,7 +85,7 @@ pub struct ToolCategorizeResult {
/// The main goose Agent
pub struct Agent {
pub(super) provider: Mutex<Option<Arc<dyn Provider>>>,
pub extension_manager: ExtensionManager,
pub extension_manager: Arc<ExtensionManager>,
pub(super) sub_recipe_manager: Mutex<SubRecipeManager>,
pub(super) tasks_manager: TasksManager,
pub(super) final_output_tool: Arc<Mutex<Option<FinalOutputTool>>>,
@@ -100,7 +97,7 @@ pub struct Agent {
pub(super) tool_result_tx: mpsc::Sender<(String, ToolResult<Vec<Content>>)>,
pub(super) tool_result_rx: ToolResultReceiver,
pub(super) tool_route_manager: ToolRouteManager,
pub tool_route_manager: Arc<ToolRouteManager>,
pub(super) scheduler_service: Mutex<Option<Arc<dyn SchedulerTrait>>>,
pub(super) retry_manager: RetryManager,
pub(super) tool_inspection_manager: ToolInspectionManager,
@@ -163,7 +160,7 @@ impl Agent {
Self {
provider: Mutex::new(None),
extension_manager: ExtensionManager::new(),
extension_manager: Arc::new(ExtensionManager::new()),
sub_recipe_manager: Mutex::new(SubRecipeManager::new()),
tasks_manager: TasksManager::new(),
final_output_tool: Arc::new(Mutex::new(None)),
@@ -174,7 +171,7 @@ impl Agent {
confirmation_rx: Mutex::new(confirm_rx),
tool_result_tx: tool_tx,
tool_result_rx: Arc::new(Mutex::new(tool_rx)),
tool_route_manager: ToolRouteManager::new(),
tool_route_manager: Arc::new(ToolRouteManager::new()),
scheduler_service: Mutex::new(None),
retry_manager: RetryManager::new(),
tool_inspection_manager: Self::create_default_tool_inspection_manager(),
@@ -404,28 +401,6 @@ impl Agent {
return (request_id, Ok(ToolCallResult::from(result)));
}
if tool_call.name == PLATFORM_MANAGE_EXTENSIONS_TOOL_NAME {
let extension_name = tool_call
.arguments
.as_ref()
.and_then(|args| args.get("extension_name"))
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let action = tool_call
.arguments
.as_ref()
.and_then(|args| args.get("action"))
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let (request_id, result) = self
.manage_extensions(action, extension_name, request_id)
.await;
return (request_id, Ok(ToolCallResult::from(result)));
}
if tool_call.name == FINAL_OUTPUT_TOOL_NAME {
return if let Some(final_output_tool) = self.final_output_tool.lock().await.as_mut() {
let result = final_output_tool.execute_tool_call(tool_call.clone()).await;
@@ -488,12 +463,12 @@ impl Agent {
let parent_session_id = session.id.to_string();
let parent_working_dir = session.working_dir.clone();
let task_config = TaskConfig::new(
provider,
parent_session_id,
parent_working_dir,
get_enabled_extensions(),
);
// Get extensions from the agent's runtime state rather than global config
// This ensures subagents inherit extensions that were dynamically enabled by the parent
let extensions = self.get_extension_configs().await;
let task_config =
TaskConfig::new(provider, parent_session_id, parent_working_dir, extensions);
let arguments = match tool_call.arguments.clone() {
Some(args) => Value::Object(args),
@@ -560,31 +535,6 @@ impl Agent {
.map(Value::Object)
.unwrap_or(Value::Object(serde_json::Map::new()));
create_dynamic_task(arguments, &self.tasks_manager, loaded_extensions).await
} else if tool_call.name == PLATFORM_READ_RESOURCE_TOOL_NAME {
// Check if the tool is read_resource and handle it separately
let arguments = tool_call
.arguments
.clone()
.map(Value::Object)
.unwrap_or(Value::Object(serde_json::Map::new()));
ToolCallResult::from(
self.extension_manager
.read_resource(arguments, cancellation_token.unwrap_or_default())
.await,
)
} else if tool_call.name == PLATFORM_LIST_RESOURCES_TOOL_NAME {
let arguments = tool_call
.arguments
.clone()
.map(Value::Object)
.unwrap_or(Value::Object(serde_json::Map::new()));
ToolCallResult::from(
self.extension_manager
.list_resources(arguments, cancellation_token.unwrap_or_default())
.await,
)
} else if tool_call.name == PLATFORM_SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME {
ToolCallResult::from(self.extension_manager.search_available_extensions().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(ErrorData::new(
@@ -653,109 +603,6 @@ impl Agent {
Ok(())
}
#[allow(clippy::too_many_lines)]
pub(super) async fn manage_extensions(
&self,
action: String,
extension_name: String,
request_id: String,
) -> (String, Result<Vec<Content>, ErrorData>) {
if self.tool_route_manager.is_router_functional().await {
let selector = self.tool_route_manager.get_router_tool_selector().await;
if let Some(selector) = selector {
let selector_action = if action == "disable" { "remove" } else { "add" };
let selector = Arc::new(selector);
if let Err(e) = ToolRouterIndexManager::update_extension_tools(
&selector,
&self.extension_manager,
&extension_name,
selector_action,
)
.await
{
return (
request_id,
Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
format!("Failed to update LLM index: {}", e),
None,
)),
);
}
}
}
if action == "disable" {
let result = self
.extension_manager
.remove_extension(&extension_name)
.await
.map(|_| {
vec![Content::text(format!(
"The extension '{}' has been disabled successfully",
extension_name
))]
})
.map_err(|e| ErrorData::new(ErrorCode::INTERNAL_ERROR, e.to_string(), None));
return (request_id, result);
}
let config = match get_extension_by_name(&extension_name) {
Some(config) => config,
None => {
return (
request_id,
Err(ErrorData::new(
ErrorCode::RESOURCE_NOT_FOUND,
format!(
"Extension '{}' not found. Please check the extension name and try again.",
extension_name
),
None,
)),
)
}
};
let result = self
.extension_manager
.add_extension(config)
.await
.map(|_| {
vec![Content::text(format!(
"The extension '{}' has been installed successfully",
extension_name
))]
})
.map_err(|e| ErrorData::new(ErrorCode::INTERNAL_ERROR, e.to_string(), None));
// Update LLM index if operation was successful and LLM routing is functional
if result.is_ok() && self.tool_route_manager.is_router_functional().await {
let selector = self.tool_route_manager.get_router_tool_selector().await;
if let Some(selector) = selector {
let llm_action = if action == "disable" { "remove" } else { "add" };
let selector = Arc::new(selector);
if let Err(e) = ToolRouterIndexManager::update_extension_tools(
&selector,
&self.extension_manager,
&extension_name,
llm_action,
)
.await
{
return (
request_id,
Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
format!("Failed to update LLM index: {}", e),
None,
)),
);
}
}
}
(request_id, result)
}
pub async fn add_extension(&self, extension: ExtensionConfig) -> ExtensionResult<()> {
match &extension {
ExtensionConfig::Frontend {
@@ -824,21 +671,10 @@ impl Agent {
if extension_name.is_none() || extension_name.as_deref() == Some("platform") {
// Add platform tools
prefixed_tools.extend([
platform_tools::search_available_extensions_tool(),
platform_tools::manage_extensions_tool(),
platform_tools::manage_schedule_tool(),
]);
// TODO: migrate the manage schedule tool as well
prefixed_tools.extend([platform_tools::manage_schedule_tool()]);
// Dynamic task tool
prefixed_tools.push(create_dynamic_task_tool());
// Add resource tools if supported
if self.extension_manager.supports_resources().await {
prefixed_tools.extend([
platform_tools::read_resource_tool(),
platform_tools::list_resources_tool(),
]);
}
}
if extension_name.is_none() {
@@ -1217,7 +1053,7 @@ impl Agent {
let mut enable_extension_request_ids = vec![];
for request in &remaining_requests {
if let Ok(tool_call) = &request.tool_call {
if tool_call.name == PLATFORM_MANAGE_EXTENSIONS_TOOL_NAME {
if tool_call.name == MANAGE_EXTENSIONS_TOOL_NAME_COMPLETE {
enable_extension_request_ids.push(request.id.clone());
}
}