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

refactor: consolidate nvim crates, extract server library, and update workspace dependencies

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-10-04 01:42:59 +01:00
1 parent a083719cf1
commit 533adfd41b
53 files changed
+5967 -1230
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/main.rs
+1 -584
View File
@@ -1,586 +1,3 @@
#![cfg_attr(
not(target_os = "windows"),
allow(dead_code, unused_imports, unreachable_code)
)]
mod api;
mod clipboard_watcher;
pub mod db;
pub mod embedding;
pub mod error;
mod handlers;
pub mod indexer;
mod mcp;
mod models;
pub mod ollama;
mod router;
mod search;
mod state;
mod store;
mod tools;
pub mod vector_db;
mod watcher;
use crate::api::rest::GateSetReq;
use crate::router::MemoryHandler;
use crate::state::MemoryState;
use clap::{Parser, Subcommand};
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::atomic::AtomicUsize;
use std::sync::{Arc, RwLock};
use std::time::Duration;
use tokio::sync::mpsc;
#[derive(Parser)]
#[command(author, version = env!("APP_VERSION"), about = "Antigravity MCP Memory Server", long_about = None)]
struct Cli {
#[command(subcommand)]
command: Option<Commands>,
/// Target URL for the proxy to connect to (e.g., http://127.0.0.1:3000)
#[arg(long)]
target: Option<String>,
/// Send a shutdown request to the currently running server
#[arg(long)]
exit: bool,
/// Send a shutdown request to the existing server and wait for it to exit
#[arg(long)]
restart: bool,
}
#[derive(Subcommand)]
enum Commands {
/// Manage authorization gates and verification for actions
Gate {
#[command(subcommand)]
subcmd: GateCommands,
},
}
#[derive(Subcommand)]
enum GateCommands {
Set {
#[arg(long)]
action: String,
#[arg(long)]
target: String,
#[arg(long)]
namespace: Option<String>,
#[arg(short = 'p', long = "param")]
params: Vec<String>,
#[arg(long, conflicts_with = "block")]
authorize: bool,
#[arg(long, conflicts_with = "authorize")]
block: bool,
#[arg(long)]
reason: Option<String>,
},
Verify {
#[arg(long)]
action: String,
#[arg(long)]
target: String,
#[arg(long)]
namespace: Option<String>,
#[arg(short = 'p', long = "param")]
params: Vec<String>,
#[arg(long)]
consume: bool,
},
}
pub struct AppState {
handler: Arc<MemoryHandler>,
clients: RwLock<HashMap<String, mpsc::Sender<String>>>,
next_id: AtomicUsize,
pub shutdown_tx: std::sync::Mutex<Option<tokio::sync::oneshot::Sender<()>>>,
}
async fn ttl_sweeper_worker(state: Arc<MemoryState>) {
loop {
tokio::time::sleep(Duration::from_secs(3600)).await;
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
state.project.tasks.modify(|tasks| {
tasks.retain(|t| t.expires_at.is_none_or(|exp| exp > now));
});
state.code.sticky.modify(|notes| {
notes.retain(|n| n.expires_at.is_none_or(|exp| exp > now));
});
state.telemetry.handoff_memos.modify(|memos| {
memos.retain(|m| m.expires_at.is_none_or(|exp| exp > now));
});
state.telemetry.session_summaries.modify(|summaries| {
summaries.retain(|s| s.expires_at.is_none_or(|exp| exp > now));
});
}
}
async fn index_committer_worker(state: Arc<MemoryState>) {
loop {
tokio::time::sleep(Duration::from_secs(5)).await;
// Periodically commit the search index to persist inline indexing operations
let idx_opt = state.search_index.read().ok().map(|idx| idx.clone());
if let Some(idx) = idx_opt {
let _ = idx.commit().await;
}
}
}
async fn condense_graph_worker(state: Arc<MemoryState>) {
loop {
tokio::time::sleep(Duration::from_secs(3600)).await;
let threshold: usize = std::env::var("MCP_MEMORY_CONDENSE_THRESHOLD")
.unwrap_or_else(|_| "100".to_string())
.parse()
.unwrap_or(100);
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
// Condense sticky notes
let mut condensed_sticky_content = String::new();
state.code.sticky.modify(|notes| {
if notes.len() > threshold {
notes.sort_by_key(|n| n.timestamp);
let to_remove = notes.len() - (threshold / 2);
let removed: Vec<_> = notes.drain(0..to_remove).collect();
for r in removed {
condensed_sticky_content.push_str(&format!("{}\n", r.content));
}
}
});
if !condensed_sticky_content.is_empty() {
state.modify_graph(|graph| {
let name = format!("StickyNote History {}", now);
graph.entities.insert(
name.clone(),
crate::models::Entity {
name: name.clone(),
entity_type: "Historical Summary".to_string(),
observations: vec![condensed_sticky_content],
namespace: crate::models::default_namespace(),
git_branch: None,
},
);
});
tracing::info!("Condensed sticky notes into Historical Summary.");
}
// Condense snippets
let mut condensed_snippet_content = String::new();
state.code.snippets.modify(|snippets| {
if snippets.len() > threshold {
snippets.sort_by_key(|s| s.updated_at);
let to_remove = snippets.len() - (threshold / 2);
let removed: Vec<_> = snippets.drain(0..to_remove).collect();
for r in removed {
condensed_snippet_content.push_str(&format!(
"Name: {}\nDesc: {}\nCode: {}\n",
r.name, r.description, r.code
));
}
}
});
if !condensed_snippet_content.is_empty() {
state.modify_graph(|graph| {
let name = format!("Snippet History {}", now);
graph.entities.insert(
name.clone(),
crate::models::Entity {
name: name.clone(),
entity_type: "Historical Summary".to_string(),
observations: vec![condensed_snippet_content],
namespace: crate::models::default_namespace(),
git_branch: None,
},
);
});
tracing::info!("Condensed snippets into Historical Summary.");
}
}
}
async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::Error>> {
let state_for_index = Arc::clone(&state);
tokio::spawn(async move {
state_for_index.rebuild_index().await;
tracing::info!("Index rebuild complete.");
});
// Start the global codebase indexer
crate::indexer::start_background_indexer(Arc::clone(&state)).await;
tokio::spawn(index_committer_worker(Arc::clone(&state)));
tokio::spawn(ttl_sweeper_worker(Arc::clone(&state)));
tokio::spawn(condense_graph_worker(Arc::clone(&state)));
crate::clipboard_watcher::spawn_watcher(Arc::clone(&state));
crate::watcher::spawn_watcher(Arc::clone(&state));
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
let app_state = Arc::new(AppState {
handler: Arc::new(MemoryHandler::new(Arc::clone(&state))),
clients: RwLock::new(HashMap::new()),
next_id: AtomicUsize::new(1),
shutdown_tx: std::sync::Mutex::new(Some(shutdown_tx)),
});
let app_state_clone = Arc::clone(&app_state);
let mut rx = state.activity_tx.subscribe();
tokio::spawn(async move {
loop {
match rx.recv().await {
Ok(msg) => {
let senders: Vec<_> = app_state_clone
.clients
.read()
.unwrap_or_else(|e| e.into_inner())
.values()
.cloned()
.collect();
for client_tx in senders {
let _ = client_tx.try_send(msg.clone());
}
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
}
}
});
// UDP Telemetry Listener
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());
if let Ok(socket) = tokio::net::UdpSocket::bind(format!("127.0.0.1:{}", port1)).await {
let mut buf = [0; 4096];
loop {
if let Ok((len, _addr)) = socket.recv_from(&mut buf).await
&& let Ok(payload) =
serde_json::from_slice::<crate::models::TerminalHistory>(&buf[..len])
{
udp_state
.handler
.state
.telemetry
.terminal_history
.modify(|history| {
history.push_front(payload.clone());
if history.len() > 100 {
history.pop_back();
}
});
let ws_msg = serde_json::json!({
"type": "terminal_telemetry",
"data": payload
});
let msg_str = ws_msg.to_string();
let senders: Vec<_> = udp_state
.clients
.read()
.unwrap_or_else(|e| e.into_inner())
.values()
.cloned()
.collect();
for tx in senders {
let _ = tx.try_send(msg_str.clone());
}
}
}
}
});
// UDP Neovim Telemetry Listener
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());
if let Ok(socket) = tokio::net::UdpSocket::bind(format!("127.0.0.1:{}", port2)).await {
let mut buf = [0; 4096];
loop {
if let Ok((len, _addr)) = socket.recv_from(&mut buf).await
&& let Ok(payload) =
serde_json::from_slice::<crate::api::telemetry::NvimTelemetry>(&buf[..len])
{
// 1. Legacy disk write for active_nvim.txt
if payload.event == "FocusGained"
|| payload.event == "BufEnter"
|| payload.event == "VimEnter"
{
let session = &payload.session_id;
let is_unix_socket = session.starts_with('/') || session.starts_with('~');
if is_unix_socket {
let wsl_path =
"\\\\wsl.localhost\\Ubuntu\\home\\riz\\.gemini\\active_nvim.txt";
let _ = tokio::fs::write(wsl_path, session).await;
} else {
let profile = std::env::var("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;
}
}
// 2. Broadcast to UI
let ws_msg = serde_json::json!({
"type": "nvim_telemetry",
"data": payload
});
let msg_str = ws_msg.to_string();
let senders: Vec<_> = nvim_udp_state
.clients
.read()
.unwrap_or_else(|e| e.into_inner())
.values()
.cloned()
.collect();
for tx in senders {
let _ = tx.try_send(msg_str.clone());
}
// 3. Event bus trigger for auto-save hook
if payload.event == "BufWritePost"
&& let Some(ref file_path) = payload.file
{
let normalized_file = file_path.replace("\\", "/");
let topic = format!("nvim:save:{}", normalized_file);
let event = crate::state::GenericEvent {
topic,
session_id: Some(payload.session_id.clone()),
payload: serde_json::json!(&payload),
};
let _ = nvim_udp_state.handler.state.event_bus_tx.send(event);
}
// 4. Interactive Agent UI Events
if payload.event.starts_with("agent_") {
let topic = format!("nvim:ui:{}", payload.event);
let event = crate::state::GenericEvent {
topic,
session_id: Some(payload.session_id.clone()),
payload: serde_json::json!(&payload),
};
let _ = nvim_udp_state.handler.state.event_bus_tx.send(event);
}
}
}
}
});
let app = api::setup::create_router(app_state);
let port_str = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
tracing::info!(
"MCP Memory Server running on http://127.0.0.1:{}/ws",
port_str
);
let addr: std::net::SocketAddr = format!("127.0.0.1:{}", port_str)
.parse()
.expect("Invalid bind address");
let listener = match tokio::net::TcpListener::bind(&addr).await {
Ok(l) => l,
Err(e) => {
let log_path = dirs::home_dir()
.unwrap_or_default()
.join(".gemini/mcp_memory/daemon_error.log");
let _ =
tokio::fs::write(&log_path, format!("Failed to bind to {}: {}\n", addr, e)).await;
return Ok(());
}
};
if let Err(e) = axum::serve(listener, app.into_make_service())
.with_graceful_shutdown(async move {
let _ = shutdown_rx.await;
})
.await
{
let log_path = dirs::home_dir()
.unwrap_or_default()
.join(".gemini/mcp_memory/daemon_error.log");
let _ = tokio::fs::write(&log_path, format!("Server crashed: {}\n", e)).await;
}
Ok(())
}
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(|_| {
dirs::home_dir()
.map(|mut h| {
h.push(".gemini/mcp_memory");
h.to_string_lossy().to_string()
})
.unwrap_or_else(|| ".gemini/mcp_memory".into())
});
let log_dir = std::path::PathBuf::from(base_dir).join("logs");
std::fs::create_dir_all(&log_dir).unwrap_or_default();
let file_appender = tracing_appender::rolling::daily(log_dir, format!("{}.log", app_name));
let (non_blocking, guard) = tracing_appender::non_blocking(file_appender);
let _ = tracing_subscriber::fmt()
.with_writer(non_blocking)
.with_ansi(false)
.with_max_level(tracing::Level::INFO)
.with_thread_ids(true)
.with_thread_names(true)
.try_init();
Some(guard)
}
fn main() -> 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(|_| {
dirs::home_dir()
.map(|mut h| {
h.push(".gemini/mcp_memory");
h.to_string_lossy().into_owned()
})
.unwrap_or_else(|| ".gemini/mcp_memory".into())
});
let base = PathBuf::from(base_dir);
if cli.exit {
let port = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
let token = std::fs::read_to_string(base.join("admin.token")).unwrap_or_default();
let mut cmd = std::process::Command::new("curl");
cmd.arg("-k").arg("-X").arg("POST");
if !token.is_empty() {
cmd.arg("-H")
.arg(format!("Authorization: Bearer {}", token.trim()));
}
let _ = cmd
.arg(format!("http://127.0.0.1:{}/shutdown", port))
.output();
if cli.restart {
std::thread::sleep(Duration::from_secs(2));
} else {
return Ok(());
}
}
if let Some(Commands::Gate { subcmd }) = cli.command {
let port = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
let rt = tokio::runtime::Runtime::new()?;
match subcmd {
GateCommands::Set {
action,
target,
namespace,
params,
authorize,
block,
reason,
} => {
let mut pmap = HashMap::new();
for p in params {
if let Some((k, v)) = p.split_once('=') {
pmap.insert(k.to_string(), v.to_string());
}
}
let req = GateSetReq {
action,
target,
namespace,
params: pmap,
authorize: if authorize { Some(true) } else { None },
block: if block { Some(true) } else { None },
reason,
};
rt.block_on(async {
let client = reqwest::Client::new();
let res = client
.post(format!("http://127.0.0.1:{}/gate/set", port))
.json(&req)
.send()
.await;
match res {
Ok(r) if r.status().is_success() => println!("Gate updated successfully"),
Ok(r) => println!("Failed to update gate: {}", r.status()),
Err(e) => println!("Error connecting to server: {}", e),
}
});
}
GateCommands::Verify {
action,
target,
namespace,
params: _,
consume,
} => {
let mut url = format!(
"http://127.0.0.1:{}/gate/verify?action={}&target={}&consume={}",
port, action, target, consume
);
if let Some(ns) = namespace {
url.push_str(&format!("&namespace={}", ns));
}
rt.block_on(async {
let res = reqwest::get(&url).await;
match res {
Ok(r) if r.status().is_success() => std::process::exit(0),
Ok(r) if r.status() == reqwest::StatusCode::FORBIDDEN => {
let text = r.text().await.unwrap_or_default();
eprintln!("{}", text);
std::process::exit(1);
}
Ok(r) if r.status() == reqwest::StatusCode::NOT_FOUND => {
eprintln!("Action not yet authorized.");
std::process::exit(2);
}
Ok(r) => {
eprintln!("Unexpected status: {}", r.status());
std::process::exit(3);
}
Err(e) => {
eprintln!("Error connecting to server: {}", e);
std::process::exit(4);
}
}
});
}
}
return Ok(());
}
let token = uuid::Uuid::new_v4().to_string();
std::fs::write(base.join("admin.token"), &token).unwrap_or_default();
let rt = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()
.unwrap();
rt.block_on(async {
let state = Arc::new(MemoryState::new(&base.to_string_lossy()));
// Initialize Qdrant VectorDB (default local URL)
match crate::vector_db::VectorDB::new("http://localhost:6334", "mcp_memory").await {
Ok(vdb) => {
tracing::info!("Successfully connected to Qdrant vector database");
*state.vector_db.write().await = Some(vdb);
}
Err(e) => {
tracing::warn!("Failed to initialize Qdrant vector database: {}. Vector search will fallback to manual embedding loop. (Is Qdrant running on localhost:6334?)", e);
}
}
if let Err(e) = run_server(state).await {
tracing::error!("Server error: {}", e);
}
});
Ok(())
mcp_memory_server::run_cli()
}