perf: fix memory leaks, optimize tracing logs, fix networking eof fragmentation bug in nvim-core and stub, and deduplicate json-rpc handlers

server/Cargo.toml

@@ -15,7 +15,7 @@ glob = "0.3.4"
 git2 = { version = "0.19.0", default-features = false }
 notify = "6.1.1"
 redb = "4.2.0"
-reqwest = { version = "0.12", default-features = false, features = ["stream", "rustls-tls"] }
+reqwest = { version = "0.12", default-features = false, features = ["stream", "rustls-tls", "json"] }
 schemars = "1.2.2"
 serde = { version = "1.0.229", features = ["derive"] }
 serde_json = "1.0.151"
@@ -28,6 +28,7 @@ tracing-subscriber = "0.3.23"
 uuid = { version = "1.26.0", features = ["v4"] }
 tracing-appender = "0.2.5"
 rmcp = { version = "3.4.0", features = ["server"] }
+thiserror = "2.0.20"
 
 [build-dependencies]
 chrono = "0.4.45"
@@ -35,8 +36,3 @@ winres = "0.1.12"
 
 [dev-dependencies]
 tempfile = "3.27.0"
-
-[[bin]]
-name = "test_rmcp"
-path = "src/bin_test.rs"
-

server/src/api/mod.rs

@@ -0,0 +1,4 @@
+pub mod rest;
+pub mod setup;
+pub mod telemetry;
+pub mod ws;

server/src/api/rest.rs

@@ -0,0 +1,107 @@
+use crate::AppState;
+use crate::error::AppError;
+use crate::models::GateRecord;
+use axum::{
+    Json,
+    extract::{Query, State},
+    response::IntoResponse,
+};
+use std::collections::HashMap;
+use std::sync::Arc;
+
+#[derive(serde::Deserialize, serde::Serialize)]
+pub struct GateVerifyReq {
+    pub action: String,
+    pub target: String,
+    pub namespace: Option<String>,
+    #[serde(default)]
+    pub params: HashMap<String, String>,
+    #[serde(default)]
+    pub consume: bool,
+}
+
+#[derive(serde::Deserialize, serde::Serialize)]
+pub struct GateSetReq {
+    pub action: String,
+    pub target: String,
+    pub namespace: Option<String>,
+    #[serde(default)]
+    pub params: HashMap<String, String>,
+    pub authorize: Option<bool>,
+    pub block: Option<bool>,
+    pub reason: Option<String>,
+}
+
+pub async fn gate_verify_handler(
+    State(app_state): State<Arc<AppState>>,
+    Query(q): Query<GateVerifyReq>,
+) -> Result<impl IntoResponse, AppError> {
+    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" {
+                Ok((axum::http::StatusCode::OK, "Authorized"))
+            } else {
+                let msg = if let Some(r) = record.reason {
+                    format!("Action blocked. Reason: {}", r)
+                } else {
+                    "Action blocked.".to_string()
+                };
+                Err(AppError::Forbidden(msg))
+            }
+        }
+        None => Err(AppError::NotFound("Action not yet authorized (no gate record found).".to_string())),
+    }
+}
+
+pub async fn gate_set_handler(
+    State(app_state): State<Arc<AppState>>,
+    Json(body): Json<GateSetReq>,
+) -> Result<impl IntoResponse, AppError> {
+    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::utils::now_secs(),
+    };
+    app_state.handler.state.gates.modify(|gates| {
+        gates.retain(|g| !(g.action == record.action && g.target == record.target));
+        gates.push(record);
+    });
+    
+    Ok((axum::http::StatusCode::OK, "Gate state updated."))
+}
+
+pub async fn health_handler() -> &'static str {
+    "OK"
+}

server/src/api/routes.txt

@@ -0,0 +1,301 @@
+        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::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)));
+    let app_state = Arc::new(AppState {
+        handler: Arc::new(MemoryHandler::new(Arc::clone(&state))),
+        clients: RwLock::new(HashMap::new()),
+        next_id: AtomicUsize::new(1),
+    });
+
+    let app_state_clone = Arc::clone(&app_state);
+    let mut rx = state.activity_tx.subscribe();
+    tokio::spawn(async move {
+        while let Ok(msg) = rx.recv().await {
+            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());
+            }
+        }
+    });
+
+    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);
+
+    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());
+    let addr: std::net::SocketAddr = format!("127.0.0.1:{}", addr)
+        .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()).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;
+    }

