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

refactor: apply zero-unwrap policy and optimize locks in store.rs and handlers

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-09-21 11:34:21 +01:00
1 parent 9f24e66d88
commit 8afbf97b11
33 files changed
+3127 -3207
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

stub/src/main.rs
+31 -10
View File
@@ -10,8 +10,6 @@ struct Cli {
target: String,
}
mod logger;
fn init_logging(app_name: &str) -> Option<tracing_appender::non_blocking::WorkerGuard> {
@@ -53,7 +51,9 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
let port = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
format!("http://127.0.0.1:{}", port)
};
let ws_url = target_url.replace("http://", "ws://").replace("https://", "wss://");
let ws_url = target_url
.replace("http://", "ws://")
.replace("https://", "wss://");
let ws_url = format!("{}/ws?client=proxy", ws_url);
loop {
@@ -63,7 +63,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
}
tracing::info!("Attempting to connect to {}", ws_url);
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
let request = match ws_url.clone().into_client_request() {
Ok(req) => req,
@@ -83,8 +83,21 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
let mut send_task = tokio::spawn(async move {
while let Ok(msg) = rx.recv().await {
let log_prefix = logger::extract_log_prefix(&msg, false);
tracing::info!(">>> [Stub] Forwarding {} to server (length: {}): {}", log_prefix, 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::info!(
">>> [Stub] Forwarding {} to server (length: {}): {}",
log_prefix,
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;
}
@@ -95,7 +108,16 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
while let Some(Ok(msg)) = read.next().await {
if let tokio_tungstenite::tungstenite::Message::Text(text) = msg {
let log_prefix = logger::extract_log_prefix(&text, true);
tracing::info!("<<< [Stub] Received {} from server (length: {}): {}", log_prefix, text.len(), if text.len() > 1000 { format!("{}...", &text[..1000]) } else { text.clone() });
tracing::info!(
"<<< [Stub] Received {} from server (length: {}): {}",
log_prefix,
text.len(),
if text.len() > 1000 {
format!("{}...", &text[..1000])
} else {
text.clone()
}
);
use tokio::io::AsyncWriteExt;
let mut stdout = tokio::io::stdout();
let _ = stdout.write_all(text.as_bytes()).await;
@@ -107,14 +129,14 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
});
tokio::select! {
_ = shutdown_rx.recv() => {
_ = shutdown_rx.recv() => {
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
tracing::info!("Shutdown received while connected");
break;
}
_ = &mut send_task => {
tracing::error!("Send task exited");
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
recv_task.abort();
break;
}
@@ -134,4 +156,3 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
Ok(())
})
}