use anyhow::{anyhow, Result};
use std::process::Stdio;
use std::time::Duration;
use tokio::process::{Child, Command, ChildStdin, ChildStdout};
use tokio::time::timeout;
use tokio::io::AsyncWriteExt;
use tracing::{debug, error, info, warn};
pub struct LspServerProcess {
child: Child,
command: String,
args: Vec<String>,
timeout_ms: u64,
}
impl LspServerProcess {
pub async fn new(command: &str, args: &[&str], timeout_ms: u64) -> Result<Self> {
info!("Starting LSP server: {} {:?}", command, args);
if !Self::command_exists(command).await {
return Err(anyhow!("LSP server command not found: {}", command));
}
let args_vec: Vec<String> = args.iter().map(|s| s.to_string()).collect();
let mut cmd = Command::new(command);
cmd.args(&args_vec)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true);
let child = cmd.spawn()?;
debug!("LSP server started with PID: {:?}", child.id());
Ok(Self {
child,
command: command.to_string(),
args: args_vec,
timeout_ms,
})
}
async fn command_exists(command: &str) -> bool {
let output = Command::new("which")
.arg(command)
.output()
.await;
match output {
Ok(output) => output.status.success(),
Err(_) => {
let fallback = Command::new(command)
.arg("--version")
.stdout(Stdio::null())
.stderr(Stdio::null())
.status()
.await;
fallback.map(|status| status.success()).unwrap_or(false)
}
}
}
pub fn stdin(&mut self) -> ChildStdin {
self.child.stdin.take().expect("Failed to get stdin")
}
pub fn stdout(&mut self) -> ChildStdout {
self.child.stdout.take().expect("Failed to get stdout")
}
pub fn is_running(&mut self) -> bool {
match self.child.try_wait() {
Ok(Some(_)) => false, Ok(None) => true, Err(_) => false, }
}
pub fn pid(&self) -> Option<u32> {
self.child.id()
}
pub async fn wait_with_timeout(&mut self, timeout_ms: u64) -> Result<std::process::ExitStatus> {
Ok(timeout(Duration::from_millis(timeout_ms), self.child.wait()).await??)
}
pub async fn shutdown(&mut self) -> Result<()> {
if !self.is_running() {
debug!("LSP server already stopped");
return Ok(());
}
info!("Shutting down LSP server: {} (PID: {:?})", self.command, self.pid());
if let Some(stdin) = self.child.stdin.as_mut() {
let shutdown_msg = r#"Content-Length: 56
{"jsonrpc":"2.0","method":"shutdown","id":999999999}"#;
let exit_msg = r#"Content-Length: 43
{"jsonrpc":"2.0","method":"exit","params":null}"#;
let _ = stdin.write_all(shutdown_msg.as_bytes()).await;
let _ = stdin.write_all(exit_msg.as_bytes()).await;
let _ = stdin.flush().await;
}
let graceful_timeout = Duration::from_millis(3000);
match timeout(graceful_timeout, self.child.wait()).await {
Ok(Ok(exit_status)) => {
info!("LSP server exited gracefully: {:?}", exit_status);
return Ok(());
}
Ok(Err(e)) => {
warn!("Error waiting for graceful exit: {:?}", e);
}
Err(_) => {
warn!("LSP server did not exit gracefully within timeout");
}
}
if self.is_running() {
warn!("Force killing LSP server: {}", self.command);
match self.child.kill().await {
Ok(_) => {
info!("LSP server force killed");
let _ = timeout(Duration::from_millis(1000), self.child.wait()).await;
}
Err(e) => {
error!("Failed to force kill LSP server: {:?}", e);
return Err(anyhow!("Failed to kill LSP server: {:?}", e));
}
}
}
Ok(())
}
pub async fn restart(&mut self) -> Result<()> {
info!("Restarting LSP server: {}", self.command);
self.shutdown().await?;
let mut cmd = Command::new(&self.command);
cmd.args(&self.args)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true);
self.child = cmd.spawn()?;
info!("LSP server restarted with new PID: {:?}", self.child.id());
Ok(())
}
pub async fn health_check(&mut self) -> ServerHealth {
if !self.is_running() {
return ServerHealth {
is_running: false,
pid: None,
uptime_ms: 0,
memory_usage_kb: 0,
cpu_usage_percent: 0.0,
};
}
let pid = self.pid();
let (memory_kb, cpu_percent) = match pid {
Some(pid) => Self::get_process_stats(pid).await.unwrap_or((0, 0.0)),
None => (0, 0.0),
};
ServerHealth {
is_running: true,
pid,
uptime_ms: 0, memory_usage_kb: memory_kb,
cpu_usage_percent: cpu_percent,
}
}
async fn get_process_stats(pid: u32) -> Result<(u64, f64)> {
let output = Command::new("ps")
.args(&["-p", &pid.to_string(), "-o", "rss,pcpu", "--no-headers"])
.output()
.await?;
if !output.status.success() {
return Err(anyhow!("Failed to get process stats"));
}
let stdout = String::from_utf8_lossy(&output.stdout);
let parts: Vec<&str> = stdout.trim().split_whitespace().collect();
if parts.len() >= 2 {
let memory_kb = parts[0].parse::<u64>().unwrap_or(0);
let cpu_percent = parts[1].parse::<f64>().unwrap_or(0.0);
Ok((memory_kb, cpu_percent))
} else {
Ok((0, 0.0))
}
}
}
#[derive(Debug, Clone)]
pub struct ServerHealth {
pub is_running: bool,
pub pid: Option<u32>,
pub uptime_ms: u64,
pub memory_usage_kb: u64,
pub cpu_usage_percent: f64,
}
impl Drop for LspServerProcess {
fn drop(&mut self) {
if self.is_running() {
warn!("LSP server process dropped while still running, killing: {}", self.command);
let _ = self.child.start_kill();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::time::sleep;
#[tokio::test]
async fn test_command_exists() {
assert!(LspServerProcess::command_exists("echo").await);
assert!(!LspServerProcess::command_exists("definitely_not_a_command_12345").await);
}
#[tokio::test]
async fn test_process_lifecycle() {
let mut process = LspServerProcess::new("cat", &[], 5000).await.unwrap();
assert!(process.is_running());
assert!(process.pid().is_some());
let health = process.health_check().await;
assert!(health.is_running);
assert!(health.pid.is_some());
process.shutdown().await.unwrap();
sleep(Duration::from_millis(100)).await;
assert!(!process.is_running());
}
}