mur_core/mcp/

client.rs

//! MCP stdio client — connects to MCP servers via stdin/stdout.

use super::types::*;
use anyhow::{Context, Result};
use std::sync::atomic::{AtomicU64, Ordering};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, Command};
use tokio::sync::Mutex;
use tracing::{debug, info, warn};

/// MCP client that communicates with a server process via stdio.
pub struct McpClient {
    child: Mutex<Child>,
    stdin: Mutex<tokio::process::ChildStdin>,
    stdout: Mutex<BufReader<tokio::process::ChildStdout>>,
    next_id: AtomicU64,
    server_name: String,
    capabilities: Mutex<Option<McpServerCapabilities>>,
}

impl McpClient {
    /// Spawn an MCP server process and perform the initialize handshake.
    pub async fn connect(config: &McpConfig) -> Result<Self> {
        let mut cmd = Command::new(&config.command);
        cmd.args(&config.args)
            .stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped());

        for (k, v) in &config.env {
            cmd.env(k, v);
        }

        if let Some(cwd) = &config.cwd {
            cmd.current_dir(cwd);
        }

        let mut child = cmd.spawn().context("Failed to spawn MCP server")?;

        let stdin = child.stdin.take().context("No stdin on MCP server")?;
        let stdout = child.stdout.take().context("No stdout on MCP server")?;

        let client = Self {
            child: Mutex::new(child),
            stdin: Mutex::new(stdin),
            stdout: Mutex::new(BufReader::new(stdout)),
            next_id: AtomicU64::new(1),
            server_name: config.command.clone(),
            capabilities: Mutex::new(None),
        };

        // Perform initialize handshake
        client.initialize().await?;

        Ok(client)
    }

    /// Send the initialize request and initialized notification.
    async fn initialize(&self) -> Result<()> {
        let params = serde_json::json!({
            "protocolVersion": "2024-11-05",
            "capabilities": {},
            "clientInfo": {
                "name": "mur-commander",
                "version": env!("CARGO_PKG_VERSION")
            }
        });

        let response = self.request("initialize", Some(params)).await?;
        let init_result: InitializeResult = serde_json::from_value(
            response.context("Empty initialize response")?
        ).context("Invalid initialize response")?;

        info!(
            "MCP server initialized: {} (protocol {})",
            init_result.server_info.as_ref().map(|s| s.name.as_str()).unwrap_or("unknown"),
            init_result.protocol_version
        );

        *self.capabilities.lock().await = Some(init_result.capabilities);

        // Send initialized notification (no id, no response expected)
        self.notify("notifications/initialized", None).await?;

        Ok(())
    }

    /// List available tools from the server.
    pub async fn list_tools(&self) -> Result<Vec<McpTool>> {
        let response = self.request("tools/list", None).await?;
        let result: ToolsListResult = serde_json::from_value(
            response.context("Empty tools/list response")?
        ).context("Invalid tools/list response")?;
        Ok(result.tools)
    }

    /// Call a tool with the given arguments.
    pub async fn call_tool(&self, name: &str, arguments: serde_json::Value) -> Result<McpToolResult> {
        let params = serde_json::json!({
            "name": name,
            "arguments": arguments
        });

        let response = self.request("tools/call", Some(params)).await?;
        let result: McpToolResult = serde_json::from_value(
            response.context("Empty tools/call response")?
        ).context("Invalid tools/call response")?;

        Ok(result)
    }

    /// Get server capabilities.
    pub async fn capabilities(&self) -> Option<McpServerCapabilities> {
        self.capabilities.lock().await.clone()
    }

    /// Send a JSON-RPC request and wait for the response.
    async fn request(&self, method: &str, params: Option<serde_json::Value>) -> Result<Option<serde_json::Value>> {
        let id = self.next_id.fetch_add(1, Ordering::SeqCst);

        let request = JsonRpcRequest {
            jsonrpc: "2.0",
            id,
            method: method.to_string(),
            params,
        };

        let mut line = serde_json::to_string(&request)?;
        line.push('\n');

        // NOTE: Only log the method name, never the full request payload,
        // because params may contain secrets (API keys, tokens, etc.).
        debug!("MCP -> {}: method={} id={}", self.server_name, method, id);

        {
            let mut stdin = self.stdin.lock().await;
            stdin.write_all(line.as_bytes()).await?;
            stdin.flush().await?;
        }

        // Read response lines until we get one with our id (with timeout)
        let mut stdout = self.stdout.lock().await;
        let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(60);
        loop {
            let mut buf = String::new();
            let read = tokio::time::timeout_at(deadline, stdout.read_line(&mut buf)).await;
            let n = match read {
                Ok(result) => result?,
                Err(_) => anyhow::bail!("MCP request '{}' timed out after 60s", method),
            };
            if n == 0 {
                anyhow::bail!("MCP server closed stdout");
            }

            let buf = buf.trim();
            if buf.is_empty() {
                continue;
            }

            debug!("MCP ← {}: [response received]", self.server_name);

            let response: JsonRpcResponse = match serde_json::from_str(buf) {
                Ok(r) => r,
                Err(_) => {
                    // Could be a notification, skip
                    warn!("Skipping non-response line from MCP server");
                    continue;
                }
            };

            // Check if this is our response
            if response.id == Some(id) {
                if let Some(error) = response.error {
                    anyhow::bail!("MCP error ({}): {}", error.code, error.message);
                }
                return Ok(response.result);
            }

            // Not our response, skip (could be notification or out-of-order)
        }
    }

    /// Send a JSON-RPC notification (no id, no response).
    async fn notify(&self, method: &str, params: Option<serde_json::Value>) -> Result<()> {
        let notification = serde_json::json!({
            "jsonrpc": "2.0",
            "method": method,
            "params": params.unwrap_or(serde_json::json!({}))
        });

        let mut line = serde_json::to_string(&notification)?;
        line.push('\n');

        let mut stdin = self.stdin.lock().await;
        stdin.write_all(line.as_bytes()).await?;
        stdin.flush().await?;

        Ok(())
    }

    /// Shutdown the MCP server gracefully.
    pub async fn shutdown(&self) -> Result<()> {
        // Try to send shutdown request
        let _ = self.request("shutdown", None).await;
        let _ = self.notify("exit", None).await;

        // Give it a moment, then kill
        let mut child = self.child.lock().await;
        tokio::select! {
            _ = child.wait() => {},
            _ = tokio::time::sleep(std::time::Duration::from_secs(5)) => {
                let _ = child.kill().await;
            }
        }

        Ok(())
    }
}

