scryer-mcp 0.2.1

//! The long-lived daemon that owns the database and serves MCP sessions over a Unix socket.

use std::fs::File;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::time::{Duration, Instant};

use rmcp::ServiceExt;
use scryer_db::{ProjectEvent, ProjectRegistry, ScryerDb, is_registrable_root};
use scryer_engine::{EngineService, WorkspaceWatcher};
use tokio::io::BufReader;
use tokio::net::unix::{OwnedReadHalf, OwnedWriteHalf};
use tokio::net::{UnixListener, UnixStream};
use tokio::sync::{broadcast, watch};

use super::handshake::{
    ControlCommand, DaemonStatus, Hello, HelloMode, HelloReply, PROTOCOL_VERSION, WatchedProject,
    read_json_line, write_json_line,
};
use super::paths::DaemonPaths;
use super::prepare_run_dir;
use crate::ScryerMcpServer;

/// How long a new connection has to send its handshake.
const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);

/// How long a starting daemon retries the PID lock before deciding another daemon owns it.
const PID_LOCK_RETRY: Duration = Duration::from_millis(500);

/// Exclusive lock proving this process is the daemon for a database.
///
/// Held for the daemon's lifetime; the file contains the daemon's PID.
pub struct DaemonLock {
    _file: File,
}

impl DaemonLock {
    /// Acquire the daemon lock, or return `None` if another daemon holds it.
    pub async fn acquire(paths: &DaemonPaths) -> anyhow::Result<Option<Self>> {
        prepare_run_dir(paths.run_dir())?;
        let mut file = File::options()
            .read(true)
            .write(true)
            .create(true)
            .truncate(false)
            .open(&paths.pid_file)?;

        // Short retries absorb a status probe briefly holding the lock.
        let deadline = Instant::now() + PID_LOCK_RETRY;
        loop {
            match file.try_lock() {
                Ok(()) => break,
                Err(std::fs::TryLockError::WouldBlock) if Instant::now() < deadline => {
                    tokio::time::sleep(Duration::from_millis(25)).await;
                }
                Err(std::fs::TryLockError::WouldBlock) => return Ok(None),
                Err(std::fs::TryLockError::Error(e)) => return Err(e.into()),
            }
        }

        file.set_len(0)?;
        writeln!(file, "{}", std::process::id())?;
        file.flush()?;
        Ok(Some(Self { _file: file }))
    }
}

/// Whether a daemon currently holds the lock for `paths`.
pub fn daemon_running(paths: &DaemonPaths) -> bool {
    let Ok(file) = File::options().read(true).open(&paths.pid_file) else {
        return false;
    };
    matches!(file.try_lock(), Err(std::fs::TryLockError::WouldBlock))
}

/// Open the database at `paths.db_path` and run the daemon until it is stopped or idles out.
///
/// Returns `Ok(())` without serving if another daemon already owns this database.
pub async fn run_daemon(paths: DaemonPaths, idle_timeout: Option<Duration>) -> anyhow::Result<()> {
    let Some(lock) = DaemonLock::acquire(&paths).await? else {
        tracing::info!(
            "A Scryer daemon is already running for {}",
            paths.db_path.display()
        );
        return Ok(());
    };

    let db_url = paths.db_path.to_string_lossy().to_string();
    tracing::info!("Scryer daemon opening database {db_url}");
    let db = ScryerDb::connect_with_auth(&db_url, None).await?;
    let registry = Arc::new(ProjectRegistry::new(db.inner()).await?);
    let engine = EngineService::new(db.clone());

    Daemon::new(paths, db, engine, registry)?
        .serve(lock, idle_timeout)
        .await
}

/// Shared daemon state; every connection gets a handle to it.
pub struct Daemon {
    paths: DaemonPaths,
    db: ScryerDb,
    engine: EngineService,
    registry: Arc<ProjectRegistry>,
    /// Watches every registered project while the daemon runs.
    watcher: WorkspaceWatcher,
    sessions: watch::Sender<usize>,
    shutdown: watch::Sender<bool>,
    /// Set once an idle-only stop was accepted; new sessions are refused from then on.
    stopping: AtomicBool,
    next_session: AtomicU64,
    started: Instant,
}

