use anyhow::{Context, Result};
use clap::{CommandFactory, Parser};
use clap_complete::generate;
use fs4::FileExt;
use std::io;
use trusty_memory::cli;
use trusty_memory::cli::output::OutputConfig;
use trusty_memory::cli::palace_resolver::{detect_serve_palace, resolve_palace};
use trusty_memory::cli::{Cli, Commands};
#[tokio::main]
async fn main() -> Result<()> {
dotenvy::from_filename(".env.local").ok();
let cli = Cli::parse();
trusty_common::init_tracing(cli.verbose);
trusty_common::maybe_disable_color(cli.no_color);
let palace = resolve_palace(cli.palace.as_deref());
let out = OutputConfig {
json: cli.json,
quiet: cli.quiet,
no_color: cli.no_color,
};
if !matches!(
&cli.command,
Commands::Serve { .. } | Commands::Service(_) | Commands::Setup { .. } | Commands::Hooks(_)
) {
ensure_daemon().await;
}
match cli.command {
Commands::Remember {
text,
room,
tags,
importance,
} => {
cli::handle_remember(&palace, text, room, tags, importance, &out).await?;
}
Commands::Recall {
query,
top_k,
room,
deep,
decay: _,
all_palaces,
} => {
if all_palaces {
cli::handle_recall_all(query, top_k, deep, &out).await?;
} else {
cli::handle_recall(&palace, query, top_k, room, deep, &out).await?;
}
}
Commands::Forget { id } => {
cli::handle_forget(&palace, &id, &out).await?;
}
Commands::List { limit, room, sort } => {
cli::handle_list(&palace, limit, room, sort, &out).await?;
}
Commands::Palace(sub) => cli::palace::handle(sub, &palace, &out).await?,
Commands::Kg(sub) => cli::kg::handle(sub, &palace, &out).await?,
Commands::Git(sub) => cli::git::handle(sub, &palace, &out).await?,
Commands::Kuzu(sub) => cli::kuzu::handle(sub, &palace, &out).await?,
Commands::Analytics(sub) => cli::analytics::handle(sub, &palace, &out).await?,
Commands::Decay(sub) => cli::decay::handle(sub, &palace, &out).await?,
Commands::Dream(sub) => cli::dream::handle(sub, &palace, &out).await?,
Commands::Serve {
http,
mcp: _,
palace: serve_palace,
} => {
let auto_detected = serve_palace.is_none();
let default_palace = resolve_palace_for_serve(serve_palace.as_deref());
if auto_detected {
if let Some(name) = default_palace.as_deref() {
eprintln!("info: auto-detected palace '{name}' from working directory");
}
}
tracing::info!(?http, ?default_palace, "Starting trusty-memory MCP server");
let data_root_for_state = cli::palace::data_root()?;
let service_root = trusty_common::resolve_data_dir("trusty-memory")
.context("resolve trusty-memory service root for lock file")?;
let lock_path = service_root.join("trusty-memory.lock");
let lock_file = std::fs::OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(&lock_path)
.with_context(|| format!("open lock file at {}", lock_path.display()))?;
match FileExt::try_lock(&lock_file) {
Ok(()) => {
use std::io::Write as _;
let mut f = &lock_file;
let _ = f.set_len(0);
let _ = writeln!(f, "{}", std::process::id());
}
Err(_) => {
eprintln!(
"Another trusty-memory instance is already running. \
Use 'trusty-memory stop' to stop it."
);
std::process::exit(1);
}
}
let _lock_file = lock_file;
if let Err(e) = cli::stop::write_pid_file(std::process::id()) {
tracing::warn!("could not write daemon pid file: {e:#}");
}
if let Some(name) = default_palace.as_deref() {
let pid = trusty_memory_core::PalaceId::new(name);
let palace_dir = data_root_for_state.join(pid.as_str());
if !palace_dir.join("palace.json").exists() {
tracing::info!(palace = %name, "auto-creating default palace");
let registry = trusty_memory_core::PalaceRegistry::new();
let palace = trusty_memory_core::Palace {
id: pid,
name: name.to_string(),
description: Some(
"Auto-created by `trusty-memory serve --palace`".to_string(),
),
created_at: chrono::Utc::now(),
data_dir: palace_dir,
};
if let Err(e) = registry.create_palace(&data_root_for_state, palace) {
tracing::warn!(palace = %name, "failed to auto-create palace: {e:#}");
}
}
}
let state = trusty_memory_mcp::AppState::new(data_root_for_state)
.with_default_palace(default_palace);
let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false);
let dream_handles: std::sync::Arc<
tokio::sync::Mutex<Vec<tokio::task::JoinHandle<()>>>,
> = std::sync::Arc::new(tokio::sync::Mutex::new(Vec::new()));
let mut addr_written = false;
let serve_result = match http {
None => {
tokio::select! {
r = trusty_memory_mcp::run_stdio(state) => r,
_ = tokio::signal::ctrl_c() => {
tracing::info!("ctrl-c received, shutting down");
Ok(())
}
}
}
Some(http_addr) => {
let listener = trusty_common::bind_with_auto_port(http_addr, 20).await?;
let bound_addr = listener.local_addr()?;
println!(
"trusty-memory v{} — HTTP admin panel: http://{}",
env!("CARGO_PKG_VERSION"),
bound_addr
);
tracing::info!(%bound_addr, "HTTP server bound");
match trusty_common::write_daemon_addr("trusty-memory", &bound_addr.to_string())
{
Ok(()) => addr_written = true,
Err(e) => tracing::warn!("could not write daemon addr file: {e:#}"),
}
let dream_handles_bg = dream_handles.clone();
let shutdown_rx_bg = shutdown_rx.clone();
tokio::spawn(async move {
let root = match cli::palace::data_root() {
Ok(r) => r,
Err(e) => {
tracing::warn!("failed to resolve data root for Dreamer: {e:#}");
return;
}
};
let palaces = match trusty_memory_core::PalaceRegistry::list_palaces(&root)
{
Ok(ps) => ps,
Err(e) => {
tracing::warn!("failed to enumerate palaces for Dreamer: {e:#}");
return;
}
};
let max_eager: usize = std::env::var("TRUSTY_MAX_STARTUP_PALACES")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.unwrap_or(5);
let total = palaces.len();
let eager_count = palaces.len().min(max_eager);
tracing::info!(
count = total,
eager = eager_count,
max_eager,
"Dreamer background init starting (capped eager open)"
);
let registry = trusty_memory_core::PalaceRegistry::new();
let mut opened = 0usize;
for p in palaces.into_iter().take(max_eager) {
match registry.open_palace(&root, &p.id) {
Ok(handle) => {
let dreamer = std::sync::Arc::new(
trusty_memory_core::dream::Dreamer::new(
trusty_memory_core::dream::DreamConfig::default(),
),
);
let jh =
dreamer.start_with_shutdown(handle, shutdown_rx_bg.clone());
dream_handles_bg.lock().await.push(jh);
opened += 1;
}
Err(e) => tracing::warn!(
palace = %p.id,
"failed to open palace for dreamer: {e:#}"
),
}
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
tokio::task::yield_now().await;
}
tracing::info!(
opened,
total,
"Dreamer background init complete — consolidation active"
);
});
tokio::spawn(trusty_memory_mcp::run_stdio(state.clone()));
tokio::select! {
r = trusty_memory_mcp::run_http_on(state, listener) => r,
_ = tokio::signal::ctrl_c() => {
tracing::info!("ctrl-c received, shutting down");
Ok(())
}
}
}
};
let _ = shutdown_tx.send(true);
let handles = {
let mut guard = dream_handles.lock().await;
std::mem::take(&mut *guard)
};
for jh in handles {
let _ = tokio::time::timeout(std::time::Duration::from_secs(2), jh).await;
}
if addr_written {
if let Ok(dir) = trusty_common::resolve_data_dir("trusty-memory") {
let _ = std::fs::remove_file(dir.join("http_addr"));
}
}
let _ = cli::stop::remove_pid_file();
serve_result?;
}
Commands::Service(sub) => {
cli::service::handle(sub)?;
}
Commands::Setup {
non_interactive,
skip_migration,
migrate_only,
} => {
let opts = cli::setup::SetupOpts {
non_interactive,
skip_migration,
migrate_only,
};
cli::setup::handle_setup(opts, &out).await?;
}
Commands::Chat {
message,
top_k,
remember,
} => {
let opts = cli::chat::ChatOpts {
message,
remember,
top_k,
};
cli::chat::handle_chat(&palace, opts, &out).await?;
}
Commands::Config(sub) => match sub {
cli::ConfigCommands::Show => {
let cfg = cli::config::UserConfig::load()?;
let path = cli::config::default_config_path()?;
println!("config: {}", path.display());
println!("{}", toml::to_string_pretty(&cfg)?);
}
cli::ConfigCommands::Set { key, value } => {
let mut cfg = cli::config::UserConfig::load().unwrap_or_default();
cfg.set_dotted(&key, &value)?;
cfg.save()?;
let path = cli::config::default_config_path()?;
println!("set {key} in {}", path.display());
}
},
Commands::Convert(args) => {
cli::convert::handle_convert(args).await?;
}
Commands::Bench(sub) => match sub {
cli::BenchCommands::Compare(args) => {
let opts = cli::bench::BenchCompareOpts {
corpus: args.corpus,
top_k: args.top_k,
mempalace: args.mempalace,
kuzu: args.kuzu,
json: args.json,
};
cli::bench::handle_bench_compare(opts).await?;
}
},
Commands::Hooks(args) => {
cli::hooks::handle(args, &palace, &out).await?;
}
Commands::Backup(args) => {
cli::backup::handle_backup(args, &out).await?;
}
Commands::Restore(args) => {
cli::backup::handle_restore(args, &out).await?;
}
Commands::Status => {
let binary = std::env::current_exe()?;
let root = cli::palace::data_root()?;
let root_clone = root.clone();
let palaces = tokio::task::spawn_blocking(move || {
trusty_memory_core::PalaceRegistry::list_palaces(&root_clone)
})
.await??;
println!("trusty-memory v{}", env!("CARGO_PKG_VERSION"));
println!("binary: {}", binary.display());
println!("data_root: {}", root.display());
println!("palaces: {}", palaces.len());
println!("active palace: {palace}");
let pid_proc = cli::daemon_probe::probe_pid_file();
match cli::daemon_probe::probe_daemon() {
Some(found)
if !matches!(found.source, cli::daemon_probe::AddrSource::CandidatePort)
|| pid_proc.is_none() =>
{
let tag = match found.source {
cli::daemon_probe::AddrSource::EnvVar => {
format!(" [via ${}]", cli::daemon_probe::HTTP_PORT_ENV)
}
cli::daemon_probe::AddrSource::DiscoveryFile => String::new(),
cli::daemon_probe::AddrSource::CandidatePort => {
" [discovery file missing/stale — found via port scan]".to_string()
}
};
println!("HTTP: http://{}{tag}", found.addr);
}
_ => match pid_proc {
Some(proc) => println!(
"daemon: running (PID {}, stdio-only — no HTTP listener)",
proc.pid
),
None => println!("daemon: not running (serve not started)"),
},
}
}
Commands::Dashboard => cli::dashboard::handle(&out).await?,
Commands::Start { http } => cli::start::handle(http, &out).await?,
Commands::Stop => cli::stop::handle(&out).await?,
Commands::Doctor => cli::doctor::handle(&out).await?,
Commands::Completions { shell } => {
let mut cmd = Cli::command();
let name = cmd.get_name().to_string();
generate(shell, &mut cmd, name, &mut io::stdout());
}
}
Ok(())
}
fn resolve_palace_for_serve(explicit: Option<&str>) -> Option<String> {
detect_serve_palace(explicit)
}
async fn ensure_daemon() {
if daemon_alive() {
return;
}
if lock_file_held() {
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
while std::time::Instant::now() < deadline {
tokio::time::sleep(std::time::Duration::from_millis(200)).await;
if daemon_alive() {
return;
}
}
eprintln!(
"[warn] another trusty-memory instance appears to hold the daemon lock; \
not spawning a duplicate"
);
return;
}
if let Ok(exe) = std::env::current_exe() {
let _ = std::process::Command::new(&exe)
.arg("serve")
.arg("--http")
.stdin(std::process::Stdio::null())
.stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::null())
.spawn();
}
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
while std::time::Instant::now() < deadline {
tokio::time::sleep(std::time::Duration::from_millis(200)).await;
if daemon_alive() {
return;
}
}
eprintln!("[warn] daemon did not start within 5 s; proceeding without HTTP server");
}
fn daemon_alive() -> bool {
cli::daemon_probe::probe_daemon().is_some() || cli::daemon_probe::probe_pid_file().is_some()
}
fn lock_file_held() -> bool {
let Ok(root) = trusty_common::resolve_data_dir("trusty-memory") else {
return false;
};
let path = root.join("trusty-memory.lock");
let Ok(file) = std::fs::OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(&path)
else {
return false;
};
match FileExt::try_lock(&file) {
Ok(()) => {
let _ = FileExt::unlock(&file);
false
}
Err(_) => true,
}
}