zclaw_openfang/crates/zclaw-protocols/src/mcp_tool_adapter.rs

//! MCP Tool Adapter — bridges MCP server tools into zclaw-runtime's ToolRegistry
//!
//! Each MCP tool is wrapped as a `dyn Tool` implementation that:
//! 1. Receives the tool call from the agent loop
//! 2. Forwards it to the MCP server via McpClient
//! 3. Returns the result back to the agent loop

use std::collections::HashMap;
use std::sync::Arc;
use serde_json::Value;
use tracing::{debug, warn};
use zclaw_types::{Result, ZclawError};

use crate::mcp::{McpClient, McpTool, McpToolCallRequest};

/// Adapter wrapping an MCP tool as a zclaw-runtime `Tool`.
///
/// This struct is intentionally decoupled from `zclaw-runtime::Tool` to avoid
/// a circular dependency. The runtime crate depends on protocols for type sharing,
/// so we expose a simple trait here that mirrors the essential Tool interface.
/// The runtime side will wrap this in a thin `Tool` impl.
pub struct McpToolAdapter {
    /// Service name this tool belongs to
    service_name: String,
    /// Tool name (original from MCP server, NOT prefixed)
    name: String,
    /// Tool description
    description: String,
    /// JSON schema for input parameters
    input_schema: Value,
    /// Reference to the MCP client for forwarding calls
    client: Arc<dyn McpClient>,
}

impl Clone for McpToolAdapter {
    fn clone(&self) -> Self {
        Self {
            service_name: self.service_name.clone(),
            name: self.name.clone(),
            description: self.description.clone(),
            input_schema: self.input_schema.clone(),
            client: self.client.clone(),
        }
    }
}

impl McpToolAdapter {
    pub fn new(service_name: String, tool: McpTool, client: Arc<dyn McpClient>) -> Self {
        Self {
            service_name,
            name: tool.name,
            description: tool.description,
            input_schema: tool.input_schema,
            client,
        }
    }

    /// Create adapters for all tools from an MCP server
    pub async fn from_server(service_name: String, client: Arc<dyn McpClient>) -> Result<Vec<Self>> {
        let tools = client.list_tools().await?;
        debug!(count = tools.len(), "Discovered MCP tools");
        Ok(tools.into_iter().map(|t| Self::new(service_name.clone(), t, client.clone())).collect())
    }

    pub fn name(&self) -> &str {
        &self.name
    }

    /// Full qualified name: service_name.tool_name (for ToolRegistry to avoid collisions)
    pub fn qualified_name(&self) -> String {
        format!("{}.{}", self.service_name, self.name)
    }

    pub fn service_name(&self) -> &str {
        &self.service_name
    }

    pub fn tool_name(&self) -> &str {
        &self.name
    }

    pub fn description(&self) -> &str {
        &self.description
    }

    pub fn input_schema(&self) -> &Value {
        &self.input_schema
    }

    /// Execute the MCP tool call
    pub async fn execute(&self, input: Value) -> Result<Value> {
        debug!(tool = %self.name, "Executing MCP tool");

        let arguments = match input {
            Value::Object(map) => map.into_iter().collect(),
            other => {
                // If input is not an object, wrap it as {"input": ...}
                HashMap::from([("input".to_string(), other)])
            }
        };

        let request = McpToolCallRequest {
            name: self.name.clone(),
            arguments,
        };

        let response = self.client.call_tool(request).await?;

        if response.is_error {
            // Extract error text from content blocks
            let error_text: String = response.content.iter()
                .filter_map(|c| match c {
                    crate::mcp::McpContent::Text { text } => Some(text.as_str()),
                    _ => None,
                })
                .collect::<Vec<_>>()
                .join("\n");

            warn!(tool = %self.name, error = %error_text, "MCP tool returned error");
            return Err(ZclawError::McpError(format!("MCP tool '{}' failed: {}", self.name, error_text)));
        }

        // Convert content blocks to JSON
        let result = response.content.into_iter()
            .filter_map(|c| match c {
                crate::mcp::McpContent::Text { text } => Some(Value::String(text)),
                _ => None,
            })
            .collect::<Vec<_>>();

        match result.len() {
            0 => Ok(Value::Null),
            1 => Ok(result.into_iter().next().unwrap_or(Value::Null)),
            _ => Ok(Value::Array(result)),
        }
    }
}

/// MCP Service Manager — manages lifecycle of MCP server connections
#[derive(Default)]
pub struct McpServiceManager {
    clients: HashMap<String, Arc<dyn McpClient>>,
    adapters: HashMap<String, Vec<McpToolAdapter>>,
}

impl McpServiceManager {
    pub fn new() -> Self {
        Self {
            clients: HashMap::new(),
            adapters: HashMap::new(),
        }
    }

    /// Register a connected MCP client and discover its tools
    pub async fn register_service(
        &mut self,
        name: String,
        client: Arc<dyn McpClient>,
    ) -> Result<Vec<&McpToolAdapter>> {
        let adapters = McpToolAdapter::from_server(name.clone(), client.clone()).await?;
        self.clients.insert(name.clone(), client);
        self.adapters.insert(name.clone(), adapters);
        Ok(self.adapters.get(&name).map(|v| v.iter().collect()).unwrap_or_default())
    }

    /// Get all registered tool adapters from all services
    pub fn all_adapters(&self) -> Vec<&McpToolAdapter> {
        self.adapters.values().flat_map(|v| v.iter()).collect()
    }

    /// Remove a service by name
    pub fn remove_service(&mut self, name: &str) {
        self.clients.remove(name);
        self.adapters.remove(name);
    }

    /// List registered service names
    pub fn service_names(&self) -> Vec<&String> {
        self.clients.keys().collect()
    }
}