llm-worker-rs/worker/src/mcp/tool.rs

use super::protocol::{CallToolResult, McpClient, McpToolDefinition};
use crate::types::{Tool, ToolResult};
use async_trait::async_trait;
use serde_json::Value;
use std::sync::Arc;
use tokio::sync::{Mutex, RwLock};
use tracing::{debug, error, info, warn};

/// Convert MCP CallToolResult to JSON format
fn convert_mcp_result_to_json(result: &CallToolResult) -> Value {
    if result.content.is_empty() {
        serde_json::json!({
            "success": true,
            "content": []
        })
    } else {
        let content_json: Vec<Value> = result
            .content
            .iter()
            .map(|content| {
                serde_json::json!({
                    "type": content.content_type,
                    "text": content.text,
                    "data": content.data
                })
            })
            .collect();

        serde_json::json!({
            "success": true,
            "content": content_json
        })
    }
}

/// MCPサーバー設定
#[derive(Debug, Clone)]
pub struct McpServerConfig {
    pub command: String,
    pub args: Vec<String>,
    pub name: String,
}

impl McpServerConfig {
    pub fn new(command: impl Into<String>, args: Vec<impl Into<String>>) -> Self {
        let command = command.into();
        let args: Vec<String> = args.into_iter().map(|s| s.into()).collect();
        let name = format!("{}({})", command, args.join(" "));

        Self {
            command,
            args,
            name,
        }
    }
}

/// 実際に動作するMCP統合ツール
pub struct McpDynamicTool {
    config: McpServerConfig,
    client: Arc<Mutex<Option<McpClient>>>,
    tools_cache: Arc<RwLock<Vec<McpToolDefinition>>>,
}

/// 単一のMCPツールを表すDynamicTool
pub struct SingleMcpTool {
    tool_name: String,
    tool_description: String,
    tool_schema: Value,
    client: Arc<Mutex<Option<McpClient>>>,
}

impl McpDynamicTool {
    /// 新しいMCPツールを作成
    pub fn new(config: McpServerConfig) -> Self {
        Self {
            config,
            client: Arc::new(Mutex::new(None)),
            tools_cache: Arc::new(RwLock::new(Vec::new())),
        }
    }

    /// MCPサーバーに接続
    async fn ensure_connected(&self) -> ToolResult<()> {
        let mut client_guard = self.client.lock().await;

        if client_guard.is_none() {
            info!("Connecting to MCP server: {}", self.config.name);

            let mut mcp_client = McpClient::new();
            mcp_client
                .connect(self.config.command.clone(), self.config.args.clone())
                .await
                .map_err(|e| {
                    crate::WorkerError::tool_execution_with_source(
                        &self.config.name,
                        format!("Failed to connect to MCP server '{}'", self.config.name),
                        e,
                    )
                })?;

            *client_guard = Some(mcp_client);
            info!("Successfully connected to MCP server: {}", self.config.name);
        }

        Ok(())
    }

    /// 利用可能なツール一覧を取得
    async fn fetch_tools(&self) -> ToolResult<Vec<McpToolDefinition>> {
        self.ensure_connected().await?;

        let mut client_guard = self.client.lock().await;
        let client = client_guard.as_mut().ok_or_else(|| {
            crate::WorkerError::tool_execution(&self.config.name, "MCP client not connected")
        })?;

        let tools = client.list_tools().await.map_err(|e| {
            crate::WorkerError::tool_execution_with_source(
                &self.config.name,
                format!(
                    "Failed to list tools from MCP server '{}'",
                    self.config.name
                ),
                e,
            )
        })?;

        debug!(
            "Retrieved {} tools from MCP server '{}'",
            tools.len(),
            self.config.name
        );
        Ok(tools)
    }

    /// ツールキャッシュを更新
    async fn update_tools_cache(&self) -> ToolResult<()> {
        let tools = self.fetch_tools().await?;
        let mut cache_guard = self.tools_cache.write().await;
        *cache_guard = tools;
        Ok(())
    }

    /// 特定のツールを名前で検索
    async fn find_tool_by_name(&self, tool_name: &str) -> ToolResult<Option<McpToolDefinition>> {
        let cache_guard = self.tools_cache.read().await;

        // キャッシュが空の場合は更新
        if cache_guard.is_empty() {
            drop(cache_guard);
            self.update_tools_cache().await?;
            let cache_guard = self.tools_cache.read().await;
            let result = cache_guard
                .iter()
                .find(|tool| tool.name == tool_name)
                .cloned();
            Ok(result)
        } else {
            let result = cache_guard
                .iter()
                .find(|tool| tool.name == tool_name)
                .cloned();
            Ok(result)
        }
    }

