mcp-memory/server/src/main.rs

mod handlers;
mod mcp;
mod models;
mod search;
mod state;
mod store;
mod tools;

use crate::handlers::MemoryHandler;
use crate::models::*;
use crate::state::MemoryState;
use crate::store::Store;

use std::fs;
use std::path::PathBuf;
use std::sync::{Arc, RwLock};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use redb::ReadableTable;
use tokio::time::sleep;

use clap::{Parser, Subcommand};
use std::collections::HashMap;

#[derive(Parser)]
#[command(author, version, about = "Antigravity MCP Memory Server", long_about = None)]
struct Cli {
    #[command(subcommand)]
    command: Option<Commands>,
    /// Target URL for the proxy to connect to (e.g., http://127.0.0.1:3000)
    #[arg(long)]
    target: Option<String>,
    /// Run the server as a background daemon process (Windows only)
    #[arg(long)]
    daemon: bool,
    /// Send a shutdown request to the currently running server
    #[arg(long)]
    exit: bool,
    /// Send a shutdown request to the existing server and wait for it to exit
    #[arg(long)]
    restart: bool,
}

#[derive(Subcommand)]
enum Commands {
    /// Manage authorization gates and verification for actions
    Gate {
        #[command(subcommand)]
        subcmd: GateCommands,
    },
}

#[derive(Subcommand)]
enum GateCommands {
    Set {
        #[arg(long)]
        action: String,
        #[arg(long)]
        target: String,
        #[arg(long)]
        namespace: Option<String>,
        #[arg(short = 'p', long = "param")]
        params: Vec<String>,
        #[arg(long, conflicts_with = "block")]
        authorize: bool,
        #[arg(long, conflicts_with = "authorize")]
        block: bool,
        #[arg(long)]
        reason: Option<String>,
    },
    Verify {
        #[arg(long)]
        action: String,
        #[arg(long)]
        target: String,
        #[arg(long)]
        namespace: Option<String>,
        #[arg(short = 'p', long = "param")]
        params: Vec<String>,
        #[arg(long)]
        consume: bool,
    },
}

async fn garbage_collector_worker(state: Arc<MemoryState>) {
    loop {
        // Run every 6 hours
        tokio::time::sleep(tokio::time::Duration::from_secs(6 * 3600)).await;
        
        let now = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs();
        
        // 1. Task GC (14 days)
        let fourteen_days = 14 * 24 * 3600;
        let task_cutoff = now.saturating_sub(fourteen_days);
        state.tasks.modify(|tasks| {
            let initial_len = tasks.len();
            tasks.retain(|task| !(task.status.to_lowercase() == "completed" && task.created_at < task_cutoff));
            if tasks.len() < initial_len {
                eprintln!("GC: Removed {} old completed tasks", initial_len - tasks.len());
            }
        });

        // 2. Ledger GC (7 days or max 1000 items)
        state.ledger.modify(|ledger| {
            let seven_days = now.saturating_sub(7 * 24 * 3600);
            ledger.retain(|c| c.timestamp >= seven_days);
            if ledger.len() > 1000 {
                let excess = ledger.len() - 1000;
                ledger.drain(0..excess);
            }
        });

        // 3. Sticky Notes GC (24 hours)
        state.sticky.modify(|notes| {
            notes.retain(|note| note.timestamp >= now.saturating_sub(24 * 3600));
        });
    }
}

async fn git_sync_worker(state: Arc<MemoryState>) {
    let repo_path = std::env::current_dir().unwrap_or_else(|_| ".".into());
    let mut last_commit_id = String::new();
    
    loop {
        tokio::time::sleep(tokio::time::Duration::from_secs(30)).await;
        
        let repo_path_clone = repo_path.clone();
        let commit_data = tokio::task::spawn_blocking(move || {
            if let Ok(repo) = git2::Repository::discover(&repo_path_clone) {
                if let Ok(head) = repo.head() {
                    if let Ok(commit) = head.peel_to_commit() {
                        let current_id = commit.id().to_string();
                        let msg = commit.message().unwrap_or("").to_string();
                        let branch = head.shorthand().unwrap_or("unknown").to_string();
                        return Some((current_id, msg, branch));
                    }
                }
            }
            None
        })
        .await
        .unwrap_or(None);

        if let Some((current_id, msg, branch)) = commit_data {
            if current_id != last_commit_id && !last_commit_id.is_empty() {
                state.ledger.modify(|changes| {
                    changes.push(crate::models::CodeChange {
                        git_commit: Some(current_id.clone()),
                        git_branch: Some(branch),
                        description: format!("Auto-synced commit: {}", msg.trim()),
                        timestamp: std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs(),
                        file_path: "".to_string(),
                    });
                });
                tracing::info!("Git Sync: Logged new commit {}", current_id);
                
                state.tasks.modify(|tasks| {
                    for task in tasks.iter_mut() {
                        if task.status != "completed" && msg.to_lowercase().contains(&task.title.to_lowercase()) {
                            task.status = "completed".to_string();
                            tracing::info!("Git Sync: Auto-completed task '{}'", task.title);
                        }
                    }
                });
            }
            last_commit_id = current_id;
        }
    }
}

