feat(lsp): cancellable heavy requests (#14851)

<!--
if this PR closes one or more issues, you can automatically link the PR
with
them by using one of the [*linking
keywords*](https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword),
e.g.
- this PR should close #xxxx
- fixes #xxxx

you can also mention related issues, PRs or discussions!
-->

# Description
<!--
Thank you for improving Nushell. Please, check our [contributing
guide](../CONTRIBUTING.md) and talk to the core team before making major
changes.

Description of your pull request goes here. **Provide examples and/or
screenshots** if your changes affect the user experience.
-->

`tower-lsp` seems not well-maintained, I ended up with a dedicated
thread for heavy computing and message passing to cancel it on any new
request.

During the progress, interrupting with edits or new requests.

<img width="522" alt="image"
src="https://github.com/user-attachments/assets/b263d73d-8ea3-4b26-a7b7-e0b30462d1af"
/>

Goto references are still blocking, with a hard timeout of 5 seconds.
Only locations found within the time limit are returned. Technically,
reference requests allow for responses with partial results, which means
instant responsiveness. However, the `lsp_types` crate hasn’t enabled
this. I believe I can still enable it with some JSON manipulation, but
I’ll leave it for future work.

# User-Facing Changes
<!-- List of all changes that impact the user experience here. This
helps us keep track of breaking changes. -->

# Tests + Formatting
<!--
Don't forget to add tests that cover your changes.

Make sure you've run and fixed any issues with these commands:

- `cargo fmt --all -- --check` to check standard code formatting (`cargo
fmt --all` applies these changes)
- `cargo clippy --workspace -- -D warnings -D clippy::unwrap_used` to
check that you're using the standard code style
- `cargo test --workspace` to check that all tests pass (on Windows make
sure to [enable developer
mode](https://learn.microsoft.com/en-us/windows/apps/get-started/developer-mode-features-and-debugging))
- `cargo run -- -c "use toolkit.nu; toolkit test stdlib"` to run the
tests for the standard library

> **Note**
> from `nushell` you can also use the `toolkit` as follows
> ```bash
> use toolkit.nu # or use an `env_change` hook to activate it
automatically
> toolkit check pr
> ```
-->

Need some clever way to test the cancellation, no test cases added yet.

# After Submitting
<!-- If your PR had any user-facing changes, update [the
documentation](https://github.com/nushell/nushell.github.io) after the
PR is merged, if necessary. This will help us keep the docs up to date.
-->

crates/nu-lsp/src/diagnostics.rs

@@ -3,24 +3,30 @@ use lsp_types::{
     notification::{Notification, PublishDiagnostics},
     Diagnostic, DiagnosticSeverity, PublishDiagnosticsParams, Uri,
 };
-use miette::{IntoDiagnostic, Result};
+use miette::{miette, IntoDiagnostic, Result};
 
 impl LanguageServer {
     pub(crate) fn publish_diagnostics_for_file(&mut self, uri: Uri) -> Result<()> {
         let mut engine_state = self.new_engine_state();
         engine_state.generate_nu_constant();
 
-        let Some((_, offset, working_set, file)) = self.parse_file(&mut engine_state, &uri, true)
-        else {
+        let Some((_, offset, working_set)) = self.parse_file(&mut engine_state, &uri, true) else {
             return Ok(());
         };
 
         let mut diagnostics = PublishDiagnosticsParams {
-            uri,
+            uri: uri.clone(),
             diagnostics: Vec::new(),
             version: None,
         };
 
+        let docs = match self.docs.lock() {
+            Ok(it) => it,
+            Err(err) => return Err(miette!(err.to_string())),
+        };
+        let file = docs
+            .get_document(&uri)
+            .ok_or_else(|| miette!("\nFailed to get document"))?;
         for err in working_set.parse_errors.iter() {
             let message = err.to_string();
 

crates/nu-lsp/src/goto.rs

@@ -33,7 +33,7 @@ impl LanguageServer {
             .text_document
             .uri
             .to_owned();
-        let (working_set, id, _, _, _) = self
+        let (working_set, id, _, _) = self
             .parse_and_find(
                 &mut engine_state,
                 &path_uri,

crates/nu-lsp/src/lib.rs

@@ -1,5 +1,6 @@
 #![doc = include_str!("../README.md")]
 use ast::find_id;
+use crossbeam_channel::{Receiver, Sender};
 use lsp_server::{Connection, IoThreads, Message, Response, ResponseError};
 use lsp_textdocument::{FullTextDocument, TextDocuments};
 use lsp_types::{
@@ -18,7 +19,7 @@ use nu_protocol::{
     engine::{CachedFile, EngineState, Stack, StateWorkingSet},
     DeclId, ModuleId, Span, Type, Value, VarId,
 };
-use std::collections::BTreeMap;
+use std::{collections::BTreeMap, sync::Mutex};
 use std::{
     path::{Path, PathBuf},
     str::FromStr,
@@ -27,6 +28,7 @@ use std::{
 };
 use symbols::SymbolCache;
 use url::Url;
+use workspace::{InternalMessage, RangePerDoc};
 
 mod ast;
 mod diagnostics;
@@ -47,13 +49,14 @@ pub enum Id {
 pub struct LanguageServer {
     connection: Connection,
     io_threads: Option<IoThreads>,
-    docs: TextDocuments,
+    docs: Arc<Mutex<TextDocuments>>,
     engine_state: EngineState,
     symbol_cache: SymbolCache,
     inlay_hints: BTreeMap<Uri, Vec<InlayHint>>,
     workspace_folders: BTreeMap<String, WorkspaceFolder>,
     // for workspace wide requests
     occurrences: BTreeMap<Uri, Vec<Range>>,
+    channels: Option<(Sender<bool>, Arc<Receiver<InternalMessage>>)>,
 }
 
 pub fn path_to_uri(path: impl AsRef<Path>) -> Uri {
@@ -92,12 +95,13 @@ impl LanguageServer {
         Ok(Self {
             connection,
             io_threads,
-            docs: TextDocuments::new(),
+            docs: Arc::new(Mutex::new(TextDocuments::new())),
             engine_state,
             symbol_cache: SymbolCache::new(),
             inlay_hints: BTreeMap::new(),
             workspace_folders: BTreeMap::new(),
             occurrences: BTreeMap::new(),
+            channels: None,
         })
     }
 
@@ -141,12 +145,19 @@ impl LanguageServer {
         self.initialize_workspace_folders(init_params)?;
 
         while !self.engine_state.signals().interrupted() {
+            // first check new messages from child thread
+            self.handle_internal_messages()?;
+
             let msg = match self
                 .connection
                 .receiver
                 .recv_timeout(Duration::from_secs(1))
             {
-                Ok(msg) => msg,
+                Ok(msg) => {
+                    // cancel execution if other messages received before job done
+                    self.cancel_background_thread();
+                    msg
+                }
                 Err(crossbeam_channel::RecvTimeoutError::Timeout) => {
                     continue;
                 }
@@ -177,7 +188,9 @@ impl LanguageServer {
                             Self::handle_lsp_request(request, |params| self.document_symbol(params))
                         }
                         request::References::METHOD => {
-                            Self::handle_lsp_request(request, |params| self.references(params))
+                            Self::handle_lsp_request(request, |params| {
+                                self.references(params, 5000)
+                            })
                         }
                         request::WorkspaceSymbolRequest::METHOD => {
                             Self::handle_lsp_request(request, |params| {
@@ -224,6 +237,42 @@ impl LanguageServer {
         Ok(())
     }
 
+    /// Send a cancel message to a running bg thread
+    pub fn cancel_background_thread(&mut self) {
+        if let Some((sender, _)) = &self.channels {
+            sender.send(true).ok();
+        }
+    }
+
+    /// Check results from background thread
+    pub fn handle_internal_messages(&mut self) -> Result<bool> {
+        let mut reset = false;
+        if let Some((_, receiver)) = &self.channels {
+            for im in receiver.try_iter() {
+                match im {
+                    InternalMessage::RangeMessage(RangePerDoc { uri, ranges }) => {
+                        self.occurrences.insert(uri, ranges);
+                    }
+                    InternalMessage::OnGoing(token, progress) => {
+                        self.send_progress_report(token, progress, None)?;
+                    }
+                    InternalMessage::Finished(token) => {
+                        reset = true;
+                        self.send_progress_end(token, Some("Finished.".to_string()))?;
+                    }
+                    InternalMessage::Cancelled(token) => {
+                        reset = true;
+                        self.send_progress_end(token, Some("interrupted.".to_string()))?;
+                    }
+                }
+            }
+        }
+        if reset {
+            self.channels = None;
+        }
+        Ok(reset)
+    }
+
     pub fn new_engine_state(&self) -> EngineState {
         let mut engine_state = self.engine_state.clone();
         let cwd = std::env::current_dir().expect("Could not get current working directory.");
@@ -236,27 +285,35 @@ impl LanguageServer {
         engine_state: &'a mut EngineState,
         uri: &Uri,
         pos: Position,
-    ) -> Result<(StateWorkingSet<'a>, Id, Span, usize, &FullTextDocument)> {
-        let (block, file_offset, mut working_set, file) = self
+    ) -> Result<(StateWorkingSet<'a>, Id, Span, usize)> {
+        let (block, file_offset, mut working_set) = self
             .parse_file(engine_state, uri, false)
             .ok_or_else(|| miette!("\nFailed to parse current file"))?;
 
+        let docs = match self.docs.lock() {
+            Ok(it) => it,
+            Err(err) => return Err(miette!(err.to_string())),
+        };
+        let file = docs
+            .get_document(uri)
+            .ok_or_else(|| miette!("\nFailed to get document"))?;
         let location = file.offset_at(pos) as usize + file_offset;
         let (id, span) = find_id(&block, &working_set, &location)
             .ok_or_else(|| miette!("\nFailed to find current name"))?;
         // add block to working_set for later references
         working_set.add_block(block);
-        Ok((working_set, id, span, file_offset, file))
+        Ok((working_set, id, span, file_offset))
     }
 
     pub fn parse_file<'a>(
         &mut self,
         engine_state: &'a mut EngineState,
         uri: &Uri,
         need_hints: bool,
-    ) -> Option<(Arc<Block>, usize, StateWorkingSet<'a>, &FullTextDocument)> {
+    ) -> Option<(Arc<Block>, usize, StateWorkingSet<'a>)> {
         let mut working_set = StateWorkingSet::new(engine_state);
-        let file = self.docs.get_document(uri)?;
+        let docs = self.docs.lock().ok()?;
+        let file = docs.get_document(uri)?;
         let file_path = uri_to_path(uri);
         let file_path_str = file_path.to_str()?;
         let contents = file.get_content(None).as_bytes();
@@ -270,7 +327,7 @@ impl LanguageServer {
             let file_inlay_hints = self.extract_inlay_hints(&working_set, &block, offset, file);
             self.inlay_hints.insert(uri.clone(), file_inlay_hints);
         }
-        Some((block, offset, working_set, file))
+        Some((block, offset, working_set))
     }
 
     fn get_location_by_span<'a>(
@@ -285,10 +342,10 @@ impl LanguageServer {
                     return None;
                 }
                 let target_uri = path_to_uri(path);
-                if let Some(doc) = self.docs.get_document(&target_uri) {
+                if let Some(file) = self.docs.lock().ok()?.get_document(&target_uri) {
                     return Some(Location {
                         uri: target_uri,
-                        range: span_to_range(span, doc, cached_file.covered_span.start),
+                        range: span_to_range(span, file, cached_file.covered_span.start),
                     });
                 } else {
                     // in case where the document is not opened yet, typically included by `nu -I`
@@ -344,7 +401,7 @@ impl LanguageServer {
             .text_document
             .uri
             .to_owned();
-        let (working_set, id, _, _, _) = self
+        let (working_set, id, _, _) = self
             .parse_and_find(
                 &mut engine_state,
                 &path_uri,
@@ -525,7 +582,8 @@ impl LanguageServer {
 
     fn complete(&mut self, params: &CompletionParams) -> Option<CompletionResponse> {
         let path_uri = params.text_document_position.text_document.uri.to_owned();
-        let file = self.docs.get_document(&path_uri)?;
+        let docs = self.docs.lock().ok()?;
+        let file = docs.get_document(&path_uri)?;
 
         let mut completer =
             NuCompleter::new(Arc::new(self.engine_state.clone()), Arc::new(Stack::new()));

crates/nu-lsp/src/notification.rs

@@ -17,8 +17,8 @@ impl LanguageServer {
         &mut self,
         notification: lsp_server::Notification,
     ) -> Option<Uri> {
-        self.docs
-            .listen(notification.method.as_str(), &notification.params);
+        let mut docs = self.docs.lock().ok()?;
+        docs.listen(notification.method.as_str(), &notification.params);
         match notification.method.as_str() {
             DidOpenTextDocument::METHOD => {
                 let params: DidOpenTextDocumentParams =
@@ -57,7 +57,7 @@ impl LanguageServer {
         }
     }
 
-    fn send_progress_notification(
+    pub fn send_progress_notification(
         &self,
         token: ProgressToken,
         value: WorkDoneProgress,
@@ -74,12 +74,13 @@ impl LanguageServer {
             .into_diagnostic()
     }
 
-    pub fn send_progress_begin(&self, token: ProgressToken, title: &str) -> Result<()> {
+    pub fn send_progress_begin(&self, token: ProgressToken, title: String) -> Result<()> {
         self.send_progress_notification(
             token,
             WorkDoneProgress::Begin(WorkDoneProgressBegin {
-                title: title.to_string(),
+                title,
                 percentage: Some(0),
+                cancellable: Some(true),
                 ..Default::default()
             }),
         )
@@ -95,8 +96,8 @@ impl LanguageServer {
             token,
             WorkDoneProgress::Report(WorkDoneProgressReport {
                 message,
+                cancellable: Some(true),
                 percentage: Some(percentage),
-                ..Default::default()
             }),
         )
     }

crates/nu-lsp/src/symbols.rs

@@ -272,7 +272,8 @@ impl LanguageServer {
     ) -> Option<DocumentSymbolResponse> {
         let engine_state = self.new_engine_state();
         let uri = params.text_document.uri.to_owned();
-        self.symbol_cache.update(&uri, &engine_state, &self.docs);
+        let docs = self.docs.lock().ok()?;
+        self.symbol_cache.update(&uri, &engine_state, &docs);
         Some(DocumentSymbolResponse::Flat(
             self.symbol_cache.get_symbols_by_uri(&uri)?,
         ))
@@ -284,7 +285,8 @@ impl LanguageServer {
     ) -> Option<WorkspaceSymbolResponse> {
         if self.symbol_cache.any_dirty() {
             let engine_state = self.new_engine_state();
-            self.symbol_cache.update_all(&engine_state, &self.docs);
+            let docs = self.docs.lock().ok()?;
+            self.symbol_cache.update_all(&engine_state, &docs);
         }
         Some(WorkspaceSymbolResponse::Flat(
             self.symbol_cache