server/src/api/setup.rs

@@ -0,0 +1,199 @@
+use crate::AppState;
+use crate::api::rest::{gate_set_handler, gate_verify_handler, health_handler};
+use crate::api::telemetry::nvim_telemetry_handler;
+use crate::api::ws::ws_handler;
+use axum::{
+    Router,
+    extract::State,
+    response::IntoResponse,
+    routing::{get, post},
+};
+use std::sync::Arc;
+
+pub fn create_router(app_state: Arc<AppState>) -> Router {
+    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)
+}

server/src/api/telemetry.rs

@@ -0,0 +1,49 @@
+use crate::AppState;
+use axum::extract::State;
+use std::sync::Arc;
+
+#[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>,
+}
+
+pub 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"}))
+}

server/src/api/ws.rs

@@ -0,0 +1,144 @@
+use crate::AppState;
+use axum::extract::{
+    Query, State,
+    ws::{Message, WebSocket},
+};
+use axum::response::IntoResponse;
+use futures_util::{SinkExt, StreamExt};
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::sync::atomic::Ordering;
+use tokio::sync::mpsc;
+
+pub async fn ws_handler(
+    ws: axum::extract::ws::WebSocketUpgrade,
+    _headers: axum::http::HeaderMap,
+    State(state): State<Arc<AppState>>,
+    Query(query): Query<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()
+}
+
+pub struct SessionCleanup {
+    pub session_id: String,
+    pub state: Arc<AppState>,
+    pub send_task: tokio::task::JoinHandle<()>,
+    pub recv_task: tokio::task::JoinHandle<()>,
+}
+
+impl Drop for SessionCleanup {
+    fn drop(&mut self) {
+        tracing::info!("Dropping session {}", self.session_id);
+        self.state
+            .clients
+            .write()
+            .unwrap_or_else(|e| e.into_inner())
+            .remove(&self.session_id);
+        self.send_task.abort();
+        self.recv_task.abort();
+    }
+}
+
+pub 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;
+            }
+        }
+    });
+
+    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
+                                );
+                            }
+                        }
+                    } else {
+                        tracing::warn!(
+                            "Failed to parse payload as JSON from websocket message: {}",
+                            text
+                        );
+                    }
+                }
+                Ok(other) => {
+                    tracing::info!("Received non-text message from websocket: {:?}", other);
+                }
+                Err(e) => {
+                    tracing::error!("Websocket receive error: {}", e);
+                    break;
+                }
+            }
+        }
+    });
+
+    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);
+        },
+    };
+}

server/src/db.rs

@@ -0,0 +1,64 @@
+use std::path::Path;
+use std::sync::Arc;
+use redb::{Database, ReadableTable};
+use crate::store::STORE_TABLE;
+
+pub fn init_redb(base: &Path) -> Arc<Database> {
+    let redb_path = base.join("mcp_store.redb");
+    let db = Arc::new(redb::Database::create(&redb_path).expect("Failed to create redb database"));
+
+    // Ensure the table exists and migrate legacy JSON files
+    let write_txn = db.begin_write().expect("Failed to begin write txn on redb");
+    {
+        let mut table = write_txn
+            .open_table(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) = std::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 _ = std::fs::rename(
+                        &json_path,
+                        json_path.with_extension("json.migrated"),
+                    );
+                }
+            }
+        }
+    }
+    write_txn.commit().expect("Failed to commit db migration");
+
+    db
+}

server/src/error.rs

