Skip to main content

rust_analyzer_cli/lsp/
client.rs

1use crate::lsp::types::*;
2use anyhow::{Context, Result, anyhow};
3use lsp_types::*;
4use serde_json::{Value, json};
5use std::collections::HashMap;
6use std::path::{Path, PathBuf};
7use std::process::Stdio;
8use std::sync::Arc;
9use std::sync::atomic::{AtomicI64, Ordering};
10use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
11use tokio::sync::{Mutex, mpsc, oneshot};
12use url::Url;
13
14pub struct LspClient {
15    tx_request: mpsc::Sender<Value>,
16    pending_requests: Arc<Mutex<HashMap<i64, oneshot::Sender<Result<Value>>>>>,
17    workspace_root: PathBuf,
18    pub process_id: u32,
19}
20
21impl LspClient {
22    pub async fn start(workspace_root: &Path) -> Result<Self> {
23        let canonical_root = workspace_root
24            .canonicalize()
25            .unwrap_or_else(|_| workspace_root.to_path_buf());
26
27        // Spawn rust-analyzer process
28        let mut child = tokio::process::Command::new("rust-analyzer")
29            .stdin(Stdio::piped())
30            .stdout(Stdio::piped())
31            .stderr(Stdio::null())
32            .spawn()
33            .context("Failed to spawn rust-analyzer. Please make sure `rust-analyzer` is installed and in PATH.")?;
34
35        let pid = child.id().unwrap_or(0);
36        let stdin = child
37            .stdin
38            .take()
39            .ok_or_else(|| anyhow!("Failed to open rust-analyzer stdin"))?;
40        let stdout = child
41            .stdout
42            .take()
43            .ok_or_else(|| anyhow!("Failed to open rust-analyzer stdout"))?;
44
45        let (tx_request, mut rx_request) = mpsc::channel::<Value>(100);
46        let pending_requests: Arc<Mutex<HashMap<i64, oneshot::Sender<Result<Value>>>>> =
47            Arc::new(Mutex::new(HashMap::new()));
48
49        // Stdio Writer Task
50        let mut writer_stdin = stdin;
51        tokio::spawn(async move {
52            while let Some(req) = rx_request.recv().await {
53                let body = serde_json::to_string(&req).unwrap_or_default();
54                let header = format!("Content-Length: {}\r\n\r\n", body.len());
55                let _ = writer_stdin.write_all(header.as_bytes()).await;
56                let _ = writer_stdin.write_all(body.as_bytes()).await;
57                let _ = writer_stdin.flush().await;
58            }
59        });
60
61        // Stdio Reader Task
62        let pending_map = pending_requests.clone();
63        tokio::spawn(async move {
64            let mut reader = BufReader::new(stdout);
65            loop {
66                let mut line = String::new();
67                match reader.read_line(&mut line).await {
68                    Ok(0) => break, // EOF
69                    Ok(_) => {
70                        let line_trimmed = line.trim();
71                        if line_trimmed.starts_with("Content-Length:") {
72                            let len_str = line_trimmed.trim_start_matches("Content-Length:").trim();
73                            if let Ok(len) = len_str.parse::<usize>() {
74                                // Read blank line
75                                let mut blank = String::new();
76                                let _ = reader.read_line(&mut blank).await;
77
78                                // Read payload
79                                let mut buf = vec![0u8; len];
80                                if reader.read_exact(&mut buf).await.is_ok()
81                                    && let Ok(v) = serde_json::from_slice::<Value>(&buf)
82                                    && let Some(id) = v.get("id").and_then(|i| i.as_i64())
83                                {
84                                    let mut map = pending_map.lock().await;
85                                    if let Some(sender) = map.remove(&id) {
86                                        if let Some(err) = v.get("error") {
87                                            let _ = sender.send(Err(anyhow!("LSP error: {}", err)));
88                                        } else {
89                                            let result =
90                                                v.get("result").cloned().unwrap_or(Value::Null);
91                                            let _ = sender.send(Ok(result));
92                                        }
93                                    }
94                                }
95                            }
96                        }
97                    }
98                    Err(_) => break,
99                }
100            }
101        });
102
103        let client = Self {
104            tx_request,
105            pending_requests,
106            workspace_root: canonical_root,
107            process_id: pid,
108        };
109
110        // Initialize handshake
111        client.initialize().await?;
112
113        Ok(client)
114    }
115
116    async fn send_request(&self, method: &str, params: Value) -> Result<Value> {
117        static REQ_ID: AtomicI64 = AtomicI64::new(1);
118        let id = REQ_ID.fetch_add(1, Ordering::SeqCst);
119        let req = json!({
120            "jsonrpc": "2.0",
121            "id": id,
122            "method": method,
123            "params": params
124        });
125
126        let (tx, rx) = oneshot::channel();
127        {
128            let mut map = self.pending_requests.lock().await;
129            map.insert(id, tx);
130        }
131
132        self.tx_request
133            .send(req)
134            .await
135            .map_err(|_| anyhow!("LSP writer task channel closed"))?;
136
137        rx.await
138            .map_err(|_| anyhow!("LSP response channel dropped"))?
139    }
140
141    async fn send_notification(&self, method: &str, params: Value) -> Result<()> {
142        let req = json!({
143            "jsonrpc": "2.0",
144            "method": method,
145            "params": params
146        });
147        let _ = self.tx_request.send(req).await;
148        Ok(())
149    }
150
151    async fn initialize(&self) -> Result<()> {
152        let root_uri = Url::from_directory_path(&self.workspace_root)
153            .map_err(|_| anyhow!("Failed to convert workspace path to URI"))?;
154
155        let init_params = json!({
156            "processId": self.process_id,
157            "rootUri": root_uri.as_str(),
158            "capabilities": {
159                "textDocument": {
160                    "documentSymbol": { "hierarchicalDocumentSymbolSupport": true },
161                    "definition": { "dynamicRegistration": false },
162                    "references": { "dynamicRegistration": false },
163                    "callHierarchy": { "dynamicRegistration": false },
164                    "typeHierarchy": { "dynamicRegistration": false }
165                },
166                "workspace": {
167                    "symbol": { "dynamicRegistration": false }
168                }
169            }
170        });
171
172        let _res = self.send_request("initialize", init_params).await?;
173        self.send_notification("initialized", json!({})).await?;
174        Ok(())
175    }
176
177    pub async fn query_symbol(
178        &self,
179        name: &str,
180        kind_filter: &str,
181        exact: bool,
182    ) -> Result<Vec<SymbolItem>> {
183        let params = json!({ "query": name });
184        let res = self.send_request("workspace/symbol", params).await?;
185        let symbols: Vec<SymbolInformation> =
186            match serde_json::from_value::<WorkspaceSymbolResponse>(res.clone()) {
187                Ok(WorkspaceSymbolResponse::Flat(syms)) => syms,
188                Ok(WorkspaceSymbolResponse::Nested(syms)) => syms
189                    .into_iter()
190                    .map(|s| SymbolInformation {
191                        name: s.name,
192                        kind: s.kind,
193                        tags: s.tags,
194                        #[allow(deprecated)]
195                        deprecated: None,
196
197                        location: match s.location {
198                            lsp_types::OneOf::Left(loc) => loc,
199                            lsp_types::OneOf::Right(w_loc) => Location {
200                                uri: w_loc.uri,
201                                range: Range::default(),
202                            },
203                        },
204                        container_name: s.container_name,
205                    })
206                    .collect(),
207                Err(_) => serde_json::from_value(res).unwrap_or_default(),
208            };
209
210        let mut result = Vec::new();
211        for sym in symbols {
212            if exact && sym.name != name {
213                continue;
214            }
215            let kind_str = format!("{:?}", sym.kind).to_lowercase();
216            if kind_filter != "any" && !kind_str.contains(&kind_filter.to_lowercase()) {
217                continue;
218            }
219
220            let file_path = uri_to_file_path(&sym.location.uri);
221            let relative_file = file_path
222                .strip_prefix(&self.workspace_root)
223                .unwrap_or(&file_path)
224                .to_string_lossy()
225                .to_string();
226
227            result.push(SymbolItem {
228                name: sym.name,
229                kind: kind_str,
230                file: relative_file,
231                line: sym.location.range.start.line + 1,
232                col: sym.location.range.start.character + 1,
233                container_name: sym.container_name,
234            });
235        }
236
237        Ok(result)
238    }
239
240    async fn ensure_file_open(&self, canonical_file: &Path) -> Result<Url> {
241        let uri = Url::from_file_path(canonical_file)
242            .map_err(|_| anyhow!("Failed to convert file path to URI"))?;
243        if let Ok(text) = tokio::fs::read_to_string(canonical_file).await {
244            let _ = self
245                .send_notification(
246                    "textDocument/didOpen",
247                    json!({
248                        "textDocument": {
249                            "uri": uri.as_str(),
250                            "languageId": "rust",
251                            "version": 1,
252                            "text": text
253                        }
254                    }),
255                )
256                .await;
257        }
258        Ok(uri)
259    }
260
261    pub async fn query_outline(&self, file: &Path) -> Result<Vec<OutlineItem>> {
262        let canonical_file = file.canonicalize().unwrap_or_else(|_| file.to_path_buf());
263        let uri = self.ensure_file_open(&canonical_file).await?;
264
265        let params = json!({
266            "textDocument": { "uri": uri.as_str() }
267        });
268
269        let res = self
270            .send_request("textDocument/documentSymbol", params)
271            .await?;
272
273        fn convert_symbols(symbols: Vec<DocumentSymbol>) -> Vec<OutlineItem> {
274            symbols
275                .into_iter()
276                .map(|s| OutlineItem {
277                    name: s.name,
278                    kind: format!("{:?}", s.kind).to_lowercase(),
279                    detail: s.detail,
280                    line: s.range.start.line + 1,
281                    col: s.range.start.character + 1,
282                    end_line: s.range.end.line + 1,
283                    children: convert_symbols(s.children.unwrap_or_default()),
284                })
285                .collect()
286        }
287
288        let items = match serde_json::from_value::<DocumentSymbolResponse>(res.clone()) {
289            Ok(DocumentSymbolResponse::Nested(symbols)) => convert_symbols(symbols),
290            Ok(DocumentSymbolResponse::Flat(syms)) => syms
291                .into_iter()
292                .map(|s| OutlineItem {
293                    name: s.name,
294                    kind: format!("{:?}", s.kind).to_lowercase(),
295                    detail: s.container_name,
296                    line: s.location.range.start.line + 1,
297                    col: s.location.range.start.character + 1,
298                    end_line: s.location.range.end.line + 1,
299                    children: vec![],
300                })
301                .collect(),
302            Err(_) => {
303                if let Ok(syms) = serde_json::from_value::<Vec<DocumentSymbol>>(res) {
304                    convert_symbols(syms)
305                } else {
306                    vec![]
307                }
308            }
309        };
310
311        Ok(items)
312    }
313
314    pub async fn query_definition(
315        &self,
316        file: &Path,
317        line: u32,
318        col: u32,
319    ) -> Result<Vec<DefinitionItem>> {
320        let canonical_file = file.canonicalize().unwrap_or_else(|_| file.to_path_buf());
321        let uri = self.ensure_file_open(&canonical_file).await?;
322
323        let params = json!({
324            "textDocument": { "uri": uri.as_str() },
325            "position": { "line": line - 1, "character": col - 1 }
326        });
327
328        let res = self.send_request("textDocument/definition", params).await?;
329        let locations: Vec<Location> = match serde_json::from_value::<GotoDefinitionResponse>(res) {
330            Ok(GotoDefinitionResponse::Scalar(loc)) => vec![loc],
331            Ok(GotoDefinitionResponse::Array(locs)) => locs,
332            Ok(GotoDefinitionResponse::Link(links)) => links
333                .into_iter()
334                .map(|l| Location {
335                    uri: l.target_uri,
336                    range: l.target_selection_range,
337                })
338                .collect(),
339            Err(_) => vec![],
340        };
341
342        let mut items = Vec::new();
343        for loc in locations {
344            let loc_file = uri_to_file_path(&loc.uri);
345            let rel_file = loc_file
346                .strip_prefix(&self.workspace_root)
347                .unwrap_or(&loc_file)
348                .to_string_lossy()
349                .to_string();
350
351            items.push(DefinitionItem {
352                file: rel_file,
353                line: loc.range.start.line + 1,
354                col: loc.range.start.character + 1,
355                end_line: loc.range.end.line + 1,
356                end_col: loc.range.end.character + 1,
357                snippet: None,
358            });
359        }
360
361        Ok(items)
362    }
363
364    pub async fn query_cursor(
365        &self,
366        file: &Path,
367        line: u32,
368        col: u32,
369        mode: &str,
370        _depth: u32,
371    ) -> Result<Vec<CursorItem>> {
372        let canonical_file = file.canonicalize().unwrap_or_else(|_| file.to_path_buf());
373        let uri = self.ensure_file_open(&canonical_file).await?;
374
375        if mode == "references" {
376            let params = json!({
377                "textDocument": { "uri": uri.as_str() },
378                "position": { "line": line - 1, "character": col - 1 },
379                "context": { "includeDeclaration": true }
380            });
381            let res = self.send_request("textDocument/references", params).await?;
382            let locs: Vec<Location> = serde_json::from_value(res).unwrap_or_default();
383
384            let mut items = Vec::new();
385            for loc in locs {
386                let loc_file = uri_to_file_path(&loc.uri);
387                let rel_file = loc_file
388                    .strip_prefix(&self.workspace_root)
389                    .unwrap_or(&loc_file)
390                    .to_string_lossy()
391                    .to_string();
392                items.push(CursorItem {
393                    name: "reference".to_string(),
394                    kind: "reference".to_string(),
395                    file: rel_file,
396                    line: loc.range.start.line + 1,
397                    col: loc.range.start.character + 1,
398                    caller_or_callee: None,
399                });
400            }
401            return Ok(items);
402        }
403
404        // Call Hierarchy
405        let prep_params = json!({
406            "textDocument": { "uri": uri.as_str() },
407            "position": { "line": line - 1, "character": col - 1 }
408        });
409        let prep_res = self
410            .send_request("textDocument/prepareCallHierarchy", prep_params)
411            .await?;
412        let items: Vec<CallHierarchyItem> = serde_json::from_value(prep_res).unwrap_or_default();
413
414        let mut results = Vec::new();
415        if let Some(item) = items.first() {
416            let item_json = serde_json::to_value(item)?;
417            if mode == "outgoing" {
418                let out_res = self
419                    .send_request("callHierarchy/outgoingCalls", json!({ "item": item_json }))
420                    .await?;
421                let out_calls: Vec<CallHierarchyOutgoingCall> =
422                    serde_json::from_value(out_res).unwrap_or_default();
423                for call in out_calls {
424                    let loc_file = uri_to_file_path(&call.to.uri);
425                    let rel_file = loc_file
426                        .strip_prefix(&self.workspace_root)
427                        .unwrap_or(&loc_file)
428                        .to_string_lossy()
429                        .to_string();
430                    results.push(CursorItem {
431                        name: call.to.name,
432                        kind: format!("{:?}", call.to.kind).to_lowercase(),
433                        file: rel_file,
434                        line: call.to.range.start.line + 1,
435                        col: call.to.range.start.character + 1,
436                        caller_or_callee: Some("callee".to_string()),
437                    });
438                }
439            } else {
440                let in_res = self
441                    .send_request("callHierarchy/incomingCalls", json!({ "item": item_json }))
442                    .await?;
443                let in_calls: Vec<CallHierarchyIncomingCall> =
444                    serde_json::from_value(in_res).unwrap_or_default();
445                for call in in_calls {
446                    let loc_file = uri_to_file_path(&call.from.uri);
447                    let rel_file = loc_file
448                        .strip_prefix(&self.workspace_root)
449                        .unwrap_or(&loc_file)
450                        .to_string_lossy()
451                        .to_string();
452                    results.push(CursorItem {
453                        name: call.from.name,
454                        kind: format!("{:?}", call.from.kind).to_lowercase(),
455                        file: rel_file,
456                        line: call.from.range.start.line + 1,
457                        col: call.from.range.start.character + 1,
458                        caller_or_callee: Some("caller".to_string()),
459                    });
460                }
461            }
462        }
463
464        Ok(results)
465    }
466
467    pub async fn query_type_hierarchy(
468        &self,
469        file: &Path,
470        line: u32,
471        col: u32,
472        mode: &str,
473        _depth: u32,
474    ) -> Result<Vec<TypeHierarchyItemResult>> {
475        let canonical_file = file.canonicalize().unwrap_or_else(|_| file.to_path_buf());
476        let uri = self.ensure_file_open(&canonical_file).await?;
477
478        let prep_params = json!({
479            "textDocument": { "uri": uri.as_str() },
480            "position": { "line": line - 1, "character": col - 1 }
481        });
482        let prep_res = self
483            .send_request("textDocument/prepareTypeHierarchy", prep_params)
484            .await?;
485        let items: Vec<lsp_types::TypeHierarchyItem> =
486            serde_json::from_value(prep_res).unwrap_or_default();
487        let mut results = Vec::new();
488
489        if let Some(item) = items.first() {
490            let item_json = serde_json::to_value(item)?;
491            let endpoint = if mode == "subtypes" {
492                "typeHierarchy/subtypes"
493            } else {
494                "typeHierarchy/supertypes"
495            };
496            let type_res = self
497                .send_request(endpoint, json!({ "item": item_json }))
498                .await?;
499            let type_items: Vec<lsp_types::TypeHierarchyItem> =
500                serde_json::from_value(type_res).unwrap_or_default();
501            for ti in type_items {
502                let file_path = uri_to_file_path(&ti.uri);
503                let rel_file = file_path
504                    .strip_prefix(&self.workspace_root)
505                    .unwrap_or(&file_path)
506                    .to_string_lossy()
507                    .to_string();
508                results.push(TypeHierarchyItemResult {
509                    name: ti.name,
510                    kind: format!("{:?}", ti.kind).to_lowercase(),
511                    file: rel_file,
512                    line: ti.range.start.line + 1,
513                    col: ti.range.start.character + 1,
514                    detail: ti.detail,
515                });
516            }
517        }
518
519        Ok(results)
520    }
521}
522
523fn uri_to_file_path(uri: &lsp_types::Uri) -> PathBuf {
524    Url::parse(uri.as_str())
525        .ok()
526        .and_then(|u| u.to_file_path().ok())
527        .unwrap_or_default()
528}