zarz 0.3.1-alpha

package/src/mcp/client.rs

@@ -0,0 +1,316 @@
+use anyhow::{anyhow, Context, Result};
+use serde_json::{json, Value};
+use std::collections::HashMap;
+use std::process::Stdio;
+use std::sync::atomic::{AtomicU64, Ordering};
+use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
+use tokio::process::{Child, ChildStdin, ChildStdout, Command};
+use tokio::sync::Mutex;
+
+use super::config::McpServerConfig;
+use super::types::*;
+
+/// MCP Client for communicating with MCP servers
+pub struct McpClient {
+    #[allow(dead_code)]
+    name: String,
+    config: McpServerConfig,
+    process: Option<Mutex<Child>>,
+    stdin: Option<Mutex<ChildStdin>>,
+    stdout: Option<Mutex<BufReader<ChildStdout>>>,
+    request_id: AtomicU64,
+    initialized: bool,
+    server_info: Option<ServerInfo>,
+    capabilities: Option<ServerCapabilities>,
+}
+
+impl McpClient {
+    /// Create a new MCP client
+    pub fn new(name: String, config: McpServerConfig) -> Self {
+        Self {
+            name,
+            config,
+            process: None,
+            stdin: None,
+            stdout: None,
+            request_id: AtomicU64::new(1),
+            initialized: false,
+            server_info: None,
+            capabilities: None,
+        }
+    }
+
+    /// Start the MCP server process (for STDIO servers)
+    pub async fn start(&mut self) -> Result<()> {
+        match &self.config {
+            McpServerConfig::Stdio { command, args, env } => {
+                // On Windows, wrap in cmd /c for proper PATH resolution
+                let mut cmd = if cfg!(target_os = "windows") {
+                    let mut win_cmd = Command::new("cmd");
+                    win_cmd.arg("/c");
+                    win_cmd.arg(command);
+                    if let Some(args) = args {
+                        win_cmd.args(args);
+                    }
+                    win_cmd
+                } else {
+                    let mut unix_cmd = Command::new(command);
+                    if let Some(args) = args {
+                        unix_cmd.args(args);
+                    }
+                    unix_cmd
+                };
+
+                if let Some(env_vars) = env {
+                    cmd.envs(env_vars);
+                }
+
+                cmd.stdin(Stdio::piped())
+                    .stdout(Stdio::piped())
+                    .stderr(Stdio::inherit());
+
+                let mut child = cmd.spawn()
+                    .with_context(|| format!("Failed to start MCP server: {}", command))?;
+
+                let stdin = child.stdin.take()
+                    .context("Failed to open stdin")?;
+                let stdout = child.stdout.take()
+                    .context("Failed to open stdout")?;
+
+                self.stdin = Some(Mutex::new(stdin));
+                self.stdout = Some(Mutex::new(BufReader::new(stdout)));
+                self.process = Some(Mutex::new(child));
+
+                self.initialize().await?;
+
+                Ok(())
+            }
+            _ => Err(anyhow!("Only STDIO servers are currently supported")),
+        }
+    }
+
+    /// Initialize MCP connection
+    async fn initialize(&mut self) -> Result<()> {
+        let params = json!({
+            "protocolVersion": "2024-11-05",
+            "capabilities": {
+                "tools": {},
+                "resources": {},
+                "prompts": {}
+            },
+            "clientInfo": {
+                "name": "ZarzCLI",
+                "version": "0.1.0"
+            }
+        });
+
+        let response = self.send_request("initialize", Some(params)).await?;
+
+        let result: InitializeResult = serde_json::from_value(response)
+            .context("Failed to parse initialize response")?;
+
+        self.server_info = Some(result.server_info);
+        self.capabilities = Some(result.capabilities);
+        self.initialized = true;
+
+        self.send_notification("notifications/initialized", None).await?;
+
+        Ok(())
+    }
+
+    /// List available tools from the MCP server
+    pub async fn list_tools(&self) -> Result<Vec<McpTool>> {
+        if !self.initialized {
+            return Err(anyhow!("MCP client not initialized"));
+        }
+
+        let response = self.send_request("tools/list", None).await?;
+        let result: ToolsListResult = serde_json::from_value(response)
+            .context("Failed to parse tools/list response")?;
+
+        Ok(result.tools)
+    }
+
+    /// Call a tool on the MCP server
+    #[allow(dead_code)]
+    pub async fn call_tool(&self, name: String, arguments: Option<HashMap<String, Value>>) -> Result<CallToolResult> {
+        if !self.initialized {
+            return Err(anyhow!("MCP client not initialized"));
+        }
+
+        let params = CallToolParams { name, arguments };
+        let params_value = serde_json::to_value(params)?;
+
+        let response = self.send_request("tools/call", Some(params_value)).await?;
+        let result: CallToolResult = serde_json::from_value(response)
+            .context("Failed to parse tools/call response")?;
+
+        Ok(result)
+    }
+
+    /// List available resources from the MCP server
+    #[allow(dead_code)]
+    pub async fn list_resources(&self) -> Result<Vec<McpResource>> {
+        if !self.initialized {
+            return Err(anyhow!("MCP client not initialized"));
+        }
+
+        let response = self.send_request("resources/list", None).await?;
+        let result: ResourcesListResult = serde_json::from_value(response)
+            .context("Failed to parse resources/list response")?;
+
+        Ok(result.resources)
+    }
+
+    /// List available prompts from the MCP server
+    #[allow(dead_code)]
+    pub async fn list_prompts(&self) -> Result<Vec<McpPrompt>> {
+        if !self.initialized {
+            return Err(anyhow!("MCP client not initialized"));
+        }
+
+        let response = self.send_request("prompts/list", None).await?;
+        let result: PromptsListResult = serde_json::from_value(response)
+            .context("Failed to parse prompts/list response")?;
+
+        Ok(result.prompts)
+    }
+
+    /// Send a JSON-RPC request and wait for response
+    async fn send_request(&self, method: &str, params: Option<Value>) -> Result<Value> {
+        let id = self.request_id.fetch_add(1, Ordering::SeqCst);
+
+        let request = JsonRpcRequest {
+            jsonrpc: "2.0".to_string(),
+            id,
+            method: method.to_string(),
+            params,
+        };
+
+        let request_json = serde_json::to_string(&request)?;
+
+        if let Some(stdin) = &self.stdin {
+            let mut stdin = stdin.lock().await;
+            stdin.write_all(request_json.as_bytes()).await?;
+            stdin.write_all(b"\n").await?;
+            stdin.flush().await?;
+        } else {
+            return Err(anyhow!("STDIN not available"));
+        }
+
+        if let Some(stdout) = &self.stdout {
+            let mut stdout = stdout.lock().await;
+
+            loop {
+                let mut line = String::new();
+                let bytes_read = stdout.read_line(&mut line).await?;
+
+                if bytes_read == 0 {
+                    return Err(anyhow!("MCP server closed the connection unexpectedly"));
+                }
+
+                if line.trim().is_empty() {
+                    continue;
+                }
+
+                let value: Value = serde_json::from_str(&line)
+                    .with_context(|| format!("Failed to parse JSON-RPC message: {}", line.trim()))?;
+
+                // Notifications do not include an `id`, so we skip them (surface useful info when present)
+                if value.get("id").is_none() {
+                    if let Some(method) = value.get("method").and_then(|m| m.as_str()) {
+                        if method == "notifications/message" {
+                            if let Some(msg) = value
+                                .get("params")
+                                .and_then(|p| p.get("data"))
+                                .and_then(|d| d.get("message"))
+                                .and_then(|m| m.as_str())
+                            {
+                                eprintln!("MCP notification: {}", msg);
+                            }
+                        }
+                    }
+                    continue;
+                }
+
+                let response: JsonRpcResponse = serde_json::from_value(value)
+                    .with_context(|| format!("Failed to parse JSON-RPC response: {}", line.trim()))?;
+
+                if let Some(error) = response.error {
+                    return Err(anyhow!("MCP error: {} (code: {})", error.message, error.code));
+                }
+
+                if let Some(result) = response.result {
+                    return Ok(result);
+                } else {
+                    return Err(anyhow!("No result in response"));
+                }
+            }
+        } else {
+            Err(anyhow!("STDOUT not available"))
+        }
+    }
+
+    /// Send a JSON-RPC notification (no response expected)
+    async fn send_notification(&self, method: &str, params: Option<Value>) -> Result<()> {
+        let notification = json!({
+            "jsonrpc": "2.0",
+            "method": method,
+            "params": params
+        });
+
+        let notification_json = serde_json::to_string(&notification)?;
+
+        if let Some(stdin) = &self.stdin {
+            let mut stdin = stdin.lock().await;
+            stdin.write_all(notification_json.as_bytes()).await?;
+            stdin.write_all(b"\n").await?;
+            stdin.flush().await?;
+        }
+
+        Ok(())
+    }
+
+    /// Get server info
+    pub fn server_info(&self) -> Option<&ServerInfo> {
+        self.server_info.as_ref()
+    }
+
+    /// Get server capabilities
+    #[allow(dead_code)]
+    pub fn capabilities(&self) -> Option<&ServerCapabilities> {
+        self.capabilities.as_ref()
+    }
+
+    /// Check if client is initialized
+    #[allow(dead_code)]
+    pub fn is_initialized(&self) -> bool {
+        self.initialized
+    }
+
+    /// Get server name
+    #[allow(dead_code)]
+    pub fn name(&self) -> &str {
+        &self.name
+    }
+
+    /// Stop the MCP server process
+    pub async fn stop(&mut self) -> Result<()> {
+        if let Some(process) = &self.process {
+            let mut process = process.lock().await;
+            process.kill().await?;
+        }
+        Ok(())
+    }
+}
+
+impl Drop for McpClient {
+    fn drop(&mut self) {
+        // Note: We can't await in drop, so we just kill the process synchronously
+        if let Some(process) = &self.process {
+            if let Ok(mut process) = process.try_lock() {
+                let _ = process.start_kill();
+            }
+        }
+    }
+}

