Skip to main content

codineer_lsp/
manager.rs

1use std::collections::{BTreeMap, BTreeSet};
2use std::path::Path;
3use std::sync::Arc;
4use std::time::Duration;
5
6use lsp_types::{Position, TextEdit};
7use serde_json::Value;
8use tokio::sync::Mutex;
9
10use crate::client::LspClient;
11use crate::error::LspError;
12use crate::types::{
13    normalize_extension, CompletionItem, DocumentSymbolInfo, FileDiagnostics, HoverResult,
14    LspContextEnrichment, LspServerConfig, SymbolLocation, WorkspaceDiagnostics,
15};
16
17pub struct LspManager {
18    server_configs: BTreeMap<String, LspServerConfig>,
19    extension_map: BTreeMap<String, String>,
20    clients: Mutex<BTreeMap<String, Arc<LspClient>>>,
21}
22
23impl LspManager {
24    pub fn new(server_configs: Vec<LspServerConfig>) -> Result<Self, LspError> {
25        let mut configs_by_name = BTreeMap::new();
26        let mut extension_map = BTreeMap::new();
27
28        for config in server_configs {
29            for extension in config.extension_to_language.keys() {
30                let normalized = normalize_extension(extension);
31                if let Some(existing_server) =
32                    extension_map.insert(normalized.clone(), config.name.clone())
33                {
34                    return Err(LspError::DuplicateExtension {
35                        extension: normalized,
36                        existing_server,
37                        new_server: config.name.clone(),
38                    });
39                }
40            }
41            configs_by_name.insert(config.name.clone(), config);
42        }
43
44        Ok(Self {
45            server_configs: configs_by_name,
46            extension_map,
47            clients: Mutex::new(BTreeMap::new()),
48        })
49    }
50
51    /// Build an `LspManager` from a JSON array of server config objects.
52    /// Each element must be deserializable as [`LspServerConfig`].
53    pub fn from_json_config(configs: &Value) -> Result<Self, LspError> {
54        let configs: Vec<LspServerConfig> =
55            serde_json::from_value(configs.clone()).map_err(|e| LspError::Protocol {
56                message: format!("invalid LSP config JSON: {e}"),
57            })?;
58        Self::new(configs)
59    }
60
61    #[must_use]
62    pub fn supports_path(&self, path: &Path) -> bool {
63        path.extension().is_some_and(|extension| {
64            let normalized = normalize_extension(extension.to_string_lossy().as_ref());
65            self.extension_map.contains_key(&normalized)
66        })
67    }
68
69    pub async fn open_document(&self, path: &Path, text: &str) -> Result<(), LspError> {
70        self.client_for_path(path)
71            .await?
72            .open_document(path, text)
73            .await
74    }
75
76    /// Sync a document from disk and notify the server.
77    pub async fn sync_document_from_disk(&self, path: &Path) -> Result<(), LspError> {
78        let contents = tokio::fs::read_to_string(path).await?;
79        self.change_document(path, &contents).await?;
80        self.save_document(path).await
81    }
82
83    /// Sync a document from disk and wait for fresh diagnostics (up to `timeout`).
84    pub async fn sync_and_await_diagnostics(
85        &self,
86        path: &Path,
87        timeout: Duration,
88    ) -> Result<WorkspaceDiagnostics, LspError> {
89        let client = self.client_for_path(path).await?;
90        let min_version = client.diagnostics_version() + 1;
91        let contents = tokio::fs::read_to_string(path).await?;
92        client.change_document(path, &contents).await?;
93        client.save_document(path).await?;
94        client
95            .wait_for_diagnostics_update(min_version, timeout)
96            .await;
97        self.collect_workspace_diagnostics().await
98    }
99
100    pub async fn change_document(&self, path: &Path, text: &str) -> Result<(), LspError> {
101        self.client_for_path(path)
102            .await?
103            .change_document(path, text)
104            .await
105    }
106
107    pub async fn save_document(&self, path: &Path) -> Result<(), LspError> {
108        self.client_for_path(path).await?.save_document(path).await
109    }
110
111    pub async fn close_document(&self, path: &Path) -> Result<(), LspError> {
112        self.client_for_path(path).await?.close_document(path).await
113    }
114
115    pub async fn go_to_definition(
116        &self,
117        path: &Path,
118        position: Position,
119    ) -> Result<Vec<SymbolLocation>, LspError> {
120        let mut locations = self
121            .client_for_path(path)
122            .await?
123            .go_to_definition(path, position)
124            .await?;
125        dedupe_locations(&mut locations);
126        Ok(locations)
127    }
128
129    pub async fn find_references(
130        &self,
131        path: &Path,
132        position: Position,
133        include_declaration: bool,
134    ) -> Result<Vec<SymbolLocation>, LspError> {
135        let mut locations = self
136            .client_for_path(path)
137            .await?
138            .find_references(path, position, include_declaration)
139            .await?;
140        dedupe_locations(&mut locations);
141        Ok(locations)
142    }
143
144    /// Request hover information at the given position.
145    pub async fn hover(
146        &self,
147        path: &Path,
148        position: Position,
149    ) -> Result<Option<HoverResult>, LspError> {
150        self.client_for_path(path)
151            .await?
152            .hover(path, position)
153            .await
154    }
155
156    /// Request code completion at the given position.
157    pub async fn completion(
158        &self,
159        path: &Path,
160        position: Position,
161    ) -> Result<Vec<CompletionItem>, LspError> {
162        self.client_for_path(path)
163            .await?
164            .completion(path, position)
165            .await
166    }
167
168    /// Request document symbol outline.
169    pub async fn document_symbols(&self, path: &Path) -> Result<Vec<DocumentSymbolInfo>, LspError> {
170        self.client_for_path(path)
171            .await?
172            .document_symbols(path)
173            .await
174    }
175
176    /// Search workspace symbols matching `query`.
177    /// Aggregates results from every connected LSP server.
178    pub async fn workspace_symbols(&self, query: &str) -> Result<Vec<SymbolLocation>, LspError> {
179        let clients = self
180            .clients
181            .lock()
182            .await
183            .values()
184            .cloned()
185            .collect::<Vec<_>>();
186
187        let mut all = Vec::new();
188        for client in clients {
189            let mut locs = client.workspace_symbols(query).await?;
190            all.append(&mut locs);
191        }
192        dedupe_locations(&mut all);
193        Ok(all)
194    }
195
196    /// Rename a symbol at the given position.
197    pub async fn rename(
198        &self,
199        path: &Path,
200        position: Position,
201        new_name: &str,
202    ) -> Result<BTreeMap<std::path::PathBuf, Vec<TextEdit>>, LspError> {
203        self.client_for_path(path)
204            .await?
205            .rename(path, position, new_name)
206            .await
207    }
208
209    /// Format a document using the server's formatter.
210    pub async fn formatting(
211        &self,
212        path: &Path,
213        tab_size: u32,
214        insert_spaces: bool,
215    ) -> Result<Vec<TextEdit>, LspError> {
216        self.client_for_path(path)
217            .await?
218            .formatting(path, tab_size, insert_spaces)
219            .await
220    }
221
222    /// Returns the server capabilities for the LSP server that handles `path`.
223    /// Useful for clients that want to check what the server supports.
224    pub async fn server_capabilities(
225        &self,
226        path: &Path,
227    ) -> Result<lsp_types::ServerCapabilities, LspError> {
228        Ok(self
229            .client_for_path(path)
230            .await?
231            .server_capabilities()
232            .await)
233    }
234
235    pub async fn collect_workspace_diagnostics(&self) -> Result<WorkspaceDiagnostics, LspError> {
236        let clients = self
237            .clients
238            .lock()
239            .await
240            .values()
241            .cloned()
242            .collect::<Vec<_>>();
243        let mut files = Vec::new();
244
245        for client in clients {
246            for (uri, diagnostics) in client.diagnostics_snapshot().await {
247                let Ok(path) = url::Url::parse(&uri).and_then(|url| {
248                    url.to_file_path()
249                        .map_err(|()| url::ParseError::RelativeUrlWithoutBase)
250                }) else {
251                    continue;
252                };
253                if diagnostics.is_empty() {
254                    continue;
255                }
256                files.push(FileDiagnostics {
257                    path,
258                    uri,
259                    diagnostics,
260                });
261            }
262        }
263
264        files.sort_by(|left, right| left.path.cmp(&right.path));
265        Ok(WorkspaceDiagnostics { files })
266    }
267
268    pub async fn context_enrichment(
269        &self,
270        path: &Path,
271        position: Position,
272    ) -> Result<LspContextEnrichment, LspError> {
273        Ok(LspContextEnrichment {
274            file_path: path.to_path_buf(),
275            diagnostics: self.collect_workspace_diagnostics().await?,
276            definitions: self.go_to_definition(path, position).await?,
277            references: self.find_references(path, position, true).await?,
278        })
279    }
280
281    pub async fn shutdown(&self) -> Result<(), LspError> {
282        let mut clients = self.clients.lock().await;
283        let drained = clients.values().cloned().collect::<Vec<_>>();
284        clients.clear();
285        drop(clients);
286
287        for client in drained {
288            client.shutdown().await?;
289        }
290        Ok(())
291    }
292
293    async fn client_for_path(&self, path: &Path) -> Result<Arc<LspClient>, LspError> {
294        let extension = path
295            .extension()
296            .map(|ext| normalize_extension(ext.to_string_lossy().as_ref()))
297            .ok_or_else(|| LspError::UnsupportedDocument {
298                path: path.to_path_buf(),
299            })?;
300        let server_name =
301            self.extension_map
302                .get(&extension)
303                .ok_or_else(|| LspError::UnsupportedDocument {
304                    path: path.to_path_buf(),
305                })?;
306        let mut clients = self.clients.lock().await;
307        if let Some(client) = clients.get(server_name) {
308            return Ok(Arc::clone(client));
309        }
310        let config = self
311            .server_configs
312            .get(server_name)
313            .ok_or_else(|| LspError::UnknownServer {
314                name: server_name.clone(),
315            })?
316            .clone();
317        let client = Arc::new(LspClient::connect(config).await?);
318        clients.insert(server_name.clone(), Arc::clone(&client));
319        Ok(client)
320    }
321}
322
323fn dedupe_locations(locations: &mut Vec<SymbolLocation>) {
324    let mut seen = BTreeSet::new();
325    locations.retain(|loc| {
326        let key = (
327            loc.path.clone(),
328            loc.range.start.line,
329            loc.range.start.character,
330        );
331        seen.insert(key)
332    });
333}