lspz 0.11.4

//! Daemon server — long-lived LSP session manager.
//!
//! Accepts connections on a Unix domain socket, multiplexes LSP requests
//! over the internal [`LspPool`], and returns results as JSON-RPC responses.

use std::future::Future;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};

use anyhow::Context;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::{UnixListener, UnixStream};
use tokio::sync::{Mutex, watch};
use tracing::{debug, error, info, warn};

use super::protocol::DaemonResponse;
use super::status::DaemonStatus;
use crate::mcp::LspPool;

/// The daemon server.
pub struct DaemonServer {
    socket_path: PathBuf,
    pool: Arc<Mutex<LspPool>>,
    status: Arc<Mutex<DaemonStatus>>,
    /// Count of currently-connected clients. The idle reaper consults this to
    /// avoid exiting while a client (e.g. a live `lspz mcp` process) is still
    /// attached, even if the pool is momentarily empty.
    active_connections: Arc<AtomicU64>,
}

impl DaemonServer {
    /// Create a new daemon server.
    ///
    /// `workspace_root` should be a canonicalized absolute path uniquely
    /// identifying the project workspace.
    pub fn new(socket_path: PathBuf) -> Self {
        Self {
            socket_path,
            pool: Arc::new(Mutex::new(LspPool::new())),
            status: Arc::new(Mutex::new(DaemonStatus::new())),
            active_connections: Arc::new(AtomicU64::new(0)),
        }
    }

    /// Start the daemon, listening on the Unix socket.
    ///
    /// Runs until SIGINT, idle timeout, or `daemon/shutdown`. On exit the
    /// session pool is cleared (killing LSP children) and the socket file is
    /// removed — never via bare `process::exit` that would skip `Drop`.
    pub async fn start(self) -> Result<(), anyhow::Error> {
        // Remove stale socket file
        if self.socket_path.exists() {
            std::fs::remove_file(&self.socket_path).with_context(|| {
                format!("Failed to remove stale socket: {:?}", self.socket_path)
            })?;
        }

        // Ensure parent directory exists
        if let Some(parent) = self.socket_path.parent() {
            std::fs::create_dir_all(parent)?;
        }

        let listener = UnixListener::bind(&self.socket_path)
            .with_context(|| format!("Failed to bind to {:?}", self.socket_path))?;

        let (shutdown_tx, mut shutdown_rx) = watch::channel(false);

        {
            let tx = shutdown_tx.clone();
            tokio::spawn(async move {
                tokio::signal::ctrl_c().await.ok();
                info!("Daemon received shutdown signal");
                let _ = tx.send(true);
            });
        }

        info!(path = %self.socket_path.display(), "Daemon listening");

        let pool = self.pool.clone();
        let status = self.status.clone();
        let active_connections = self.active_connections.clone();

        spawn_idle_reaper(
            pool.clone(),
            status.clone(),
            active_connections.clone(),
            shutdown_tx.clone(),
        );

        loop {
            tokio::select! {
                _ = shutdown_rx.changed() => {
                    if *shutdown_rx.borrow() {
                        info!("Daemon shutdown requested, tearing down");
                        break;
                    }
                }
                accepted = listener.accept() => {
                    match accepted {
                        Ok((stream, _addr)) => {
                            active_connections.fetch_add(1, Ordering::Relaxed);
                            status.lock().await.record_connection();
                            debug!(
                                active = active_connections.load(Ordering::Relaxed),
                                "Daemon: new client connection"
                            );
                            let pool = pool.clone();
                            let status = status.clone();
                            let conns = active_connections.clone();
                            let shutdown_tx = shutdown_tx.clone();
                            tokio::spawn(async move {
                                let _guard = ConnectionGuard(conns);
                                if let Err(e) =
                                    handle_client(stream, pool, status, shutdown_tx).await
                                {
                                    warn!(error = %e, "Client handler exited with error");
                                }
                            });
                        }
                        Err(e) => {
                            error!(error = %e, "Daemon accept error");
                        }
                    }
                }
            }
        }

        // Graceful teardown: drop sessions (kills children), then remove socket.
        pool.lock().await.clear();
        status.lock().await.sessions.clear();
        if self.socket_path.exists() {
            let _ = std::fs::remove_file(&self.socket_path);
        }
        info!(path = %self.socket_path.display(), "Daemon stopped cleanly");
        Ok(())
    }
}

