chore: migrate store to redb
This commit is contained in:
1 parent
7f0286dc40
commit
9a8b6e52e9
12 files changed
+1954
-1365
No files matched your search
+26
-16
@@ -3,8 +3,8 @@ use futures_util::StreamExt;
|
||||
use std::sync::Arc;
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
use tokio::sync::RwLock;
|
||||
use tokio_util::io::StreamReader;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::io::StreamReader;
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(name = "mcp-memory-stub")]
|
||||
@@ -27,7 +27,9 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
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; }
|
||||
if bytes == 0 {
|
||||
break;
|
||||
}
|
||||
let _ = msg_tx.blocking_send(buffer.clone());
|
||||
buffer.clear();
|
||||
}
|
||||
@@ -37,7 +39,7 @@ 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_clone = Arc::clone(&post_url);
|
||||
|
||||
|
||||
let client = reqwest::Client::builder().build()?;
|
||||
|
||||
tokio::spawn(async move {
|
||||
@@ -46,14 +48,16 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
loop {
|
||||
let url = post_url_clone.read().await.clone();
|
||||
if !url.is_empty() {
|
||||
let res = client.post(&url)
|
||||
let res = client
|
||||
.post(&url)
|
||||
.header("Accept", "application/json, text/event-stream")
|
||||
.header("Content-Type", "application/json")
|
||||
.body(msg.clone())
|
||||
.send().await;
|
||||
|
||||
.send()
|
||||
.await;
|
||||
|
||||
if res.is_ok() {
|
||||
break;
|
||||
break;
|
||||
}
|
||||
}
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
|
||||
@@ -73,20 +77,27 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
|
||||
let sse_url = format!("{}/sse", target_url);
|
||||
let client = reqwest::Client::builder().build()?;
|
||||
|
||||
match client.get(&sse_url).header("Accept", "text/event-stream").send().await {
|
||||
|
||||
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(|e| std::io::Error::new(std::io::ErrorKind::Other, e)));
|
||||
|
||||
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() => {
|
||||
@@ -103,14 +114,13 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
} else if trimmed.starts_with("event: endpoint") {
|
||||
is_endpoint = true;
|
||||
is_message = false;
|
||||
} else if trimmed.starts_with("data: ") {
|
||||
} else if let Some(stripped) = trimmed.strip_prefix("data: ") {
|
||||
if is_message {
|
||||
println!("{}", &trimmed[6..]);
|
||||
println!("{}", stripped);
|
||||
is_message = false;
|
||||
} else if is_endpoint {
|
||||
let ep = &trimmed[6..];
|
||||
let mut p = post_url.write().await;
|
||||
*p = format!("{}{}", target_url, ep);
|
||||
*p = format!("{}{}", target_url, stripped);
|
||||
is_endpoint = false;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user