async fn reconcile_worker(state: Arc<MemoryState>) {
    loop {
        sleep(Duration::from_secs(5)).await;
        
        let has_local = {
            let session = state.session_graph.read().unwrap();
            !session.entities.is_empty() || !session.relations.is_empty()
        };

        let base_dir = state.base_dir.clone();
        let has_files = tokio::task::spawn_blocking(move || {
            let pattern = format!("{}/delta_*.json", base_dir.display());
            glob::glob(&pattern).map(|p| p.count() > 0).unwrap_or(false)
        })
        .await
        .unwrap_or(false);

        if has_local || has_files {
            state.apply_sync_write(|_master| {}).await;
            let state_clone = state.clone();
            let _ = tokio::task::spawn_blocking(move || {
                state_clone.rebuild_index();
            }).await;
        }
    }
}

use axum::{
    Json, Router,
    extract::{Query, State, ws::{WebSocketUpgrade, WebSocket, Message}},
    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: SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_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()
}

fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::Error>> {
    let rt = tokio::runtime::Runtime::new().unwrap();
    rt.block_on(async {
        tokio::spawn(reconcile_worker(Arc::clone(&state)));
        let app_state = Arc::new(AppState {
            handler: Arc::new(MemoryHandler { state: Arc::clone(&state) }),
            clients: RwLock::new(HashMap::new()),
            next_id: AtomicUsize::new(1),
        });

        let app = Router::new()
            .route("/api/version", get(|| async move {
                axum::Json(serde_json::json!({
                    "version": env!("BUILD_DATE"),
                    "git_hash": option_env!("GIT_HASH").unwrap_or("unknown")
                }))
            }))
            .route("/ws", get(ws_handler))
            .route("/sse", get(sse_handler))
            .route("/messages", post(message_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(|| async move {
                    std::thread::spawn(|| {
                        std::thread::sleep(std::time::Duration::from_millis(100));
                        std::process::exit(0);
                    });
                    "Shutting down..."
                }),
            )
            .route(
                "/",
                get(|| async move { axum::response::Html(include_str!("dashboard.html")) }),
            )
                        .route("/api/graph", get({
                let state_clone = app_state.handler.state.clone();
                move || async move {
                    let graph = state_clone.get_full_graph();
                    axum::Json(graph)
                }
            }))
            
            .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 = state_clone.tasks.read();
                    axum::Json(tasks.clone())
                }
            }))
            .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") {
                        if let Ok(idx) = state_clone.search_index.read() {
                            if let Ok(results) = idx.search(q, None) {
                                let mut formatted_results = Vec::new();
                                for (type_name, content) in results {
                                    formatted_results.push(serde_json::json!({
                                        "type_name": type_name,
                                        "content": content,
                                        "score": 1.0
                                    }));
                                }
                                return axum::Json(serde_json::json!({ "results": formatted_results }));
                            }
                        }
                    }
                    axum::Json(serde_json::json!({ "results": [] }))
                }
            }))
            .route("/api/stats",
                get({
                    let state_clone = app_state.handler.state.clone();
                    move || async move {
                        let (entities, relations) = {
                            let graph = state_clone.get_full_graph();
                            (graph.entities.len(), graph.relations.len())
                        };
                        let tasks = state_clone.tasks.read().len();
                        let snippets = state_clone.snippets.read().len();
                        let tech_debts = state_clone.tech_debts.read().len();
                        let adrs = state_clone.adrs.read().len();

                        let ledger = state_clone.ledger.read().len();
                        let sticky = state_clone.sticky.read().len();
                        let error_fixes = state_clone.error_fixes.read().len();
                        let pinned_files = state_clone.pinned_files.read().len();
                        let session_summaries = state_clone.session_summaries.read().len();
                        let handoff_memos = state_clone.handoff_memos.read().len();
                        let env_fingerprints = state_clone.env_fingerprints.read().len();
                        let env_requirements = state_clone.env_requirements.read().len();
                        let milestones = state_clone.milestones.read().len();
                        let environments = state_clone.environments.read().len();
                        let pr_checklists = state_clone.pr_checklists.read().len();
                        let gates = state_clone.gates.read().len();
                        let context_workspaces = state_clone.context_workspaces.read().len();

                        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 listener = match tokio::net::TcpListener::bind(std::env::var("MCP_PORT").map(|p| format!("127.0.0.1:{}", p)).unwrap_or_else(|_| std::env::var("MCP_PORT").map(|p| format!("127.0.0.1:{}", p)).unwrap_or_else(|_| "127.0.0.1:3000".to_string()).to_string())).await {
            Ok(l) => l,
            Err(e) => {
                tracing::info!("Port 3000 is already in use ({}). Assuming server is already running and exiting gracefully.", e);
                std::process::exit(0);
            }
        };
        
        tokio::spawn(garbage_collector_worker(Arc::clone(&state)));

        tokio::spawn(git_sync_worker(Arc::clone(&state)));
    tracing::info!("MCP Memory Server running on http://127.0.0.1:3000/sse");
    if let Err(e) = axum::serve(listener, app).await {
            let log_path = dirs::home_dir()
                .unwrap_or_default()
                .join(".gemini/mcp_memory/daemon_error.log");
            let _ = std::fs::write(&log_path, format!("Server crashed: {}\n", e));
        }
        Ok(())
    })
}

#[derive(serde::Deserialize)]
struct MsgQuery {
    session_id: String,
}

async fn message_handler(
    State(state): State<Arc<AppState>>,
    Query(q): Query<MsgQuery>,
    Json(payload): Json<serde_json::Value>,
) -> impl axum::response::IntoResponse {
    let session_id = q.session_id;
    if let Some(response) = state.handler.handle_request(payload).await {
        let res_str = serde_json::to_string(&response).unwrap();
        let tx_opt = state.clients.read().unwrap().get(&session_id).cloned();
        if let Some(tx) = tx_opt {
            let _ = tx.send(res_str).await;
        }
    }
    (axum::http::StatusCode::ACCEPTED, "Accepted").into_response()
}

async fn sse_handler(
    State(state): State<Arc<AppState>>,
) -> axum::response::sse::Sse<impl tokio_stream::Stream<Item = Result<axum::response::sse::Event, std::convert::Infallible>>> {
    let session_id = format!("{}", state.next_id.fetch_add(1, Ordering::SeqCst));
    let (tx, rx) = mpsc::channel::<String>(100);
    state.clients.write().unwrap().insert(session_id.clone(), tx.clone());

    let endpoint = format!("/messages?session_id={}", session_id);
    let _ = tx.send(format!("endpoint|{}", endpoint)).await;

    let rx_stream = tokio_stream::wrappers::ReceiverStream::new(rx);
    let event_stream = rx_stream.map(|msg| {
        if let Some(ep) = msg.strip_prefix("endpoint|") {
            Ok(axum::response::sse::Event::default().event("endpoint").data(ep))
        } else {
            Ok(axum::response::sse::Event::default().event("message").data(msg))
        }
    });

    axum::response::sse::Sse::new(event_stream).keep_alive(axum::response::sse::KeepAlive::new())
}

async fn ws_handler(
    ws: WebSocketUpgrade,
    State(state): State<Arc<AppState>>,
    Query(query): Query<std::collections::HashMap<String, String>>,
) -> impl axum::response::IntoResponse {
    let client_type = query.get("client").cloned().unwrap_or_else(|| "unknown".to_string());
    ws.on_upgrade(move |socket| handle_socket(socket, state, client_type))
}

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().insert(session_id.clone(), tx.clone());

    let (mut sender, mut receiver) = socket.split();

    let mut send_task = tokio::spawn(async move {
        while let Some(msg) = rx.recv().await {
            if sender.send(Message::Text(msg.into())).await.is_err() {
                break;
            }
        }
    });

    let handler = Arc::clone(&state.handler);
    let state_clone = Arc::clone(&state);
    let session_id_clone = session_id.clone();

    let mut recv_task = tokio::spawn(async move {
        while let Some(Ok(Message::Text(text))) = receiver.next().await {
            if let Ok(payload) = serde_json::from_str::<serde_json::Value>(&text) {
                if client_type == "proxy" {
                    // Send activity broadcast to UI clients
                    if let Some(method) = payload.get("method").and_then(|m| m.as_str()) {
                        if method == "tools/call" {
                            let name = payload.get("params").and_then(|p| p.get("name")).and_then(|n| n.as_str()).unwrap_or("unknown_tool");
                            let activity_msg = format!("Agent executed tool: {}", name);
                            
                            let event = serde_json::json!({
                                "type": "activity",
                                "data": activity_msg
                            });
                            
                            let clients_map = state_clone.clients.read().unwrap().clone();
                            for (id, client_tx) in clients_map.iter() {
                                if id != &session_id_clone {
                                    let _ = client_tx.send(event.to_string()).await;
                                }
                            }
                        }
                    }

                    // Process MCP request
                    if let Some(response) = handler.handle_request(payload).await {
                        let res_str = serde_json::to_string(&response).unwrap();
                        let tx_opt = state_clone.clients.read().unwrap().get(&session_id_clone).cloned();
                        if let Some(client_tx) = tx_opt {
                            let _ = client_tx.send(res_str).await;
                        }
                    }
                }
            }
        }
    });

    tokio::select! {
        _ = (&mut send_task) => recv_task.abort(),
        _ = (&mut recv_task) => send_task.abort(),
    };

    state.clients.write().unwrap().remove(&session_id);
}


