zesdex/crates/zesdex-backend/src/app/mcp/transport.rs

//! MCP transport layer: stdio child process management and HTTP client calls.
//! This module handles the low-level protocol details of communicating with
//! MCP servers (both spawned subprocesses and remote HTTP endpoints).
//!
//! Flow: `spawn_stdio_child` → `StdioChild::call` for JSON-RPC messages;
//! `call_via_stdio` / `call_via_http` are convenience wrappers for
//! `tools/call` that reuse a persistent child handle or spawn a fresh one.

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

// ---------------------------------------------------------------------------
// Constants
// ---------------------------------------------------------------------------

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

// ---------------------------------------------------------------------------
// Static string cache
// ---------------------------------------------------------------------------

/// 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).
pub(super) fn mcp_static_str(s: &str) -> &'static str {
    static CACHE: OnceLock<Mutex<Vec<&'static str>>> = OnceLock::new();
    let mut cache = match CACHE.get_or_init(|| Mutex::new(Vec::new())).lock() {
        Ok(c) => c,
        Err(poisoned) => {
            warn!("[mcp] static string cache mutex poisoned, recovering");
            poisoned.into_inner()
        }
    };
    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
}

// ---------------------------------------------------------------------------
// Core transport types
// ---------------------------------------------------------------------------

/// How an MCP server is reached: a spawned child process talking
/// newline-delimited JSON-RPC over stdio, or a remote HTTP endpoint.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum McpTransport {
    Stdio { command: String, args: Vec<String> },
    StreamableHttp { url: String },
}

/// A single tool advertised by an MCP server, as returned by `tools/list`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpToolInfo {
    pub name: String,
    pub description: String,
    pub input_schema: Value,
}

// ---------------------------------------------------------------------------
// Stdio child process handle
// ---------------------------------------------------------------------------

/// Live handle to an MCP server child process communicating over stdio
/// via newline-delimited JSON-RPC 2.0.
#[derive(Debug)]
pub struct StdioChild {
    stdin: std::process::ChildStdin,
    stdout: BufReader<std::process::ChildStdout>,
    next_id: u64,
}

impl StdioChild {
    /// Send a JSON-RPC request to the child and block for its matching response.
    ///
    /// Flow: assign the next request id → write request + newline to stdin →
    /// loop reading lines from stdout until one has a matching `id` or the
    /// timeout elapses → return its `result` (or error out on an `error` field).
    ///
    /// Why: the child may interleave unrelated/malformed lines, so blank
    /// lines are skipped and non-matching ids are ignored rather than
    /// treated as a protocol violation.
    ///
    /// Return: the `result` value of the matching response, or `Err` on
    /// timeout, EOF, JSON-RPC error, or I/O failure.
    pub fn call(&mut self, method: &str, params: &Value) -> anyhow::Result<Value> {
        const MAX_LINE_LENGTH: usize = 1_048_576; // 1 MiB
        self.next_id += 1;
        let id = self.next_id;
        debug!(method = method, id = id, "MCP stdio call");
        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 {MCP_CALL_TIMEOUT_MS}ms");
            }
            // Read one byte at a time up to MAX_LINE_LENGTH to prevent
            // OOM from a malicious server (CWE-400). BufReader already
            // buffers reads, so byte-by-byte over a buffered reader is
            // cheap (hits the in-memory buffer).
            response_line.clear();
            let mut line_truncated = false;
            loop {
                let byte = match self.stdout.fill_buf() {
                    Ok([]) => {
                        // EOF without newline
                        anyhow::bail!("MCP stdio child process closed unexpectedly");
                    }
                    Ok(buf) => {
                        let b = buf[0];
                        self.stdout.consume(1);
                        b
                    }
                    Err(e) => anyhow::bail!("MCP stdio read error: {e}"),
                };
                if byte == b'\n' {
                    break;
                }
                if response_line.len() >= MAX_LINE_LENGTH {
                    line_truncated = true;
                    // Consume rest of line to keep stream in sync
                    loop {
                        let buf = self
                            .stdout
                            .fill_buf()
                            .map_err(|e| anyhow::anyhow!("MCP stdio read error: {e}"))?;
                        if buf.is_empty() {
                            anyhow::bail!("MCP stdio child closed mid-line");
                        }
                        if buf[0] == b'\n' {
                            self.stdout.consume(1);
                            break;
                        }
                        self.stdout.consume(1);
                    }
                    break;
                }
                response_line.push(byte as char);
            }
            if line_truncated {
                anyhow::bail!("MCP response line exceeded {MAX_LINE_LENGTH} byte limit");
            }
            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_else(|| {
                    warn!("MCP stdio response missing 'result' field: {}", trimmed);
                    Value::Null
                }));
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Spawning and connecting
// ---------------------------------------------------------------------------

pub(crate) fn spawn_stdio_child(
    command: &str,
    extra_args: &[String],
) -> anyhow::Result<StdioChild> {
    info!(command = command, "MCP spawn stdio child");
    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());
    // Pipe stderr so diagnostics from MCP servers are surfaced via tracing
    // rather than discarded silently, making connectivity issues debugable.
    cmd.stderr(std::process::Stdio::piped());

    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)
}

// ---------------------------------------------------------------------------
// Tool-call helpers
// ---------------------------------------------------------------------------

pub(super) fn call_via_stdio(
    existing_handle: Option<&Mutex<StdioChild>>,
    command: &str,
    extra_args: &[String],
    tool_name: &str,
    tool_args: &Value,
) -> anyhow::Result<String> {
    debug!(tool = tool_name, has_handle = existing_handle.is_some(), "MCP call_via_stdio");
    // 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 {
        // No persistent handle — spawn a fresh child for this one call.
        let mut fresh = spawn_stdio_child(command, extra_args)?;
        let result = fresh.call(
            "tools/call",
            &json!({
                "name": tool_name,
                "arguments": tool_args
            }),
        )?;
        return Ok(extract_text_content(&result));
    };

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

    Ok(extract_text_content(&result))
}

pub(super) fn call_via_http(url: &str, tool_name: &str, tool_args: &Value) -> anyhow::Result<String> {
    debug!(tool = tool_name, url = url, "MCP call_via_http");
    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(|e| {
            warn!(
                "MCP HTTP client builder failed with connect timeout: {}. \
                 retrying without connect timeout",
                e,
            );
            reqwest::blocking::Client::builder()
                .timeout(std::time::Duration::from_millis(MCP_CALL_TIMEOUT_MS))
                .build()
                .unwrap_or_else(|e2| {
                    warn!(
                        "MCP also failed: {}. using default client (no configured timeouts)",
                        e2,
                    );
                    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_else(|e| {
            warn!("MCP failed to read HTTP response body: {}", e);
            String::new()
        });
        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_else(|| {
        warn!("MCP HTTP response missing 'result' field");
        Value::Null
    });
    Ok(extract_text_content(&result))
}

pub(super) fn extract_text_content(result: &Value) -> 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(std::string::ToString::to_string)
                    } else {
                        None
                    }
                })
                .collect();
            if !text.is_empty() {
                return text.join("\n");
            }
        }
    }
    serde_json::to_string_pretty(result).unwrap_or_else(|e| {
        warn!("MCP failed to pretty-print result: {}", e);
        result.to_string()
    })
}