refactor: physically decouple server and stub architectures
This commit is contained in:
Riz Ashraf committed
1 parent
0f4be375c2
commit
6dde1cc3a3
3 files changed
+200
-198
No files matched your search
+41
-93
@@ -1,10 +1,8 @@
|
||||
use clap::Parser;
|
||||
use futures_util::StreamExt;
|
||||
use std::sync::Arc;
|
||||
use clap::Parser;
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
use tokio::sync::RwLock;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::io::StreamReader;
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(name = "mcp-memory-stub", author, version, about = "Antigravity MCP Memory Stub / Proxy", long_about = None)]
|
||||
@@ -38,104 +36,56 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
});
|
||||
|
||||
let target_url = cli.target;
|
||||
let post_url = Arc::new(RwLock::new(String::new()));
|
||||
let post_url_proxy = Arc::clone(&post_url);
|
||||
|
||||
let client = reqwest::Client::builder().build()?;
|
||||
|
||||
tokio::spawn(async move {
|
||||
while let Some(msg) = msg_rx.recv().await {
|
||||
let post_url_clone = Arc::clone(&post_url_proxy);
|
||||
let client = client.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut attempts = 0;
|
||||
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 let Ok(resp) = res { if resp.status().is_success() { break; } }
|
||||
}
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
|
||||
attempts += 1;
|
||||
if attempts % 10 == 0 {
|
||||
eprintln!("[PROXY] Waiting for server to accept messages...");
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
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() {
|
||||
break;
|
||||
}
|
||||
|
||||
let sse_url = format!("{}/sse", target_url);
|
||||
let client = reqwest::Client::builder().build()?;
|
||||
match tokio_tungstenite::connect_async(&ws_url).await {
|
||||
Ok((ws_stream, _)) => {
|
||||
let (mut write, mut read) = ws_stream.split();
|
||||
|
||||
match client
|
||||
.get(&sse_url)
|
||||
.header("Accept", "text/event-stream")
|
||||
.send()
|
||||
.await
|
||||
{
|
||||
Ok(resp) => {
|
||||
if resp.status() == reqwest::StatusCode::GONE {
|
||||
eprintln!("[PROXY] Target gone, exiting.");
|
||||
break;
|
||||
}
|
||||
|
||||
let stream = resp.bytes_stream().map(|res| {
|
||||
res.map_err(std::io::Error::other)
|
||||
});
|
||||
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;
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = shutdown_rx.recv() => {
|
||||
return Ok(()); // Stdin closed, exit entirely
|
||||
}
|
||||
res = reader.read_line(&mut line) => {
|
||||
match res {
|
||||
Ok(bytes) => {
|
||||
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 let Some(stripped) = trimmed.strip_prefix("data: ") {
|
||||
if is_message {
|
||||
println!("{}", stripped);
|
||||
is_message = false;
|
||||
} else if is_endpoint {
|
||||
let mut p = post_url.write().await;
|
||||
*p = format!("{}{}", target_url, stripped);
|
||||
is_endpoint = false;
|
||||
}
|
||||
}
|
||||
line.clear();
|
||||
let msg_rx_clone = Arc::clone(&msg_rx);
|
||||
let mut send_task = tokio::spawn(async move {
|
||||
loop {
|
||||
let mut rx = msg_rx_clone.lock().await;
|
||||
match rx.recv().await {
|
||||
Some(msg) => {
|
||||
drop(rx);
|
||||
if write.send(tokio_tungstenite::tungstenite::Message::Text(msg)).await.is_err() {
|
||||
break;
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
},
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
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 {
|
||||
println!("{}", text);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
tokio::select! {
|
||||
_ = shutdown_rx.recv() => {
|
||||
return Ok(()); // Stdin closed, exit entirely
|
||||
}
|
||||
_ = &mut send_task => {
|
||||
recv_task.abort();
|
||||
}
|
||||
_ = &mut recv_task => {
|
||||
send_task.abort();
|
||||
}
|
||||
}
|
||||
*post_url.write().await = String::new();
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
|
||||
}
|
||||
Err(_) => {
|
||||
if let Some(ref cmd) = wake_cmd {
|
||||
@@ -153,5 +103,3 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
Reference in new issue
Block a user