package/src/mcp/config.rs

@@ -0,0 +1,133 @@
+use anyhow::{Context, Result};
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::fs;
+use std::path::PathBuf;
+
+/// MCP server configuration
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(untagged)]
+pub enum McpServerConfig {
+    Stdio {
+        command: String,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        args: Option<Vec<String>>,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        env: Option<HashMap<String, String>>,
+    },
+    Http {
+        url: String,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        headers: Option<HashMap<String, String>>,
+    },
+    Sse {
+        url: String,
+        #[serde(skip_serializing_if = "Option::is_none")]
+        headers: Option<HashMap<String, String>>,
+    },
+}
+
+/// Root MCP configuration file structure
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+pub struct McpConfig {
+    #[serde(rename = "mcpServers")]
+    pub mcp_servers: HashMap<String, McpServerConfig>,
+}
+
+impl McpConfig {
+    /// Get the path to the MCP config file (~/.zarz/mcp.json)
+    pub fn config_path() -> Result<PathBuf> {
+        let home = dirs::home_dir()
+            .context("Could not determine home directory")?;
+        Ok(home.join(".zarz").join("mcp.json"))
+    }
+
+    /// Load MCP config from file, or return default if file doesn't exist
+    pub fn load() -> Result<Self> {
+        let path = Self::config_path()?;
+
+        if !path.exists() {
+            return Ok(Self::default());
+        }
+
+        let content = fs::read_to_string(&path)
+            .context("Failed to read MCP config file")?;
+
+        let config: McpConfig = serde_json::from_str(&content)
+            .context("Failed to parse MCP config file")?;
+
+        Ok(config)
+    }
+
+    /// Save MCP config to file
+    pub fn save(&self) -> Result<()> {
+        let path = Self::config_path()?;
+
+        // Create parent directory if it doesn't exist
+        if let Some(parent) = path.parent() {
+            fs::create_dir_all(parent)
+                .context("Failed to create MCP config directory")?;
+        }
+
+        let content = serde_json::to_string_pretty(self)
+            .context("Failed to serialize MCP config")?;
+
+        fs::write(&path, content)
+            .context("Failed to write MCP config file")?;
+
+        Ok(())
+    }
+
+    /// Add a new MCP server
+    pub fn add_server(&mut self, name: String, config: McpServerConfig) {
+        self.mcp_servers.insert(name, config);
+    }
+
+    /// Remove an MCP server
+    pub fn remove_server(&mut self, name: &str) -> bool {
+        self.mcp_servers.remove(name).is_some()
+    }
+
+    /// Get an MCP server config by name
+    pub fn get_server(&self, name: &str) -> Option<&McpServerConfig> {
+        self.mcp_servers.get(name)
+    }
+
+    /// List all configured servers
+    #[allow(dead_code)]
+    pub fn list_servers(&self) -> Vec<String> {
+        self.mcp_servers.keys().cloned().collect()
+    }
+
+    /// Check if config has any servers
+    #[allow(dead_code)]
+    pub fn has_servers(&self) -> bool {
+        !self.mcp_servers.is_empty()
+    }
+}
+
+impl McpServerConfig {
+    /// Create a new STDIO server config
+    pub fn stdio(command: String, args: Option<Vec<String>>, env: Option<HashMap<String, String>>) -> Self {
+        McpServerConfig::Stdio { command, args, env }
+    }
+
+    /// Create a new HTTP server config
+    pub fn http(url: String, headers: Option<HashMap<String, String>>) -> Self {
+        McpServerConfig::Http { url, headers }
+    }
+
+    /// Create a new SSE server config
+    pub fn sse(url: String, headers: Option<HashMap<String, String>>) -> Self {
+        McpServerConfig::Sse { url, headers }
+    }
+
+    /// Get server type as string
+    pub fn server_type(&self) -> &'static str {
+        match self {
+            McpServerConfig::Stdio { .. } => "stdio",
+            McpServerConfig::Http { .. } => "http",
+            McpServerConfig::Sse { .. } => "sse",
+        }
+    }
+}

