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

Signed-off-by: Angela Ning <aning@squareup.com>

crates/goose/tests/agent.rs

@@ -3,6 +3,7 @@ use std::sync::Arc;
 use anyhow::Result;
 use futures::StreamExt;
 use goose::agents::{Agent, AgentEvent};
+use goose::config::extensions::{set_extension, ExtensionEntry};
 use goose::conversation::message::Message;
 use goose::conversation::Conversation;
 use goose::model::ModelConfig;
@@ -1139,3 +1140,271 @@ mod max_turns_tests {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod extension_manager_tests {
+    use super::*;
+    use goose::agents::extension::{ExtensionConfig, PlatformExtensionContext};
+    use goose::agents::extension_manager_extension::{
+        MANAGE_EXTENSIONS_TOOL_NAME, SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME,
+    };
+    use rmcp::model::CallToolRequestParam;
+    use rmcp::object;
+
+    async fn setup_agent_with_extension_manager() -> Agent {
+        // Add the TODO extension to the config so it can be discovered by search_available_extensions
+        // Set it as disabled initially so tests can enable it
+        let todo_extension_entry = ExtensionEntry {
+            enabled: false,
+            config: ExtensionConfig::Platform {
+                name: "todo".to_string(),
+                description:
+                    "Enable a todo list for Goose so it can keep track of what it is doing"
+                        .to_string(),
+                bundled: Some(true),
+                available_tools: vec![],
+            },
+        };
+        set_extension(todo_extension_entry);
+
+        let agent = Agent::new();
+
+        agent
+            .extension_manager
+            .set_context(PlatformExtensionContext {
+                session_id: Some("test_session".to_string()),
+                extension_manager: Some(Arc::downgrade(&agent.extension_manager)),
+                tool_route_manager: Some(Arc::downgrade(&agent.tool_route_manager)),
+            })
+            .await;
+
+        // Now add the extension manager platform extension
+        let ext_config = ExtensionConfig::Platform {
+            name: "extensionmanager".to_string(),
+            description: "Extension Manager".to_string(),
+            bundled: Some(true),
+            available_tools: vec![],
+        };
+
+        agent
+            .add_extension(ext_config)
+            .await
+            .expect("Failed to add extension manager");
+        agent
+    }
+
+    #[tokio::test]
+    async fn test_extension_manager_tools_available() {
+        let agent = setup_agent_with_extension_manager().await;
+        let tools = agent.list_tools(None).await;
+
+        // Note: Tool names are prefixed with the normalized extension name "extensionmanager"
+        // not the display name "Extension Manager"
+        let search_tool = tools.iter().find(|tool| {
+            tool.name == format!("extensionmanager__{SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME}")
+        });
+        assert!(
+            search_tool.is_some(),
+            "search_available_extensions tool should be available"
+        );
+
+        let manage_tool = tools
+            .iter()
+            .find(|tool| tool.name == format!("extensionmanager__{MANAGE_EXTENSIONS_TOOL_NAME}"));
+        assert!(
+            manage_tool.is_some(),
+            "manage_extensions tool should be available"
+        );
+    }
+
+    #[tokio::test]
+    async fn test_search_available_extensions_tool_call() {
+        let agent = setup_agent_with_extension_manager().await;
+
+        let tool_call = CallToolRequestParam {
+            name: format!("extensionmanager__{SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME}").into(),
+            arguments: Some(object!({})),
+        };
+
+        let (_, result) = agent
+            .dispatch_tool_call(tool_call, "request_1".to_string(), None, None)
+            .await;
+
+        assert!(result.is_ok(), "search_available_extensions should succeed");
+        let call_result = result.unwrap().result.await;
+        assert!(
+            call_result.is_ok(),
+            "search_available_extensions execution should succeed"
+        );
+
+        let content = call_result.unwrap();
+        assert!(!content.is_empty(), "Should return some content");
+
+        // Verify the content contains expected text
+        let text = content.first().unwrap().as_text().unwrap();
+        assert!(
+            text.text.contains("Extensions available to enable:"),
+            "Content should contain 'Extensions available to enable:'"
+        );
+    }
+
+    #[tokio::test]
+    async fn test_manage_extensions_enable_disable_success() {
+        let agent = setup_agent_with_extension_manager().await;
+
+        let enable_call = CallToolRequestParam {
+            name: format!("extensionmanager__{MANAGE_EXTENSIONS_TOOL_NAME}").into(),
+            arguments: Some(object!({
+                "action": "enable",
+                "extension_name": "todo"
+            })),
+        };
+        let (_, result) = agent
+            .dispatch_tool_call(enable_call, "request_3a".to_string(), None, None)
+            .await;
+        assert!(result.is_ok());
+        let call_result = result.unwrap().result.await;
+        assert!(
+            call_result.is_ok(),
+            "manage_extensions enable execution should succeed"
+        );
+
+        let content = call_result.unwrap();
+        let text = content.first().unwrap().as_text().unwrap();
+        assert!(
+            text.text.contains("todo") && text.text.contains("installed successfully"),
+            "Response should indicate success, got: {}",
+            text.text
+        );
+
+        // Now disable it
+        let disable_call = CallToolRequestParam {
+            name: format!("extensionmanager__{MANAGE_EXTENSIONS_TOOL_NAME}").into(),
+            arguments: Some(object!({
+                "action": "disable",
+                "extension_name": "todo"
+            })),
+        };
+
+        let (_, result) = agent
+            .dispatch_tool_call(disable_call, "request_3b".to_string(), None, None)
+            .await;
+
+        assert!(result.is_ok(), "manage_extensions disable should succeed");
+        let call_result = result.unwrap().result.await;
+        assert!(
+            call_result.is_ok(),
+            "manage_extensions disable execution should succeed"
+        );
+
+        let content = call_result.unwrap();
+        assert!(!content.is_empty(), "Should return confirmation message");
+
+        // Verify the message indicates success
+        let text = content.first().unwrap().as_text().unwrap();
+        assert!(
+            text.text.contains("successfully") || text.text.contains("disabled"),
+            "Response should indicate success, got: {}",
+            text.text
+        );
+    }
+
+    #[tokio::test]
+    async fn test_manage_extensions_missing_parameters() {
+        let agent = setup_agent_with_extension_manager().await;
+
+        // Call manage_extensions without action parameter
+        let tool_call = CallToolRequestParam {
+            name: format!("extensionmanager__{MANAGE_EXTENSIONS_TOOL_NAME}").into(),
+            arguments: Some(object!({
+                "extension_name": "todo"
+            })),
+        };
+
+        let (_, result) = agent
+            .dispatch_tool_call(tool_call, "request_4".to_string(), None, None)
+            .await;
+
+        assert!(result.is_ok(), "Tool call should return a result");
+        let call_result = result.unwrap().result.await;
+        assert!(
+            call_result.is_ok(),
+            "Tool execution should return an error result"
+        );
+
+        let content = call_result.unwrap();
+        let text = content.first().unwrap().as_text().unwrap();
+        assert!(
+            text.text.contains("action") || text.text.contains("parameter"),
+            "Error should mention missing action parameter"
+        );
+    }
+
+    #[tokio::test]
+    async fn test_manage_extensions_invalid_action() {
+        let agent = setup_agent_with_extension_manager().await;
+
+        let tool_call = CallToolRequestParam {
+            name: format!("extensionmanager__{MANAGE_EXTENSIONS_TOOL_NAME}").into(),
+            arguments: Some(object!({
+                "action": "invalid_action",
+                "extension_name": "todo"
+            })),
+        };
+
+        let (_, result) = agent
+            .dispatch_tool_call(tool_call, "request_6".to_string(), None, None)
+            .await;
+
+        assert!(result.is_ok(), "Tool call should return a result");
+        let call_result = result.unwrap().result.await;
+        assert!(
+            call_result.is_ok(),
+            "Tool execution should return an error result"
+        );
+
+        let content = call_result.unwrap();
+        let text = content.first().unwrap().as_text().unwrap();
+        assert!(
+            text.text.contains("Invalid action")
+                || text.text.contains("enable")
+                || text.text.contains("disable"),
+            "Error should mention invalid action, got: {}",
+            text.text
+        );
+    }
+
+    #[tokio::test]
+    async fn test_manage_extensions_extension_not_found() {
+        let agent = setup_agent_with_extension_manager().await;
+
+        // Try to enable a non-existent extension
+        let tool_call = CallToolRequestParam {
+            name: format!("extensionmanager__{MANAGE_EXTENSIONS_TOOL_NAME}").into(),
+            arguments: Some(object!({
+                "action": "enable",
+                "extension_name": "nonexistent_extension_12345"
+            })),
+        };
+
+        let (_, result) = agent
+            .dispatch_tool_call(tool_call, "request_7".to_string(), None, None)
+            .await;
+
+        assert!(result.is_ok(), "Tool call should return a result");
+        let call_result = result.unwrap().result.await;
+        assert!(
+            call_result.is_ok(),
+            "Tool execution should return an error result"
+        );
+
+        // Check that the error message indicates extension not found
+        let content = call_result.unwrap();
+        let text = content.first().unwrap().as_text().unwrap();
+        assert!(
+            text.text.contains("not found") || text.text.contains("Extension"),
+            "Error should mention extension not found, got: {}",
+            text.text
+        );
+    }
+}