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.
This commit is contained in:
1 parent
3533c243d2
commit
ffb316a6e3
12 files changed
+228
-45
No files matched your search
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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.");
|
||||
}
|
||||
@@ -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.");
|
||||
}
|
||||
@@ -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");
|
||||
}
|
||||
+23
-20
@@ -1,28 +1,22 @@
|
||||
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<HashMap<String, String>> = 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 config_path.exists() {
|
||||
if let Ok(config_str) = fs::read_to_string(&config_path) {
|
||||
if let Ok(parsed) = serde_json::from_str::<Value>(&config_str) {
|
||||
if let Some(env_obj) = parsed
|
||||
.get("mcpServers")
|
||||
.and_then(|v| v.get("mcp-memory"))
|
||||
@@ -30,14 +24,23 @@ pub fn load_mcp_config_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 let Some(val_str) = val.as_str() {
|
||||
map.insert(key.clone(), val_str.to_string());
|
||||
info!("Loaded env var {} from mcp_config.json", key);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let _ = CONFIG_ENV.set(map);
|
||||
}
|
||||
|
||||
pub fn get_env(key: &str) -> Result<String, env::VarError> {
|
||||
if let Some(map) = CONFIG_ENV.get() {
|
||||
if let Some(val) = map.get(key) {
|
||||
return Ok(val.clone());
|
||||
}
|
||||
}
|
||||
env::var(key)
|
||||
}
|
||||
+2
-2
@@ -47,7 +47,7 @@ pub fn init_redb(base: &Path) -> Arc<Database> {
|
||||
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<Database> {
|
||||
}
|
||||
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!(
|
||||
|
||||
@@ -14,7 +14,7 @@ struct CandleEmbeddingModel {
|
||||
|
||||
impl CandleEmbeddingModel {
|
||||
fn new() -> Result<Self, String> {
|
||||
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())?;
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -131,7 +131,7 @@ pub async fn start_background_indexer(state: Arc<MemoryState>) {
|
||||
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);
|
||||
|
||||
+8
-8
@@ -220,7 +220,7 @@ pub async fn condense_graph_worker(state: Arc<MemoryState>) {
|
||||
_ = 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<MemoryState>) -> Result<(), Box<dyn std::erro
|
||||
|
||||
let udp_state = Arc::clone(&app_state);
|
||||
tokio::spawn(async move {
|
||||
let port1 = std::env::var("MCP_UDP_PORT1").unwrap_or_else(|_| "3001".to_string());
|
||||
let port1 = crate::config::get_env("MCP_UDP_PORT1").unwrap_or_else(|_| "3001".to_string());
|
||||
if let Ok(socket) = tokio::net::UdpSocket::bind(format!("127.0.0.1:{}", port1)).await {
|
||||
let socket = Arc::new(socket);
|
||||
let socket_rx = socket.clone();
|
||||
@@ -470,7 +470,7 @@ pub async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::erro
|
||||
|
||||
let nvim_udp_state = Arc::clone(&app_state);
|
||||
tokio::spawn(async move {
|
||||
let port2 = std::env::var("MCP_UDP_PORT2").unwrap_or_else(|_| "3002".to_string());
|
||||
let port2 = crate::config::get_env("MCP_UDP_PORT2").unwrap_or_else(|_| "3002".to_string());
|
||||
if let Ok(socket) = tokio::net::UdpSocket::bind(format!("127.0.0.1:{}", port2)).await {
|
||||
let mut buf = vec![0u8; 65536];
|
||||
loop {
|
||||
@@ -586,7 +586,7 @@ pub async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::erro
|
||||
|
||||
let app = api::setup::create_router(app_state);
|
||||
|
||||
let port_str = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
|
||||
let port_str = crate::config::get_env("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
|
||||
tracing::info!(
|
||||
"MCP Memory Server running on http://127.0.0.1:{}/ws",
|
||||
port_str
|
||||
@@ -623,7 +623,7 @@ pub async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::erro
|
||||
}
|
||||
|
||||
pub fn init_logging(app_name: &str) -> Option<tracing_appender::non_blocking::WorkerGuard> {
|
||||
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<dyn std::error::Error>> {
|
||||
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<dyn std::error::Error>> {
|
||||
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<dyn std::error::Error>> {
|
||||
}
|
||||
|
||||
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 {
|
||||
|
||||
@@ -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::<u64>()
|
||||
.unwrap_or(60);
|
||||
|
||||
+1
-1
@@ -235,7 +235,7 @@ impl MemoryState {
|
||||
|
||||
pub fn modify_graph<F: FnOnce(&mut KnowledgeGraph)>(&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);
|
||||
|
||||
Reference in new issue
Block a user