#[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 _ = std::fs::write(&win_path, &payload.session_id);
        
        let wsl_path = "\\\\wsl.localhost\\Ubuntu\\home\\riz\\.gemini\\active_nvim.txt";
        let _ = std::fs::write(wsl_path, &payload.session_id);
    }

    // 2. Broadcast to UI WebSockets
    let ws_msg = serde_json::json!({
        "type": "nvim_telemetry",
        "data": payload
    });
    
    let msg_str = ws_msg.to_string();
    let clients = state.clients.read().unwrap().clone();
    for (_, tx) in clients.iter() {
        let _ = tx.send(msg_str.clone()).await;
    }

    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()
            .map(|mut h| {
                h.push(".gemini/mcp_memory");
                h.to_string_lossy().to_string()
            })
            .unwrap_or_else(|| ".gemini/mcp_memory".into())
    });
    let log_dir = std::path::PathBuf::from(base_dir).join("logs");
    std::fs::create_dir_all(&log_dir).unwrap_or_default();
    
    let file_appender = tracing_appender::rolling::daily(log_dir, format!("{}.log", app_name));
    let (non_blocking, guard) = tracing_appender::non_blocking(file_appender);
    
    let _ = tracing_subscriber::fmt()
        .with_writer(non_blocking)
        .with_ansi(false)
        .with_max_level(tracing::Level::INFO)
        .try_init();
    
    Some(guard)
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let _guard = init_logging("server");
    let cli = Cli::parse();

    if cli.exit {
        if let Ok(mut stream) = std::net::TcpStream::connect(std::env::var("MCP_PORT").map(|p| format!("127.0.0.1:{}", p)).unwrap_or_else(|_| "127.0.0.1:3000".to_string())) {
            use std::io::Write;
            let _ = stream.write_all(
                b"POST /shutdown HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n",
            );
        }
        println!("Sent shutdown request to server.");
        return Ok(());
    }

    if cli.restart {
        if let Ok(mut stream) = std::net::TcpStream::connect(std::env::var("MCP_PORT").map(|p| format!("127.0.0.1:{}", p)).unwrap_or_else(|_| "127.0.0.1:3000".to_string())) {
            use std::io::Write;
            let _ = stream.write_all(
                b"POST /shutdown HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n",
            );
            println!("Sent shutdown request to existing server. Waiting for it to exit...");
            std::thread::sleep(std::time::Duration::from_millis(1500));
        }
        return Ok(());
    }

    #[cfg(target_os = "windows")]
    {
        use std::os::windows::process::CommandExt;
        if !cli.daemon {
            // 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().unwrap())
                .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");
            return Ok(());
        }
    }

    #[cfg(not(target_os = "windows"))]
    {
        // Linux no longer executes server logic natively due to workspace split
        return Ok(());
    }

    let base_dir = std::env::var("MCP_MEMORY_STORE_DIR").unwrap_or_else(|_| {
        dirs::home_dir()
            .map(|mut h| {
                h.push(".gemini/mcp_memory");
                h.to_string_lossy().into_owned()
            })
            .unwrap_or_else(|| ".gemini/mcp_memory".into())
    });
    let base = PathBuf::from(base_dir);
    fs::create_dir_all(&base).expect("Failed to create store dir");

    let redb_path = base.join("mcp_store.redb");
    let db = Arc::new(redb::Database::create(&redb_path).unwrap());
    
    // Ensure table exists and migrate old JSON files
    {
        let write_txn = db.begin_write().unwrap();
        {
            let mut table = write_txn.open_table(crate::store::STORE_TABLE).unwrap();
            
            let stores = [
                ("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).unwrap().is_none() {
                    let json_path = base.join(file_name);
                    if json_path.exists() {
                        if let Ok(data) = fs::read(&json_path) {
                            if serde_json::from_slice::<serde_json::Value>(&data).is_ok() {
                                table.insert(*key, data.as_slice()).unwrap();
                            }
                        }
                    }
                }
            }
        }
        write_txn.commit().unwrap();
    }

    let state = Arc::new(MemoryState {
        master_path: base.join("knowledge_graph_master.json"),
        session_graph: RwLock::new(KnowledgeGraph::default()),
        base_dir: base.clone(),
        master_cache: RwLock::new((KnowledgeGraph::default(), SystemTime::UNIX_EPOCH)),
        search_index: RwLock::new(crate::search::MemoryIndex::new(&base).unwrap()),
        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()),
    });

    state.recover_wal();
    state.rebuild_index();

    run_server(state)
}