    /// MCPサーバーのツールを実行
    async fn call_mcp_tool(&self, tool_name: &str, args: Value) -> ToolResult<Value> {
        self.ensure_connected().await?;

        let mut client_guard = self.client.lock().await;
        let client = client_guard.as_mut().ok_or_else(|| {
            crate::WorkerError::tool_execution(tool_name, "MCP client not connected")
        })?;

        debug!("Calling MCP tool '{}' with args: {}", tool_name, args);

        let result = client.call_tool(tool_name, Some(args)).await.map_err(|e| {
            crate::WorkerError::tool_execution_with_source(
                tool_name,
                format!("Failed to call MCP tool '{}'", tool_name),
                e,
            )
        })?;

        debug!("MCP tool '{}' returned: {:?}", tool_name, result);

        // Convert MCP result to JSON
        Ok(convert_mcp_result_to_json(&result))
    }
}

impl SingleMcpTool {
    /// 新しい単一MCPツールを作成
    pub fn new(
        tool_name: String,
        tool_description: String,
        tool_schema: Value,
        client: Arc<Mutex<Option<McpClient>>>,
    ) -> Self {
        Self {
            tool_name,
            tool_description,
            tool_schema,
            client,
        }
    }

    /// MCPサーバーのツールを実行
    async fn call_mcp_tool(&self, args: Value) -> ToolResult<Value> {
        let mut client_guard = self.client.lock().await;
        let client = client_guard.as_mut().ok_or_else(|| {
            crate::WorkerError::tool_execution(&self.tool_name, "MCP client not connected")
        })?;

        debug!("Calling MCP tool '{}' with args: {}", self.tool_name, args);

        let result = client
            .call_tool(&self.tool_name, Some(args))
            .await
            .map_err(|e| {
                crate::WorkerError::tool_execution_with_source(
                    &self.tool_name,
                    format!("Failed to call MCP tool '{}'", self.tool_name),
                    e,
                )
            })?;

        debug!("MCP tool '{}' returned: {:?}", self.tool_name, result);

        // Convert MCP result to JSON
        Ok(convert_mcp_result_to_json(&result))
    }
}

#[async_trait]
impl Tool for SingleMcpTool {
    fn name(&self) -> &str {
        &self.tool_name
    }

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

    fn parameters_schema(&self) -> Value {
        self.tool_schema.clone()
    }

    async fn execute(&self, args: Value) -> ToolResult<Value> {
        self.call_mcp_tool(args).await
    }
}

#[async_trait]
impl Tool for McpDynamicTool {
    fn name(&self) -> &str {
        "mcp_proxy"
    }

    fn description(&self) -> &str {
        "Execute tools from external MCP servers"
    }

    fn parameters_schema(&self) -> Value {
        serde_json::json!({
            "type": "object",
            "properties": {
                "tool_name": {
                    "type": "string",
                    "description": "Name of the MCP tool to execute"
                },
                "tool_args": {
                    "type": "object",
                    "description": "Arguments to pass to the MCP tool",
                    "additionalProperties": true
                }
            },
            "required": ["tool_name", "tool_args"]
        })
    }

    async fn execute(&self, args: Value) -> ToolResult<Value> {
        let tool_name = args
            .get("tool_name")
            .and_then(|v| v.as_str())
            .ok_or_else(|| {
                crate::WorkerError::tool_execution(
                    "mcp_proxy",
                    "Missing required parameter 'tool_name'",
                )
            })?;

        let tool_args = args
            .get("tool_args")
            .ok_or_else(|| {
                crate::WorkerError::tool_execution(
                    "mcp_proxy",
                    "Missing required parameter 'tool_args'",
                )
            })?
            .clone();

        // ツールが存在するか確認
        match self.find_tool_by_name(tool_name).await? {
            Some(_tool) => {
                // ツールを実行
                let result = self.call_mcp_tool(tool_name, tool_args).await?;
                Ok(serde_json::json!({
                    "success": true,
                    "tool_name": tool_name,
                    "result": result
                }))
            }
            None => {
                // ツールキャッシュを更新して再試行
                warn!("Tool '{}' not found in cache, refreshing...", tool_name);
                self.update_tools_cache().await?;

                match self.find_tool_by_name(tool_name).await? {
                    Some(_tool) => {
                        let result = self.call_mcp_tool(tool_name, tool_args).await?;
                        Ok(serde_json::json!({
                            "success": true,
                            "tool_name": tool_name,
                            "result": result
                        }))
                    }
                    None => Err(Box::new(crate::WorkerError::tool_execution(
                        tool_name,
                        format!(
                            "Tool '{}' not found in MCP server '{}'",
                            tool_name, self.config.name
                        ),
                    ))
                        as Box<dyn std::error::Error + Send + Sync>),
                }
            }
        }
    }
}

