PicoBot/src/mcp/tool_adapter.rs

//! MCP Tool Adapter - wraps MCP tools as PicoBot tools

use async_trait::async_trait;
use std::sync::Arc;

use rmcp::model::Tool;

use crate::mcp::client::McpClientManager;
use crate::tools::traits::{Tool as PicoBotTool, ToolResult};

/// Wrapper that adapts an MCP tool to PicoBot's Tool trait
#[derive(Clone)]
pub struct McpToolWrapper {
    /// The MCP client manager
    manager: Arc<McpClientManager>,
    /// The server key this tool belongs to (from mcpServers map key)
    server_key: String,
    /// The original tool name on the MCP server
    tool_name: String,
    /// The full tool name with namespace (mcp_{key}_{tool})
    full_name: String,
    /// Tool information from MCP server
    tool_info: Tool,
}

impl McpToolWrapper {
    /// Create a new tool wrapper
    pub fn new(
        manager: Arc<McpClientManager>,
        server_key: String,
        tool_info: Tool,
    ) -> Self {
        let tool_name = tool_info.name.clone().into_owned();
        let full_name = format!("mcp_{}_{}", server_key, tool_name);
        Self {
            manager,
            server_key,
            tool_name,
            full_name,
            tool_info,
        }
    }

    /// Get the server key
    pub fn server_key(&self) -> &str {
        &self.server_key
    }

    /// Get the original tool name
    pub fn original_name(&self) -> &str {
        &self.tool_name
    }
}

#[async_trait]
impl PicoBotTool for McpToolWrapper {
    fn name(&self) -> &str {
        &self.full_name
    }

    fn description(&self) -> &str {
        self.tool_info.description.as_deref().unwrap_or("MCP tool")
    }

    fn parameters_schema(&self) -> serde_json::Value {
        // Convert Arc<JsonObject> to serde_json::Value
        let schema = (*self.tool_info.input_schema).clone();
        serde_json::Value::Object(schema)
    }

    async fn execute(&self, args: serde_json::Value) -> anyhow::Result<ToolResult> {
        tracing::debug!(
            server_key = %self.server_key,
            tool = %self.tool_name,
            "Calling MCP tool"
        );

        let result = self
            .manager
            .call_tool(&self.server_key, &self.tool_name, args)
            .await?;

        // Convert MCP CallToolResult to PicoBot ToolResult
        let output = extract_text_content(&result);
        let is_error = result.is_error.unwrap_or(false);

        Ok(ToolResult {
            success: !is_error,
            output,
            error: if is_error {
                Some("MCP tool returned error".to_string())
            } else {
                None
            },
        })
    }

    fn read_only(&self) -> bool {
        // MCP tools may or may not be read-only; we default to false
        // This could be enhanced if MCP servers provide this info via annotations
        false
    }
}

/// Extract text content from MCP CallToolResult
fn extract_text_content(result: &rmcp::model::CallToolResult) -> String {
    let mut text_parts = Vec::new();

    for content in &result.content {
        if let Some(text) = content.as_text() {
            text_parts.push(text.text.clone());
        }
    }

    if text_parts.is_empty() {
        // No text content found, try to serialize the whole result
        serde_json::to_string_pretty(&result).unwrap_or_else(|_| "Empty result".to_string())
    } else {
        text_parts.join("\n")
    }
}

/// Register all MCP tools from connected servers into a tool registry
pub async fn register_mcp_tools(
    manager: Arc<McpClientManager>,
    registry: &mut crate::tools::registry::ToolRegistry,
) -> anyhow::Result<()> {
    let all_tools = manager.all_tools().await;

    for (server_key, tool_info) in all_tools {
        let wrapper = McpToolWrapper::new(
            manager.clone(),
            server_key.clone(),
            tool_info,
        );

        tracing::info!(
            name = %wrapper.name(),
            server_key = %server_key,
            "Registering MCP tool"
        );

        registry.register(wrapper);
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use rmcp::model::{CallToolResult, Content};

    #[test]
    fn test_extract_text_content_from_text() {
        let result = CallToolResult::success(vec![
            Content::text("Hello"),
            Content::text("World"),
        ]);

        let text = extract_text_content(&result);
        assert_eq!(text, "Hello\nWorld");
    }

    #[test]
    fn test_extract_text_content_empty() {
        let result = CallToolResult::success(vec![]);
        let text = extract_text_content(&result);
        // When content is empty, the function serializes the result to JSON
        // which contains an empty content array
        assert!(text.contains("content") || text.contains("Empty result"));
    }

    #[test]
    fn test_mcp_tool_wrapper_name() {
        let manager = Arc::new(McpClientManager::new());
        // Create a minimal tool info using rmcp's Tool constructor
        let schema: serde_json::Map<String, serde_json::Value> = serde_json::json!({"type": "object"})
            .as_object()
            .unwrap()
            .clone();
        let tool_info = Tool::new("echo", "Echo tool", schema);

        let wrapper = McpToolWrapper::new(manager, "filesystem".to_string(), tool_info);
        assert_eq!(wrapper.name(), "mcp_filesystem_echo");
        assert_eq!(wrapper.original_name(), "echo");
        assert_eq!(wrapper.server_key(), "filesystem");
    }
}