riz/mcp-memory
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Add resilient MPSC queues to proxy stubs to fix thread leaks and message loss on reconnect

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-09-09 04:32:07 +01:00
1 parent 51fc3692c8
commit 05e84e98bd
2 files changed
+169 -198
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/proxy.rs
+79 -83
View File
@@ -3,105 +3,101 @@ use tokio::io::AsyncBufReadExt;
use futures_util::StreamExt;
use std::sync::Arc;
use tokio::sync::RwLock;
use tokio::sync::mpsc;
pub fn run_proxy(target_url: &str) -> Result<bool, Box<dyn std::error::Error + Send + Sync>> {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
let client = reqwest::Client::builder().build().unwrap();
let sse_url = format!("{}/sse", target_url);
eprintln!("[PROXY] Connecting to Streamable HTTP: {}", sse_url);
let resp = match client.get(&sse_url)
.header("Accept", "text/event-stream")
.send().await {
Ok(r) => r,
Err(e) => {
eprintln!("[PROXY] Connection failed: {}", e);
return Ok(true);
}
};
let (msg_tx, mut msg_rx) = mpsc::channel::<String>(100);
let (shutdown_tx, mut shutdown_rx) = mpsc::channel::<()>(1);
if resp.status() == reqwest::StatusCode::GONE {
eprintln!("[PROXY] Streamable HTTP endpoint is gone.");
return Ok(true);
}
let post_url = Arc::new(RwLock::new(format!("{}/messages", target_url))); // Will be updated by endpoint event
let post_url_clone = Arc::clone(&post_url);
let (tx, mut rx) = tokio::sync::mpsc::channel(1);
let tx_clone = tx.clone();
let client_clone = client.clone();
tokio::task::spawn_blocking(move || {
let stdin = std::io::stdin();
let mut handle = stdin.lock();
let mut buffer = String::new();
while let Ok(bytes) = std::io::BufRead::read_line(&mut handle, &mut buffer) {
if bytes == 0 { break; }
let body = buffer.clone();
let _ = msg_tx.blocking_send(buffer.clone());
buffer.clear();
let client = client_clone.clone();
let url_arc = Arc::clone(&post_url_clone);
tokio::spawn(async move {
let mut url = "".to_string();
for _ in 0..50 {
let u = url_arc.read().await.clone();
if u.contains("sessionId") {
url = u;
break;
}
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
}
if url.is_empty() {
url = url_arc.read().await.clone();
}
let _ = client.post(&url)
.header("Accept", "application/json, text/event-stream")
.header("Content-Type", "application/json")
.body(body)
.send().await;
});
}
let _ = tx_clone.blocking_send(false); // Stdin EOF
let _ = shutdown_tx.blocking_send(());
});
let target_url = target_url.to_string();
let tx_clone2 = tx.clone();
tokio::spawn(async move {
let stream = resp.bytes_stream().map(|res| res.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e)));
let mut reader = tokio::io::BufReader::new(StreamReader::new(stream));
let mut line = String::new();
let mut is_message = false;
let mut is_endpoint = false;
while let Ok(bytes) = reader.read_line(&mut line).await {
if bytes == 0 { break; }
let trimmed = line.trim();
if trimmed.starts_with("event: message") {
is_message = true;
is_endpoint = false;
} else if trimmed.starts_with("event: endpoint") {
is_endpoint = true;
is_message = false;
} else if trimmed.starts_with("data: ") {
if is_message {
println!("{}", &trimmed[6..]);
is_message = false;
} else if is_endpoint {
let ep = &trimmed[6..];
let mut p = post_url.write().await;
*p = format!("{}{}", target_url, ep);
is_endpoint = false;
}
}
line.clear();
}
let _ = tx_clone2.send(true).await; // Stream dropped (Leader dead)
});
let post_url = Arc::new(RwLock::new(String::new()));
let post_url_clone = Arc::clone(&post_url);
let dropped = rx.recv().await.unwrap_or(true);
Ok(dropped)
let client = reqwest::Client::builder().build().unwrap();
tokio::spawn(async move {
while let Some(msg) = msg_rx.recv().await {
loop {
let url = post_url_clone.read().await.clone();
if !url.is_empty() {
let res = client.post(&url)
.header("Accept", "application/json, text/event-stream")
.header("Content-Type", "application/json")
.body(msg.clone())
.send().await;
if res.is_ok() {
break;
}
}
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
}
}
});
loop {
if shutdown_rx.try_recv().is_ok() {
return Ok(false);
}
let sse_url = format!("{}/sse", target_url);
let client = reqwest::Client::builder().build().unwrap();
match client.get(&sse_url).header("Accept", "text/event-stream").send().await {
Ok(resp) => {
if resp.status() == reqwest::StatusCode::GONE {
return Ok(true);
}
let stream = resp.bytes_stream().map(|res| res.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e)));
let mut reader = tokio::io::BufReader::new(StreamReader::new(stream));
let mut line = String::new();
let mut is_message = false;
let mut is_endpoint = false;
while let Ok(bytes) = reader.read_line(&mut line).await {
if bytes == 0 { break; }
let trimmed = line.trim();
if trimmed.starts_with("event: message") {
is_message = true;
is_endpoint = false;
} else if trimmed.starts_with("event: endpoint") {
is_endpoint = true;
is_message = false;
} else if trimmed.starts_with("data: ") {
if is_message {
println!("{}", &trimmed[6..]);
is_message = false;
} else if is_endpoint {
let ep = &trimmed[6..];
let mut p = post_url.write().await;
*p = format!("{}{}", target_url, ep);
is_endpoint = false;
}
}
line.clear();
}
*post_url.write().await = String::new();
return Ok(true);
}
Err(_) => {
return Ok(true);
}
}
}
})
}