[ty] Support for notebooks in VS Code (#21175)

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,18 +41,17 @@ pub fn source_text(db: &dyn Db, file: File) -> SourceText {
     }
 }
 
-fn is_notebook(path: &FilePath) -> bool {
+fn is_notebook(system: &dyn System, path: &FilePath) -> bool {
-    match path {
+    let source_type = match path {
-        FilePath::System(system) => system.extension().is_some_and(|extension| {
+        FilePath::System(path) => system.source_type(path),
-            PySourceType::try_from_extension(extension) == Some(PySourceType::Ipynb)
+        FilePath::SystemVirtual(system_virtual) => system.virtual_path_source_type(system_virtual),
-        }),
+        FilePath::Vendored(_) => return false,
-        FilePath::SystemVirtual(system_virtual) => {
+    };
-            system_virtual.extension().is_some_and(|extension| {
+
-                PySourceType::try_from_extension(extension) == Some(PySourceType::Ipynb)
+    let with_extension_fallback =
-            })
+        source_type.or_else(|| PySourceType::try_from_extension(path.extension()?));
-        }
+
-        FilePath::Vendored(_) => false,
+    with_extension_fallback == Some(PySourceType::Ipynb)
-    }
 }
 
 /// The source text of a file containing python code.

crates/ruff_db/src/system.rs

@@ -9,6 +9,7 @@ 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};
@@ -16,12 +17,11 @@ 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,35 @@ 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`.
+    ///
+    /// The default is to always return `None`, assuming the system
+    /// has no additional information and that the caller should
+    /// rely on the file extension instead.
+    ///
+    /// 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`.
+    ///
+    /// The default is to always return `None`, assuming the system
+    /// has no additional information and that the caller should
+    /// rely on the file extension instead.
+    ///
+    /// 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

@@ -39,7 +39,7 @@ impl NotebookIndex {
 
     /// Returns an iterator over the starting rows of each cell (1-based).
     ///
-    /// This yields one entry per Python cell (skipping over Makrdown cell).
+    /// This yields one entry per Python cell (skipping over Markdown cell).
     pub fn iter(&self) -> impl Iterator<Item = CellStart> + '_ {
         self.cell_starts.iter().copied()
     }
