[ty] Support publishing diagnostics in the server

dhruvmanila · dhruvmanila · commit 4ed53ac39601 · 2025-05-26T09:53:53.000+05:30
diff --git a/crates/ty_server/src/document.rs b/crates/ty_server/src/document.rs
@@ -55,6 +55,15 @@ impl DocumentKey {
             | DocumentKey::Text(url) => url,
         }
     }
+
+    /// Converts the key back into its original URL.
+    pub(crate) fn into_url(self) -> Url {
+        match self {
+            DocumentKey::NotebookCell(url)
+            | DocumentKey::Notebook(url)
+            | DocumentKey::Text(url) => url,
+        }
+    }
 }
 
 impl std::fmt::Display for DocumentKey {
diff --git a/crates/ty_server/src/document/notebook.rs b/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)
     }
diff --git a/crates/ty_server/src/server/api/diagnostics.rs b/crates/ty_server/src/server/api/diagnostics.rs
@@ -1,23 +1,29 @@
 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::{DocumentSnapshot, PositionEncoding};
 
 use super::LSPResult;
 
+/// A series of diagnostics across a single text document or an arbitrary number of notebook cells.
+pub(super) type DiagnosticsMap = FxHashMap<Url, Vec<Diagnostic>>;
+
+/// 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,31 +35,91 @@ pub(super) fn clear_diagnostics(uri: &Url, notifier: &Notifier) -> Result<()> {
     Ok(())
 }
 
+/// Publishes the diagnostics for the given document snapshot using the [publish diagnostics
+/// notification].
+///
+/// [publish diagnostics notification]: https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_publishDiagnostics
+pub(super) fn publish_diagnostics_for_document(
+    db: &ProjectDatabase,
+    snapshot: &DocumentSnapshot,
+    notifier: &Notifier,
+) -> Result<()> {
+    for (uri, diagnostics) in compute_diagnostics(db, snapshot) {
+        notifier
+            .notify::<PublishDiagnostics>(PublishDiagnosticsParams {
+                uri,
+                diagnostics,
+                version: Some(snapshot.query().version()),
+            })
+            .with_failure_code(lsp_server::ErrorCode::InternalError)?;
+    }
+
+    Ok(())
+}
+
 pub(super) fn compute_diagnostics(
     db: &ProjectDatabase,
     snapshot: &DocumentSnapshot,
-) -> Vec<Diagnostic> {
+) -> DiagnosticsMap {
     let Some(file) = snapshot.file(db) else {
         tracing::info!(
             "No file found for snapshot for `{}`",
             snapshot.query().file_url()
         );
-        return vec![];
+        return DiagnosticsMap::default();
     };
 
     let diagnostics = match db.check_file(file) {
         Ok(diagnostics) => diagnostics,
         Err(cancelled) => {
             tracing::info!("Diagnostics computation {cancelled}");
-            return vec![];
+            return DiagnosticsMap::default();
         }
     };
 
-    diagnostics
-        .as_slice()
-        .iter()
-        .map(|message| to_lsp_diagnostic(db, message, snapshot.encoding()))
-        .collect()
+    let mut diagnostics_map = DiagnosticsMap::default();
+    let query = snapshot.query();
+    // let source_kind = query.make_source_kind();
+
+    // Populates all relevant URLs with an empty diagnostic list.
+    // This ensures that documents without diagnostics still get updated.
+    if let Some(notebook) = query.as_notebook() {
+        for url in notebook.urls() {
+            diagnostics_map.entry(url.clone()).or_default();
+        }
+    } else {
+        diagnostics_map
+            .entry(query.make_key().into_url())
+            .or_default();
+    }
+
+    let lsp_diagnostics = diagnostics.as_slice().iter().map(|diagnostic| {
+        (
+            // TODO: Use the cell index instead using `source_kind`
+            usize::default(),
+            to_lsp_diagnostic(db, diagnostic, snapshot.encoding()),
+        )
+    });
+
+    if let Some(notebook) = query.as_notebook() {
+        for (index, diagnostic) in lsp_diagnostics {
+            let Some(uri) = notebook.cell_uri_by_index(index) else {
+                tracing::warn!("Unable to find notebook cell at index {index}");
+                continue;
+            };
+            diagnostics_map
+                .entry(uri.clone())
+                .or_default()
+                .push(diagnostic);
+        }
+    } else {
+        diagnostics_map
+            .entry(query.make_key().into_url())
+            .or_default()
+            .extend(lsp_diagnostics.map(|(_, diagnostic)| diagnostic));
+    }
+
+    diagnostics_map
 }
 
 /// Converts the tool specific [`Diagnostic`][ruff_db::diagnostic::Diagnostic] to an LSP
@@ -97,9 +163,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();
diff --git a/crates/ty_server/src/server/api/notifications/did_change.rs b/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_for_document;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::server::client::{Notifier, Requester};
 use crate::session::Session;
