[ty] Support publishing diagnostics in the server (#18309)

## Summary

This PR adds support for [publishing
diagnostics](https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_publishDiagnostics)
from the ty language server.

It only adds support for it for text documents and not notebook
documents because the server doesn't have full notebook support yet.

Closes: astral-sh/ty#79

## Test Plan

Testing this out in Helix and Zed since those are the two editors that I
know of that doesn't support pull diagnostics:

### Helix


https://github.com/user-attachments/assets/e193f804-0b32-4f7e-8b83-6f9307e3d2d4



### Zed



https://github.com/user-attachments/assets/93ec7169-ce2b-4521-b009-a82d8afb9eaa

crates/ty_server/src/document/notebook.rs

@@ -199,7 +199,6 @@ impl NotebookDocument {
     }
 
     /// Get the URI for a cell by its index within the cell array.
-    #[expect(dead_code)]
     pub(crate) fn cell_uri_by_index(&self, index: CellId) -> Option<&lsp_types::Url> {
         self.cells.get(index).map(|cell| &cell.url)
     }
@@ -212,7 +211,7 @@ impl NotebookDocument {
     }
 
     /// Returns a list of cell URIs in the order they appear in the array.
-    pub(crate) fn urls(&self) -> impl Iterator<Item = &lsp_types::Url> {
+    pub(crate) fn cell_urls(&self) -> impl Iterator<Item = &lsp_types::Url> {
         self.cells.iter().map(|cell| &cell.url)
     }
 

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

@@ -1,23 +1,51 @@
 use lsp_server::ErrorCode;
 use lsp_types::notification::PublishDiagnostics;
 use lsp_types::{
-    Diagnostic, DiagnosticRelatedInformation, DiagnosticSeverity, DiagnosticTag, NumberOrString,
-    PublishDiagnosticsParams, Range, Url,
+    CodeDescription, Diagnostic, DiagnosticRelatedInformation, DiagnosticSeverity, DiagnosticTag,
+    NumberOrString, PublishDiagnosticsParams, Range, Url,
 };
+use rustc_hash::FxHashMap;
 
 use ruff_db::diagnostic::{Annotation, Severity, SubDiagnostic};
 use ruff_db::files::FileRange;
 use ruff_db::source::{line_index, source_text};
 use ty_project::{Db, ProjectDatabase};
 
-use crate::DocumentSnapshot;
-use crate::PositionEncoding;
 use crate::document::{FileRangeExt, ToRangeExt};
 use crate::server::Result;
 use crate::server::client::Notifier;
+use crate::system::url_to_any_system_path;
+use crate::{DocumentSnapshot, PositionEncoding, Session};
 
 use super::LSPResult;
 
+/// Represents the diagnostics for a text document or a notebook document.
+pub(super) enum Diagnostics {
+    TextDocument(Vec<Diagnostic>),
+
+    /// A map of cell URLs to the diagnostics for that cell.
+    NotebookDocument(FxHashMap<Url, Vec<Diagnostic>>),
+}
+
+impl Diagnostics {
+    /// Returns the diagnostics for a text document.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the diagnostics are for a notebook document.
+    pub(super) fn expect_text_document(self) -> Vec<Diagnostic> {
+        match self {
+            Diagnostics::TextDocument(diagnostics) => diagnostics,
+            Diagnostics::NotebookDocument(_) => {
+                panic!("Expected a text document diagnostics, but got notebook diagnostics")
+            }
+        }
+    }
+}
+
+/// Clears the diagnostics for the document at `uri`.
+///
+/// This is done by notifying the client with an empty list of diagnostics for the document.
 pub(super) fn clear_diagnostics(uri: &Url, notifier: &Notifier) -> Result<()> {
     notifier
         .notify::<PublishDiagnostics>(PublishDiagnosticsParams {
@@ -29,25 +57,106 @@ pub(super) fn clear_diagnostics(uri: &Url, notifier: &Notifier) -> Result<()> {
     Ok(())
 }
 
+/// Publishes the diagnostics for the given document snapshot using the [publish diagnostics
+/// notification].
+///
+/// This function is a no-op if the client supports pull diagnostics.
+///
+/// [publish diagnostics notification]: https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_publishDiagnostics
+pub(super) fn publish_diagnostics(session: &Session, url: Url, notifier: &Notifier) -> Result<()> {
+    if session.client_capabilities().pull_diagnostics {
+        return Ok(());
+    }
+
+    let Ok(path) = url_to_any_system_path(&url) else {
+        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)?;
+
+    let db = session.project_db_or_default(&path);
+
+    let Some(diagnostics) = compute_diagnostics(db, &snapshot) else {
+        return Ok(());
+    };
+
+    // Sends a notification to the client with the diagnostics for the document.
+    let publish_diagnostics_notification = |uri: Url, diagnostics: Vec<Diagnostic>| {
+        notifier
+            .notify::<PublishDiagnostics>(PublishDiagnosticsParams {
+                uri,
+                diagnostics,
+                version: Some(snapshot.query().version()),
+            })
+            .with_failure_code(lsp_server::ErrorCode::InternalError)
+    };
+
+    match diagnostics {
+        Diagnostics::TextDocument(diagnostics) => {
+            publish_diagnostics_notification(url, diagnostics)?;
+        }
+        Diagnostics::NotebookDocument(cell_diagnostics) => {
+            for (cell_url, diagnostics) in cell_diagnostics {
+                publish_diagnostics_notification(cell_url, diagnostics)?;
+            }
+        }
+    }
+
+    Ok(())
+}
+
 pub(super) fn compute_diagnostics(
     db: &ProjectDatabase,
     snapshot: &DocumentSnapshot,
-) -> Vec<Diagnostic> {
+) -> Option<Diagnostics> {
     let Some(file) = snapshot.file(db) else {
         tracing::info!(
             "No file found for snapshot for `{}`",
             snapshot.query().file_url()
         );
-        return vec![];
+        return None;
     };
 
     let diagnostics = db.check_file(file);
 
-    diagnostics
-        .as_slice()
-        .iter()
-        .map(|message| to_lsp_diagnostic(db, message, snapshot.encoding()))
-        .collect()
+    if let Some(notebook) = snapshot.query().as_notebook() {
+        let mut cell_diagnostics: FxHashMap<Url, Vec<Diagnostic>> = FxHashMap::default();
+
+        // Populates all relevant URLs with an empty diagnostic list. This ensures that documents
+        // without diagnostics still get updated.
+        for cell_url in notebook.cell_urls() {
+            cell_diagnostics.entry(cell_url.clone()).or_default();
+        }
+
+        for (cell_index, diagnostic) in diagnostics.iter().map(|diagnostic| {
+            (
+                // TODO: Use the cell index instead using `SourceKind`
+                usize::default(),
+                to_lsp_diagnostic(db, diagnostic, snapshot.encoding()),
+            )
+        }) {
+            let Some(cell_uri) = notebook.cell_uri_by_index(cell_index) else {
+                tracing::warn!("Unable to find notebook cell at index {cell_index}");
+                continue;
+            };
+            cell_diagnostics
+                .entry(cell_uri.clone())
+                .or_default()
+                .push(diagnostic);
+        }
+
+        Some(Diagnostics::NotebookDocument(cell_diagnostics))
+    } else {
+        Some(Diagnostics::TextDocument(
+            diagnostics
+                .iter()
+                .map(|diagnostic| to_lsp_diagnostic(db, diagnostic, snapshot.encoding()))
+                .collect(),
+        ))
+    }
 }
 
 /// Converts the tool specific [`Diagnostic`][ruff_db::diagnostic::Diagnostic] to an LSP
@@ -91,9 +200,8 @@ fn to_lsp_diagnostic(
         .id()
         .is_lint()
         .then(|| {
-            Some(lsp_types::CodeDescription {
-                href: lsp_types::Url::parse(&format!("https://ty.dev/rules#{}", diagnostic.id()))
-                    .ok()?,
+            Some(CodeDescription {
+                href: Url::parse(&format!("https://ty.dev/rules#{}", diagnostic.id())).ok()?,
             })
         })
         .flatten();

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

@@ -1,11 +1,12 @@
 use lsp_server::ErrorCode;
-use lsp_types::DidChangeTextDocumentParams;
 use lsp_types::notification::DidChangeTextDocument;
+use lsp_types::{DidChangeTextDocumentParams, VersionedTextDocumentIdentifier};
 
 use ty_project::watch::ChangeEvent;
 
 use crate::server::Result;
 use crate::server::api::LSPResult;
+use crate::server::api::diagnostics::publish_diagnostics;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::server::client::{Notifier, Requester};
 use crate::session::Session;
@@ -20,18 +21,23 @@ impl NotificationHandler for DidChangeTextDocumentHandler {
 impl SyncNotificationHandler for DidChangeTextDocumentHandler {
     fn run(
         session: &mut Session,
-        _notifier: Notifier,
+        notifier: Notifier,
         _requester: &mut Requester,
         params: DidChangeTextDocumentParams,
     ) -> Result<()> {
-        let Ok(path) = url_to_any_system_path(&params.text_document.uri) else {
+        let DidChangeTextDocumentParams {
+            text_document: VersionedTextDocumentIdentifier { uri, version },
+            content_changes,
+        } = params;
+
+        let Ok(path) = url_to_any_system_path(&uri) else {
             return Ok(());
         };
 
-        let key = session.key_from_url(params.text_document.uri);
+        let key = session.key_from_url(uri.clone());
 
         session
-            .update_text_document(&key, params.content_changes, params.text_document.version)
+            .update_text_document(&key, content_changes, version)
             .with_failure_code(ErrorCode::InternalError)?;
 
         match path {
@@ -48,8 +54,6 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
             }
         }
 
-        // TODO(dhruvmanila): Publish diagnostics if the client doesn't support pull diagnostics
-
-        Ok(())
+        publish_diagnostics(session, uri, &notifier)
     }
 }

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

@@ -1,5 +1,6 @@
 use crate::server::Result;
 use crate::server::api::LSPResult;
+use crate::server::api::diagnostics::publish_diagnostics;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::server::client::{Notifier, Requester};
 use crate::server::schedule::Task;
@@ -20,7 +21,7 @@ impl NotificationHandler for DidChangeWatchedFiles {
 impl SyncNotificationHandler for DidChangeWatchedFiles {
     fn run(
         session: &mut Session,
-        _notifier: Notifier,
+        notifier: Notifier,
         requester: &mut Requester,
         params: types::DidChangeWatchedFilesParams,
     ) -> Result<()> {
@@ -85,19 +86,35 @@ impl SyncNotificationHandler for DidChangeWatchedFiles {
             return Ok(());
         }
 
+        let mut project_changed = false;
+
         for (root, changes) in events_by_db {
             tracing::debug!("Applying changes to `{root}`");
+
+            // SAFETY: Only paths that are part of the workspace are registered for file watching.
+            // So, virtual paths and paths that are outside of a workspace does not trigger this
+            // notification.
             let db = session.project_db_for_path_mut(&*root).unwrap();
 
-            db.apply_changes(changes, None);
+            let result = db.apply_changes(changes, None);
+
+            project_changed |= result.project_changed();
         }
 
         let client_capabilities = session.client_capabilities();
 
-        if client_capabilities.diagnostics_refresh {
-            requester
-                .request::<types::request::WorkspaceDiagnosticRefresh>((), |()| Task::nothing())
-                .with_failure_code(lsp_server::ErrorCode::InternalError)?;
+        if project_changed {
+            if client_capabilities.diagnostics_refresh {
+                requester
+                    .request::<types::request::WorkspaceDiagnosticRefresh>((), |()| Task::nothing())
+                    .with_failure_code(lsp_server::ErrorCode::InternalError)?;
+            } else {
+                for url in session.text_document_urls() {
+                    publish_diagnostics(session, url.clone(), &notifier)?;
+                }
+            }
+
+            // TODO: always publish diagnostics for notebook files (since they don't use pull diagnostics)
         }
 
         if client_capabilities.inlay_refresh {

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

@@ -6,6 +6,7 @@ use ty_project::watch::ChangeEvent;
 
 use crate::TextDocument;
 use crate::server::Result;
+use crate::server::api::diagnostics::publish_diagnostics;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::server::client::{Notifier, Requester};
 use crate::session::Session;
@@ -20,7 +21,7 @@ impl NotificationHandler for DidOpenTextDocumentHandler {
 impl SyncNotificationHandler for DidOpenTextDocumentHandler {
     fn run(
         session: &mut Session,
-        _notifier: Notifier,
+        notifier: Notifier,
         _requester: &mut Requester,
         DidOpenTextDocumentParams {
             text_document:
@@ -37,24 +38,22 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
         };
 
         let document = TextDocument::new(text, version).with_language_id(&language_id);
-        session.open_text_document(uri, document);
+        session.open_text_document(uri.clone(), document);
 
-        match path {
+        match &path {
             AnySystemPath::System(path) => {
                 let db = match session.project_db_for_path_mut(path.as_std_path()) {
                     Some(db) => db,
                     None => session.default_project_db_mut(),
                 };
-                db.apply_changes(vec![ChangeEvent::Opened(path)], None);
+                db.apply_changes(vec![ChangeEvent::Opened(path.clone())], None);
             }
             AnySystemPath::SystemVirtual(virtual_path) => {
                 let db = session.default_project_db_mut();
-                db.files().virtual_file(db, &virtual_path);
+                db.files().virtual_file(db, virtual_path);
             }
         }
 
-        // TODO(dhruvmanila): Publish diagnostics if the client doesn't support pull diagnostics
-
-        Ok(())
+        publish_diagnostics(session, uri, &notifier)
     }
 }

crates/ty_server/src/server/api/requests/diagnostic.rs

@@ -6,7 +6,7 @@ use lsp_types::{
     FullDocumentDiagnosticReport, RelatedFullDocumentDiagnosticReport,
 };
 
-use crate::server::api::diagnostics::compute_diagnostics;
+use crate::server::api::diagnostics::{Diagnostics, compute_diagnostics};
 use crate::server::api::traits::{BackgroundDocumentRequestHandler, RequestHandler};
 use crate::server::{Result, client::Notifier};
 use crate::session::DocumentSnapshot;
@@ -29,14 +29,15 @@ impl BackgroundDocumentRequestHandler for DocumentDiagnosticRequestHandler {
         _notifier: Notifier,
         _params: DocumentDiagnosticParams,
     ) -> Result<DocumentDiagnosticReportResult> {
-        let diagnostics = compute_diagnostics(db, &snapshot);
-
         Ok(DocumentDiagnosticReportResult::Report(
             DocumentDiagnosticReport::Full(RelatedFullDocumentDiagnosticReport {
                 related_documents: None,
                 full_document_diagnostic_report: FullDocumentDiagnosticReport {
                     result_id: None,
-                    items: diagnostics,
+                    // SAFETY: Pull diagnostic requests are only called for text documents, not for
+                    // notebook documents.
+                    items: compute_diagnostics(db, &snapshot)
+                        .map_or_else(Vec::new, Diagnostics::expect_text_document),
                 },
             }),
         ))

crates/ty_server/src/session.rs

@@ -46,6 +46,7 @@ pub struct Session {
 
     /// The global position encoding, negotiated during LSP initialization.
     position_encoding: PositionEncoding,
+
     /// Tracks what LSP features the client supports and doesn't support.
     resolved_client_capabilities: Arc<ResolvedClientCapabilities>,
 }
@@ -90,6 +91,14 @@ impl Session {
     // and `default_workspace_db_mut` but the borrow checker doesn't allow that.
     // https://github.com/astral-sh/ruff/pull/13041#discussion_r1726725437
 
+    /// Returns a reference to the project's [`ProjectDatabase`] corresponding to the given path,
+    /// or the default project if no project is found for the path.
+    pub(crate) fn project_db_or_default(&self, path: &AnySystemPath) -> &ProjectDatabase {
+        path.as_system()
+            .and_then(|path| self.project_db_for_path(path.as_std_path()))
+            .unwrap_or_else(|| self.default_project_db())
+    }
+
     /// Returns a reference to the project's [`ProjectDatabase`] corresponding to the given path, if
     /// any.
     pub(crate) fn project_db_for_path(&self, path: impl AsRef<Path>) -> Option<&ProjectDatabase> {
@@ -141,6 +150,11 @@ impl Session {
         })
     }
 
+    /// Iterates over the LSP URLs for all open text documents. These URLs are valid file paths.
+    pub(super) fn text_document_urls(&self) -> impl Iterator<Item = &Url> + '_ {
+        self.index().text_document_urls()
+    }
+
     /// Registers a notebook document at the provided `url`.
     /// If a document is already open here, it will be overwritten.
     pub fn open_notebook_document(&mut self, url: Url, document: NotebookDocument) {

crates/ty_server/src/session/capabilities.rs

@@ -8,7 +8,12 @@ pub(crate) struct ResolvedClientCapabilities {
     pub(crate) document_changes: bool,
     pub(crate) diagnostics_refresh: bool,
     pub(crate) inlay_refresh: bool,
+
+    /// Whether [pull diagnostics] is supported.
+    ///
+    /// [pull diagnostics]: https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_pullDiagnostics
     pub(crate) pull_diagnostics: bool,
+
     /// Whether `textDocument.typeDefinition.linkSupport` is `true`
     pub(crate) type_definition_link_support: bool,
 

crates/ty_server/src/session/index.rs

@@ -34,7 +34,6 @@ impl Index {
         }
     }
 
-    #[expect(dead_code)]
     pub(super) fn text_document_urls(&self) -> impl Iterator<Item = &Url> + '_ {
         self.documents
             .iter()
@@ -135,7 +134,7 @@ impl Index {
     }
 
     pub(super) fn open_notebook_document(&mut self, notebook_url: Url, document: NotebookDocument) {
-        for cell_url in document.urls() {
+        for cell_url in document.cell_urls() {
             self.notebook_cells
                 .insert(cell_url.clone(), notebook_url.clone());
         }

crates/ty_server/src/system.rs

@@ -47,12 +47,21 @@ pub(crate) fn file_to_url(db: &dyn Db, file: File) -> Option<Url> {
 }
 
 /// Represents either a [`SystemPath`] or a [`SystemVirtualPath`].
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub(crate) enum AnySystemPath {
     System(SystemPathBuf),
     SystemVirtual(SystemVirtualPathBuf),
 }
 
+impl AnySystemPath {
+    pub(crate) const fn as_system(&self) -> Option<&SystemPathBuf> {
+        match self {
+            AnySystemPath::System(system_path_buf) => Some(system_path_buf),
+            AnySystemPath::SystemVirtual(_) => None,
+        }
+    }
+}
+
 #[derive(Debug)]
 pub(crate) struct LSPSystem {
     /// A read-only copy of the index where the server stores all the open documents and settings.