@@ -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,
+    InlayHintServerCapabilities, MarkupKind, NotebookCellSelector, NotebookSelector, OneOf,
-    SelectionRangeProviderCapability, SemanticTokensFullOptions, SemanticTokensLegend,
+    RenameOptions, SelectionRangeProviderCapability, SemanticTokensFullOptions,
-    SemanticTokensOptions, SemanticTokensServerCapabilities, ServerCapabilities,
+    SemanticTokensLegend, SemanticTokensOptions, SemanticTokensServerCapabilities,
-    SignatureHelpOptions, TextDocumentSyncCapability, TextDocumentSyncKind,
+    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 {

crates/ty_server/src/document/notebook.rs

@@ -1,26 +1,29 @@
-use anyhow::Ok;
 use lsp_types::NotebookCellKind;
 use ruff_notebook::CellMetadata;
-use rustc_hash::{FxBuildHasher, FxHashMap};
+use ruff_source_file::OneIndexed;
+use rustc_hash::FxHashMap;
 
-use super::DocumentVersion;
+use super::{DocumentKey, DocumentVersion};
-use crate::{PositionEncoding, TextDocument};
+use crate::session::index::Index;
 
-pub(super) type CellId = usize;
+/// A notebook document.
-
+///
-/// The state of a notebook document in the server. Contains an array of cells whose
+/// This notebook document only stores the metadata about the notebook
-/// contents are internally represented by [`TextDocument`]s.
+/// and the cell metadata. The cell contents are stored as separate
+/// [`super::TextDocument`]s (they can be looked up by the Cell's URL).
 #[derive(Clone, Debug)]
 pub struct NotebookDocument {
     url: lsp_types::Url,
     cells: Vec<NotebookCell>,
     metadata: ruff_notebook::RawNotebookMetadata,
     version: DocumentVersion,
-    // Used to quickly find the index of a cell for a given URL.
+    /// Map from Cell URL to their index in `cells`
-    cell_index: FxHashMap<String, CellId>,
+    cell_index: FxHashMap<lsp_types::Url, usize>,
 }
 
-/// A single cell within a notebook, which has text contents represented as a `TextDocument`.
+/// The metadata of a single cell within a notebook.
+///
+/// The cell's content are stored as a [`TextDocument`] and can be looked up by the Cell's URL.
 #[derive(Clone, Debug)]
 struct NotebookCell {
     /// The URL uniquely identifying the cell.
@@ -33,7 +36,7 @@ struct NotebookCell {
     /// > <https://microsoft.github.io/language-server-protocol/specifications/lsp/3.18/specification/#notebookDocument_synchronization>
     url: lsp_types::Url,
     kind: NotebookCellKind,
-    document: TextDocument,
+    execution_summary: Option<lsp_types::ExecutionSummary>,
 }
 
 impl NotebookDocument {
@@ -42,32 +45,18 @@ impl NotebookDocument {
         notebook_version: DocumentVersion,
         cells: Vec<lsp_types::NotebookCell>,
         metadata: serde_json::Map<String, serde_json::Value>,
-        cell_documents: Vec<lsp_types::TextDocumentItem>,
     ) -> crate::Result<Self> {
-        let mut cells: Vec<_> = cells.into_iter().map(NotebookCell::empty).collect();
+        let cells: Vec<_> = cells.into_iter().map(NotebookCell::new).collect();
-
+        let index = cells
-        let cell_index = Self::make_cell_index(&cells);
+            .iter()
-
+            .enumerate()
-        for cell_document in cell_documents {
+            .map(|(index, cell)| (cell.url.clone(), index))
-            let index = cell_index
+            .collect();
-                .get(cell_document.uri.as_str())
-                .copied()
-                .ok_or_else(|| {
-                    anyhow::anyhow!(
-                        "Received content for cell `{}` that isn't present in the metadata",
-                        cell_document.uri
-                    )
-                })?;
-
-            cells[index].document =
-                TextDocument::new(cell_document.uri, cell_document.text, cell_document.version)
-                    .with_language_id(&cell_document.language_id);
-        }
 
         Ok(Self {
+            cell_index: index,
             url,
             version: notebook_version,
-            cell_index,
             cells,
             metadata: serde_json::from_value(serde_json::Value::Object(metadata))?,
         })
@@ -79,30 +68,46 @@ impl NotebookDocument {
 
     /// Generates a pseudo-representation of a notebook that lacks per-cell metadata and contextual information
     /// but should still work with Ruff's linter.
-    pub fn make_ruff_notebook(&self) -> ruff_notebook::Notebook {
+    pub(crate) fn to_ruff_notebook(&self, index: &Index) -> ruff_notebook::Notebook {
         let cells = self
             .cells
             .iter()
-            .map(|cell| match cell.kind {
+            .map(|cell| {
+                let cell_text =
+                    if let Ok(document) = index.document(&DocumentKey::from_url(&cell.url)) {
+                        if let Some(text_document) = document.as_text() {
+                            Some(text_document.contents().to_string())
+                        } else {
+                            tracing::warn!("Non-text document found for cell `{}`", cell.url);
+                            None
+                        }
+                    } else {
+                        tracing::warn!("Text document not found for cell `{}`", cell.url);
+                        None
+                    }
+                    .unwrap_or_default();
+
+                let source = ruff_notebook::SourceValue::String(cell_text);
+                match cell.kind {
                     NotebookCellKind::Code => ruff_notebook::Cell::Code(ruff_notebook::CodeCell {
-                    execution_count: None,
+                        execution_count: cell
+                            .execution_summary
+                            .as_ref()
+                            .map(|summary| i64::from(summary.execution_order)),
                         id: None,
                         metadata: CellMetadata::default(),
                         outputs: vec![],
-                    source: ruff_notebook::SourceValue::String(
+                        source,
-                        cell.document.contents().to_string(),
-                    ),
                     }),
                     NotebookCellKind::Markup => {
                         ruff_notebook::Cell::Markdown(ruff_notebook::MarkdownCell {
                             attachments: None,
                             id: None,
                             metadata: CellMetadata::default(),
-                        source: ruff_notebook::SourceValue::String(
+                            source,
-                            cell.document.contents().to_string(),
-                        ),
                         })
                     }
+                }
             })
             .collect();
         let raw_notebook = ruff_notebook::RawNotebook {
@@ -118,93 +123,38 @@ impl NotebookDocument {
 
     pub(crate) fn update(
         &mut self,
-        cells: Option<lsp_types::NotebookDocumentCellChange>,
+        array: lsp_types::NotebookCellArrayChange,
+        updated_cells: Vec<lsp_types::NotebookCell>,
         metadata_change: Option<serde_json::Map<String, serde_json::Value>>,
         version: DocumentVersion,
-        encoding: PositionEncoding,
     ) -> crate::Result<()> {
         self.version = version;
 
-        if let Some(lsp_types::NotebookDocumentCellChange {
+        let new_cells = array.cells.unwrap_or_default();
-            structure,
+        let start = array.start as usize;
-            data,
-            text_content,
-        }) = cells
-        {
-            // The structural changes should be done first, as they may affect the cell index.
-            if let Some(structure) = structure {
-                let start = structure.array.start as usize;
-                let delete = structure.array.delete_count as usize;
 
-                // This is required because of the way the `NotebookCell` is modelled. We include
+        let added = new_cells.len();
-                // the `TextDocument` within the `NotebookCell` so when it's deleted, the
+        let deleted_range = start..start + array.delete_count as usize;
-                // corresponding `TextDocument` is removed as well. But, when cells are
-                // re-ordered, the change request doesn't provide the actual contents of the cell.
-                // Instead, it only provides that (a) these cell URIs were removed, and (b) these
-                // cell URIs were added.
-                // https://github.com/astral-sh/ruff/issues/12573
-                let mut deleted_cells = FxHashMap::default();
 
-                // First, delete the cells and remove them from the index.
+        self.cells.splice(
-                if delete > 0 {
+            deleted_range.clone(),
-                    for cell in self.cells.drain(start..start + delete) {
+            new_cells.into_iter().map(NotebookCell::new),
-                        self.cell_index.remove(cell.url.as_str());
+        );
-                        deleted_cells.insert(cell.url, cell.document);
-                    }
-                }
 
-                // Second, insert the new cells with the available information. This array does not
+        // Re-build the cell-index if new cells were added, deleted or removed
-                // provide the actual contents of the cells, so we'll initialize them with empty
+        if !deleted_range.is_empty() || added > 0 {
-                // contents.
+            self.cell_index.clear();
-                for cell in structure.array.cells.into_iter().flatten().rev() {
+            self.cell_index.extend(
-                    let (content, version) =
-                        if let Some(text_document) = deleted_cells.remove(&cell.document) {
-                            let version = text_document.version();
-                            (text_document.into_contents(), version)
-                        } else {
-                            (String::new(), 0)
-                        };
                 self.cells
-                        .insert(start, NotebookCell::new(cell, content, version));
+                    .iter()
-                }
+                    .enumerate()
-
+                    .map(|(i, cell)| (cell.url.clone(), i)),
-                // Third, register the new cells in the index and update existing ones that came
-                // after the insertion.
-                for (index, cell) in self.cells.iter().enumerate().skip(start) {
-                    self.cell_index.insert(cell.url.to_string(), index);
-                }
-
-                // Finally, update the text document that represents the cell with the actual
-                // contents. This should be done at the end so that both the `cells` and
-                // `cell_index` are updated before we start applying the changes to the cells.
-                if let Some(did_open) = structure.did_open {
-                    for cell_text_document in did_open {
-                        if let Some(cell) = self.cell_by_uri_mut(cell_text_document.uri.as_str()) {
-                            cell.document = TextDocument::new(
-                                cell_text_document.uri,
-                                cell_text_document.text,
-                                cell_text_document.version,
             );
         }
-                    }
-                }
-            }
 
-            if let Some(cell_data) = data {
+        for cell in updated_cells {
-                for cell in cell_data {
+            if let Some(existing_cell_index) = self.cell_index.get(&cell.document).copied() {
-                    if let Some(existing_cell) = self.cell_by_uri_mut(cell.document.as_str()) {
+                self.cells[existing_cell_index].kind = cell.kind;
-                        existing_cell.kind = cell.kind;
-                    }
-                }
-            }
-
-            if let Some(content_changes) = text_content {
-                for content_change in content_changes {
-                    if let Some(cell) = self.cell_by_uri_mut(content_change.document.uri.as_str()) {
-                        cell.document
-                            .apply_changes(content_change.changes, version, encoding);
-                    }
-                }
             }
         }
 
@@ -221,16 +171,10 @@ impl NotebookDocument {
     }
 
     /// Get the URI for a cell by its index within the cell array.
-    pub(crate) fn cell_uri_by_index(&self, index: CellId) -> Option<&lsp_types::Url> {
+    pub(crate) fn cell_uri_by_index(&self, index: OneIndexed) -> Option<&lsp_types::Url> {
-        self.cells.get(index).map(|cell| &cell.url)
-    }
-
-    /// Get the text document representing the contents of a cell by the cell URI.
-    #[expect(unused)]
-    pub(crate) fn cell_document_by_uri(&self, uri: &str) -> Option<&TextDocument> {
         self.cells
-            .get(*self.cell_index.get(uri)?)
+            .get(index.to_zero_indexed())
-            .map(|cell| &cell.document)
+            .map(|cell| &cell.url)
     }
 
     /// Returns a list of cell URIs in the order they appear in the array.
@@ -238,160 +182,19 @@ impl NotebookDocument {
         self.cells.iter().map(|cell| &cell.url)
     }
 
-    fn cell_by_uri_mut(&mut self, uri: &str) -> Option<&mut NotebookCell> {
+    pub(crate) fn cell_index_by_uri(&self, cell_url: &lsp_types::Url) -> Option<OneIndexed> {
-        self.cells.get_mut(*self.cell_index.get(uri)?)
+        Some(OneIndexed::from_zero_indexed(
-    }
+            self.cell_index.get(cell_url).copied()?,
-
+        ))
-    fn make_cell_index(cells: &[NotebookCell]) -> FxHashMap<String, CellId> {
-        let mut index = FxHashMap::with_capacity_and_hasher(cells.len(), FxBuildHasher);
-        for (i, cell) in cells.iter().enumerate() {
-            index.insert(cell.url.to_string(), i);
-        }
-        index
     }
 }
 
 impl NotebookCell {
-    pub(crate) fn empty(cell: lsp_types::NotebookCell) -> Self {
+    pub(crate) fn new(cell: lsp_types::NotebookCell) -> Self {
         Self {
-            kind: cell.kind,
-            document: TextDocument::new(
-                cell.document.clone(),
-                String::new(),
-                DocumentVersion::default(),
-            ),
-            url: cell.document,
-        }
-    }
-
-    pub(crate) fn new(
-        cell: lsp_types::NotebookCell,
-        contents: String,
-        version: DocumentVersion,
-    ) -> Self {
-        Self {
-            document: TextDocument::new(cell.document.clone(), contents, version),
             url: cell.document,
             kind: cell.kind,
+            execution_summary: cell.execution_summary,
         }
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::NotebookDocument;
-
-    enum TestCellContent {
-        #[expect(dead_code)]
-        Markup(String),
-        Code(String),
-    }
-
-    fn create_test_url(index: usize) -> lsp_types::Url {
-        lsp_types::Url::parse(&format!("cell:/test.ipynb#{index}")).unwrap()
-    }
-
-    fn create_test_notebook(test_cells: Vec<TestCellContent>) -> NotebookDocument {
-        let mut cells = Vec::with_capacity(test_cells.len());
-        let mut cell_documents = Vec::with_capacity(test_cells.len());
-
-        for (index, test_cell) in test_cells.into_iter().enumerate() {
-            let url = create_test_url(index);
-            match test_cell {
-                TestCellContent::Markup(content) => {
-                    cells.push(lsp_types::NotebookCell {
-                        kind: lsp_types::NotebookCellKind::Markup,
-                        document: url.clone(),
-                        metadata: None,
-                        execution_summary: None,
-                    });
-                    cell_documents.push(lsp_types::TextDocumentItem {
-                        uri: url,
-                        language_id: "markdown".to_owned(),
-                        version: 0,
-                        text: content,
-                    });
-                }
-                TestCellContent::Code(content) => {
-                    cells.push(lsp_types::NotebookCell {
-                        kind: lsp_types::NotebookCellKind::Code,
-                        document: url.clone(),
-                        metadata: None,
-                        execution_summary: None,
-                    });
-                    cell_documents.push(lsp_types::TextDocumentItem {
-                        uri: url,
-                        language_id: "python".to_owned(),
-                        version: 0,
-                        text: content,
-                    });
-                }
-            }
-        }
-
-        NotebookDocument::new(
-            lsp_types::Url::parse("file://test.ipynb").unwrap(),
-            0,
-            cells,
-            serde_json::Map::default(),
-            cell_documents,
-        )
-        .unwrap()
-    }
-
-    /// This test case checks that for a notebook with three code cells, when the client sends a
-    /// change request to swap the first two cells, the notebook document is updated correctly.
-    ///
-    /// The swap operation as a change request is represented as deleting the first two cells and
-    /// adding them back in the reverse order.
-    #[test]
-    fn swap_cells() {
-        let mut notebook = create_test_notebook(vec![
-            TestCellContent::Code("cell = 0".to_owned()),
-            TestCellContent::Code("cell = 1".to_owned()),
-            TestCellContent::Code("cell = 2".to_owned()),
-        ]);
-
-        notebook
-            .update(
-                Some(lsp_types::NotebookDocumentCellChange {
-                    structure: Some(lsp_types::NotebookDocumentCellChangeStructure {
-                        array: lsp_types::NotebookCellArrayChange {
-                            start: 0,
-                            delete_count: 2,
-                            cells: Some(vec![
-                                lsp_types::NotebookCell {
-                                    kind: lsp_types::NotebookCellKind::Code,
-                                    document: create_test_url(1),
-                                    metadata: None,
-                                    execution_summary: None,
-                                },
-                                lsp_types::NotebookCell {
-                                    kind: lsp_types::NotebookCellKind::Code,
-                                    document: create_test_url(0),
-                                    metadata: None,
-                                    execution_summary: None,
-                                },
-                            ]),
-                        },
-                        did_open: None,
-                        did_close: None,
-                    }),
-                    data: None,
-                    text_content: None,
-                }),
-                None,
-                1,
-                crate::PositionEncoding::default(),
-            )
-            .unwrap();
-
-        assert_eq!(
-            notebook.make_ruff_notebook().source_code(),
-            "cell = 1
-cell = 0
-cell = 2
-"
-        );
-    }
-}

crates/ty_server/src/document/range.rs

@@ -78,7 +78,7 @@ impl LspPosition {
 }
 
 pub(crate) trait RangeExt {
-    /// Convert an LSP Range to internal [`TextRange`].
+    /// Convert an LSP Range to a [`TextRange`].
     ///
     /// Returns `None` if `file` is a notebook and the
     /// cell identified by `url` can't be looked up or if the notebook
@@ -110,6 +110,10 @@ impl RangeExt for lsp_types::Range {
 pub(crate) trait PositionExt {
     /// Convert an LSP Position to internal `TextSize`.
     ///
+    /// For notebook support, this uses the URI to determine which cell the position
+    /// refers to, and maps the cell-relative position to the absolute position in the
+    /// concatenated notebook file.
+    ///
     /// Returns `None` if `file` is a notebook and the
     /// cell identified by `url` can't be looked up or if the notebook
     /// isn't open in the editor.
@@ -127,18 +131,39 @@ impl PositionExt for lsp_types::Position {
         &self,
         db: &dyn Db,
         file: File,
-        _url: &lsp_types::Url,
+        url: &lsp_types::Url,
         encoding: PositionEncoding,
     ) -> Option<TextSize> {
         let source = source_text(db, file);
         let index = line_index(db, file);
 
+        if let Some(notebook) = source.as_notebook() {
+            let notebook_document = db.notebook_document(file)?;
+            let cell_index = notebook_document.cell_index_by_uri(url)?;
+
+            let cell_start_offset = notebook.cell_offset(cell_index).unwrap_or_default();
+            let cell_relative_line = OneIndexed::from_zero_indexed(u32_index_to_usize(self.line));
+
+            let cell_start_location =
+                index.source_location(cell_start_offset, source.as_str(), encoding.into());
+            assert_eq!(cell_start_location.character_offset, OneIndexed::MIN);
+
+            // Absolute position into the concatenated notebook source text.
+            let absolute_position = SourceLocation {
+                line: cell_start_location
+                    .line
+                    .saturating_add(cell_relative_line.to_zero_indexed()),
+                character_offset: OneIndexed::from_zero_indexed(u32_index_to_usize(self.character)),
+            };
+            return Some(index.offset(absolute_position, &source, encoding.into()));
+        }
+
         Some(lsp_position_to_text_size(*self, &source, &index, encoding))
     }
 }
 
 pub(crate) trait TextSizeExt {
-    /// Converts self into a position into an LSP text document (can be a cell or regular document).
+    /// Converts `self` into a position in an LSP text document (can be a cell or regular document).
     ///
     /// Returns `None` if the position can't be converted:
     ///
@@ -165,6 +190,19 @@ impl TextSizeExt for TextSize {
         let source = source_text(db, file);
         let index = line_index(db, file);
 
+        if let Some(notebook) = source.as_notebook() {
+            let notebook_document = db.notebook_document(file)?;
+            let start = index.source_location(*self, source.as_str(), encoding.into());
+            let cell = notebook.index().cell(start.line)?;
+
+            let cell_relative_start = notebook.index().translate_source_location(&start);
+
+            return Some(LspPosition {
+                uri: Some(notebook_document.cell_uri_by_index(cell)?.clone()),
+                position: source_location_to_position(&cell_relative_start),
+            });
+        }
+
         let uri = file_to_url(db, file);
         let position = text_size_to_lsp_position(*self, &source, &index, encoding);
 
@@ -252,6 +290,34 @@ impl ToRangeExt for TextRange {
     ) -> Option<LspRange> {
         let source = source_text(db, file);
         let index = line_index(db, file);
+
+        if let Some(notebook) = source.as_notebook() {
+            let notebook_index = notebook.index();
+            let notebook_document = db.notebook_document(file)?;
+
+            let start_in_concatenated =
+                index.source_location(self.start(), &source, encoding.into());
+            let cell_index = notebook_index.cell(start_in_concatenated.line)?;
+
+            let end_in_concatenated = index.source_location(self.end(), &source, encoding.into());
+
+            let start_in_cell = source_location_to_position(
+                &notebook_index.translate_source_location(&start_in_concatenated),
+            );
+            let end_in_cell = source_location_to_position(
+                &notebook_index.translate_source_location(&end_in_concatenated),
+            );
+
+            let cell_uri = notebook_document
+                .cell_uri_by_index(cell_index)
+                .expect("Index to contain an URI for every cell");
+
+            return Some(LspRange {
+                uri: Some(cell_uri.clone()),
+                range: lsp_types::Range::new(start_in_cell, end_in_cell),
+            });
+        }
+
         let range = text_range_to_lsp_range(*self, &source, &index, encoding);
 
         let uri = file_to_url(db, file);

crates/ty_server/src/document/text_document.rs

@@ -2,11 +2,13 @@ use lsp_types::{TextDocumentContentChangeEvent, Url};
 use ruff_source_file::LineIndex;
 
 use crate::PositionEncoding;
-
+use crate::document::range::lsp_range_to_text_range;
-use super::range::lsp_range_to_text_range;
+use crate::system::AnySystemPath;
 
 pub(crate) type DocumentVersion = i32;
 
+/// A regular text file or the content of a notebook cell.
+///
 /// The state of an individual document in the server. Stays up-to-date
 /// with changes made by the user, including unsaved changes.
 #[derive(Debug, Clone)]
@@ -16,15 +18,16 @@ pub struct TextDocument {
 
     /// The string contents of the document.
     contents: String,
-    /// A computed line index for the document. This should always reflect
+
-    /// the current version of `contents`. Using a function like [`Self::modify`]
-    /// will re-calculate the line index automatically when the `contents` value is updated.
-    index: LineIndex,
     /// The latest version of the document, set by the LSP client. The server will panic in
     /// debug mode if we attempt to update the document with an 'older' version.
     version: DocumentVersion,
+
     /// The language ID of the document as provided by the client.
     language_id: Option<LanguageId>,
+
+    /// For cells, the path to the notebook document.
+    notebook: Option<AnySystemPath>,
 }
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
@@ -44,13 +47,12 @@ impl From<&str> for LanguageId {
 
 impl TextDocument {
     pub fn new(url: Url, contents: String, version: DocumentVersion) -> Self {
-        let index = LineIndex::from_source_text(&contents);
         Self {
             url,
             contents,
-            index,
             version,
             language_id: None,
+            notebook: None,
         }
     }
 
@@ -60,6 +62,12 @@ impl TextDocument {
         self
     }
 
+    #[must_use]
+    pub(crate) fn with_notebook(mut self, notebook: AnySystemPath) -> Self {
+        self.notebook = Some(notebook);
+        self
+    }
+
     pub fn into_contents(self) -> String {
         self.contents
     }
@@ -72,10 +80,6 @@ impl TextDocument {
         &self.contents
     }
 
-    pub fn index(&self) -> &LineIndex {
-        &self.index
-    }
-
     pub fn version(&self) -> DocumentVersion {
         self.version
     }
@@ -84,6 +88,10 @@ impl TextDocument {
         self.language_id
     }
 
+    pub(crate) fn notebook(&self) -> Option<&AnySystemPath> {
+        self.notebook.as_ref()
+    }
+
     pub fn apply_changes(
         &mut self,
         changes: Vec<lsp_types::TextDocumentContentChangeEvent>,
@@ -105,7 +113,7 @@ impl TextDocument {
         }
 
         let mut new_contents = self.contents().to_string();
-        let mut active_index = self.index().clone();
+        let mut active_index = LineIndex::from_source_text(&new_contents);
 
         for TextDocumentContentChangeEvent {
             range,
@@ -127,34 +135,22 @@ impl TextDocument {
             active_index = LineIndex::from_source_text(&new_contents);
         }
 
-        self.modify_with_manual_index(|contents, version, index| {
+        self.modify(|contents, version| {
-            *index = active_index;
             *contents = new_contents;
             *version = new_version;
         });
     }
 
     pub fn update_version(&mut self, new_version: DocumentVersion) {
-        self.modify_with_manual_index(|_, version, _| {
+        self.modify(|_, version| {
             *version = new_version;
         });
     }
 
-    // A private function for modifying the document's internal state
-    fn modify(&mut self, func: impl FnOnce(&mut String, &mut DocumentVersion)) {
-        self.modify_with_manual_index(|c, v, i| {
-            func(c, v);
-            *i = LineIndex::from_source_text(c);
-        });
-    }
-
     // A private function for overriding how we update the line index by default.
-    fn modify_with_manual_index(
+    fn modify(&mut self, func: impl FnOnce(&mut String, &mut DocumentVersion)) {
-        &mut self,
-        func: impl FnOnce(&mut String, &mut DocumentVersion, &mut LineIndex),
-    ) {
         let old_version = self.version;
-        func(&mut self.contents, &mut self.version, &mut self.index);
+        func(&mut self.contents, &mut self.version);
         debug_assert!(self.version >= old_version);
     }
 }

crates/ty_server/src/server/api.rs

@@ -147,6 +147,9 @@ pub(super) fn notification(notif: server::Notification) -> Task {
         notifications::DidOpenNotebookHandler::METHOD => {
             sync_notification_task::<notifications::DidOpenNotebookHandler>(notif)
         }
+        notifications::DidChangeNotebookHandler::METHOD => {
+            sync_notification_task::<notifications::DidChangeNotebookHandler>(notif)
+        }
         notifications::DidCloseNotebookHandler::METHOD => {
             sync_notification_task::<notifications::DidCloseNotebookHandler>(notif)
         }
@@ -273,8 +276,8 @@ where
             });
         };
 
-        let path = document.to_file_path();
+        let path = document.notebook_or_file_path();
-        let db = session.project_db(&path).clone();
+        let db = session.project_db(path).clone();
 
         Box::new(move |client| {
             let _span = tracing::debug_span!("request", %id, method = R::METHOD).entered();

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

@@ -3,29 +3,30 @@ use std::hash::{DefaultHasher, Hash as _, Hasher as _};
 use lsp_types::notification::PublishDiagnostics;
 use lsp_types::{
     CodeDescription, Diagnostic, DiagnosticRelatedInformation, DiagnosticSeverity, DiagnosticTag,
-    NumberOrString, PublishDiagnosticsParams, Range, Url,
+    NumberOrString, PublishDiagnosticsParams, Url,
 };
+use ruff_db::source::source_text;
 use rustc_hash::FxHashMap;
 
 use ruff_db::diagnostic::{Annotation, Severity, SubDiagnostic};
-use ruff_db::files::FileRange;
+use ruff_db::files::{File, FileRange};
 use ruff_db::system::SystemPathBuf;
 use ty_project::{Db as _, ProjectDatabase};
 
 use crate::Db;
 use crate::document::{FileRangeExt, ToRangeExt};
-use crate::session::DocumentSnapshot;
+use crate::session::DocumentHandle;
 use crate::session::client::Client;
 use crate::system::{AnySystemPath, file_to_url};
-use crate::{NotebookDocument, PositionEncoding, Session};
+use crate::{PositionEncoding, Session};
 
-pub(super) struct Diagnostics<'a> {
+pub(super) struct Diagnostics {
     items: Vec<ruff_db::diagnostic::Diagnostic>,
     encoding: PositionEncoding,
-    notebook: Option<&'a NotebookDocument>,
+    file_or_notebook: File,
 }
 
-impl Diagnostics<'_> {
+impl Diagnostics {
     /// Computes the result ID for `diagnostics`.
     ///
     /// Returns `None` if there are no diagnostics.
@@ -53,30 +54,27 @@ impl Diagnostics<'_> {
     }
 
     pub(super) fn to_lsp_diagnostics(&self, db: &ProjectDatabase) -> LspDiagnostics {
-        if let Some(notebook) = self.notebook {
+        if let Some(notebook_document) = db.notebook_document(self.file_or_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() {
+            for cell_url in notebook_document.cell_urls() {
                 cell_diagnostics.entry(cell_url.clone()).or_default();
             }
 
-            for (cell_index, diagnostic) in self.items.iter().map(|diagnostic| {
+            for diagnostic in &self.items {
-                (
+                let (url, lsp_diagnostic) = to_lsp_diagnostic(db, diagnostic, self.encoding);
-                    // TODO: Use the cell index instead using `SourceKind`
+
-                    usize::default(),
+                let Some(url) = url else {
-                    to_lsp_diagnostic(db, diagnostic, self.encoding),
+                    tracing::warn!("Unable to find notebook cell");
-                )
-            }) {
-                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())
+                    .entry(url)
                     .or_default()
-                    .push(diagnostic);
+                    .push(lsp_diagnostic);
             }
 
             LspDiagnostics::NotebookDocument(cell_diagnostics)
@@ -84,7 +82,7 @@ impl Diagnostics<'_> {
             LspDiagnostics::TextDocument(
                 self.items
                     .iter()
-                    .map(|diagnostic| to_lsp_diagnostic(db, diagnostic, self.encoding))
+                    .map(|diagnostic| to_lsp_diagnostic(db, diagnostic, self.encoding).1)
                     .collect(),
             )
         }
@@ -115,16 +113,25 @@ impl LspDiagnostics {
     }
 }
 
+pub(super) fn clear_diagnostics_if_needed(
+    document: &DocumentHandle,
+    session: &Session,
+    client: &Client,
+) {
+    if session.client_capabilities().supports_pull_diagnostics() && !document.is_cell_or_notebook()
+    {
+        return;
+    }
+
+    clear_diagnostics(document.url(), client);
+}
+
 /// Clears the diagnostics for the document identified by `uri`.
 ///
 /// This is done by notifying the client with an empty list of diagnostics for the document.
 /// For notebook cells, this clears diagnostics for the specific cell.
 /// For other document types, this clears diagnostics for the main document.
-pub(super) fn clear_diagnostics(session: &Session, uri: &lsp_types::Url, client: &Client) {
+pub(super) fn clear_diagnostics(uri: &lsp_types::Url, client: &Client) {
-    if session.client_capabilities().supports_pull_diagnostics() {
-        return;
-    }
-
     client.send_notification::<PublishDiagnostics>(PublishDiagnosticsParams {
         uri: uri.clone(),
         diagnostics: vec![],
@@ -133,27 +140,32 @@ pub(super) fn clear_diagnostics(session: &Session, uri: &lsp_types::Url, client:
 }
 
 /// Publishes the diagnostics for the given document snapshot using the [publish diagnostics
-/// notification].
+/// notification] .
 ///
-/// This function is a no-op if the client supports pull diagnostics.
+/// Unlike [`publish_diagnostics`], this function only publishes diagnostics if a client doesn't support
+/// pull diagnostics and `document` is not a notebook or cell (VS Code
+/// does not support pull diagnostics for notebooks or cells (as of 2025-11-12).
 ///
 /// [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: &lsp_types::Url, client: &Client) {
+pub(super) fn publish_diagnostics_if_needed(
-    if session.client_capabilities().supports_pull_diagnostics() {
+    document: &DocumentHandle,
+    session: &Session,
+    client: &Client,
+) {
+    if !document.is_cell_or_notebook() && session.client_capabilities().supports_pull_diagnostics()
+    {
         return;
     }
 
-    let snapshot = match session.snapshot_document(url) {
+    publish_diagnostics(document, session, client);
-        Ok(document) => document,
+}
-        Err(err) => {
-            tracing::debug!("Failed to resolve document for URL `{}`: {}", url, err);
-            return;
-        }
-    };
 
-    let db = session.project_db(&snapshot.to_file_path());
+/// Publishes the diagnostics for the given document snapshot using the [publish diagnostics
+/// notification].
+pub(super) fn publish_diagnostics(document: &DocumentHandle, session: &Session, client: &Client) {
+    let db = session.project_db(document.notebook_or_file_path());
 
-    let Some(diagnostics) = compute_diagnostics(db, &snapshot) else {
+    let Some(diagnostics) = compute_diagnostics(db, document, session.position_encoding()) else {
         return;
     };
 
@@ -162,13 +174,13 @@ pub(super) fn publish_diagnostics(session: &Session, url: &lsp_types::Url, clien
         client.send_notification::<PublishDiagnostics>(PublishDiagnosticsParams {
             uri,
             diagnostics,
-            version: Some(snapshot.document().version()),
+            version: Some(document.version()),
         });
     };
 
     match diagnostics.to_lsp_diagnostics(db) {
         LspDiagnostics::TextDocument(diagnostics) => {
-            publish_diagnostics_notification(url.clone(), diagnostics);
+            publish_diagnostics_notification(document.url().clone(), diagnostics);
         }
         LspDiagnostics::NotebookDocument(cell_diagnostics) => {
             for (cell_url, diagnostics) in cell_diagnostics {
@@ -238,7 +250,7 @@ pub(crate) fn publish_settings_diagnostics(
         // Convert diagnostics to LSP format
         let lsp_diagnostics = file_diagnostics
             .into_iter()
-            .map(|diagnostic| to_lsp_diagnostic(db, &diagnostic, session_encoding))
+            .map(|diagnostic| to_lsp_diagnostic(db, &diagnostic, session_encoding).1)
             .collect::<Vec<_>>();
 
         client.send_notification::<PublishDiagnostics>(PublishDiagnosticsParams {
@@ -249,24 +261,26 @@ pub(crate) fn publish_settings_diagnostics(
     }
 }
 
-pub(super) fn compute_diagnostics<'a>(
+pub(super) fn compute_diagnostics(
     db: &ProjectDatabase,
-    snapshot: &'a DocumentSnapshot,
+    document: &DocumentHandle,
-) -> Option<Diagnostics<'a>> {
+    encoding: PositionEncoding,
-    let Some(file) = snapshot.to_file(db) else {
+) -> Option<Diagnostics> {
+    let Some(file) = document.notebook_or_file(db) else {
         tracing::info!(
             "No file found for snapshot for `{}`",
-            snapshot.to_file_path()
+            document.notebook_or_file_path()
         );
         return None;
     };
 
+    tracing::debug!("source text: {}", source_text(db, file).as_str());
     let diagnostics = db.check_file(file);
 
     Some(Diagnostics {
         items: diagnostics,
-        encoding: snapshot.encoding(),
+        encoding,
-        notebook: snapshot.notebook(),
+        file_or_notebook: file,
     })
 }
 
@@ -276,16 +290,18 @@ pub(super) fn to_lsp_diagnostic(
     db: &dyn Db,
     diagnostic: &ruff_db::diagnostic::Diagnostic,
     encoding: PositionEncoding,
-) -> Diagnostic {
+) -> (Option<lsp_types::Url>, Diagnostic) {
-    let range = if let Some(span) = diagnostic.primary_span() {
+    let location = diagnostic.primary_span().and_then(|span| {
         let file = span.expect_ty_file();
-
+        span.range()?
-        span.range()
+            .to_lsp_range(db, file, encoding)
-            .and_then(|range| range.to_lsp_range(db, file, encoding))
             .unwrap_or_default()
-            .local_range()
+            .to_location()
-    } else {
+    });
-        Range::default()
+
+    let (range, url) = match location {
+        Some(location) => (location.range, Some(location.uri)),
+        None => (lsp_types::Range::default(), None),
     };
 
     let severity = match diagnostic.severity() {
@@ -341,6 +357,8 @@ pub(super) fn to_lsp_diagnostic(
         );
     }
 
+    (
+        url,
         Diagnostic {
             range,
             severity: Some(severity),
@@ -351,7 +369,8 @@ pub(super) fn to_lsp_diagnostic(
             message: diagnostic.concise_message().to_string(),
             related_information: Some(related_information),
             data: None,
-    }
+        },
+    )
 }
 
 /// Converts an [`Annotation`] to a [`DiagnosticRelatedInformation`].

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

@@ -1,5 +1,6 @@
 mod cancel;
 mod did_change;
+mod did_change_notebook;
 mod did_change_watched_files;
 mod did_close;
 mod did_close_notebook;
@@ -8,6 +9,7 @@ mod did_open_notebook;
 
 pub(super) use cancel::CancelNotificationHandler;
 pub(super) use did_change::DidChangeTextDocumentHandler;
+pub(super) use did_change_notebook::DidChangeNotebookHandler;
 pub(super) use did_change_watched_files::DidChangeWatchedFiles;
 pub(super) use did_close::DidCloseTextDocumentHandler;
 pub(super) use did_close_notebook::DidCloseNotebookHandler;

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

@@ -4,12 +4,10 @@ use lsp_types::{DidChangeTextDocumentParams, VersionedTextDocumentIdentifier};
 
 use crate::server::Result;
 use crate::server::api::LSPResult;
-use crate::server::api::diagnostics::publish_diagnostics;
+use crate::server::api::diagnostics::publish_diagnostics_if_needed;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::session::Session;
 use crate::session::client::Client;
-use crate::system::AnySystemPath;
-use ty_project::watch::ChangeEvent;
 
 pub(crate) struct DidChangeTextDocumentHandler;
 
@@ -36,19 +34,7 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
             .update_text_document(session, content_changes, version)
             .with_failure_code(ErrorCode::InternalError)?;
 
-        let path = document.to_file_path();
+        publish_diagnostics_if_needed(&document, session, client);
-        let changes = match &*path {
-            AnySystemPath::System(system_path) => {
-                vec![ChangeEvent::file_content_changed(system_path.clone())]
-            }
-            AnySystemPath::SystemVirtual(virtual_path) => {
-                vec![ChangeEvent::ChangedVirtual(virtual_path.clone())]
-            }
-        };
-
-        session.apply_changes(&path, changes);
-
-        publish_diagnostics(session, document.url(), client);
 
         Ok(())
     }

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

@@ -0,0 +1,40 @@
+use lsp_server::ErrorCode;
+use lsp_types as types;
+use lsp_types::notification as notif;
+
+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::session::Session;
+use crate::session::client::Client;
+
+pub(crate) struct DidChangeNotebookHandler;
+
+impl NotificationHandler for DidChangeNotebookHandler {
+    type NotificationType = notif::DidChangeNotebookDocument;
+}
+
+impl SyncNotificationHandler for DidChangeNotebookHandler {
+    fn run(
+        session: &mut Session,
+        client: &Client,
+        types::DidChangeNotebookDocumentParams {
+            notebook_document: types::VersionedNotebookDocumentIdentifier { uri, version },
+            change: types::NotebookDocumentChangeEvent { cells, metadata },
+        }: types::DidChangeNotebookDocumentParams,
+    ) -> Result<()> {
+        let document = session
+            .document_handle(&uri)
+            .with_failure_code(ErrorCode::InternalError)?;
+
+        document
+            .update_notebook_document(session, cells, metadata, version)
+            .with_failure_code(ErrorCode::InternalError)?;
+
+        // Always publish diagnostics because notebooks only support publish diagnostics.
+        publish_diagnostics(&document, session, client);
+
+        Ok(())
+    }
+}

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

@@ -26,8 +26,7 @@ impl SyncNotificationHandler for DidChangeWatchedFiles {
         let mut events_by_db: FxHashMap<_, Vec<ChangeEvent>> = FxHashMap::default();
 
         for change in params.changes {
-            let key = DocumentKey::from_url(&change.uri);
+            let path = DocumentKey::from_url(&change.uri).into_file_path();
-            let path = key.to_file_path();
 
             let system_path = match path {
                 AnySystemPath::System(system) => system,
@@ -93,10 +92,9 @@ impl SyncNotificationHandler for DidChangeWatchedFiles {
             );
         } else {
             for key in session.text_document_handles() {
-                publish_diagnostics(session, key.url(), client);
+                publish_diagnostics(&key, session, client);
             }
         }
-        // TODO: always publish diagnostics for notebook files (since they don't use pull diagnostics)
 
         if client_capabilities.supports_inlay_hint_refresh() {
             client.send_request::<types::request::InlayHintRefreshRequest>(session, (), |_, ()| {});

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

@@ -1,15 +1,13 @@
-use crate::server::Result;
-use crate::server::api::LSPResult;
-use crate::server::api::diagnostics::clear_diagnostics;
-use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
-use crate::session::Session;
-use crate::session::client::Client;
-use crate::system::AnySystemPath;
 use lsp_server::ErrorCode;
 use lsp_types::notification::DidCloseTextDocument;
 use lsp_types::{DidCloseTextDocumentParams, TextDocumentIdentifier};
-use ruff_db::Db as _;
+
-use ty_project::Db as _;
+use crate::server::Result;
+use crate::server::api::LSPResult;
+use crate::server::api::diagnostics::clear_diagnostics_if_needed;
+use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
+use crate::session::Session;
+use crate::session::client::Client;
 
 pub(crate) struct DidCloseTextDocumentHandler;
 
@@ -31,53 +29,12 @@ impl SyncNotificationHandler for DidCloseTextDocumentHandler {
             .document_handle(&uri)
             .with_failure_code(ErrorCode::InternalError)?;
 
-        let path = document.to_file_path().into_owned();
+        let should_clear_diagnostics = document
-        let url = document.url().clone();
-
-        document
             .close(session)
             .with_failure_code(ErrorCode::InternalError)?;
 
-        let db = session.project_db_mut(&path);
+        if should_clear_diagnostics {
-
+            clear_diagnostics_if_needed(&document, session, client);
-        match &path {
-            AnySystemPath::System(system_path) => {
-                if let Some(file) = db.files().try_system(db, system_path) {
-                    db.project().close_file(db, file);
-                } else {
-                    // This can only fail when the path is a directory or it doesn't exists but the
-                    // file should exists for this handler in this branch. This is because every
-                    // close call is preceded by an open call, which ensures that the file is
-                    // interned in the lookup table (`Files`).
-                    tracing::warn!("Salsa file does not exists for {}", system_path);
-                }
-
-                // For non-virtual files, we clear diagnostics if:
-                //
-                // 1. The file does not belong to any workspace e.g., opening a random file from
-                //    outside the workspace because closing it acts like the file doesn't exists
-                // 2. The diagnostic mode is set to open-files only
-                if session.workspaces().for_path(system_path).is_none()
-                    || session
-                        .global_settings()
-                        .diagnostic_mode()
-                        .is_open_files_only()
-                {
-                    clear_diagnostics(session, &url, client);
-                }
-            }
-            AnySystemPath::SystemVirtual(virtual_path) => {
-                if let Some(virtual_file) = db.files().try_virtual_file(virtual_path) {
-                    db.project().close_file(db, virtual_file.file());
-                    virtual_file.close(db);
-                } else {
-                    tracing::warn!("Salsa virtual file does not exists for {}", virtual_path);
-                }
-
-                // Always clear diagnostics for virtual files, as they don't really exist on disk
-                // which means closing them is like deleting the file.
-                clear_diagnostics(session, &url, client);
-            }
         }
 
         Ok(())

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

@@ -6,8 +6,6 @@ use crate::server::api::LSPResult;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::session::Session;
 use crate::session::client::Client;
-use crate::system::AnySystemPath;
-use ty_project::watch::ChangeEvent;
 
 pub(crate) struct DidCloseNotebookHandler;
 
@@ -30,19 +28,12 @@ impl SyncNotificationHandler for DidCloseNotebookHandler {
             .document_handle(&uri)
             .with_failure_code(lsp_server::ErrorCode::InternalError)?;
 
-        let path = document.to_file_path().into_owned();
+        // We don't need to call publish any diagnostics because we clear
-
+        // the diagnostics when closing the corresponding cell documents.
-        document
+        let _ = document
             .close(session)
             .with_failure_code(lsp_server::ErrorCode::InternalError)?;
 
-        if let AnySystemPath::SystemVirtual(virtual_path) = &path {
-            session.apply_changes(
-                &path,
-                vec![ChangeEvent::DeletedVirtual(virtual_path.clone())],
-            );
-        }
-
         Ok(())
     }
 }

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

@@ -1,17 +1,12 @@
 use lsp_types::notification::DidOpenTextDocument;
 use lsp_types::{DidOpenTextDocumentParams, TextDocumentItem};
-use ruff_db::Db as _;
-use ruff_db::files::system_path_to_file;
-use ty_project::Db as _;
-use ty_project::watch::{ChangeEvent, CreatedKind};
 
 use crate::TextDocument;
 use crate::server::Result;
-use crate::server::api::diagnostics::publish_diagnostics;
+use crate::server::api::diagnostics::publish_diagnostics_if_needed;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::session::Session;
 use crate::session::client::Client;
-use crate::system::AnySystemPath;
 
 pub(crate) struct DidOpenTextDocumentHandler;
 
@@ -39,44 +34,7 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
             TextDocument::new(uri, text, version).with_language_id(&language_id),
         );
 
-        let path = document.to_file_path();
+        publish_diagnostics_if_needed(&document, session, client);
-
-        // This is a "maybe" because the `File` might've not been interned yet i.e., the
-        // `try_system` call will return `None` which doesn't mean that the file is new, it's just
-        // that the server didn't need the file yet.
-        let is_maybe_new_system_file = path.as_system().is_some_and(|system_path| {
-            let db = session.project_db(&path);
-            db.files()
-                .try_system(db, system_path)
-                .is_none_or(|file| !file.exists(db))
-        });
-
-        match &*path {
-            AnySystemPath::System(system_path) => {
-                let event = if is_maybe_new_system_file {
-                    ChangeEvent::Created {
-                        path: system_path.clone(),
-                        kind: CreatedKind::File,
-                    }
-                } else {
-                    ChangeEvent::Opened(system_path.clone())
-                };
-                session.apply_changes(&path, vec![event]);
-
-                let db = session.project_db_mut(&path);
-                match system_path_to_file(db, system_path) {
-                    Ok(file) => db.project().open_file(db, file),
-                    Err(err) => tracing::warn!("Failed to open file {system_path}: {err}"),
-                }
-            }
-            AnySystemPath::SystemVirtual(virtual_path) => {
-                let db = session.project_db_mut(&path);
-                let virtual_file = db.files().virtual_file(db, virtual_path);
-                db.project().open_file(db, virtual_file.file());
-            }
-        }
-
-        publish_diagnostics(session, document.url(), client);
 
         Ok(())
     }

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

@@ -2,16 +2,14 @@ use lsp_server::ErrorCode;
 use lsp_types::DidOpenNotebookDocumentParams;
 use lsp_types::notification::DidOpenNotebookDocument;
 
-use ruff_db::Db;
+use crate::TextDocument;
-use ty_project::watch::ChangeEvent;
-
 use crate::document::NotebookDocument;
 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::session::Session;
 use crate::session::client::Client;
-use crate::system::AnySystemPath;
 
 pub(crate) struct DidOpenNotebookHandler;
 
@@ -22,7 +20,7 @@ impl NotificationHandler for DidOpenNotebookHandler {
 impl SyncNotificationHandler for DidOpenNotebookHandler {
     fn run(
         session: &mut Session,
-        _client: &Client,
+        client: &Client,
         params: DidOpenNotebookDocumentParams,
     ) -> Result<()> {
         let lsp_types::NotebookDocument {
@@ -33,29 +31,22 @@ impl SyncNotificationHandler for DidOpenNotebookHandler {
             ..
         } = params.notebook_document;
 
-        let notebook = NotebookDocument::new(
+        let notebook =
-            notebook_uri,
+            NotebookDocument::new(notebook_uri, version, cells, metadata.unwrap_or_default())
-            version,
-            cells,
-            metadata.unwrap_or_default(),
-            params.cell_text_documents,
-        )
                 .with_failure_code(ErrorCode::InternalError)?;
 
         let document = session.open_notebook_document(notebook);
-        let path = document.to_file_path();
+        let notebook_path = document.notebook_or_file_path();
 
-        match &*path {
+        for cell in params.cell_text_documents {
-            AnySystemPath::System(system_path) => {
+            let cell_document = TextDocument::new(cell.uri, cell.text, cell.version)
-                session.apply_changes(&path, vec![ChangeEvent::Opened(system_path.clone())]);
+                .with_language_id(&cell.language_id)
-            }
+                .with_notebook(notebook_path.clone());
-            AnySystemPath::SystemVirtual(virtual_path) => {
+            session.open_text_document(cell_document);
-                let db = session.project_db_mut(&path);
-                db.files().virtual_file(db, virtual_path);
-            }
         }
 
-        // TODO(dhruvmanila): Publish diagnostics if the client doesn't support pull diagnostics
+        // Always publish diagnostics because notebooks only support publish diagnostics.
+        publish_diagnostics(&document, session, client);
 
         Ok(())
     }

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

@@ -44,7 +44,7 @@ impl BackgroundDocumentRequestHandler for CompletionRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.to_file(db) else {
+        let Some(file) = snapshot.to_notebook_or_file(db) else {
             return Ok(None);
         };
 
@@ -56,7 +56,6 @@ impl BackgroundDocumentRequestHandler for CompletionRequestHandler {
         ) else {
             return Ok(None);
         };
-
         let settings = CompletionSettings {
             auto_import: snapshot.global_settings().is_auto_import_enabled(),
         };

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

@@ -33,7 +33,7 @@ impl BackgroundDocumentRequestHandler for DocumentDiagnosticRequestHandler {
         _client: &Client,
         params: DocumentDiagnosticParams,
     ) -> Result<DocumentDiagnosticReportResult> {
-        let diagnostics = compute_diagnostics(db, snapshot);
+        let diagnostics = compute_diagnostics(db, snapshot.document(), snapshot.encoding());
 
         let Some(diagnostics) = diagnostics else {
             return Ok(DocumentDiagnosticReportResult::Report(

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

@@ -36,7 +36,7 @@ impl BackgroundDocumentRequestHandler for DocumentHighlightRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.to_file(db) else {
+        let Some(file) = snapshot.to_notebook_or_file(db) else {
             return Ok(None);
         };
 

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

@@ -39,7 +39,7 @@ impl BackgroundDocumentRequestHandler for DocumentSymbolRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.to_file(db) else {
+        let Some(file) = snapshot.to_notebook_or_file(db) else {
             return Ok(None);
         };
 

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

@@ -36,7 +36,7 @@ impl BackgroundDocumentRequestHandler for GotoDeclarationRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.to_file(db) else {
+        let Some(file) = snapshot.to_notebook_or_file(db) else {
             return Ok(None);
         };
 

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

@@ -36,7 +36,7 @@ impl BackgroundDocumentRequestHandler for GotoDefinitionRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.to_file(db) else {
+        let Some(file) = snapshot.to_notebook_or_file(db) else {
             return Ok(None);
         };
 

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

@@ -36,7 +36,7 @@ impl BackgroundDocumentRequestHandler for ReferencesRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.to_file(db) else {
+        let Some(file) = snapshot.to_notebook_or_file(db) else {
             return Ok(None);
         };
 

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

@@ -36,7 +36,7 @@ impl BackgroundDocumentRequestHandler for GotoTypeDefinitionRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.to_file(db) else {
+        let Some(file) = snapshot.to_notebook_or_file(db) else {
             return Ok(None);
         };
 

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

@@ -35,7 +35,7 @@ impl BackgroundDocumentRequestHandler for HoverRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.to_file(db) else {
+        let Some(file) = snapshot.to_notebook_or_file(db) else {
             return Ok(None);
         };
 

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

@@ -1,15 +1,16 @@
 use std::borrow::Cow;
 
+use lsp_types::request::InlayHintRequest;
+use lsp_types::{InlayHintParams, Url};
+use ty_ide::{InlayHintKind, InlayHintLabel, inlay_hints};
+use ty_project::ProjectDatabase;
+
 use crate::document::{RangeExt, TextSizeExt};
 use crate::server::api::traits::{
     BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler,
 };
 use crate::session::DocumentSnapshot;
 use crate::session::client::Client;
-use lsp_types::request::InlayHintRequest;
-use lsp_types::{InlayHintParams, Url};
-use ty_ide::{InlayHintKind, InlayHintLabel, inlay_hints};
-use ty_project::ProjectDatabase;
 
 pub(crate) struct InlayHintRequestHandler;
 
@@ -35,7 +36,7 @@ impl BackgroundDocumentRequestHandler for InlayHintRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.to_file(db) else {
+        let Some(file) = snapshot.to_notebook_or_file(db) else {
             return Ok(None);
         };
 

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

@@ -36,7 +36,7 @@ impl BackgroundDocumentRequestHandler for PrepareRenameRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.to_file(db) else {
+        let Some(file) = snapshot.to_notebook_or_file(db) else {
             return Ok(None);
         };
 

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

@@ -37,7 +37,7 @@ impl BackgroundDocumentRequestHandler for RenameRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.to_file(db) else {
+        let Some(file) = snapshot.to_notebook_or_file(db) else {
             return Ok(None);
         };
 

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

@@ -36,7 +36,7 @@ impl BackgroundDocumentRequestHandler for SelectionRangeRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.to_file(db) else {
+        let Some(file) = snapshot.to_notebook_or_file(db) else {
             return Ok(None);
         };
 

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

@@ -1,13 +1,16 @@
 use std::borrow::Cow;
 
+use lsp_types::{SemanticTokens, SemanticTokensParams, SemanticTokensResult, Url};
+use ruff_db::source::source_text;
+use ty_project::ProjectDatabase;
+
+use crate::db::Db;
 use crate::server::api::semantic_tokens::generate_semantic_tokens;
 use crate::server::api::traits::{
     BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler,
 };
 use crate::session::DocumentSnapshot;
 use crate::session::client::Client;
-use lsp_types::{SemanticTokens, SemanticTokensParams, SemanticTokensResult, Url};
-use ty_project::ProjectDatabase;
 
 pub(crate) struct SemanticTokensRequestHandler;
 
@@ -33,14 +36,29 @@ impl BackgroundDocumentRequestHandler for SemanticTokensRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.to_file(db) else {
+        let Some(file) = snapshot.to_notebook_or_file(db) else {
             return Ok(None);
         };
 
+        // If this document is a notebook cell, limit the highlighting range
+        // to the lines of this cell (instead of highlighting the entire notebook).
+        // Not only avoids this unnecessary work, this is also required
+        // because all ranges in the response must be within this **this document**.
+        let mut cell_range = None;
+
+        if snapshot.document().is_cell()
+            && let Some(notebook_document) = db.notebook_document(file)
+            && let Some(notebook) = source_text(db, file).as_notebook()
+        {
+            let cell_index = notebook_document.cell_index_by_uri(snapshot.url());
+
+            cell_range = cell_index.and_then(|index| notebook.cell_range(index));
+        }
+
         let lsp_tokens = generate_semantic_tokens(
             db,
             file,
-            None,
+            cell_range,
             snapshot.encoding(),
             snapshot
                 .resolved_client_capabilities()

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

@@ -1,5 +1,8 @@
 use std::borrow::Cow;
 
+use lsp_types::{SemanticTokens, SemanticTokensRangeParams, SemanticTokensRangeResult, Url};
+use ty_project::ProjectDatabase;
+
 use crate::document::RangeExt;
 use crate::server::api::semantic_tokens::generate_semantic_tokens;
 use crate::server::api::traits::{
@@ -7,8 +10,6 @@ use crate::server::api::traits::{
 };
 use crate::session::DocumentSnapshot;
 use crate::session::client::Client;
-use lsp_types::{SemanticTokens, SemanticTokensRangeParams, SemanticTokensRangeResult, Url};
-use ty_project::ProjectDatabase;
 
 pub(crate) struct SemanticTokensRangeRequestHandler;
 
@@ -34,7 +35,7 @@ impl BackgroundDocumentRequestHandler for SemanticTokensRangeRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.to_file(db) else {
+        let Some(file) = snapshot.to_notebook_or_file(db) else {
             return Ok(None);
         };
 

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

@@ -38,7 +38,7 @@ impl BackgroundDocumentRequestHandler for SignatureHelpRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.to_file(db) else {
+        let Some(file) = snapshot.to_notebook_or_file(db) else {
             return Ok(None);
         };
 

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

@@ -1,3 +1,23 @@
+use std::collections::BTreeMap;
+use std::sync::Mutex;
+use std::time::{Duration, Instant};
+
+use lsp_server::RequestId;
+use lsp_types::request::WorkspaceDiagnosticRequest;
+use lsp_types::{
+    FullDocumentDiagnosticReport, PreviousResultId, ProgressToken,
+    UnchangedDocumentDiagnosticReport, Url, WorkspaceDiagnosticParams, WorkspaceDiagnosticReport,
+    WorkspaceDiagnosticReportPartialResult, WorkspaceDiagnosticReportResult,
+    WorkspaceDocumentDiagnosticReport, WorkspaceFullDocumentDiagnosticReport,
+    WorkspaceUnchangedDocumentDiagnosticReport, notification::Notification,
+};
+use ruff_db::diagnostic::Diagnostic;
+use ruff_db::files::File;
+use ruff_db::source::source_text;
+use rustc_hash::FxHashMap;
+use serde::{Deserialize, Serialize};
+use ty_project::{ProgressReporter, ProjectDatabase};
+
 use crate::PositionEncoding;
 use crate::document::DocumentKey;
 use crate::server::api::diagnostics::{Diagnostics, to_lsp_diagnostic};
@@ -10,23 +30,6 @@ use crate::session::client::Client;
 use crate::session::index::Index;
 use crate::session::{SessionSnapshot, SuspendedWorkspaceDiagnosticRequest};
 use crate::system::file_to_url;
-use lsp_server::RequestId;
-use lsp_types::request::WorkspaceDiagnosticRequest;
-use lsp_types::{
-    FullDocumentDiagnosticReport, PreviousResultId, ProgressToken,
-    UnchangedDocumentDiagnosticReport, Url, WorkspaceDiagnosticParams, WorkspaceDiagnosticReport,
-    WorkspaceDiagnosticReportPartialResult, WorkspaceDiagnosticReportResult,
-    WorkspaceDocumentDiagnosticReport, WorkspaceFullDocumentDiagnosticReport,
-    WorkspaceUnchangedDocumentDiagnosticReport, notification::Notification,
-};
-use ruff_db::diagnostic::Diagnostic;
-use ruff_db::files::File;
-use rustc_hash::FxHashMap;
-use serde::{Deserialize, Serialize};
-use std::collections::BTreeMap;
-use std::sync::Mutex;
-use std::time::{Duration, Instant};
-use ty_project::{ProgressReporter, ProjectDatabase};
 
 /// Handler for [Workspace diagnostics](workspace-diagnostics)
 ///
@@ -368,6 +371,15 @@ impl<'a> ResponseWriter<'a> {
             tracing::debug!("Failed to convert file path to URL at {}", file.path(db));
             return;
         };
+
+        if source_text(db, file).is_notebook() {
+            // Notebooks only support publish diagnostics.
+            // and we can't convert text ranges to notebook ranges unless
+            // the document is open in the editor, in which case
+            // we publish the diagnostics already.
+            return;
+        }
+
         let key = DocumentKey::from_url(&url);
         let version = self
             .index
@@ -394,7 +406,7 @@ impl<'a> ResponseWriter<'a> {
             new_id => {
                 let lsp_diagnostics = diagnostics
                     .iter()
-                    .map(|diagnostic| to_lsp_diagnostic(db, diagnostic, self.position_encoding))
+                    .map(|diagnostic| to_lsp_diagnostic(db, diagnostic, self.position_encoding).1)
                     .collect::<Vec<_>>();
 
                 WorkspaceDocumentDiagnosticReport::Full(WorkspaceFullDocumentDiagnosticReport {

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

@@ -26,6 +26,9 @@ pub(crate) fn convert_symbol_kind(kind: ty_ide::SymbolKind) -> SymbolKind {
 }
 
 /// Convert a `ty_ide` `SymbolInfo` to LSP `SymbolInformation`
+///
+/// Returns `None` if the symbol's range cannot be converted to a location
+/// (e.g., if the file cannot be converted to a URL).
 pub(crate) fn convert_to_lsp_symbol_information(
     db: &dyn Db,
     file: ruff_db::files::File,
@@ -33,12 +36,10 @@ pub(crate) fn convert_to_lsp_symbol_information(
     encoding: PositionEncoding,
 ) -> Option<SymbolInformation> {
     let symbol_kind = convert_symbol_kind(symbol.kind);
-
     let location = symbol
         .full_range
         .to_lsp_range(db, file, encoding)?
         .to_location()?;
-
     Some(SymbolInformation {
         name: symbol.name.into_owned(),
         kind: symbol_kind,

crates/ty_server/src/session.rs

@@ -1,7 +1,11 @@
 //! Data model, state management, and configuration resolution.
 
+use std::collections::{BTreeMap, HashSet, VecDeque};
+use std::ops::{Deref, DerefMut};
+use std::panic::RefUnwindSafe;
+use std::sync::Arc;
+
 use anyhow::{Context, anyhow};
-use index::DocumentError;
 use lsp_server::{Message, RequestId};
 use lsp_types::notification::{DidChangeWatchedFiles, Exit, Notification};
 use lsp_types::request::{
@@ -13,20 +17,17 @@ use lsp_types::{
     DidChangeWatchedFilesRegistrationOptions, FileSystemWatcher, Registration, RegistrationParams,
     TextDocumentContentChangeEvent, Unregistration, UnregistrationParams, Url,
 };
-use options::GlobalOptions;
 use ruff_db::Db;
 use ruff_db::files::{File, system_path_to_file};
 use ruff_db::system::{System, SystemPath, SystemPathBuf};
-use std::borrow::Cow;
-use std::collections::{BTreeMap, HashSet, VecDeque};
-use std::ops::{Deref, DerefMut};
-use std::panic::RefUnwindSafe;
-use std::sync::Arc;
 use ty_combine::Combine;
 use ty_project::metadata::Options;
-use ty_project::watch::ChangeEvent;
+use ty_project::watch::{ChangeEvent, CreatedKind};
 use ty_project::{ChangeResult, CheckMode, Db as _, ProjectDatabase, ProjectMetadata};
 
+use index::DocumentError;
+use options::GlobalOptions;
+
 pub(crate) use self::options::InitializationOptions;
 pub use self::options::{ClientOptions, DiagnosticMode};
 pub(crate) use self::settings::{GlobalSettings, WorkspaceSettings};
@@ -36,6 +37,7 @@ use crate::capabilities::{
 use crate::document::{DocumentKey, DocumentVersion, NotebookDocument};
 use crate::server::{Action, publish_settings_diagnostics};
 use crate::session::client::Client;
+use crate::session::index::Document;
 use crate::session::request_queue::RequestQueue;
 use crate::system::{AnySystemPath, LSPSystem};
 use crate::{PositionEncoding, TextDocument};
@@ -816,25 +818,16 @@ impl Session {
         let index = self.index();
         let document_handle = index.document_handle(url)?;
 
-        let notebook = if let Some(notebook_path) = &document_handle.notebook_path {
-            index
-                .notebook_arc(&DocumentKey::from(notebook_path.clone()))
-                .ok()
-        } else {
-            None
-        };
-
         Ok(DocumentSnapshot {
             resolved_client_capabilities: self.resolved_client_capabilities,
             global_settings: self.global_settings.clone(),
             workspace_settings: document_handle
-                .to_file_path()
+                .notebook_or_file_path()
                 .as_system()
                 .and_then(|path| self.workspaces.settings_for_path(path))
                 .unwrap_or_else(|| Arc::new(WorkspaceSettings::default())),
             position_encoding: self.position_encoding,
             document: document_handle,
-            notebook,
         })
     }
 
@@ -860,13 +853,7 @@ impl Session {
     pub(super) fn text_document_handles(&self) -> impl Iterator<Item = DocumentHandle> + '_ {
         self.index()
             .text_documents()
-            .map(|(key, document)| DocumentHandle {
+            .map(|(_, document)| DocumentHandle::from_text_document(document))
-                key: key.clone(),
-                url: document.url().clone(),
-                version: document.version(),
-                // TODO: Set notebook path if text document is part of a notebook
-                notebook_path: None,
-            })
     }
 
     /// Returns a handle to the document specified by its URL.
@@ -887,7 +874,7 @@ impl Session {
     /// Returns a handle to the opened document.
     pub(crate) fn open_notebook_document(&mut self, document: NotebookDocument) -> DocumentHandle {
         let handle = self.index_mut().open_notebook_document(document);
-        self.bump_revision();
+        self.open_document_in_db(&handle);
         handle
     }
 
@@ -897,9 +884,49 @@ impl Session {
     /// Returns a handle to the opened document.
     pub(crate) fn open_text_document(&mut self, document: TextDocument) -> DocumentHandle {
         let handle = self.index_mut().open_text_document(document);
+        self.open_document_in_db(&handle);
+        handle
+    }
+
+    fn open_document_in_db(&mut self, document: &DocumentHandle) {
+        let path = document.notebook_or_file_path();
+
+        // This is a "maybe" because the `File` might've not been interned yet i.e., the
+        // `try_system` call will return `None` which doesn't mean that the file is new, it's just
+        // that the server didn't need the file yet.
+        let is_maybe_new_system_file = path.as_system().is_some_and(|system_path| {
+            let db = self.project_db(path);
+            db.files()
+                .try_system(db, system_path)
+                .is_none_or(|file| !file.exists(db))
+        });
+
+        match path {
+            AnySystemPath::System(system_path) => {
+                let event = if is_maybe_new_system_file {
+                    ChangeEvent::Created {
+                        path: system_path.clone(),
+                        kind: CreatedKind::File,
+                    }
+                } else {
+                    ChangeEvent::Opened(system_path.clone())
+                };
+                self.apply_changes(path, vec![event]);
+
+                let db = self.project_db_mut(path);
+                match system_path_to_file(db, system_path) {
+                    Ok(file) => db.project().open_file(db, file),
+                    Err(err) => tracing::warn!("Failed to open file {system_path}: {err}"),
+                }
+            }
+            AnySystemPath::SystemVirtual(virtual_path) => {
+                let db = self.project_db_mut(path);
+                let virtual_file = db.files().virtual_file(db, virtual_path);
+                db.project().open_file(db, virtual_file.file());
+            }
+        }
 
         self.bump_revision();
-        handle
     }
 
     /// Returns a reference to the index.
@@ -999,7 +1026,6 @@ pub(crate) struct DocumentSnapshot {
     workspace_settings: Arc<WorkspaceSettings>,
     position_encoding: PositionEncoding,
     document: DocumentHandle,
-    notebook: Option<Arc<NotebookDocument>>,
 }
 
 impl DocumentSnapshot {
@@ -1028,17 +1054,12 @@ impl DocumentSnapshot {
         &self.document
     }
 
-    /// Returns the URL of the document.
     pub(crate) fn url(&self) -> &lsp_types::Url {
         self.document.url()
     }
 
-    pub(crate) fn notebook(&self) -> Option<&NotebookDocument> {
+    pub(crate) fn to_notebook_or_file(&self, db: &dyn Db) -> Option<File> {
-        self.notebook.as_deref()
+        let file = self.document.notebook_or_file(db);
-    }
-
-    pub(crate) fn to_file(&self, db: &dyn Db) -> Option<File> {
-        let file = self.document.to_file(db);
         if file.is_none() {
             tracing::debug!(
                 "Failed to resolve file: file not found for `{}`",
@@ -1048,8 +1069,8 @@ impl DocumentSnapshot {
         file
     }
 
-    pub(crate) fn to_file_path(&self) -> Cow<'_, AnySystemPath> {
+    pub(crate) fn notebook_or_file_path(&self) -> &AnySystemPath {
-        self.document.to_file_path()
+        self.document.notebook_or_file_path()
     }
 }
 
@@ -1330,34 +1351,99 @@ impl SuspendedWorkspaceDiagnosticRequest {
 ///
 /// It also exposes methods to get the file-path of the corresponding ty-file.
 #[derive(Clone, Debug)]
-pub(crate) struct DocumentHandle {
+pub(crate) enum DocumentHandle {
-    /// The key that uniquely identifies this document in the index.
+    Text {
-    key: DocumentKey,
         url: lsp_types::Url,
-    /// The path to the enclosing notebook file if this document is a notebook or a notebook cell.
+        path: AnySystemPath,
-    notebook_path: Option<AnySystemPath>,
         version: DocumentVersion,
+    },
+    Notebook {
+        url: lsp_types::Url,
+        path: AnySystemPath,
+        version: DocumentVersion,
+    },
+    Cell {
+        url: lsp_types::Url,
+        version: DocumentVersion,
+        notebook_path: AnySystemPath,
+    },
 }
 
 impl DocumentHandle {
+    fn from_text_document(document: &TextDocument) -> Self {
+        match document.notebook() {
+            None => Self::Text {
+                version: document.version(),
+                url: document.url().clone(),
+                path: DocumentKey::from_url(document.url()).into_file_path(),
+            },
+            Some(notebook) => Self::Cell {
+                notebook_path: notebook.clone(),
+                version: document.version(),
+                url: document.url().clone(),
+            },
+        }
+    }
+
+    fn from_notebook_document(document: &NotebookDocument) -> Self {
+        Self::Notebook {
+            path: DocumentKey::from_url(document.url()).into_file_path(),
+            url: document.url().clone(),
+            version: document.version(),
+        }
+    }
+
+    fn from_document(document: &Document) -> Self {
+        match document {
+            Document::Text(text) => Self::from_text_document(text),
+            Document::Notebook(notebook) => Self::from_notebook_document(notebook),
+        }
+    }
+
+    fn key(&self) -> DocumentKey {
+        DocumentKey::from_url(self.url())
+    }
+
     pub(crate) const fn version(&self) -> DocumentVersion {
-        self.version
+        match self {
+            Self::Text { version, .. }
+            | Self::Notebook { version, .. }
+            | Self::Cell { version, .. } => *version,
+        }
     }
 
     /// The URL as used by the client to reference this document.
     pub(crate) fn url(&self) -> &lsp_types::Url {
-        &self.url
+        match self {
+            Self::Text { url, .. } | Self::Notebook { url, .. } | Self::Cell { url, .. } => url,
+        }
     }
 
     /// The path to the enclosing file for this document.
     ///
     /// This is the path corresponding to the URL, except for notebook cells where the
     /// path corresponds to the notebook file.
-    pub(crate) fn to_file_path(&self) -> Cow<'_, AnySystemPath> {
+    pub(crate) fn notebook_or_file_path(&self) -> &AnySystemPath {
-        if let Some(path) = self.notebook_path.as_ref() {
+        match self {
-            Cow::Borrowed(path)
+            Self::Text { path, .. } | Self::Notebook { path, .. } => path,
-        } else {
+            Self::Cell { notebook_path, .. } => notebook_path,
-            Cow::Owned(self.key.to_file_path())
+        }
+    }
+
+    #[expect(unused)]
+    pub(crate) fn file_path(&self) -> Option<&AnySystemPath> {
+        match self {
+            Self::Text { path, .. } | Self::Notebook { path, .. } => Some(path),
+            Self::Cell { .. } => None,
+        }
+    }
+
+    #[expect(unused)]
+    pub(crate) fn notebook_path(&self) -> Option<&AnySystemPath> {
+        match self {
+            DocumentHandle::Notebook { path, .. } => Some(path),
+            DocumentHandle::Cell { notebook_path, .. } => Some(notebook_path),
+            DocumentHandle::Text { .. } => None,
         }
     }
 
@@ -1366,8 +1452,8 @@ impl DocumentHandle {
     /// It returns [`None`] for the following cases:
     /// - For virtual file, if it's not yet opened
     /// - For regular file, if it does not exists or is a directory
-    pub(crate) fn to_file(&self, db: &dyn Db) -> Option<File> {
+    pub(crate) fn notebook_or_file(&self, db: &dyn Db) -> Option<File> {
-        match &*self.to_file_path() {
+        match &self.notebook_or_file_path() {
             AnySystemPath::System(path) => system_path_to_file(db, path).ok(),
             AnySystemPath::SystemVirtual(virtual_path) => db
                 .files()
@@ -1376,6 +1462,14 @@ impl DocumentHandle {
         }
     }
 
+    pub(crate) fn is_cell(&self) -> bool {
+        matches!(self, Self::Cell { .. })
+    }
+
+    pub(crate) fn is_cell_or_notebook(&self) -> bool {
+        matches!(self, Self::Cell { .. } | Self::Notebook { .. })
+    }
+
     pub(crate) fn update_text_document(
         &self,
         session: &mut Session,
@@ -1383,9 +1477,10 @@ impl DocumentHandle {
         new_version: DocumentVersion,
     ) -> crate::Result<()> {
         let position_encoding = session.position_encoding();
+        {
             let mut index = session.index_mut();
 
-        let document_mut = index.document_mut(&self.key)?;
+            let document_mut = index.document_mut(&self.key())?;
 
             let Some(document) = document_mut.as_text_mut() else {
                 anyhow::bail!("Text document path does not point to a text document");
@@ -1393,19 +1488,110 @@ impl DocumentHandle {
 
             if content_changes.is_empty() {
                 document.update_version(new_version);
-            return Ok(());
+            } else {
+                document.apply_changes(content_changes, new_version, position_encoding);
+            }
         }
 
-        document.apply_changes(content_changes, new_version, position_encoding);
+        self.update_in_db(session);
 
         Ok(())
     }
 
+    pub(crate) fn update_notebook_document(
+        &self,
+        session: &mut Session,
+        cells: Option<lsp_types::NotebookDocumentCellChange>,
+        metadata: Option<lsp_types::LSPObject>,
+        new_version: DocumentVersion,
+    ) -> crate::Result<()> {
+        let position_encoding = session.position_encoding();
+        {
+            let mut index = session.index_mut();
+
+            index.update_notebook_document(
+                &self.key(),
+                cells,
+                metadata,
+                new_version,
+                position_encoding,
+            )?;
+        }
+
+        self.update_in_db(session);
+        Ok(())
+    }
+
+    fn update_in_db(&self, session: &mut Session) {
+        let path = self.notebook_or_file_path();
+        let changes = match path {
+            AnySystemPath::System(system_path) => {
+                vec![ChangeEvent::file_content_changed(system_path.clone())]
+            }
+            AnySystemPath::SystemVirtual(virtual_path) => {
+                vec![ChangeEvent::ChangedVirtual(virtual_path.clone())]
+            }
+        };
+
+        session.apply_changes(path, changes);
+    }
+
     /// De-registers a document, specified by its key.
     /// Calling this multiple times for the same document is a logic error.
-    pub(crate) fn close(self, session: &mut Session) -> crate::Result<()> {
+    ///
-        session.index_mut().close_document(&self.key)?;
+    /// Returns `true` if the client needs to clear the diagnostics for this document.
+    pub(crate) fn close(&self, session: &mut Session) -> crate::Result<bool> {
+        let is_cell = self.is_cell();
+        let path = self.notebook_or_file_path();
+        session.index_mut().close_document(&self.key())?;
+
+        // Close the text or notebook file in the database but skip this
+        // step for cells because closing a cell doesn't close its notebook.
+        let requires_clear_diagnostics = if is_cell {
+            true
+        } else {
+            let db = session.project_db_mut(path);
+
+            match path {
+                AnySystemPath::System(system_path) => {
+                    if let Some(file) = db.files().try_system(db, system_path) {
+                        db.project().close_file(db, file);
+                    } else {
+                        // This can only fail when the path is a directory or it doesn't exists but the
+                        // file should exists for this handler in this branch. This is because every
+                        // close call is preceded by an open call, which ensures that the file is
+                        // interned in the lookup table (`Files`).
+                        tracing::warn!("Salsa file does not exists for {}", system_path);
+                    }
+
+                    // For non-virtual files, we clear diagnostics if:
+                    //
+                    // 1. The file does not belong to any workspace e.g., opening a random file from
+                    //    outside the workspace because closing it acts like the file doesn't exists
+                    // 2. The diagnostic mode is set to open-files only
+                    session.workspaces().for_path(system_path).is_none()
+                        || session
+                            .global_settings()
+                            .diagnostic_mode()
+                            .is_open_files_only()
+                }
+                AnySystemPath::SystemVirtual(virtual_path) => {
+                    if let Some(virtual_file) = db.files().try_virtual_file(virtual_path) {
+                        db.project().close_file(db, virtual_file.file());
+                        virtual_file.close(db);
+                    } else {
+                        tracing::warn!("Salsa virtual file does not exists for {}", virtual_path);
+                    }
+
+                    // Always clear diagnostics for virtual files, as they don't really exist on disk
+                    // which means closing them is like deleting the file.
+                    true
+                }
+            }
+        };
+
         session.bump_revision();
-        Ok(())
+
+        Ok(requires_clear_diagnostics)
     }
 }

crates/ty_server/src/session/index.rs

@@ -1,3 +1,4 @@
+use rustc_hash::FxHashMap;
 use std::sync::Arc;
 
 use crate::document::DocumentKey;
@@ -5,27 +6,19 @@ use crate::session::DocumentHandle;
 use crate::{
     PositionEncoding, TextDocument,
     document::{DocumentVersion, NotebookDocument},
-    system::AnySystemPath,
 };
 
-use ruff_db::system::SystemVirtualPath;
-use rustc_hash::FxHashMap;
-
 /// Stores and tracks all open documents in a session, along with their associated settings.
 #[derive(Debug)]
 pub(crate) struct Index {
     /// Maps all document file paths to the associated document controller
     documents: FxHashMap<DocumentKey, Document>,
-
-    /// Maps opaque cell URLs to a notebook path (document)
-    notebook_cells: FxHashMap<String, AnySystemPath>,
 }
 
 impl Index {
     pub(super) fn new() -> Self {
         Self {
             documents: FxHashMap::default(),
-            notebook_cells: FxHashMap::default(),
         }
     }
 
@@ -47,23 +40,7 @@ impl Index {
             return Err(DocumentError::NotFound(key));
         };
 
-        if let Some(path) = key.as_opaque() {
+        Ok(DocumentHandle::from_document(document))
-            if let Some(notebook_path) = self.notebook_cells.get(path) {
-                return Ok(DocumentHandle {
-                    key: key.clone(),
-                    notebook_path: Some(notebook_path.clone()),
-                    url: url.clone(),
-                    version: document.version(),
-                });
-            }
-        }
-
-        Ok(DocumentHandle {
-            key: key.clone(),
-            notebook_path: None,
-            url: url.clone(),
-            version: document.version(),
-        })
     }
 
     #[expect(dead_code)]
@@ -74,7 +51,6 @@ impl Index {
             .map(|(key, _)| key)
     }
 
-    #[expect(dead_code)]
     pub(super) fn update_notebook_document(
         &mut self,
         notebook_key: &DocumentKey,
@@ -83,26 +59,100 @@ impl Index {
         new_version: DocumentVersion,
         encoding: PositionEncoding,
     ) -> crate::Result<()> {
-        // update notebook cell index
-        if let Some(lsp_types::NotebookDocumentCellChangeStructure {
-            did_open: Some(did_open),
-            ..
-        }) = cells.as_ref().and_then(|cells| cells.structure.as_ref())
-        {
-            for opened_cell in did_open {
-                let cell_path = SystemVirtualPath::new(opened_cell.uri.as_str());
-                self.notebook_cells
-                    .insert(cell_path.to_string(), notebook_key.to_file_path());
-            }
-            // deleted notebook cells are closed via textDocument/didClose - we don't close them here.
-        }
-
         let document = self.document_mut(notebook_key)?;
         let Some(notebook) = document.as_notebook_mut() else {
             anyhow::bail!("Notebook document path does not point to a notebook document");
         };
 
-        notebook.update(cells, metadata, new_version, encoding)?;
+        let (structure, data, text_content) = cells
+            .map(|cells| {
+                let lsp_types::NotebookDocumentCellChange {
+                    structure,
+                    data,
+                    text_content,
+                } = cells;
+                (structure, data, text_content)
+            })
+            .unwrap_or_default();
+
+        let (array, did_open, did_close) = structure
+            .map(|structure| {
+                let lsp_types::NotebookDocumentCellChangeStructure {
+                    array,
+                    did_open,
+                    did_close,
+                } = structure;
+
+                (array, did_open, did_close)
+            })
+            .unwrap_or_else(|| {
+                (
+                    lsp_types::NotebookCellArrayChange {
+                        start: 0,
+                        delete_count: 0,
+                        cells: None,
+                    },
+                    None,
+                    None,
+                )
+            });
+
+        tracing::info!(
+            "version: {}, new_version: {}",
+            notebook.version(),
+            new_version
+        );
+        notebook.update(array, data.unwrap_or_default(), metadata, new_version)?;
+
+        let notebook_path = notebook_key.to_file_path();
+
+        for opened_cell in did_open.into_iter().flatten() {
+            self.documents.insert(
+                DocumentKey::from_url(&opened_cell.uri),
+                Document::Text(
+                    TextDocument::new(opened_cell.uri, opened_cell.text, opened_cell.version)
+                        .with_language_id(&opened_cell.language_id)
+                        .with_notebook(notebook_path.clone())
+                        .into(),
+                ),
+            );
+        }
+
+        for updated_cell in text_content.into_iter().flatten() {
+            let Ok(document_mut) =
+                self.document_mut(&DocumentKey::from_url(&updated_cell.document.uri))
+            else {
+                tracing::warn!(
+                    "Could not find document for cell {}",
+                    updated_cell.document.uri
+                );
+                continue;
+            };
+
+            let Some(document) = document_mut.as_text_mut() else {
+                continue;
+            };
+
+            if updated_cell.changes.is_empty() {
+                document.update_version(updated_cell.document.version);
+            } else {
+                document.apply_changes(
+                    updated_cell.changes,
+                    updated_cell.document.version,
+                    encoding,
+                );
+            }
+        }
+
+        // VS Code sends a separate `didClose` request for every cell
+        // and they're removed from the metadata (notebook document)
+        // because they get deleted as part of `change.cells.structure.array`
+        let _ = did_close;
+
+        let notebook = self.document(notebook_key).unwrap().as_notebook().unwrap();
+        let ruff_notebook = notebook.to_ruff_notebook(self);
+        tracing::debug!("Updated notebook: {:?}", ruff_notebook.source_code());
+
         Ok(())
     }
 
@@ -117,31 +167,10 @@ impl Index {
         Ok(document)
     }
 
-    pub(crate) fn notebook_arc(
-        &self,
-        key: &DocumentKey,
-    ) -> Result<Arc<NotebookDocument>, DocumentError> {
-        let Some(document) = self.documents.get(key) else {
-            return Err(DocumentError::NotFound(key.clone()));
-        };
-
-        if let Document::Notebook(notebook) = document {
-            Ok(notebook.clone())
-        } else {
-            Err(DocumentError::NotFound(key.clone()))
-        }
-    }
-
     pub(super) fn open_text_document(&mut self, document: TextDocument) -> DocumentHandle {
         let key = DocumentKey::from_url(document.url());
 
-        // TODO: Fix file path for notebook cells
+        let handle = DocumentHandle::from_text_document(&document);
-        let handle = DocumentHandle {
-            key: key.clone(),
-            notebook_path: None,
-            url: document.url().clone(),
-            version: document.version(),
-        };
 
         self.documents.insert(key, Document::new_text(document));
 
@@ -149,38 +178,18 @@ impl Index {
     }
 
     pub(super) fn open_notebook_document(&mut self, document: NotebookDocument) -> DocumentHandle {
+        let handle = DocumentHandle::from_notebook_document(&document);
         let notebook_key = DocumentKey::from_url(document.url());
-        let url = document.url().clone();
-        let version = document.version();
-
-        for cell_url in document.cell_urls() {
-            self.notebook_cells
-                .insert(cell_url.to_string(), notebook_key.to_file_path());
-        }
 
         self.documents
-            .insert(notebook_key.clone(), Document::new_notebook(document));
+            .insert(notebook_key, Document::new_notebook(document));
 
-        DocumentHandle {
+        handle
-            notebook_path: Some(notebook_key.to_file_path()),
-            key: notebook_key,
-            url,
-            version,
-        }
-    }
-
-    pub(super) fn close_document(&mut self, key: &DocumentKey) -> crate::Result<()> {
-        // Notebook cells URIs are removed from the index here, instead of during
-        // `update_notebook_document`. This is because a notebook cell, as a text document,
-        // is requested to be `closed` by VS Code after the notebook gets updated.
-        // This is not documented in the LSP specification explicitly, and this assumption
-        // may need revisiting in the future as we support more editors with notebook support.
-        if let DocumentKey::Opaque(uri) = key {
-            self.notebook_cells.remove(uri);
     }
 
+    pub(super) fn close_document(&mut self, key: &DocumentKey) -> Result<(), DocumentError> {
         let Some(_) = self.documents.remove(key) else {
-            anyhow::bail!("tried to close document that didn't exist at {key}")
+            return Err(DocumentError::NotFound(key.clone()));
         };
 
         Ok(())

crates/ty_server/src/system.rs

@@ -1,11 +1,12 @@
 use std::any::Any;
 use std::fmt;
 use std::fmt::Display;
+use std::hash::{DefaultHasher, Hash, Hasher as _};
 use std::panic::RefUnwindSafe;
 use std::sync::Arc;
 
 use crate::Db;
-use crate::document::DocumentKey;
+use crate::document::{DocumentKey, LanguageId};
 use crate::session::index::{Document, Index};
 use lsp_types::Url;
 use ruff_db::file_revision::FileRevision;
@@ -16,6 +17,7 @@ use ruff_db::system::{
     SystemPath, SystemPathBuf, SystemVirtualPath, SystemVirtualPathBuf, WritableSystem,
 };
 use ruff_notebook::{Notebook, NotebookError};
+use ruff_python_ast::PySourceType;
 use ty_ide::cached_vendored_path;
 
 /// Returns a [`Url`] for the given [`File`].
@@ -117,6 +119,23 @@ impl LSPSystem {
         index.document(&DocumentKey::from(path)).ok()
     }
 
+    fn source_type_from_document(
+        document: &Document,
+        extension: Option<&str>,
+    ) -> Option<PySourceType> {
+        match document {
+            Document::Text(text) => match text.language_id()? {
+                LanguageId::Python => Some(
+                    extension
+                        .and_then(PySourceType::try_from_extension)
+                        .unwrap_or(PySourceType::Python),
+                ),
+                LanguageId::Other => None,
+            },
+            Document::Notebook(_) => Some(PySourceType::Ipynb),
+        }
+    }
+
     pub(crate) fn system_path_to_document(&self, path: &SystemPath) -> Option<&Document> {
         let any_path = AnySystemPath::System(path.to_path_buf());
         self.document(any_path)
@@ -137,7 +156,7 @@ impl System for LSPSystem {
 
         if let Some(document) = document {
             Ok(Metadata::new(
-                document_revision(document),
+                document_revision(document, self.index()),
                 None,
                 FileType::File,
             ))
@@ -154,6 +173,16 @@ impl System for LSPSystem {
         self.native_system.path_exists_case_sensitive(path, prefix)
     }
 
+    fn source_type(&self, path: &SystemPath) -> Option<PySourceType> {
+        let document = self.system_path_to_document(path)?;
+        Self::source_type_from_document(document, path.extension())
+    }
+
+    fn virtual_path_source_type(&self, path: &SystemVirtualPath) -> Option<PySourceType> {
+        let document = self.system_virtual_path_to_document(path)?;
+        Self::source_type_from_document(document, path.extension())
+    }
+
     fn read_to_string(&self, path: &SystemPath) -> Result<String> {
         let document = self.system_path_to_document(path);
 
@@ -168,7 +197,7 @@ impl System for LSPSystem {
 
         match document {
             Some(Document::Text(document)) => Notebook::from_source_code(document.contents()),
-            Some(Document::Notebook(notebook)) => Ok(notebook.make_ruff_notebook()),
+            Some(Document::Notebook(notebook)) => Ok(notebook.to_ruff_notebook(self.index())),
             None => self.native_system.read_to_notebook(path),
         }
     }
@@ -195,7 +224,7 @@ impl System for LSPSystem {
 
         match document {
             Document::Text(document) => Notebook::from_source_code(document.contents()),
-            Document::Notebook(notebook) => Ok(notebook.make_ruff_notebook()),
+            Document::Notebook(notebook) => Ok(notebook.to_ruff_notebook(self.index())),
         }
     }
 
@@ -272,9 +301,33 @@ fn virtual_path_not_found(path: impl Display) -> std::io::Error {
 }
 
 /// Helper function to get the [`FileRevision`] of the given document.
-fn document_revision(document: &Document) -> FileRevision {
+fn document_revision(document: &Document, index: &Index) -> FileRevision {
     // The file revision is just an opaque number which doesn't have any significant meaning other
     // than that the file has changed if the revisions are different.
     #[expect(clippy::cast_sign_loss)]
-    FileRevision::new(document.version() as u128)
+    match document {
+        Document::Text(text) => FileRevision::new(text.version() as u128),
+        Document::Notebook(notebook) => {
+            // VS Code doesn't always bump the notebook version when the cell content changes.
+            // Specifically, I noticed that VS Code re-uses the same version when:
+            // 1. Adding a new cell
+            // 2. Pasting some code that has an error
+            //
+            // The notification updating the cell content on paste re-used the same version as when the cell was added.
+            // Because of that, hash all cell versions and the notebook versions together.
+            let mut hasher = DefaultHasher::new();
+            for cell_url in notebook.cell_urls() {
+                if let Ok(cell) = index.document(&DocumentKey::from_url(cell_url)) {
+                    cell.version().hash(&mut hasher);
+                }
+            }
+
+            // Use higher 64 bits for notebook version and lower 64 bits for cell revisions
+            let notebook_version_high = (notebook.version() as u128) << 64;
+            let cell_versions_low = u128::from(hasher.finish()) & 0xFFFF_FFFF_FFFF_FFFF;
+            let combined_revision = notebook_version_high | cell_versions_low;
+
+            FileRevision::new(combined_revision)
+        }
+    }
 }

crates/ty_server/tests/e2e/main.rs

@@ -30,11 +30,12 @@
 mod commands;
 mod initialize;
 mod inlay_hints;
+mod notebook;
 mod publish_diagnostics;
 mod pull_diagnostics;
 
 use std::collections::hash_map::Entry;
-use std::collections::{HashMap, VecDeque};
+use std::collections::{BTreeMap, HashMap, VecDeque};
 use std::num::NonZeroUsize;
 use std::sync::{Arc, OnceLock};
 use std::thread::JoinHandle;
@@ -433,6 +434,33 @@ impl TestServer {
         ))
     }
 
+    /// Collects `N` publish diagnostic notifications into a map, indexed by the document url.
+    ///
+    /// ## Panics
+    /// If there are multiple publish diagnostics notifications for the same document.
+    pub(crate) fn collect_publish_diagnostic_notifications(
+        &mut self,
+        count: usize,
+    ) -> Result<BTreeMap<lsp_types::Url, Vec<lsp_types::Diagnostic>>> {
+        let mut results = BTreeMap::default();
+
+        for _ in 0..count {
+            let notification =
+                self.await_notification::<lsp_types::notification::PublishDiagnostics>()?;
+
+            if let Some(existing) =
+                results.insert(notification.uri.clone(), notification.diagnostics)
+            {
+                panic!(
+                    "Received multiple publish diagnostic notifications for {url}: ({existing:#?})",
+                    url = &notification.uri
+                );
+            }
+        }
+
+        Ok(results)
+    }
+
     /// Wait for a request of the specified type from the server and return the request ID and
     /// parameters.
     ///
@@ -774,6 +802,7 @@ impl Drop for TestServer {
             match self.await_response::<Shutdown>(&shutdown_id) {
                 Ok(()) => {
                     self.send_notification::<Exit>(());
+
                     None
                 }
                 Err(err) => Some(format!("Failed to get shutdown response: {err:?}")),
@@ -782,9 +811,32 @@ impl Drop for TestServer {
             None
         };
 
-        if let Some(_client_connection) = self.client_connection.take() {
         // Drop the client connection before joining the server thread to avoid any hangs
         // in case the server didn't respond to the shutdown request.
+        if let Some(client_connection) = self.client_connection.take() {
+            if !std::thread::panicking() {
+                // Wait for the client sender to drop (confirmation that it processed the exit notification).
+
+                match client_connection
+                    .receiver
+                    .recv_timeout(Duration::from_secs(20))
+                {
+                    Err(RecvTimeoutError::Disconnected) => {
+                        // Good, the server terminated
+                    }
+                    Err(RecvTimeoutError::Timeout) => {
+                        tracing::warn!(
+                            "The server didn't exit within 20ms after receiving the EXIT notification"
+                        );
+                    }
+                    Ok(message) => {
+                        // Ignore any errors: A duplicate pending message
+                        // won't matter that much because `assert_no_pending_messages` will
+                        // panic anyway.
+                        let _ = self.handle_message(message);
+                    }
+                }
+            }
         }
 
         if std::thread::panicking() {

crates/ty_server/tests/e2e/notebook.rs

@@ -0,0 +1,361 @@
+use insta::assert_json_snapshot;
+use lsp_types::{NotebookCellKind, Position, Range};
+
+use crate::{TestServer, TestServerBuilder};
+
+#[test]
+fn publish_diagnostics_open() -> anyhow::Result<()> {
+    let mut server = TestServerBuilder::new()?
+        .build()?
+        .wait_until_workspaces_are_initialized()?;
+
+    server.initialization_result().unwrap();
+
+    let mut builder = NotebookBuilder::virtual_file("test.ipynb");
+    builder.add_python_cell(
+        r#"from typing import Literal
+
+type Style = Literal["italic", "bold", "underline"]"#,
+    );
+
+    builder.add_python_cell(
+        r#"def with_style(line: str, word, style: Style) -> str:
+    if style == "italic":
+        return line.replace(word, f"*{word}*")
+    elif style == "bold":
+        return line.replace(word, f"__{word}__")
+
+    position = line.find(word)
+    output = line + "\n"
+    output += " " * position
+    output += "-" * len(word)
+"#,
+    );
+
+    builder.add_python_cell(
+        r#"print(with_style("ty is a fast type checker for Python.", "fast", "underlined"))
+"#,
+    );
+
+    builder.open(&mut server);
+
+    let cell1_diagnostics =
+        server.await_notification::<lsp_types::notification::PublishDiagnostics>()?;
+    let cell2_diagnostics =
+        server.await_notification::<lsp_types::notification::PublishDiagnostics>()?;
+    let cell3_diagnostics =
+        server.await_notification::<lsp_types::notification::PublishDiagnostics>()?;
+
+    assert_json_snapshot!([cell1_diagnostics, cell2_diagnostics, cell3_diagnostics]);
+
+    Ok(())
+}
+
+#[test]
+fn diagnostic_end_of_file() -> anyhow::Result<()> {
+    let mut server = TestServerBuilder::new()?
+        .build()?
+        .wait_until_workspaces_are_initialized()?;
+
+    server.initialization_result().unwrap();
+
+    let mut builder = NotebookBuilder::virtual_file("test.ipynb");
+    builder.add_python_cell(
+        r#"from typing import Literal
+
+type Style = Literal["italic", "bold", "underline"]"#,
+    );
+
+    builder.add_python_cell(
+        r#"def with_style(line: str, word, style: Style) -> str:
+    if style == "italic":
+        return line.replace(word, f"*{word}*")
+    elif style == "bold":
+        return line.replace(word, f"__{word}__")
+
+    position = line.find(word)
+    output = line + "\n"
+    output += " " * position
+    output += "-" * len(word)
+"#,
+    );
+
+    let cell_3 = builder.add_python_cell(
+        r#"with_style("test", "word", "underline")
+
+IOError"#,
+    );
+
+    let notebook_url = builder.open(&mut server);
+
+    server.collect_publish_diagnostic_notifications(3)?;
+
+    server.send_notification::<lsp_types::notification::DidChangeNotebookDocument>(
+        lsp_types::DidChangeNotebookDocumentParams {
+            notebook_document: lsp_types::VersionedNotebookDocumentIdentifier {
+                version: 0,
+                uri: notebook_url,
+            },
+            change: lsp_types::NotebookDocumentChangeEvent {
+                metadata: None,
+                cells: Some(lsp_types::NotebookDocumentCellChange {
+                    structure: None,
+                    data: None,
+                    text_content: Some(vec![lsp_types::NotebookDocumentChangeTextContent {
+                        document: lsp_types::VersionedTextDocumentIdentifier {
+                            uri: cell_3,
+                            version: 0,
+                        },
+
+                        changes: {
+                            vec![lsp_types::TextDocumentContentChangeEvent {
+                                range: Some(Range::new(Position::new(0, 16), Position::new(0, 17))),
+                                range_length: Some(1),
+                                text: String::new(),
+                            }]
+                        },
+                    }]),
+                }),
+            },
+        },
+    );
+
+    let diagnostics = server.collect_publish_diagnostic_notifications(3)?;
+    assert_json_snapshot!(diagnostics);
+
+    Ok(())
+}
+
+#[test]
+fn semantic_tokens() -> anyhow::Result<()> {
+    let mut server = TestServerBuilder::new()?
+        .build()?
+        .wait_until_workspaces_are_initialized()?;
+
+    server.initialization_result().unwrap();
+
+    let mut builder = NotebookBuilder::virtual_file("src/test.ipynb");
+    let first_cell = builder.add_python_cell(
+        r#"from typing import Literal
+
+type Style = Literal["italic", "bold", "underline"]"#,
+    );
+
+    let second_cell = builder.add_python_cell(
+        r#"def with_style(line: str, word, style: Style) -> str:
+    if style == "italic":
+        return line.replace(word, f"*{word}*")
+    elif style == "bold":
+        return line.replace(word, f"__{word}__")
+
+    position = line.find(word)
+    output = line + "\n"
+    output += " " * position
+    output += "-" * len(word)
+"#,
+    );
+
+    let third_cell = builder.add_python_cell(
+        r#"print(with_style("ty is a fast type checker for Python.", "fast", "underlined"))
+"#,
+    );
+
+    builder.open(&mut server);
+
+    let cell1_tokens = semantic_tokens_full_for_cell(&mut server, &first_cell)?;
+    let cell2_tokens = semantic_tokens_full_for_cell(&mut server, &second_cell)?;
+    let cell3_tokens = semantic_tokens_full_for_cell(&mut server, &third_cell)?;
+
+    assert_json_snapshot!([cell1_tokens, cell2_tokens, cell3_tokens]);
+
+    server.collect_publish_diagnostic_notifications(3)?;
+
+    Ok(())
+}
+
+#[test]
+fn swap_cells() -> anyhow::Result<()> {
+    let mut server = TestServerBuilder::new()?
+        .build()?
+        .wait_until_workspaces_are_initialized()?;
+
+    server.initialization_result().unwrap();
+
+    let mut builder = NotebookBuilder::virtual_file("src/test.ipynb");
+
+    let first_cell = builder.add_python_cell(
+        r#"b = a
+"#,
+    );
+
+    let second_cell = builder.add_python_cell(r#"a = 10"#);
+
+    builder.add_python_cell(r#"c = b"#);
+
+    let notebook = builder.open(&mut server);
+
+    let diagnostics = server.collect_publish_diagnostic_notifications(3)?;
+    assert_json_snapshot!(diagnostics, @r###"
+    {
+      "vscode-notebook-cell://src/test.ipynb#0": [
+        {
+          "range": {
+            "start": {
+              "line": 0,
+              "character": 4
+            },
+            "end": {
+              "line": 0,
+              "character": 5
+            }
+          },
+          "severity": 1,
+          "code": "unresolved-reference",
+          "codeDescription": {
+            "href": "https://ty.dev/rules#unresolved-reference"
+          },
+          "source": "ty",
+          "message": "Name `a` used when not defined",
+          "relatedInformation": []
+        }
+      ],
+      "vscode-notebook-cell://src/test.ipynb#1": [],
+      "vscode-notebook-cell://src/test.ipynb#2": []
+    }
+    "###);
+
+    // Re-order the cells from `b`, `a`, `c` to `a`, `b`, `c` (swapping cell 1 and 2)
+    server.send_notification::<lsp_types::notification::DidChangeNotebookDocument>(
+        lsp_types::DidChangeNotebookDocumentParams {
+            notebook_document: lsp_types::VersionedNotebookDocumentIdentifier {
+                version: 1,
+                uri: notebook,
+            },
+            change: lsp_types::NotebookDocumentChangeEvent {
+                metadata: None,
+                cells: Some(lsp_types::NotebookDocumentCellChange {
+                    structure: Some(lsp_types::NotebookDocumentCellChangeStructure {
+                        array: lsp_types::NotebookCellArrayChange {
+                            start: 0,
+                            delete_count: 2,
+                            cells: Some(vec![
+                                lsp_types::NotebookCell {
+                                    kind: NotebookCellKind::Code,
+                                    document: second_cell,
+                                    metadata: None,
+                                    execution_summary: None,
+                                },
+                                lsp_types::NotebookCell {
+                                    kind: NotebookCellKind::Code,
+                                    document: first_cell,
+                                    metadata: None,
+                                    execution_summary: None,
+                                },
+                            ]),
+                        },
+                        did_open: None,
+                        did_close: None,
+                    }),
+                    data: None,
+                    text_content: None,
+                }),
+            },
+        },
+    );
+
+    let diagnostics = server.collect_publish_diagnostic_notifications(3)?;
+
+    assert_json_snapshot!(diagnostics, @r###"
+    {
+      "vscode-notebook-cell://src/test.ipynb#0": [],
+      "vscode-notebook-cell://src/test.ipynb#1": [],
+      "vscode-notebook-cell://src/test.ipynb#2": []
+    }
+    "###);
+
+    Ok(())
+}
+
+fn semantic_tokens_full_for_cell(
+    server: &mut TestServer,
+    cell_uri: &lsp_types::Url,
+) -> crate::Result<Option<lsp_types::SemanticTokensResult>> {
+    let cell1_tokens_req_id = server.send_request::<lsp_types::request::SemanticTokensFullRequest>(
+        lsp_types::SemanticTokensParams {
+            work_done_progress_params: lsp_types::WorkDoneProgressParams::default(),
+            partial_result_params: lsp_types::PartialResultParams::default(),
+            text_document: lsp_types::TextDocumentIdentifier {
+                uri: cell_uri.clone(),
+            },
+        },
+    );
+
+    server.await_response::<lsp_types::request::SemanticTokensFullRequest>(&cell1_tokens_req_id)
+}
+
+#[derive(Debug)]
+pub(crate) struct NotebookBuilder {
+    notebook_url: lsp_types::Url,
+    // The cells: (cell_metadata, content, language_id)
+    cells: Vec<(lsp_types::NotebookCell, String, String)>,
+}
+
+impl NotebookBuilder {
+    pub(crate) fn virtual_file(name: &str) -> Self {
+        let url: lsp_types::Url = format!("vs-code:/{name}").parse().unwrap();
+        Self {
+            notebook_url: url,
+            cells: Vec::new(),
+        }
+    }
+
+    pub(crate) fn add_python_cell(&mut self, content: &str) -> lsp_types::Url {
+        let index = self.cells.len();
+        let id = format!(
+            "vscode-notebook-cell:/{}#{}",
+            self.notebook_url.path(),
+            index
+        );
+
+        let url: lsp_types::Url = id.parse().unwrap();
+
+        self.cells.push((
+            lsp_types::NotebookCell {
+                kind: NotebookCellKind::Code,
+                document: url.clone(),
+                metadata: None,
+                execution_summary: None,
+            },
+            content.to_string(),
+            "python".to_string(),
+        ));
+
+        url
+    }
+
+    pub(crate) fn open(self, server: &mut TestServer) -> lsp_types::Url {
+        server.send_notification::<lsp_types::notification::DidOpenNotebookDocument>(
+            lsp_types::DidOpenNotebookDocumentParams {
+                notebook_document: lsp_types::NotebookDocument {
+                    uri: self.notebook_url.clone(),
+                    notebook_type: "jupyter-notebook".to_string(),
+                    version: 0,
+                    metadata: None,
+                    cells: self.cells.iter().map(|(cell, _, _)| cell.clone()).collect(),
+                },
+                cell_text_documents: self
+                    .cells
+                    .iter()
+                    .map(|(cell, content, language_id)| lsp_types::TextDocumentItem {
+                        uri: cell.document.clone(),
+                        language_id: language_id.clone(),
+                        version: 0,
+                        text: content.clone(),
+                    })
+                    .collect(),
+            },
+        );
+
+        self.notebook_url
+    }
+}

crates/ty_server/tests/e2e/snapshots/e2e__initialize__initialization.snap

@@ -1,6 +1,5 @@
 ---
 source: crates/ty_server/tests/e2e/initialize.rs
-assertion_line: 17
 expression: initialization_result
 ---
 {
@@ -10,6 +9,18 @@ expression: initialization_result
       "openClose": true,
       "change": 2
     },
+    "notebookDocumentSync": {
+      "notebookSelector": [
+        {
+          "cells": [
+            {
+              "language": "python"
+            }
+          ]
+        }
+      ],
+      "save": false
+    },
     "selectionRangeProvider": true,
     "hoverProvider": true,
     "completionProvider": {

crates/ty_server/tests/e2e/snapshots/e2e__initialize__initialization_with_workspace.snap

@@ -1,6 +1,5 @@
 ---
 source: crates/ty_server/tests/e2e/initialize.rs
-assertion_line: 32
 expression: initialization_result
 ---
 {
@@ -10,6 +9,18 @@ expression: initialization_result
       "openClose": true,
       "change": 2
     },
+    "notebookDocumentSync": {
+      "notebookSelector": [
+        {
+          "cells": [
+            {
+              "language": "python"
+            }
+          ]
+        }
+      ],
+      "save": false
+    },
     "selectionRangeProvider": true,
     "hoverProvider": true,
     "completionProvider": {

crates/ty_server/tests/e2e/snapshots/e2e__notebook__diagnostic_end_of_file.snap

@@ -0,0 +1,263 @@
+---
+source: crates/ty_server/tests/e2e/notebook.rs
+expression: diagnostics
+---
+{
+  "vscode-notebook-cell://test.ipynb#0": [],
+  "vscode-notebook-cell://test.ipynb#1": [
+    {
+      "range": {
+        "start": {
+          "line": 0,
+          "character": 49
+        },
+        "end": {
+          "line": 0,
+          "character": 52
+        }
+      },
+      "severity": 1,
+      "code": "invalid-return-type",
+      "codeDescription": {
+        "href": "https://ty.dev/rules#invalid-return-type"
+      },
+      "source": "ty",
+      "message": "Function can implicitly return `None`, which is not assignable to return type `str`",
+      "relatedInformation": []
+    }
+  ],
+  "vscode-notebook-cell://test.ipynb#2": [
+    {
+      "range": {
+        "start": {
+          "line": 0,
+          "character": 19
+        },
+        "end": {
+          "line": 0,
+          "character": 23
+        }
+      },
+      "severity": 1,
+      "code": "invalid-syntax",
+      "source": "ty",
+      "message": "Expected `,`, found name",
+      "relatedInformation": []
+    },
+    {
+      "range": {
+        "start": {
+          "line": 0,
+          "character": 19
+        },
+        "end": {
+          "line": 0,
+          "character": 23
+        }
+      },
+      "severity": 1,
+      "code": "unresolved-reference",
+      "codeDescription": {
+        "href": "https://ty.dev/rules#unresolved-reference"
+      },
+      "source": "ty",
+      "message": "Name `word` used when not defined",
+      "relatedInformation": []
+    },
+    {
+      "range": {
+        "start": {
+          "line": 0,
+          "character": 23
+        },
+        "end": {
+          "line": 0,
+          "character": 27
+        }
+      },
+      "severity": 1,
+      "code": "invalid-syntax",
+      "source": "ty",
+      "message": "Expected `,`, found string",
+      "relatedInformation": []
+    },
+    {
+      "range": {
+        "start": {
+          "line": 0,
+          "character": 23
+        },
+        "end": {
+          "line": 0,
+          "character": 27
+        }
+      },
+      "severity": 1,
+      "code": "invalid-argument-type",
+      "codeDescription": {
+        "href": "https://ty.dev/rules#invalid-argument-type"
+      },
+      "source": "ty",
+      "message": "Argument to function `with_style` is incorrect: Expected `Style`, found `Literal[/", /"]`",
+      "relatedInformation": [
+        {
+          "location": {
+            "uri": "vscode-notebook-cell://test.ipynb#1",
+            "range": {
+              "start": {
+                "line": 0,
+                "character": 4
+              },
+              "end": {
+                "line": 0,
+                "character": 14
+              }
+            }
+          },
+          "message": "Function defined here"
+        },
+        {
+          "location": {
+            "uri": "vscode-notebook-cell://test.ipynb#1",
+            "range": {
+              "start": {
+                "line": 0,
+                "character": 32
+              },
+              "end": {
+                "line": 0,
+                "character": 44
+              }
+            }
+          },
+          "message": "Parameter declared here"
+        }
+      ]
+    },
+    {
+      "range": {
+        "start": {
+          "line": 0,
+          "character": 27
+        },
+        "end": {
+          "line": 0,
+          "character": 36
+        }
+      },
+      "severity": 1,
+      "code": "invalid-syntax",
+      "source": "ty",
+      "message": "Expected `,`, found name",
+      "relatedInformation": []
+    },
+    {
+      "range": {
+        "start": {
+          "line": 0,
+          "character": 27
+        },
+        "end": {
+          "line": 0,
+          "character": 36
+        }
+      },
+      "severity": 1,
+      "code": "unresolved-reference",
+      "codeDescription": {
+        "href": "https://ty.dev/rules#unresolved-reference"
+      },
+      "source": "ty",
+      "message": "Name `underline` used when not defined",
+      "relatedInformation": []
+    },
+    {
+      "range": {
+        "start": {
+          "line": 0,
+          "character": 27
+        },
+        "end": {
+          "line": 0,
+          "character": 36
+        }
+      },
+      "severity": 1,
+      "code": "too-many-positional-arguments",
+      "codeDescription": {
+        "href": "https://ty.dev/rules#too-many-positional-arguments"
+      },
+      "source": "ty",
+      "message": "Too many positional arguments to function `with_style`: expected 3, got 6",
+      "relatedInformation": [
+        {
+          "location": {
+            "uri": "vscode-notebook-cell://test.ipynb#1",
+            "range": {
+              "start": {
+                "line": 0,
+                "character": 4
+              },
+              "end": {
+                "line": 0,
+                "character": 52
+              }
+            }
+          },
+          "message": "Function signature here"
+        }
+      ]
+    },
+    {
+      "range": {
+        "start": {
+          "line": 0,
+          "character": 36
+        },
+        "end": {
+          "line": 0,
+          "character": 38
+        }
+      },
+      "severity": 1,
+      "code": "invalid-syntax",
+      "source": "ty",
+      "message": "missing closing quote in string literal",
+      "relatedInformation": []
+    },
+    {
+      "range": {
+        "start": {
+          "line": 2,
+          "character": 0
+        },
+        "end": {
+          "line": 2,
+          "character": 7
+        }
+      },
+      "severity": 1,
+      "code": "invalid-syntax",
+      "source": "ty",
+      "message": "Expected `,`, found name",
+      "relatedInformation": []
+    },
+    {
+      "range": {
+        "start": {
+          "line": 3,
+          "character": 0
+        },
+        "end": {
+          "line": 3,
+          "character": 0
+        }
+      },
+      "severity": 1,
+      "code": "invalid-syntax",
+      "source": "ty",
+      "message": "unexpected EOF while parsing",
+      "relatedInformation": []
+    }
+  ]
+}

crates/ty_server/tests/e2e/snapshots/e2e__notebook__publish_diagnostics_open.snap

@@ -0,0 +1,96 @@
+---
+source: crates/ty_server/tests/e2e/notebook.rs
+expression: "[cell1_diagnostics, cell2_diagnostics, cell3_diagnostics]"
+---
+[
+  {
+    "uri": "vscode-notebook-cell://test.ipynb#1",
+    "diagnostics": [
+      {
+        "range": {
+          "start": {
+            "line": 0,
+            "character": 49
+          },
+          "end": {
+            "line": 0,
+            "character": 52
+          }
+        },
+        "severity": 1,
+        "code": "invalid-return-type",
+        "codeDescription": {
+          "href": "https://ty.dev/rules#invalid-return-type"
+        },
+        "source": "ty",
+        "message": "Function can implicitly return `None`, which is not assignable to return type `str`",
+        "relatedInformation": []
+      }
+    ],
+    "version": 0
+  },
+  {
+    "uri": "vscode-notebook-cell://test.ipynb#2",
+    "diagnostics": [
+      {
+        "range": {
+          "start": {
+            "line": 0,
+            "character": 66
+          },
+          "end": {
+            "line": 0,
+            "character": 78
+          }
+        },
+        "severity": 1,
+        "code": "invalid-argument-type",
+        "codeDescription": {
+          "href": "https://ty.dev/rules#invalid-argument-type"
+        },
+        "source": "ty",
+        "message": "Argument to function `with_style` is incorrect: Expected `Style`, found `Literal[/"underlined/"]`",
+        "relatedInformation": [
+          {
+            "location": {
+              "uri": "vscode-notebook-cell://test.ipynb#1",
+              "range": {
+                "start": {
+                  "line": 0,
+                  "character": 4
+                },
+                "end": {
+                  "line": 0,
+                  "character": 14
+                }
+              }
+            },
+            "message": "Function defined here"
+          },
+          {
+            "location": {
+              "uri": "vscode-notebook-cell://test.ipynb#1",
+              "range": {
+                "start": {
+                  "line": 0,
+                  "character": 32
+                },
+                "end": {
+                  "line": 0,
+                  "character": 44
+                }
+              }
+            },
+            "message": "Parameter declared here"
+          }
+        ]
+      }
+    ],
+    "version": 0
+  },
+  {
+    "uri": "vscode-notebook-cell://test.ipynb#0",
+    "diagnostics": [],
+    "version": 0
+  }
+]

crates/ty_server/tests/e2e/snapshots/e2e__notebook__semantic_tokens.snap

@@ -0,0 +1,263 @@
+---
+source: crates/ty_server/tests/e2e/notebook.rs
+expression: "[cell1_tokens, cell2_tokens, cell3_tokens]"
+---
+[
+  {
+    "data": [
+      0,
+      5,
+      6,
+      0,
+      0,
+      0,
+      14,
+      7,
+      5,
+      0,
+      2,
+      5,
+      5,
+      14,
+      0,
+      0,
+      8,
+      7,
+      5,
+      0,
+      0,
+      8,
+      8,
+      10,
+      0,
+      0,
+      10,
+      6,
+      10,
+      0,
+      0,
+      8,
+      11,
+      10,
+      0
+    ]
+  },
+  {
+    "data": [
+      0,
+      4,
+      10,
+      7,
+      1,
+      0,
+      11,
+      4,
+      2,
+      1,
+      0,
+      6,
+      3,
+      1,
+      0,
+      0,
+      5,
+      4,
+      2,
+      1,
+      0,
+      6,
+      5,
+      2,
+      1,
+      0,
+      7,
+      5,
+      14,
+      0,
+      0,
+      10,
+      3,
+      1,
+      0,
+      1,
+      7,
+      5,
+      2,
+      0,
+      0,
+      9,
+      8,
+      10,
+      0,
+      1,
+      15,
+      4,
+      2,
+      0,
+      0,
+      5,
+      7,
+      8,
+      0,
+      0,
+      8,
+      4,
+      2,
+      0,
+      0,
+      8,
+      1,
+      10,
+      0,
+      0,
+      2,
+      4,
+      2,
+      0,
+      0,
+      5,
+      1,
+      10,
+      0,
+      1,
+      9,
+      5,
+      2,
+      0,
+      0,
+      9,
+      6,
+      10,
+      0,
+      1,
+      15,
+      4,
+      2,
+      0,
+      0,
+      5,
+      7,
+      8,
+      0,
+      0,
+      8,
+      4,
+      2,
+      0,
+      0,
+      8,
+      2,
+      10,
+      0,
+      0,
+      3,
+      4,
+      2,
+      0,
+      0,
+      5,
+      2,
+      10,
+      0,
+      2,
+      4,
+      8,
+      5,
+      1,
+      0,
+      11,
+      4,
+      2,
+      0,
+      0,
+      5,
+      4,
+      8,
+      0,
+      0,
+      5,
+      4,
+      2,
+      0,
+      1,
+      4,
+      6,
+      5,
+      1,
+      0,
+      9,
+      4,
+      2,
+      0,
+      0,
+      7,
+      4,
+      10,
+      0,
+      1,
+      4,
+      6,
+      5,
+      0,
+      0,
+      10,
+      3,
+      10,
+      0,
+      0,
+      6,
+      8,
+      5,
+      0,
+      1,
+      4,
+      6,
+      5,
+      0,
+      0,
+      10,
+      3,
+      10,
+      0,
+      0,
+      6,
+      3,
+      7,
+      0,
+      0,
+      4,
+      4,
+      2,
+      0
+    ]
+  },
+  {
+    "data": [
+      0,
+      0,
+      5,
+      7,
+      0,
+      0,
+      6,
+      10,
+      7,
+      0,
+      0,
+      11,
+      39,
+      10,
+      0,
+      0,
+      41,
+      6,
+      10,
+      0,
+      0,
+      8,
+      12,
+      10,
+      0
+    ]
+  }
+]