[ty] Make sure to always respond to client requests (#19277)

## Summary

This PR fixes a bug that didn't return a response to the client if the
document snapshotting failed.

This is resolved by making sure that the server always creates the
document snapshot and embed the any failures inside the snapshot.

Closes: astral-sh/ty#798

## Test Plan

Using the test case as described in the linked issue:



https://github.com/user-attachments/assets/f32833f8-03e5-4641-8c7f-2a536fe2e270

Cargo.lock

@@ -4272,6 +4272,7 @@ dependencies = [
  "serde",
  "serde_json",
  "shellexpand",
+ "thiserror 2.0.12",
  "tracing",
  "tracing-subscriber",
  "ty_ide",

crates/ty_server/Cargo.toml

@@ -31,6 +31,7 @@ salsa = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 shellexpand = { workspace = true }
+thiserror = { workspace = true }
 tracing = { workspace = true }
 tracing-subscriber = { workspace = true, features = ["chrono"] }
 

crates/ty_server/src/lib.rs

@@ -1,7 +1,7 @@
 use crate::server::{ConnectionInitializer, Server};
 use anyhow::Context;
 pub use document::{NotebookDocument, PositionEncoding, TextDocument};
-pub use session::{DocumentQuery, DocumentSnapshot, Session};
+pub(crate) use session::{DocumentQuery, Session};
 use std::num::NonZeroUsize;
 
 mod document;

crates/ty_server/src/server/api.rs

@@ -218,8 +218,22 @@ where
         let url = R::document_url(&params).into_owned();
 
         let Ok(path) = AnySystemPath::try_from_url(&url) else {
-            tracing::warn!("Ignoring request for invalid `{url}`");
-            return Box::new(|_| {});
+            let reason = format!("URL `{url}` isn't a valid system path");
+            tracing::warn!(
+                "Ignoring request id={id} method={} because {reason}",
+                R::METHOD
+            );
+            return Box::new(|client| {
+                respond_silent_error(
+                    id,
+                    client,
+                    lsp_server::ResponseError {
+                        code: lsp_server::ErrorCode::InvalidParams as i32,
+                        message: reason,
+                        data: None,
+                    },
+                );
+            });
         };
 
         let db = match &path {
@@ -230,10 +244,7 @@ where
             AnySystemPath::SystemVirtual(_) => session.default_project_db().clone(),
         };
 
-        let Some(snapshot) = session.take_document_snapshot(url) else {
-            tracing::warn!("Ignoring request because snapshot for path `{path:?}` doesn't exist");
-            return Box::new(|_| {});
-        };
+        let snapshot = session.take_document_snapshot(url);
 
         Box::new(move |client| {
             let _span = tracing::debug_span!("request", %id, method = R::METHOD).entered();
@@ -331,12 +342,7 @@ where
     let (id, params) = cast_notification::<N>(req)?;
     Ok(Task::background(schedule, move |session: &Session| {
         let url = N::document_url(&params);
-        let Some(snapshot) = session.take_document_snapshot((*url).clone()) else {
-            tracing::debug!(
-                "Ignoring notification because snapshot for url `{url}` doesn't exist."
-            );
-            return Box::new(|_| {});
-        };
+        let snapshot = session.take_document_snapshot((*url).clone());
         Box::new(move |client| {
             let _span = tracing::debug_span!("notification", method = N::METHOD).entered();
 

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

@@ -10,11 +10,10 @@ use ruff_db::files::FileRange;
 use ruff_db::source::{line_index, source_text};
 use ty_project::{Db, ProjectDatabase};
 
-use super::LSPResult;
 use crate::document::{DocumentKey, FileRangeExt, ToRangeExt};
-use crate::server::Result;
+use crate::session::DocumentSnapshot;
 use crate::session::client::Client;
-use crate::{DocumentSnapshot, PositionEncoding, Session};
+use crate::{PositionEncoding, Session};
 
 /// Represents the diagnostics for a text document or a notebook document.
 pub(super) enum Diagnostics {
@@ -64,30 +63,29 @@ pub(super) fn clear_diagnostics(key: &DocumentKey, client: &Client) {
 /// This function is a no-op if 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,
-    key: &DocumentKey,
-    client: &Client,
-) -> Result<()> {
+pub(super) fn publish_diagnostics(session: &Session, key: &DocumentKey, client: &Client) {
     if session.client_capabilities().pull_diagnostics {
-        return Ok(());
+        return;
     }
 
     let Some(url) = key.to_url() else {
-        return Ok(());
+        return;
     };
 
-    let path = key.path();
+    let snapshot = session.take_document_snapshot(url.clone());
 
-    let snapshot = session
-        .take_document_snapshot(url.clone())
-        .ok_or_else(|| anyhow::anyhow!("Unable to take snapshot for document with URL {url}"))
-        .with_failure_code(lsp_server::ErrorCode::InternalError)?;
+    let document = match snapshot.document() {
+        Ok(document) => document,
+        Err(err) => {
+            tracing::debug!("Failed to resolve document for URL `{}`: {}", url, err);
+            return;
+        }
+    };
 
-    let db = session.project_db_or_default(path);
+    let db = session.project_db_or_default(key.path());
 
     let Some(diagnostics) = compute_diagnostics(db, &snapshot) else {
-        return Ok(());
+        return;
     };
 
     // Sends a notification to the client with the diagnostics for the document.
@@ -95,7 +93,7 @@ pub(super) fn publish_diagnostics(
         client.send_notification::<PublishDiagnostics>(PublishDiagnosticsParams {
             uri,
             diagnostics,
-            version: Some(snapshot.query().version()),
+            version: Some(document.version()),
         });
     };
 
@@ -109,25 +107,28 @@ pub(super) fn publish_diagnostics(
             }
         }
     }
-
-    Ok(())
 }
 
 pub(super) fn compute_diagnostics(
     db: &ProjectDatabase,
     snapshot: &DocumentSnapshot,
 ) -> Option<Diagnostics> {
-    let Some(file) = snapshot.file(db) else {
-        tracing::info!(
-            "No file found for snapshot for `{}`",
-            snapshot.query().file_url()
-        );
+    let document = match snapshot.document() {
+        Ok(document) => document,
+        Err(err) => {
+            tracing::info!("Failed to resolve document for snapshot: {}", err);
+            return None;
+        }
+    };
+
+    let Some(file) = document.file(db) else {
+        tracing::info!("No file found for snapshot for `{}`", document.file_path());
         return None;
     };
 
     let diagnostics = db.check_file(file);
 
-    if let Some(notebook) = snapshot.query().as_notebook() {
+    if let Some(notebook) = document.as_notebook() {
         let mut cell_diagnostics: FxHashMap<Url, Vec<Diagnostic>> = FxHashMap::default();
 
         // Populates all relevant URLs with an empty diagnostic list. This ensures that documents

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

@@ -28,9 +28,12 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
             content_changes,
         } = params;
 
-        let Ok(key) = session.key_from_url(uri.clone()) else {
+        let key = match session.key_from_url(uri) {
+            Ok(key) => key,
+            Err(uri) => {
                 tracing::debug!("Failed to create document key from URI: {}", uri);
                 return Ok(());
+            }
         };
 
         session
@@ -54,6 +57,8 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
             }
         }
 
-        publish_diagnostics(session, &key, client)
+        publish_diagnostics(session, &key, client);
+
+        Ok(())
     }
 }

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

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

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

@@ -6,8 +6,8 @@ use crate::session::Session;
 use crate::session::client::Client;
 use crate::system::AnySystemPath;
 use lsp_server::ErrorCode;
-use lsp_types::DidCloseTextDocumentParams;
 use lsp_types::notification::DidCloseTextDocument;
+use lsp_types::{DidCloseTextDocumentParams, TextDocumentIdentifier};
 use ty_project::watch::ChangeEvent;
 
 pub(crate) struct DidCloseTextDocumentHandler;
@@ -22,13 +22,18 @@ impl SyncNotificationHandler for DidCloseTextDocumentHandler {
         client: &Client,
         params: DidCloseTextDocumentParams,
     ) -> Result<()> {
-        let Ok(key) = session.key_from_url(params.text_document.uri.clone()) else {
-            tracing::debug!(
-                "Failed to create document key from URI: {}",
-                params.text_document.uri
-            );
+        let DidCloseTextDocumentParams {
+            text_document: TextDocumentIdentifier { uri },
+        } = params;
+
+        let key = match session.key_from_url(uri) {
+            Ok(key) => key,
+            Err(uri) => {
+                tracing::debug!("Failed to create document key from URI: {}", uri);
                 return Ok(());
+            }
         };
+
         session
             .close_document(&key)
             .with_failure_code(ErrorCode::InternalError)?;

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

@@ -1,5 +1,5 @@
-use lsp_types::DidCloseNotebookDocumentParams;
 use lsp_types::notification::DidCloseNotebookDocument;
+use lsp_types::{DidCloseNotebookDocumentParams, NotebookDocumentIdentifier};
 
 use crate::server::Result;
 use crate::server::api::LSPResult;
@@ -21,13 +21,19 @@ impl SyncNotificationHandler for DidCloseNotebookHandler {
         _client: &Client,
         params: DidCloseNotebookDocumentParams,
     ) -> Result<()> {
-        let Ok(key) = session.key_from_url(params.notebook_document.uri.clone()) else {
-            tracing::debug!(
-                "Failed to create document key from URI: {}",
-                params.notebook_document.uri
-            );
+        let DidCloseNotebookDocumentParams {
+            notebook_document: NotebookDocumentIdentifier { uri },
+            ..
+        } = params;
+
+        let key = match session.key_from_url(uri) {
+            Ok(key) => key,
+            Err(uri) => {
+                tracing::debug!("Failed to create document key from URI: {}", uri);
                 return Ok(());
+            }
         };
+
         session
             .close_document(&key)
             .with_failure_code(lsp_server::ErrorCode::InternalError)?;

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

@@ -21,7 +21,9 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
     fn run(
         session: &mut Session,
         client: &Client,
-        DidOpenTextDocumentParams {
+        params: DidOpenTextDocumentParams,
+    ) -> Result<()> {
+        let DidOpenTextDocumentParams {
             text_document:
                 TextDocumentItem {
                     uri,
@@ -29,11 +31,14 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
                     version,
                     language_id,
                 },
-        }: DidOpenTextDocumentParams,
-    ) -> Result<()> {
-        let Ok(key) = session.key_from_url(uri.clone()) else {
+        } = params;
+
+        let key = match session.key_from_url(uri) {
+            Ok(key) => key,
+            Err(uri) => {
                 tracing::debug!("Failed to create document key from URI: {}", uri);
                 return Ok(());
+            }
         };
 
         let document = TextDocument::new(text, version).with_language_id(&language_id);
@@ -53,6 +58,8 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
             }
         }
 
-        publish_diagnostics(session, &key, client)
+        publish_diagnostics(session, &key, client);
+
+        Ok(())
     }
 }

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

@@ -8,11 +8,11 @@ use ty_ide::completion;
 use ty_project::ProjectDatabase;
 use ty_python_semantic::CompletionKind;
 
-use crate::DocumentSnapshot;
 use crate::document::PositionExt;
 use crate::server::api::traits::{
     BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler,
 };
+use crate::session::DocumentSnapshot;
 use crate::session::client::Client;
 
 pub(crate) struct CompletionRequestHandler;
@@ -38,8 +38,7 @@ impl BackgroundDocumentRequestHandler for CompletionRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
-            tracing::debug!("Failed to resolve file for {:?}", params);
+        let Some(file) = snapshot.file_ok(db) else {
             return Ok(None);
         };
 

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

@@ -6,11 +6,11 @@ use ruff_db::source::{line_index, source_text};
 use ty_ide::goto_type_definition;
 use ty_project::ProjectDatabase;
 
-use crate::DocumentSnapshot;
 use crate::document::{PositionExt, ToLink};
 use crate::server::api::traits::{
     BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler,
 };
+use crate::session::DocumentSnapshot;
 use crate::session::client::Client;
 
 pub(crate) struct GotoTypeDefinitionRequestHandler;
@@ -34,8 +34,7 @@ impl BackgroundDocumentRequestHandler for GotoTypeDefinitionRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
-            tracing::debug!("Failed to resolve file for {:?}", params);
+        let Some(file) = snapshot.file_ok(db) else {
             return Ok(None);
         };
 

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

@@ -1,10 +1,10 @@
 use std::borrow::Cow;
 
-use crate::DocumentSnapshot;
 use crate::document::{PositionExt, ToRangeExt};
 use crate::server::api::traits::{
     BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler,
 };
+use crate::session::DocumentSnapshot;
 use crate::session::client::Client;
 use lsp_types::request::HoverRequest;
 use lsp_types::{HoverContents, HoverParams, MarkupContent, Url};
@@ -34,8 +34,7 @@ impl BackgroundDocumentRequestHandler for HoverRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
-            tracing::debug!("Failed to resolve file for {:?}", params);
+        let Some(file) = snapshot.file_ok(db) else {
             return Ok(None);
         };
 

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

@@ -1,10 +1,10 @@
 use std::borrow::Cow;
 
-use crate::DocumentSnapshot;
 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};
@@ -33,8 +33,7 @@ impl BackgroundDocumentRequestHandler for InlayHintRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
-            tracing::debug!("Failed to resolve file for {:?}", params);
+        let Some(file) = snapshot.file_ok(db) else {
             return Ok(None);
         };
 

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

@@ -1,10 +1,10 @@
 use std::borrow::Cow;
 
-use crate::DocumentSnapshot;
 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;
@@ -24,14 +24,13 @@ impl BackgroundDocumentRequestHandler for SemanticTokensRequestHandler {
         db: &ProjectDatabase,
         snapshot: DocumentSnapshot,
         _client: &Client,
-        params: SemanticTokensParams,
+        _params: SemanticTokensParams,
     ) -> crate::server::Result<Option<SemanticTokensResult>> {
         if snapshot.client_settings().is_language_services_disabled() {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
-            tracing::debug!("Failed to resolve file for {:?}", params);
+        let Some(file) = snapshot.file_ok(db) else {
             return Ok(None);
         };
 

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

@@ -1,11 +1,11 @@
 use std::borrow::Cow;
 
-use crate::DocumentSnapshot;
 use crate::document::RangeExt;
 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, SemanticTokensRangeParams, SemanticTokensRangeResult, Url};
 use ruff_db::source::{line_index, source_text};
@@ -32,8 +32,7 @@ impl BackgroundDocumentRequestHandler for SemanticTokensRangeRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
-            tracing::debug!("Failed to resolve file for {:?}", params);
+        let Some(file) = snapshot.file_ok(db) else {
             return Ok(None);
         };
 

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

@@ -1,10 +1,10 @@
 use std::borrow::Cow;
 
-use crate::DocumentSnapshot;
 use crate::document::{PositionEncoding, PositionExt};
 use crate::server::api::traits::{
     BackgroundDocumentRequestHandler, RequestHandler, RetriableRequestHandler,
 };
+use crate::session::DocumentSnapshot;
 use crate::session::client::Client;
 use lsp_types::request::SignatureHelpRequest;
 use lsp_types::{
@@ -36,8 +36,7 @@ impl BackgroundDocumentRequestHandler for SignatureHelpRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
-            tracing::debug!("Failed to resolve file for {:?}", params);
+        let Some(file) = snapshot.file_ok(db) else {
             return Ok(None);
         };
 

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

@@ -65,7 +65,7 @@ impl BackgroundRequestHandler for WorkspaceDiagnosticRequestHandler {
                 let version = index
                     .key_from_url(url.clone())
                     .ok()
-                    .and_then(|key| index.make_document_ref(&key))
+                    .and_then(|key| index.make_document_ref(key).ok())
                     .map(|doc| i64::from(doc.version()));
 
                 // Convert diagnostics to LSP format

crates/ty_server/src/session.rs

@@ -5,23 +5,25 @@ use std::ops::{Deref, DerefMut};
 use std::sync::Arc;
 
 use anyhow::{Context, anyhow};
+use index::DocumentQueryError;
 use lsp_server::Message;
 use lsp_types::{ClientCapabilities, TextDocumentContentChangeEvent, Url};
 use options::GlobalOptions;
 use ruff_db::Db;
-use ruff_db::files::{File, system_path_to_file};
+use ruff_db::files::File;
 use ruff_db::system::{System, SystemPath, SystemPathBuf};
 use ty_project::metadata::Options;
 use ty_project::{ProjectDatabase, ProjectMetadata};
 
 pub(crate) use self::capabilities::ResolvedClientCapabilities;
-pub use self::index::DocumentQuery;
+pub(crate) use self::index::DocumentQuery;
 pub(crate) use self::options::{AllOptions, ClientOptions, DiagnosticMode};
 pub(crate) use self::settings::ClientSettings;
 use crate::document::{DocumentKey, DocumentVersion, NotebookDocument};
 use crate::session::request_queue::RequestQueue;
 use crate::system::{AnySystemPath, LSPSystem};
 use crate::{PositionEncoding, TextDocument};
+use index::Index;
 
 mod capabilities;
 pub(crate) mod client;
@@ -31,7 +33,7 @@ mod request_queue;
 mod settings;
 
 /// The global state for the LSP
-pub struct Session {
+pub(crate) struct Session {
     /// Used to retrieve information about open documents and settings.
     ///
     /// This will be [`None`] when a mutable reference is held to the index via [`index_mut`]
@@ -39,7 +41,7 @@ pub struct Session {
     /// when the mutable reference ([`MutIndexGuard`]) is dropped.
     ///
     /// [`index_mut`]: Session::index_mut
-    index: Option<Arc<index::Index>>,
+    index: Option<Arc<Index>>,
 
     /// Maps workspace folders to their respective workspace.
     workspaces: Workspaces,
@@ -71,7 +73,7 @@ impl Session {
         global_options: GlobalOptions,
         workspace_folders: Vec<(Url, ClientOptions)>,
     ) -> crate::Result<Self> {
-        let index = Arc::new(index::Index::new(global_options.into_settings()));
+        let index = Arc::new(Index::new(global_options.into_settings()));
 
         let mut workspaces = Workspaces::default();
         for (url, options) in workspace_folders {
@@ -219,7 +221,10 @@ impl Session {
             .chain(std::iter::once(&mut self.default_project))
     }
 
-    pub(crate) fn key_from_url(&self, url: Url) -> crate::Result<DocumentKey> {
+    /// Returns the [`DocumentKey`] for the given URL.
+    ///
+    /// Refer to [`Index::key_from_url`] for more details.
+    pub(crate) fn key_from_url(&self, url: Url) -> Result<DocumentKey, Url> {
         self.index().key_from_url(url)
     }
 
@@ -278,16 +283,17 @@ impl Session {
     }
 
     /// Creates a document snapshot with the URL referencing the document to snapshot.
-    ///
-    /// Returns `None` if the url can't be converted to a document key or if the document isn't open.
-    pub(crate) fn take_document_snapshot(&self, url: Url) -> Option<DocumentSnapshot> {
-        let key = self.key_from_url(url).ok()?;
-        Some(DocumentSnapshot {
+    pub(crate) fn take_document_snapshot(&self, url: Url) -> DocumentSnapshot {
+        let index = self.index();
+        DocumentSnapshot {
             resolved_client_capabilities: self.resolved_client_capabilities.clone(),
-            client_settings: self.index().global_settings(),
-            document_ref: self.index().make_document_ref(&key)?,
+            client_settings: index.global_settings(),
             position_encoding: self.position_encoding,
-        })
+            document_query_result: self
+                .key_from_url(url)
+                .map_err(DocumentQueryError::InvalidUrl)
+                .and_then(|key| index.make_document_ref(key)),
+        }
     }
 
     /// Creates a snapshot of the current state of the [`Session`].
@@ -350,7 +356,7 @@ impl Session {
     /// Panics if there's a mutable reference to the index via [`index_mut`].
     ///
     /// [`index_mut`]: Session::index_mut
-    fn index(&self) -> &index::Index {
+    fn index(&self) -> &Index {
         self.index.as_ref().unwrap()
     }
 
@@ -394,11 +400,11 @@ impl Session {
 /// When dropped, this guard restores all references to the index.
 struct MutIndexGuard<'a> {
     session: &'a mut Session,
-    index: Option<index::Index>,
+    index: Option<Index>,
 }
 
 impl Deref for MutIndexGuard<'_> {
-    type Target = index::Index;
+    type Target = Index;
 
     fn deref(&self) -> &Self::Target {
         self.index.as_ref().unwrap()
@@ -428,48 +434,69 @@ impl Drop for MutIndexGuard<'_> {
     }
 }
 
-/// An immutable snapshot of `Session` that references
-/// a specific document.
+/// An immutable snapshot of [`Session`] that references a specific document.
 #[derive(Debug)]
-pub struct DocumentSnapshot {
+pub(crate) struct DocumentSnapshot {
     resolved_client_capabilities: Arc<ResolvedClientCapabilities>,
     client_settings: Arc<ClientSettings>,
-    document_ref: index::DocumentQuery,
     position_encoding: PositionEncoding,
+    document_query_result: Result<DocumentQuery, DocumentQueryError>,
 }
 
 impl DocumentSnapshot {
+    /// Returns the resolved client capabilities that were captured during initialization.
     pub(crate) fn resolved_client_capabilities(&self) -> &ResolvedClientCapabilities {
         &self.resolved_client_capabilities
     }
 
-    pub(crate) fn query(&self) -> &index::DocumentQuery {
-        &self.document_ref
-    }
-
+    /// Returns the position encoding that was negotiated during initialization.
     pub(crate) fn encoding(&self) -> PositionEncoding {
         self.position_encoding
     }
 
+    /// Returns the client settings for this document.
     pub(crate) fn client_settings(&self) -> &ClientSettings {
         &self.client_settings
     }
 
-    pub(crate) fn file(&self, db: &dyn Db) -> Option<File> {
-        match AnySystemPath::try_from_url(self.document_ref.file_url()).ok()? {
-            AnySystemPath::System(path) => system_path_to_file(db, path).ok(),
-            AnySystemPath::SystemVirtual(virtual_path) => db
-                .files()
-                .try_virtual_file(&virtual_path)
-                .map(|virtual_file| virtual_file.file()),
+    /// Returns the result of the document query for this snapshot.
+    pub(crate) fn document(&self) -> Result<&DocumentQuery, &DocumentQueryError> {
+        self.document_query_result.as_ref()
+    }
+
+    pub(crate) fn file_ok(&self, db: &dyn Db) -> Option<File> {
+        match self.file(db) {
+            Ok(file) => Some(file),
+            Err(err) => {
+                tracing::debug!("Failed to resolve file: {}", err);
+                None
             }
         }
     }
 
+    fn file(&self, db: &dyn Db) -> Result<File, FileLookupError> {
+        let document = match self.document() {
+            Ok(document) => document,
+            Err(err) => return Err(FileLookupError::DocumentQuery(err.clone())),
+        };
+        document
+            .file(db)
+            .ok_or_else(|| FileLookupError::NotFound(document.file_path().clone()))
+    }
+}
+
+#[derive(Debug, thiserror::Error)]
+pub(crate) enum FileLookupError {
+    #[error("file not found for path `{0}`")]
+    NotFound(AnySystemPath),
+    #[error(transparent)]
+    DocumentQuery(DocumentQueryError),
+}
+
 /// An immutable snapshot of the current state of [`Session`].
 pub(crate) struct SessionSnapshot {
     projects: Vec<ProjectDatabase>,
-    index: Arc<index::Index>,
+    index: Arc<Index>,
     position_encoding: PositionEncoding,
 }
 
@@ -478,7 +505,7 @@ impl SessionSnapshot {
         &self.projects
     }
 
-    pub(crate) fn index(&self) -> &index::Index {
+    pub(crate) fn index(&self) -> &Index {
         &self.index
     }
 

crates/ty_server/src/session/index.rs

@@ -1,6 +1,8 @@
 use std::sync::Arc;
 
 use lsp_types::Url;
+use ruff_db::Db;
+use ruff_db::files::{File, system_path_to_file};
 use rustc_hash::FxHashMap;
 
 use crate::session::settings::ClientSettings;
@@ -68,16 +70,17 @@ impl Index {
         Ok(())
     }
 
-    pub(crate) fn key_from_url(&self, url: Url) -> crate::Result<DocumentKey> {
+    /// Returns the [`DocumentKey`] corresponding to the given URL.
+    ///
+    /// It returns [`Err`] with the original URL if it cannot be converted to a [`AnySystemPath`].
+    pub(crate) fn key_from_url(&self, url: Url) -> Result<DocumentKey, Url> {
         if let Some(notebook_path) = self.notebook_cells.get(&url) {
             Ok(DocumentKey::NotebookCell {
                 cell_url: url,
                 notebook_path: notebook_path.clone(),
             })
         } else {
-            let path = AnySystemPath::try_from_url(&url)
-                .map_err(|()| anyhow::anyhow!("Failed to convert URL to system path: {}", url))?;
-
+            let path = AnySystemPath::try_from_url(&url).map_err(|()| url)?;
             if path
                 .extension()
                 .is_some_and(|ext| ext.eq_ignore_ascii_case("ipynb"))
@@ -122,17 +125,27 @@ impl Index {
         Ok(())
     }
 
-    pub(crate) fn make_document_ref(&self, key: &DocumentKey) -> Option<DocumentQuery> {
+    /// Create a document reference corresponding to the given document key.
+    ///
+    /// Returns an error if the document is not found or if the path cannot be converted to a URL.
+    pub(crate) fn make_document_ref(
+        &self,
+        key: DocumentKey,
+    ) -> Result<DocumentQuery, DocumentQueryError> {
         let path = key.path();
-        let controller = self.documents.get(path)?;
-        let (cell_url, file_url) = match &key {
+        let Some(controller) = self.documents.get(path) else {
+            return Err(DocumentQueryError::NotFound(key));
+        };
+        // TODO: The `to_url` conversion shouldn't be an error because the paths themselves are
+        // constructed from the URLs but the `Index` APIs don't maintain this invariant.
+        let (cell_url, file_path) = match key {
             DocumentKey::NotebookCell {
                 cell_url,
                 notebook_path,
-            } => (Some(cell_url.clone()), notebook_path.to_url()?),
-            DocumentKey::Notebook(path) | DocumentKey::Text(path) => (None, path.to_url()?),
+            } => (Some(cell_url), notebook_path),
+            DocumentKey::Notebook(path) | DocumentKey::Text(path) => (None, path),
         };
-        Some(controller.make_ref(cell_url, file_url))
+        Ok(controller.make_ref(cell_url, file_path))
     }
 
     pub(super) fn open_text_document(&mut self, path: &AnySystemPath, document: TextDocument) {
@@ -207,15 +220,15 @@ impl DocumentController {
         Self::Notebook(Arc::new(document))
     }
 
-    fn make_ref(&self, cell_url: Option<Url>, file_url: Url) -> DocumentQuery {
+    fn make_ref(&self, cell_url: Option<Url>, file_path: AnySystemPath) -> DocumentQuery {
         match &self {
             Self::Notebook(notebook) => DocumentQuery::Notebook {
                 cell_url,
-                file_url,
+                file_path,
                 notebook: notebook.clone(),
             },
             Self::Text(document) => DocumentQuery::Text {
-                file_url,
+                file_path,
                 document: document.clone(),
             },
         }
@@ -251,26 +264,27 @@ impl DocumentController {
 }
 
 /// A read-only query to an open document.
+///
 /// This query can 'select' a text document, full notebook, or a specific notebook cell.
 /// It also includes document settings.
 #[derive(Debug, Clone)]
-pub enum DocumentQuery {
+pub(crate) enum DocumentQuery {
     Text {
-        file_url: Url,
+        file_path: AnySystemPath,
         document: Arc<TextDocument>,
     },
     Notebook {
         /// The selected notebook cell, if it exists.
         cell_url: Option<Url>,
-        /// The URL of the notebook.
-        file_url: Url,
+        /// The path to the notebook.
+        file_path: AnySystemPath,
         notebook: Arc<NotebookDocument>,
     },
 }
 
 impl DocumentQuery {
     /// Attempts to access the underlying notebook document that this query is selecting.
-    pub fn as_notebook(&self) -> Option<&NotebookDocument> {
+    pub(crate) fn as_notebook(&self) -> Option<&NotebookDocument> {
         match self {
             Self::Notebook { notebook, .. } => Some(notebook),
             Self::Text { .. } => None,
@@ -285,10 +299,10 @@ impl DocumentQuery {
         }
     }
 
-    /// Get the URL for the document selected by this query.
-    pub(crate) fn file_url(&self) -> &Url {
+    /// Get the system path for the document selected by this query.
+    pub(crate) fn file_path(&self) -> &AnySystemPath {
         match self {
-            Self::Text { file_url, .. } | Self::Notebook { file_url, .. } => file_url,
+            Self::Text { file_path, .. } | Self::Notebook { file_path, .. } => file_path,
         }
     }
 
@@ -307,4 +321,27 @@ impl DocumentQuery {
                 .and_then(|cell_uri| notebook.cell_document_by_uri(cell_uri)),
         }
     }
+
+    /// Returns the salsa interned [`File`] for the document selected by this query.
+    ///
+    /// 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 file(&self, db: &dyn Db) -> Option<File> {
+        match self.file_path() {
+            AnySystemPath::System(path) => system_path_to_file(db, path).ok(),
+            AnySystemPath::SystemVirtual(virtual_path) => db
+                .files()
+                .try_virtual_file(virtual_path)
+                .map(|virtual_file| virtual_file.file()),
+        }
+    }
+}
+
+#[derive(Debug, Clone, thiserror::Error)]
+pub(crate) enum DocumentQueryError {
+    #[error("invalid URL: {0}")]
+    InvalidUrl(Url),
+    #[error("document not found for key: {0}")]
+    NotFound(DocumentKey),
 }

crates/ty_server/src/system.rs

@@ -1,4 +1,5 @@
 use std::any::Any;
+use std::fmt;
 use std::fmt::Display;
 use std::sync::Arc;
 
@@ -97,6 +98,15 @@ impl AnySystemPath {
     }
 }
 
+impl fmt::Display for AnySystemPath {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            AnySystemPath::System(system_path) => write!(f, "{system_path}"),
+            AnySystemPath::SystemVirtual(virtual_path) => write!(f, "{virtual_path}"),
+        }
+    }
+}
+
 #[derive(Debug)]
 pub(crate) struct LSPSystem {
     /// A read-only copy of the index where the server stores all the open documents and settings.
@@ -145,7 +155,7 @@ impl LSPSystem {
     fn make_document_ref(&self, path: AnySystemPath) -> Option<DocumentQuery> {
         let index = self.index();
         let key = DocumentKey::from_path(path);
-        index.make_document_ref(&key)
+        index.make_document_ref(key).ok()
     }
 
     fn system_path_to_document_ref(&self, path: &SystemPath) -> Option<DocumentQuery> {