Move db updating into Session

Author: MichaReiser

crates/ty_server/src/document/notebook.rs

@@ -141,14 +141,14 @@ impl NotebookDocument {
             new_cells.into_iter().map(NotebookCell::new),
         );
 
-        // Re-build the cell-index if new cells were added, deleted or removed
+        // Re-build the cell-index if new cells were added or deleted
         if !deleted_range.is_empty() || added > 0 {
             self.cell_index.clear();
             self.cell_index.extend(
                 self.cells
                     .iter()
                     .enumerate()
-                    .map(|(i, cell)| (cell.url.clone(), i)),
+                    .map(|(cell_index, cell)| (cell.url.clone(), cell_index)),
             );
         }
 

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
@@ -148,21 +148,22 @@ impl PositionExt for lsp_types::Position {
                 index.source_location(cell_start_offset, source.as_str(), encoding.into());
             assert_eq!(cell_start_location.character_offset, OneIndexed::MIN);
 
-            let absolute_start = SourceLocation {
+            // 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_start, &source, encoding.into()));
+            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:
     ///
@@ -297,12 +298,8 @@ impl ToRangeExt for TextRange {
             let start_in_concatenated =
                 index.source_location(self.start(), &source, encoding.into());
             let cell_index = notebook_index.cell(start_in_concatenated.line)?;
-            let cell_range = notebook.cell_range(cell_index)?;
 
-            // Clamp the end offset to the end of the cell.
-            let end_offset = self.end().min(cell_range.end());
-
-            let end_in_concatenated = index.source_location(end_offset, &source, encoding.into());
+            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),

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

@@ -113,15 +113,16 @@ impl LspDiagnostics {
 }
 
 pub(super) fn clear_diagnostics_if_needed(
+    document: &DocumentHandle,
     session: &Session,
-    uri: &lsp_types::Url,
     client: &Client,
 ) {
-    if session.client_capabilities().supports_pull_diagnostics() {
+    if session.client_capabilities().supports_pull_diagnostics() && !document.is_cell_or_notebook()
+    {
         return;
     }
 
-    clear_diagnostics(uri, client);
+    clear_diagnostics(document.url(), client);
 }
 
 /// Clears the diagnostics for the document identified by `uri`.
@@ -138,21 +139,29 @@ pub(super) fn clear_diagnostics(uri: &lsp_types::Url, client: &Client) {
 }
 
 /// Publishes the diagnostics for the given document snapshot using the [publish diagnostics
-/// notification].
+/// notification] .
 ///
-/// This function publishes the diagnostics for the entire notebook if `url` points to a notebook or a cell,
-/// even if the client supports pull diagnostics. This is because VS Code
+/// 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).
 ///
-/// This function is a no-op for simple text files when the client supports pull diagnostics.
-///
 /// [publish diagnostics notification]: https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_publishDiagnostics
-pub(super) fn publish_diagnostics(session: &Session, document: &DocumentHandle, client: &Client) {
+pub(super) fn publish_diagnostics_if_needed(
+    document: &DocumentHandle,
+    session: &Session,
+    client: &Client,
+) {
     if !document.is_cell_or_notebook() && session.client_capabilities().supports_pull_diagnostics()
     {
         return;
     }
 
+    publish_diagnostics(document, session, client);
+}
+
+/// 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, document, session.position_encoding()) else {

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::session::{DocumentHandle, Session};
-use crate::system::AnySystemPath;
-use ty_project::watch::ChangeEvent;
 
 pub(crate) struct DidChangeTextDocumentHandler;
 
@@ -36,24 +34,8 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
             .update_text_document(session, content_changes, version)
             .with_failure_code(ErrorCode::InternalError)?;
 
-        file_changed(&document, session, client);
+        publish_diagnostics_if_needed(&document, session, client);
 
         Ok(())
     }
 }
-
-pub(super) fn file_changed(document: &DocumentHandle, session: &mut Session, client: &Client) {
-    let path = document.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);
-
-    publish_diagnostics(session, document, client);
-}

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

@@ -4,7 +4,7 @@ use lsp_types::notification as notif;
 
 use crate::server::Result;
 use crate::server::api::LSPResult;
-use crate::server::api::notifications::did_change::file_changed;
+use crate::server::api::diagnostics::publish_diagnostics;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::session::Session;
 use crate::session::client::Client;
@@ -32,7 +32,7 @@ impl SyncNotificationHandler for DidChangeNotebookHandler {
             .update_notebook_document(session, cells, metadata, version)
             .with_failure_code(ErrorCode::InternalError)?;
 
-        file_changed(&document, session, client);
+        publish_diagnostics(&document, session, client);
 
         Ok(())
     }

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

@@ -92,7 +92,7 @@ impl SyncNotificationHandler for DidChangeWatchedFiles {
             );
         } else {
             for key in session.text_document_handles() {
-                publish_diagnostics(session, &key, client);
+                publish_diagnostics(&key, session, client);
             }
         }
 

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

