zesdex/src/app/mcp/manager.rs

use serde_json::{json, Value};
use serde::{Deserialize, Serialize};
use std::io::{BufRead, BufReader, Write};
use std::sync::{Arc, Mutex, OnceLock};


const MCP_CONNECT_TIMEOUT_MS: u64 = 20_000;
const MCP_CALL_TIMEOUT_MS: u64 = 60_000;

/// Global cache for `&'static str` names/descriptions of MCP tools, so we
/// never need `Box::leak`.  Entries are never removed (small, bounded by the
/// number of MCP tools ever registered in a session).
fn mcp_static_str(s: &str) -> &'static str {
    static CACHE: OnceLock<Mutex<Vec<&'static str>>> = OnceLock::new();
    let mut cache = CACHE.get_or_init(|| Mutex::new(Vec::new())).lock().unwrap();
    if let Some(&existing) = cache.iter().find(|e| **e == s) {
        return existing;
    }
    let leaked: &'static str = Box::leak(s.to_string().into_boxed_str());
    cache.push(leaked);
    leaked
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum McpTransport {
    Stdio {
        command: String,
        args: Vec<String>,
    },
    StreamableHttp {
        url: String,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpToolInfo {
    pub name: String,
    pub description: String,
    pub input_schema: Value,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpServer {
    pub name: String,
    pub transport: McpTransport,
    pub tools: Vec<McpToolInfo>,
    /// Held child-process handle so subsequent tool calls reuse the same
    /// connection instead of spawning a new child each time.  Not serialized
    /// because the child only lives in this process.
    #[serde(skip)]
    pub child_handle: Option<Arc<Mutex<StdioChild>>>,
}

#[derive(Debug)]
pub struct StdioChild {
    stdin: std::process::ChildStdin,
    stdout: BufReader<std::process::ChildStdout>,
    next_id: u64,
}

impl StdioChild {
    pub fn call(&mut self, method: &str, params: Value) -> anyhow::Result<Value> {
        self.next_id += 1;
        let id = self.next_id;
        let req = json!({
            "jsonrpc": "2.0",
            "id": id,
            "method": method,
            "params": params
        });
        let mut line = serde_json::to_string(&req)?;
        line.push('\n');
        self.stdin.write_all(line.as_bytes())?;
        self.stdin.flush()?;

        let mut response_line = String::new();
        let deadline = std::time::Instant::now()
            + std::time::Duration::from_millis(MCP_CALL_TIMEOUT_MS);
        loop {
            if std::time::Instant::now() > deadline {
                anyhow::bail!("MCP call timed out after {}ms", MCP_CALL_TIMEOUT_MS);
            }
            response_line.clear();
            match self.stdout.read_line(&mut response_line) {
                Ok(0) => anyhow::bail!("MCP stdio child process closed unexpectedly"),
                Ok(_) => {
                    let trimmed = response_line.trim();
                    if trimmed.is_empty() {
                        continue;
                    }
                    let resp: Value = serde_json::from_str(trimmed)
                        .map_err(|e| anyhow::anyhow!("invalid JSON from MCP server: {}", e))?;
                    if resp.get("id") == Some(&json!(id)) {
                        if let Some(err) = resp.get("error") {
                            anyhow::bail!("MCP error: {}", err);
                        }
                        return Ok(resp.get("result").cloned().unwrap_or(Value::Null));
                    }
                }
                Err(e) => anyhow::bail!("MCP stdio read error: {}", e),
            }
        }
    }
}

pub(crate) fn spawn_stdio_child(command: &str, extra_args: &[String]) -> anyhow::Result<StdioChild> {
    let parts: Vec<&str> = command.split_whitespace().collect();
    let (prog, prog_args) = parts.split_first()
        .ok_or_else(|| anyhow::anyhow!("MCP stdio command is empty"))?;

    let mut cmd = std::process::Command::new(prog);
    cmd.args(prog_args);
    cmd.args(extra_args);
    cmd.stdin(std::process::Stdio::piped());
    cmd.stdout(std::process::Stdio::piped());
    cmd.stderr(std::process::Stdio::null());

    let mut child = cmd.spawn()
        .map_err(|e| anyhow::anyhow!("failed to spawn MCP stdio server '{}': {}", command, e))?;

    let stdin = child.stdin.take()
        .ok_or_else(|| anyhow::anyhow!("failed to get stdin for MCP server"))?;
    let stdout = child.stdout.take()
        .ok_or_else(|| anyhow::anyhow!("failed to get stdout for MCP server"))?;

    let mut mcp = StdioChild {
        stdin,
        stdout: BufReader::new(stdout),
        next_id: 0,
    };

    let deadline = std::time::Instant::now()
        + std::time::Duration::from_millis(MCP_CONNECT_TIMEOUT_MS);

    let init_result = mcp.call("initialize", json!({
        "protocolVersion": "2024-11-05",
        "capabilities": {},
        "clientInfo": {
            "name": "zesdex",
            "version": "0.1.0"
        }
    }));

    if std::time::Instant::now() > deadline {
        anyhow::bail!("MCP initialize timed out");
    }

    init_result.map_err(|e| anyhow::anyhow!("MCP initialize failed: {}", e))?;

    let _ = mcp.call("notifications/initialized", json!({}));

    Ok(mcp)
}

fn call_via_stdio(
    existing_handle: Option<&Mutex<StdioChild>>,
    command: &str,
    extra_args: &[String],
    tool_name: &str,
    tool_args: &Value,
) -> anyhow::Result<String> {
    // Reuse the persistent child handle if available; otherwise spawn a new one.
    let mut guard;
    let child: &mut StdioChild = if let Some(mtx) = existing_handle {
        guard = mtx.lock().map_err(|e| anyhow::anyhow!("MCP handle lock: {}", e))?;
        &mut guard
    } else {
        let mut fresh = spawn_stdio_child(command, extra_args)?;
        let result = fresh.call("tools/call", json!({
            "name": tool_name,
            "arguments": tool_args
        }))?;
        return extract_text_content(&result);
    };

    let result = child.call("tools/call", json!({
        "name": tool_name,
        "arguments": tool_args
    }))?;

    extract_text_content(&result)
}

fn call_via_http(url: &str, tool_name: &str, tool_args: &Value) -> anyhow::Result<String> {
    let client = reqwest::blocking::Client::builder()
        .timeout(std::time::Duration::from_millis(MCP_CALL_TIMEOUT_MS))
        .connect_timeout(std::time::Duration::from_millis(MCP_CONNECT_TIMEOUT_MS))
        .build()
        .unwrap_or_else(|_| reqwest::blocking::Client::new());

    let request_id: u64 = 1;
    let body = json!({
        "jsonrpc": "2.0",
        "id": request_id,
        "method": "tools/call",
        "params": {
            "name": tool_name,
            "arguments": tool_args
        }
    });

    let resp = client.post(url)
        .header("Content-Type", "application/json")
        .json(&body)
        .send()
        .map_err(|e| anyhow::anyhow!("MCP HTTP request failed: {}", e))?;

    if !resp.status().is_success() {
        let status = resp.status();
        let text = resp.text().unwrap_or_default();
        anyhow::bail!("MCP HTTP server returned {}: {}", status, text);
    }

    let response: Value = resp.json()
        .map_err(|e| anyhow::anyhow!("invalid JSON from MCP HTTP server: {}", e))?;

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

    let result = response.get("result").cloned().unwrap_or(Value::Null);
    extract_text_content(&result)
}

fn extract_text_content(result: &Value) -> anyhow::Result<String> {
    if let Some(content) = result.get("content") {
        if let Some(arr) = content.as_array() {
            let text: Vec<String> = arr.iter().filter_map(|item| {
                if item.get("type").and_then(|t| t.as_str()) == Some("text") {
                    item.get("text").and_then(|t| t.as_str()).map(|s| s.to_string())
                } else {
                    None
                }
            }).collect();
            if !text.is_empty() {
                return Ok(text.join("\n"));
            }
        }
    }
    Ok(serde_json::to_string_pretty(result).unwrap_or_else(|_| result.to_string()))
}

#[derive(Debug, Clone)]
pub struct McpManager {
    pub servers: Vec<McpServer>,
}

pub struct McpToolAdapter {
    pub tool_name: String,
    pub server_name: String,
    pub transport: McpTransport,
    pub description: String,
    pub parameters: Value,
    /// Shared handle to a persistent child process (stdio transport only).
    pub child_handle: Option<Arc<Mutex<StdioChild>>>,
}

impl crate::tool::Tool for McpToolAdapter {
    fn name(&self) -> &'static str {
        mcp_static_str(&format!("mcp__{}__{}", self.server_name, self.tool_name))
    }

    fn description(&self) -> &'static str {
        mcp_static_str(&self.description)
    }

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

    fn run(&self, _ctx: &crate::tool::ToolCtx, args: &Value) -> anyhow::Result<String> {
        match &self.transport {
            McpTransport::Stdio { command, args: extra_args } => {
                call_via_stdio(self.child_handle.as_ref().map(|h| h.as_ref()), command, extra_args, &self.tool_name, args)
            }
            McpTransport::StreamableHttp { url } => {
                call_via_http(url, &self.tool_name, args)
            }
        }
    }
}

impl McpManager {
    pub fn new() -> Self {
        McpManager {
            servers: Vec::new(),
        }
    }

    pub fn as_tools(&self) -> Vec<Box<dyn crate::tool::Tool>> {
        self.servers.iter().flat_map(|server| {
            let handle = server.child_handle.clone();
            server.tools.iter().map(move |info| {
                Box::new(McpToolAdapter {
                    tool_name: info.name.clone(),
                    server_name: server.name.clone(),
                    transport: server.transport.clone(),
                    description: info.description.clone(),
                    parameters: info.input_schema.clone(),
                    child_handle: handle.clone(),
                }) as Box<dyn crate::tool::Tool>
            })
        }).collect()
    }

    /// Connects to an MCP server via stdio by spawning the child process, running
    /// the `initialize` handshake, calling `tools/list`, and registering the server
    /// with its advertised tools in `self.servers`. The child process stays alive
    /// for subsequent `tools/call` invocations via the stored `McpServer.tools`.
    pub fn connect_stdio(&mut self, name: &str, command: &str, extra_args: &[String]) -> anyhow::Result<()> {
        let transport = McpTransport::Stdio {
            command: command.to_string(),
            args: extra_args.to_vec(),
        };

        let mut child = spawn_stdio_child(command, extra_args)?;
        let result = child.call("tools/list", json!({}))?;

        let tools = if let Some(tool_list) = result.get("tools").and_then(|v| v.as_array()) {
            tool_list.iter().filter_map(|t| {
                Some(McpToolInfo {
                    name: t.get("name")?.as_str()?.to_string(),
                    description: t.get("description").and_then(|v| v.as_str()).unwrap_or("").to_string(),
                    input_schema: t.get("inputSchema").cloned().unwrap_or(serde_json::Value::Null),
                })
            }).collect()
        } else {
            Vec::new()
        };

        let handle = Arc::new(Mutex::new(child));

        self.servers.push(McpServer {
            name: name.to_string(),
            transport,
            tools,
            child_handle: Some(handle),
        });

        Ok(())
    }

    /// Removes a server by name. Returns `true` if a server was found and removed.
    #[allow(dead_code)]
    pub fn disconnect(&mut self, name: &str) -> bool {
        let len = self.servers.len();
        self.servers.retain(|s| s.name != name);
        self.servers.len() < len
    }
}