enact-mcp 0.0.2

//! MCP (Model Context Protocol) Client for Enact
//!
//! This module provides a client for the Model Context Protocol,
//! allowing Enact to connect to MCP servers and use their tools.

pub mod adapter;
pub mod config;

use anyhow::Result;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::process::Stdio;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, Command};
use tracing::debug;

pub use adapter::{discover_mcp_tools, McpToolAdapter};

/// MCP Tool definition
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpTool {
    pub name: String,
    pub description: String,
    pub parameters: serde_json::Value,
}

/// MCP Client for stdio transport
pub struct McpStdioClient {
    process: Child,
    stdin: tokio::process::ChildStdin,
    stdout: BufReader<tokio::process::ChildStdout>,
}

/// MCP Client for HTTP transport
pub struct McpHttpClient {
    client: Client,
    url: String,
}

impl McpHttpClient {
    pub fn new(url: impl Into<String>) -> Self {
        Self {
            client: Client::new(),
            url: url.into(),
        }
    }

    async fn send_request(&self, request: &serde_json::Value) -> Result<serde_json::Value> {
        let response = self
            .client
            .post(&self.url)
            .json(request)
            .send()
            .await?
            .error_for_status()?
            .json::<serde_json::Value>()
            .await?;
        Ok(response)
    }

    pub async fn list_tools(&self) -> Result<Vec<McpTool>> {
        let request = serde_json::json!({
            "jsonrpc": "2.0",
            "id": 2,
            "method": "tools/list"
        });
        let response = self.send_request(&request).await?;
        let tools = response
            .get("result")
            .and_then(|r| r.get("tools"))
            .and_then(|t| t.as_array())
            .map(|arr| {
                arr.iter()
                    .filter_map(|tool| {
                        Some(McpTool {
                            name: tool.get("name")?.as_str()?.to_string(),
                            description: tool.get("description")?.as_str()?.to_string(),
                            parameters: tool.get("parameters")?.clone(),
                        })
                    })
                    .collect()
            })
            .unwrap_or_default();
        Ok(tools)
    }

    pub async fn call_tool(&self, name: &str, arguments: serde_json::Value) -> Result<String> {
        let request = serde_json::json!({
            "jsonrpc": "2.0",
            "id": 3,
            "method": "tools/call",
            "params": { "name": name, "arguments": arguments }
        });
        let response = self.send_request(&request).await?;
        if let Some(error) = response.get("error") {
            anyhow::bail!("MCP tool error: {:?}", error);
        }
        let content = response
            .get("result")
            .and_then(|r| r.get("content"))
            .and_then(|c| c.as_array())
            .and_then(|arr| arr.first())
            .and_then(|item| item.get("text"))
            .and_then(|t| t.as_str())
            .unwrap_or("No content returned");
        Ok(content.to_string())
    }
}

impl McpStdioClient {
    /// Create a new MCP client connected to a command via stdio.
    /// `envs` are injected into the child process environment in addition to the inherited env.
    pub async fn new(
        command: &str,
        args: &[&str],
        envs: &std::collections::HashMap<String, String>,
    ) -> Result<Self> {
        let mut process = Command::new(command)
            .args(args)
            .envs(envs)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()?;

        let stdin = process.stdin.take().unwrap();
        let stdout = BufReader::new(process.stdout.take().unwrap());

        let mut client = Self {
            process,
            stdin,
            stdout,
        };

        // Initialize connection
        client.initialize().await?;

        Ok(client)
    }

    async fn initialize(&mut self) -> Result<()> {
        let init_request = serde_json::json!({
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {
                "protocolVersion": "2024-11-05",
                "capabilities": {},
                "clientInfo": {
                    "name": "enact-mcp",
                    "version": "0.1.0"
                }
            }
        });

        self.send_request(&init_request).await?;
        let response = self.read_response().await?;
        debug!("MCP initialized: {:?}", response);

        Ok(())
    }

    async fn send_request(&mut self, request: &serde_json::Value) -> Result<()> {
        let request_str = request.to_string();
        debug!("Sending MCP request: {}", request_str);

        self.stdin.write_all(request_str.as_bytes()).await?;
        self.stdin.write_all(b"\n").await?;
        self.stdin.flush().await?;

        Ok(())
    }

    async fn read_response(&mut self) -> Result<serde_json::Value> {
        let mut line = String::new();
        self.stdout.read_line(&mut line).await?;

        debug!("Received MCP response: {}", line);
        let response: serde_json::Value = serde_json::from_str(&line)?;
        Ok(response)
    }

    /// List available tools from the MCP server
    pub async fn list_tools(&mut self) -> Result<Vec<McpTool>> {
        let request = serde_json::json!({
            "jsonrpc": "2.0",
            "id": 2,
            "method": "tools/list"
        });

        self.send_request(&request).await?;
        let response = self.read_response().await?;

        let tools = response
            .get("result")
            .and_then(|r| r.get("tools"))
            .and_then(|t| t.as_array())
            .map(|arr| {
                arr.iter()
                    .filter_map(|tool| {
                        Some(McpTool {
                            name: tool.get("name")?.as_str()?.to_string(),
                            description: tool.get("description")?.as_str()?.to_string(),
                            parameters: tool.get("parameters")?.clone(),
                        })
                    })
                    .collect()
            })
            .unwrap_or_default();

        Ok(tools)
    }

    /// Call a tool on the MCP server
    pub async fn call_tool(&mut self, name: &str, arguments: serde_json::Value) -> Result<String> {
        let request = serde_json::json!({
            "jsonrpc": "2.0",
            "id": 3,
            "method": "tools/call",
            "params": {
                "name": name,
                "arguments": arguments
            }
        });

        self.send_request(&request).await?;
        let response = self.read_response().await?;

        if let Some(error) = response.get("error") {
            anyhow::bail!("MCP tool error: {:?}", error);
        }

        let content = response
            .get("result")
            .and_then(|r| r.get("content"))
            .and_then(|c| c.as_array())
            .and_then(|arr| arr.first())
            .and_then(|item| item.get("text"))
            .and_then(|t| t.as_str())
            .unwrap_or("No content returned");

        Ok(content.to_string())
    }
}

impl Drop for McpStdioClient {
    fn drop(&mut self) {
        let _ = self.process.start_kill();
    }
}

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

    #[test]
    fn test_mcp_tool_creation() {
        let tool = McpTool {
            name: "test".to_string(),
            description: "Test tool".to_string(),
            parameters: serde_json::json!({}),
        };
        assert_eq!(tool.name, "test");
    }
}