mcp-pocket-memory 0.5.0

mod db;
mod tool;

use std::{env, path::PathBuf};

use db::Database;
use rmcp::ServiceExt;
use tool::PocketMemoryServer;

const DB_PATH_ENV: &str = "POCKET_MEMORY_DB";
const DEFAULT_DATA_DIR: &str = ".mcp-pocket-memory";
const DEFAULT_DB_FILE: &str = "pocket_memory.sqlite3";

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut args = env::args().skip(1);
    let command = args.next();
    let db_path = db_path(args.next());

    match command.as_deref() {
        Some("init") => {
            Database::open(&db_path)?;
            println!("initialized database at {}", db_path.display());
        }
        Some("serve") | None => {
            Database::open(&db_path)?;
            let server = PocketMemoryServer::new(&db_path)
                .serve(rmcp::transport::stdio())
                .await?;

            server.waiting().await?;
        }
        Some(command) => {
            eprintln!("unknown command: {command}");
            eprintln!("usage: mcp-pocket-memory [init|serve] [db-path]");
            std::process::exit(2);
        }
    }

    Ok(())
}

fn db_path(arg_path: Option<String>) -> PathBuf {
    arg_path
        .or_else(|| env::var(DB_PATH_ENV).ok())
        .map(PathBuf::from)
        .unwrap_or_else(default_db_path)
}

fn default_db_path() -> PathBuf {
    env::var("HOME")
        .map(PathBuf::from)
        .unwrap_or_else(|_| PathBuf::from("."))
        .join(DEFAULT_DATA_DIR)
        .join(DEFAULT_DB_FILE)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn db_path_prefers_arg_over_env_and_default() {
        assert_eq!(
            db_path(Some("/tmp/custom.sqlite3".into())),
            PathBuf::from("/tmp/custom.sqlite3")
        );
    }
}