diff --git a/crates/ruff_server/src/server/api/diagnostics.rs b/crates/ruff_server/src/server/api/diagnostics.rs
index 6f8efe47e8..2c8faab6db 100644
--- a/crates/ruff_server/src/server/api/diagnostics.rs
+++ b/crates/ruff_server/src/server/api/diagnostics.rs
@@ -1,4 +1,7 @@
+use lsp_types::Url;
+
 use crate::{
+    Session,
     lint::DiagnosticsMap,
     session::{Client, DocumentQuery, DocumentSnapshot},
 };
@@ -19,10 +22,21 @@ pub(super) fn generate_diagnostics(snapshot: &DocumentSnapshot) -> DiagnosticsMa
 }
 
 pub(super) fn publish_diagnostics_for_document(
-    snapshot: &DocumentSnapshot,
+    session: &Session,
+    url: &Url,
     client: &Client,
 ) -> crate::server::Result<()> {
-    for (uri, diagnostics) in generate_diagnostics(snapshot) {
+    // Publish diagnostics if the client doesn't support pull diagnostics
+    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 {
@@ -38,9 +52,14 @@ 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 {
diff --git a/crates/ruff_server/src/server/api/notifications/did_change.rs b/crates/ruff_server/src/server/api/notifications/did_change.rs
index 8e77cb593f..5ac7a1f606 100644
--- a/crates/ruff_server/src/server/api/notifications/did_change.rs
+++ b/crates/ruff_server/src/server/api/notifications/did_change.rs
@@ -31,11 +31,7 @@ impl super::SyncNotificationHandler for DidChange {
             .update_text_document(&key, content_changes, new_version)
             .with_failure_code(ErrorCode::InternalError)?;
 
-        // 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)?;
-        }
+        publish_diagnostics_for_document(session, &key.into_url(), client)?;
 
         Ok(())
     }
diff --git a/crates/ruff_server/src/server/api/notifications/did_change_notebook.rs b/crates/ruff_server/src/server/api/notifications/did_change_notebook.rs
index d092ccacb8..da11755d71 100644
--- a/crates/ruff_server/src/server/api/notifications/did_change_notebook.rs
+++ b/crates/ruff_server/src/server/api/notifications/did_change_notebook.rs
@@ -27,10 +27,7 @@ impl super::SyncNotificationHandler for DidChangeNotebook {
             .with_failure_code(ErrorCode::InternalError)?;
 
         // publish new diagnostics
-        let snapshot = session
-            .take_snapshot(key.into_url())
-            .expect("snapshot should be available");
-        publish_diagnostics_for_document(&snapshot, client)?;
+        publish_diagnostics_for_document(session, &key.into_url(), client)?;
 
         Ok(())
     }
diff --git a/crates/ruff_server/src/server/api/notifications/did_change_watched_files.rs b/crates/ruff_server/src/server/api/notifications/did_change_watched_files.rs
index bc97231411..cb157d81f9 100644
--- a/crates/ruff_server/src/server/api/notifications/did_change_watched_files.rs
+++ b/crates/ruff_server/src/server/api/notifications/did_change_watched_files.rs
@@ -31,19 +31,13 @@ impl super::SyncNotificationHandler for DidChangeWatchedFiles {
             } else {
                 // publish diagnostics for text documents
                 for url in session.text_document_urls() {
-                    let snapshot = session
-                        .take_snapshot(url.clone())
-                        .expect("snapshot should be available");
-                    publish_diagnostics_for_document(&snapshot, client)?;
+                    publish_diagnostics_for_document(session, url, client)?;
                 }
             }
 
             // always publish diagnostics for notebook files (since they don't use pull diagnostics)
             for url in session.notebook_document_urls() {
-                let snapshot = session
-                    .take_snapshot(url.clone())
-                    .expect("snapshot should be available");
-                publish_diagnostics_for_document(&snapshot, client)?;
+                publish_diagnostics_for_document(session, url, client)?;
             }
         }
 
diff --git a/crates/ruff_server/src/server/api/notifications/did_close.rs b/crates/ruff_server/src/server/api/notifications/did_close.rs
index a3075a4846..5a482c4fcc 100644
--- a/crates/ruff_server/src/server/api/notifications/did_close.rs
+++ b/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(snapshot.query(), client)?;
+        clear_diagnostics_for_document(session, snapshot.query(), client)?;
 
         session
             .close_document(&key)
diff --git a/crates/ruff_server/src/server/api/notifications/did_open.rs b/crates/ruff_server/src/server/api/notifications/did_open.rs
index 41a6fb6cf8..fa5f6b92df 100644
--- a/crates/ruff_server/src/server/api/notifications/did_open.rs
+++ b/crates/ruff_server/src/server/api/notifications/did_open.rs
@@ -1,6 +1,5 @@
 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;
@@ -30,16 +29,7 @@ impl super::SyncNotificationHandler for DidOpen {
 
         session.open_text_document(uri.clone(), document);
 
-        // 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)?;
-        }
+        publish_diagnostics_for_document(session, &uri, client)?;
 
         Ok(())
     }
diff --git a/crates/ruff_server/src/server/api/notifications/did_open_notebook.rs b/crates/ruff_server/src/server/api/notifications/did_open_notebook.rs
index a75e88ecc5..3ce27168e4 100644
--- a/crates/ruff_server/src/server/api/notifications/did_open_notebook.rs
+++ b/crates/ruff_server/src/server/api/notifications/did_open_notebook.rs
@@ -40,10 +40,7 @@ impl super::SyncNotificationHandler for DidOpenNotebook {
         session.open_notebook_document(uri.clone(), notebook);
 
         // publish diagnostics
-        let snapshot = session
-            .take_snapshot(uri)
-            .expect("snapshot should be available");
-        publish_diagnostics_for_document(&snapshot, client)?;
+        publish_diagnostics_for_document(session, &uri, client)?;
 
         Ok(())
     }