/// Decrements the live session count when a session ends.
struct SessionGuard<'a>(&'a watch::Sender<usize>);

impl Drop for SessionGuard<'_> {
    fn drop(&mut self) {
        self.0.send_modify(|count| *count -= 1);
    }
}

/// Removes the socket file when the daemon stops listening.
struct SocketCleanup(PathBuf);

impl Drop for SocketCleanup {
    fn drop(&mut self) {
        let _ = std::fs::remove_file(&self.0);
    }
}

impl Daemon {
    pub fn new(
        paths: DaemonPaths,
        db: ScryerDb,
        engine: EngineService,
        registry: Arc<ProjectRegistry>,
    ) -> anyhow::Result<Arc<Self>> {
        Ok(Arc::new(Self {
            paths,
            db,
            watcher: engine.watcher()?,
            engine,
            registry,
            sessions: watch::Sender::new(0),
            shutdown: watch::Sender::new(false),
            stopping: AtomicBool::new(false),
            next_session: AtomicU64::new(1),
            started: Instant::now(),
        }))
    }

    /// Bind the socket and serve connections until stopped, signalled, or idle.
    ///
    /// `idle_timeout` of `None` keeps the daemon alive with no sessions.
    pub async fn serve(
        self: Arc<Self>,
        _lock: DaemonLock,
        idle_timeout: Option<Duration>,
    ) -> anyhow::Result<()> {
        // Holding the daemon lock means any existing socket file is stale.
        let _ = std::fs::remove_file(&self.paths.socket);
        let listener = UnixListener::bind(&self.paths.socket)?;
        let _cleanup = SocketCleanup(self.paths.socket.clone());
        tracing::info!(
            "Scryer daemon {} (pid {}) listening on {}",
            crate::mcp_version(),
            std::process::id(),
            self.paths.socket.display()
        );

        tokio::spawn(Arc::clone(&self).maintain_watches());

        let mut shutdown = self.shutdown.subscribe();
        let idle = idle_monitor(self.sessions.subscribe(), idle_timeout);
        tokio::pin!(idle);
        let terminate = terminate_signal();
        tokio::pin!(terminate);

        loop {
            tokio::select! {
                accepted = listener.accept() => match accepted {
                    Ok((stream, _)) => {
                        let daemon = Arc::clone(&self);
                        tokio::spawn(async move {
                            if let Err(e) = daemon.handle_connection(stream).await {
                                tracing::debug!("Daemon connection ended with error: {e}");
                            }
                        });
                    }
                    Err(e) => tracing::warn!("Daemon accept failed: {e}"),
                },
                _ = &mut idle => {
                    let idle_secs = idle_timeout.unwrap_or_default().as_secs();
                    tracing::info!("No sessions for {idle_secs}s; daemon exiting");
                    break;
                }
                _ = shutdown.wait_for(|stop| *stop) => {
                    tracing::info!("Daemon stop requested");
                    break;
                }
                _ = &mut terminate => {
                    tracing::info!("Daemon received termination signal");
                    break;
                }
            }
        }

        // Release inotify watches now; tasks may keep the daemon alive briefly.
        self.stop();
        for watched in self.watcher.projects() {
            self.watcher.remove_project(watched.project_id);
        }
        Ok(())
    }

    /// Watch every registered project, then follow registrations and removals until shutdown.
    ///
    /// After watching starts, each project gets a background incremental index so edits made
    /// while the daemon was not running are picked up.
    async fn maintain_watches(self: Arc<Self>) {
        // Subscribe before listing so no registration falls between the two.
        let mut events = self.registry.subscribe();
        let mut shutdown = self.shutdown.subscribe();

        let mut catch_up = Vec::new();
        for project in self.registry.list_projects().await {
            if self.watch_project(project.id, &project.root_path) {
                catch_up.push((project.id, PathBuf::from(project.root_path)));
            }
        }
        let engine = self.engine.clone();
        tokio::spawn(async move {
            for (project_id, root) in catch_up {
                if let Err(e) = engine.index_project(project_id, &root).await {
                    tracing::warn!("Catch-up index of {} failed: {e}", root.display());
                }
            }
        });

        loop {
            let event = tokio::select! {
                event = events.recv() => event,
                _ = shutdown.wait_for(|stop| *stop) => return,
            };
            match event {
                Ok(ProjectEvent::Registered(project)) => {
                    self.watch_project(project.id, &project.root_path);
                }
                Ok(ProjectEvent::Removed(project_id)) => self.watcher.remove_project(project_id),
                Err(broadcast::error::RecvError::Lagged(_)) => self.reconcile_watches().await,
                Err(broadcast::error::RecvError::Closed) => return,
            }
        }
    }

