Add notebook support to LSP server

Implement comprehensive notebook document support:

- Refactor NotebookDocument to simplify cell management
  - Remove embedded TextDocuments, defer to index
  - Add to_ruff_notebook() for converting to linter format

- Implement LSP notification handlers for notebook lifecycle:
  - did_open_notebook: Open notebook and cell documents
  - did_close_notebook: Close notebook and cell documents
  - did_change_notebook: Handle cell structure and content changes

- Update all LSP request handlers to use new conversion APIs:
  - Navigation (goto definition, references, declaration, type def)
  - Information (hover, completion, inlay hints, signature help)
  - Symbols (document and workspace symbols)
  - All handlers now use to_local_range()/to_location() appropriately

- Add notebook diagnostics support:
  - Per-cell diagnostic publishing
  - Workspace diagnostics for notebooks

- Update session and index management:
  - Track notebook documents and their cells
  - Map cell URIs to parent notebooks
  - Handle notebook file changes

- Extend system path handling for notebook URIs
- Enable notebook capabilities in server initialization
- Add comprehensive E2E tests for notebooks

This enables the LSP server to work with Jupyter notebooks by treating
each cell as a separate text document while maintaining the notebook
context for type checking and cross-cell references.

Discard changes to crates/ty_ide/src/importer.rs

Discard changes to crates/ty_server/src/server/api/requests/diagnostic.rs

Discard changes to crates/ty_server/src/server/api/semantic_tokens.rs

Author: MichaReiser

crates/ruff_db/src/source.rs

@@ -7,6 +7,7 @@ use ruff_source_file::LineIndex;
 
 use crate::Db;
 use crate::files::{File, FilePath};
+use crate::system::System;
 
 /// Reads the source text of a python text file (must be valid UTF8) or notebook.
 #[salsa::tracked(heap_size=ruff_memory_usage::heap_size)]