@@ -20,21 +21,26 @@ 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 {
+        match path.clone() {
             AnySystemPath::System(path) => {
                 let db = match session.project_db_for_path_mut(path.as_std_path()) {
                     Some(db) => db,
@@ -48,7 +54,20 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
             }
         }
 
-        // TODO(dhruvmanila): Publish diagnostics if the client doesn't support pull diagnostics
+        // Publish diagnostics if the client doesn't support pull diagnostics
+        if !session.resolved_client_capabilities().pull_diagnostics {
+            let db = path
+                .as_system()
+                .and_then(|path| session.project_db_for_path(path.as_std_path()))
+                .unwrap_or_else(|| session.default_project_db());
+            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(db, &snapshot, &notifier)?;
+        }
 
         Ok(())
     }
diff --git a/crates/ty_server/src/server/api/notifications/did_open.rs b/crates/ty_server/src/server/api/notifications/did_open.rs
@@ -6,6 +6,8 @@ use ty_project::watch::ChangeEvent;
 
 use crate::TextDocument;
 use crate::server::Result;
+use crate::server::api::LSPResult;
+use crate::server::api::diagnostics::publish_diagnostics_for_document;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::server::client::{Notifier, Requester};
 use crate::session::Session;
@@ -20,7 +22,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,23 +39,36 @@ 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
+        // Publish diagnostics if the client doesn't support pull diagnostics
+        if !session.resolved_client_capabilities().pull_diagnostics {
+            let db = path
+                .as_system()
+                .and_then(|path| session.project_db_for_path(path.as_std_path()))
+                .unwrap_or_else(|| session.default_project_db());
+            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(db, &snapshot, &notifier)?;
+        }
 
         Ok(())
     }
diff --git a/crates/ty_server/src/server/api/requests/diagnostic.rs b/crates/ty_server/src/server/api/requests/diagnostic.rs
@@ -29,14 +29,18 @@ 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,
+                    // Pull diagnostic requests are only called for text documents.
+                    // Since diagnostic requests generate
+                    items: compute_diagnostics(db, &snapshot)
+                        .into_iter()
+                        .next()
+                        .map(|(_, diagnostics)| diagnostics)
+                        .unwrap_or_default(),
                 },
             }),
         ))
diff --git a/crates/ty_server/src/session.rs b/crates/ty_server/src/session.rs
@@ -86,6 +86,10 @@ impl Session {
         })
     }
 
+    pub(crate) fn resolved_client_capabilities(&self) -> &ResolvedClientCapabilities {
+        &self.resolved_client_capabilities
+    }
+
     // TODO(dhruvmanila): Ideally, we should have a single method for `workspace_db_for_path_mut`
     // and `default_workspace_db_mut` but the borrow checker doesn't allow that.
     // https://github.com/astral-sh/ruff/pull/13041#discussion_r1726725437
diff --git a/crates/ty_server/src/session/capabilities.rs b/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,
 
diff --git a/crates/ty_server/src/session/index.rs b/crates/ty_server/src/session/index.rs
@@ -265,7 +265,6 @@ pub enum DocumentQuery {
 
 impl DocumentQuery {
     /// Retrieve the original key that describes this document query.
-    #[expect(dead_code)]
     pub(crate) fn make_key(&self) -> DocumentKey {
         match self {
             Self::Text { file_url, .. } => DocumentKey::Text(file_url.clone()),
@@ -300,6 +299,16 @@ impl DocumentQuery {
         }
     }
 
+    // /// Generate a source kind used by the linter.
+    // pub(crate) fn make_source_kind(&self) -> SourceKind {
+    //     match self {
+    //         Self::Text { document, .. } => SourceKind::Python(document.contents().to_string()),
+    //         Self::Notebook { notebook, .. } => {
+    //             SourceKind::ipy_notebook(notebook.make_ruff_notebook())
+    //         }
+    //     }
+    // }
+
     /// Attempt to access the single inner text document selected by the query.
     /// If this query is selecting an entire notebook document, this will return `None`.
     #[expect(dead_code)]
diff --git a/crates/ty_server/src/system.rs b/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.

Original file line number	Diff line number	Diff line change
`@@ -55,6 +55,15 @@ impl DocumentKey {`
`55`	`55`	`\| DocumentKey::Text(url) => url,`
`56`	`56`	`}`
`57`	`57`	`}`
	`58`	`+`
	`59`	`+ /// Converts the key back into its original URL.`
	`60`	`+ pub(crate) fn into_url(self) -> Url {`
	`61`	`+ match self {`
	`62`	`+ DocumentKey::NotebookCell(url)`
	`63`	`+ \| DocumentKey::Notebook(url)`
	`64`	`+ \| DocumentKey::Text(url) => url,`
	`65`	`+ }`
	`66`	`+ }`
`58`	`67`	`}`
`59`	`68`
`60`	`69`	`impl std::fmt::Display for DocumentKey {`
Original file line number	Diff line number	Diff line change
`@@ -199,7 +199,6 @@ impl NotebookDocument {`
`199`	`199`	`}`
`200`	`200`
`201`	`201`	`/// Get the URI for a cell by its index within the cell array.`
`202`		`- #[expect(dead_code)]`
`203`	`202`	`pub(crate) fn cell_uri_by_index(&self, index: CellId) -> Option<&lsp_types::Url> {`
`204`	`203`	`self.cells.get(index).map(\|cell\| &cell.url)`
`205`	`204`	`}`