    /// Start watching a project, logging failures (e.g. a deleted root). Returns success.
    fn watch_project(&self, project_id: u64, root: &str) -> bool {
        match self.watcher.add_project(project_id, Path::new(root)) {
            Ok(_) => true,
            Err(e) => {
                tracing::warn!("Not watching project {project_id} at {root}: {e}");
                false
            }
        }
    }

    /// Make the watched set match the registry, after missing registry events.
    async fn reconcile_watches(&self) {
        let registered = self.registry.list_projects().await;
        for watched in self.watcher.projects() {
            if !registered.iter().any(|p| p.id == watched.project_id) {
                self.watcher.remove_project(watched.project_id);
            }
        }
        for project in registered {
            self.watch_project(project.id, &project.root_path);
        }
    }

    /// Ask a running [`Daemon::serve`] loop to exit.
    pub fn stop(&self) {
        self.shutdown.send_replace(true);
    }

    /// Current status snapshot.
    pub fn status(&self) -> DaemonStatus {
        DaemonStatus {
            pid: std::process::id(),
            version: crate::mcp_version().to_string(),
            build: Some(crate::build_info::build_id().to_string()),
            db_path: self.paths.db_path.clone(),
            sessions: *self.sessions.borrow(),
            uptime_secs: self.started.elapsed().as_secs(),
            watched: self
                .watcher
                .projects()
                .into_iter()
                .map(|w| WatchedProject {
                    project_id: w.project_id,
                    root: w.root,
                    directories: w.directories,
                })
                .collect(),
        }
    }

    async fn handle_connection(self: Arc<Self>, stream: UnixStream) -> anyhow::Result<()> {
        let (read, mut write) = stream.into_split();
        let mut reader = BufReader::new(read);

        let hello: Hello =
            match tokio::time::timeout(HANDSHAKE_TIMEOUT, read_json_line(&mut reader)).await {
                Ok(Ok(hello)) => hello,
                Ok(Err(e)) => {
                    let reply = HelloReply::error(format!("invalid handshake: {e}"));
                    let _ = write_json_line(&mut write, &reply).await;
                    return Err(e);
                }
                Err(_) => anyhow::bail!("handshake timed out"),
            };

        if hello.scryer != PROTOCOL_VERSION {
            let reply = HelloReply::error(format!(
                "unsupported handshake protocol {} (daemon speaks {PROTOCOL_VERSION})",
                hello.scryer
            ));
            return write_json_line(&mut write, &reply).await;
        }
        if hello.version != crate::mcp_version() {
            tracing::warn!(
                "Client version {} differs from daemon version {}",
                hello.version,
                crate::mcp_version()
            );
        }

        match hello.mode {
            HelloMode::Control {
                cmd: ControlCommand::Status,
            } => {
                let reply = HelloReply {
                    status: Some(self.status()),
                    ..HelloReply::ok()
                };
                write_json_line(&mut write, &reply).await
            }
            HelloMode::Control {
                cmd: ControlCommand::Stop,
            } => {
                write_json_line(&mut write, &HelloReply::ok()).await?;
                self.stop();
                Ok(())
            }
            HelloMode::Control {
                cmd: ControlCommand::StopIfIdle,
            } => {
                // Checked under the session counter's lock, the same lock `serve_session`
                // takes to admit a session, so a session can't slip in between.
                let mut idle = false;
                self.sessions.send_if_modified(|count| {
                    idle = *count == 0;
                    if idle {
                        self.stopping.store(true, Ordering::SeqCst);
                    }
                    false
                });
                if idle {
                    write_json_line(&mut write, &HelloReply::ok()).await?;
                    self.stop();
                } else {
                    let reply = HelloReply {
                        status: Some(self.status()),
                        ..HelloReply::ok()
                    };
                    write_json_line(&mut write, &reply).await?;
                }
                Ok(())
            }
            HelloMode::Mcp { cwd, watch } => self.serve_session(reader, write, cwd, watch).await,
        }
    }