@@ -0,0 +1,39 @@
+use axum::{
+    http::StatusCode,
+    response::{IntoResponse, Response},
+    Json,
+};
+use serde_json::json;
+use thiserror::Error;
+
+#[derive(Error, Debug)]
+pub enum AppError {
+    #[error("Not Found: {0}")]
+    NotFound(String),
+
+    #[error("Forbidden: {0}")]
+    Forbidden(String),
+
+    #[error("Internal Server Error: {0}")]
+    Internal(String),
+
+    #[error("Bad Request: {0}")]
+    BadRequest(String),
+}
+
+impl IntoResponse for AppError {
+    fn into_response(self) -> Response {
+        let (status, error_message) = match &self {
+            AppError::NotFound(msg) => (StatusCode::NOT_FOUND, msg.clone()),
+            AppError::Forbidden(msg) => (StatusCode::FORBIDDEN, msg.clone()),
+            AppError::Internal(msg) => (StatusCode::INTERNAL_SERVER_ERROR, msg.clone()),
+            AppError::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg.clone()),
+        };
+
+        let body = Json(json!({
+            "error": error_message,
+        }));
+
+        (status, body).into_response()
+    }
+}

server/src/handlers/env.rs

@@ -31,7 +31,7 @@ impl McpTool for UpdateEnvFingerprintHandler {
                     os: std::env::consts::OS.to_string(),
                     shell: std::env::var("SHELL").unwrap_or_else(|_| "unknown".to_string()),
                     tool_versions: req.tool_versions,
-                    updated_at: crate::handlers_v2::utils::now_secs(),
+                    updated_at: crate::handlers::utils::now_secs(),
                 },
             );
         });
@@ -125,7 +125,7 @@ impl McpTool for RegisterEnvironmentHandler {
                 url: req.url,
                 description: req.description,
                 requires_vpn: req.requires_vpn,
-                updated_at: crate::handlers_v2::utils::now_secs(),
+                updated_at: crate::handlers::utils::now_secs(),
             });
         });
         Ok("Environment registered".to_string())

server/src/handlers/graph.rs

@@ -576,4 +576,4 @@ impl McpTool for FindOrphansHandler {
     }
 }
 
-use crate::handlers_v2::utils::*;
+use crate::handlers::utils::*;

server/src/handlers/meta.rs

@@ -36,7 +36,7 @@ impl McpTool for LogDecisionHandler {
                 context: req.context,
                 decision: req.decision,
                 consequence: req.consequence,
-                timestamp: crate::handlers_v2::utils::now_secs(),
+                timestamp: crate::handlers::utils::now_secs(),
             };
 
             drop(idx.index_adr(&a));
@@ -98,7 +98,7 @@ impl McpTool for LogErrorFixHandler {
             fixes.push(crate::models::ErrorFix {
                 signature: req.signature,
                 solution: req.solution,
-                timestamp: crate::handlers_v2::utils::now_secs(),
+                timestamp: crate::handlers::utils::now_secs(),
                 git_commit: req.git_commit,
                 git_branch: req.git_branch,
             })
@@ -155,7 +155,7 @@ impl McpTool for LogCodeChangeHandler {
         let req: LogCodeChangeTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
         state.ledger.modify(|ledger| {
             ledger.push(CodeChange {
-                timestamp: crate::handlers_v2::utils::now_secs(),
+                timestamp: crate::handlers::utils::now_secs(),
                 file_path: req.file_path,
                 description: req.description,
                 git_commit: req.git_commit,
@@ -209,7 +209,7 @@ impl McpTool for LearnPreferenceHandler {
                 crate::models::Preference {
                     key: req.key.clone(),
                     value: req.value,
-                    updated_at: crate::handlers_v2::utils::now_secs(),
+                    updated_at: crate::handlers::utils::now_secs(),
                 },
             );
         });
@@ -258,7 +258,7 @@ impl McpTool for LogTechDebtHandler {
                 description: req.description,
                 ideal_solution: req.ideal_solution,
                 is_resolved: false,
-                created_at: crate::handlers_v2::utils::now_secs(),
+                created_at: crate::handlers::utils::now_secs(),
                 git_commit: req.git_commit,
                 git_branch: req.git_branch,
             })
@@ -501,4 +501,4 @@ impl McpTool for GetProjectHealthHandler {
     }
 }
 
-use crate::handlers_v2::utils::*;
+use crate::handlers::utils::*;

server/src/handlers/notes.rs

@@ -23,7 +23,7 @@ impl McpTool for AddStickyNoteHandler {
         let req: AddStickyNoteTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
         state.sticky.modify(|notes| {
             notes.push(StickyNote {
-                timestamp: crate::handlers_v2::utils::now_secs(),
+                timestamp: crate::handlers::utils::now_secs(),
                 content: req.content,
             });
         });
@@ -132,7 +132,7 @@ impl McpTool for LeaveHandoffMemoHandler {
                 author: "agy".to_string(),
                 content: req.content,
                 namespace: req.namespace,
-                timestamp: crate::handlers_v2::utils::now_secs(),
+                timestamp: crate::handlers::utils::now_secs(),
             })
         });
         Ok("Handoff memo left".to_string())
