From c0042af553c9b25509cf5c4e7a33339206749cec Mon Sep 17 00:00:00 2001
From: Wendy Tang <wendytang@squareup.com>
Date: Fri, 20 Jun 2025 11:07:42 -0700
Subject: [PATCH] fix: disable extension  (#3000)

---
 crates/goose/src/agents/agent.rs | 59 +++++++++++++++-----------------
 1 file changed, 28 insertions(+), 31 deletions(-)

diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs
index 3f56dacc..c9d1e182 100644
--- a/crates/goose/src/agents/agent.rs
+++ b/crates/goose/src/agents/agent.rs
@@ -321,6 +321,31 @@ impl Agent {
     ) -> (String, Result<Vec<Content>, 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)
@@ -368,34 +393,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)
     }
 
@@ -521,9 +518,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) {
@@ -539,6 +533,9 @@ impl Agent {
             }
         }
 
+        let mut extension_manager = self.extension_manager.lock().await;
+        extension_manager.remove_extension(name).await?;
+
         Ok(())
     }