Fix missing diagnostics for notebooks (#21156)

crates/ruff_server/src/server/api/diagnostics.rs

@@ -1,7 +1,4 @@
-use lsp_types::Url;
-
 use crate::{
-    Session,
     lint::DiagnosticsMap,
     session::{Client, DocumentQuery, DocumentSnapshot},
 };
@@ -22,21 +19,10 @@ pub(super) fn generate_diagnostics(snapshot: &DocumentSnapshot) -> DiagnosticsMa
 }
 
 pub(super) fn publish_diagnostics_for_document(
-    session: &Session,
+    snapshot: &DocumentSnapshot,
-    url: &Url,
     client: &Client,
 ) -> crate::server::Result<()> {
-    // Publish diagnostics if the client doesn't support pull diagnostics
+    for (uri, diagnostics) in generate_diagnostics(snapshot) {
-    if session.resolved_client_capabilities().pull_diagnostics {
-        return Ok(());
-    }
-
-    let snapshot = session
-        .take_snapshot(url.clone())
-        .ok_or_else(|| anyhow::anyhow!("Unable to take snapshot for document with URL {url}"))
-        .with_failure_code(lsp_server::ErrorCode::InternalError)?;
-
-    for (uri, diagnostics) in generate_diagnostics(&snapshot) {
         client
             .send_notification::<lsp_types::notification::PublishDiagnostics>(
                 lsp_types::PublishDiagnosticsParams {
@@ -52,14 +38,9 @@ pub(super) fn publish_diagnostics_for_document(
 }
 
 pub(super) fn clear_diagnostics_for_document(
-    session: &Session,
     query: &DocumentQuery,
     client: &Client,
 ) -> crate::server::Result<()> {
-    if session.resolved_client_capabilities().pull_diagnostics {
-        return Ok(());
-    }
-
     client
         .send_notification::<lsp_types::notification::PublishDiagnostics>(
             lsp_types::PublishDiagnosticsParams {

crates/ruff_server/src/server/api/notifications/did_change.rs

@@ -31,7 +31,11 @@ impl super::SyncNotificationHandler for DidChange {
             .update_text_document(&key, content_changes, new_version)
             .with_failure_code(ErrorCode::InternalError)?;
 
-        publish_diagnostics_for_document(session, &key.into_url(), client)?;
+        // Publish diagnostics if the client doesn't support pull diagnostics
+        if !session.resolved_client_capabilities().pull_diagnostics {
+            let snapshot = session.take_snapshot(key.into_url()).unwrap();
+            publish_diagnostics_for_document(&snapshot, client)?;
+        }
 
         Ok(())
     }

crates/ruff_server/src/server/api/notifications/did_change_notebook.rs

@@ -27,7 +27,10 @@ impl super::SyncNotificationHandler for DidChangeNotebook {
             .with_failure_code(ErrorCode::InternalError)?;
 
         // publish new diagnostics
-        publish_diagnostics_for_document(session, &key.into_url(), client)?;
+        let snapshot = session
+            .take_snapshot(key.into_url())
+            .expect("snapshot should be available");
+        publish_diagnostics_for_document(&snapshot, client)?;
 
         Ok(())
     }

crates/ruff_server/src/server/api/notifications/did_change_watched_files.rs

@@ -31,13 +31,19 @@ impl super::SyncNotificationHandler for DidChangeWatchedFiles {
             } else {
                 // publish diagnostics for text documents
                 for url in session.text_document_urls() {
-                    publish_diagnostics_for_document(session, url, client)?;
+                    let snapshot = session
+                        .take_snapshot(url.clone())
+                        .expect("snapshot should be available");
+                    publish_diagnostics_for_document(&snapshot, client)?;
                 }
             }
 
             // always publish diagnostics for notebook files (since they don't use pull diagnostics)
             for url in session.notebook_document_urls() {
-                publish_diagnostics_for_document(session, url, client)?;
+                let snapshot = session
+                    .take_snapshot(url.clone())
+                    .expect("snapshot should be available");
+                publish_diagnostics_for_document(&snapshot, client)?;
             }
         }
 

crates/ruff_server/src/server/api/notifications/did_close.rs

@@ -27,7 +27,7 @@ impl super::SyncNotificationHandler for DidClose {
             );
             return Ok(());
         };
-        clear_diagnostics_for_document(session, snapshot.query(), client)?;
+        clear_diagnostics_for_document(snapshot.query(), client)?;
 
         session
             .close_document(&key)

crates/ruff_server/src/server/api/notifications/did_open.rs

@@ -1,5 +1,6 @@
 use crate::TextDocument;
 use crate::server::Result;
+use crate::server::api::LSPResult;
 use crate::server::api::diagnostics::publish_diagnostics_for_document;
 use crate::session::{Client, Session};
 use lsp_types as types;
@@ -29,7 +30,16 @@ impl super::SyncNotificationHandler for DidOpen {
 
         session.open_text_document(uri.clone(), document);
 
-        publish_diagnostics_for_document(session, &uri, client)?;
+        // Publish diagnostics if the client doesn't support pull diagnostics
+        if !session.resolved_client_capabilities().pull_diagnostics {
+            let snapshot = session
+                .take_snapshot(uri.clone())
+                .ok_or_else(|| {
+                    anyhow::anyhow!("Unable to take snapshot for document with URL {uri}")
+                })
+                .with_failure_code(lsp_server::ErrorCode::InternalError)?;
+            publish_diagnostics_for_document(&snapshot, client)?;
+        }
 
         Ok(())
     }

crates/ruff_server/src/server/api/notifications/did_open_notebook.rs

@@ -40,7 +40,10 @@ impl super::SyncNotificationHandler for DidOpenNotebook {
         session.open_notebook_document(uri.clone(), notebook);
 
         // publish diagnostics
-        publish_diagnostics_for_document(session, &uri, client)?;
+        let snapshot = session
+            .take_snapshot(uri)
+            .expect("snapshot should be available");
+        publish_diagnostics_for_document(&snapshot, client)?;
 
         Ok(())
     }