chore: migrate store to redb
This commit is contained in:
Riz Ashraf committed
1 parent
7f0286dc40
commit
9a8b6e52e9
12 files changed
+1954
-1365
No files matched your search
+244
-138
@@ -1,10 +1,10 @@
|
||||
mod handlers;
|
||||
mod models;
|
||||
mod mcp;
|
||||
mod models;
|
||||
mod search;
|
||||
mod state;
|
||||
mod store;
|
||||
mod tools;
|
||||
mod search;
|
||||
|
||||
use crate::handlers::MemoryHandler;
|
||||
use crate::models::*;
|
||||
@@ -29,6 +29,10 @@ struct Cli {
|
||||
target: Option<String>,
|
||||
#[arg(long)]
|
||||
daemon: bool,
|
||||
#[arg(long)]
|
||||
exit: bool,
|
||||
#[arg(long)]
|
||||
restart: bool,
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
@@ -71,7 +75,6 @@ enum GateCommands {
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
async fn reconcile_worker(state: Arc<MemoryState>) {
|
||||
loop {
|
||||
sleep(Duration::from_secs(5)).await;
|
||||
@@ -104,21 +107,17 @@ async fn reconcile_worker(state: Arc<MemoryState>) {
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
use axum::{
|
||||
extract::{State, Query},
|
||||
Json, Router,
|
||||
extract::{Query, State},
|
||||
response::sse::{Event, Sse},
|
||||
routing::{get, post},
|
||||
Json, Router,
|
||||
};
|
||||
use futures_util::stream::Stream;
|
||||
use std::convert::Infallible;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_stream::wrappers::ReceiverStream;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
struct AppState {
|
||||
handler: Arc<MemoryHandler>,
|
||||
@@ -140,58 +139,72 @@ fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::Error>>
|
||||
.route("/sse", get(sse_handler))
|
||||
.route("/messages", post(message_handler))
|
||||
.route("/health", get(health_handler))
|
||||
.route("/", get(|| async move {
|
||||
axum::response::Html(include_str!("dashboard.html"))
|
||||
}))
|
||||
.route("/api/stats", get({
|
||||
let state_clone = app_state.handler.state.clone();
|
||||
move || async move {
|
||||
let (entities, relations) = {
|
||||
let graph = state_clone.get_full_graph();
|
||||
(graph.entities.len(), graph.relations.len())
|
||||
};
|
||||
let tasks = state_clone.tasks.read().len();
|
||||
let snippets = state_clone.snippets.read().len();
|
||||
let tech_debts = state_clone.tech_debts.read().len();
|
||||
let adrs = state_clone.adrs.read().len();
|
||||
|
||||
let ledger = state_clone.ledger.read().len();
|
||||
let sticky = state_clone.sticky.read().len();
|
||||
let error_fixes = state_clone.error_fixes.read().len();
|
||||
let pinned_files = state_clone.pinned_files.read().len();
|
||||
let session_summaries = state_clone.session_summaries.read().len();
|
||||
let handoff_memos = state_clone.handoff_memos.read().len();
|
||||
let env_fingerprints = state_clone.env_fingerprints.read().len();
|
||||
let env_requirements = state_clone.env_requirements.read().len();
|
||||
let milestones = state_clone.milestones.read().len();
|
||||
let environments = state_clone.environments.read().len();
|
||||
let pr_checklists = state_clone.pr_checklists.read().len();
|
||||
let gates = state_clone.gates.read().len();
|
||||
let context_workspaces = state_clone.context_workspaces.read().len();
|
||||
.route(
|
||||
"/shutdown",
|
||||
post(|| async move {
|
||||
std::thread::spawn(|| {
|
||||
std::thread::sleep(std::time::Duration::from_millis(100));
|
||||
std::process::exit(0);
|
||||
});
|
||||
"Shutting down..."
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/",
|
||||
get(|| async move { axum::response::Html(include_str!("dashboard.html")) }),
|
||||
)
|
||||
.route(
|
||||
"/api/stats",
|
||||
get({
|
||||
let state_clone = app_state.handler.state.clone();
|
||||
move || async move {
|
||||
let (entities, relations) = {
|
||||
let graph = state_clone.get_full_graph();
|
||||
(graph.entities.len(), graph.relations.len())
|
||||
};
|
||||
let tasks = state_clone.tasks.read().len();
|
||||
let snippets = state_clone.snippets.read().len();
|
||||
let tech_debts = state_clone.tech_debts.read().len();
|
||||
let adrs = state_clone.adrs.read().len();
|
||||
|
||||
axum::Json(serde_json::json!({
|
||||
"entities": entities,
|
||||
"relations": relations,
|
||||
"tasks": tasks,
|
||||
"snippets": snippets,
|
||||
"tech_debts": tech_debts,
|
||||
"adrs": adrs,
|
||||
"ledger": ledger,
|
||||
"sticky": sticky,
|
||||
"error_fixes": error_fixes,
|
||||
"pinned_files": pinned_files,
|
||||
"session_summaries": session_summaries,
|
||||
"handoff_memos": handoff_memos,
|
||||
"env_fingerprints": env_fingerprints,
|
||||
"env_requirements": env_requirements,
|
||||
"milestones": milestones,
|
||||
"environments": environments,
|
||||
"pr_checklists": pr_checklists,
|
||||
"gates": gates,
|
||||
"context_workspaces": context_workspaces
|
||||
}))
|
||||
}
|
||||
}))
|
||||
let ledger = state_clone.ledger.read().len();
|
||||
let sticky = state_clone.sticky.read().len();
|
||||
let error_fixes = state_clone.error_fixes.read().len();
|
||||
let pinned_files = state_clone.pinned_files.read().len();
|
||||
let session_summaries = state_clone.session_summaries.read().len();
|
||||
let handoff_memos = state_clone.handoff_memos.read().len();
|
||||
let env_fingerprints = state_clone.env_fingerprints.read().len();
|
||||
let env_requirements = state_clone.env_requirements.read().len();
|
||||
let milestones = state_clone.milestones.read().len();
|
||||
let environments = state_clone.environments.read().len();
|
||||
let pr_checklists = state_clone.pr_checklists.read().len();
|
||||
let gates = state_clone.gates.read().len();
|
||||
let context_workspaces = state_clone.context_workspaces.read().len();
|
||||
|
||||
axum::Json(serde_json::json!({
|
||||
"entities": entities,
|
||||
"relations": relations,
|
||||
"tasks": tasks,
|
||||
"snippets": snippets,
|
||||
"tech_debts": tech_debts,
|
||||
"adrs": adrs,
|
||||
"ledger": ledger,
|
||||
"sticky": sticky,
|
||||
"error_fixes": error_fixes,
|
||||
"pinned_files": pinned_files,
|
||||
"session_summaries": session_summaries,
|
||||
"handoff_memos": handoff_memos,
|
||||
"env_fingerprints": env_fingerprints,
|
||||
"env_requirements": env_requirements,
|
||||
"milestones": milestones,
|
||||
"environments": environments,
|
||||
"pr_checklists": pr_checklists,
|
||||
"gates": gates,
|
||||
"context_workspaces": context_workspaces
|
||||
}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.with_state(app_state);
|
||||
|
||||
let mut retries = 0;
|
||||
@@ -199,16 +212,48 @@ fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::Error>>
|
||||
match tokio::net::TcpListener::bind("127.0.0.1:3000").await {
|
||||
Ok(l) => break l,
|
||||
Err(e) => {
|
||||
// Check if it's already running and healthy
|
||||
if let Ok(mut stream) = std::net::TcpStream::connect("127.0.0.1:3000") {
|
||||
use std::io::{Read, Write};
|
||||
let _ = stream.write_all(
|
||||
b"GET /health HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n",
|
||||
);
|
||||
let mut response = String::new();
|
||||
let _ = stream.read_to_string(&mut response);
|
||||
if response.contains("200 OK") {
|
||||
// Already healthy! Just exit cleanly instead of panicking/retrying loop.
|
||||
std::process::exit(0);
|
||||
}
|
||||
}
|
||||
|
||||
retries += 1;
|
||||
if retries > 15 {
|
||||
let log_path = dirs::home_dir().unwrap_or_default().join(".gemini/mcp_memory/daemon_fatal.log");
|
||||
let _ = std::fs::write(&log_path, format!("FATAL: Could not bind to 127.0.0.1:3000 after 15 seconds: {}\n", e));
|
||||
let log_path = dirs::home_dir()
|
||||
.unwrap_or_default()
|
||||
.join(".gemini/mcp_memory/daemon_fatal.log");
|
||||
let _ = std::fs::write(
|
||||
&log_path,
|
||||
format!(
|
||||
"FATAL: Could not bind to 127.0.0.1:3000 after 15 seconds: {}\n",
|
||||
e
|
||||
),
|
||||
);
|
||||
std::process::exit(1);
|
||||
}
|
||||
let log_path = dirs::home_dir().unwrap_or_default().join(".gemini/mcp_memory/daemon_error.log");
|
||||
if let Ok(mut file) = std::fs::OpenOptions::new().create(true).append(true).open(&log_path) {
|
||||
let log_path = dirs::home_dir()
|
||||
.unwrap_or_default()
|
||||
.join(".gemini/mcp_memory/daemon_error.log");
|
||||
if let Ok(mut file) = std::fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&log_path)
|
||||
{
|
||||
use std::io::Write;
|
||||
let _ = writeln!(file, "Failed to bind to 127.0.0.1:3000 (attempt {}): {}. Retrying in 1s...", retries, e);
|
||||
let _ = writeln!(
|
||||
file,
|
||||
"Failed to bind to 127.0.0.1:3000 (attempt {}): {}. Retrying in 1s...",
|
||||
retries, e
|
||||
);
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
|
||||
}
|
||||
@@ -216,7 +261,9 @@ fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::Error>>
|
||||
};
|
||||
eprintln!("MCP Memory Server running on http://127.0.0.1:3000/sse");
|
||||
if let Err(e) = axum::serve(listener, app).await {
|
||||
let log_path = dirs::home_dir().unwrap_or_default().join(".gemini/mcp_memory/daemon_error.log");
|
||||
let log_path = dirs::home_dir()
|
||||
.unwrap_or_default()
|
||||
.join(".gemini/mcp_memory/daemon_error.log");
|
||||
let _ = std::fs::write(&log_path, format!("Server crashed: {}\n", e));
|
||||
}
|
||||
Ok(())
|
||||
@@ -228,11 +275,19 @@ async fn sse_handler(
|
||||
) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
|
||||
let session_id = format!("{}", state.next_id.fetch_add(1, Ordering::SeqCst));
|
||||
let (tx, rx) = mpsc::channel::<Result<Event, Infallible>>(100);
|
||||
|
||||
state.clients.write().unwrap().insert(session_id.clone(), tx.clone());
|
||||
|
||||
let _ = tx.send(Ok(Event::default().event("endpoint").data(format!("/messages?sessionId={}", session_id)))).await;
|
||||
|
||||
|
||||
state
|
||||
.clients
|
||||
.write()
|
||||
.unwrap()
|
||||
.insert(session_id.clone(), tx.clone());
|
||||
|
||||
let _ = tx
|
||||
.send(Ok(Event::default()
|
||||
.event("endpoint")
|
||||
.data(format!("/messages?sessionId={}", session_id))))
|
||||
.await;
|
||||
|
||||
let stream = ReceiverStream::new(rx);
|
||||
Sse::new(stream).keep_alive(axum::response::sse::KeepAlive::new())
|
||||
}
|
||||
@@ -255,32 +310,56 @@ async fn message_handler(
|
||||
let handler = Arc::clone(&state.handler);
|
||||
let session_id = query.session_id.clone();
|
||||
let clients = Arc::clone(&state);
|
||||
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Some(response) = handler.handle_request(payload).await {
|
||||
let tx_opt = clients.clients.read().unwrap().get(&session_id).cloned();
|
||||
if let Some(tx) = tx_opt {
|
||||
let data = serde_json::to_string(&response).unwrap();
|
||||
let _ = tx.send(Ok(Event::default().event("message").data(data))).await;
|
||||
let _ = tx
|
||||
.send(Ok(Event::default().event("message").data(data)))
|
||||
.await;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
|
||||
axum::http::StatusCode::ACCEPTED
|
||||
}
|
||||
|
||||
mod proxy;
|
||||
|
||||
|
||||
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let cli = Cli::parse();
|
||||
|
||||
if cli.exit {
|
||||
if let Ok(mut stream) = std::net::TcpStream::connect("127.0.0.1:3000") {
|
||||
use std::io::Write;
|
||||
let _ = stream.write_all(
|
||||
b"POST /shutdown HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n",
|
||||
);
|
||||
}
|
||||
println!("Sent shutdown request to server.");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if cli.restart {
|
||||
if let Ok(mut stream) = std::net::TcpStream::connect("127.0.0.1:3000") {
|
||||
use std::io::Write;
|
||||
let _ = stream.write_all(
|
||||
b"POST /shutdown HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n",
|
||||
);
|
||||
println!("Sent shutdown request to existing server. Waiting for it to exit...");
|
||||
std::thread::sleep(std::time::Duration::from_millis(1500));
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
#[cfg(target_os = "windows")]
|
||||
{
|
||||
use std::os::windows::process::CommandExt;
|
||||
if !cli.daemon {
|
||||
loop {
|
||||
if let Err(_) = std::net::TcpListener::bind("127.0.0.1:3000") {
|
||||
if std::net::TcpListener::bind("127.0.0.1:3000").is_err() {
|
||||
// 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) {
|
||||
@@ -293,7 +372,8 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
}
|
||||
} else {
|
||||
// Port is free. We must spawn the daemon, then loop again to become proxy
|
||||
std::process::Command::new(std::env::current_exe().unwrap())
|
||||
#[allow(clippy::zombie_processes)]
|
||||
let _ = std::process::Command::new(std::env::current_exe().unwrap())
|
||||
.arg("--daemon")
|
||||
.stdin(std::process::Stdio::null())
|
||||
.stdout(std::process::Stdio::null())
|
||||
@@ -332,7 +412,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
session_graph: RwLock::new(KnowledgeGraph::default()),
|
||||
base_dir: base.clone(),
|
||||
master_cache: RwLock::new((KnowledgeGraph::default(), SystemTime::UNIX_EPOCH)),
|
||||
search_index: RwLock::new(crate::search::MemoryIndex::new().unwrap()),
|
||||
search_index: RwLock::new(crate::search::MemoryIndex::new(&base).unwrap()),
|
||||
ledger: Store::new(base.join("audit_ledger.json")),
|
||||
sticky: Store::new(base.join("sticky_notes.json")),
|
||||
tasks: Store::new(base.join("tasks.json")),
|
||||
@@ -357,75 +437,101 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
|
||||
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());
|
||||
}
|
||||
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 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);
|
||||
}
|
||||
}
|
||||
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);
|
||||
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 {
|
||||
if let Some(r) = record.reason {
|
||||
eprintln!("❌ Action blocked. Reason: {}", r);
|
||||
} else {
|
||||
eprintln!("❌ Action blocked.");
|
||||
}
|
||||
std::process::exit(1);
|
||||
eprintln!("❌ Action blocked.");
|
||||
}
|
||||
std::process::exit(1);
|
||||
}
|
||||
None => {
|
||||
eprintln!("❌ Action not yet authorized (no gate record found).");
|
||||
std::process::exit(2);
|
||||
}
|
||||
}
|
||||
None => {
|
||||
eprintln!("❌ Action not yet authorized (no gate record found).");
|
||||
std::process::exit(2);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in new issue
Block a user