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        include_body: bool,
183        max_lines: usize,
184    ) -> Result<Vec<SymbolItem>> {
185        let params = json!({ "query": name });
186        let res = self.send_request("workspace/symbol", params).await?;
187        let symbols: Vec<SymbolInformation> =
188            match serde_json::from_value::<WorkspaceSymbolResponse>(res.clone()) {
189                Ok(WorkspaceSymbolResponse::Flat(syms)) => syms,
190                Ok(WorkspaceSymbolResponse::Nested(syms)) => syms
191                    .into_iter()
192                    .map(|s| SymbolInformation {
193                        name: s.name,
194                        kind: s.kind,
195                        tags: s.tags,
196                        #[allow(deprecated)]
197                        deprecated: None,
198
199                        location: match s.location {
200                            lsp_types::OneOf::Left(loc) => loc,
201                            lsp_types::OneOf::Right(w_loc) => Location {
202                                uri: w_loc.uri,
203                                range: Range::default(),
204                            },
205                        },
206                        container_name: s.container_name,
207                    })
208                    .collect(),
209                Err(_) => serde_json::from_value(res).unwrap_or_default(),
210            };
211
212        let mut result = Vec::new();
213        for sym in symbols {
214            if exact && sym.name != name {
215                continue;
216            }
217            let kind_str = format!("{:?}", sym.kind).to_lowercase();
218            if kind_filter != "any" && !kind_str.contains(&kind_filter.to_lowercase()) {
219                continue;
220            }
221
222            let file_path = uri_to_file_path(&sym.location.uri);
223            let relative_file = file_path
224                .strip_prefix(&self.workspace_root)
225                .unwrap_or(&file_path)
226                .to_string_lossy()
227                .to_string();
228
229            let (snippet, _, _) = if include_body {
230                extract_body_snippet(
231                    &file_path,
232                    sym.location.range.start.line + 1,
233                    sym.location.range.end.line + 1,
234                    max_lines,
235                )
236            } else {
237                (String::new(), 0, false)
238            };
239
240            let body_opt = if include_body && !snippet.is_empty() {
241                Some(snippet)
242            } else {
243                None
244            };
245
246            result.push(SymbolItem {
247                name: sym.name,
248                kind: kind_str,
249                file: relative_file,
250                line: sym.location.range.start.line + 1,
251                col: sym.location.range.start.character + 1,
252                container_name: sym.container_name,
253                body: body_opt,
254            });
255        }
256
257        Ok(result)
258    }
259
260
261    async fn ensure_file_open(&self, canonical_file: &Path) -> Result<Url> {
262        let uri = Url::from_file_path(canonical_file)
263            .map_err(|_| anyhow!("Failed to convert file path to URI"))?;
264        if let Ok(text) = tokio::fs::read_to_string(canonical_file).await {
265            let _ = self
266                .send_notification(
267                    "textDocument/didOpen",
268                    json!({
269                        "textDocument": {
270                            "uri": uri.as_str(),
271                            "languageId": "rust",
272                            "version": 1,
273                            "text": text
274                        }
275                    }),
276                )
277                .await;
278        }
279        Ok(uri)
280    }
281
282    pub async fn query_outline(
283        &self,
284        file: &Path,
285        include_body: bool,
286        max_lines: usize,
287    ) -> Result<Vec<OutlineItem>> {
288        let canonical_file = file.canonicalize().unwrap_or_else(|_| file.to_path_buf());
289        let uri = self.ensure_file_open(&canonical_file).await?;
290
291        let params = json!({
292            "textDocument": { "uri": uri.as_str() }
293        });
294
295        let res = self
296            .send_request("textDocument/documentSymbol", params)
297            .await?;
298
299        fn convert_symbols(
300            symbols: Vec<DocumentSymbol>,
301            file_path: &Path,
302            include_body: bool,
303            max_lines: usize,
304        ) -> Vec<OutlineItem> {
305            symbols
306                .into_iter()
307                .map(|s| {
308                    let (snippet, _, _) = if include_body {
309                        extract_body_snippet(
310                            file_path,
311                            s.range.start.line + 1,
312                            s.range.end.line + 1,
313                            max_lines,
314                        )
315                    } else {
316                        (String::new(), 0, false)
317                    };
318                    let body_opt = if include_body && !snippet.is_empty() {
319                        Some(snippet)
320                    } else {
321                        None
322                    };
323
324                    OutlineItem {
325                        name: s.name,
326                        kind: format!("{:?}", s.kind).to_lowercase(),
327                        detail: s.detail,
328                        line: s.range.start.line + 1,
329                        col: s.range.start.character + 1,
330                        end_line: s.range.end.line + 1,
331                        children: convert_symbols(
332                            s.children.unwrap_or_default(),
333                            file_path,
334                            include_body,
335                            max_lines,
336                        ),
337                        body: body_opt,
338                    }
339                })
340                .collect()
341        }
342
343        let items = match serde_json::from_value::<DocumentSymbolResponse>(res.clone()) {
344            Ok(DocumentSymbolResponse::Nested(symbols)) => {
345                convert_symbols(symbols, &canonical_file, include_body, max_lines)
346            }
347            Ok(DocumentSymbolResponse::Flat(syms)) => syms
348                .into_iter()
349                .map(|s| {
350                    let (snippet, _, _) = if include_body {
351                        extract_body_snippet(
352                            &canonical_file,
353                            s.location.range.start.line + 1,
354                            s.location.range.end.line + 1,
355                            max_lines,
356                        )
357                    } else {
358                        (String::new(), 0, false)
359                    };
360                    let body_opt = if include_body && !snippet.is_empty() {
361                        Some(snippet)
362                    } else {
363                        None
364                    };
365
366                    OutlineItem {
367                        name: s.name,
368                        kind: format!("{:?}", s.kind).to_lowercase(),
369                        detail: s.container_name,
370                        line: s.location.range.start.line + 1,
371                        col: s.location.range.start.character + 1,
372                        end_line: s.location.range.end.line + 1,
373                        children: vec![],
374                        body: body_opt,
375                    }
376                })
377                .collect(),
378            Err(_) => {
379                if let Ok(syms) = serde_json::from_value::<Vec<DocumentSymbol>>(res) {
380                    convert_symbols(syms, &canonical_file, include_body, max_lines)
381                } else {
382                    vec![]
383                }
384            }
385        };
386
387        Ok(items)
388    }
389
390    pub async fn query_definition(
391        &self,
392        file: &Path,
393        line: u32,
394        col: u32,
395        include_body: bool,
396        max_lines: usize,
397    ) -> Result<Vec<DefinitionItem>> {
398        let canonical_file = file.canonicalize().unwrap_or_else(|_| file.to_path_buf());
399        let uri = self.ensure_file_open(&canonical_file).await?;
400
401        let params = json!({
402            "textDocument": { "uri": uri.as_str() },
403            "position": { "line": line - 1, "character": col - 1 }
404        });
405
406        let res = self.send_request("textDocument/definition", params).await?;
407        let locations: Vec<Location> = match serde_json::from_value::<GotoDefinitionResponse>(res) {
408            Ok(GotoDefinitionResponse::Scalar(loc)) => vec![loc],
409            Ok(GotoDefinitionResponse::Array(locs)) => locs,
410            Ok(GotoDefinitionResponse::Link(links)) => links
411                .into_iter()
412                .map(|l| Location {
413                    uri: l.target_uri,
414                    range: l.target_selection_range,
415                })
416                .collect(),
417            Err(_) => vec![],
418        };
419
420        let mut items = Vec::new();
421        for loc in locations {
422            let loc_file = uri_to_file_path(&loc.uri);
423            let rel_file = loc_file
424                .strip_prefix(&self.workspace_root)
425                .unwrap_or(&loc_file)
426                .to_string_lossy()
427                .to_string();
428
429            let (snippet, _, _) = if include_body {
430                extract_body_snippet(
431                    &loc_file,
432                    loc.range.start.line + 1,
433                    loc.range.end.line + 1,
434                    max_lines,
435                )
436            } else {
437                (String::new(), 0, false)
438            };
439
440            let body_opt = if include_body && !snippet.is_empty() {
441                Some(snippet.clone())
442            } else {
443                None
444            };
445
446            items.push(DefinitionItem {
447                file: rel_file,
448                line: loc.range.start.line + 1,
449                col: loc.range.start.character + 1,
450                end_line: loc.range.end.line + 1,
451                end_col: loc.range.end.character + 1,
452                snippet: body_opt.clone(),
453                body: body_opt,
454            });
455        }
456
457        Ok(items)
458    }
459
460
461    pub async fn query_cursor(
462        &self,
463        file: &Path,
464        line: u32,
465        col: u32,
466        mode: &str,
467        _depth: u32,
468    ) -> Result<Vec<CursorItem>> {
469        let canonical_file = file.canonicalize().unwrap_or_else(|_| file.to_path_buf());
470        let uri = self.ensure_file_open(&canonical_file).await?;
471
472        if mode == "references" {
473            let params = json!({
474                "textDocument": { "uri": uri.as_str() },
475                "position": { "line": line - 1, "character": col - 1 },
476                "context": { "includeDeclaration": true }
477            });
478            let res = self.send_request("textDocument/references", params).await?;
479            let locs: Vec<Location> = serde_json::from_value(res).unwrap_or_default();
480
481            let mut items = Vec::new();
482            for loc in locs {
483                let loc_file = uri_to_file_path(&loc.uri);
484                let rel_file = loc_file
485                    .strip_prefix(&self.workspace_root)
486                    .unwrap_or(&loc_file)
487                    .to_string_lossy()
488                    .to_string();
489                items.push(CursorItem {
490                    name: "reference".to_string(),
491                    kind: "reference".to_string(),
492                    file: rel_file,
493                    line: loc.range.start.line + 1,
494                    col: loc.range.start.character + 1,
495                    caller_or_callee: None,
496                });
497            }
498            return Ok(items);
499        }
500
501        // Call Hierarchy
502        let prep_params = json!({
503            "textDocument": { "uri": uri.as_str() },
504            "position": { "line": line - 1, "character": col - 1 }
505        });
506        let prep_res = self
507            .send_request("textDocument/prepareCallHierarchy", prep_params)
508            .await?;
509        let items: Vec<CallHierarchyItem> = serde_json::from_value(prep_res).unwrap_or_default();
510
511        let mut results = Vec::new();
512        if let Some(item) = items.first() {
513            let item_json = serde_json::to_value(item)?;
514            if mode == "outgoing" {
515                let out_res = self
516                    .send_request("callHierarchy/outgoingCalls", json!({ "item": item_json }))
517                    .await?;
518                let out_calls: Vec<CallHierarchyOutgoingCall> =
519                    serde_json::from_value(out_res).unwrap_or_default();
520                for call in out_calls {
521                    let loc_file = uri_to_file_path(&call.to.uri);
522                    let rel_file = loc_file
523                        .strip_prefix(&self.workspace_root)
524                        .unwrap_or(&loc_file)
525                        .to_string_lossy()
526                        .to_string();
527                    results.push(CursorItem {
528                        name: call.to.name,
529                        kind: format!("{:?}", call.to.kind).to_lowercase(),
530                        file: rel_file,
531                        line: call.to.range.start.line + 1,
532                        col: call.to.range.start.character + 1,
533                        caller_or_callee: Some("callee".to_string()),
534                    });
535                }
536            } else {
537                let in_res = self
538                    .send_request("callHierarchy/incomingCalls", json!({ "item": item_json }))
539                    .await?;
540                let in_calls: Vec<CallHierarchyIncomingCall> =
541                    serde_json::from_value(in_res).unwrap_or_default();
542                for call in in_calls {
543                    let loc_file = uri_to_file_path(&call.from.uri);
544                    let rel_file = loc_file
545                        .strip_prefix(&self.workspace_root)
546                        .unwrap_or(&loc_file)
547                        .to_string_lossy()
548                        .to_string();
549                    results.push(CursorItem {
550                        name: call.from.name,
551                        kind: format!("{:?}", call.from.kind).to_lowercase(),
552                        file: rel_file,
553                        line: call.from.range.start.line + 1,
554                        col: call.from.range.start.character + 1,
555                        caller_or_callee: Some("caller".to_string()),
556                    });
557                }
558            }
559        }
560
561        Ok(results)
562    }
563
564    pub async fn query_type_hierarchy(
565        &self,
566        file: &Path,
567        line: u32,
568        col: u32,
569        mode: &str,
570        _depth: u32,
571    ) -> Result<Vec<TypeHierarchyItemResult>> {
572        let canonical_file = file.canonicalize().unwrap_or_else(|_| file.to_path_buf());
573        let uri = self.ensure_file_open(&canonical_file).await?;
574
575        let prep_params = json!({
576            "textDocument": { "uri": uri.as_str() },
577            "position": { "line": line - 1, "character": col - 1 }
578        });
579        let prep_res = self
580            .send_request("textDocument/prepareTypeHierarchy", prep_params)
581            .await?;
582        let items: Vec<lsp_types::TypeHierarchyItem> =
583            serde_json::from_value(prep_res).unwrap_or_default();
584        let mut results = Vec::new();
585
586        if let Some(item) = items.first() {
587            let item_json = serde_json::to_value(item)?;
588            let endpoint = if mode == "subtypes" {
589                "typeHierarchy/subtypes"
590            } else {
591                "typeHierarchy/supertypes"
592            };
593            let type_res = self
594                .send_request(endpoint, json!({ "item": item_json }))
595                .await?;
596            let type_items: Vec<lsp_types::TypeHierarchyItem> =
597                serde_json::from_value(type_res).unwrap_or_default();
598            for ti in type_items {
599                let file_path = uri_to_file_path(&ti.uri);
600                let rel_file = file_path
601                    .strip_prefix(&self.workspace_root)
602                    .unwrap_or(&file_path)
603                    .to_string_lossy()
604                    .to_string();
605                results.push(TypeHierarchyItemResult {
606                    name: ti.name,
607                    kind: format!("{:?}", ti.kind).to_lowercase(),
608                    file: rel_file,
609                    line: ti.range.start.line + 1,
610                    col: ti.range.start.character + 1,
611                    detail: ti.detail,
612                });
613            }
614        }
615
616        Ok(results)
617    }
618
619    pub async fn query_body(&self, file: &Path, line: u32, col: u32, max_lines: usize) -> Result<BodyItem> {
620        let canonical_file = file.canonicalize().unwrap_or_else(|_| file.to_path_buf());
621        let uri = self.ensure_file_open(&canonical_file).await?;
622
623        let params = json!({
624            "textDocument": { "uri": uri.as_str() },
625            "position": { "line": line - 1, "character": col - 1 }
626        });
627
628        let res = self.send_request("textDocument/definition", params).await?;
629        let locations: Vec<Location> = match serde_json::from_value::<GotoDefinitionResponse>(res) {
630            Ok(GotoDefinitionResponse::Scalar(loc)) => vec![loc],
631            Ok(GotoDefinitionResponse::Array(locs)) => locs,
632            Ok(GotoDefinitionResponse::Link(links)) => links
633                .into_iter()
634                .map(|l| Location {
635                    uri: l.target_uri,
636                    range: l.target_selection_range,
637                })
638                .collect(),
639            Err(_) => vec![],
640        };
641
642        if let Some(loc) = locations.first() {
643            let loc_file = uri_to_file_path(&loc.uri);
644            let rel_file = loc_file
645                .strip_prefix(&self.workspace_root)
646                .unwrap_or(&loc_file)
647                .to_string_lossy()
648                .to_string();
649
650            let (snippet, total_lines, is_truncated) = extract_body_snippet(
651                &loc_file,
652                loc.range.start.line + 1,
653                loc.range.end.line + 1,
654                max_lines,
655            );
656
657            Ok(BodyItem {
658                file: rel_file,
659                line: loc.range.start.line + 1,
660                col: loc.range.start.character + 1,
661                end_line: loc.range.end.line + 1,
662                end_col: loc.range.end.character + 1,
663                body: snippet,
664                total_lines,
665                is_truncated,
666            })
667        } else {
668            Err(anyhow::anyhow!("No symbol/entity definition found at specified position"))
669        }
670    }
671}
672
673fn uri_to_file_path(uri: &lsp_types::Uri) -> PathBuf {
674    Url::parse(uri.as_str())
675        .ok()
676        .and_then(|u| u.to_file_path().ok())
677        .unwrap_or_default()
678}
679
680pub fn extract_body_snippet(file: &Path, start_line: u32, end_line: u32, max_lines: usize) -> (String, usize, bool) {
681    let content = match std::fs::read_to_string(file) {
682        Ok(c) => c,
683        Err(_) => return (String::new(), 0, false),
684    };
685    let lines: Vec<&str> = content.lines().collect();
686    let start_idx = (start_line.saturating_sub(1)) as usize;
687    let end_idx = (end_line as usize).min(lines.len());
688    if start_idx >= lines.len() || start_idx > end_idx {
689        return (String::new(), 0, false);
690    }
691    let total = end_idx - start_idx;
692    let limit = if max_lines == 0 { total } else { max_lines.min(total) };
693    let is_truncated = limit < total;
694    let slice = &lines[start_idx..start_idx + limit];
695    (slice.join("\n"), total, is_truncated)
696}
697
698pub fn format_body_with_line_numbers(snippet: &str, start_line: u32) -> String {
699    snippet
700        .lines()
701        .enumerate()
702        .map(|(idx, line)| format!("{:4} | {}", start_line as usize + idx, line))
703        .collect::<Vec<_>>()
704        .join("\n")
705}
706