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