/// RAII guard that decrements the daemon's active-connection counter when
/// dropped. Ensures the count stays accurate even if a handler task panics or
/// is aborted, which the idle reaper relies on to decide when the daemon is
/// truly unused.
struct ConnectionGuard(Arc<AtomicU64>);

impl Drop for ConnectionGuard {
    fn drop(&mut self) {
        self.0.fetch_sub(1, Ordering::Relaxed);
    }
}

/// Reclaim an idle LSP session after this long without I/O.
///
/// Dropping the session kills its child language-server process, freeing the
/// bulk of the memory.
const SESSION_IDLE_TTL: Duration = Duration::from_secs(10 * 60);

/// Once the pool is empty **and** no client is connected, wait this long before
/// the daemon shuts itself down. This is the only safety net for daemons that
/// short-lived MCP clients detached via `setsid()` and then crashed: there is
/// no one left to send `daemon/shutdown`.
const DAEMON_IDLE_TTL: Duration = Duration::from_secs(10 * 60);

/// How often the reaper rechecks idle state.
const REAPER_CHECK_INTERVAL: Duration = Duration::from_secs(60);

/// Spawn the background idle reaper.
///
/// Periodically reaps quiet LSP sessions and, once the daemon has had no
/// sessions and no connections for [`DAEMON_IDLE_TTL`], signals shutdown so
/// the main loop can clear the pool and remove the socket cleanly.
fn spawn_idle_reaper(
    pool: Arc<Mutex<LspPool>>,
    status: Arc<Mutex<DaemonStatus>>,
    active_connections: Arc<AtomicU64>,
    shutdown_tx: watch::Sender<bool>,
) {
    tokio::spawn(async move {
        let mut daemon_idle_since: Option<Instant> = None;
        loop {
            tokio::time::sleep(REAPER_CHECK_INTERVAL).await;

            // Stop if shutdown already requested.
            if *shutdown_tx.borrow() {
                break;
            }

            // Reap idle sessions first — this may empty the pool.
            let reaped = pool.lock().await.reap_idle(SESSION_IDLE_TTL);
            if reaped > 0 {
                info!(reaped, "Reaped idle LSP sessions");
                // Drop status entries whose backing session no longer exists,
                // so `daemon list` does not show ghost sessions.
                let live_keys = pool.lock().await.session_keys();
                status
                    .lock()
                    .await
                    .sessions
                    .retain(|info| live_keys.contains(&info.key));
            }

            let busy = {
                let pool_guard = pool.lock().await;
                !pool_guard.is_empty() || active_connections.load(Ordering::Relaxed) > 0
            };

            if busy {
                if daemon_idle_since.is_some() {
                    debug!("Daemon active again, cancelling pending self-exit");
                }
                daemon_idle_since = None;
            } else if daemon_idle_since.is_none() {
                daemon_idle_since = Some(Instant::now());
                info!(
                    ttl_secs = DAEMON_IDLE_TTL.as_secs(),
                    "Daemon is idle; will self-exit if it stays unused"
                );
            } else if daemon_idle_since.unwrap().elapsed() >= DAEMON_IDLE_TTL {
                info!("Daemon idle timeout reached, requesting shutdown");
                let _ = shutdown_tx.send(true);
                break;
            }
        }
    });
}