@@ -1,16 +1,13 @@
 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, clear_diagnostics_if_needed};
+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;
-use crate::session::{DocumentHandle, Session};
-use crate::system::AnySystemPath;
 
 pub(crate) struct DidCloseTextDocumentHandler;
 
@@ -31,68 +28,15 @@ impl SyncNotificationHandler for DidCloseTextDocumentHandler {
         let document = session
             .document_handle(&uri)
             .with_failure_code(ErrorCode::InternalError)?;
-        let is_cell = document.is_cell();
 
-        if !is_cell {
-            close_document(&document, session, client);
-        }
-
-        let url = document.url().clone();
-
-        document
+        let should_clear_diagnostics = document
             .close(session)
             .with_failure_code(ErrorCode::InternalError)?;
 
-        if is_cell {
-            // Notebooks only support publish diagnostics.
-            clear_diagnostics(&url, client);
+        if should_clear_diagnostics {
+            clear_diagnostics_if_needed(&document, session, client);
         }
 
         Ok(())
     }
 }
-
-pub(super) fn close_document(document: &DocumentHandle, session: &mut Session, client: &Client) {
-    let path = document.notebook_or_file_path();
-    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
-            if session.workspaces().for_path(system_path).is_none()
-                || session
-                    .global_settings()
-                    .diagnostic_mode()
-                    .is_open_files_only()
-            {
-                clear_diagnostics_if_needed(session, document.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_if_needed(session, document.url(), client);
-        }
-    }
-}

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

@@ -3,7 +3,6 @@ use lsp_types::{DidCloseNotebookDocumentParams, NotebookDocumentIdentifier};
 
 use crate::server::Result;
 use crate::server::api::LSPResult;
-use crate::server::api::notifications::did_close::close_document;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::session::Session;
 use crate::session::client::Client;
@@ -17,7 +16,7 @@ impl NotificationHandler for DidCloseNotebookHandler {
 impl SyncNotificationHandler for DidCloseNotebookHandler {
     fn run(
         session: &mut Session,
-        client: &Client,
+        _client: &Client,
         params: DidCloseNotebookDocumentParams,
     ) -> Result<()> {
         let DidCloseNotebookDocumentParams {
@@ -29,9 +28,9 @@ impl SyncNotificationHandler for DidCloseNotebookHandler {
             .document_handle(&uri)
             .with_failure_code(lsp_server::ErrorCode::InternalError)?;
 
-        close_document(&document, session, client);
-
-        document
+        // We don't need to call publish any diagnostics because we clear
+        // the diagnostics when closing the corresponding cell documents.
+        let _ = document
             .close(session)
             .with_failure_code(lsp_server::ErrorCode::InternalError)?;
 

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::session::{DocumentHandle, Session};
-use crate::system::AnySystemPath;
 
 pub(crate) struct DidOpenTextDocumentHandler;
 
@@ -39,49 +34,8 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
             TextDocument::new(uri, text, version).with_language_id(&language_id),
         );
 
-        opened_document(&document, session, client);
+        publish_diagnostics_if_needed(&document, session, client);
 
         Ok(())
     }
 }
-
-pub(super) fn opened_document(document: &DocumentHandle, session: &mut Session, client: &Client) {
-    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 = 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, client);
-}