zesdex/src/app/mcp/manager.rs

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


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

#[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>,
}

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

impl StdioChild {
    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),
            }
        }
    }
}

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(command: &str, extra_args: &[String], tool_name: &str, tool_args: &Value) -> anyhow::Result<String> {
    let mut child = spawn_stdio_child(command, extra_args)?;
    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,
}

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

    fn description(&self) -> &'static str {
        Box::leak(self.description.clone().into_boxed_str())
    }

    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(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| {
            server.tools.iter().map(|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(),
                }) as Box<dyn crate::tool::Tool>
            })
        }).collect()
    }
}