refactor(mcp): strip sse fallback architecture in favor of pure websockets and fix nvim NDJSON bug
This commit is contained in:
1 parent
0e29b12ac8
commit
3716c3e698
33 files changed
+2082
-1756
No files matched your search
+5
-1
@@ -4,16 +4,20 @@ version = "0.1.0"
|
||||
edition = "2024"
|
||||
|
||||
[dependencies]
|
||||
#rustls-tls = "0.2"
|
||||
clap = { version = "4.6.6", features = ["derive"] }
|
||||
reqwest = { version = "0.12", default-features = false, features = ["stream", "rustls-tls"] }
|
||||
tokio = { version = "1.53.1", features = ["full"] }
|
||||
tokio-util = { version = "0.7.19", features = ["io"] }
|
||||
futures-util = "0.3.34"
|
||||
tokio-tungstenite = "0.21.0"
|
||||
tokio-tungstenite = { version = "0.21.0" }
|
||||
tracing-appender = "0.2.5"
|
||||
tracing = "0.1.44"
|
||||
tracing-subscriber = "0.3.23"
|
||||
dirs = "7.0.0"
|
||||
serde_json = "1.0.151"
|
||||
|
||||
|
||||
|
||||
[dev-dependencies]
|
||||
serde_json = "1.0.151"
|
||||
@@ -0,0 +1,116 @@
|
||||
use reqwest::Client;
|
||||
use std::env;
|
||||
use std::time::Duration;
|
||||
use futures_util::StreamExt;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
tracing_subscriber::fmt::init();
|
||||
|
||||
let target = env::var("MCP_TARGET").unwrap_or_else(|_| "https://127.0.0.1:3000".to_string());
|
||||
let token = env::var("MCP_AUTH_TOKEN").unwrap_or_else(|_| "jP76lUJ5DtFRZmcvXH8LKdCTIkp29eAf".to_string());
|
||||
|
||||
tracing::info!("Starting skeletal client to {}", target);
|
||||
|
||||
let client = Client::builder()
|
||||
.danger_accept_invalid_certs(true)
|
||||
.build()?;
|
||||
|
||||
let sse_url = format!("{}/sse", target);
|
||||
|
||||
tracing::info!("Connecting to SSE: {}", sse_url);
|
||||
|
||||
let res = client.get(&sse_url)
|
||||
.bearer_auth(&token)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !res.status().is_success() {
|
||||
tracing::error!("Failed to connect to SSE: {}", res.status());
|
||||
return Err("SSE connection failed".into());
|
||||
}
|
||||
|
||||
tracing::info!("SSE Connected. Reading stream...");
|
||||
|
||||
let mut stream = res.bytes_stream();
|
||||
let mut buffer = Vec::new();
|
||||
let mut post_endpoint = None;
|
||||
|
||||
// Read the initial event containing the POST endpoint
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let bytes = chunk?;
|
||||
buffer.extend_from_slice(&bytes);
|
||||
|
||||
while let Some(pos) = buffer.windows(2).position(|w| w == b"\n\n" || w == b"\r\n") {
|
||||
let msg_bytes = buffer.drain(..pos).collect::<Vec<_>>();
|
||||
buffer.drain(..2);
|
||||
|
||||
let text = String::from_utf8_lossy(&msg_bytes);
|
||||
let mut is_endpoint = false;
|
||||
let mut data_content = String::new();
|
||||
|
||||
for line in text.lines() {
|
||||
if line.starts_with("event: endpoint") {
|
||||
is_endpoint = true;
|
||||
} else if let Some(data) = line.strip_prefix("data: ") {
|
||||
data_content.push_str(data);
|
||||
}
|
||||
}
|
||||
|
||||
if is_endpoint && !data_content.is_empty() {
|
||||
tracing::info!("Received POST endpoint: {}", data_content);
|
||||
post_endpoint = Some(data_content);
|
||||
break;
|
||||
} else {
|
||||
tracing::info!("Received early SSE data: {}", text);
|
||||
}
|
||||
}
|
||||
if post_endpoint.is_some() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let post_endpoint = post_endpoint.ok_or("Did not receive endpoint from SSE stream")?;
|
||||
let post_url = format!("{}{}", target, post_endpoint);
|
||||
|
||||
let payload = r#"{"jsonrpc":"2.0","id":999,"method":"server/discover","params":{}}"#;
|
||||
|
||||
tracing::info!("Sending test payload to {}", post_url);
|
||||
tracing::info!("Payload: {}", payload);
|
||||
|
||||
let post_res = client.post(&post_url)
|
||||
.bearer_auth(&token)
|
||||
.header("Content-Type", "application/json")
|
||||
.body(payload.to_string())
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
tracing::info!("POST Response Status: {}", post_res.status());
|
||||
let post_body = post_res.text().await?;
|
||||
tracing::info!("POST Response Body: {}", post_body);
|
||||
|
||||
// Wait for the SSE stream to deliver the response
|
||||
tracing::info!("Waiting 2 seconds for SSE response delivery...");
|
||||
|
||||
let mut timeout = tokio::time::interval(Duration::from_secs(2));
|
||||
timeout.tick().await; // first tick is immediate
|
||||
|
||||
tokio::select! {
|
||||
_ = timeout.tick() => {
|
||||
tracing::warn!("Timed out waiting for SSE response.");
|
||||
}
|
||||
_ = async {
|
||||
while let Some(chunk) = stream.next().await {
|
||||
if let Ok(bytes) = chunk {
|
||||
tracing::info!("Received SSE Chunk: {}", String::from_utf8_lossy(&bytes));
|
||||
break;
|
||||
}
|
||||
}
|
||||
} => {
|
||||
tracing::info!("Successfully read SSE response from stream.");
|
||||
}
|
||||
}
|
||||
|
||||
tracing::info!("Skeletal client test complete.");
|
||||
Ok(())
|
||||
}
|
||||
+35
-37
@@ -10,9 +10,6 @@ struct Cli {
|
||||
/// Target URL for the stub to proxy messages to
|
||||
#[arg(long, default_value = "http://localhost:3000")]
|
||||
target: String,
|
||||
/// Optional command to execute if the target server is unreachable
|
||||
#[arg(long)]
|
||||
wake_cmd: Option<String>,
|
||||
}
|
||||
|
||||
async fn read_mcp_message(stdin: &mut tokio::io::BufReader<tokio::io::Stdin>) -> Option<String> {
|
||||
@@ -26,6 +23,11 @@ async fn read_mcp_message(stdin: &mut tokio::io::BufReader<tokio::io::Stdin>) ->
|
||||
return None;
|
||||
}
|
||||
tracing::info!("Read {} bytes from stdin: {:?}", bytes_read, line);
|
||||
|
||||
if line.starts_with('{') {
|
||||
return Some(line.trim_end().to_string());
|
||||
}
|
||||
|
||||
let line = line.trim_end();
|
||||
if line.is_empty() {
|
||||
break;
|
||||
@@ -46,24 +48,17 @@ async fn read_mcp_message(stdin: &mut tokio::io::BufReader<tokio::io::Stdin>) ->
|
||||
}
|
||||
|
||||
fn init_logging(app_name: &str) -> Option<tracing_appender::non_blocking::WorkerGuard> {
|
||||
let base_dir = std::env::var("MCP_MEMORY_STORE_DIR").unwrap_or_else(|_| {
|
||||
dirs::home_dir()
|
||||
.map(|mut h| {
|
||||
h.push(".gemini/mcp_memory");
|
||||
h.to_string_lossy().to_string()
|
||||
})
|
||||
.unwrap_or_else(|| ".gemini/mcp_memory".into())
|
||||
});
|
||||
let log_dir = std::path::PathBuf::from(base_dir).join("logs");
|
||||
std::fs::create_dir_all(&log_dir).unwrap_or_default();
|
||||
let mut base_dir = dirs::home_dir().unwrap_or_else(|| std::path::PathBuf::from("."));
|
||||
base_dir.push(".gemini/mcp_memory/logs");
|
||||
std::fs::create_dir_all(&base_dir).unwrap_or_default();
|
||||
|
||||
let file_appender = tracing_appender::rolling::daily(log_dir, format!("{}.log", app_name));
|
||||
let file_appender = tracing_appender::rolling::daily(base_dir, format!("{}.log", app_name));
|
||||
let (non_blocking, guard) = tracing_appender::non_blocking(file_appender);
|
||||
|
||||
let _ = tracing_subscriber::fmt()
|
||||
.with_writer(non_blocking)
|
||||
.with_ansi(false)
|
||||
.with_max_level(tracing::Level::INFO)
|
||||
.with_max_level(tracing::Level::TRACE)
|
||||
.try_init();
|
||||
|
||||
Some(guard)
|
||||
@@ -89,7 +84,6 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let ws_url = target_url.replace("http://", "ws://").replace("https://", "wss://");
|
||||
let ws_url = format!("{}/ws?client=proxy", ws_url);
|
||||
let msg_rx = Arc::new(tokio::sync::Mutex::new(msg_rx));
|
||||
let wake_cmd = cli.wake_cmd;
|
||||
|
||||
loop {
|
||||
if shutdown_rx.try_recv().is_ok() {
|
||||
@@ -98,7 +92,18 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
}
|
||||
|
||||
tracing::info!("Attempting to connect to {}", ws_url);
|
||||
match tokio_tungstenite::connect_async(&ws_url).await {
|
||||
|
||||
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
|
||||
let mut request = match ws_url.clone().into_client_request() {
|
||||
Ok(req) => req,
|
||||
Err(e) => {
|
||||
tracing::error!("Failed to parse target URL {}: {}", ws_url, e);
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
match tokio_tungstenite::connect_async(request).await {
|
||||
Ok((ws_stream, _)) => {
|
||||
tracing::info!("Successfully connected to target server");
|
||||
let (mut write, mut read) = ws_stream.split();
|
||||
@@ -110,7 +115,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
match rx.recv().await {
|
||||
Some(msg) => {
|
||||
drop(rx);
|
||||
tracing::info!("Forwarding message to target server");
|
||||
tracing::info!("Forwarding message to target server (length: {}): {}", msg.len(), if msg.len() > 1000 { format!("{}...", &msg[..1000]) } else { msg.clone() });
|
||||
if write.send(tokio_tungstenite::tungstenite::Message::Text(msg)).await.is_err() {
|
||||
tracing::error!("Failed to write to websocket");
|
||||
break;
|
||||
@@ -124,25 +129,25 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let mut recv_task = tokio::spawn(async move {
|
||||
while let Some(Ok(msg)) = read.next().await {
|
||||
if let tokio_tungstenite::tungstenite::Message::Text(text) = msg {
|
||||
tracing::info!("Received message from target server, proxying to stdout");
|
||||
let payload = format!("Content-Length: {}\r\n\r\n{}", text.len(), text);
|
||||
use std::io::Write;
|
||||
let mut stdout = std::io::stdout();
|
||||
let _ = stdout.write_all(payload.as_bytes());
|
||||
let _ = stdout.flush();
|
||||
tracing::info!("Received message from target server (length: {}): {}", text.len(), if text.len() > 1000 { format!("{}...", &text[..1000]) } else { text.clone() });
|
||||
let payload = format!("{}\n", text);
|
||||
use tokio::io::AsyncWriteExt;
|
||||
let mut stdout = tokio::io::stdout();
|
||||
let _ = stdout.write_all(payload.as_bytes()).await;
|
||||
let _ = stdout.flush().await;
|
||||
}
|
||||
}
|
||||
tracing::error!("Websocket read loop exited");
|
||||
});
|
||||
|
||||
tokio::select! {
|
||||
_ = shutdown_rx.recv() => {
|
||||
_ = shutdown_rx.recv() => { tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
|
||||
tracing::info!("Shutdown received while connected");
|
||||
return Ok(()); // Stdin closed, exit entirely
|
||||
|
||||
}
|
||||
_ = &mut send_task => {
|
||||
tracing::error!("Send task exited");
|
||||
recv_task.abort();
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; recv_task.abort();
|
||||
}
|
||||
_ = &mut recv_task => {
|
||||
tracing::error!("Recv task exited");
|
||||
@@ -152,16 +157,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("Failed to connect to target server: {}", e);
|
||||
if let Some(ref cmd) = wake_cmd {
|
||||
tracing::info!("Executing wake command: {}", cmd);
|
||||
let parts: Vec<&str> = cmd.split_whitespace().collect();
|
||||
if !parts.is_empty() {
|
||||
let _ = std::process::Command::new(parts[0])
|
||||
.args(&parts[1..])
|
||||
.spawn();
|
||||
}
|
||||
}
|
||||
tracing::error!("Failed to connect via WSS: {}", e);
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
||||
}
|
||||
}
|
||||
@@ -169,3 +165,5 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+92
-77
@@ -5,107 +5,104 @@ use std::time::Duration;
|
||||
|
||||
fn send_message(stdin: &mut std::process::ChildStdin, msg: Value) {
|
||||
let s = serde_json::to_string(&msg).unwrap();
|
||||
let payload = format!("Content-Length: {}\r\n\r\n{}", s.len(), s);
|
||||
let payload = format!("{}\n", s);
|
||||
stdin.write_all(payload.as_bytes()).unwrap();
|
||||
stdin.flush().unwrap();
|
||||
}
|
||||
|
||||
fn read_message(stdout: &mut std::process::ChildStdout) -> Option<Value> {
|
||||
let mut reader = BufReader::new(stdout);
|
||||
let mut length = 0;
|
||||
|
||||
loop {
|
||||
let mut line = String::new();
|
||||
if reader.read_line(&mut line).unwrap_or(0) == 0 {
|
||||
return None;
|
||||
}
|
||||
let line = line.trim_end();
|
||||
if line.is_empty() {
|
||||
break;
|
||||
}
|
||||
if let Some(len_str) = line.strip_prefix("Content-Length: ") {
|
||||
length = len_str.parse().unwrap_or(0);
|
||||
}
|
||||
}
|
||||
|
||||
if length == 0 {
|
||||
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;
|
||||
}
|
||||
|
||||
let mut buf = vec![0u8; length];
|
||||
reader.read_exact(&mut buf).unwrap();
|
||||
let body_str = String::from_utf8_lossy(&buf);
|
||||
|
||||
Some(serde_json::from_str(&body_str).unwrap())
|
||||
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 = "3031"; // Use a distinct port
|
||||
let test_port = "3042"; // Use a distinct port
|
||||
let test_auth_token = "test-token-12345";
|
||||
|
||||
let mut exe_dir = std::env::current_exe().unwrap();
|
||||
exe_dir.pop(); // pop test executable name
|
||||
exe_dir.pop(); // pop deps/
|
||||
|
||||
let mut server_exe = exe_dir.join("mcp-memory-server.exe");
|
||||
if !server_exe.exists() {
|
||||
let mut target_dir = std::env::current_dir().unwrap();
|
||||
if target_dir.ends_with("stub") {
|
||||
target_dir.pop();
|
||||
}
|
||||
server_exe = target_dir.join("target").join("debug").join("mcp-memory-server.exe");
|
||||
}
|
||||
|
||||
let mut nvim_exe = exe_dir.join("mcp-memory-win-nvim.exe");
|
||||
if !nvim_exe.exists() {
|
||||
let mut target_dir = std::env::current_dir().unwrap();
|
||||
if target_dir.ends_with("stub") {
|
||||
target_dir.pop();
|
||||
}
|
||||
nvim_exe = target_dir.join("target").join("debug").join("mcp-memory-win-nvim.exe");
|
||||
}
|
||||
// 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 mut server = Command::new(&server_exe)
|
||||
.env("MCP_PORT", test_port)
|
||||
let mut server = 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())
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.env("MCP_AUTH_TOKEN", test_auth_token).env("RUST_LOG", "debug")
|
||||
.stdout(Stdio::inherit())
|
||||
.stderr(Stdio::inherit())
|
||||
.spawn()
|
||||
.expect("Failed to start server");
|
||||
|
||||
tokio::time::sleep(Duration::from_secs(2)).await;
|
||||
// Give server time to generate TLS cert and start
|
||||
let client = reqwest::Client::builder()
|
||||
.danger_accept_invalid_certs(true)
|
||||
.build()
|
||||
.unwrap();
|
||||
let mut started = false;
|
||||
for _ in 0..30 {
|
||||
if let Ok(resp) = client.get(format!("http://127.0.0.1:{}/health", test_port)).send().await {
|
||||
if 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 stub_exe = env!("CARGO_BIN_EXE_mcp-memory-stub");
|
||||
let mut stub = Command::new(stub_exe)
|
||||
let mut stub = 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::null())
|
||||
.stderr(Stdio::inherit())
|
||||
.spawn()
|
||||
.expect("Failed to start stub");
|
||||
|
||||
let mut stub_stdin = stub.stdin.take().unwrap();
|
||||
let mut stub_stdout = stub.stdout.take().unwrap();
|
||||
let mut stub_stdout = BufReader::new(stub.stdout.take().unwrap());
|
||||
|
||||
// 3. Start Win-Nvim
|
||||
let mut win_nvim = Command::new(&nvim_exe)
|
||||
// 3. Start Nvim Bridge
|
||||
let mut nvim = Command::new(&nvim_exe)
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::null())
|
||||
.stderr(Stdio::inherit())
|
||||
.spawn()
|
||||
.expect("Failed to start win-nvim");
|
||||
.expect("Failed to start nvim bridge");
|
||||
|
||||
let mut nvim_stdin = win_nvim.stdin.take().unwrap();
|
||||
let mut nvim_stdout = win_nvim.stdout.take().unwrap();
|
||||
let mut nvim_stdin = nvim.stdin.take().unwrap();
|
||||
let mut nvim_stdout = BufReader::new(nvim.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!({
|
||||
@@ -114,13 +111,27 @@ async fn test_full_system_e2e_performance() {
|
||||
"params": {},
|
||||
"id": i
|
||||
});
|
||||
send_message(&mut stub_stdin, tools_req);
|
||||
let resp = read_message(&mut stub_stdout).expect("Failed to read rapid response from stub");
|
||||
|
||||
// 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!("{}\n", s).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!({
|
||||
@@ -129,8 +140,19 @@ async fn test_full_system_e2e_performance() {
|
||||
"params": {},
|
||||
"id": i
|
||||
});
|
||||
send_message(&mut nvim_stdin, tools_req);
|
||||
let resp = read_message(&mut nvim_stdout).expect("Failed to read rapid response from win-nvim");
|
||||
|
||||
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!("{}\n", s).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();
|
||||
@@ -139,15 +161,8 @@ async fn test_full_system_e2e_performance() {
|
||||
println!("Win-Nvim 100 requests: {:?}", nvim_duration);
|
||||
|
||||
// Cleanup
|
||||
let _ = stub.kill();
|
||||
let _ = win_nvim.kill();
|
||||
let _ = server.kill();
|
||||
let _ = stub.kill();
|
||||
let _ = nvim.kill();
|
||||
let _ = std::fs::remove_dir_all(temp_dir);
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,88 @@
|
||||
use std::process::Stdio;
|
||||
use std::time::{Duration, Instant};
|
||||
use tokio::process::Command;
|
||||
|
||||
fn get_stub_exe() -> std::path::PathBuf {
|
||||
let manifest_dir = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"));
|
||||
let debug_dir = manifest_dir.parent().unwrap().join("target").join("debug");
|
||||
debug_dir.join(format!("mcp-memory-stub{}", std::env::consts::EXE_SUFFIX))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_stub_connection_refused() {
|
||||
let _ = std::process::Command::new("cargo").arg("build").arg("--bin").arg("mcp-memory-stub").status();
|
||||
|
||||
let target = "http://127.0.0.1:49999";
|
||||
|
||||
let start = Instant::now();
|
||||
let mut child = Command::new(get_stub_exe())
|
||||
.arg("--target")
|
||||
.arg(target)
|
||||
.stdin(Stdio::null()) // close stdin immediately to simulate EOF
|
||||
.spawn()
|
||||
.expect("Failed to execute stub");
|
||||
|
||||
let res = tokio::time::timeout(Duration::from_secs(5), child.wait()).await;
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
assert!(res.is_ok(), "Stub hung on connection refused! Took {:?}", elapsed);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_stub_handles_eof_cleanly() {
|
||||
let _ = std::process::Command::new("cargo").arg("build").arg("--bin").arg("mcp-memory-stub").status();
|
||||
|
||||
let target = "http://127.0.0.1:49998";
|
||||
|
||||
let mut child = Command::new(get_stub_exe())
|
||||
.arg("--target")
|
||||
.arg(target)
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.spawn()
|
||||
.expect("Failed to execute stub");
|
||||
|
||||
if let Some(mut stdin) = child.stdin.take() {
|
||||
use tokio::io::AsyncWriteExt;
|
||||
let msg = "Content-Length: 51\r\n\r\n{\"jsonrpc\":\"2.0\",\"method\":\"tools/list\",\"params\":{},\"id\":1}";
|
||||
stdin.write_all(msg.as_bytes()).await.unwrap();
|
||||
} // stdin dropped here
|
||||
|
||||
let start = Instant::now();
|
||||
let res = tokio::time::timeout(Duration::from_secs(5), child.wait()).await;
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
assert!(res.is_ok(), "Stub hung after EOF! Took {:?}", elapsed);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_stub_sse_fallback_failure() {
|
||||
let _ = std::process::Command::new("cargo").arg("build").arg("--bin").arg("mcp-memory-stub").status();
|
||||
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let local_addr = listener.local_addr().unwrap();
|
||||
let target = format!("http://127.0.0.1:{}", local_addr.port());
|
||||
|
||||
tokio::spawn(async move {
|
||||
while let Ok((mut socket, _)) = listener.accept().await {
|
||||
use tokio::io::AsyncReadExt;
|
||||
let mut buf = [0; 1024];
|
||||
let _ = socket.read(&mut buf).await;
|
||||
drop(socket);
|
||||
}
|
||||
});
|
||||
|
||||
let start = Instant::now();
|
||||
let mut child = Command::new(get_stub_exe())
|
||||
.arg("--target")
|
||||
.arg(target)
|
||||
.stdin(Stdio::null())
|
||||
.spawn()
|
||||
.expect("Failed to execute stub");
|
||||
|
||||
let res = tokio::time::timeout(Duration::from_secs(5), child.wait()).await;
|
||||
let elapsed = start.elapsed();
|
||||
|
||||
assert!(res.is_ok(), "Stub hung on fallback failure! Took {:?}", elapsed);
|
||||
}
|
||||
Reference in new issue
Block a user