use std::path::{Path, PathBuf};
use std::sync::Arc;
use axum::extract::State;
use axum::http::{HeaderMap, StatusCode};
use axum::routing::{get, post};
use axum::{Json, Router};
use serde::{Deserialize, Serialize};
use serde_json::json;
use crate::config::ServerConfig;
use crate::tools::{MenteDbServer, ProcessTurnRequest};
use rmcp::handler::server::wrapper::Parameters;
const DAEMON_FILE: &str = "daemon.json";
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DaemonInfo {
pub port: u16,
pub pid: u32,
pub token: String,
}
pub fn info_path(data_dir: &Path) -> PathBuf {
data_dir.join(DAEMON_FILE)
}
pub fn read_info(data_dir: &Path) -> Option<DaemonInfo> {
let raw = std::fs::read_to_string(info_path(data_dir)).ok()?;
serde_json::from_str(&raw).ok()
}
#[derive(Clone)]
struct DaemonState {
server: Arc<MenteDbServer>,
token: String,
}
fn authorized(state: &DaemonState, headers: &HeaderMap) -> bool {
headers
.get("x-mentedb-token")
.and_then(|v| v.to_str().ok())
.is_some_and(|t| t == state.token)
}
async fn health() -> Json<serde_json::Value> {
Json(json!({ "ok": true, "version": env!("CARGO_PKG_VERSION") }))
}
#[derive(Deserialize)]
struct ContextRequest {
prompt: String,
#[serde(default)]
limit: Option<usize>,
}
async fn context(
State(state): State<DaemonState>,
headers: HeaderMap,
Json(req): Json<ContextRequest>,
) -> Result<Json<serde_json::Value>, StatusCode> {
if !authorized(&state, &headers) {
return Err(StatusCode::UNAUTHORIZED);
}
let ctx = state
.server
.hook_context(&req.prompt, req.limit.unwrap_or(8));
Ok(Json(ctx))
}
#[derive(Deserialize)]
struct TurnRequest {
user_message: String,
#[serde(default)]
assistant_response: Option<String>,
turn_id: u64,
#[serde(default)]
project_context: Option<String>,
}
async fn turn(
State(state): State<DaemonState>,
headers: HeaderMap,
Json(req): Json<TurnRequest>,
) -> Result<Json<serde_json::Value>, StatusCode> {
if !authorized(&state, &headers) {
return Err(StatusCode::UNAUTHORIZED);
}
let result = state
.server
.process_turn(Parameters(ProcessTurnRequest {
user_message: req.user_message,
assistant_response: req.assistant_response,
turn_id: req.turn_id,
project_context: req.project_context,
agent_id: None,
}))
.await
.map_err(|e| {
tracing::error!(error = %e, "daemon process_turn failed");
StatusCode::INTERNAL_SERVER_ERROR
})?;
if let Err(e) = state.server.db_ref().flush() {
tracing::warn!(error = %e, "post-turn flush failed");
}
let text = result
.content
.first()
.and_then(|c| c.as_text())
.map(|t| t.text.clone())
.unwrap_or_default();
let parsed: serde_json::Value =
serde_json::from_str(&text).unwrap_or_else(|_| json!({ "raw": text }));
Ok(Json(parsed))
}
async fn session_context(
State(state): State<DaemonState>,
headers: HeaderMap,
) -> Result<Json<serde_json::Value>, StatusCode> {
if !authorized(&state, &headers) {
return Err(StatusCode::UNAUTHORIZED);
}
Ok(Json(state.server.hook_session_context()))
}
async fn daemon_health_ok(port: u16) -> bool {
let url = format!("http://127.0.0.1:{port}/health");
match reqwest::Client::builder()
.timeout(std::time::Duration::from_secs(1))
.build()
{
Ok(client) => client
.get(&url)
.send()
.await
.map(|r| r.status().is_success())
.unwrap_or(false),
Err(_) => false,
}
}
pub async fn run(config: ServerConfig) -> anyhow::Result<()> {
let data_dir = config.data_dir.clone();
std::fs::create_dir_all(&data_dir)?;
if let Some(existing) = read_info(&data_dir)
&& daemon_health_ok(existing.port).await
{
anyhow::bail!(
"daemon already running for {} (pid {}, port {})",
data_dir.display(),
existing.pid,
existing.port
);
}
let db = mentedb::MenteDb::open(&data_dir)
.map_err(|e| anyhow::anyhow!("failed to open database: {e}"))?;
let server = Arc::new(MenteDbServer::new(db, config));
let token = uuid::Uuid::new_v4().to_string();
let state = DaemonState {
server: Arc::clone(&server),
token: token.clone(),
};
let app = Router::new()
.route("/health", get(health))
.route("/v1/context", post(context))
.route("/v1/turn", post(turn))
.route("/v1/session-context", post(session_context))
.with_state(state);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
let port = listener.local_addr()?.port();
let info = DaemonInfo {
port,
pid: std::process::id(),
token,
};
let info_file = info_path(&data_dir);
std::fs::write(&info_file, serde_json::to_string(&info)?)?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(&info_file, std::fs::Permissions::from_mode(0o600)).ok();
}
tracing::info!(port, data_dir = %data_dir.display(), "hook daemon listening");
let shutdown_server = Arc::clone(&server);
let shutdown_info = info_file.clone();
axum::serve(listener, app)
.with_graceful_shutdown(async move {
let ctrl_c = tokio::signal::ctrl_c();
#[cfg(unix)]
{
let mut sigterm =
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
.expect("failed to install SIGTERM handler");
tokio::select! {
_ = ctrl_c => {},
_ = sigterm.recv() => {},
}
}
#[cfg(not(unix))]
{
let _ = ctrl_c.await;
}
})
.await?;
if let Err(e) = shutdown_server.db_ref().close() {
tracing::warn!(error = %e, "database close failed during daemon shutdown");
}
std::fs::remove_file(&shutdown_info).ok();
tracing::info!("hook daemon stopped");
Ok(())
}