impl Drop for McpClient {
    fn drop(&mut self) {
        // Best-effort kill on drop
        if let Ok(mut child) = self.child.try_lock() {
            let _ = child.start_kill();
        }
    }
}

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

    #[test]
    fn test_json_rpc_request_serialization() {
        let req = JsonRpcRequest {
            jsonrpc: "2.0",
            id: 1,
            method: "tools/list".to_string(),
            params: None,
        };
        let json = serde_json::to_string(&req).unwrap();
        assert!(json.contains("\"jsonrpc\":\"2.0\""));
        assert!(json.contains("\"method\":\"tools/list\""));
        assert!(!json.contains("params")); // skip_serializing_if
    }

    #[test]
    fn test_mcp_tool_deserialization() {
        let json = r#"{
            "name": "read_file",
            "description": "Read a file",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "path": {"type": "string"}
                }
            }
        }"#;
        let tool: McpTool = serde_json::from_str(json).unwrap();
        assert_eq!(tool.name, "read_file");
        assert_eq!(tool.description.as_deref(), Some("Read a file"));
    }

    #[test]
    fn test_mcp_content_text() {
        let json = r#"{"type": "text", "text": "hello"}"#;
        let content: McpContent = serde_json::from_str(json).unwrap();
        match content {
            McpContent::Text { text } => assert_eq!(text, "hello"),
            _ => panic!("Expected text content"),
        }
    }

    #[test]
    fn test_mcp_tool_result_deserialization() {
        let json = r#"{
            "content": [{"type": "text", "text": "result"}],
            "isError": false
        }"#;
        let result: McpToolResult = serde_json::from_str(json).unwrap();
        assert_eq!(result.content.len(), 1);
        assert!(!result.is_error);
    }
}