From 0c9165fc3aa628b7dd975dade291bb978318e92b Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Mon, 25 Nov 2024 15:14:30 +0530
Subject: [PATCH] Use `Result` for failed text document retrieval in LSP
 requests (#14579)

## Summary

Ref:
https://github.com/astral-sh/ruff-vscode/issues/644#issuecomment-2496588452

## Test Plan

Not sure how to test this as this is mainly to get more context on the
panic that the server is raising.
---
 .../src/server/api/requests/format.rs         |  4 ++-
 .../src/server/api/requests/format_range.rs   |  4 ++-
 .../src/server/api/requests/hover.rs          |  4 ++-
 crates/ruff_server/src/session/index.rs       | 27 +++++++++++++++----
 4 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/crates/ruff_server/src/server/api/requests/format.rs b/crates/ruff_server/src/server/api/requests/format.rs
index 853be16618a6a..77bc1fc8c3dfc 100644
--- a/crates/ruff_server/src/server/api/requests/format.rs
+++ b/crates/ruff_server/src/server/api/requests/format.rs
@@ -1,3 +1,4 @@
+use anyhow::Context;
 use lsp_types::{self as types, request as req};
 use types::TextEdit;
 
@@ -64,7 +65,8 @@ pub(super) fn format_document(snapshot: &DocumentSnapshot) -> Result<super::Form
     let text_document = snapshot
         .query()
         .as_single_document()
-        .expect("format should only be called on text documents or notebook cells");
+        .context("Failed to get text document for the format request")
+        .unwrap();
     let query = snapshot.query();
     format_text_document(
         text_document,
diff --git a/crates/ruff_server/src/server/api/requests/format_range.rs b/crates/ruff_server/src/server/api/requests/format_range.rs
index 2eb8ea5ab749c..7cbc42116e678 100644
--- a/crates/ruff_server/src/server/api/requests/format_range.rs
+++ b/crates/ruff_server/src/server/api/requests/format_range.rs
@@ -1,3 +1,4 @@
+use anyhow::Context;
 use lsp_types::{self as types, request as req, Range};
 
 use crate::edit::{RangeExt, ToRangeExt};
@@ -32,7 +33,8 @@ fn format_document_range(
     let text_document = snapshot
         .query()
         .as_single_document()
-        .expect("format should only be called on text documents or notebook cells");
+        .context("Failed to get text document for the format range request")
+        .unwrap();
     let query = snapshot.query();
     format_text_document_range(text_document, range, query, snapshot.encoding())
 }
diff --git a/crates/ruff_server/src/server/api/requests/hover.rs b/crates/ruff_server/src/server/api/requests/hover.rs
index d982b497a312e..04aeaa53f0a97 100644
--- a/crates/ruff_server/src/server/api/requests/hover.rs
+++ b/crates/ruff_server/src/server/api/requests/hover.rs
@@ -1,5 +1,6 @@
 use crate::server::{client::Notifier, Result};
 use crate::session::DocumentSnapshot;
+use anyhow::Context;
 use lsp_types::{self as types, request as req};
 use regex::Regex;
 use ruff_diagnostics::FixAvailability;
@@ -33,7 +34,8 @@ pub(crate) fn hover(
     let document = snapshot
         .query()
         .as_single_document()
-        .expect("hover should only be called on text documents or notebook cells");
+        .context("Failed to get text document for the hover request")
+        .unwrap();
     let line_number: usize = position
         .position
         .line
diff --git a/crates/ruff_server/src/session/index.rs b/crates/ruff_server/src/session/index.rs
index a4adbd4a17ccb..25129e73fe6bb 100644
--- a/crates/ruff_server/src/session/index.rs
+++ b/crates/ruff_server/src/session/index.rs
@@ -6,6 +6,7 @@ use std::{collections::BTreeMap, path::Path, sync::Arc};
 use anyhow::anyhow;
 use lsp_types::Url;
 use rustc_hash::FxHashMap;
+use thiserror::Error;
 
 pub(crate) use ruff_settings::RuffSettings;
 
@@ -597,16 +598,24 @@ impl DocumentQuery {
 
     /// Attempt to access the single inner text document selected by the query.
     /// If this query is selecting an entire notebook document, this will return `None`.
-    pub(crate) fn as_single_document(&self) -> Option<&TextDocument> {
+    pub(crate) fn as_single_document(&self) -> Result<&TextDocument, SingleDocumentError> {
         match self {
-            Self::Text { document, .. } => Some(document),
+            Self::Text { document, .. } => Ok(document),
             Self::Notebook {
                 notebook,
+                file_url,
                 cell_url: cell_uri,
                 ..
-            } => cell_uri
-                .as_ref()
-                .and_then(|cell_uri| notebook.cell_document_by_uri(cell_uri)),
+            } => {
+                if let Some(cell_uri) = cell_uri {
+                    let cell = notebook
+                        .cell_document_by_uri(cell_uri)
+                        .ok_or_else(|| SingleDocumentError::CellDoesNotExist(cell_uri.clone()))?;
+                    Ok(cell)
+                } else {
+                    Err(SingleDocumentError::Notebook(file_url.clone()))
+                }
+            }
         }
     }
 
@@ -618,3 +627,11 @@ impl DocumentQuery {
         }
     }
 }
+
+#[derive(Debug, Error)]
+pub(crate) enum SingleDocumentError {
+    #[error("Expected a single text document, but found a notebook document: {0}")]
+    Notebook(Url),
+    #[error("Cell with URL {0} does not exist in the internal notebook document")]
+    CellDoesNotExist(Url),
+}