fix(diagnostics): refresh dependent documents after file changes

crates/emmylua_ls/Cargo.toml

@@ -53,6 +53,7 @@ optional = true
 
 [dev-dependencies]
 googletest.workspace = true
+tokio = { workspace = true, features = ["test-util"] }
 
 
 [[bin]]

crates/emmylua_ls/src/context/diagnostic_service.rs

@@ -18,6 +18,8 @@ pub struct DiagnosticService {
     client: Arc<ClientProxy>,
     status_bar: Arc<StatusBar>,
     diagnostic_tokens: Arc<Mutex<HashMap<FileId, CancellationToken>>>,
+    refresh_token: Arc<Mutex<Option<CancellationToken>>>,
+    refresh_shutdown: CancellationToken,
     workspace_diagnostic_token: Arc<Mutex<Option<CancellationToken>>>,
     workspace_run_lock: Arc<Mutex<()>>,
     workspace_diagnostic_done: AtomicBool,
@@ -34,6 +36,8 @@ impl DiagnosticService {
             analysis,
             client,
             diagnostic_tokens: Arc::new(Mutex::new(HashMap::new())),
+            refresh_token: Arc::new(Mutex::new(None)),
+            refresh_shutdown: CancellationToken::new(),
             workspace_diagnostic_token: Arc::new(Mutex::new(None)),
             workspace_run_lock: Arc::new(Mutex::new(())),
             workspace_diagnostic_done: AtomicBool::new(false),
@@ -42,6 +46,52 @@ impl DiagnosticService {
         }
     }
 
+    /// Coalesce cross-file invalidations after analysis has incorporated an edit.
+    /// `interval` is in milliseconds; zero requests a refresh without a debounce delay.
+    pub async fn schedule_diagnostic_refresh(&self, interval: u64) {
+        let mut pending = self.refresh_token.lock().await;
+        if self.refresh_shutdown.is_cancelled() {
+            return;
+        }
+        if let Some(token) = pending.take() {
+            token.cancel();
+        }
+        let token = CancellationToken::new();
+        *pending = Some(token.clone());
+        drop(pending);
+        let client = Arc::downgrade(&self.client);
+        let pending = self.refresh_token.clone();
+        let shutdown = self.refresh_shutdown.clone();
+        let deadline = tokio::time::Instant::now() + Duration::from_millis(interval);
+        tokio::spawn(async move {
+            tokio::select! {
+                biased;
+                _ = shutdown.cancelled() => {},
+                _ = token.cancelled() => {},
+                _ = tokio::time::sleep_until(deadline) => {
+                    // Serialize the final send with replacement and shutdown, so no
+                    // refresh can be sent after shutdown_diagnostic_refresh returns.
+                    let mut pending = pending.lock().await;
+                    if !shutdown.is_cancelled() && !token.is_cancelled() {
+                        if let Some(client) = client.upgrade() {
+                            client.refresh_workspace_diagnostics();
+                        }
+                        pending.take();
+                    }
+                }
+            }
+        });
+    }
+
+    /// Stop pending refreshes and reject scheduling by outstanding update tasks.
+    pub async fn shutdown_diagnostic_refresh(&self) {
+        let mut pending = self.refresh_token.lock().await;
+        self.refresh_shutdown.cancel();
+        if let Some(token) = pending.take() {
+            token.cancel();
+        }
+    }
+
     pub async fn add_diagnostic_task(&self, file_id: FileId, interval: u64) {
         let mut tokens = self.diagnostic_tokens.lock().await;
 
@@ -385,3 +435,106 @@ async fn run_workspace_batch(
 
     (result, count == valid_file_count)
 }
+
+impl Drop for DiagnosticService {
+    fn drop(&mut self) {
+        self.refresh_shutdown.cancel();
+    }
+}
+
+#[cfg(test)]
+mod refresh_tests {
+    use super::*;
+    use crate::context::ServerContext;
+    use lsp_server::{Connection, Message};
+
+    fn diagnostic() -> (Arc<DiagnosticService>, Connection) {
+        let (server, client_connection) = Connection::memory();
+        let client = Arc::new(ClientProxy::new(Arc::new(server)));
+        let status = Arc::new(StatusBar::new(client.clone(), false));
+        (
+            Arc::new(DiagnosticService::new(
+                Arc::new(AnalysisState::new()),
+                status,
+                client,
+            )),
+            client_connection,
+        )
+    }
+
+    async fn advance(milliseconds: u64) {
+        tokio::task::yield_now().await;
+        tokio::time::advance(Duration::from_millis(milliseconds)).await;
+        tokio::task::yield_now().await;
+    }
+
+    fn expect_refresh(client: &Connection) {
+        let Message::Request(request) = client.receiver.try_recv().unwrap() else {
+            panic!("expected request")
+        };
+        assert_eq!(request.method, "workspace/diagnostic/refresh");
+        assert!(client.receiver.try_recv().is_err());
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn refresh_coalesces_and_is_scoped_to_its_client() {
+        let (first, client) = diagnostic();
+        let (_second, other_client) = diagnostic();
+        for _ in 0..10 {
+            first.schedule_diagnostic_refresh(1000).await;
+        }
+        advance(999).await;
+        assert!(client.receiver.try_recv().is_err());
+        advance(1).await;
+        expect_refresh(&client);
+        assert!(other_client.receiver.try_recv().is_err());
+        assert!(first.refresh_token.lock().await.is_none());
+        first.schedule_diagnostic_refresh(0).await;
+        advance(0).await;
+        expect_refresh(&client);
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn a_later_edit_restarts_the_debounce_interval() {
+        let (diagnostic, client) = diagnostic();
+        diagnostic.schedule_diagnostic_refresh(1000).await;
+        advance(900).await;
+        diagnostic.schedule_diagnostic_refresh(1000).await;
+        advance(100).await;
+        assert!(client.receiver.try_recv().is_err());
+        advance(900).await;
+        expect_refresh(&client);
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn dropping_diagnostics_cancels_without_retaining_the_client() {
+        let (diagnostic, client) = diagnostic();
+        let weak_client = Arc::downgrade(&diagnostic.client);
+        let references = Arc::strong_count(&diagnostic.client);
+        diagnostic.schedule_diagnostic_refresh(1000).await;
+        assert_eq!(Arc::strong_count(&diagnostic.client), references);
+        drop(diagnostic);
+        assert!(weak_client.upgrade().is_none());
+        advance(1000).await;
+        assert!(client.receiver.try_recv().is_err());
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn closing_context_stops_refreshes_while_snapshots_remain_alive() {
+        let (server, client) = Connection::memory();
+        let context = ServerContext::new(Arc::new(server), Default::default());
+        let snapshot = context.snapshot();
+        snapshot
+            .file_diagnostic()
+            .schedule_diagnostic_refresh(1000)
+            .await;
+        context.close().await;
+        // An already queued update must not restart refreshes during shutdown.
+        snapshot
+            .file_diagnostic()
+            .schedule_diagnostic_refresh(0)
+            .await;
+        advance(1000).await;
+        assert!(client.receiver.try_recv().is_err());
+    }
+}

crates/emmylua_ls/src/context/mod.rs

@@ -53,7 +53,6 @@ impl ServerContext {
             client.clone(),
         ));
         let workspace_manager = Arc::new(Mutex::new(WorkspaceManager::new(
-            client.clone(),
             file_diagnostic.clone(),
             lsp_features.clone(),
         )));
@@ -136,6 +135,10 @@ impl ServerContext {
     }
 
     pub async fn close(&self) {
+        self.inner
+            .file_diagnostic
+            .shutdown_diagnostic_refresh()
+            .await;
         let mut workspace_manager = self.inner.workspace_manager.lock().await;
         workspace_manager.watcher = None;
     }

crates/emmylua_ls/src/context/snapshot.rs

@@ -30,6 +30,17 @@ impl ServerContextSnapshot {
         &self.inner.file_diagnostic
     }
 
+    /// Notify pull clients after a change that can affect other documents.
+    pub async fn schedule_diagnostic_refresh(&self, interval: u64) {
+        if self.lsp_features().supports_pull_diagnostic()
+            && self.lsp_features().supports_refresh_diagnostic()
+        {
+            self.file_diagnostic()
+                .schedule_diagnostic_refresh(interval)
+                .await;
+        }
+    }
+
     pub fn workspace_manager(&self) -> &Mutex<WorkspaceManager> {
         &self.inner.workspace_manager
     }

crates/emmylua_ls/src/context/workspace_manager.rs

@@ -7,7 +7,7 @@ use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU8, AtomicU64, Ordering};
 use std::{path::PathBuf, sync::Arc, time::Duration};
 
 use super::workspace_state::{OpenFilesSnapshot, WorkspaceState};
-use super::{AnalysisState, ClientProxy, DiagnosticService};
+use super::{AnalysisState, DiagnosticService};
 use crate::context::ServerContextSnapshot;
 use crate::context::lsp_features::LspFeatures;
 use crate::handlers::{ClientConfig, init_analysis, register_files_watch};
@@ -19,7 +19,6 @@ use tokio::sync::Mutex as AsyncMutex;
 use tokio_util::sync::CancellationToken;
 
 pub struct WorkspaceManager {
-    client: Arc<ClientProxy>,
     config_reload_token: Arc<PendingTask>,
     reload_lock: Arc<AsyncMutex<()>>,
     reload_generation: Arc<AtomicU64>,
@@ -31,13 +30,8 @@ pub struct WorkspaceManager {
 }
 
 impl WorkspaceManager {
-    pub fn new(
-        client: Arc<ClientProxy>,
-        file_diagnostic: Arc<DiagnosticService>,
-        lsp_features: Arc<LspFeatures>,
-    ) -> Self {
+    pub fn new(file_diagnostic: Arc<DiagnosticService>, lsp_features: Arc<LspFeatures>) -> Self {
         Self {
-            client,
             config_reload_token: Arc::new(PendingTask::default()),
             reload_lock: Arc::new(AsyncMutex::new(())),
             reload_generation: Arc::new(AtomicU64::new(0)),
@@ -137,7 +131,6 @@ impl WorkspaceManager {
     //     }
 
     //     let analysis = self.analysis.clone();
-    //     let client = self.client.clone();
     //     let file_diagnostic = self.file_diagnostic.clone();
     //     let lsp_features = self.lsp_features.clone();
     //     let reindex_token = self.reindex_token.clone();
@@ -161,7 +154,6 @@ impl WorkspaceManager {
     //         refresh_workspace_diagnostics(
     //             file_diagnostic,
     //             lsp_features,
-    //             client,
     //             workspace_diagnostic_level,
     //         )
     //         .await;
@@ -171,7 +163,6 @@ impl WorkspaceManager {
 
     fn reload_task_handles(&self) -> ReloadTaskHandles {
         ReloadTaskHandles {
-            client: self.client.clone(),
             file_diagnostic: self.file_diagnostic.clone(),
             lsp_features: self.lsp_features.clone(),
             reload_lock: self.reload_lock.clone(),
@@ -305,14 +296,15 @@ impl PendingTask {
 async fn refresh_workspace_diagnostics(
     file_diagnostic: Arc<DiagnosticService>,
     lsp_features: Arc<LspFeatures>,
-    client: Arc<ClientProxy>,
     workspace_diagnostic_level: Arc<AtomicU8>,
 ) {
     file_diagnostic.cancel_workspace_diagnostic().await;
     workspace_diagnostic_level.store(WorkspaceDiagnosticLevel::Fast.to_u8(), Ordering::Release);
 
     if lsp_features.supports_workspace_diagnostic() {
-        client.refresh_workspace_diagnostics();
+        if lsp_features.supports_refresh_diagnostic() {
+            file_diagnostic.schedule_diagnostic_refresh(0).await;
+        }
     } else {
         file_diagnostic
             .add_workspace_diagnostic_task(500, true)
@@ -322,7 +314,6 @@ async fn refresh_workspace_diagnostics(
 
 #[derive(Clone)]
 struct ReloadTaskHandles {
-    client: Arc<ClientProxy>,
     file_diagnostic: Arc<DiagnosticService>,
     lsp_features: Arc<LspFeatures>,
     reload_lock: Arc<AsyncMutex<()>>,
@@ -351,7 +342,6 @@ fn spawn_workspace_reload_task(
         refresh_workspace_diagnostics(
             handles.file_diagnostic,
             handles.lsp_features,
-            handles.client,
             handles.workspace_diagnostic_level,
         )
         .await;

crates/emmylua_ls/src/handlers/diagnostic/mod.rs

@@ -15,7 +15,7 @@ impl RegisterCapabilities for DiagnosticCapabilities {
         server_capabilities.diagnostic_provider =
             Some(DiagnosticServerCapabilities::Options(DiagnosticOptions {
                 identifier: Some("EmmyLua".to_string()),
-                inter_file_dependencies: false,
+                inter_file_dependencies: true,
                 workspace_diagnostics: true,
                 ..Default::default()
             }))