/// MCPサーバーから利用可能なツールをDynamicToolDefinitionとして取得
pub async fn get_mcp_tools_as_definitions(
    config: &McpServerConfig,
) -> ToolResult<Vec<crate::types::DynamicToolDefinition>> {
    let mcp_tool = McpDynamicTool::new(config.clone());
    let tools = mcp_tool.fetch_tools().await?;

    let mut definitions = Vec::new();

    for tool in tools {
        let definition = crate::types::DynamicToolDefinition {
            name: tool.name.clone(),
            description: tool
                .description
                .unwrap_or_else(|| format!("MCP tool: {}", tool.name)),
            parameters_schema: tool.input_schema,
        };
        definitions.push(definition);
    }

    info!(
        "Converted {} MCP tools to DynamicToolDefinitions",
        definitions.len()
    );
    Ok(definitions)
}

/// MCPサーバーから単一のツールを取得してSingleMcpToolを作成
pub async fn create_single_mcp_tools(config: &McpServerConfig) -> ToolResult<Vec<Box<dyn Tool>>> {
    let mcp_tool = McpDynamicTool::new(config.clone());
    let tools = mcp_tool.fetch_tools().await?;

    // 共有クライアントを作成
    let shared_client = mcp_tool.client.clone();

    let mut single_tools: Vec<Box<dyn Tool>> = Vec::new();

    for tool in tools {
        let tool_name = tool.name;
        let tool_description = tool
            .description
            .unwrap_or_else(|| format!("MCP tool: {}", tool_name));
        let tool_schema = tool.input_schema;

        let single_tool = SingleMcpTool::new(
            tool_name,
            tool_description,
            tool_schema,
            shared_client.clone(),
        );

        single_tools.push(Box::new(single_tool));
    }

    info!(
        "Created {} SingleMcpTools from MCP server '{}'",
        single_tools.len(),
        config.name
    );
    Ok(single_tools)
}

/// MCPサーバーとの接続をテストする
pub async fn test_mcp_connection(
    config: &McpServerConfig,
) -> Result<bool, Box<dyn std::error::Error + Send + Sync>> {
    info!("Testing MCP connection to server: {}", config.name);

    let mut client = McpClient::new();
    match client
        .connect(config.command.clone(), config.args.clone())
        .await
    {
        Ok(()) => {
            info!("Successfully connected to MCP server: {}", config.name);

            // Test listing tools
            match client.list_tools().await {
                Ok(tools) => {
                    info!(
                        "MCP server '{}' provides {} tools",
                        config.name,
                        tools.len()
                    );
                    for tool in &tools {
                        debug!(
                            "Available tool: {} - {}",
                            tool.name,
                            tool.description.as_deref().unwrap_or("No description")
                        );
                    }
                }
                Err(e) => {
                    warn!(
                        "Failed to list tools from MCP server '{}': {}",
                        config.name, e
                    );
                }
            }

            // Close connection
            let _ = client.close().await;
            Ok(true)
        }
        Err(e) => {
            error!("Failed to connect to MCP server '{}': {}", config.name, e);
            Ok(false)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_mcp_server_config() {
        let config =
            McpServerConfig::new("npx", vec!["-y", "@modelcontextprotocol/server-everything"]);
        assert_eq!(config.command, "npx");
        assert_eq!(
            config.args,
            vec!["-y", "@modelcontextprotocol/server-everything"]
        );
        assert_eq!(
            config.name,
            "npx(-y @modelcontextprotocol/server-everything)"
        );
    }

    #[tokio::test]
    async fn test_mcp_tool_creation() {
        let config = McpServerConfig::new("echo", vec!["test"]);
        let tool = McpDynamicTool::new(config);

        assert_eq!(tool.name(), "mcp_proxy");
        assert!(!tool.description().is_empty());

        let schema = tool.parameters_schema();
        assert!(schema.is_object());
        assert!(schema.get("properties").is_some());
    }
}