From 791f6a1820cb9312622bca110ec47c351ced1e21 Mon Sep 17 00:00:00 2001
From: Jane Lewis <jane@astral.sh>
Date: Fri, 21 Jun 2024 10:53:30 -0700
Subject: [PATCH] `ruff server`: Closing an untitled, unsaved notebook document
 no longer throws an error (#11942)

## Summary

Fixes #11651.
Fixes #11851.

We were double-closing a notebook document from the index, once in
`textDocument/didClose` and then in the `notebookDocument/didClose`
handler. The second time this happens, taking a snapshot fails.

I've rewritten how we handle snapshots for closing notebooks / notebook
cells so that any failure is simply logged instead of propagating
upwards. This implementation works consistently even if we don't receive
`textDocument/didClose` notifications for each specific cell, since they
get closed (and the diagnostics get cleared) in the notebook document
removal process.

## Test Plan

1. Open an untitled, unsaved notebook with the `Create: New Jupyter
Notebook` command from the VS Code command palette (`Ctrl/Cmd + Shift +
P`)
2. Without saving the document, close it.
3. No error popup should appear.
4. Run the debug command (`Ruff: print debug information`) to confirm
that there are no open documents
---
 .../src/server/api/notifications/did_close.rs       | 13 +++++++------
 .../server/api/notifications/did_close_notebook.rs  |  9 +++------
 crates/ruff_server/src/session/index.rs             |  9 +--------
 3 files changed, 11 insertions(+), 20 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 e8837327fd..491fa06c3d 100644
--- a/crates/ruff_server/src/server/api/notifications/did_close.rs
+++ b/crates/ruff_server/src/server/api/notifications/did_close.rs
@@ -21,15 +21,16 @@ impl super::SyncNotificationHandler for DidClose {
             text_document: types::TextDocumentIdentifier { uri },
         }: types::DidCloseTextDocumentParams,
     ) -> Result<()> {
+        let key = session.key_from_url(uri);
         // Publish an empty diagnostic report for the document. This will de-register any existing diagnostics.
-        let snapshot = session
-            .take_snapshot(uri.clone())
-            .ok_or_else(|| anyhow::anyhow!("Unable to take snapshot for document with URL {uri}"))
-            .with_failure_code(lsp_server::ErrorCode::InternalError)?;
+        let Some(snapshot) = session.take_snapshot(key.clone().into_url()) else {
+            tracing::debug!(
+                "Unable to close document with key {key} - the snapshot was unavailable"
+            );
+            return Ok(());
+        };
         clear_diagnostics_for_document(snapshot.query(), &notifier)?;
 
-        let key = snapshot.query().make_key();
-
         session
             .close_document(&key)
             .with_failure_code(lsp_server::ErrorCode::InternalError)
diff --git a/crates/ruff_server/src/server/api/notifications/did_close_notebook.rs b/crates/ruff_server/src/server/api/notifications/did_close_notebook.rs
index 561f2d8e68..913ccf5d4a 100644
--- a/crates/ruff_server/src/server/api/notifications/did_close_notebook.rs
+++ b/crates/ruff_server/src/server/api/notifications/did_close_notebook.rs
@@ -2,9 +2,8 @@ use crate::server::api::LSPResult;
 use crate::server::client::{Notifier, Requester};
 use crate::server::Result;
 use crate::session::Session;
-use lsp_server::ErrorCode;
-use lsp_types as types;
 use lsp_types::notification as notif;
+use lsp_types::{self as types, NotebookDocumentIdentifier};
 
 pub(crate) struct DidCloseNotebook;
 
@@ -18,16 +17,14 @@ impl super::SyncNotificationHandler for DidCloseNotebook {
         _notifier: Notifier,
         _requester: &mut Requester,
         types::DidCloseNotebookDocumentParams {
-            notebook_document: types::NotebookDocumentIdentifier { uri },
+            notebook_document: NotebookDocumentIdentifier { uri },
             ..
         }: types::DidCloseNotebookDocumentParams,
     ) -> Result<()> {
         let key = session.key_from_url(uri);
-
         session
             .close_document(&key)
-            .with_failure_code(ErrorCode::InternalError)?;
-
+            .with_failure_code(lsp_server::ErrorCode::InternalError)?;
         Ok(())
     }
 }
diff --git a/crates/ruff_server/src/session/index.rs b/crates/ruff_server/src/session/index.rs
index 4b5fdadbea..64e6333a07 100644
--- a/crates/ruff_server/src/session/index.rs
+++ b/crates/ruff_server/src/session/index.rs
@@ -352,16 +352,9 @@ impl Index {
             anyhow::bail!("Tried to close unavailable document `{key}`");
         };
 
-        let Some(controller) = self.documents.remove(&url) else {
+        let Some(_) = self.documents.remove(&url) else {
             anyhow::bail!("tried to close document that didn't exist at {}", url)
         };
-        if let Some(notebook) = controller.as_notebook() {
-            for url in notebook.urls() {
-                self.notebook_cells.remove(url).ok_or_else(|| {
-                    anyhow!("tried to de-register notebook cell with URL {url} that didn't exist")
-                })?;
-            }
-        }
         Ok(())
     }