@@ -219,7 +219,7 @@ impl McpTool for AddSessionSummaryHandler {
             summaries.push(crate::models::SessionSummary {
                 summary: req.summary,
                 namespace: req.namespace,
-                timestamp: crate::handlers_v2::utils::now_secs(),
+                timestamp: crate::handlers::utils::now_secs(),
             })
         });
         Ok("Session summary added".to_string())
@@ -245,7 +245,7 @@ impl McpTool for GenerateStandupReportHandler {
         let req: GenerateStandupReportTool =
             serde_json::from_value(args).map_err(|e| e.to_string())?;
         let cutoff =
-            crate::handlers_v2::utils::now_secs().saturating_sub(req.hours_lookback * 3600);
+            crate::handlers::utils::now_secs().saturating_sub(req.hours_lookback * 3600);
 
         let report_str = state.tasks.read_with(|items| {
             state.ledger.read_with(|changes| {

server/src/handlers/tasks.rs

@@ -20,7 +20,7 @@ impl McpTool for AddTaskHandler {
 
     async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
         let req: AddTaskTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
-        let now = crate::handlers_v2::utils::now_secs();
+        let now = crate::handlers::utils::now_secs();
         let task_id = uuid::Uuid::new_v4().to_string();
 
         let deps = req.dependencies.unwrap_or_default();
@@ -214,7 +214,7 @@ impl McpTool for UpdateTaskStatusHandler {
             if !blocked {
                 // Apply update
                 tasks[target_idx].status = target_status.clone();
-                tasks[target_idx].updated_at = crate::handlers_v2::utils::now_secs();
+                tasks[target_idx].updated_at = crate::handlers::utils::now_secs();
 
                 // Cascade cancellation to children
                 if target_status == "cancelled" || target_status == "abandoned" {
@@ -337,7 +337,7 @@ impl McpTool for SetAcceptanceCriteriaHandler {
                         is_met: false,
                     })
                     .collect();
-                task.updated_at = crate::handlers_v2::utils::now_secs();
+                task.updated_at = crate::handlers::utils::now_secs();
                 success = true;
             }
         });
@@ -381,7 +381,7 @@ impl McpTool for VerifyAcceptanceCriteriaHandler {
                 } else {
                     ac.is_met = true;
                     success = true;
-                    task.updated_at = crate::handlers_v2::utils::now_secs();
+                    task.updated_at = crate::handlers::utils::now_secs();
                 }
             }
         });

server/src/handlers/workspaces.rs

@@ -25,7 +25,7 @@ impl McpTool for PinFileHandler {
             pinned.push(crate::models::PinnedFile {
                 namespace: req.namespace,
                 file_path: req.file_path,
-                timestamp: crate::handlers_v2::utils::now_secs(),
+                timestamp: crate::handlers::utils::now_secs(),
                 git_branch: req.git_branch,
             });
         });
@@ -115,7 +115,7 @@ impl McpTool for StoreSnippetHandler {
             language: req.language,
             code: req.code,
             description: req.description,
-            updated_at: crate::handlers_v2::utils::now_secs(),
+            updated_at: crate::handlers::utils::now_secs(),
         };
 
         let idx = state
@@ -223,7 +223,7 @@ impl McpTool for SaveContextWorkspaceHandler {
                 name: req.name,
                 pinned_files: req.pinned_files,
                 active_task_ids: req.active_task_ids,
-                saved_at: crate::handlers_v2::utils::now_secs(),
+                saved_at: crate::handlers::utils::now_secs(),
             });
         });
         Ok("Context workspace saved".to_string())
@@ -363,4 +363,4 @@ impl McpTool for ClearPrChecklistHandler {
     }
 }
 
-use crate::handlers_v2::utils::*;
+use crate::handlers::utils::*;

server/src/main.rs

@@ -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(())
 }

server/src/router.rs