/// Handle a single client connection.
async fn handle_client(
    stream: UnixStream,
    pool: Arc<Mutex<LspPool>>,
    status: Arc<Mutex<DaemonStatus>>,
    shutdown_tx: watch::Sender<bool>,
) -> Result<(), anyhow::Error> {
    let (reader, mut writer) = stream.into_split();
    let mut buf_reader = BufReader::new(reader);
    let mut line = String::new();

    loop {
        line.clear();
        let n = buf_reader.read_line(&mut line).await?;
        if n == 0 {
            break; // EOF
        }

        let request: super::protocol::DaemonRequest = match serde_json::from_str(line.trim()) {
            Ok(r) => r,
            Err(e) => {
                warn!(error = %e, "Failed to parse client request");
                let resp = DaemonResponse::err(0, format!("Parse error: {e}"));
                let json = serde_json::to_string(&resp)?;
                writer.write_all(format!("{json}\n").as_bytes()).await?;
                continue;
            }
        };

        let response = dispatch(request, &pool, &status, &shutdown_tx).await;

        let json = serde_json::to_string(&response)?;
        writer.write_all(format!("{json}\n").as_bytes()).await?;
    }

    Ok(())
}

/// Dispatch a client request to the appropriate handler.
async fn dispatch(
    req: super::protocol::DaemonRequest,
    pool: &Arc<Mutex<LspPool>>,
    status: &Arc<Mutex<DaemonStatus>>,
    shutdown_tx: &watch::Sender<bool>,
) -> DaemonResponse {
    let id = req.id;
    status.lock().await.record_request();

    match req.method.as_str() {
        "lsp/spawn" => handle_spawn(id, &req.params, pool, status).await,
        "lsp/request" => handle_lsp_request(id, &req.params, pool, status).await,
        "lsp/notify" => handle_lsp_notify(id, &req.params, pool, status).await,
        "lsp/wait_notify" => handle_wait_notify(id, &req.params, pool, status).await,
        "lsp/sync_document" => handle_sync_document(id, &req.params, pool, status).await,
        "daemon/status" => handle_status(id, status).await,
        "daemon/shutdown" => {
            info!("Daemon shutdown requested by client");
            let _ = shutdown_tx.send(true);
            DaemonResponse::ok(id, serde_json::json!({"ok": true}))
        }
        _ => DaemonResponse::err(id, format!("Unknown method: {}", req.method)),
    }
}

/// Handle `lsp/spawn` — get or create an LSP session.
async fn handle_spawn(
    id: u64,
    params: &serde_json::Value,
    pool: &Arc<Mutex<LspPool>>,
    status: &Arc<Mutex<DaemonStatus>>,
) -> DaemonResponse {
    let spawn: super::protocol::SpawnParams = match serde_json::from_value(params.clone()) {
        Ok(s) => s,
        Err(e) => return DaemonResponse::err(id, format!("Invalid params: {e}")),
    };

    let key = crate::mcp::pool_key(&spawn.language, &spawn.backend, spawn.root_path.as_deref());

    match LspPool::get_or_spawn(
        pool,
        &spawn.language,
        &spawn.backend,
        spawn.root_path.as_deref(),
        &spawn.extra_args,
    )
    .await
    {
        Ok(_session) => {
            let mut s = status.lock().await;
            s.touch_session(
                &key,
                &spawn.language,
                &spawn.backend,
                spawn.root_path.clone(),
            );
            DaemonResponse::ok(
                id,
                serde_json::json!({
                    "session_key": key,
                    "status": "ready",
                }),
            )
        }
        Err(e) => DaemonResponse::err(id, format!("Failed to spawn LSP session: {e}")),
    }
}

/// Handle `lsp/request` — send an LSP request through a session.
///
/// Lock the pool mutex here to get mutable access to the session.
async fn handle_lsp_request(
    id: u64,
    params: &serde_json::Value,
    pool: &Arc<Mutex<LspPool>>,
    status: &Arc<Mutex<DaemonStatus>>,
) -> DaemonResponse {
    let req: super::protocol::LspRequestParams = match serde_json::from_value(params.clone()) {
        Ok(r) => r,
        Err(e) => return DaemonResponse::err(id, format!("Invalid params: {e}")),
    };
    status.lock().await.touch_by_key(&req.session_key);

    let session = match pool.lock().await.get_by_key(&req.session_key) {
        Ok(s) => s,
        Err(e) => return DaemonResponse::err(id, e.to_string()),
    };
    let mut session = session.lock().await;

    match session.send_request(&req.method, req.params).await {
        Ok(result) => DaemonResponse::ok(id, result),
        Err(e) => DaemonResponse::err(id, format!("LSP request failed: {e}")),
    }
}

