mcp-memory/nvim-core/src/headless.rs

use std::sync::Arc;
use std::sync::LazyLock;
use tokio::io::AsyncBufReadExt;

pub static ACTIVE_SOCKET: LazyLock<Arc<tokio::sync::Mutex<Option<String>>>> =
    LazyLock::new(|| Arc::new(tokio::sync::Mutex::new(None)));

pub static HEADLESS_PROC: LazyLock<Arc<tokio::sync::Mutex<Option<tokio::process::Child>>>> =
    LazyLock::new(|| Arc::new(tokio::sync::Mutex::new(None)));

unsafe extern "C" {
    fn atexit(f: unsafe extern "C" fn()) -> i32;
}

static HEADLESS_PID: std::sync::atomic::AtomicU32 = std::sync::atomic::AtomicU32::new(0);
static ONCE_INIT: std::sync::Once = std::sync::Once::new();

unsafe extern "C" fn cleanup_on_exit() {
    let pid = HEADLESS_PID.swap(0, std::sync::atomic::Ordering::SeqCst);
    if pid != 0 {
        #[cfg(windows)]
        {
            let _ = std::process::Command::new("taskkill")
                .args(["/F", "/T", "/PID", &pid.to_string()])
                .output();
        }
        #[cfg(unix)]
        {
            let _ = std::process::Command::new("kill")
                .args(["-9", &pid.to_string()])
                .output();
        }
    }
}

pub async fn spawn_headless_nvim() -> Result<String, String> {
    if std::process::Command::new("nvim")
        .arg("--version")
        .output()
        .is_err()
    {
        return Err("Error: Neovim is not installed or not in PATH. Cannot execute nvim tool. Fallback to basic terminal tools.".to_string());
    }

    ONCE_INIT.call_once(|| unsafe {
        atexit(cleanup_on_exit);
    });

    #[cfg(windows)]
    let socket_name = format!(
        r"\\.\pipe\agy-headless-nvim-{}-{}",
        std::process::id(),
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_nanos())
            .unwrap_or(0)
    );
    #[cfg(unix)]
    let socket_name = format!(
        "/tmp/agy-headless-nvim-{}-{}.sock",
        std::process::id(),
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_nanos())
            .unwrap_or(0)
    );

    let mut child = tokio::process::Command::new("nvim")
        .arg("--headless")
        .arg("--listen")
        .arg(&socket_name)
        .arg("--cmd")
        .arg("lua vim.schedule(function() io.stdout:write('READY\\n') io.stdout:flush() end)")
        .stdin(std::process::Stdio::piped())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::null())
        .kill_on_drop(true)
        .spawn()
        .map_err(|e| format!("Failed to spawn headless Neovim: {}", e))?;

    if let Some(id) = child.id() {
        HEADLESS_PID.store(id, std::sync::atomic::Ordering::SeqCst);
    }

    if let Some(stdout) = child.stdout.take() {
        let mut reader = tokio::io::BufReader::new(stdout);
        let mut ready_line = String::new();
        match tokio::time::timeout(
            tokio::time::Duration::from_millis(2000),
            reader.read_line(&mut ready_line),
        )
        .await
        {
            Ok(Ok(_)) => {
                if ready_line.trim() != "READY" {
                    HEADLESS_PID.store(0, std::sync::atomic::Ordering::SeqCst);
                    let _ = child.kill().await;
                    return Err(format!(
                        "Headless Neovim child output unexpected banner: {:?}",
                        ready_line
                    ));
                }
            }
            Ok(Err(e)) => {
                HEADLESS_PID.store(0, std::sync::atomic::Ordering::SeqCst);
                let _ = child.kill().await;
                return Err(format!(
                    "Failed to read readiness from headless Neovim: {}",
                    e
                ));
            }
            Err(_) => {
                HEADLESS_PID.store(0, std::sync::atomic::Ordering::SeqCst);
                let _ = child.kill().await;
                return Err("Timeout waiting for headless Neovim to become ready".to_string());
            }
        }
    }

    let old_child = {
        let mut lock = HEADLESS_PROC.lock().await;
        let old = lock.take();
        *lock = Some(child);
        old
    };

    if let Some(mut old) = old_child {
        let _ = old.kill().await;
    }

    Ok(socket_name)
}

pub async fn kill_headless_nvim() {
    HEADLESS_PID.store(0, std::sync::atomic::Ordering::SeqCst);
    let child_to_kill = {
        let mut proc_lock = HEADLESS_PROC.lock().await;
        proc_lock.take()
    };
    if let Some(mut child) = child_to_kill {
        let _ = child.kill().await;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_headless_nvim_lifecycle() {
        let res = spawn_headless_nvim().await;
        if let Ok(socket_name) = res {
            assert!(socket_name.contains("agy-headless-nvim"));
            kill_headless_nvim().await;
        }
    }
}