feat: implement dual-transport WAL architecture with permanent background leader and lightweight stubs
This commit is contained in:
Riz Ashraf committed
1 parent
a2febd1b00
commit
e4ff476b6d
16 files changed
+1910
-52
No files matched your search
@@ -0,0 +1,309 @@
|
||||
mod handlers;
|
||||
mod models;
|
||||
mod state;
|
||||
mod store;
|
||||
mod tools;
|
||||
|
||||
use crate::handlers::MemoryHandler;
|
||||
use crate::models::*;
|
||||
use crate::state::MemoryState;
|
||||
use crate::store::Store;
|
||||
use rust_mcp_sdk::{
|
||||
ServerDetails, ToMcpServerHandler,
|
||||
error::SdkResult,
|
||||
|
||||
schema::{Implementation, ServerCapabilities, ServerCapabilitiesTools},
|
||||
};
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
use tokio::time::sleep;
|
||||
|
||||
use clap::{Parser, Subcommand};
|
||||
use std::collections::HashMap;
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(author, version, about, long_about = None)]
|
||||
struct Cli {
|
||||
#[command(subcommand)]
|
||||
command: Option<Commands>,
|
||||
#[arg(long)]
|
||||
target: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
enum Commands {
|
||||
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,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
async fn reconcile_worker(state: Arc<MemoryState>) {
|
||||
loop {
|
||||
sleep(Duration::from_secs(5)).await;
|
||||
let pattern = format!("{}/delta_*.json", state.base_dir.display());
|
||||
let has_local = {
|
||||
let session = state.session_graph.read().unwrap();
|
||||
!session.entities.is_empty() || !session.relations.is_empty()
|
||||
};
|
||||
let has_files = glob::glob(&pattern).map(|p| p.count() > 0).unwrap_or(false);
|
||||
if has_local || has_files {
|
||||
state.apply_sync_write(|_master| {});
|
||||
}
|
||||
|
||||
let now = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs();
|
||||
state.ledger.modify(|ledger| {
|
||||
let seven_days = now.saturating_sub(7 * 24 * 60 * 60);
|
||||
ledger.retain(|c| c.timestamp >= seven_days);
|
||||
if ledger.len() > 1000 {
|
||||
let excess = ledger.len() - 1000;
|
||||
ledger.drain(0..excess);
|
||||
}
|
||||
});
|
||||
state.sticky.modify(|notes| {
|
||||
notes.retain(|note| note.timestamp >= now.saturating_sub(24 * 60 * 60));
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
use rust_mcp_axum::{create_axum_server, AxumServerOptions};
|
||||
|
||||
fn run_server(state: Arc<MemoryState>) -> SdkResult<()> {
|
||||
let rt = tokio::runtime::Runtime::new().unwrap();
|
||||
rt.block_on(async {
|
||||
let state_clone = state.clone();
|
||||
tokio::spawn(async move {
|
||||
reconcile_worker(state_clone).await;
|
||||
});
|
||||
|
||||
let server_details = ServerDetails {
|
||||
server_info: Implementation {
|
||||
name: "gemini-mcp-memory".into(),
|
||||
version: "3.0.0".into(),
|
||||
title: Some("MCP Memory Server".into()),
|
||||
description: Some("Persistent memory MCP server".into()),
|
||||
icons: vec![],
|
||||
website_url: None,
|
||||
},
|
||||
capabilities: ServerCapabilities {
|
||||
tools: Some(ServerCapabilitiesTools {
|
||||
list_changed: Some(true),
|
||||
}),
|
||||
..Default::default()
|
||||
},
|
||||
instructions: None,
|
||||
meta: None,
|
||||
};
|
||||
|
||||
let options = AxumServerOptions {
|
||||
host: "0.0.0.0".to_string(),
|
||||
port: 3000,
|
||||
custom_sse_endpoint: Some("/sse".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let server = create_axum_server(
|
||||
server_details,
|
||||
MemoryHandler { state: state.clone() }.to_mcp_server_handler(),
|
||||
options,
|
||||
);
|
||||
|
||||
let server = server.with_route(
|
||||
"/gate/verify",
|
||||
axum::routing::get(|| async move {
|
||||
axum::http::StatusCode::OK
|
||||
}),
|
||||
);
|
||||
|
||||
println!("MCP Memory Server running on http://0.0.0.0:3000/sse");
|
||||
server.start().await.unwrap();
|
||||
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
mod proxy;
|
||||
|
||||
fn main() -> SdkResult<()> {
|
||||
let cli = Cli::parse();
|
||||
|
||||
#[cfg(target_os = "windows")]
|
||||
{
|
||||
loop {
|
||||
if let Err(_) = std::net::TcpListener::bind("0.0.0.0:3000") {
|
||||
// Port in use, become a stub proxy!
|
||||
let target_url = cli.target.as_deref().unwrap_or("http://127.0.0.1:3000");
|
||||
match proxy::run_proxy(target_url) {
|
||||
Ok(true) => {
|
||||
std::thread::sleep(std::time::Duration::from_millis(50));
|
||||
continue; // Leader died, race to bind 3000
|
||||
}
|
||||
Ok(false) => return Ok(()), // Stdin closed, user exited
|
||||
Err(_) => std::thread::sleep(std::time::Duration::from_millis(1000)),
|
||||
}
|
||||
} else {
|
||||
break; // Proceed as Leader
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "windows"))]
|
||||
{
|
||||
// This shouldn't be executed on linux natively anymore due to workspace split,
|
||||
// but keeping it as a fallback.
|
||||
let target_url = cli.target.as_deref().unwrap_or("http://host.docker.internal:3000");
|
||||
let _ = proxy::run_proxy(target_url);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
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);
|
||||
fs::create_dir_all(&base).expect("Failed to create store dir");
|
||||
|
||||
let state = Arc::new(MemoryState {
|
||||
master_path: base.join("knowledge_graph_master.json"),
|
||||
session_graph: RwLock::new(KnowledgeGraph::default()),
|
||||
base_dir: base.clone(),
|
||||
master_cache: RwLock::new((KnowledgeGraph::default(), SystemTime::UNIX_EPOCH)),
|
||||
ledger: Store::new(base.join("audit_ledger.json")),
|
||||
sticky: Store::new(base.join("sticky_notes.json")),
|
||||
tasks: Store::new(base.join("tasks.json")),
|
||||
snippets: Store::new(base.join("snippets.json")),
|
||||
adrs: Store::new(base.join("adrs.json")),
|
||||
prefs: Store::new(base.join("preferences.json")),
|
||||
error_fixes: Store::new(base.join("error_fixes.json")),
|
||||
pinned_files: Store::new(base.join("pinned_files.json")),
|
||||
session_summaries: Store::new(base.join("session_summaries.json")),
|
||||
handoff_memos: Store::new(base.join("handoff_memos.json")),
|
||||
env_fingerprints: Store::new(base.join("env_fingerprints.json")),
|
||||
env_requirements: Store::new(base.join("env_requirements.json")),
|
||||
milestones: Store::new(base.join("milestones.json")),
|
||||
environments: Store::new(base.join("environments.json")),
|
||||
pr_checklists: Store::new(base.join("pr_checklists.json")),
|
||||
tech_debts: Store::new(base.join("tech_debts.json")),
|
||||
gates: Store::new(base.join("gates.json")),
|
||||
context_workspaces: Store::new(base.join("context_workspaces.json")),
|
||||
});
|
||||
|
||||
if let Some(command) = cli.command {
|
||||
match command {
|
||||
Commands::Gate { subcmd } => {
|
||||
match subcmd {
|
||||
GateCommands::Set { action, target, namespace, params, authorize, block, reason } => {
|
||||
let status = if authorize { "authorized".to_string() } else if block { "blocked".to_string() } else { "pending".to_string() };
|
||||
let mut param_map = HashMap::new();
|
||||
for p in params {
|
||||
if let Some((k, v)) = p.split_once('=') {
|
||||
param_map.insert(k.to_string(), v.to_string());
|
||||
}
|
||||
}
|
||||
let record = GateRecord {
|
||||
id: uuid::Uuid::new_v4().to_string(),
|
||||
action: action.clone(),
|
||||
target: target.clone(),
|
||||
namespace,
|
||||
params: param_map,
|
||||
status,
|
||||
reason,
|
||||
timestamp: SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs(),
|
||||
};
|
||||
state.gates.modify(|gates| {
|
||||
gates.retain(|g| !(g.action == record.action && g.target == record.target));
|
||||
gates.push(record);
|
||||
});
|
||||
println!("Gate state updated.");
|
||||
std::process::exit(0);
|
||||
}
|
||||
GateCommands::Verify { action, target, namespace, params, consume } => {
|
||||
let mut param_map = HashMap::new();
|
||||
for p in params {
|
||||
if let Some((k, v)) = p.split_once('=') {
|
||||
param_map.insert(k.to_string(), v.to_string());
|
||||
}
|
||||
}
|
||||
let mut found = None;
|
||||
let mut to_remove = None;
|
||||
state.gates.modify(|gates| {
|
||||
if let Some(idx) = gates.iter().position(|g| g.action == action && g.target == target && g.namespace == namespace && g.params == param_map) {
|
||||
found = Some(gates[idx].clone());
|
||||
if consume {
|
||||
to_remove = Some(idx);
|
||||
}
|
||||
}
|
||||
if let Some(idx) = to_remove {
|
||||
gates.remove(idx);
|
||||
}
|
||||
});
|
||||
|
||||
match found {
|
||||
Some(record) => {
|
||||
if record.status == "authorized" {
|
||||
std::process::exit(0);
|
||||
} else {
|
||||
if let Some(r) = record.reason {
|
||||
eprintln!("❌ Action blocked. Reason: {}", r);
|
||||
} else {
|
||||
eprintln!("❌ Action blocked.");
|
||||
}
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
None => {
|
||||
eprintln!("❌ Action not yet authorized (no gate record found).");
|
||||
std::process::exit(2);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
run_server(state)
|
||||
}
|
||||
Reference in new issue
Block a user