Adrian Cole
496 lines
18 KiB
Rust
496 lines
18 KiB
Rust
use crate::agents::extension::PlatformExtensionContext;
|
|
use crate::agents::mcp_client::{Error, McpClientTrait};
|
|
use crate::config::get_extension_by_name;
|
|
use anyhow::Result;
|
|
use async_trait::async_trait;
|
|
use indoc::indoc;
|
|
use rmcp::model::{
|
|
CallToolResult, Content, ErrorCode, ErrorData, GetPromptResult, Implementation,
|
|
InitializeResult, JsonObject, ListPromptsResult, ListResourcesResult, ListToolsResult,
|
|
ProtocolVersion, ReadResourceResult, ServerCapabilities, ServerNotification, Tool,
|
|
ToolAnnotations, ToolsCapability,
|
|
};
|
|
use schemars::{schema_for, JsonSchema};
|
|
use serde::{Deserialize, Serialize};
|
|
use serde_json::Value;
|
|
use std::sync::Arc;
|
|
use tokio::sync::mpsc;
|
|
use tokio_util::sync::CancellationToken;
|
|
use tracing::error;
|
|
|
|
pub static EXTENSION_NAME: &str = "Extension Manager";
|
|
// pub static DISPLAY_NAME: &str = "Extension Manager";
|
|
|
|
#[derive(Debug, thiserror::Error)]
|
|
pub enum ExtensionManagerToolError {
|
|
#[error("Unknown tool: {tool_name}")]
|
|
UnknownTool { tool_name: String },
|
|
|
|
#[error("Extension manager not available")]
|
|
ManagerUnavailable,
|
|
|
|
#[error("Missing required parameter: {param_name}")]
|
|
MissingParameter { param_name: String },
|
|
|
|
#[error("Invalid action: {action}. Must be 'enable' or 'disable'")]
|
|
InvalidAction { action: String },
|
|
|
|
#[error("Extension operation failed: {message}")]
|
|
OperationFailed { message: String },
|
|
|
|
#[error("Failed to deserialize parameters: {0}")]
|
|
DeserializationError(#[from] serde_json::Error),
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
|
|
#[serde(rename_all = "lowercase")]
|
|
pub enum ManageExtensionAction {
|
|
Enable,
|
|
Disable,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
|
|
pub struct ManageExtensionsParams {
|
|
pub action: ManageExtensionAction,
|
|
pub extension_name: String,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
|
|
pub struct ReadResourceParams {
|
|
pub uri: String,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub extension_name: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
|
|
pub struct ListResourcesParams {
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub extension_name: Option<String>,
|
|
}
|
|
|
|
pub const READ_RESOURCE_TOOL_NAME: &str = "read_resource";
|
|
pub const LIST_RESOURCES_TOOL_NAME: &str = "list_resources";
|
|
pub const SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME: &str = "search_available_extensions";
|
|
pub const MANAGE_EXTENSIONS_TOOL_NAME: &str = "manage_extensions";
|
|
pub const MANAGE_EXTENSIONS_TOOL_NAME_COMPLETE: &str = "extensionmanager__manage_extensions";
|
|
|
|
pub struct ExtensionManagerClient {
|
|
info: InitializeResult,
|
|
#[allow(dead_code)]
|
|
context: PlatformExtensionContext,
|
|
}
|
|
|
|
impl ExtensionManagerClient {
|
|
pub fn new(context: PlatformExtensionContext) -> Result<Self> {
|
|
let info = InitializeResult {
|
|
protocol_version: ProtocolVersion::V_2025_03_26,
|
|
capabilities: ServerCapabilities {
|
|
tasks: None,
|
|
tools: Some(ToolsCapability {
|
|
list_changed: Some(false),
|
|
}),
|
|
resources: None,
|
|
prompts: None,
|
|
completions: None,
|
|
experimental: None,
|
|
logging: None,
|
|
},
|
|
server_info: Implementation {
|
|
name: EXTENSION_NAME.to_string(),
|
|
title: Some(EXTENSION_NAME.to_string()),
|
|
version: "1.0.0".to_string(),
|
|
icons: None,
|
|
website_url: None,
|
|
},
|
|
instructions: Some(indoc! {r#"
|
|
Extension Management
|
|
|
|
Use these tools to discover, enable, and disable extensions, as well as review resources.
|
|
|
|
Available tools:
|
|
- search_available_extensions: Find extensions available to enable/disable
|
|
- manage_extensions: Enable or disable extensions
|
|
- list_resources: List resources from extensions
|
|
- read_resource: Read specific resources from extensions
|
|
|
|
When you lack the tools needed to complete a task, use search_available_extensions first
|
|
to discover what extensions can help.
|
|
|
|
Use manage_extensions to enable or disable specific extensions by name.
|
|
Use list_resources and read_resource to work with extension data and resources.
|
|
"#}.to_string()),
|
|
};
|
|
|
|
Ok(Self { info, context })
|
|
}
|
|
|
|
async fn handle_search_available_extensions(
|
|
&self,
|
|
) -> Result<Vec<Content>, ExtensionManagerToolError> {
|
|
if let Some(weak_ref) = &self.context.extension_manager {
|
|
if let Some(extension_manager) = weak_ref.upgrade() {
|
|
match extension_manager.search_available_extensions().await {
|
|
Ok(content) => Ok(content),
|
|
Err(e) => Err(ExtensionManagerToolError::OperationFailed {
|
|
message: format!("Failed to search available extensions: {}", e.message),
|
|
}),
|
|
}
|
|
} else {
|
|
Err(ExtensionManagerToolError::ManagerUnavailable)
|
|
}
|
|
} else {
|
|
Err(ExtensionManagerToolError::ManagerUnavailable)
|
|
}
|
|
}
|
|
|
|
async fn handle_manage_extensions(
|
|
&self,
|
|
arguments: Option<JsonObject>,
|
|
) -> Result<Vec<Content>, ExtensionManagerToolError> {
|
|
let arguments = arguments.ok_or(ExtensionManagerToolError::MissingParameter {
|
|
param_name: "arguments".to_string(),
|
|
})?;
|
|
|
|
let params: ManageExtensionsParams =
|
|
serde_json::from_value(serde_json::Value::Object(arguments))?;
|
|
|
|
match self
|
|
.manage_extensions_impl(params.action, params.extension_name)
|
|
.await
|
|
{
|
|
Ok(content) => Ok(content),
|
|
Err(error_data) => Err(ExtensionManagerToolError::OperationFailed {
|
|
message: error_data.message.to_string(),
|
|
}),
|
|
}
|
|
}
|
|
|
|
async fn manage_extensions_impl(
|
|
&self,
|
|
action: ManageExtensionAction,
|
|
extension_name: String,
|
|
) -> Result<Vec<Content>, ErrorData> {
|
|
let extension_manager = self
|
|
.context
|
|
.extension_manager
|
|
.as_ref()
|
|
.and_then(|weak| weak.upgrade())
|
|
.ok_or_else(|| {
|
|
ErrorData::new(
|
|
ErrorCode::INTERNAL_ERROR,
|
|
"Extension manager is no longer available".to_string(),
|
|
None,
|
|
)
|
|
})?;
|
|
|
|
if action == ManageExtensionAction::Disable {
|
|
return 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));
|
|
}
|
|
|
|
let config = match get_extension_by_name(&extension_name) {
|
|
Some(config) => config,
|
|
None => {
|
|
return Err(ErrorData::new(
|
|
ErrorCode::RESOURCE_NOT_FOUND,
|
|
format!(
|
|
"Extension '{}' not found. Please check the extension name and try again.",
|
|
extension_name
|
|
),
|
|
None,
|
|
));
|
|
}
|
|
};
|
|
|
|
extension_manager
|
|
.add_extension_with_working_dir(config, None)
|
|
.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))
|
|
}
|
|
|
|
async fn handle_list_resources(
|
|
&self,
|
|
session_id: &str,
|
|
arguments: Option<JsonObject>,
|
|
) -> Result<Vec<Content>, ExtensionManagerToolError> {
|
|
if let Some(weak_ref) = &self.context.extension_manager {
|
|
if let Some(extension_manager) = weak_ref.upgrade() {
|
|
let params = arguments
|
|
.map(serde_json::Value::Object)
|
|
.unwrap_or(serde_json::Value::Object(serde_json::Map::new()));
|
|
|
|
match extension_manager
|
|
.list_resources(
|
|
session_id,
|
|
params,
|
|
tokio_util::sync::CancellationToken::default(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(content) => Ok(content),
|
|
Err(e) => Err(ExtensionManagerToolError::OperationFailed {
|
|
message: format!("Failed to list resources: {}", e.message),
|
|
}),
|
|
}
|
|
} else {
|
|
Err(ExtensionManagerToolError::ManagerUnavailable)
|
|
}
|
|
} else {
|
|
Err(ExtensionManagerToolError::ManagerUnavailable)
|
|
}
|
|
}
|
|
|
|
async fn handle_read_resource(
|
|
&self,
|
|
session_id: &str,
|
|
arguments: Option<JsonObject>,
|
|
) -> Result<Vec<Content>, ExtensionManagerToolError> {
|
|
if let Some(weak_ref) = &self.context.extension_manager {
|
|
if let Some(extension_manager) = weak_ref.upgrade() {
|
|
let params = arguments
|
|
.map(serde_json::Value::Object)
|
|
.unwrap_or(serde_json::Value::Object(serde_json::Map::new()));
|
|
|
|
match extension_manager
|
|
.read_resource_tool(
|
|
session_id,
|
|
params,
|
|
tokio_util::sync::CancellationToken::default(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(content) => Ok(content),
|
|
Err(e) => Err(ExtensionManagerToolError::OperationFailed {
|
|
message: format!("Failed to read resource: {}", e.message),
|
|
}),
|
|
}
|
|
} else {
|
|
Err(ExtensionManagerToolError::ManagerUnavailable)
|
|
}
|
|
} else {
|
|
Err(ExtensionManagerToolError::ManagerUnavailable)
|
|
}
|
|
}
|
|
|
|
#[allow(clippy::too_many_lines)]
|
|
async fn get_tools(&self) -> Vec<Tool> {
|
|
let mut tools = vec![
|
|
Tool::new(
|
|
SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME.to_string(),
|
|
"Searches for additional extensions available to help complete tasks.
|
|
Use this tool when you're unable to find a specific feature or functionality you need to complete your task, or when standard approaches aren't working.
|
|
These extensions might provide the exact tools needed to solve your problem.
|
|
If you find a relevant one, consider using your tools to enable it.".to_string(),
|
|
Arc::new(
|
|
serde_json::json!({
|
|
"type": "object",
|
|
"required": [],
|
|
"properties": {}
|
|
})
|
|
.as_object()
|
|
.expect("Schema must be an object")
|
|
.clone()
|
|
),
|
|
).annotate(ToolAnnotations {
|
|
title: Some("Discover extensions".to_string()),
|
|
read_only_hint: Some(true),
|
|
destructive_hint: Some(false),
|
|
idempotent_hint: Some(false),
|
|
open_world_hint: Some(false),
|
|
}),
|
|
Tool::new(
|
|
MANAGE_EXTENSIONS_TOOL_NAME.to_string(),
|
|
"Tool to manage extensions and tools in goose context.
|
|
Enable or disable extensions to help complete tasks.
|
|
Enable or disable an extension by providing the extension name.
|
|
".to_string(),
|
|
Arc::new(
|
|
serde_json::to_value(schema_for!(ManageExtensionsParams))
|
|
.expect("Failed to serialize schema")
|
|
.as_object()
|
|
.expect("Schema must be an object")
|
|
.clone()
|
|
),
|
|
).annotate(ToolAnnotations {
|
|
title: Some("Enable or disable an extension".to_string()),
|
|
read_only_hint: Some(false),
|
|
destructive_hint: Some(false),
|
|
idempotent_hint: Some(false),
|
|
open_world_hint: Some(false),
|
|
}),
|
|
];
|
|
|
|
// Only add resource tools if extension manager supports resources
|
|
if let Some(weak_ref) = &self.context.extension_manager {
|
|
if let Some(extension_manager) = weak_ref.upgrade() {
|
|
if extension_manager.supports_resources().await {
|
|
tools.extend([
|
|
Tool::new(
|
|
LIST_RESOURCES_TOOL_NAME.to_string(),
|
|
indoc! {r#"
|
|
List resources from an extension(s).
|
|
|
|
Resources allow extensions to share data that provide context to LLMs, such as
|
|
files, database schemas, or application-specific information. This tool lists resources
|
|
in the provided extension, and returns a list for the user to browse. If no extension
|
|
is provided, the tool will search all extensions for the resource.
|
|
"#}.to_string(),
|
|
Arc::new(
|
|
serde_json::to_value(schema_for!(ListResourcesParams))
|
|
.expect("Failed to serialize schema")
|
|
.as_object()
|
|
.expect("Schema must be an object")
|
|
.clone()
|
|
),
|
|
).annotate(ToolAnnotations {
|
|
title: Some("List resources".to_string()),
|
|
read_only_hint: Some(true),
|
|
destructive_hint: Some(false),
|
|
idempotent_hint: Some(false),
|
|
open_world_hint: Some(false),
|
|
}),
|
|
Tool::new(
|
|
READ_RESOURCE_TOOL_NAME.to_string(),
|
|
indoc! {r#"
|
|
Read a resource from an extension.
|
|
|
|
Resources allow extensions to share data that provide context to LLMs, such as
|
|
files, database schemas, or application-specific information. This tool searches for the
|
|
resource URI in the provided extension, and reads in the resource content. If no extension
|
|
is provided, the tool will search all extensions for the resource.
|
|
"#}.to_string(),
|
|
Arc::new(
|
|
serde_json::to_value(schema_for!(ReadResourceParams))
|
|
.expect("Failed to serialize schema")
|
|
.as_object()
|
|
.expect("Schema must be an object")
|
|
.clone()
|
|
),
|
|
).annotate(ToolAnnotations {
|
|
title: Some("Read a resource".to_string()),
|
|
read_only_hint: Some(true),
|
|
destructive_hint: Some(false),
|
|
idempotent_hint: Some(false),
|
|
open_world_hint: Some(false),
|
|
}),
|
|
]);
|
|
}
|
|
}
|
|
}
|
|
|
|
tools
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl McpClientTrait for ExtensionManagerClient {
|
|
async fn list_resources(
|
|
&self,
|
|
_session_id: &str,
|
|
_next_cursor: Option<String>,
|
|
_cancellation_token: CancellationToken,
|
|
) -> Result<ListResourcesResult, Error> {
|
|
Err(Error::TransportClosed)
|
|
}
|
|
|
|
async fn read_resource(
|
|
&self,
|
|
_session_id: &str,
|
|
_uri: &str,
|
|
_cancellation_token: CancellationToken,
|
|
) -> Result<ReadResourceResult, Error> {
|
|
// Extension manager doesn't expose resources directly
|
|
Err(Error::TransportClosed)
|
|
}
|
|
|
|
async fn list_tools(
|
|
&self,
|
|
_session_id: &str,
|
|
_next_cursor: Option<String>,
|
|
_cancellation_token: CancellationToken,
|
|
) -> Result<ListToolsResult, Error> {
|
|
Ok(ListToolsResult {
|
|
tools: self.get_tools().await,
|
|
next_cursor: None,
|
|
meta: None,
|
|
})
|
|
}
|
|
|
|
async fn call_tool(
|
|
&self,
|
|
session_id: &str,
|
|
name: &str,
|
|
arguments: Option<JsonObject>,
|
|
_cancellation_token: CancellationToken,
|
|
) -> Result<CallToolResult, Error> {
|
|
let result = match name {
|
|
SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME => {
|
|
self.handle_search_available_extensions().await
|
|
}
|
|
MANAGE_EXTENSIONS_TOOL_NAME => self.handle_manage_extensions(arguments).await,
|
|
LIST_RESOURCES_TOOL_NAME => self.handle_list_resources(session_id, arguments).await,
|
|
READ_RESOURCE_TOOL_NAME => self.handle_read_resource(session_id, arguments).await,
|
|
_ => Err(ExtensionManagerToolError::UnknownTool {
|
|
tool_name: name.to_string(),
|
|
}),
|
|
};
|
|
|
|
match result {
|
|
Ok(content) => Ok(CallToolResult::success(content)),
|
|
Err(error) => {
|
|
// Log the error for debugging
|
|
error!("Extension manager tool '{}' failed: {}", name, error);
|
|
|
|
// Return proper error result with is_error flag set
|
|
Ok(CallToolResult {
|
|
content: vec![Content::text(error.to_string())],
|
|
is_error: Some(true), // ✅ Properly mark as error
|
|
structured_content: None,
|
|
meta: None,
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn list_prompts(
|
|
&self,
|
|
_session_id: &str,
|
|
_next_cursor: Option<String>,
|
|
_cancellation_token: CancellationToken,
|
|
) -> Result<ListPromptsResult, Error> {
|
|
Err(Error::TransportClosed)
|
|
}
|
|
|
|
async fn get_prompt(
|
|
&self,
|
|
_session_id: &str,
|
|
_name: &str,
|
|
_arguments: Value,
|
|
_cancellation_token: CancellationToken,
|
|
) -> Result<GetPromptResult, Error> {
|
|
Err(Error::TransportClosed)
|
|
}
|
|
|
|
async fn subscribe(&self) -> mpsc::Receiver<ServerNotification> {
|
|
mpsc::channel(1).1
|
|
}
|
|
|
|
fn get_info(&self) -> Option<&InitializeResult> {
|
|
Some(&self.info)
|
|
}
|
|
}
|