/// Handle `lsp/notify` — send an LSP notification through a session.
async fn handle_lsp_notify(
    id: u64,
    params: &serde_json::Value,
    pool: &Arc<Mutex<LspPool>>,
    status: &Arc<Mutex<DaemonStatus>>,
) -> DaemonResponse {
    let req: super::protocol::LspNotifyParams = match serde_json::from_value(params.clone()) {
        Ok(r) => r,
        Err(e) => return DaemonResponse::err(id, format!("Invalid params: {e}")),
    };
    status.lock().await.touch_by_key(&req.session_key);

    let session = match pool.lock().await.get_by_key(&req.session_key) {
        Ok(s) => s,
        Err(e) => return DaemonResponse::err(id, e.to_string()),
    };
    let mut session = session.lock().await;

    match session.send_notification(&req.method, req.params).await {
        Ok(()) => DaemonResponse::ok(id, serde_json::json!({"ok": true})),
        Err(e) => DaemonResponse::err(id, format!("LSP notify failed: {e}")),
    }
}

/// Handle `lsp/sync_document` — open or update via session SSOT helper.
async fn handle_sync_document(
    id: u64,
    params: &serde_json::Value,
    pool: &Arc<Mutex<LspPool>>,
    status: &Arc<Mutex<DaemonStatus>>,
) -> DaemonResponse {
    let req: super::protocol::SyncDocumentParams = match serde_json::from_value(params.clone()) {
        Ok(r) => r,
        Err(e) => return DaemonResponse::err(id, format!("Invalid params: {e}")),
    };
    status.lock().await.touch_by_key(&req.session_key);

    let session = match pool.lock().await.get_by_key(&req.session_key) {
        Ok(s) => s,
        Err(e) => return DaemonResponse::err(id, e.to_string()),
    };
    let mut session = session.lock().await;

    match session
        .open_or_update_document(&req.uri, &req.language_id, &req.content)
        .await
    {
        Ok(()) => DaemonResponse::ok(id, serde_json::json!({"ok": true})),
        Err(e) => DaemonResponse::err(id, format!("LSP sync_document failed: {e}")),
    }
}

/// Handle `lsp/wait_notify` — wait for a notification from the LSP server.
async fn handle_wait_notify(
    id: u64,
    params: &serde_json::Value,
    pool: &Arc<Mutex<LspPool>>,
    status: &Arc<Mutex<DaemonStatus>>,
) -> DaemonResponse {
    let req: super::protocol::WaitNotifyParams = match serde_json::from_value(params.clone()) {
        Ok(r) => r,
        Err(e) => return DaemonResponse::err(id, format!("Invalid params: {e}")),
    };
    status.lock().await.touch_by_key(&req.session_key);

    let session = match pool.lock().await.get_by_key(&req.session_key) {
        Ok(s) => s,
        Err(e) => return DaemonResponse::err(id, e.to_string()),
    };
    let mut session = session.lock().await;

    // The client may pass a deadline (`timeout_ms`). The daemon enforces it so
    // it always writes a response *before* the client gives up. Without this,
    // a client that cancels its `lsp_wait_notify` future at its own (shorter)
    // deadline would leave an orphan response on the wire and desync the
    // newline-delimited protocol — see the analysis of the spawn/desync bug.
    let deadline = req.timeout_ms.map(Duration::from_millis);

    let result = if let Some(uri) = &req.filter_uri {
        let uri_clone = uri.clone();
        let fut = session.wait_for_notification_where(&req.method, move |p| {
            p.get("uri").and_then(serde_json::Value::as_str) == Some(&uri_clone)
        });
        apply_wait_deadline(fut, deadline, &req.method).await
    } else {
        let fut = session.wait_for_notification(&req.method);
        apply_wait_deadline(fut, deadline, &req.method).await
    };

    match result {
        Ok(params) => DaemonResponse::ok(id, params),
        Err(e) => DaemonResponse::err(id, format!("Wait for notification failed: {e}")),
    }
}

