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