958 lines
35 KiB
Rust
958 lines
35 KiB
Rust
#![cfg_attr(
|
|
not(target_os = "windows"),
|
|
allow(dead_code, unused_imports, unreachable_code)
|
|
)]
|
|
|
|
pub mod api;
|
|
pub mod config;
|
|
pub mod db;
|
|
pub mod embedding;
|
|
pub mod error;
|
|
pub mod handlers;
|
|
pub mod indexer;
|
|
pub mod mcp;
|
|
pub mod models;
|
|
pub mod ollama;
|
|
pub mod router;
|
|
pub mod search;
|
|
pub mod state;
|
|
pub mod store;
|
|
pub mod tools;
|
|
pub 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)]
|
|
pub struct Cli {
|
|
#[command(subcommand)]
|
|
pub command: Option<Commands>,
|
|
/// Target URL for the proxy to connect to (e.g., http://127.0.0.1:3000)
|
|
#[arg(long)]
|
|
pub target: Option<String>,
|
|
/// Send a shutdown request to the currently running server
|
|
#[arg(long)]
|
|
pub exit: bool,
|
|
/// Send a shutdown request to the existing server and wait for it to exit
|
|
#[arg(long)]
|
|
pub restart: bool,
|
|
}
|
|
|
|
#[derive(Subcommand)]
|
|
pub enum Commands {
|
|
/// Manage authorization gates and verification for actions
|
|
Gate {
|
|
#[command(subcommand)]
|
|
subcmd: GateCommands,
|
|
},
|
|
}
|
|
|
|
#[derive(Subcommand)]
|
|
pub 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 {
|
|
pub handler: Arc<MemoryHandler>,
|
|
pub clients: RwLock<HashMap<String, mpsc::Sender<String>>>,
|
|
pub next_id: AtomicUsize,
|
|
pub shutdown_tx: std::sync::Mutex<Option<tokio::sync::oneshot::Sender<()>>>,
|
|
}
|
|
|
|
pub async fn ttl_sweeper_worker(state: Arc<MemoryState>) {
|
|
loop {
|
|
let now = std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.unwrap_or_default()
|
|
.as_secs();
|
|
|
|
let mut next_expiry: Option<u64> = None;
|
|
|
|
state.project.tasks.read_with(|tasks| {
|
|
for t in tasks.iter() {
|
|
if let Some(exp) = t.expires_at
|
|
&& t.is_active() {
|
|
next_expiry = Some(next_expiry.map_or(exp, |curr| curr.min(exp)));
|
|
}
|
|
}
|
|
});
|
|
|
|
state.telemetry.handoff_memos.read_with(|memos| {
|
|
for m in memos.iter() {
|
|
if let Some(exp) = m.expires_at {
|
|
next_expiry = Some(next_expiry.map_or(exp, |curr| curr.min(exp)));
|
|
}
|
|
}
|
|
});
|
|
|
|
state.telemetry.session_summaries.read_with(|summaries| {
|
|
for s in summaries.iter() {
|
|
if let Some(exp) = s.expires_at {
|
|
next_expiry = Some(next_expiry.map_or(exp, |curr| curr.min(exp)));
|
|
}
|
|
}
|
|
});
|
|
|
|
state.env.gates.read_with(|gates| {
|
|
for g in gates.iter() {
|
|
if let Some(exp) = g.expires_at {
|
|
next_expiry = Some(next_expiry.map_or(exp, |curr| curr.min(exp)));
|
|
}
|
|
}
|
|
});
|
|
|
|
let sleep_duration = match next_expiry {
|
|
Some(exp) if exp > now => {
|
|
let diff = exp - now;
|
|
std::time::Duration::from_secs(diff.min(60).max(1))
|
|
}
|
|
Some(_) => std::time::Duration::from_millis(50),
|
|
None => std::time::Duration::from_secs(60),
|
|
};
|
|
|
|
tokio::select! {
|
|
_ = state.shutdown_notify.notified() => break,
|
|
_ = state.ttl_notify.notified() => {},
|
|
_ = tokio::time::sleep(sleep_duration) => {},
|
|
}
|
|
|
|
let now = std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.unwrap_or_default()
|
|
.as_secs();
|
|
|
|
let mut expired_tasks = Vec::new();
|
|
state.project.tasks.modify(|tasks| {
|
|
for t in tasks.iter_mut() {
|
|
if let Some(exp) = t.expires_at
|
|
&& exp <= now && t.is_active() {
|
|
t.status = "expired".to_string();
|
|
t.updated_at = now;
|
|
expired_tasks.push(t.id.clone());
|
|
}
|
|
}
|
|
});
|
|
|
|
for tid in expired_tasks {
|
|
state.record_activity(
|
|
"task_expired",
|
|
&format!("Task {} expired due to TTL", tid),
|
|
Some("expired"),
|
|
);
|
|
state.broadcast_task_event(crate::models::TaskEvent {
|
|
task_id: tid,
|
|
status: "expired".to_string(),
|
|
action: Some("ttl_expire".to_string()),
|
|
result: Some(serde_json::json!({ "status": "expired" })),
|
|
error: None,
|
|
timestamp: now,
|
|
session_id: None,
|
|
..Default::default()
|
|
});
|
|
}
|
|
|
|
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));
|
|
});
|
|
state.env.gates.modify(|gates| {
|
|
gates.retain(|g| g.expires_at.is_none_or(|exp| exp > now));
|
|
});
|
|
}
|
|
}
|
|
|
|
pub async fn index_committer_worker(state: Arc<MemoryState>) {
|
|
loop {
|
|
tokio::select! {
|
|
_ = state.shutdown_notify.notified() => break,
|
|
_ = state.index_commit_notify.notified() => {},
|
|
}
|
|
let idx = state.search_index.read().await.clone();
|
|
let _ = idx.commit().await;
|
|
}
|
|
}
|
|
|
|
pub async fn condense_graph_worker(state: Arc<MemoryState>) {
|
|
loop {
|
|
tokio::select! {
|
|
_ = state.shutdown_notify.notified() => break,
|
|
_ = state.condense_notify.notified() => {},
|
|
}
|
|
|
|
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();
|
|
|
|
let snippet_condensation = state.code.snippets.read_with(|snippets| {
|
|
if snippets.len() > threshold {
|
|
let mut sorted = snippets.clone();
|
|
sorted.sort_by_key(|s| s.updated_at);
|
|
let to_remove = sorted.len() - (threshold / 2);
|
|
let removed: Vec<_> = sorted.into_iter().take(to_remove).collect();
|
|
let mut content = String::new();
|
|
let mut names = Vec::new();
|
|
for r in &removed {
|
|
content.push_str(&format!(
|
|
"Name: {}\nDesc: {}\nCode: {}\n",
|
|
r.name, r.description, r.code
|
|
));
|
|
names.push(r.name.clone());
|
|
}
|
|
Some((content, names))
|
|
} else {
|
|
None
|
|
}
|
|
});
|
|
|
|
if let Some((content, names)) = snippet_condensation
|
|
&& !content.is_empty() {
|
|
let name = format!("Snippet History {}", now);
|
|
state.modify_graph(|graph| {
|
|
graph.entities.insert(
|
|
name.clone(),
|
|
crate::models::Entity {
|
|
name: name.clone(),
|
|
entity_type: "Historical Summary".to_string(),
|
|
observations: vec![content],
|
|
namespace: crate::models::default_namespace(),
|
|
git_branch: None,
|
|
..Default::default()
|
|
},
|
|
);
|
|
});
|
|
let name_set: std::collections::HashSet<String> = names.into_iter().collect();
|
|
state.code.snippets.modify(|snippets| {
|
|
snippets.retain(|s| !name_set.contains(&s.name));
|
|
});
|
|
tracing::info!("Condensed snippets into Historical Summary.");
|
|
}
|
|
}
|
|
}
|
|
|
|
pub async fn memory_consolidation_worker(state: Arc<MemoryState>) {
|
|
let mut interval = tokio::time::interval(std::time::Duration::from_secs(300));
|
|
loop {
|
|
tokio::select! {
|
|
_ = state.shutdown_notify.notified() => break,
|
|
_ = interval.tick() => {},
|
|
}
|
|
|
|
let entities: Vec<_> = state.graph.read_with(|g| {
|
|
g.entities
|
|
.values()
|
|
.map(|e| (e.name.clone(), e.entity_type.clone()))
|
|
.collect()
|
|
});
|
|
|
|
if entities.len() < 2 {
|
|
continue;
|
|
}
|
|
|
|
let mut entity_summaries = String::new();
|
|
for (name, e_type) in entities.iter().take(50) {
|
|
entity_summaries.push_str(&format!("- [{}] {}\n", e_type, name));
|
|
}
|
|
|
|
let prompt = format!(
|
|
"Analyze the following list of entities and identify exactly TWO that represent the exact same concept or item but have slightly different names (e.g. 'auth_service' and 'AuthService'). Return ONLY a valid JSON array containing exactly two strings: the two names to merge. If no obvious duplicates exist, return an empty array []. Do not output any markdown formatting or extra text.\n\nEntities:\n{}",
|
|
entity_summaries
|
|
);
|
|
|
|
if let Ok(response) = state
|
|
.ollama
|
|
.generate(
|
|
&prompt,
|
|
None,
|
|
Some("You are a helpful JSON-only data deduplication assistant. Output only JSON."),
|
|
)
|
|
.await
|
|
{
|
|
let cleaned = response
|
|
.trim()
|
|
.trim_start_matches("```json")
|
|
.trim_start_matches("```")
|
|
.trim_end_matches("```")
|
|
.trim();
|
|
if let Ok(duplicates) = serde_json::from_str::<Vec<String>>(cleaned)
|
|
&& duplicates.len() == 2 {
|
|
let e1_name = &duplicates[0];
|
|
let e2_name = &duplicates[1];
|
|
|
|
if e1_name != e2_name {
|
|
tracing::info!(
|
|
"Memory Consolidation Daemon: Merging '{}' into '{}'",
|
|
e2_name,
|
|
e1_name
|
|
);
|
|
state.modify_graph(|g| {
|
|
if let Some(mut e2) = g.entities.remove(e2_name) {
|
|
if let Some(e1) = g.entities.get_mut(e1_name) {
|
|
e1.observations.append(&mut e2.observations);
|
|
} else {
|
|
g.entities.insert(e2_name.clone(), e2);
|
|
}
|
|
}
|
|
|
|
for rel in g.relations.iter_mut() {
|
|
if rel.from == *e2_name {
|
|
rel.from = e1_name.clone();
|
|
}
|
|
if rel.to == *e2_name {
|
|
rel.to = e1_name.clone();
|
|
}
|
|
}
|
|
});
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub 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.");
|
|
});
|
|
|
|
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)));
|
|
tokio::spawn(memory_consolidation_worker(Arc::clone(&state)));
|
|
crate::watcher::spawn_watcher(Arc::clone(&state));
|
|
crate::handlers::vision::spawn_clipboard_listener(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 mut closed_ids = Vec::new();
|
|
{
|
|
let clients_guard = app_state_clone
|
|
.clients
|
|
.read()
|
|
.unwrap_or_else(|e| e.into_inner());
|
|
for (id, tx) in clients_guard.iter() {
|
|
if tx.try_send(msg.clone()).is_err() && tx.is_closed() {
|
|
closed_ids.push(id.clone());
|
|
}
|
|
}
|
|
}
|
|
if !closed_ids.is_empty() {
|
|
let mut write_guard = app_state_clone
|
|
.clients
|
|
.write()
|
|
.unwrap_or_else(|e| e.into_inner());
|
|
for id in closed_ids {
|
|
write_guard.remove(&id);
|
|
}
|
|
}
|
|
}
|
|
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
|
|
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
|
|
}
|
|
}
|
|
});
|
|
|
|
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 socket = Arc::new(socket);
|
|
let socket_rx = socket.clone();
|
|
|
|
let mut subscribers: HashMap<(String, String), std::net::SocketAddr> = HashMap::new();
|
|
let mut buf = vec![0u8; 65536];
|
|
let mut event_rx = udp_state.handler.state.event_bus_tx.subscribe();
|
|
|
|
loop {
|
|
tokio::select! {
|
|
recv_res = socket_rx.recv_from(&mut buf) => {
|
|
if let Ok((len, addr)) = recv_res {
|
|
if let Ok(payload) = serde_json::from_slice::<crate::models::TerminalHistory>(&buf[..len]) {
|
|
udp_state
|
|
.handler
|
|
.state
|
|
.record_terminal_history(payload.clone());
|
|
|
|
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());
|
|
}
|
|
} else if let Ok(json_payload) = serde_json::from_slice::<serde_json::Value>(&buf[..len]) {
|
|
if json_payload.get("type").and_then(|t| t.as_str()) == Some("ping") {
|
|
let _ = socket.send_to(b"pong", addr).await;
|
|
} else if json_payload.get("type").and_then(|t| t.as_str()) == Some("gate_wait")
|
|
&& let (Some(action), Some(target)) = (
|
|
json_payload.get("action").and_then(|a| a.as_str()),
|
|
json_payload.get("target").and_then(|t| t.as_str())
|
|
) {
|
|
subscribers.insert((action.to_string(), target.to_string()), addr);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Ok(event) = event_rx.recv() => {
|
|
if event.topic == "gate:event"
|
|
&& let (Some(action), Some(target), Some(status)) = (
|
|
event.payload.get("action").and_then(|a| a.as_str()),
|
|
event.payload.get("target").and_then(|t| t.as_str()),
|
|
event.payload.get("status").and_then(|s| s.as_str()),
|
|
) {
|
|
if (status == "authorized" || status == "blocked")
|
|
&& let Some(addr) = subscribers.remove(&(action.to_string(), target.to_string())) {
|
|
let response = if status == "authorized" { b"APPROVED" } else { b"REJECTED" };
|
|
let _ = socket.send_to(response, addr).await;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
});
|
|
|
|
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 = vec![0u8; 65536];
|
|
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])
|
|
{
|
|
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 let Some(home) = dirs::home_dir() {
|
|
let nvim_dir = home.join(".gemini");
|
|
let nvim_txt_path = nvim_dir.join("active_nvim.txt");
|
|
let tmp_path =
|
|
nvim_dir.join(format!("active_nvim_{}.tmp", std::process::id()));
|
|
if tokio::fs::create_dir_all(&nvim_dir).await.is_ok()
|
|
&& tokio::fs::write(&tmp_path, session).await.is_ok()
|
|
{
|
|
let _ = tokio::fs::rename(&tmp_path, &nvim_txt_path).await;
|
|
}
|
|
}
|
|
}
|
|
|
|
let (tech_debts, adrs) = if let Some(ref f) = payload.file {
|
|
crate::api::telemetry::find_projected_knowledge(
|
|
&nvim_udp_state.handler.state,
|
|
f,
|
|
)
|
|
} else {
|
|
(Vec::new(), Vec::new())
|
|
};
|
|
|
|
let ws_msg = serde_json::json!({
|
|
"type": "nvim_telemetry",
|
|
"data": payload,
|
|
"tech_debts": tech_debts,
|
|
"adrs": adrs
|
|
});
|
|
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());
|
|
}
|
|
|
|
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);
|
|
}
|
|
|
|
if payload.event.starts_with("agent_") || payload.event.starts_with("diff_") {
|
|
let payload_val = serde_json::json!(&payload);
|
|
// 1. General event topic (e.g. nvim:ui:agent_prompt_response, nvim:ui:agent_diff_accepted)
|
|
let _ = nvim_udp_state.handler.state.event_bus_tx.send(
|
|
crate::state::GenericEvent {
|
|
topic: format!("nvim:ui:{}", payload.event),
|
|
session_id: Some(payload.session_id.clone()),
|
|
payload: payload_val.clone(),
|
|
},
|
|
);
|
|
|
|
// 2. Correlated request_id topic (e.g. nvim:ui:agent_prompt_response:REQ_ID)
|
|
if let Some(ref req_id) = payload.request_id {
|
|
let _ = nvim_udp_state.handler.state.event_bus_tx.send(
|
|
crate::state::GenericEvent {
|
|
topic: format!("nvim:ui:{}:{}", payload.event, req_id),
|
|
session_id: Some(payload.session_id.clone()),
|
|
payload: payload_val.clone(),
|
|
},
|
|
);
|
|
}
|
|
|
|
// 3. Correlated diff_id topics
|
|
if let Some(ref diff_id) = payload.diff_id {
|
|
let _ = nvim_udp_state.handler.state.event_bus_tx.send(
|
|
crate::state::GenericEvent {
|
|
topic: format!("nvim:ui:{}:{}", payload.event, diff_id),
|
|
session_id: Some(payload.session_id.clone()),
|
|
payload: payload_val.clone(),
|
|
},
|
|
);
|
|
// General diff decision topic
|
|
let _ = nvim_udp_state.handler.state.event_bus_tx.send(
|
|
crate::state::GenericEvent {
|
|
topic: format!("nvim:ui:diff_decision:{}", diff_id),
|
|
session_id: Some(payload.session_id.clone()),
|
|
payload: payload_val.clone(),
|
|
},
|
|
);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
});
|
|
|
|
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;
|
|
}
|
|
|
|
tracing::info!("axum::serve graceful shutdown complete.");
|
|
Ok(())
|
|
}
|
|
|
|
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(|_| {
|
|
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)
|
|
}
|
|
|
|
pub fn run_cli() -> Result<(), Box<dyn std::error::Error>> {
|
|
crate::config::load_mcp_config_env();
|
|
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 rt = tokio::runtime::Runtime::new()?;
|
|
rt.block_on(async {
|
|
let client = reqwest::Client::builder().build().unwrap_or_default();
|
|
let mut req = client.post(format!("http://127.0.0.1:{}/shutdown", port));
|
|
if !token.trim().is_empty() {
|
|
req = req.header("Authorization", format!("Bearer {}", token.trim()));
|
|
}
|
|
let _ = req.send().await;
|
|
});
|
|
|
|
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()));
|
|
|
|
if let Err(e) = run_server(state).await {
|
|
tracing::error!("Server error: {}", e);
|
|
}
|
|
});
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use clap::Parser;
|
|
|
|
#[test]
|
|
fn test_cli_parsing_default() {
|
|
let cli = Cli::try_parse_from(&["mcp-memory-server"]).unwrap();
|
|
assert!(cli.command.is_none());
|
|
assert!(!cli.exit);
|
|
assert!(!cli.restart);
|
|
}
|
|
|
|
#[test]
|
|
fn test_cli_parsing_exit_and_target() {
|
|
let cli = Cli::try_parse_from(&[
|
|
"mcp-memory-server",
|
|
"--exit",
|
|
"--target",
|
|
"http://localhost:3000",
|
|
])
|
|
.unwrap();
|
|
assert!(cli.exit);
|
|
assert_eq!(cli.target.as_deref(), Some("http://localhost:3000"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_cli_parsing_gate_set() {
|
|
let cli = Cli::try_parse_from(&[
|
|
"mcp-memory-server",
|
|
"gate",
|
|
"set",
|
|
"--action",
|
|
"git_push",
|
|
"--target",
|
|
"master",
|
|
"--authorize",
|
|
"--reason",
|
|
"Approved by lead",
|
|
])
|
|
.unwrap();
|
|
|
|
if let Some(Commands::Gate {
|
|
subcmd:
|
|
GateCommands::Set {
|
|
action,
|
|
target,
|
|
authorize,
|
|
reason,
|
|
..
|
|
},
|
|
}) = cli.command
|
|
{
|
|
assert_eq!(action, "git_push");
|
|
assert_eq!(target, "master");
|
|
assert!(authorize);
|
|
assert_eq!(reason.as_deref(), Some("Approved by lead"));
|
|
} else {
|
|
panic!("Expected Gate Set subcommand");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_cli_parsing_gate_verify() {
|
|
let cli = Cli::try_parse_from(&[
|
|
"mcp-memory-server",
|
|
"gate",
|
|
"verify",
|
|
"--action",
|
|
"deploy",
|
|
"--target",
|
|
"prod",
|
|
"--consume",
|
|
])
|
|
.unwrap();
|
|
|
|
if let Some(Commands::Gate {
|
|
subcmd:
|
|
GateCommands::Verify {
|
|
action,
|
|
target,
|
|
consume,
|
|
..
|
|
},
|
|
}) = cli.command
|
|
{
|
|
assert_eq!(action, "deploy");
|
|
assert_eq!(target, "prod");
|
|
assert!(consume);
|
|
} else {
|
|
panic!("Expected Gate Verify subcommand");
|
|
}
|
|
}
|
|
|
|
static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
|
|
|
|
#[test]
|
|
fn test_init_logging_helper() {
|
|
let _lock = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
|
let temp_dir = tempfile::tempdir().unwrap();
|
|
unsafe {
|
|
std::env::set_var("MCP_MEMORY_STORE_DIR", temp_dir.path().to_str().unwrap());
|
|
}
|
|
let guard = init_logging("test_app");
|
|
assert!(guard.is_some());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_background_workers_one_tick() {
|
|
let temp_dir = tempfile::tempdir().unwrap();
|
|
let state = Arc::new(MemoryState::new(temp_dir.path().to_str().unwrap()));
|
|
|
|
// Test worker functions by spawning them briefly
|
|
let handle1 = tokio::spawn(ttl_sweeper_worker(state.clone()));
|
|
let handle2 = tokio::spawn(index_committer_worker(state.clone()));
|
|
let handle3 = tokio::spawn(condense_graph_worker(state.clone()));
|
|
|
|
state.ttl_notify.notify_one();
|
|
state.index_commit_notify.notify_one();
|
|
state.condense_notify.notify_one();
|
|
|
|
tokio::time::sleep(Duration::from_millis(50)).await;
|
|
|
|
handle1.abort();
|
|
handle2.abort();
|
|
handle3.abort();
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_run_server_graceful_shutdown() {
|
|
let _lock = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
|
|
let temp_dir = tempfile::tempdir().unwrap();
|
|
let state = Arc::new(MemoryState::new(temp_dir.path().to_str().unwrap()));
|
|
|
|
// Bind to a free port to avoid conflicts
|
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
|
let port = listener.local_addr().unwrap().port();
|
|
drop(listener);
|
|
|
|
unsafe {
|
|
std::env::set_var("MCP_PORT", port.to_string());
|
|
std::env::set_var("MCP_UDP_PORT1", (port + 1).to_string());
|
|
std::env::set_var("MCP_UDP_PORT2", (port + 2).to_string());
|
|
}
|
|
|
|
let server_handle = tokio::spawn(async move {
|
|
let _ = run_server(state).await;
|
|
});
|
|
|
|
tokio::time::sleep(Duration::from_millis(300)).await;
|
|
|
|
let _ = reqwest::Client::new()
|
|
.get(format!("http://127.0.0.1:{}/ping", port))
|
|
.send()
|
|
.await;
|
|
|
|
server_handle.abort();
|
|
}
|
|
}
|