diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 9b7a8394e1ba..214e704c1adc 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -304,6 +304,31 @@ impl Agent { ) -> (String, Result, ToolError>) { let mut extension_manager = self.extension_manager.lock().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 selector = Arc::new(selector); + if let Err(e) = ToolRouterIndexManager::update_extension_tools( + &selector, + &extension_manager, + &extension_name, + selector_action, + ) + .await + { + return ( + request_id, + Err(ToolError::ExecutionError(format!( + "Failed to update vector index: {}", + e + ))), + ); + } + } + } + if action == "disable" { let result = extension_manager .remove_extension(&extension_name) @@ -351,34 +376,6 @@ 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.lock().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) } @@ -503,9 +500,6 @@ impl Agent { } pub async fn remove_extension(&self, name: &str) -> Result<()> { - let mut extension_manager = self.extension_manager.lock().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) { @@ -521,6 +515,9 @@ impl Agent { } } + let mut extension_manager = self.extension_manager.lock().await; + extension_manager.remove_extension(name).await?; + Ok(()) }