@@ -14,3 +14,197 @@ pub trait McpTool: Send + Sync {
     /// Execute the tool with the given arguments
     async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String>;
 }
+
+pub struct MemoryHandler {
+    pub state: Arc<MemoryState>,
+    pub tools: std::collections::HashMap<String, Box<dyn McpTool>>,
+}
+
+impl MemoryHandler {
+    pub fn new(state: Arc<MemoryState>) -> Self {
+        let mut tools: std::collections::HashMap<String, Box<dyn McpTool>> =
+            std::collections::HashMap::new();
+
+        macro_rules! register {
+            ($module:ident::$handler:ident) => {
+                let h = crate::handlers::$module::$handler;
+                tools.insert(h.name().to_string(), Box::new(h));
+            };
+        }
+
+        register!(graph::QueryGraphPathHandler);
+        register!(graph::CreateEntitiesHandler);
+        register!(graph::CreateRelationsHandler);
+        register!(graph::AddObservationsHandler);
+        register!(graph::DeleteEntitiesHandler);
+        register!(graph::DeleteObservationsHandler);
+        register!(graph::DeleteRelationsHandler);
+        register!(graph::ReadGraphHandler);
+        register!(graph::SearchNodesHandler);
+        register!(graph::OpenNodesHandler);
+        register!(graph::VisualizeGraphHandler);
+        register!(graph::CondenseEntityHandler);
+        register!(graph::MergeEntitiesHandler);
+        register!(graph::FindOrphansHandler);
+
+        register!(tasks::AddTaskHandler);
+        register!(tasks::DeleteTaskHandler);
+        register!(tasks::UpdateTaskStatusHandler);
+        register!(tasks::ListActiveTasksHandler);
+        register!(tasks::SetAcceptanceCriteriaHandler);
+        register!(tasks::VerifyAcceptanceCriteriaHandler);
+        register!(tasks::AddMilestoneHandler);
+        register!(tasks::UpdateMilestoneHandler);
+        register!(tasks::ListMilestonesHandler);
+
+        register!(notes::AddStickyNoteHandler);
+        register!(notes::ReadStickyNotesHandler);
+        register!(notes::DeleteStickyNoteHandler);
+        register!(notes::ClearStickyNotesHandler);
+        register!(notes::LeaveHandoffMemoHandler);
+        register!(notes::ReadHandoffMemosHandler);
+        register!(notes::ClearHandoffMemosHandler);
+        register!(notes::AddSessionSummaryHandler);
+        register!(notes::GenerateStandupReportHandler);
+
+        register!(meta::LogDecisionHandler);
+        register!(meta::QueryDecisionsHandler);
+        register!(meta::LogErrorFixHandler);
+        register!(meta::SearchErrorFixesHandler);
+        register!(meta::LogCodeChangeHandler);
+        register!(meta::QueryRecentChangesHandler);
+        register!(meta::LearnPreferenceHandler);
+        register!(meta::ReadPreferencesHandler);
+        register!(meta::LogTechDebtHandler);
+        register!(meta::ResolveTechDebtHandler);
+        register!(meta::ListTechDebtHandler);
+        register!(meta::OmniSearchHandler);
+        register!(meta::GetProjectHealthHandler);
+
+        register!(env::UpdateEnvFingerprintHandler);
+        register!(env::ReadEnvFingerprintHandler);
+        register!(env::LogEnvRequirementHandler);
+        register!(env::RegisterEnvironmentHandler);
+        register!(env::GetEnvironmentDetailsHandler);
+
+        register!(workspaces::PinFileHandler);
+        register!(workspaces::UnpinFileHandler);
+        register!(workspaces::ListPinnedFilesHandler);
+        register!(workspaces::StoreSnippetHandler);
+        register!(workspaces::SearchSnippetsHandler);
+        register!(workspaces::DeleteSnippetHandler);
+        register!(workspaces::SaveContextWorkspaceHandler);
+        register!(workspaces::LoadContextWorkspaceHandler);
+        register!(workspaces::ListContextWorkspacesHandler);
+        register!(workspaces::AddPrChecklistItemHandler);
+        register!(workspaces::GetPrChecklistHandler);
+        register!(workspaces::ClearPrChecklistHandler);
+
+        Self { state, tools }
+    }
+
+    pub async fn handle_request(&self, req: serde_json::Value) -> Option<serde_json::Value> {
+        let id = req.get("id").cloned().unwrap_or(serde_json::Value::Null);
+        let id_clone = id.clone();
+        let method = req.get("method").and_then(|m| m.as_str()).unwrap_or("");
+
+        match method {
+            "server/discover" => {
+                let payload = serde_json::json!({
+                    "resultType": "complete",
+                    "ttlMs": 0,
+                    "cacheScope": "public",
+                    "supportedVersions": ["2026-07-28", "2025-11-25", "2025-06-18", "2025-03-26", "2024-11-05"],
+                    "capabilities": {
+                        "tools": serde_json::json!({})
+                    },
+                    "_meta": {
+                        "io.modelcontextprotocol/serverInfo": {
+                            "name": "gemini-mcp-memory",
+                            "version": "3.0.0"
+                        }
+                    }
+                });
+                Some(crate::mcp::success(id, payload))
+            }
+            "initialize" => {
+                let init = rmcp::model::InitializeResult::new(
+                    rmcp::model::ServerCapabilities::builder()
+                        .enable_tools()
+                        .build(),
+                )
+                .with_server_info(rmcp::model::Implementation::new(
+                    "gemini-mcp-memory",
+                    "3.0.0",
+                ));
+                Some(crate::mcp::success(
+                    id,
+                    serde_json::to_value(&init).unwrap_or_default(),
+                ))
+            }
+            "notifications/initialized" => None,
+            "tools/list" => {
+                let mut tools: Vec<serde_json::Value> =
+                    self.tools.values().map(|t| t.schema()).collect();
+                tools.sort_by_key(|t| {
+                    t.get("name")
+                        .and_then(|n| n.as_str())
+                        .unwrap_or("")
+                        .to_string()
+                });
+                Some(crate::mcp::success(
+                    id,
+                    serde_json::json!({ "tools": tools }),
+                ))
+            }
+            "tools/call" => {
+                let params = req.get("params").unwrap_or(&serde_json::Value::Null);
+                let name = params.get("name").and_then(|n| n.as_str()).unwrap_or("");
+                let args = params
+                    .get("arguments")
+                    .cloned()
+                    .unwrap_or(serde_json::Value::Object(Default::default()));
+
+                self.state
+                    .broadcast_activity(&format!("Agent executed tool: {}", name));
+
+                let result: Result<String, String> = if let Some(tool) = self.tools.get(name) {
+                    tool.execute(args, self.state.clone()).await
+                } else {
+                    Err(format!("Unknown tool: {}", name))
+                };
+
+                match result {
+                    Ok(text) => {
+                        let payload = serde_json::json!({
+                            "content": [{"type": "text", "text": text}],
+                            "isError": false
+                        });
+                        Some(crate::mcp::success(id_clone, payload))
+                    }
+                    Err(e) => {
+                        tracing::error!("Tool {} failed: {}", name, e);
+                        let payload = serde_json::json!({
+                            "content": [{"type": "text", "text": e}],
+                            "isError": true
+                        });
+                        Some(crate::mcp::success(id_clone, payload))
+                    }
+                }
+            }
+            m if m.starts_with("notifications/") => None,
+            "ping" => Some(crate::mcp::success(id, serde_json::json!({}))),
+            _ => {
+                if id.is_null() {
+                    None
+                } else {
+                    Some(crate::mcp::error(
+                        id,
+                        -32601,
+                        &format!("Method {} not found", method),
+                    ))
+                }
+            }
+        }
+    }
+}

server/src/state.rs

@@ -32,6 +32,49 @@ pub struct MemoryState {
 }
 
 impl MemoryState {
+    pub fn new(base_dir_str: &str) -> Self {
+        let base = std::path::PathBuf::from(base_dir_str);
+        std::fs::create_dir_all(&base).expect("Failed to create store dir");
+
+        let db = crate::db::init_redb(&base);
+
+        Self {
+            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,
+        }
+    }
+
     pub fn deduplicate<T: Eq + std::hash::Hash + Clone>(input: &mut Vec<T>) {
         let mut keys = std::collections::HashSet::new();
         input.retain(|entry| keys.insert(entry.clone()));