    async fn serve_session(
        &self,
        reader: BufReader<OwnedReadHalf>,
        mut write: OwnedWriteHalf,
        cwd: Option<PathBuf>,
        watch: bool,
    ) -> anyhow::Result<()> {
        let admitted = self.sessions.send_if_modified(|count| {
            if self.stopping.load(Ordering::SeqCst) {
                return false;
            }
            *count += 1;
            true
        });
        if !admitted {
            let reply = HelloReply::error(
                "Scryer daemon is shutting down to be replaced by a newer build; retry",
            );
            return write_json_line(&mut write, &reply).await;
        }
        let _session = SessionGuard(&self.sessions);

        let session_id = format!(
            "scryer-sess-{}-{}",
            std::process::id(),
            self.next_session.fetch_add(1, Ordering::Relaxed)
        );
        let server = ScryerMcpServer::new(
            self.db.clone(),
            self.engine.clone(),
            Arc::clone(&self.registry),
            Some(session_id),
        )?;

        if let Some(cwd) = cwd {
            server.set_session_cwd(cwd.clone()).await;
            if watch && let Err(e) = self.register_cwd(&cwd).await {
                tracing::warn!("Could not register {}: {e}", cwd.display());
            }
        }

        write_json_line(&mut write, &HelloReply::ok()).await?;

        // The BufReader may already hold MCP bytes sent right after the handshake.
        let running = server.serve((reader, write)).await?;
        let _ = running.waiting().await;
        Ok(())
    }

    /// Register `cwd` as a project (which starts watching it) and index it in the background,
    /// unless a registered project already contains it.
    ///
    /// Skips directories that don't look like a project root (see [`is_registrable_root`]),
    /// e.g. an editor launched from the home directory.
    async fn register_cwd(&self, cwd: &Path) -> anyhow::Result<()> {
        if self.registry.resolve_path(cwd).await.is_some() {
            return Ok(());
        }
        if !is_registrable_root(cwd) {
            tracing::info!(
                "Not registering {}: not inside a registered project and not a project root",
                cwd.display()
            );
            return Ok(());
        }
        let project = self.registry.register_or_update(cwd, None).await?;
        let engine = self.engine.clone();
        tokio::spawn(async move {
            let root = PathBuf::from(&project.root_path);
            if let Err(e) = engine.index_project(project.id, &root).await {
                tracing::warn!("Initial index of {} failed: {e}", root.display());
            }
        });
        Ok(())
    }
}

/// Resolves once the session count has stayed at zero for `idle_timeout`.
async fn idle_monitor(mut sessions: watch::Receiver<usize>, idle_timeout: Option<Duration>) {
    let Some(idle_timeout) = idle_timeout else {
        return std::future::pending().await;
    };

    loop {
        if *sessions.borrow_and_update() == 0 {
            tokio::select! {
                _ = tokio::time::sleep(idle_timeout) => return,
                changed = sessions.changed() => {
                    if changed.is_err() {
                        return std::future::pending().await;
                    }
                }
            }
        } else if sessions.changed().await.is_err() {
            return std::future::pending().await;
        }
    }
}

/// Resolves on SIGINT or SIGTERM.
async fn terminate_signal() {
    use tokio::signal::unix::{SignalKind, signal};

    match signal(SignalKind::terminate()) {
        Ok(mut sigterm) => {
            tokio::select! {
                _ = tokio::signal::ctrl_c() => {}
                _ = sigterm.recv() => {}
            }
        }
        Err(_) => {
            let _ = tokio::signal::ctrl_c().await;
        }
    }
}