/// Run a notification-wait future, optionally bounded by a client-supplied
/// deadline. When the deadline elapses, return a timeout error instead of
/// continuing to wait on the LSP transport.
async fn apply_wait_deadline<F>(
    fut: F,
    deadline: Option<Duration>,
    method: &str,
) -> Result<serde_json::Value, anyhow::Error>
where
    F: Future<Output = Result<serde_json::Value, anyhow::Error>>,
{
    match deadline {
        Some(d) => match tokio::time::timeout(d, fut).await {
            Ok(inner) => inner,
            Err(_) => Err(anyhow::anyhow!(
                "timeout waiting for '{method}' notification"
            )),
        },
        None => fut.await,
    }
}

/// Handle `daemon/status` — return current daemon state.
async fn handle_status(id: u64, status: &Arc<Mutex<DaemonStatus>>) -> DaemonResponse {
    let mut s = status.lock().await;
    s.refresh_uptime();
    let json = serde_json::to_value(&*s).unwrap_or(serde_json::Value::Null);
    DaemonResponse::ok(id, json)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::Transport;
    use crate::error::LspzError;
    use crate::mcp::{LspPool, LspSession};

    /// A transport whose `receive()` never completes, simulating an LSP server
    /// that never publishes the waited-for notification.
    struct PendingTransport;

    #[async_trait::async_trait]
    impl Transport for PendingTransport {
        async fn receive(&mut self) -> Result<Vec<u8>, LspzError> {
            std::future::pending().await
        }
        async fn send(&mut self, _data: &[u8]) -> Result<(), LspzError> {
            Ok(())
        }
    }

    #[tokio::test]
    async fn test_daemon_shutdown_signals_watch_not_process_exit() {
        let pool = Arc::new(Mutex::new(LspPool::new()));
        pool.lock().await.insert_session_for_test(
            "rust:fake:/tmp",
            LspSession::with_transport(Box::new(PendingTransport)),
        );
        let status = Arc::new(Mutex::new(DaemonStatus::default()));
        let (shutdown_tx, shutdown_rx) = watch::channel(false);

        let req = crate::daemon::protocol::DaemonRequest {
            id: 42,
            method: "daemon/shutdown".into(),
            params: serde_json::Value::Null,
        };
        let resp = dispatch(req, &pool, &status, &shutdown_tx).await;
        assert!(resp.error.is_none());
        assert!(*shutdown_rx.borrow());
        // Pool is cleared by the main loop after signal; dispatch only signals.
        assert!(!pool.lock().await.is_empty());
    }

    /// `handle_wait_notify` must honour `timeout_ms` instead of falling back to
    /// the LSP session's internal ~30s wait. Enforcing the deadline on the
    /// daemon side is what lets the client await the response directly (no
    /// cancellation, no orphan response line). A timeout here must therefore
    /// arrive at ~`timeout_ms`, not 30s.
    #[tokio::test]
    async fn test_wait_notify_respects_client_timeout_ms() {
        let pool = Arc::new(Mutex::new(LspPool::new()));
        pool.lock().await.insert_session_for_test(
            "rust:fake:/tmp",
            LspSession::with_transport(Box::new(PendingTransport)),
        );

        let params = serde_json::json!({
            "session_key": "rust:fake:/tmp",
            "method": "textDocument/publishDiagnostics",
            "timeout_ms": 200,
        });

        let start = std::time::Instant::now();
        let status = Arc::new(Mutex::new(DaemonStatus::default()));
        let resp = handle_wait_notify(1, &params, &pool, &status).await;
        let elapsed = start.elapsed();

        assert!(
            resp.error.is_some(),
            "expected a timeout error, got success: {:?}",
            resp.result
        );
        assert!(
            elapsed >= Duration::from_millis(150),
            "waited only {elapsed:?}, expected ~200ms"
        );
        assert!(
            elapsed < Duration::from_secs(2),
            "waited {elapsed:?}; the daemon did NOT honour timeout_ms and fell \
             back to the long default wait (would orphan a cancelled client)"
        );
    }
}