mcp-memory/stub/tests/e2e.rs

use serde_json::{Value, json};
use std::io::{BufRead, BufReader, Write};
use std::process::{Child, Command, Stdio};
use std::time::Duration;

struct ChildGuard(Child);

impl Drop for ChildGuard {
    fn drop(&mut self) {
        let _ = self.0.kill();
        let _ = self.0.wait();
    }
}

fn send_message(stdin: &mut std::process::ChildStdin, msg: Value) {
    let s = serde_json::to_string(&msg).unwrap();
    let payload = format!("{s}\n");
    stdin.write_all(payload.as_bytes()).unwrap();
    stdin.flush().unwrap();
}

fn read_message(reader: &mut impl BufRead) -> Option<Value> {
    let mut line = String::new();
    if reader.read_line(&mut line).unwrap_or(0) == 0 {
        return None;
    }
    serde_json::from_str(line.trim()).ok()
}

#[tokio::test]
async fn test_full_system_e2e_performance() {
    let temp_dir = std::env::temp_dir().join(format!(
        "mcp_e2e_{}",
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs()
    ));
    std::fs::create_dir_all(&temp_dir).unwrap();

    let test_port = "3042"; // Use a distinct port
    let test_auth_token = "test-token-12345";

    // Since tests run from inside `target/debug/deps`, and `cargo test` does not guarantee
    // `env!("CARGO_BIN_EXE_name")` works correctly for binaries compiled in other crates without build dependencies,
    // we use `CARGO_MANIFEST_DIR` (which points to `stub`) to reliably locate the workspace `target/debug`.
    let manifest_dir = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    let debug_dir = manifest_dir.parent().unwrap().join("target").join("debug");

    let server_exe = debug_dir.join(format!("mcp-memory-server{}", std::env::consts::EXE_SUFFIX));
    let nvim_name = if cfg!(windows) {
        "mcp-memory-win-nvim"
    } else {
        "mcp-memory-linux-nvim"
    };
    let nvim_exe = debug_dir.join(format!("{}{}", nvim_name, std::env::consts::EXE_SUFFIX));
    let stub_exe = debug_dir.join(format!("mcp-memory-stub{}", std::env::consts::EXE_SUFFIX));

    assert!(server_exe.exists(), "Server not found at {server_exe:?}");
    assert!(nvim_exe.exists(), "Nvim not found at {nvim_exe:?}");
    assert!(stub_exe.exists(), "Stub not found at {stub_exe:?}");

    // 1. Start Server
    let _server = ChildGuard(
        Command::new(&server_exe)
            .arg("--daemon")
            .env("MCP_PORT", test_port)
            .env("RUST_LOG", "debug")
            .env("MCP_MEMORY_STORE_DIR", temp_dir.to_str().unwrap())
            .env("MCP_AUTH_TOKEN", test_auth_token)
            .env("RUST_LOG", "debug")
            .stdout(Stdio::inherit())
            .stderr(Stdio::inherit())
            .spawn()
            .expect("Failed to start server"),
    );

    // Give server time to start
    let client = reqwest::Client::new();
    let mut started = false;
    for _ in 0..30 {
        if let Ok(resp) = client
            .get(format!("http://127.0.0.1:{test_port}/health"))
            .send()
            .await
            && resp.status().is_success()
        {
            started = true;
            break;
        }
        tokio::time::sleep(Duration::from_millis(500)).await;
    }
    assert!(started, "Server failed to start in time");

    // 2. Start Stub
    let mut stub = ChildGuard(
        Command::new(&stub_exe)
            .arg("--target")
            .arg(format!("http://127.0.0.1:{test_port}"))
            .env("MCP_MEMORY_STORE_DIR", temp_dir.to_str().unwrap())
            .env("MCP_AUTH_TOKEN", test_auth_token)
            .env("RUST_LOG", "debug")
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::inherit())
            .spawn()
            .expect("Failed to start stub"),
    );

    let mut stub_stdin = stub.0.stdin.take().unwrap();
    let mut stub_stdout = BufReader::new(stub.0.stdout.take().unwrap());

    // 3. Start Nvim Bridge
    let mut nvim = ChildGuard(
        Command::new(&nvim_exe)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::inherit())
            .spawn()
            .expect("Failed to start nvim bridge"),
    );

    let mut nvim_stdin = nvim.0.stdin.take().unwrap();
    let mut nvim_stdout = BufReader::new(nvim.0.stdout.take().unwrap());

    println!("Server, stub, and nvim spawned successfully");

    // Send 100 concurrent-like sequential rapid requests to Stub
    println!("Starting 100 requests to stub...");
    let start_time = std::time::Instant::now();
    for i in 1..=100 {
        let tools_req = json!({
            "jsonrpc": "2.0",
            "method": "tools/list",
            "params": {},
            "id": i
        });

        // Alternate between LSP header format and JSONL format
        if i % 2 == 0 {
            send_message(&mut stub_stdin, tools_req);
        } else {
            let s = serde_json::to_string(&tools_req).unwrap();
            stub_stdin.write_all(format!("{s}\n").as_bytes()).unwrap();
            stub_stdin.flush().unwrap();
        }

        let mut resp =
            read_message(&mut stub_stdout).expect("Failed to read rapid response from stub");
        while resp.get("id").is_none() || resp["id"].is_null() {
            resp = read_message(&mut stub_stdout).expect("Failed to read rapid response from stub");
        }
        assert_eq!(resp["id"], i);
    }
    let stub_duration = start_time.elapsed();
    println!("Stub 100 requests: {stub_duration:?}");

    // Send 100 concurrent-like sequential rapid requests to Win-Nvim
    println!("Starting 100 requests to nvim...");
    let start_time_nvim = std::time::Instant::now();
    for i in 1..=100 {
        let tools_req = json!({
            "jsonrpc": "2.0",
            "method": "tools/list",
            "params": {},
            "id": i
        });

        if i % 2 == 0 {
            send_message(&mut nvim_stdin, tools_req);
        } else {
            let s = serde_json::to_string(&tools_req).unwrap();
            nvim_stdin.write_all(format!("{s}\n").as_bytes()).unwrap();
            nvim_stdin.flush().unwrap();
        }

        let mut resp =
            read_message(&mut nvim_stdout).expect("Failed to read rapid response from win-nvim");
        while resp.get("id").is_none() || resp["id"].is_null() {
            resp = read_message(&mut nvim_stdout)
                .expect("Failed to read rapid response from win-nvim");
        }
        assert_eq!(resp["id"], i);
    }
    let nvim_duration = start_time_nvim.elapsed();

    println!("Stub 100 requests: {stub_duration:?}");
    println!("Win-Nvim 100 requests: {nvim_duration:?}");

    let _ = std::fs::remove_dir_all(temp_dir);
}