1use std::collections::{BTreeMap, BTreeSet};
2use std::path::Path;
3use std::sync::Arc;
4
5use lsp_types::Position;
6use tokio::sync::Mutex;
7
8use crate::client::LspClient;
9use crate::error::LspError;
10use crate::types::{
11 normalize_extension, FileDiagnostics, LspContextEnrichment, LspServerConfig, SymbolLocation,
12 WorkspaceDiagnostics,
13};
14
15pub struct LspManager {
16 server_configs: BTreeMap<String, LspServerConfig>,
17 extension_map: BTreeMap<String, String>,
18 clients: Mutex<BTreeMap<String, Arc<LspClient>>>,
19}
20
21impl LspManager {
22 pub fn new(server_configs: Vec<LspServerConfig>) -> Result<Self, LspError> {
23 let mut configs_by_name = BTreeMap::new();
24 let mut extension_map = BTreeMap::new();
25
26 for config in server_configs {
27 for extension in config.extension_to_language.keys() {
28 let normalized = normalize_extension(extension);
29 if let Some(existing_server) = extension_map.insert(normalized.clone(), config.name.clone()) {
30 return Err(LspError::DuplicateExtension {
31 extension: normalized,
32 existing_server,
33 new_server: config.name.clone(),
34 });
35 }
36 }
37 configs_by_name.insert(config.name.clone(), config);
38 }
39
40 Ok(Self {
41 server_configs: configs_by_name,
42 extension_map,
43 clients: Mutex::new(BTreeMap::new()),
44 })
45 }
46
47 #[must_use]
48 pub fn supports_path(&self, path: &Path) -> bool {
49 path.extension().is_some_and(|extension| {
50 let normalized = normalize_extension(extension.to_string_lossy().as_ref());
51 self.extension_map.contains_key(&normalized)
52 })
53 }
54
55 pub async fn open_document(&self, path: &Path, text: &str) -> Result<(), LspError> {
56 self.client_for_path(path).await?.open_document(path, text).await
57 }
58
59 pub async fn sync_document_from_disk(&self, path: &Path) -> Result<(), LspError> {
60 let contents = std::fs::read_to_string(path)?;
61 self.change_document(path, &contents).await?;
62 self.save_document(path).await
63 }
64
65 pub async fn change_document(&self, path: &Path, text: &str) -> Result<(), LspError> {
66 self.client_for_path(path).await?.change_document(path, text).await
67 }
68
69 pub async fn save_document(&self, path: &Path) -> Result<(), LspError> {
70 self.client_for_path(path).await?.save_document(path).await
71 }
72
73 pub async fn close_document(&self, path: &Path) -> Result<(), LspError> {
74 self.client_for_path(path).await?.close_document(path).await
75 }
76
77 pub async fn go_to_definition(
78 &self,
79 path: &Path,
80 position: Position,
81 ) -> Result<Vec<SymbolLocation>, LspError> {
82 let mut locations = self.client_for_path(path).await?.go_to_definition(path, position).await?;
83 dedupe_locations(&mut locations);
84 Ok(locations)
85 }
86
87 pub async fn find_references(
88 &self,
89 path: &Path,
90 position: Position,
91 include_declaration: bool,
92 ) -> Result<Vec<SymbolLocation>, LspError> {
93 let mut locations = self
94 .client_for_path(path)
95 .await?
96 .find_references(path, position, include_declaration)
97 .await?;
98 dedupe_locations(&mut locations);
99 Ok(locations)
100 }
101
102 pub async fn collect_workspace_diagnostics(&self) -> Result<WorkspaceDiagnostics, LspError> {
103 let clients = self.clients.lock().await.values().cloned().collect::<Vec<_>>();
104 let mut files = Vec::new();
105
106 for client in clients {
107 for (uri, diagnostics) in client.diagnostics_snapshot().await {
108 let Ok(path) = url::Url::parse(&uri)
109 .and_then(|url| url.to_file_path().map_err(|()| url::ParseError::RelativeUrlWithoutBase))
110 else {
111 continue;
112 };
113 if diagnostics.is_empty() {
114 continue;
115 }
116 files.push(FileDiagnostics {
117 path,
118 uri,
119 diagnostics,
120 });
121 }
122 }
123
124 files.sort_by(|left, right| left.path.cmp(&right.path));
125 Ok(WorkspaceDiagnostics { files })
126 }
127
128 pub async fn context_enrichment(
129 &self,
130 path: &Path,
131 position: Position,
132 ) -> Result<LspContextEnrichment, LspError> {
133 Ok(LspContextEnrichment {
134 file_path: path.to_path_buf(),
135 diagnostics: self.collect_workspace_diagnostics().await?,
136 definitions: self.go_to_definition(path, position).await?,
137 references: self.find_references(path, position, true).await?,
138 })
139 }
140
141 pub async fn shutdown(&self) -> Result<(), LspError> {
142 let mut clients = self.clients.lock().await;
143 let drained = clients.values().cloned().collect::<Vec<_>>();
144 clients.clear();
145 drop(clients);
146
147 for client in drained {
148 client.shutdown().await?;
149 }
150 Ok(())
151 }
152
153 async fn client_for_path(&self, path: &Path) -> Result<Arc<LspClient>, LspError> {
154 let extension = path
155 .extension()
156 .map(|extension| normalize_extension(extension.to_string_lossy().as_ref()))
157 .ok_or_else(|| LspError::UnsupportedDocument(path.to_path_buf()))?;
158 let server_name = self
159 .extension_map
160 .get(&extension)
161 .cloned()
162 .ok_or_else(|| LspError::UnsupportedDocument(path.to_path_buf()))?;
163
164 let mut clients = self.clients.lock().await;
165 if let Some(client) = clients.get(&server_name) {
166 return Ok(client.clone());
167 }
168
169 let config = self
170 .server_configs
171 .get(&server_name)
172 .cloned()
173 .ok_or_else(|| LspError::UnknownServer(server_name.clone()))?;
174 let client = Arc::new(LspClient::connect(config).await?);
175 clients.insert(server_name, client.clone());
176 Ok(client)
177 }
178}
179
180fn dedupe_locations(locations: &mut Vec<SymbolLocation>) {
181 let mut seen = BTreeSet::new();
182 locations.retain(|location| {
183 seen.insert((
184 location.path.clone(),
185 location.range.start.line,
186 location.range.start.character,
187 location.range.end.line,
188 location.range.end.character,
189 ))
190 });
191}