package/src/mcp/manager.rs

@@ -0,0 +1,186 @@
+use anyhow::{anyhow, Result};
+use std::collections::HashMap;
+use tokio::sync::RwLock;
+
+use super::client::McpClient;
+use super::config::{McpConfig, McpServerConfig};
+use super::types::{McpTool, McpResource, McpPrompt};
+
+/// Manages multiple MCP clients
+pub struct McpManager {
+    clients: RwLock<HashMap<String, McpClient>>,
+}
+
+impl McpManager {
+    /// Create a new MCP manager
+    pub fn new() -> Self {
+        Self {
+            clients: RwLock::new(HashMap::new()),
+        }
+    }
+
+    /// Load and start all configured MCP servers
+    pub async fn load_from_config(&self) -> Result<()> {
+        let config = McpConfig::load()?;
+
+        for (name, server_config) in config.mcp_servers {
+            if let Err(e) = self.start_server(name.clone(), server_config).await {
+                eprintln!("Warning: Failed to start MCP server '{}': {}", name, e);
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Start a specific MCP server
+    pub async fn start_server(&self, name: String, config: McpServerConfig) -> Result<()> {
+        let mut client = McpClient::new(name.clone(), config);
+        client.start().await?;
+
+        let mut clients = self.clients.write().await;
+        clients.insert(name, client);
+
+        Ok(())
+    }
+
+    /// Stop a specific MCP server
+    #[allow(dead_code)]
+    pub async fn stop_server(&self, name: &str) -> Result<()> {
+        let mut clients = self.clients.write().await;
+
+        if let Some(mut client) = clients.remove(name) {
+            client.stop().await?;
+            Ok(())
+        } else {
+            Err(anyhow!("Server '{}' not found", name))
+        }
+    }
+
+    /// List all running servers
+    pub async fn list_servers(&self) -> Vec<String> {
+        let clients = self.clients.read().await;
+        clients.keys().cloned().collect()
+    }
+
+    /// Get all available tools from all servers
+    pub async fn get_all_tools(&self) -> Result<HashMap<String, Vec<McpTool>>> {
+        let clients = self.clients.read().await;
+        let mut all_tools = HashMap::new();
+
+        for (name, client) in clients.iter() {
+            match client.list_tools().await {
+                Ok(tools) => {
+                    all_tools.insert(name.clone(), tools);
+                }
+                Err(e) => {
+                    eprintln!("Warning: Failed to get tools from '{}': {}", name, e);
+                }
+            }
+        }
+
+        Ok(all_tools)
+    }
+
+    /// Get all available resources from all servers
+    #[allow(dead_code)]
+    pub async fn get_all_resources(&self) -> Result<HashMap<String, Vec<McpResource>>> {
+        let clients = self.clients.read().await;
+        let mut all_resources = HashMap::new();
+
+        for (name, client) in clients.iter() {
+            match client.list_resources().await {
+                Ok(resources) => {
+                    all_resources.insert(name.clone(), resources);
+                }
+                Err(e) => {
+                    eprintln!("Warning: Failed to get resources from '{}': {}", name, e);
+                }
+            }
+        }
+
+        Ok(all_resources)
+    }
+
+    /// Get all available prompts from all servers
+    #[allow(dead_code)]
+    pub async fn get_all_prompts(&self) -> Result<HashMap<String, Vec<McpPrompt>>> {
+        let clients = self.clients.read().await;
+        let mut all_prompts = HashMap::new();
+
+        for (name, client) in clients.iter() {
+            match client.list_prompts().await {
+                Ok(prompts) => {
+                    all_prompts.insert(name.clone(), prompts);
+                }
+                Err(e) => {
+                    eprintln!("Warning: Failed to get prompts from '{}': {}", name, e);
+                }
+            }
+        }
+
+        Ok(all_prompts)
+    }
+
+    /// Call a tool on a specific server
+    #[allow(dead_code)]
+    pub async fn call_tool(
+        &self,
+        server_name: &str,
+        tool_name: String,
+        mut arguments: Option<HashMap<String, serde_json::Value>>,
+    ) -> Result<super::types::CallToolResult> {
+        if let Some(args) = arguments.as_mut() {
+            if let Some(value) = args.get_mut("sources") {
+                if let serde_json::Value::Array(items) = value {
+                    let contains_strings = items.iter().all(|item| matches!(item, serde_json::Value::String(_)));
+                    if contains_strings {
+                        eprintln!("Warning: Removing incompatible 'sources' parameter from MCP tool call (expected object array). Using server defaults instead.");
+                        args.remove("sources");
+                    }
+                }
+            }
+        }
+
+        let clients = self.clients.read().await;
+
+        let client = clients.get(server_name)
+            .ok_or_else(|| anyhow!("Server '{}' not found", server_name))?;
+
+        client.call_tool(tool_name, arguments).await
+    }
+
+    /// Get server info for a specific server
+    pub async fn get_server_info(&self, name: &str) -> Option<String> {
+        let clients = self.clients.read().await;
+        clients.get(name).and_then(|c| {
+            c.server_info().map(|info| {
+                format!("{} v{}", info.name, info.version)
+            })
+        })
+    }
+
+    /// Check if any servers are running
+    pub async fn has_servers(&self) -> bool {
+        let clients = self.clients.read().await;
+        !clients.is_empty()
+    }
+
+    /// Stop all servers
+    pub async fn stop_all(&self) -> Result<()> {
+        let mut clients = self.clients.write().await;
+
+        for (name, mut client) in clients.drain() {
+            if let Err(e) = client.stop().await {
+                eprintln!("Warning: Failed to stop server '{}': {}", name, e);
+            }
+        }
+
+        Ok(())
+    }
+}
+
+impl Default for McpManager {
+    fn default() -> Self {
+        Self::new()
+    }
+}

package/src/mcp/mod.rs

@@ -0,0 +1,12 @@
+// MCP (Model Context Protocol) support for ZarzCLI
+pub mod config;
+pub mod client;
+pub mod types;
+pub mod manager;
+
+pub use config::{McpConfig, McpServerConfig};
+#[allow(unused_imports)]
+pub use client::McpClient;
+#[allow(unused_imports)]
+pub use types::{McpTool, McpResource, McpPrompt};
+pub use manager::McpManager;