1use std::collections::HashMap;
2use std::io;
3use std::io::ErrorKind;
4use std::path::{Path, PathBuf};
5use std::process::Stdio;
6use std::sync::atomic::{AtomicI64, Ordering};
7use std::sync::{Arc, Mutex, PoisonError};
8use std::time::Duration;
9
10use futures::{SinkExt, StreamExt};
11use lsp_types::request::{DocumentSymbolRequest, GotoDefinition, HoverRequest, References, Rename, Request};
12use lsp_types::{
13 DocumentSymbolParams, DocumentSymbolResponse, GotoDefinitionParams, GotoDefinitionResponse, Hover, HoverParams,
14 Location, PartialResultParams, Position, PublishDiagnosticsParams, ReferenceContext, ReferenceParams, RenameParams,
15 TextDocumentIdentifier, TextDocumentPositionParams, Uri, WorkDoneProgressParams, WorkspaceEdit,
16};
17use serde::Serialize;
18use serde::de::DeserializeOwned;
19use serde_json::Value;
20use thiserror::Error;
21use tokio::io::{ReadHalf, WriteHalf};
22use tokio::net::UnixStream;
23use tokio::process::Command;
24use tokio::sync::{Mutex as AsyncMutex, oneshot};
25
26use crate::language_catalog::LanguageId;
27use crate::protocol::{
28 DaemonCall, DaemonRequest, DaemonResponse, FrameReader, FrameWriter, InitializeRequest, frame_reader, frame_writer,
29};
30use crate::socket_path::{ensure_socket_dir, log_file_path};
31
32#[doc = include_str!("docs/client_error.md")]
33#[derive(Debug, Error)]
34pub enum ClientError {
35 #[error("Failed to connect to daemon: {0}")]
36 ConnectionFailed(#[source] io::Error),
37
38 #[error("IO error: {0}")]
39 Io(#[from] io::Error),
40
41 #[error("LSP error (code={code}): {message}")]
42 LspError { code: i32, message: String },
43
44 #[error("Failed to spawn daemon: {0}")]
45 SpawnFailed(#[source] io::Error),
46
47 #[error("Timeout waiting for daemon to start")]
48 SpawnTimeout,
49
50 #[error("Daemon binary not found: {0}")]
51 DaemonBinaryNotFound(String),
52
53 #[error("Protocol error: {0}")]
54 ProtocolError(String),
55
56 #[error("Initialization failed: {0}")]
57 InitializationFailed(String),
58}
59
60pub type ClientResult<T> = std::result::Result<T, ClientError>;
61
62#[doc = include_str!("docs/client.md")]
63pub struct LspClient {
64 writer: AsyncMutex<FrameWriter<WriteHalf<UnixStream>, DaemonRequest>>,
65 pending: PendingRequests,
66 next_id: AtomicI64,
67 reader_task: tokio::task::JoinHandle<()>,
68}
69
70impl LspClient {
71 pub async fn connect(workspace_root: &Path, language: LanguageId) -> ClientResult<Self> {
72 let socket_path = ensure_socket_dir(workspace_root, language).map_err(ClientError::Io)?;
73
74 match UnixStream::connect(&socket_path).await {
75 Ok(stream) => {
76 return Self::from_stream(stream, workspace_root, language).await;
77 }
78 Err(err) if err.kind() == ErrorKind::ConnectionRefused || err.kind() == ErrorKind::NotFound => {}
79 Err(err) => return Err(ClientError::ConnectionFailed(err)),
80 }
81
82 spawn_daemon(&socket_path).await?;
83 let stream = UnixStream::connect(&socket_path).await.map_err(ClientError::ConnectionFailed)?;
84 Self::from_stream(stream, workspace_root, language).await
85 }
86
87 pub async fn goto_definition(&self, uri: Uri, line: u32, character: u32) -> ClientResult<GotoDefinitionResponse> {
88 let params = GotoDefinitionParams {
89 text_document_position_params: position(uri, line, character),
90 work_done_progress_params: WorkDoneProgressParams::default(),
91 partial_result_params: PartialResultParams::default(),
92 };
93 self.call(GotoDefinition::METHOD, ¶ms, || GotoDefinitionResponse::Array(vec![])).await
94 }
95
96 pub async fn find_references(
97 &self,
98 uri: Uri,
99 line: u32,
100 character: u32,
101 include_declaration: bool,
102 ) -> ClientResult<Vec<Location>> {
103 let params = ReferenceParams {
104 text_document_position: position(uri, line, character),
105 work_done_progress_params: WorkDoneProgressParams::default(),
106 partial_result_params: PartialResultParams::default(),
107 context: ReferenceContext { include_declaration },
108 };
109 self.call(References::METHOD, ¶ms, Vec::new).await
110 }
111
112 pub async fn hover(&self, uri: Uri, line: u32, character: u32) -> ClientResult<Option<Hover>> {
113 let params = HoverParams {
114 text_document_position_params: position(uri, line, character),
115 work_done_progress_params: WorkDoneProgressParams::default(),
116 };
117 self.call(HoverRequest::METHOD, ¶ms, || None).await
118 }
119
120 pub async fn document_symbol(&self, uri: Uri) -> ClientResult<DocumentSymbolResponse> {
121 let params = DocumentSymbolParams {
122 text_document: TextDocumentIdentifier::new(uri),
123 work_done_progress_params: WorkDoneProgressParams::default(),
124 partial_result_params: PartialResultParams::default(),
125 };
126 self.call(DocumentSymbolRequest::METHOD, ¶ms, || DocumentSymbolResponse::Flat(vec![])).await
127 }
128
129 pub async fn rename(
130 &self,
131 uri: Uri,
132 line: u32,
133 character: u32,
134 new_name: String,
135 ) -> ClientResult<Option<WorkspaceEdit>> {
136 let params = RenameParams {
137 text_document_position: position(uri, line, character),
138 new_name,
139 work_done_progress_params: WorkDoneProgressParams::default(),
140 };
141 self.call(Rename::METHOD, ¶ms, || None).await
142 }
143
144 pub async fn get_diagnostics(&self, uri: Option<Uri>) -> ClientResult<Vec<PublishDiagnosticsParams>> {
145 self.request(DaemonCall::Diagnostics { uri })
146 .await
147 .and_then(|value| serde_json::from_value(value).map_err(|err| ClientError::ProtocolError(err.to_string())))
148 }
149
150 pub async fn queue_diagnostic_refresh(&self, uri: Uri) -> ClientResult<()> {
151 self.request(DaemonCall::QueueRefresh { uri }).await.map(|_| ())
152 }
153
154 pub fn is_connected(&self) -> bool {
156 !self.reader_task.is_finished()
157 }
158
159 pub async fn disconnect(self) -> ClientResult<()> {
160 let mut writer = self.writer.lock().await;
161 writer.send(DaemonRequest::Disconnect).await.map_err(ClientError::Io)
162 }
163
164 pub async fn call<P: Serialize, R: DeserializeOwned>(
165 &self,
166 method: &str,
167 params: &P,
168 default: impl FnOnce() -> R,
169 ) -> ClientResult<R> {
170 let params = serde_json::to_value(params).map_err(|err| ClientError::ProtocolError(err.to_string()))?;
171 let value = self.request(DaemonCall::Lsp { method: method.to_string(), params }).await?;
172
173 if value.is_null() {
174 Ok(default())
175 } else {
176 serde_json::from_value(value).map_err(|err| ClientError::ProtocolError(format!("Parse error: {err}")))
177 }
178 }
179}
180
181impl LspClient {
182 async fn from_stream(stream: UnixStream, workspace_root: &Path, language: LanguageId) -> ClientResult<Self> {
183 let (reader, writer) = tokio::io::split(stream);
184 let mut reader = frame_reader::<_, DaemonResponse>(reader);
185 let mut writer = frame_writer::<_, DaemonRequest>(writer);
186
187 let initialize =
188 DaemonRequest::Initialize(InitializeRequest { workspace_root: workspace_root.to_path_buf(), language });
189
190 writer.send(initialize).await.map_err(ClientError::Io)?;
191
192 let response = match reader.next().await {
193 Some(Ok(resp)) => resp,
194 Some(Err(err)) => return Err(ClientError::Io(err)),
195 None => {
196 return Err(ClientError::ProtocolError("Connection closed during initialization".into()));
197 }
198 };
199
200 match response {
201 DaemonResponse::Initialized => {}
202 DaemonResponse::InitializeFailed(message) => return Err(ClientError::InitializationFailed(message)),
203 DaemonResponse::Reply { .. } => {
204 return Err(ClientError::ProtocolError("Unexpected response to Initialize".into()));
205 }
206 }
207
208 let pending = Arc::new(Mutex::new(Some(HashMap::new())));
209 let reader_task = tokio::spawn(run_reader(reader, Arc::clone(&pending)));
210
211 Ok(Self { writer: AsyncMutex::new(writer), pending, next_id: AtomicI64::new(1), reader_task })
212 }
213
214 async fn request(&self, call: DaemonCall) -> ClientResult<Value> {
215 let id = self.next_id.fetch_add(1, Ordering::SeqCst);
216 let (response_tx, response_rx) = oneshot::channel();
217 {
218 let mut pending = self.pending.lock().unwrap_or_else(PoisonError::into_inner);
219 let pending = pending.as_mut().ok_or_else(|| ClientError::ProtocolError("Daemon disconnected".into()))?;
220 pending.insert(id, response_tx);
221 }
222
223 let write_result = {
224 let mut writer = self.writer.lock().await;
225 writer.send(DaemonRequest::Call { id, call }).await
226 };
227 if let Err(error) = write_result {
228 if let Some(pending) = self.pending.lock().unwrap_or_else(PoisonError::into_inner).as_mut() {
229 pending.remove(&id);
230 }
231 return Err(ClientError::Io(error));
232 }
233 response_rx.await.map_err(|_| ClientError::ProtocolError("Response channel closed".into()))?
234 }
235}
236
237impl Drop for LspClient {
238 fn drop(&mut self) {
239 self.reader_task.abort();
240 }
241}
242
243type PendingResult = Result<Value, ClientError>;
244type PendingRequests = Arc<Mutex<Option<HashMap<i64, oneshot::Sender<PendingResult>>>>>;
245
246async fn run_reader(mut reader: FrameReader<ReadHalf<UnixStream>, DaemonResponse>, pending: PendingRequests) {
247 while let Some(msg) = reader.next().await {
248 let response = match msg {
249 Ok(response) => response,
250 Err(error) => {
251 tracing::debug!(%error, "Error reading daemon response");
252 break;
253 }
254 };
255 let DaemonResponse::Reply { id, result } = response else { continue };
256 if let Some(response_tx) =
257 pending.lock().unwrap_or_else(PoisonError::into_inner).as_mut().and_then(|pending| pending.remove(&id))
258 {
259 let _ = response_tx
260 .send(result.map_err(|error| ClientError::LspError { code: error.code, message: error.message }));
261 }
262 }
263
264 if let Some(pending) = pending.lock().unwrap_or_else(PoisonError::into_inner).take() {
265 for (_, response_tx) in pending {
266 let _ = response_tx.send(Err(ClientError::ProtocolError("Daemon disconnected".into())));
267 }
268 }
269}
270
271async fn spawn_daemon(socket_path: &Path) -> ClientResult<()> {
272 let (binary, subcommand) = find_daemon_binary()?;
273 let log_file = log_file_path(socket_path);
274
275 let mut cmd = Command::new(&binary);
276 if let Some(sub) = subcommand {
277 cmd.arg(sub);
278 }
279 cmd.arg("--socket")
280 .arg(socket_path)
281 .arg("--log-file")
282 .arg(&log_file)
283 .arg("--log-level")
284 .arg("debug")
285 .stdin(Stdio::null())
286 .stdout(Stdio::null())
287 .stderr(Stdio::null());
288
289 #[cfg(unix)]
290 unsafe {
291 use std::os::unix::process::CommandExt;
292 cmd.as_std_mut()
293 .pre_exec(|| nix::unistd::setsid().map(|_| ()).map_err(|e| std::io::Error::from_raw_os_error(e as i32)));
294 }
295
296 let mut child = cmd.spawn().map_err(ClientError::SpawnFailed)?;
297
298 for _ in 0..50 {
299 match child.try_wait() {
300 Ok(Some(status)) if !status.success() => {
301 return Err(ClientError::SpawnFailed(io::Error::other(format!("Daemon exited with status: {status}"))));
302 }
303 Ok(_) => {}
304 Err(err) => return Err(ClientError::SpawnFailed(err)),
305 }
306
307 tokio::time::sleep(Duration::from_millis(100)).await;
308 if UnixStream::connect(socket_path).await.is_ok() {
309 tokio::spawn(async move {
310 match child.wait().await {
311 Ok(status) => tracing::debug!(%status, "aether-lspd launcher reaped"),
312 Err(err) => tracing::warn!(%err, "Failed to reap aether-lspd launcher"),
313 }
314 });
315 return Ok(());
316 }
317 }
318
319 let _ = child.kill().await;
320 let _ = child.wait().await;
321 Err(ClientError::SpawnTimeout)
322}
323
324fn find_daemon_binary() -> ClientResult<(PathBuf, Option<&'static str>)> {
325 let exe = std::env::current_exe().ok();
326 let exe_dir = exe.as_deref().and_then(|p| p.parent());
327
328 let standalone_candidates = [
329 exe_dir.map(|dir| dir.join("aether-lspd")),
330 exe_dir.and_then(|dir| dir.parent()).map(|dir| dir.join("aether-lspd")),
331 which_aether_lspd(),
332 Some(PathBuf::from("target/debug/aether-lspd")),
333 Some(PathBuf::from("target/release/aether-lspd")),
334 Some(PathBuf::from("../../target/debug/aether-lspd")),
335 Some(PathBuf::from("../../target/release/aether-lspd")),
336 ];
337
338 for candidate in standalone_candidates.into_iter().flatten() {
339 if candidate.exists() {
340 return Ok((candidate, None));
341 }
342 }
343
344 if let Some(exe) = exe {
345 return Ok((exe, Some("lspd")));
346 }
347
348 Err(ClientError::DaemonBinaryNotFound("aether-lspd not found".into()))
349}
350
351fn which_aether_lspd() -> Option<PathBuf> {
352 std::env::var_os("PATH")
353 .and_then(|paths| std::env::split_paths(&paths).map(|path| path.join("aether-lspd")).find(|path| path.exists()))
354}
355
356fn position(uri: Uri, line: u32, character: u32) -> TextDocumentPositionParams {
357 TextDocumentPositionParams::new(TextDocumentIdentifier::new(uri), Position::new(line, character))
358}