codelens_engine/lsp/
session.rs

1use crate::project::ProjectRoot;
2use anyhow::{Context, Result, bail};
3use serde_json::{Value, json};
4use std::collections::HashMap;
5use std::io::{BufRead, BufReader};
6use std::path::Path;
7use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio};
8use std::thread;
9use std::time::{Duration, Instant};
10use url::Url;
11
12use super::commands::is_allowed_lsp_command;
13use super::protocol::{language_id_for_path, poll_readable, read_message, send_message};
14use super::registry::resolve_lsp_binary;
15use super::types::{
16    LspCodeActionRefactorPlan, LspCodeActionRefactorResult, LspCodeActionRequest, LspDiagnostic,
17    LspDiagnosticRequest, LspReference, LspRenamePlan, LspRenamePlanRequest, LspRenameRequest,
18    LspRequest, LspResolveTargetRequest, LspResolvedTarget, LspTypeHierarchyRequest,
19    LspWorkspaceEditTransaction, LspWorkspaceSymbol, LspWorkspaceSymbolRequest,
20};
21
22#[derive(Debug, Clone, PartialEq, Eq, Hash)]
23struct SessionKey {
24    command: String,
25    args: Vec<String>,
26}
27
28#[derive(Debug, Clone)]
29struct OpenDocumentState {
30    version: i32,
31    text: String,
32}
33
34pub struct LspSessionPool {
35    project: ProjectRoot,
36    sessions: std::sync::Mutex<HashMap<SessionKey, LspSession>>,
37}
38
39pub(super) struct LspSession {
40    pub(super) project: ProjectRoot,
41    child: Child,
42    stdin: ChildStdin,
43    reader: BufReader<ChildStdout>,
44    next_request_id: u64,
45    documents: HashMap<String, OpenDocumentState>,
46    #[allow(dead_code)] // retained for future stderr diagnostics
47    stderr_buffer: std::sync::Arc<std::sync::Mutex<String>>,
48}
49
50fn ensure_session<'a>(
51    sessions: &'a mut HashMap<SessionKey, LspSession>,
52    project: &ProjectRoot,
53    command: &str,
54    args: &[String],
55) -> Result<&'a mut LspSession> {
56    if !is_allowed_lsp_command(command) {
57        bail!(
58            "Blocked: '{command}' is not a known LSP server. Only whitelisted LSP binaries are allowed."
59        );
60    }
61
62    let key = SessionKey {
63        command: command.to_owned(),
64        args: args.to_owned(),
65    };
66
67    // Check for dead sessions: if the child process has exited, remove the stale entry.
68    if let Some(session) = sessions.get_mut(&key) {
69        match session.child.try_wait() {
70            Ok(Some(_status)) => {
71                // Process exited — remove stale session so we start fresh below.
72                sessions.remove(&key);
73            }
74            Ok(None) => {} // Still running — will return it via Occupied below.
75            Err(_) => {
76                sessions.remove(&key);
77            }
78        }
79    }
80
81    match sessions.entry(key) {
82        std::collections::hash_map::Entry::Occupied(e) => Ok(e.into_mut()),
83        std::collections::hash_map::Entry::Vacant(e) => {
84            let session = LspSession::start(project, command, args)?;
85            Ok(e.insert(session))
86        }
87    }
88}
89
90impl LspSessionPool {
91    pub fn new(project: ProjectRoot) -> Self {
92        Self {
93            project,
94            sessions: std::sync::Mutex::new(HashMap::new()),
95        }
96    }
97
98    /// Replace the project root and close all existing sessions.
99    pub fn reset(&self, project: ProjectRoot) -> Self {
100        // Drop existing sessions so LSP processes are killed.
101        self.sessions
102            .lock()
103            .unwrap_or_else(|p| p.into_inner())
104            .clear();
105        Self::new(project)
106    }
107
108    pub fn session_count(&self) -> usize {
109        self.sessions
110            .lock()
111            .unwrap_or_else(|p| p.into_inner())
112            .len()
113    }
114
115    pub fn find_referencing_symbols(&self, request: &LspRequest) -> Result<Vec<LspReference>> {
116        let mut sessions = self.sessions.lock().unwrap_or_else(|p| p.into_inner());
117        let session = ensure_session(
118            &mut sessions,
119            &self.project,
120            &request.command,
121            &request.args,
122        )?;
123        session.find_references(request)
124    }
125
126    pub fn get_diagnostics(&self, request: &LspDiagnosticRequest) -> Result<Vec<LspDiagnostic>> {
127        let mut sessions = self.sessions.lock().unwrap_or_else(|p| p.into_inner());
128        let session = ensure_session(
129            &mut sessions,
130            &self.project,
131            &request.command,
132            &request.args,
133        )?;
134        session.get_diagnostics(request)
135    }
136
137    pub fn search_workspace_symbols(
138        &self,
139        request: &LspWorkspaceSymbolRequest,
140    ) -> Result<Vec<LspWorkspaceSymbol>> {
141        let mut sessions = self.sessions.lock().unwrap_or_else(|p| p.into_inner());
142        let session = ensure_session(
143            &mut sessions,
144            &self.project,
145            &request.command,
146            &request.args,
147        )?;
148        session.search_workspace_symbols(request)
149    }
150
151    pub fn get_type_hierarchy(
152        &self,
153        request: &LspTypeHierarchyRequest,
154    ) -> Result<HashMap<String, Value>> {
155        let mut sessions = self.sessions.lock().unwrap_or_else(|p| p.into_inner());
156        let session = ensure_session(
157            &mut sessions,
158            &self.project,
159            &request.command,
160            &request.args,
161        )?;
162        session.get_type_hierarchy(request)
163    }
164
165    pub fn resolve_symbol_target(
166        &self,
167        request: &LspResolveTargetRequest,
168    ) -> Result<Vec<LspResolvedTarget>> {
169        let mut sessions = self.sessions.lock().unwrap_or_else(|p| p.into_inner());
170        let session = ensure_session(
171            &mut sessions,
172            &self.project,
173            &request.command,
174            &request.args,
175        )?;
176        session.resolve_symbol_target(request)
177    }
178
179    pub fn get_rename_plan(&self, request: &LspRenamePlanRequest) -> Result<LspRenamePlan> {
180        let mut sessions = self.sessions.lock().unwrap_or_else(|p| p.into_inner());
181        let session = ensure_session(
182            &mut sessions,
183            &self.project,
184            &request.command,
185            &request.args,
186        )?;
187        session.get_rename_plan(request)
188    }
189
190    pub fn rename_symbol(&self, request: &LspRenameRequest) -> Result<crate::rename::RenameResult> {
191        let mut sessions = self.sessions.lock().unwrap_or_else(|p| p.into_inner());
192        let session = ensure_session(
193            &mut sessions,
194            &self.project,
195            &request.command,
196            &request.args,
197        )?;
198        session.rename_symbol(request)
199    }
200
201    pub fn rename_symbol_transaction(
202        &self,
203        request: &LspRenameRequest,
204    ) -> Result<LspWorkspaceEditTransaction> {
205        let mut sessions = self.sessions.lock().unwrap_or_else(|p| p.into_inner());
206        let session = ensure_session(
207            &mut sessions,
208            &self.project,
209            &request.command,
210            &request.args,
211        )?;
212        session.rename_symbol_transaction(request)
213    }
214
215    pub fn code_action_refactor(
216        &self,
217        request: &LspCodeActionRequest,
218    ) -> Result<LspCodeActionRefactorResult> {
219        let mut sessions = self.sessions.lock().unwrap_or_else(|p| p.into_inner());
220        let session = ensure_session(
221            &mut sessions,
222            &self.project,
223            &request.command,
224            &request.args,
225        )?;
226        session.code_action_refactor(request)
227    }
228
229    pub fn code_action_refactor_plan(
230        &self,
231        request: &LspCodeActionRequest,
232    ) -> Result<LspCodeActionRefactorPlan> {
233        let mut sessions = self.sessions.lock().unwrap_or_else(|p| p.into_inner());
234        let session = ensure_session(
235            &mut sessions,
236            &self.project,
237            &request.command,
238            &request.args,
239        )?;
240        session.code_action_refactor_plan(request)
241    }
242}
243
244impl LspSession {
245    fn start(project: &ProjectRoot, command: &str, args: &[String]) -> Result<Self> {
246        let command_path = resolve_lsp_binary(command).unwrap_or_else(|| command.into());
247        let mut child = Command::new(&command_path)
248            .args(args)
249            .stdin(Stdio::piped())
250            .stdout(Stdio::piped())
251            .stderr(Stdio::piped())
252            .spawn()
253            .with_context(|| format!("failed to spawn LSP server {}", command_path.display()))?;
254
255        let stdin = child.stdin.take().context("failed to open LSP stdin")?;
256        let stdout = child.stdout.take().context("failed to open LSP stdout")?;
257
258        // Capture stderr in a background thread (bounded 4KB ring buffer).
259        let stderr_buffer = std::sync::Arc::new(std::sync::Mutex::new(String::new()));
260        if let Some(stderr) = child.stderr.take() {
261            let buf = std::sync::Arc::clone(&stderr_buffer);
262            thread::spawn(move || {
263                let mut reader = BufReader::new(stderr);
264                let mut line = String::new();
265                while reader.read_line(&mut line).unwrap_or(0) > 0 {
266                    if let Ok(mut b) = buf.lock() {
267                        if b.len() > 4096 {
268                            let drain_to = b.len() - 2048;
269                            b.drain(..drain_to);
270                        }
271                        b.push_str(&line);
272                    }
273                    line.clear();
274                }
275            });
276        }
277
278        let mut session = Self {
279            project: project.clone(),
280            child,
281            stdin,
282            reader: BufReader::new(stdout),
283            next_request_id: 1,
284            documents: HashMap::new(),
285            stderr_buffer,
286        };
287        session.initialize()?;
288        if let Some(grace) = configured_startup_grace() {
289            thread::sleep(grace);
290        }
291        Ok(session)
292    }
293
294    fn initialize(&mut self) -> Result<()> {
295        let id = self.next_id();
296        let root_uri = Url::from_directory_path(self.project.as_path())
297            .ok()
298            .map(|url| url.to_string());
299        self.send_request(
300            id,
301            "initialize",
302            json!({
303                "processId":null,
304                "rootUri": root_uri,
305                "capabilities":{
306                    "workspace":{
307                        "workspaceEdit":{
308                            "documentChanges":true,
309                            "resourceOperations":["create","rename","delete"],
310                            "failureHandling":"textOnlyTransactional"
311                        },
312                        "symbol":{"dynamicRegistration":false}
313                    },
314                    "textDocument":{
315                        "declaration":{"dynamicRegistration":false},
316                        "definition":{"dynamicRegistration":false},
317                        "implementation":{"dynamicRegistration":false},
318                        "typeDefinition":{"dynamicRegistration":false},
319                        "references":{"dynamicRegistration":false},
320                        "rename":{"dynamicRegistration":false,"prepareSupport":true},
321                        "diagnostic":{"dynamicRegistration":false},
322                        "typeHierarchy":{"dynamicRegistration":false},
323                        "codeAction":{
324                            "dynamicRegistration":false,
325                            "codeActionLiteralSupport":{
326                                "codeActionKind":{
327                                    "valueSet":["quickfix","refactor","refactor.extract","refactor.inline","refactor.rewrite"]
328                                }
329                            },
330                            "resolveSupport":{"properties":["edit","command"]}
331                        }
332                    }
333                },
334                "workspaceFolders":[
335                    {
336                        "uri": Url::from_directory_path(self.project.as_path()).ok().map(|url| url.to_string()),
337                        "name": self.project.as_path().file_name().and_then(|n| n.to_str()).unwrap_or("workspace")
338                    }
339                ]
340            }),
341        )?;
342        let _ = self.read_response_for_id(id)?;
343        self.send_notification("initialized", json!({}))?;
344        Ok(())
345    }
346
347    pub(super) fn prepare_document(&mut self, absolute_path: &Path) -> Result<(String, String)> {
348        let uri = Url::from_file_path(absolute_path).map_err(|_| {
349            anyhow::anyhow!("failed to build file uri for {}", absolute_path.display())
350        })?;
351        let uri_string = uri.to_string();
352        let source = std::fs::read_to_string(absolute_path)
353            .with_context(|| format!("failed to read {}", absolute_path.display()))?;
354        let language_id = language_id_for_path(absolute_path)?;
355        self.sync_document(&uri_string, language_id, &source)?;
356        Ok((uri_string, source))
357    }
358
359    fn sync_document(&mut self, uri: &str, language_id: &str, source: &str) -> Result<()> {
360        if let Some(state) = self.documents.get(uri)
361            && state.text == source
362        {
363            return Ok(());
364        }
365
366        if let Some(state) = self.documents.get_mut(uri) {
367            state.version += 1;
368            state.text = source.to_owned();
369            let version = state.version;
370            return self.send_notification(
371                "textDocument/didChange",
372                json!({
373                    "textDocument":{"uri":uri,"version":version},
374                    "contentChanges":[{"text":source}]
375                }),
376            );
377        }
378
379        self.documents.insert(
380            uri.to_owned(),
381            OpenDocumentState {
382                version: 1,
383                text: source.to_owned(),
384            },
385        );
386        self.send_notification(
387            "textDocument/didOpen",
388            json!({
389                "textDocument":{
390                    "uri":uri,
391                    "languageId":language_id,
392                    "version":1,
393                    "text":source
394                }
395            }),
396        )
397    }
398
399    pub(super) fn next_id(&mut self) -> u64 {
400        let id = self.next_request_id;
401        self.next_request_id += 1;
402        id
403    }
404
405    pub(super) fn send_request(&mut self, id: u64, method: &str, params: Value) -> Result<()> {
406        send_message(
407            &mut self.stdin,
408            &json!({
409                "jsonrpc":"2.0",
410                "id":id,
411                "method":method,
412                "params":params
413            }),
414        )
415    }
416
417    fn send_notification(&mut self, method: &str, params: Value) -> Result<()> {
418        send_message(
419            &mut self.stdin,
420            &json!({
421                "jsonrpc":"2.0",
422                "method":method,
423                "params":params
424            }),
425        )
426    }
427
428    pub(super) fn read_response_for_id(&mut self, expected_id: u64) -> Result<Value> {
429        let deadline = Instant::now() + Duration::from_secs(30);
430        let mut discarded = 0u32;
431        const MAX_DISCARDED: u32 = 500;
432
433        loop {
434            let remaining = deadline.saturating_duration_since(Instant::now());
435            if remaining.is_zero() {
436                bail!(
437                    "LSP response timeout: no response for request id {expected_id} within 30s \
438                     ({discarded} unrelated messages discarded)"
439                );
440            }
441            if discarded >= MAX_DISCARDED {
442                bail!(
443                    "LSP response loop: discarded {MAX_DISCARDED} messages without finding id {expected_id}"
444                );
445            }
446
447            // Poll the pipe before blocking read — prevents infinite hang
448            if !poll_readable(self.reader.get_ref(), remaining.min(Duration::from_secs(5))) {
449                continue; // no data yet, re-check deadline
450            }
451
452            let message = read_message(&mut self.reader)?;
453            let matches_id = message
454                .get("id")
455                .and_then(Value::as_u64)
456                .map(|id| id == expected_id)
457                .unwrap_or(false);
458            if matches_id {
459                if let Some(error) = message.get("error") {
460                    let code = error.get("code").and_then(Value::as_i64).unwrap_or(-1);
461                    let error_message = error
462                        .get("message")
463                        .and_then(Value::as_str)
464                        .unwrap_or("unknown LSP error");
465                    bail!("LSP request failed ({code}): {error_message}");
466                }
467                return Ok(message);
468            }
469            discarded += 1;
470        }
471    }
472
473    fn shutdown(&mut self) -> Result<()> {
474        let id = self.next_id();
475        self.send_request(id, "shutdown", Value::Null)?;
476        let _ = self.read_response_for_id(id)?;
477        self.send_notification("exit", Value::Null)
478    }
479}
480
481fn configured_startup_grace() -> Option<Duration> {
482    let millis = std::env::var("CODELENS_LSP_STARTUP_GRACE_MS")
483        .ok()
484        .and_then(|value| value.trim().parse::<u64>().ok())
485        .unwrap_or(0)
486        .min(10_000);
487    (millis > 0).then(|| Duration::from_millis(millis))
488}
489
490impl Drop for LspSession {
491    fn drop(&mut self) {
492        let _ = self.shutdown();
493        let deadline = Instant::now() + Duration::from_millis(250);
494        while Instant::now() < deadline {
495            match self.child.try_wait() {
496                Ok(Some(_status)) => return,
497                Ok(None) => thread::sleep(Duration::from_millis(10)),
498                Err(_) => break,
499            }
500        }
501        let _ = self.child.kill();
502        let _ = self.child.wait();
503    }
504}
codelens_engine/lsp/session.rs

codelens_engine/lsp/
session.rs