@@ -15,7 +16,7 @@ pub fn source_text(db: &dyn Db, file: File) -> SourceText {
     let _span = tracing::trace_span!("source_text", file = %path).entered();
     let mut read_error = None;
 
-    let kind = if is_notebook(file.path(db)) {
+    let kind = if is_notebook(db.system(), path) {
         file.read_to_notebook(db)
             .unwrap_or_else(|error| {
                 tracing::debug!("Failed to read notebook '{path}': {error}");
@@ -40,17 +41,19 @@ pub fn source_text(db: &dyn Db, file: File) -> SourceText {
     }
 }
 
-fn is_notebook(path: &FilePath) -> bool {
-    match path {
-        FilePath::System(system) => system.extension().is_some_and(|extension| {
+fn is_notebook(system: &dyn System, path: &FilePath) -> bool {
+    let source_type = match path {
+        FilePath::System(path) => system.source_type(path),
+        FilePath::SystemVirtual(system_virtual) => system.virtual_path_source_type(system_virtual),
+        FilePath::Vendored(_) => return false,
+    };
+
+    if let Some(source_type) = source_type {
+        source_type.is_ipynb()
+    } else {
+        path.extension().is_some_and(|extension| {
             PySourceType::try_from_extension(extension) == Some(PySourceType::Ipynb)
-        }),
-        FilePath::SystemVirtual(system_virtual) => {
-            system_virtual.extension().is_some_and(|extension| {
-                PySourceType::try_from_extension(extension) == Some(PySourceType::Ipynb)
-            })
-        }
-        FilePath::Vendored(_) => false,
+        })
     }
 }
 

crates/ruff_db/src/system.rs

@@ -9,19 +9,19 @@ pub use os::OsSystem;
 
 use filetime::FileTime;
 use ruff_notebook::{Notebook, NotebookError};
+use ruff_python_ast::PySourceType;
 use std::error::Error;
 use std::fmt::{Debug, Formatter};
 use std::path::{Path, PathBuf};
 use std::{fmt, io};
 pub use test::{DbWithTestSystem, DbWithWritableSystem, InMemorySystem, TestSystem};
 use walk_directory::WalkDirectoryBuilder;
 
-use crate::file_revision::FileRevision;
-
 pub use self::path::{
     DeduplicatedNestedPathsIter, SystemPath, SystemPathBuf, SystemVirtualPath,
     SystemVirtualPathBuf, deduplicate_nested_paths,
 };
+use crate::file_revision::FileRevision;
 
 mod memory_fs;
 #[cfg(feature = "os")]
@@ -66,6 +66,27 @@ pub trait System: Debug + Sync + Send {
     /// See [dunce::canonicalize] for more information.
     fn canonicalize_path(&self, path: &SystemPath) -> Result<SystemPathBuf>;
 
+    /// Returns the source type for `path` if known or `None`.
+    ///
+    /// This is primarily used for the LSP integration to respect
+    /// the chosen language (or the fact that it is a notebook) in
+    /// the editor.
+    fn source_type(&self, path: &SystemPath) -> Option<PySourceType> {
+        let _ = path;
+        None
+    }
+
+    /// Returns the source type for `path` if known or `None`.
+    ///
+    /// This is primarily used for the LSP integration to respect
+    /// the chosen language (or the fact that it is a notebook) in
+    /// the editor.
+    fn virtual_path_source_type(&self, path: &SystemVirtualPath) -> Option<PySourceType> {
+        let _ = path;
+
+        None
+    }
+
     /// Reads the content of the file at `path` into a [`String`].
     fn read_to_string(&self, path: &SystemPath) -> Result<String>;
 

crates/ruff_notebook/src/index.rs

@@ -47,7 +47,7 @@ impl NotebookIndex {
     /// Translates the given [`LineColumn`] based on the indexing table.
     ///
     /// This will translate the row/column in the concatenated source code
-    /// to the row/column in the Jupyter Notebook.
+    /// to the row/column in the Jupyter Notebook cell.
     pub fn translate_line_column(&self, source_location: &LineColumn) -> LineColumn {
         LineColumn {
             line: self
@@ -60,7 +60,7 @@ impl NotebookIndex {
     /// Translates the given [`SourceLocation`] based on the indexing table.
     ///
     /// This will translate the line/character in the concatenated source code
-    /// to the line/character in the Jupyter Notebook.
+    /// to the line/character in the Jupyter Notebook cell.
     pub fn translate_source_location(&self, source_location: &SourceLocation) -> SourceLocation {
         SourceLocation {
             line: self

crates/ruff_notebook/src/notebook.rs

@@ -13,7 +13,7 @@ use thiserror::Error;
 
 use ruff_diagnostics::{SourceMap, SourceMarker};
 use ruff_source_file::{NewlineWithTrailingNewline, OneIndexed, UniversalNewlineIterator};
-use ruff_text_size::TextSize;
+use ruff_text_size::{TextRange, TextSize};
 
 use crate::cell::CellOffsets;
 use crate::index::NotebookIndex;
@@ -294,7 +294,7 @@ impl Notebook {
         }
     }
 
-    /// Build and return the [`JupyterIndex`].
+    /// Build and return the [`NotebookIndex`].
     ///
     /// ## Notes
     ///
@@ -388,6 +388,21 @@ impl Notebook {
         &self.cell_offsets
     }
 
+    /// Returns the start offset of the cell at index `cell` in the concatenated
+    /// text document.
+    pub fn cell_offset(&self, cell: OneIndexed) -> Option<TextSize> {
+        self.cell_offsets.get(cell.to_zero_indexed()).copied()
+    }
+
+    /// Returns the text range in the concatenated document of the cell
+    /// with index `cell`.
+    pub fn cell_range(&self, cell: OneIndexed) -> Option<TextRange> {
+        let start = self.cell_offsets.get(cell.to_zero_indexed()).copied()?;
+        let end = self.cell_offsets.get(cell.to_zero_indexed() + 1).copied()?;
+
+        Some(TextRange::new(start, end))
+    }
+
     /// Return `true` if the notebook has a trailing newline, `false` otherwise.
     pub fn trailing_newline(&self) -> bool {
         self.trailing_newline

crates/ty_server/src/capabilities.rs

@@ -1,10 +1,10 @@
 use lsp_types::{
     ClientCapabilities, CompletionOptions, DeclarationCapability, DiagnosticOptions,
     DiagnosticServerCapabilities, HoverProviderCapability, InlayHintOptions,
-    InlayHintServerCapabilities, MarkupKind, OneOf, RenameOptions,
-    SelectionRangeProviderCapability, SemanticTokensFullOptions, SemanticTokensLegend,
-    SemanticTokensOptions, SemanticTokensServerCapabilities, ServerCapabilities,
-    SignatureHelpOptions, TextDocumentSyncCapability, TextDocumentSyncKind,
+    InlayHintServerCapabilities, MarkupKind, NotebookCellSelector, NotebookSelector, OneOf,
+    RenameOptions, SelectionRangeProviderCapability, SemanticTokensFullOptions,
+    SemanticTokensLegend, SemanticTokensOptions, SemanticTokensServerCapabilities,
+    ServerCapabilities, SignatureHelpOptions, TextDocumentSyncCapability, TextDocumentSyncKind,
     TextDocumentSyncOptions, TypeDefinitionProviderCapability, WorkDoneProgressOptions,
 };
 
@@ -422,6 +422,16 @@ pub(crate) fn server_capabilities(
         selection_range_provider: Some(SelectionRangeProviderCapability::Simple(true)),
         document_symbol_provider: Some(OneOf::Left(true)),
         workspace_symbol_provider: Some(OneOf::Left(true)),
+        notebook_document_sync: Some(OneOf::Left(lsp_types::NotebookDocumentSyncOptions {
+            save: Some(false),
+            notebook_selector: [NotebookSelector::ByCells {
+                notebook: None,
+                cells: vec![NotebookCellSelector {
+                    language: "python".to_string(),
+                }],
+            }]
+            .to_vec(),
+        })),
         ..Default::default()
     }
 }

crates/ty_server/src/document.rs

@@ -5,15 +5,15 @@ mod notebook;
 mod range;
 mod text_document;
 
-pub(crate) use location::ToLink;
 use lsp_types::{PositionEncodingKind, Url};
+use ruff_db::system::{SystemPathBuf, SystemVirtualPath, SystemVirtualPathBuf};
 
 use crate::system::AnySystemPath;
+pub(crate) use location::ToLink;
 pub use notebook::NotebookDocument;
 pub(crate) use range::{FileRangeExt, PositionExt, RangeExt, TextSizeExt, ToRangeExt};
-use ruff_db::system::{SystemPathBuf, SystemVirtualPath};
-pub(crate) use text_document::DocumentVersion;
 pub use text_document::TextDocument;
+pub(crate) use text_document::{DocumentVersion, LanguageId};
 
 /// A convenient enumeration for supported text encodings. Can be converted to [`lsp_types::PositionEncodingKind`].
 // Please maintain the order from least to greatest priority for the derived `Ord` impl.
@@ -84,13 +84,6 @@ impl DocumentKey {
         }
     }
 
-    pub(crate) fn as_opaque(&self) -> Option<&str> {
-        match self {
-            Self::Opaque(uri) => Some(uri),
-            Self::File(_) => None,
-        }
-    }
-
     /// Returns the corresponding [`AnySystemPath`] for this document key.
     ///
     /// Note, calling this method on a `DocumentKey::Opaque` representing a cell document
@@ -104,6 +97,13 @@ impl DocumentKey {
             }
         }
     }
+
+    pub(super) fn into_file_path(self) -> AnySystemPath {
+        match self {
+            Self::File(path) => AnySystemPath::System(path),
+            Self::Opaque(uri) => AnySystemPath::SystemVirtual(SystemVirtualPathBuf::from(uri)),
+        }
+    }
 }
 
 impl From<AnySystemPath> for DocumentKey {