perf: fix memory leaks, optimize tracing logs, fix networking eof fragmentation bug in nvim-core and stub, and deduplicate json-rpc handlers
This commit is contained in:
Riz Ashraf committed
1 parent
ce88bef45e
commit
f76d8ccce8
23 files changed
+1373
-803
No files matched your search
+129
-635
@@ -3,29 +3,28 @@
|
||||
allow(dead_code, unused_imports, unreachable_code)
|
||||
)]
|
||||
|
||||
mod api;
|
||||
mod handlers;
|
||||
mod handlers_v2;
|
||||
mod mcp;
|
||||
mod models;
|
||||
mod router;
|
||||
mod search;
|
||||
pub mod db;
|
||||
pub mod error;
|
||||
mod state;
|
||||
mod store;
|
||||
mod tools;
|
||||
|
||||
use crate::handlers::MemoryHandler;
|
||||
use crate::models::*;
|
||||
use crate::api::rest::GateSetReq;
|
||||
use crate::router::MemoryHandler;
|
||||
use crate::state::MemoryState;
|
||||
use crate::store::Store;
|
||||
|
||||
use redb::ReadableTable;
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::time::Duration;
|
||||
|
||||
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)]
|
||||
@@ -87,6 +86,12 @@ enum GateCommands {
|
||||
},
|
||||
}
|
||||
|
||||
pub struct AppState {
|
||||
handler: Arc<MemoryHandler>,
|
||||
clients: RwLock<HashMap<String, mpsc::Sender<String>>>,
|
||||
next_id: AtomicUsize,
|
||||
}
|
||||
|
||||
async fn index_committer_worker(state: Arc<MemoryState>) {
|
||||
loop {
|
||||
tokio::time::sleep(Duration::from_secs(5)).await;
|
||||
@@ -98,121 +103,6 @@ async fn index_committer_worker(state: Arc<MemoryState>) {
|
||||
}
|
||||
}
|
||||
|
||||
use axum::{
|
||||
Json, Router,
|
||||
extract::{
|
||||
Query, State,
|
||||
ws::{Message, WebSocket},
|
||||
},
|
||||
response::IntoResponse,
|
||||
routing::{get, post},
|
||||
};
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
struct AppState {
|
||||
handler: Arc<MemoryHandler>,
|
||||
clients: RwLock<HashMap<String, mpsc::Sender<String>>>,
|
||||
next_id: AtomicUsize,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct GateVerifyReq {
|
||||
action: String,
|
||||
target: String,
|
||||
namespace: Option<String>,
|
||||
#[serde(default)]
|
||||
params: HashMap<String, String>,
|
||||
#[serde(default)]
|
||||
consume: bool,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct GateSetReq {
|
||||
action: String,
|
||||
target: String,
|
||||
namespace: Option<String>,
|
||||
#[serde(default)]
|
||||
params: HashMap<String, String>,
|
||||
authorize: Option<bool>,
|
||||
block: Option<bool>,
|
||||
reason: Option<String>,
|
||||
}
|
||||
|
||||
async fn gate_verify_handler(
|
||||
State(app_state): State<Arc<AppState>>,
|
||||
Query(q): Query<GateVerifyReq>,
|
||||
) -> axum::response::Response {
|
||||
let mut found = None;
|
||||
let mut to_remove = None;
|
||||
app_state.handler.state.gates.modify(|gates| {
|
||||
if let Some(idx) = gates.iter().position(|g| {
|
||||
g.action == q.action
|
||||
&& g.target == q.target
|
||||
&& g.namespace == q.namespace
|
||||
&& g.params == q.params
|
||||
}) {
|
||||
found = Some(gates[idx].clone());
|
||||
if q.consume {
|
||||
to_remove = Some(idx);
|
||||
}
|
||||
}
|
||||
if let Some(idx) = to_remove {
|
||||
gates.remove(idx);
|
||||
}
|
||||
});
|
||||
|
||||
match found {
|
||||
Some(record) => {
|
||||
if record.status == "authorized" {
|
||||
(axum::http::StatusCode::OK, "Authorized").into_response()
|
||||
} else {
|
||||
let msg = if let Some(r) = record.reason {
|
||||
format!("Action blocked. Reason: {}", r)
|
||||
} else {
|
||||
"Action blocked.".to_string()
|
||||
};
|
||||
(axum::http::StatusCode::FORBIDDEN, msg).into_response()
|
||||
}
|
||||
}
|
||||
None => (
|
||||
axum::http::StatusCode::NOT_FOUND,
|
||||
"Action not yet authorized (no gate record found).",
|
||||
)
|
||||
.into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn gate_set_handler(
|
||||
State(app_state): State<Arc<AppState>>,
|
||||
Json(body): Json<GateSetReq>,
|
||||
) -> axum::response::Response {
|
||||
let status = if body.block.unwrap_or(false) {
|
||||
"blocked".to_string()
|
||||
} else if body.authorize.unwrap_or(false) {
|
||||
"authorized".to_string()
|
||||
} else {
|
||||
"pending".to_string()
|
||||
};
|
||||
|
||||
let record = GateRecord {
|
||||
id: uuid::Uuid::new_v4().to_string(),
|
||||
action: body.action.clone(),
|
||||
target: body.target.clone(),
|
||||
namespace: body.namespace.clone(),
|
||||
params: body.params.clone(),
|
||||
status,
|
||||
reason: body.reason.clone(),
|
||||
timestamp: crate::handlers_v2::utils::now_secs(),
|
||||
};
|
||||
app_state.handler.state.gates.modify(|gates| {
|
||||
gates.retain(|g| !(g.action == record.action && g.target == record.target));
|
||||
gates.push(record);
|
||||
});
|
||||
(axum::http::StatusCode::OK, "Gate state updated.").into_response()
|
||||
}
|
||||
|
||||
async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::Error>> {
|
||||
state.rebuild_index().await;
|
||||
tokio::spawn(index_committer_worker(Arc::clone(&state)));
|
||||
@@ -239,191 +129,7 @@ async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::E
|
||||
}
|
||||
});
|
||||
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/api/version",
|
||||
get(|| async move {
|
||||
axum::Json(serde_json::json!({
|
||||
"version": env!("APP_VERSION"),
|
||||
"git_hash": option_env!("GIT_HASH").unwrap_or("unknown")
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route("/ws", get(ws_handler))
|
||||
.route("/health", get(health_handler))
|
||||
.route("/nvim/telemetry", post(nvim_telemetry_handler))
|
||||
.route("/gate/verify", get(gate_verify_handler))
|
||||
.route("/gate/set", post(gate_set_handler))
|
||||
.route(
|
||||
"/shutdown",
|
||||
post(
|
||||
|headers: axum::http::HeaderMap, State(state): State<Arc<AppState>>| async move {
|
||||
let token_path = state.handler.state.base_dir.join("admin.token");
|
||||
let expected_token = tokio::fs::read_to_string(&token_path)
|
||||
.await
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_string();
|
||||
|
||||
let auth_header = headers
|
||||
.get(axum::http::header::AUTHORIZATION)
|
||||
.and_then(|h| h.to_str().ok())
|
||||
.unwrap_or_default();
|
||||
|
||||
if expected_token.is_empty() || auth_header != format!("Bearer {}", expected_token) {
|
||||
return (axum::http::StatusCode::UNAUTHORIZED, "Unauthorized").into_response();
|
||||
}
|
||||
|
||||
std::thread::spawn(|| {
|
||||
tracing::info!(
|
||||
"Received shutdown request via /shutdown endpoint. Exiting process cleanly."
|
||||
);
|
||||
std::thread::sleep(std::time::Duration::from_millis(100));
|
||||
std::process::exit(0);
|
||||
});
|
||||
(axum::http::StatusCode::OK, "Shutting down...").into_response()
|
||||
},
|
||||
),
|
||||
)
|
||||
.route(
|
||||
"/",
|
||||
get(|| async move { axum::response::Html(include_str!("dashboard.html")) }),
|
||||
)
|
||||
.route(
|
||||
"/api/graph",
|
||||
get({
|
||||
let state_clone = app_state.handler.state.clone();
|
||||
move || async move {
|
||||
let graph_json = state_clone.read_graph(|g| serde_json::to_string(g).unwrap_or_else(|_| "{}".to_string()));
|
||||
([(axum::http::header::CONTENT_TYPE, "application/json")], graph_json)
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/tasks/{id}/complete",
|
||||
post({
|
||||
let state_clone = app_state.handler.state.clone();
|
||||
move |axum::extract::Path(id): axum::extract::Path<String>| async move {
|
||||
state_clone.tasks.modify(|tasks| {
|
||||
for t in tasks.iter_mut() {
|
||||
if t.id == id {
|
||||
t.status = "completed".to_string();
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
axum::Json(serde_json::json!({"status": "success"}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/tasks",
|
||||
get({
|
||||
let state_clone = app_state.handler.state.clone();
|
||||
move || async move {
|
||||
let tasks_json = state_clone.tasks.read_with(|t| serde_json::to_string(t).unwrap_or_else(|_| "[]".to_string()));
|
||||
([(axum::http::header::CONTENT_TYPE, "application/json")], tasks_json)
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/sticky",
|
||||
get({
|
||||
let state_clone = app_state.handler.state.clone();
|
||||
move || async move {
|
||||
let sticky_json = state_clone.sticky.read_with(|s| serde_json::to_string(s).unwrap_or_else(|_| "[]".to_string()));
|
||||
([(axum::http::header::CONTENT_TYPE, "application/json")], sticky_json)
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/search",
|
||||
get({
|
||||
let state_clone = app_state.handler.state.clone();
|
||||
move |axum::extract::Query(params): axum::extract::Query<
|
||||
std::collections::HashMap<String, String>,
|
||||
>| async move {
|
||||
if let Some(q) = params.get("q")
|
||||
&& let Ok(idx) = state_clone.search_index.read()
|
||||
&& let Ok(results) = idx.search(q, None) {
|
||||
let mut formatted_results = Vec::new();
|
||||
for (id, doc_type, title, body, score) in results {
|
||||
formatted_results.push(serde_json::json!({
|
||||
"id": id,
|
||||
"type_name": doc_type,
|
||||
"title": title,
|
||||
"content": body,
|
||||
"score": score
|
||||
}));
|
||||
}
|
||||
return axum::Json(
|
||||
serde_json::json!({ "results": formatted_results }),
|
||||
);
|
||||
}
|
||||
axum::Json(serde_json::json!({ "results": [] }))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/activity",
|
||||
get({
|
||||
let state_clone = app_state.handler.state.clone();
|
||||
move || async move {
|
||||
let activities_json = state_clone.recent_activities.read_with(|a| serde_json::to_string(a).unwrap_or_else(|_| "[]".to_string()));
|
||||
([(axum::http::header::CONTENT_TYPE, "application/json")], activities_json)
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/api/stats",
|
||||
get({
|
||||
let state_clone = app_state.handler.state.clone();
|
||||
move || async move {
|
||||
let (entities, relations) = state_clone.read_graph(|g| (g.entities.len(), g.relations.len()));
|
||||
let tasks = state_clone.tasks.read_with(|items| items.len());
|
||||
let snippets = state_clone.snippets.read_with(|items| items.len());
|
||||
let tech_debts = state_clone.tech_debts.read_with(|items| items.len());
|
||||
let adrs = state_clone.adrs.read_with(|items| items.len());
|
||||
|
||||
let ledger = state_clone.ledger.read_with(|items| items.len());
|
||||
let sticky = state_clone.sticky.read_with(|items| items.len());
|
||||
let error_fixes = state_clone.error_fixes.read_with(|items| items.len());
|
||||
let pinned_files = state_clone.pinned_files.read_with(|items| items.len());
|
||||
let session_summaries = state_clone.session_summaries.read_with(|items| items.len());
|
||||
let handoff_memos = state_clone.handoff_memos.read_with(|items| items.len());
|
||||
let env_fingerprints = state_clone.env_fingerprints.read_with(|items| items.len());
|
||||
let env_requirements = state_clone.env_requirements.read_with(|items| items.len());
|
||||
let milestones = state_clone.milestones.read_with(|items| items.len());
|
||||
let environments = state_clone.environments.read_with(|items| items.len());
|
||||
let pr_checklists = state_clone.pr_checklists.read_with(|items| items.len());
|
||||
let gates = state_clone.gates.read_with(|items| items.len());
|
||||
let context_workspaces = state_clone.context_workspaces.read_with(|items| items.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 app = api::setup::create_router(app_state);
|
||||
|
||||
tracing::info!("MCP Memory Server running on http://127.0.0.1:3000/sse");
|
||||
let addr = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
|
||||
@@ -451,201 +157,6 @@ async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::E
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn ws_handler(
|
||||
ws: axum::extract::ws::WebSocketUpgrade,
|
||||
_headers: axum::http::HeaderMap,
|
||||
axum::extract::State(state): axum::extract::State<Arc<AppState>>,
|
||||
axum::extract::Query(query): axum::extract::Query<std::collections::HashMap<String, String>>,
|
||||
) -> axum::response::Response {
|
||||
let client_type = query
|
||||
.get("client")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
ws.on_upgrade(move |socket| handle_socket(socket, state, client_type))
|
||||
.into_response()
|
||||
}
|
||||
|
||||
async fn handle_socket(socket: WebSocket, state: Arc<AppState>, _client_type: String) {
|
||||
let session_id = format!("{}", state.next_id.fetch_add(1, Ordering::SeqCst));
|
||||
let (tx, mut rx) = mpsc::channel::<String>(100);
|
||||
|
||||
state
|
||||
.clients
|
||||
.write()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.insert(session_id.clone(), tx.clone());
|
||||
|
||||
let (mut sender, mut receiver) = socket.split();
|
||||
|
||||
let send_task = tokio::spawn(async move {
|
||||
while let Some(msg) = rx.recv().await {
|
||||
tracing::trace!(
|
||||
"Sending message to websocket (length: {}): {}",
|
||||
msg.len(),
|
||||
msg
|
||||
);
|
||||
if sender.send(Message::Text(msg.into())).await.is_err() {
|
||||
tracing::error!("Failed to send message to websocket");
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Premature list_changed notification removed for MCP protocol compliance
|
||||
|
||||
let handler = Arc::clone(&state.handler);
|
||||
let state_clone = Arc::clone(&state);
|
||||
let session_id_clone = session_id.clone();
|
||||
|
||||
let recv_task = tokio::spawn(async move {
|
||||
while let Some(msg_result) = receiver.next().await {
|
||||
match msg_result {
|
||||
Ok(Message::Text(text)) => {
|
||||
tracing::info!(
|
||||
"Received text message from websocket (length: {})",
|
||||
text.len()
|
||||
);
|
||||
tracing::trace!("Message content: {}", text);
|
||||
if let Ok(payload) = serde_json::from_str::<serde_json::Value>(&text) {
|
||||
// Process MCP request
|
||||
if let Some(response) = handler.handle_request(payload).await {
|
||||
let res_str = serde_json::to_string(&response).unwrap_or_default();
|
||||
let tx_opt = state_clone
|
||||
.clients
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.get(&session_id_clone)
|
||||
.cloned();
|
||||
if let Some(client_tx) = tx_opt {
|
||||
if let Err(e) = client_tx.send(res_str).await {
|
||||
tracing::error!(
|
||||
"Failed to send response to client channel for session {}: {}",
|
||||
session_id_clone,
|
||||
e
|
||||
);
|
||||
}
|
||||
} else {
|
||||
tracing::warn!(
|
||||
"Could not find client_tx for session_id {} when trying to send response",
|
||||
session_id_clone
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
// End if let Ok(payload)
|
||||
else {
|
||||
tracing::warn!(
|
||||
"Failed to parse payload as JSON from websocket message: {}",
|
||||
text
|
||||
);
|
||||
}
|
||||
} // End Ok(Message::Text(text))
|
||||
Ok(other) => {
|
||||
tracing::info!("Received non-text message from websocket: {:?}", other);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("Websocket receive error: {}", e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
tracing::info!(
|
||||
"Websocket receiver task ended for session {}",
|
||||
session_id_clone
|
||||
);
|
||||
});
|
||||
|
||||
struct SessionCleanup {
|
||||
session_id: String,
|
||||
state: Arc<AppState>,
|
||||
send_task: tokio::task::JoinHandle<()>,
|
||||
recv_task: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
impl Drop for SessionCleanup {
|
||||
fn drop(&mut self) {
|
||||
self.state
|
||||
.clients
|
||||
.write()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.remove(&self.session_id);
|
||||
self.send_task.abort();
|
||||
self.recv_task.abort();
|
||||
tracing::info!(
|
||||
"Websocket session {} closed and cleaned up",
|
||||
self.session_id
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let mut cleanup = SessionCleanup {
|
||||
session_id: session_id.clone(),
|
||||
state: Arc::clone(&state),
|
||||
send_task,
|
||||
recv_task,
|
||||
};
|
||||
|
||||
tokio::select! {
|
||||
_ = &mut cleanup.send_task => {
|
||||
tracing::info!("Websocket send task finished for session {}", session_id);
|
||||
},
|
||||
_ = &mut cleanup.recv_task => {
|
||||
tracing::info!("Websocket recv task finished for session {}", session_id);
|
||||
},
|
||||
};
|
||||
// Drop guard automatically handles removal and aborts the other task.
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize, serde::Serialize, Debug)]
|
||||
pub struct NvimTelemetry {
|
||||
pub session_id: String,
|
||||
pub event: String,
|
||||
pub file: Option<String>,
|
||||
pub line: Option<i64>,
|
||||
pub col: Option<i64>,
|
||||
}
|
||||
|
||||
async fn nvim_telemetry_handler(
|
||||
State(state): State<Arc<AppState>>,
|
||||
axum::Json(payload): axum::Json<NvimTelemetry>,
|
||||
) -> impl axum::response::IntoResponse {
|
||||
// 1. Update active_nvim.txt if FocusGained, BufEnter, or VimEnter
|
||||
if payload.event == "FocusGained" || payload.event == "BufEnter" || payload.event == "VimEnter"
|
||||
{
|
||||
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, &payload.session_id).await;
|
||||
|
||||
let wsl_path = "\\\\wsl.localhost\\Ubuntu\\home\\riz\\.gemini\\active_nvim.txt";
|
||||
let _ = tokio::fs::write(wsl_path, &payload.session_id).await;
|
||||
}
|
||||
|
||||
// 2. Broadcast to UI WebSockets
|
||||
let ws_msg = serde_json::json!({
|
||||
"type": "nvim_telemetry",
|
||||
"data": payload
|
||||
});
|
||||
|
||||
let msg_str = ws_msg.to_string();
|
||||
let senders: Vec<_> = state
|
||||
.clients
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.values()
|
||||
.cloned()
|
||||
.collect();
|
||||
for tx in senders {
|
||||
let _ = tx.try_send(msg_str.clone());
|
||||
}
|
||||
|
||||
axum::Json(serde_json::json!({"status": "ok"}))
|
||||
}
|
||||
|
||||
async fn health_handler() -> &'static str {
|
||||
"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()
|
||||
@@ -696,148 +207,131 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
.arg(format!("Authorization: Bearer {}", token.trim()));
|
||||
}
|
||||
let _ = cmd
|
||||
.arg(format!("https://127.0.0.1:{}/shutdown", port))
|
||||
.arg(format!("http://127.0.0.1:{}/shutdown", port))
|
||||
.output();
|
||||
println!("Sent shutdown request to server.");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if cli.restart {
|
||||
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!("https://127.0.0.1:{}/shutdown", port))
|
||||
.output();
|
||||
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 {
|
||||
// Just spawn the daemon and exit. We no longer act as a proxy.
|
||||
#[allow(clippy::zombie_processes)]
|
||||
let _ = std::process::Command::new(
|
||||
std::env::current_exe().expect("Failed to get current executable path"),
|
||||
)
|
||||
.arg("--daemon")
|
||||
.stdin(std::process::Stdio::null())
|
||||
.stdout(std::process::Stdio::null())
|
||||
.stderr(std::process::Stdio::null())
|
||||
.creation_flags(0x08000000) // CREATE_NO_WINDOW
|
||||
.spawn()
|
||||
.expect("Failed to spawn daemon");
|
||||
if cli.restart {
|
||||
std::thread::sleep(Duration::from_secs(2));
|
||||
} else {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
fs::create_dir_all(&base).expect("Failed to create store dir");
|
||||
|
||||
// Generate token
|
||||
let admin_token = uuid::Uuid::new_v4().to_string();
|
||||
std::fs::write(base.join("admin.token"), &admin_token).expect("Failed to write admin token");
|
||||
|
||||
let redb_path = base.join("mcp_store.redb");
|
||||
|
||||
let db = Arc::new(redb::Database::create(&redb_path).expect("Failed to create redb database"));
|
||||
|
||||
// Ensure table exists and migrate old JSON files
|
||||
{
|
||||
let write_txn = db.begin_write().expect("Failed to begin write txn on redb");
|
||||
{
|
||||
let mut table = write_txn
|
||||
.open_table(crate::store::STORE_TABLE)
|
||||
.expect("Failed to open STORE_TABLE");
|
||||
|
||||
let stores = vec![
|
||||
("knowledge_graph_master", "knowledge_graph_master.json"),
|
||||
("audit_ledger", "audit_ledger.json"),
|
||||
("sticky_notes", "sticky_notes.json"),
|
||||
("tasks", "tasks.json"),
|
||||
("snippets", "snippets.json"),
|
||||
("adrs", "adrs.json"),
|
||||
("preferences", "preferences.json"),
|
||||
("error_fixes", "error_fixes.json"),
|
||||
("pinned_files", "pinned_files.json"),
|
||||
("session_summaries", "session_summaries.json"),
|
||||
("handoff_memos", "handoff_memos.json"),
|
||||
("env_fingerprints", "env_fingerprints.json"),
|
||||
("env_requirements", "env_requirements.json"),
|
||||
("milestones", "milestones.json"),
|
||||
("environments", "environments.json"),
|
||||
("pr_checklists", "pr_checklists.json"),
|
||||
("tech_debts", "tech_debts.json"),
|
||||
("gates", "gates.json"),
|
||||
("context_workspaces", "context_workspaces.json"),
|
||||
];
|
||||
|
||||
for (key, file_name) in stores.iter() {
|
||||
if table
|
||||
.get(*key)
|
||||
.expect("Failed to read from table")
|
||||
.is_none()
|
||||
{
|
||||
let json_path = base.join(file_name);
|
||||
if json_path.exists()
|
||||
&& let Ok(data) = fs::read(&json_path)
|
||||
&& serde_json::from_slice::<serde_json::Value>(&data).is_ok()
|
||||
{
|
||||
table
|
||||
.insert(*key, data.as_slice())
|
||||
.expect("Failed to insert migrated data");
|
||||
let _ = fs::rename(&json_path, json_path.with_extension("json.migrated"));
|
||||
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);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
write_txn.commit().expect("Failed to commit db migration");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let rt = tokio::runtime::Runtime::new().expect("Failed to create tokio runtime");
|
||||
let _guard = rt.enter();
|
||||
#[cfg(target_os = "windows")]
|
||||
if cli.daemon {
|
||||
let exe = std::env::current_exe()?;
|
||||
std::process::Command::new("powershell")
|
||||
.args([
|
||||
"-WindowStyle",
|
||||
"Hidden",
|
||||
"-Command",
|
||||
&format!(
|
||||
"Start-Process -FilePath '{}' -WindowStyle Hidden",
|
||||
exe.display()
|
||||
),
|
||||
])
|
||||
.spawn()?;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let state = Arc::new(MemoryState {
|
||||
graph: Store::new("knowledge_graph_master", db.clone()),
|
||||
base_dir: base.clone(),
|
||||
search_index: RwLock::new(match crate::search::MemoryIndex::new(&base) {
|
||||
Ok(idx) => idx,
|
||||
Err(e) => {
|
||||
let log_path = dirs::home_dir()
|
||||
.unwrap_or_default()
|
||||
.join(".gemini/mcp_memory/daemon_error.log");
|
||||
let _ = std::fs::write(&log_path, format!("Failed to create MemoryIndex: {}\n", e));
|
||||
std::process::exit(1);
|
||||
}
|
||||
}),
|
||||
ledger: Store::new("audit_ledger", db.clone()),
|
||||
sticky: Store::new("sticky_notes", db.clone()),
|
||||
tasks: Store::new("tasks", db.clone()),
|
||||
snippets: Store::new("snippets", db.clone()),
|
||||
adrs: Store::new("adrs", db.clone()),
|
||||
prefs: Store::new("preferences", db.clone()),
|
||||
error_fixes: Store::new("error_fixes", db.clone()),
|
||||
pinned_files: Store::new("pinned_files", db.clone()),
|
||||
session_summaries: Store::new("session_summaries", db.clone()),
|
||||
handoff_memos: Store::new("handoff_memos", db.clone()),
|
||||
env_fingerprints: Store::new("env_fingerprints", db.clone()),
|
||||
env_requirements: Store::new("env_requirements", db.clone()),
|
||||
milestones: Store::new("milestones", db.clone()),
|
||||
environments: Store::new("environments", db.clone()),
|
||||
pr_checklists: Store::new("pr_checklists", db.clone()),
|
||||
tech_debts: Store::new("tech_debts", db.clone()),
|
||||
gates: Store::new("gates", db.clone()),
|
||||
context_workspaces: Store::new("context_workspaces", db.clone()),
|
||||
recent_activities: Store::new("recent_activities", db.clone()),
|
||||
activity_tx: tokio::sync::broadcast::channel(100).0,
|
||||
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);
|
||||
}
|
||||
});
|
||||
|
||||
rt.block_on(run_server(state))
|
||||
Ok(())
|
||||
}
|
||||
Reference in new issue
Block a user