Skip to main content

lens_core/lsp/
server_process.rs

1//! LSP Server Process Management
2//! 
3//! Manages language server processes with proper lifecycle and error handling
4
5use anyhow::{anyhow, Result};
6use std::process::Stdio;
7use std::time::Duration;
8use tokio::process::{Child, Command, ChildStdin, ChildStdout};
9use tokio::time::timeout;
10use tokio::io::AsyncWriteExt;
11use tracing::{debug, error, info, warn};
12
13/// Manages an LSP server process
14pub struct LspServerProcess {
15    child: Child,
16    command: String,
17    args: Vec<String>,
18    timeout_ms: u64,
19}
20
21impl LspServerProcess {
22    /// Start a new LSP server process
23    pub async fn new(command: &str, args: &[&str], timeout_ms: u64) -> Result<Self> {
24        info!("Starting LSP server: {} {:?}", command, args);
25        
26        // Check if command exists
27        if !Self::command_exists(command).await {
28            return Err(anyhow!("LSP server command not found: {}", command));
29        }
30
31        let args_vec: Vec<String> = args.iter().map(|s| s.to_string()).collect();
32        
33        let mut cmd = Command::new(command);
34        cmd.args(&args_vec)
35           .stdin(Stdio::piped())
36           .stdout(Stdio::piped())
37           .stderr(Stdio::piped())
38           .kill_on_drop(true);
39
40        let child = cmd.spawn()?;
41
42        debug!("LSP server started with PID: {:?}", child.id());
43
44        Ok(Self {
45            child,
46            command: command.to_string(),
47            args: args_vec,
48            timeout_ms,
49        })
50    }
51
52    /// Check if a command exists in PATH
53    async fn command_exists(command: &str) -> bool {
54        let output = Command::new("which")
55            .arg(command)
56            .output()
57            .await;
58        
59        match output {
60            Ok(output) => output.status.success(),
61            Err(_) => {
62                // Fallback: try to run the command with --version
63                let fallback = Command::new(command)
64                    .arg("--version")
65                    .stdout(Stdio::null())
66                    .stderr(Stdio::null())
67                    .status()
68                    .await;
69                
70                fallback.map(|status| status.success()).unwrap_or(false)
71            }
72        }
73    }
74
75    /// Get stdin handle for sending requests
76    pub fn stdin(&mut self) -> ChildStdin {
77        self.child.stdin.take().expect("Failed to get stdin")
78    }
79
80    /// Get stdout handle for reading responses
81    pub fn stdout(&mut self) -> ChildStdout {
82        self.child.stdout.take().expect("Failed to get stdout")
83    }
84
85    /// Check if the process is still running
86    pub fn is_running(&mut self) -> bool {
87        match self.child.try_wait() {
88            Ok(Some(_)) => false, // Process has exited
89            Ok(None) => true,     // Process is still running
90            Err(_) => false,      // Error checking status, assume dead
91        }
92    }
93
94    /// Get process ID
95    pub fn pid(&self) -> Option<u32> {
96        self.child.id()
97    }
98
99    /// Wait for process to exit with timeout
100    pub async fn wait_with_timeout(&mut self, timeout_ms: u64) -> Result<std::process::ExitStatus> {
101        Ok(timeout(Duration::from_millis(timeout_ms), self.child.wait()).await??)
102    }
103
104    /// Gracefully shutdown the server
105    pub async fn shutdown(&mut self) -> Result<()> {
106        if !self.is_running() {
107            debug!("LSP server already stopped");
108            return Ok(());
109        }
110
111        info!("Shutting down LSP server: {} (PID: {:?})", self.command, self.pid());
112
113        // First try graceful shutdown
114        if let Some(stdin) = self.child.stdin.as_mut() {
115            // Send LSP shutdown sequence
116            let shutdown_msg = r#"Content-Length: 56
117
118{"jsonrpc":"2.0","method":"shutdown","id":999999999}"#;
119            let exit_msg = r#"Content-Length: 43
120
121{"jsonrpc":"2.0","method":"exit","params":null}"#;
122            
123            let _ = stdin.write_all(shutdown_msg.as_bytes()).await;
124            let _ = stdin.write_all(exit_msg.as_bytes()).await;
125            let _ = stdin.flush().await;
126        }
127
128        // Wait for graceful exit with timeout
129        let graceful_timeout = Duration::from_millis(3000);
130        match timeout(graceful_timeout, self.child.wait()).await {
131            Ok(Ok(exit_status)) => {
132                info!("LSP server exited gracefully: {:?}", exit_status);
133                return Ok(());
134            }
135            Ok(Err(e)) => {
136                warn!("Error waiting for graceful exit: {:?}", e);
137            }
138            Err(_) => {
139                warn!("LSP server did not exit gracefully within timeout");
140            }
141        }
142
143        // Force kill if graceful shutdown failed
144        if self.is_running() {
145            warn!("Force killing LSP server: {}", self.command);
146            match self.child.kill().await {
147                Ok(_) => {
148                    info!("LSP server force killed");
149                    // Wait a bit for the kill to take effect
150                    let _ = timeout(Duration::from_millis(1000), self.child.wait()).await;
151                }
152                Err(e) => {
153                    error!("Failed to force kill LSP server: {:?}", e);
154                    return Err(anyhow!("Failed to kill LSP server: {:?}", e));
155                }
156            }
157        }
158
159        Ok(())
160    }
161
162    /// Restart the server process
163    pub async fn restart(&mut self) -> Result<()> {
164        info!("Restarting LSP server: {}", self.command);
165        
166        // Shutdown existing process
167        self.shutdown().await?;
168
169        // Start new process
170        let mut cmd = Command::new(&self.command);
171        cmd.args(&self.args)
172           .stdin(Stdio::piped())
173           .stdout(Stdio::piped())
174           .stderr(Stdio::piped())
175           .kill_on_drop(true);
176
177        self.child = cmd.spawn()?;
178
179        info!("LSP server restarted with new PID: {:?}", self.child.id());
180        Ok(())
181    }
182
183    /// Get server health status
184    pub async fn health_check(&mut self) -> ServerHealth {
185        if !self.is_running() {
186            return ServerHealth {
187                is_running: false,
188                pid: None,
189                uptime_ms: 0,
190                memory_usage_kb: 0,
191                cpu_usage_percent: 0.0,
192            };
193        }
194
195        let pid = self.pid();
196        let (memory_kb, cpu_percent) = match pid {
197            Some(pid) => Self::get_process_stats(pid).await.unwrap_or((0, 0.0)),
198            None => (0, 0.0),
199        };
200
201        ServerHealth {
202            is_running: true,
203            pid,
204            uptime_ms: 0, // TODO: Track start time
205            memory_usage_kb: memory_kb,
206            cpu_usage_percent: cpu_percent,
207        }
208    }
209
210    /// Get process memory and CPU statistics
211    async fn get_process_stats(pid: u32) -> Result<(u64, f64)> {
212        // Use ps command to get process stats
213        let output = Command::new("ps")
214            .args(&["-p", &pid.to_string(), "-o", "rss,pcpu", "--no-headers"])
215            .output()
216            .await?;
217
218        if !output.status.success() {
219            return Err(anyhow!("Failed to get process stats"));
220        }
221
222        let stdout = String::from_utf8_lossy(&output.stdout);
223        let parts: Vec<&str> = stdout.trim().split_whitespace().collect();
224        
225        if parts.len() >= 2 {
226            let memory_kb = parts[0].parse::<u64>().unwrap_or(0);
227            let cpu_percent = parts[1].parse::<f64>().unwrap_or(0.0);
228            Ok((memory_kb, cpu_percent))
229        } else {
230            Ok((0, 0.0))
231        }
232    }
233}
234
235/// Server health information
236#[derive(Debug, Clone)]
237pub struct ServerHealth {
238    pub is_running: bool,
239    pub pid: Option<u32>,
240    pub uptime_ms: u64,
241    pub memory_usage_kb: u64,
242    pub cpu_usage_percent: f64,
243}
244
245impl Drop for LspServerProcess {
246    fn drop(&mut self) {
247        if self.is_running() {
248            warn!("LSP server process dropped while still running, killing: {}", self.command);
249            // Note: We can't await in Drop, so this is best-effort
250            let _ = self.child.start_kill();
251        }
252    }
253}
254
255#[cfg(test)]
256mod tests {
257    use super::*;
258    use tokio::time::sleep;
259
260    #[tokio::test]
261    async fn test_command_exists() {
262        // Test with a command that should exist
263        assert!(LspServerProcess::command_exists("echo").await);
264        
265        // Test with a command that shouldn't exist
266        assert!(!LspServerProcess::command_exists("definitely_not_a_command_12345").await);
267    }
268
269    #[tokio::test]
270    async fn test_process_lifecycle() {
271        // Use 'cat' as a simple long-running process
272        let mut process = LspServerProcess::new("cat", &[], 5000).await.unwrap();
273        
274        assert!(process.is_running());
275        assert!(process.pid().is_some());
276        
277        // Test health check
278        let health = process.health_check().await;
279        assert!(health.is_running);
280        assert!(health.pid.is_some());
281        
282        // Test shutdown
283        process.shutdown().await.unwrap();
284        
285        // Give it a moment to fully exit
286        sleep(Duration::from_millis(100)).await;
287        
288        assert!(!process.is_running());
289    }
290}