Skip to main content

aether_lspd/
client.rs

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, &params, || 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, &params, 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, &params, || 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, &params, || 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, &params, || 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    /// Whether the daemon connection's reader task is still running.
155    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}