From ffb316a6e39a4fecd83f6b7a181393163b4ea6fb Mon Sep 17 00:00:00 2001 From: Riz Ashraf Date: Sun, 11 Oct 2026 06:56:24 +0100 Subject: [PATCH] fix: replace unsafe env::set_var with thread-safe OnceLock Resolves intermittent STATUS_HEAP_CORRUPTION on Windows by using a static OnceLock map instead of mutating the global OS environment block. --- server/src/api/telemetry.rs | 2 +- server/src/bin/test_embedding_crash.rs | 19 ++++ server/src/bin/test_server_shutdown.rs | 115 +++++++++++++++++++++++++ server/src/bin/test_tokenizer_crash.rs | 46 ++++++++++ server/src/config.rs | 57 ++++++------ server/src/db.rs | 4 +- server/src/embedding.rs | 2 +- server/src/handlers/env.rs | 8 +- server/src/indexer.rs | 2 +- server/src/lib.rs | 16 ++-- server/src/ollama.rs | 14 +-- server/src/state.rs | 2 +- 12 files changed, 235 insertions(+), 52 deletions(-) create mode 100644 server/src/bin/test_embedding_crash.rs create mode 100644 server/src/bin/test_server_shutdown.rs create mode 100644 server/src/bin/test_tokenizer_crash.rs diff --git a/server/src/api/telemetry.rs b/server/src/api/telemetry.rs index ba4df9c..8058b6a 100644 --- a/server/src/api/telemetry.rs +++ b/server/src/api/telemetry.rs @@ -146,7 +146,7 @@ pub async fn nvim_telemetry_handler( let _ = tokio::fs::write(wsl_path, session).await; } else { let profile = - std::env::var("USERPROFILE").unwrap_or_else(|_| "C:\\Users\\reazul.ashraf".into()); + crate::config::get_env("USERPROFILE").unwrap_or_else(|_| "C:\\Users\\reazul.ashraf".into()); let win_path = format!("{}\\.gemini\\active_nvim.txt", profile); let _ = tokio::fs::write(&win_path, session).await; } diff --git a/server/src/bin/test_embedding_crash.rs b/server/src/bin/test_embedding_crash.rs new file mode 100644 index 0000000..0a7f412 --- /dev/null +++ b/server/src/bin/test_embedding_crash.rs @@ -0,0 +1,19 @@ +use mcp_memory_server::embedding::generate_embedding_async; + +#[tokio::main] +async fn main() { + println!("Starting embedding test..."); + + for i in 0..100 { + match generate_embedding_async("Test embedding initialization".to_string()).await { + Ok(_) => { + println!("Iteration {} successful", i); + } + Err(e) => { + println!("Iteration {} failed: {}", i, e); + } + } + } + + println!("Embedding test completed successfully."); +} diff --git a/server/src/bin/test_server_shutdown.rs b/server/src/bin/test_server_shutdown.rs new file mode 100644 index 0000000..79773d1 --- /dev/null +++ b/server/src/bin/test_server_shutdown.rs @@ -0,0 +1,115 @@ +use std::process::Command; +use std::time::Duration; +use std::net::TcpListener; +use reqwest::Client; +use uuid::Uuid; + +fn get_free_port() -> u16 { + TcpListener::bind("127.0.0.1:0") + .unwrap() + .local_addr() + .unwrap() + .port() +} + +#[tokio::main] +async fn main() { + println!("Starting server shutdown test..."); + + let exe_path = if cfg!(windows) { + "target/debug/mcp-memory-server.exe" + } else { + "target/debug/mcp-memory-server" + }; + + let client = Client::new(); + + for i in 0..10 { + println!("Iteration {}...", i); + + let port = get_free_port(); + let udp1 = get_free_port(); + let udp2 = get_free_port(); + + // Create a unique temporary directory for this iteration + let temp_dir = std::env::temp_dir().join(format!("mcp_test_{}", Uuid::new_v4())); + std::fs::create_dir_all(&temp_dir).unwrap(); + + let mut child = Command::new(exe_path) + .env("MCP_PORT", port.to_string()) + .env("MCP_UDP_PORT1", udp1.to_string()) + .env("MCP_UDP_PORT2", udp2.to_string()) + .env("MCP_MEMORY_STORE_DIR", temp_dir.to_str().unwrap()) + .env("MCP_OFFLINE_EMBEDDINGS", "1") // use fallback embeddings to speed up + .spawn() + .expect("Failed to start server"); + + // Poll until the server is responsive + let mut is_ready = false; + for _ in 0..20 { + if let Ok(res) = client.get(format!("http://127.0.0.1:{}/ping", port)).send().await { + if res.status().is_success() { + is_ready = true; + break; + } + } + tokio::time::sleep(Duration::from_millis(500)).await; + } + + if !is_ready { + println!("Server failed to start within 10 seconds!"); + let _ = child.kill(); + let _ = child.wait(); + continue; + } + + let token = std::fs::read_to_string(temp_dir.join("admin.token")).unwrap_or_default(); + + match client + .post(format!("http://127.0.0.1:{}/shutdown", port)) + .header("Authorization", format!("Bearer {}", token.trim())) + .send() + .await + { + Ok(res) => { + if !res.status().is_success() { + println!("Failed to send shutdown request: {}", res.status()); + let _ = child.kill(); + } + } + Err(e) => { + println!("Error sending shutdown request: {}", e); + let _ = child.kill(); + } + } + + // Wait for child to exit gracefully + let mut exited = false; + for _ in 0..10 { + match child.try_wait() { + Ok(Some(status)) => { + println!("Server exited with status: {}", status); + exited = true; + break; + } + Ok(None) => { + tokio::time::sleep(Duration::from_millis(500)).await; + } + Err(e) => { + println!("Error waiting for server: {}", e); + break; + } + } + } + + if !exited { + println!("Server didn't exit gracefully, killing it..."); + let _ = child.kill(); + let _ = child.wait(); + } + + std::fs::remove_dir_all(&temp_dir).unwrap_or_default(); + } + + println!("Server shutdown test completed."); +} diff --git a/server/src/bin/test_tokenizer_crash.rs b/server/src/bin/test_tokenizer_crash.rs new file mode 100644 index 0000000..b01bd56 --- /dev/null +++ b/server/src/bin/test_tokenizer_crash.rs @@ -0,0 +1,46 @@ +#[cfg(target_os = "windows")] +fn main() { + println!("Starting Tokenizer isolation test"); + + let client = hf_hub::HFClientSync::new().unwrap(); + let repo = client.model("sentence-transformers", "all-MiniLM-L6-v2"); + + println!("Downloading tokenizer.json..."); + let tokenizer_file = repo + .download_file() + .filename("tokenizer.json") + .send() + .expect("Failed to download tokenizer.json"); + + println!("Loading Tokenizer..."); + let mut tokenizer = tokenizers::Tokenizer::from_file(&tokenizer_file) + .expect("Failed to load tokenizer"); + + tokenizer.with_padding(Some(tokenizers::PaddingParams::default())); + + println!("Encoding sample text..."); + let texts = vec!["Hello, this is a test string for the tokenizer.".to_string(); 100]; + + // Attempt encoding + let encodings = tokenizer + .encode_batch(texts, true) + .expect("Failed to encode batch"); + + println!("Encoded {} items.", encodings.len()); + + // Force a few iterations to see if repeated allocation/deallocation triggers the heap corruption + for i in 0..100 { + let texts_batch = vec![format!("Batch item {}", i); 50]; + let _ = tokenizer.encode_batch(texts_batch, true).unwrap(); + } + + // Drop explicitly to test destruction phase + drop(tokenizer); + + println!("Tokenizer isolation test completed without crash."); +} + +#[cfg(not(target_os = "windows"))] +fn main() { + println!("Not on windows, skipping isolation test"); +} diff --git a/server/src/config.rs b/server/src/config.rs index c9ad99f..3e56af4 100644 --- a/server/src/config.rs +++ b/server/src/config.rs @@ -1,43 +1,46 @@ +use std::sync::OnceLock; use serde_json::Value; +use std::collections::HashMap; use std::env; use std::fs; use std::path::PathBuf; use tracing::info; +pub static CONFIG_ENV: OnceLock> = OnceLock::new(); + pub fn load_mcp_config_env() { + let mut map = HashMap::new(); let config_path = dirs::home_dir() .unwrap_or_else(|| PathBuf::from(".")) .join(".gemini/config/mcp_config.json"); - if !config_path.exists() { - return; - } - - let config_str = match fs::read_to_string(&config_path) { - Ok(s) => s, - Err(_) => return, - }; - - let parsed: Value = match serde_json::from_str(&config_str) { - Ok(v) => v, - Err(_) => return, - }; - - if let Some(env_obj) = parsed - .get("mcpServers") - .and_then(|v| v.get("mcp-memory")) - .and_then(|v| v.get("env")) - .and_then(|v| v.as_object()) - { - for (key, val) in env_obj { - if let Some(val_str) = val.as_str() - && env::var(key).is_err() - { - unsafe { - env::set_var(key, val_str); + if config_path.exists() { + if let Ok(config_str) = fs::read_to_string(&config_path) { + if let Ok(parsed) = serde_json::from_str::(&config_str) { + if let Some(env_obj) = parsed + .get("mcpServers") + .and_then(|v| v.get("mcp-memory")) + .and_then(|v| v.get("env")) + .and_then(|v| v.as_object()) + { + for (key, val) in env_obj { + if let Some(val_str) = val.as_str() { + map.insert(key.clone(), val_str.to_string()); + info!("Loaded env var {} from mcp_config.json", key); + } + } } - info!("Loaded env var {} from mcp_config.json", key); } } } + let _ = CONFIG_ENV.set(map); +} + +pub fn get_env(key: &str) -> Result { + if let Some(map) = CONFIG_ENV.get() { + if let Some(val) = map.get(key) { + return Ok(val.clone()); + } + } + env::var(key) } diff --git a/server/src/db.rs b/server/src/db.rs index d909007..6be3907 100644 --- a/server/src/db.rs +++ b/server/src/db.rs @@ -47,7 +47,7 @@ pub fn init_redb(base: &Path) -> Arc { match redb::Database::create(&redb_path) { Ok(db) => Arc::new(db), Err(create_err) => { - if std::env::var("MCP_ALLOW_TMP_FALLBACK").unwrap_or_default() == "1" { + if crate::config::get_env("MCP_ALLOW_TMP_FALLBACK").unwrap_or_default() == "1" { let temp_path = std::env::temp_dir() .join(format!("mcp_store_fallback_{}.redb", uuid::Uuid::new_v4())); tracing::warn!( @@ -95,7 +95,7 @@ pub fn init_redb(base: &Path) -> Arc { } if let Some(db) = db_opt { db - } else if std::env::var("MCP_ALLOW_TMP_FALLBACK").unwrap_or_default() == "1" { + } else if crate::config::get_env("MCP_ALLOW_TMP_FALLBACK").unwrap_or_default() == "1" { let temp_path = std::env::temp_dir() .join(format!("mcp_store_fallback_{}.redb", uuid::Uuid::new_v4())); tracing::warn!( diff --git a/server/src/embedding.rs b/server/src/embedding.rs index 7a20a69..673f0c7 100644 --- a/server/src/embedding.rs +++ b/server/src/embedding.rs @@ -14,7 +14,7 @@ struct CandleEmbeddingModel { impl CandleEmbeddingModel { fn new() -> Result { - if cfg!(test) || std::env::var("MCP_OFFLINE_EMBEDDINGS").is_ok() { + if cfg!(test) || crate::config::get_env("MCP_OFFLINE_EMBEDDINGS").is_ok() { return Err("Offline embedding mode active (tests/offline flag)".to_string()); } let client = hf_hub::HFClientSync::new().map_err(|e| e.to_string())?; diff --git a/server/src/handlers/env.rs b/server/src/handlers/env.rs index 91c0507..3a85aea 100644 --- a/server/src/handlers/env.rs +++ b/server/src/handlers/env.rs @@ -31,17 +31,17 @@ impl McpTool for EnvironmentHandler { let detected_os = req.os.unwrap_or_else(|| std::env::consts::OS.to_string()); let detected_shell = req .shell - .or_else(|| std::env::var("SHELL").ok()) - .or_else(|| std::env::var("STARSHIP_SHELL").ok()) + .or_else(|| crate::config::get_env("SHELL").ok()) + .or_else(|| crate::config::get_env("STARSHIP_SHELL").ok()) .or_else(|| { - if std::env::var("PSModulePath").is_ok() { + if crate::config::get_env("PSModulePath").is_ok() { Some("pwsh".to_string()) } else { None } }) .or_else(|| { - if std::env::var("COMSPEC").is_ok() { + if crate::config::get_env("COMSPEC").is_ok() { Some("cmd".to_string()) } else { None diff --git a/server/src/indexer.rs b/server/src/indexer.rs index 5e803ff..48d9012 100644 --- a/server/src/indexer.rs +++ b/server/src/indexer.rs @@ -131,7 +131,7 @@ pub async fn start_background_indexer(state: Arc) { snippets.drain(0..overflow); } }); - let threshold: usize = std::env::var("MCP_MEMORY_CONDENSE_THRESHOLD") + let threshold: usize = crate::config::get_env("MCP_MEMORY_CONDENSE_THRESHOLD") .unwrap_or_else(|_| "100".to_string()) .parse() .unwrap_or(100); diff --git a/server/src/lib.rs b/server/src/lib.rs index 04298e9..3193658 100644 --- a/server/src/lib.rs +++ b/server/src/lib.rs @@ -220,7 +220,7 @@ pub async fn condense_graph_worker(state: Arc) { _ = state.condense_notify.notified() => {}, } - let threshold: usize = std::env::var("MCP_MEMORY_CONDENSE_THRESHOLD") + let threshold: usize = crate::config::get_env("MCP_MEMORY_CONDENSE_THRESHOLD") .unwrap_or_else(|_| "100".to_string()) .parse() .unwrap_or(100); @@ -402,7 +402,7 @@ pub async fn run_server(state: Arc) -> Result<(), Box) -> Result<(), Box) -> Result<(), Box) -> Result<(), Box Option { - let base_dir = std::env::var("MCP_MEMORY_STORE_DIR").unwrap_or_else(|_| { + let base_dir = crate::config::get_env("MCP_MEMORY_STORE_DIR").unwrap_or_else(|_| { dirs::home_dir() .map(|mut h| { h.push(".gemini/mcp_memory"); @@ -653,7 +653,7 @@ pub fn run_cli() -> Result<(), Box> { let _guard = init_logging("mcp-memory-server"); let cli = Cli::parse(); - let base_dir = std::env::var("MCP_MEMORY_STORE_DIR").unwrap_or_else(|_| { + let base_dir = crate::config::get_env("MCP_MEMORY_STORE_DIR").unwrap_or_else(|_| { dirs::home_dir() .map(|mut h| { h.push(".gemini/mcp_memory"); @@ -664,7 +664,7 @@ pub fn run_cli() -> Result<(), Box> { let base = PathBuf::from(base_dir); if cli.exit { - let port = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string()); + let port = crate::config::get_env("MCP_PORT").unwrap_or_else(|_| "3000".to_string()); let token = std::fs::read_to_string(base.join("admin.token")).unwrap_or_default(); let rt = tokio::runtime::Runtime::new()?; rt.block_on(async { @@ -684,7 +684,7 @@ pub fn run_cli() -> Result<(), Box> { } if let Some(Commands::Gate { subcmd }) = cli.command { - let port = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string()); + let port = crate::config::get_env("MCP_PORT").unwrap_or_else(|_| "3000".to_string()); let rt = tokio::runtime::Runtime::new()?; match subcmd { GateCommands::Set { diff --git a/server/src/ollama.rs b/server/src/ollama.rs index 13496c3..5975266 100644 --- a/server/src/ollama.rs +++ b/server/src/ollama.rs @@ -1,6 +1,6 @@ use crate::error::AppError; use serde::{Deserialize, Serialize}; -use std::env; + use std::time::Duration; use tracing::debug; @@ -52,17 +52,17 @@ struct EmbeddingResponse { impl OllamaClient { pub fn new_from_env() -> Self { let base_url = - env::var("OLLAMA_URL").unwrap_or_else(|_| "http://192.168.1.30:11434".to_string()); + crate::config::get_env("OLLAMA_URL").unwrap_or_else(|_| "http://192.168.1.30:11434".to_string()); let coder_model = - env::var("OLLAMA_CODER_MODEL").unwrap_or_else(|_| "qwen2.5-coder:1.5b".to_string()); + crate::config::get_env("OLLAMA_CODER_MODEL").unwrap_or_else(|_| "qwen2.5-coder:1.5b".to_string()); let reasoning_model = - env::var("OLLAMA_REASONING_MODEL").unwrap_or_else(|_| "deepseek-r1:1.5b".to_string()); + crate::config::get_env("OLLAMA_REASONING_MODEL").unwrap_or_else(|_| "deepseek-r1:1.5b".to_string()); let vision_model = - env::var("OLLAMA_VISION_MODEL").unwrap_or_else(|_| "moondream".to_string()); - let embed_model = env::var("OLLAMA_EMBED_MODEL") + crate::config::get_env("OLLAMA_VISION_MODEL").unwrap_or_else(|_| "moondream".to_string()); + let embed_model = crate::config::get_env("OLLAMA_EMBED_MODEL") .unwrap_or_else(|_| "nomic-embed-text:latest".to_string()); - let timeout_sec = env::var("OLLAMA_TIMEOUT_SEC") + let timeout_sec = crate::config::get_env("OLLAMA_TIMEOUT_SEC") .unwrap_or_else(|_| "60".to_string()) .parse::() .unwrap_or(60); diff --git a/server/src/state.rs b/server/src/state.rs index d11e878..687a134 100644 --- a/server/src/state.rs +++ b/server/src/state.rs @@ -235,7 +235,7 @@ impl MemoryState { pub fn modify_graph(&self, update_fn: F) { self.graph.modify(update_fn); - let threshold: usize = std::env::var("MCP_MEMORY_CONDENSE_THRESHOLD") + let threshold: usize = crate::config::get_env("MCP_MEMORY_CONDENSE_THRESHOLD") .unwrap_or_else(|_| "100".to_string()) .parse() .unwrap_or(100);