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

fix(core): remove all unwraps, fix unpolled futures, and fix search index memory leak on restart

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-09-21 18:54:19 +01:00
1 parent ae0ea9ab22
commit 7dc0c329ad
9 files changed
+45 -59
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/main.rs
+19 -37
View File
@@ -206,7 +206,7 @@ async fn gate_set_handler(
reason: body.reason.clone(),
timestamp: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.unwrap_or_default()
.as_secs(),
};
app_state.handler.state.gates.modify(|gates| {
@@ -432,7 +432,7 @@ async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::E
tracing::info!("MCP Memory Server running on http://127.0.0.1:3000/sse");
let addr = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
let addr: std::net::SocketAddr = format!("127.0.0.1:{}", addr).parse().unwrap();
let addr: std::net::SocketAddr = format!("127.0.0.1:{}", addr).parse().expect("Invalid bind address");
let listener = match tokio::net::TcpListener::bind(&addr).await {
Ok(l) => l,
@@ -468,7 +468,7 @@ async fn ws_handler(
.into_response()
}
async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: String) {
async fn handle_socket(socket: WebSocket, state: Arc<AppState>, _client_type: String) {
let session_id = format!("{}", state.next_id.fetch_add(1, Ordering::SeqCst));
let (tx, mut rx) = mpsc::channel::<String>(100);
@@ -510,20 +510,6 @@ async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: Str
);
tracing::trace!("Message content: {}", text);
if let Ok(payload) = serde_json::from_str::<serde_json::Value>(&text) {
if client_type == "proxy" {
// Send activity broadcast to UI clients
if let Some(method) = payload.get("method").and_then(|m| m.as_str())
&& method == "tools/call"
{
let name = payload
.get("params")
.and_then(|p| p.get("name"))
.and_then(|n| n.as_str())
.unwrap_or("unknown_tool");
handler.state.broadcast_activity(&format!("Agent executed tool: {}", name));
}
} // End if proxy
// Process MCP request
if let Some(response) = handler.handle_request(payload).await {
let res_str = serde_json::to_string(&response).unwrap_or_default();
@@ -575,8 +561,8 @@ async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: Str
struct SessionCleanup {
session_id: String,
state: Arc<AppState>,
send_task: Option<tokio::task::JoinHandle<()>>,
recv_task: Option<tokio::task::JoinHandle<()>>,
send_task: tokio::task::JoinHandle<()>,
recv_task: tokio::task::JoinHandle<()>,
}
impl Drop for SessionCleanup {
@@ -586,12 +572,8 @@ async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: Str
.write()
.unwrap_or_else(|e| e.into_inner())
.remove(&self.session_id);
if let Some(task) = self.send_task.take() {
task.abort();
}
if let Some(task) = self.recv_task.take() {
task.abort();
}
self.send_task.abort();
self.recv_task.abort();
tracing::info!(
"Websocket session {} closed and cleaned up",
self.session_id
@@ -602,15 +584,15 @@ async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: Str
let mut cleanup = SessionCleanup {
session_id: session_id.clone(),
state: Arc::clone(&state),
send_task: Some(send_task),
recv_task: Some(recv_task),
send_task,
recv_task,
};
tokio::select! {
_ = cleanup.send_task.as_mut().unwrap() => {
_ = &mut cleanup.send_task => {
tracing::info!("Websocket send task finished for session {}", session_id);
},
_ = cleanup.recv_task.as_mut().unwrap() => {
_ = &mut cleanup.recv_task => {
tracing::info!("Websocket recv task finished for session {}", session_id);
},
};
@@ -746,7 +728,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
if !cli.daemon {
// Just spawn the daemon and exit. We no longer act as a proxy.
#[allow(clippy::zombie_processes)]
let _ = std::process::Command::new(std::env::current_exe().unwrap())
let _ = std::process::Command::new(std::env::current_exe().expect("Failed to get current executable path"))
.arg("--daemon")
.stdin(std::process::Stdio::null())
.stdout(std::process::Stdio::null())
@@ -766,13 +748,13 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
let redb_path = base.join("mcp_store.redb");
let db = Arc::new(redb::Database::create(&redb_path).unwrap());
let db = Arc::new(redb::Database::create(&redb_path).expect("Failed to create redb database"));
// Ensure table exists and migrate old JSON files
{
let write_txn = db.begin_write().unwrap();
let write_txn = db.begin_write().expect("Failed to begin write txn on redb");
{
let mut table = write_txn.open_table(crate::store::STORE_TABLE).unwrap();
let mut table = write_txn.open_table(crate::store::STORE_TABLE).expect("Failed to open STORE_TABLE");
let stores = vec![
("knowledge_graph_master", "knowledge_graph_master.json"),
@@ -797,22 +779,22 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
];
for (key, file_name) in stores.iter() {
if table.get(*key).unwrap().is_none() {
if table.get(*key).expect("Failed to read from table").is_none() {
let json_path = base.join(file_name);
if json_path.exists()
&& let Ok(data) = fs::read(&json_path)
&& serde_json::from_slice::<serde_json::Value>(&data).is_ok()
{
table.insert(*key, data.as_slice()).unwrap();
table.insert(*key, data.as_slice()).expect("Failed to insert migrated data");
let _ = fs::rename(&json_path, json_path.with_extension("json.migrated"));
}
}
}
}
write_txn.commit().unwrap();
write_txn.commit().expect("Failed to commit db migration");
}
let rt = tokio::runtime::Runtime::new().unwrap();
let rt = tokio::runtime::Runtime::new().expect("Failed to create tokio runtime");
let _guard = rt.enter();
let state = Arc::new(MemoryState {