From ffc98522cdbb62f6c759e845639390545212eaec Mon Sep 17 00:00:00 2001
From: Jane Lewis <jane@astral.sh>
Date: Mon, 17 Jun 2024 20:37:40 -0700
Subject: [PATCH] `ruff server`: Defer notebook cell deletion to avoid an error
 message (#11864)

## Summary

Fixes https://github.com/astral-sh/ruff-vscode/issues/496.

Cells are no longer removed from the notebook index when a notebook gets
updated, but rather when `textDocument/didClose` is called for them.
This solves an issue where their premature removal from the notebook
cell index would cause their URL to be un-queryable in the
`textDocument/didClose` handler.

## Test Plan

Create and then delete a notebook cell in VS Code. No error should
appear.
---
 .../src/server/api/notifications/did_close.rs |  8 ------
 crates/ruff_server/src/session/index.rs       | 27 ++++++++++---------
 2 files changed, 14 insertions(+), 21 deletions(-)

diff --git a/crates/ruff_server/src/server/api/notifications/did_close.rs b/crates/ruff_server/src/server/api/notifications/did_close.rs
index 4fc4bb1c5fc5b..e8837327fd47a 100644
--- a/crates/ruff_server/src/server/api/notifications/did_close.rs
+++ b/crates/ruff_server/src/server/api/notifications/did_close.rs
@@ -1,4 +1,3 @@
-use crate::edit::DocumentKey;
 use crate::server::api::diagnostics::clear_diagnostics_for_document;
 use crate::server::api::LSPResult;
 use crate::server::client::{Notifier, Requester};
@@ -31,13 +30,6 @@ impl super::SyncNotificationHandler for DidClose {
 
         let key = snapshot.query().make_key();
 
-        // Notebook cells will go through the `textDocument/didClose` path.
-        // We still want to publish empty diagnostics for them, but we
-        // shouldn't call `session.close_document` on them.
-        if matches!(key, DocumentKey::NotebookCell(_)) {
-            return Ok(());
-        }
-
         session
             .close_document(&key)
             .with_failure_code(lsp_server::ErrorCode::InternalError)
diff --git a/crates/ruff_server/src/session/index.rs b/crates/ruff_server/src/session/index.rs
index 6d60f148953a6..341e92cc73815 100644
--- a/crates/ruff_server/src/session/index.rs
+++ b/crates/ruff_server/src/session/index.rs
@@ -145,11 +145,8 @@ impl Index {
         encoding: PositionEncoding,
     ) -> crate::Result<()> {
         // update notebook cell index
-        if let Some(lsp_types::NotebookDocumentCellChangeStructure {
-            did_open,
-            did_close,
-            ..
-        }) = cells.as_ref().and_then(|cells| cells.structure.as_ref())
+        if let Some(lsp_types::NotebookDocumentCellChangeStructure { did_open, .. }) =
+            cells.as_ref().and_then(|cells| cells.structure.as_ref())
         {
             let Some(path) = self.url_for_key(key).cloned() else {
                 anyhow::bail!("Tried to open unavailable document `{key}`");
@@ -159,14 +156,7 @@ impl Index {
                 self.notebook_cells
                     .insert(opened_cell.uri.clone(), path.clone());
             }
-            for closed_cell in did_close.iter().flatten() {
-                if self.notebook_cells.remove(&closed_cell.uri).is_none() {
-                    tracing::warn!(
-                        "Tried to remove a notebook cell that does not exist: {}",
-                        closed_cell.uri
-                    );
-                }
-            }
+            // deleted notebook cells are closed via textDocument/didClose - we don't close them here.
         }
 
         let controller = self.document_controller_for_key(key)?;
@@ -347,6 +337,17 @@ impl Index {
     }
 
     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::NotebookCell(uri) = key {
+            if self.notebook_cells.remove(uri).is_none() {
+                tracing::warn!("Tried to remove a notebook cell that does not exist: {uri}",);
+            }
+            return Ok(());
+        }
         let Some(url) = self.url_for_key(key).cloned() else {
             anyhow::bail!("Tried to close unavailable document `{key}`");
         };