mod messages;
mod rpc;
mod tools;
use std::path::PathBuf;
use std::process::ExitCode;
use std::sync::Arc;
use clap::Parser;
const ENV_DB: &str = "PLUGMEM_DB";
const DEFAULT_DB: &str = "plugmem.db";
#[derive(Parser)]
#[command(name = "plugmem-mcp", version, about = messages::ABOUT_CLI)]
struct Args {
#[arg(long)]
db: Option<PathBuf>,
#[arg(long)]
config: Option<PathBuf>,
#[arg(long)]
read_only: bool,
#[arg(long)]
workers: Option<usize>,
}
fn main() -> ExitCode {
let args = Args::parse();
let path = args
.db
.or_else(|| std::env::var_os(ENV_DB).map(PathBuf::from))
.unwrap_or_else(|| PathBuf::from(DEFAULT_DB));
let table = match plugmem_host::read_config(args.config.as_deref()) {
Ok(t) => t,
Err(e) => {
eprintln!("plugmem-mcp: {e}");
return ExitCode::from(2);
}
};
let workers = args
.workers
.unwrap_or_else(|| resolve_workers(table.as_ref()));
let settings = match plugmem_host::Settings::from_table(table.as_ref()) {
Ok(s) => s,
Err(e) => {
eprintln!("plugmem-mcp: {e}");
return ExitCode::from(2);
}
};
let shared = if args.read_only {
let embedder = settings.embedder;
match plugmem_host::Database::open_readonly(&path, settings.config) {
Ok(db) => rpc::Shared::Reader(Arc::new(tools::ReaderShared::new(db, embedder))),
Err(e) => {
eprintln!("plugmem-mcp: {}: {e}", path.display());
return ExitCode::from(2);
}
}
} else {
match settings.open(&path) {
Ok(db) => rpc::Shared::Writer(db),
Err(e) => {
eprintln!("plugmem-mcp: {}: {e}", path.display());
return ExitCode::from(2);
}
}
};
rpc::serve(shared, workers);
ExitCode::SUCCESS
}
fn resolve_workers(table: Option<&toml::Table>) -> usize {
table
.and_then(|t| t.get("server"))
.and_then(toml::Value::as_table)
.and_then(|s| s.get("workers"))
.and_then(toml::Value::as_integer)
.filter(|n| *n > 0)
.map(|n| n as usize)
.unwrap_or_else(default_workers)
}
fn default_workers() -> usize {
std::thread::available_parallelism()
.map(|n| (n.get() / 2).max(1))
.unwrap_or(1)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn workers_from_config_else_default() {
let table: toml::Table = "[server]\nworkers = 3\n".parse().unwrap();
assert_eq!(resolve_workers(Some(&table)), 3);
let empty: toml::Table = "[engine]\ndim = 8\n".parse().unwrap();
assert!(resolve_workers(Some(&empty)) >= 1);
assert!(resolve_workers(None) >= 1);
let zero: toml::Table = "[server]\nworkers = 0\n".parse().unwrap();
assert!(resolve_workers(Some(&zero)) >= 1);
assert!(default_workers() >= 1);
}
}