[ty] Fix stale documents on Windows (#18544)

crates/ty_server/src/document.rs

@@ -7,6 +7,8 @@ mod text_document;
 
 pub(crate) use location::ToLink;
 use lsp_types::{PositionEncodingKind, Url};
+
+use crate::system::AnySystemPath;
 pub use notebook::NotebookDocument;
 pub(crate) use range::{FileRangeExt, PositionExt, RangeExt, TextSizeExt, ToRangeExt};
 pub(crate) use text_document::DocumentVersion;
@@ -40,19 +42,38 @@ impl From<PositionEncoding> for ruff_source_file::PositionEncoding {
 /// A unique document ID, derived from a URL passed as part of an LSP request.
 /// This document ID can point to either be a standalone Python file, a full notebook, or a cell within a notebook.
 #[derive(Clone, Debug)]
-pub enum DocumentKey {
-    Notebook(Url),
-    NotebookCell(Url),
-    Text(Url),
+pub(crate) enum DocumentKey {
+    Notebook(AnySystemPath),
+    NotebookCell {
+        cell_url: Url,
+        notebook_path: AnySystemPath,
+    },
+    Text(AnySystemPath),
 }
 
 impl DocumentKey {
-    /// Returns the URL associated with the key.
-    pub(crate) fn url(&self) -> &Url {
+    /// Returns the file path associated with the key.
+    pub(crate) fn path(&self) -> &AnySystemPath {
         match self {
-            DocumentKey::NotebookCell(url)
-            | DocumentKey::Notebook(url)
-            | DocumentKey::Text(url) => url,
+            DocumentKey::Notebook(path) | DocumentKey::Text(path) => path,
+            DocumentKey::NotebookCell { notebook_path, .. } => notebook_path,
+        }
+    }
+
+    pub(crate) fn from_path(path: AnySystemPath) -> Self {
+        // For text documents, we assume it's a text document unless it's a notebook file.
+        match path.extension() {
+            Some("ipynb") => Self::Notebook(path),
+            _ => Self::Text(path),
+        }
+    }
+
+    /// Returns the URL for this document key. For notebook cells, returns the cell URL.
+    /// For other document types, converts the path to a URL.
+    pub(crate) fn to_url(&self) -> Option<Url> {
+        match self {
+            DocumentKey::NotebookCell { cell_url, .. } => Some(cell_url.clone()),
+            DocumentKey::Notebook(path) | DocumentKey::Text(path) => path.to_url(),
         }
     }
 }
@@ -60,7 +81,11 @@ impl DocumentKey {
 impl std::fmt::Display for DocumentKey {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            Self::NotebookCell(url) | Self::Notebook(url) | Self::Text(url) => url.fmt(f),
+            Self::NotebookCell { cell_url, .. } => cell_url.fmt(f),
+            Self::Notebook(path) | Self::Text(path) => match path {
+                AnySystemPath::System(system_path) => system_path.fmt(f),
+                AnySystemPath::SystemVirtual(virtual_path) => virtual_path.fmt(f),
+            },
         }
     }
 }

crates/ty_server/src/lib.rs

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

crates/ty_server/src/server/api.rs

@@ -1,6 +1,6 @@
 use crate::server::schedule::Task;
 use crate::session::Session;
-use crate::system::{AnySystemPath, url_to_any_system_path};
+use crate::system::AnySystemPath;
 use anyhow::anyhow;
 use lsp_server as server;
 use lsp_server::RequestId;
@@ -154,7 +154,7 @@ where
 
         let url = R::document_url(&params).into_owned();
 
-        let Ok(path) = url_to_any_system_path(&url) else {
+        let Ok(path) = AnySystemPath::try_from_url(&url) else {
             tracing::warn!("Ignoring request for invalid `{url}`");
             return Box::new(|_| {});
         };

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

@@ -12,10 +12,9 @@ use ruff_db::source::{line_index, source_text};
 use ty_project::{Db, ProjectDatabase};
 
 use super::LSPResult;
-use crate::document::{FileRangeExt, ToRangeExt};
+use crate::document::{DocumentKey, FileRangeExt, ToRangeExt};
 use crate::server::Result;
 use crate::session::client::Client;
-use crate::system::url_to_any_system_path;
 use crate::{DocumentSnapshot, PositionEncoding, Session};
 
 /// Represents the diagnostics for a text document or a notebook document.
@@ -42,13 +41,20 @@ impl Diagnostics {
     }
 }
 
-/// Clears the diagnostics for the document at `uri`.
+/// Clears the diagnostics for the document identified by `key`.
 ///
 /// This is done by notifying the client with an empty list of diagnostics for the document.
-pub(super) fn clear_diagnostics(uri: &Url, client: &Client) -> Result<()> {
+/// For notebook cells, this clears diagnostics for the specific cell.
+/// For other document types, this clears diagnostics for the main document.
+pub(super) fn clear_diagnostics(key: &DocumentKey, client: &Client) -> Result<()> {
+    let Some(uri) = key.to_url() else {
+        // If we can't convert to URL, we can't clear diagnostics
+        return Ok(());
+    };
+
     client
         .send_notification::<PublishDiagnostics>(PublishDiagnosticsParams {
-            uri: uri.clone(),
+            uri,
             diagnostics: vec![],
             version: None,
         })
@@ -62,21 +68,27 @@ pub(super) fn clear_diagnostics(uri: &Url, client: &Client) -> Result<()> {
 /// 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, url: Url, client: &Client) -> Result<()> {
+pub(super) fn publish_diagnostics(
+    session: &Session,
+    key: &DocumentKey,
+    client: &Client,
+) -> Result<()> {
     if session.client_capabilities().pull_diagnostics {
         return Ok(());
     }
 
-    let Ok(path) = url_to_any_system_path(&url) else {
+    let Some(url) = key.to_url() else {
         return Ok(());
     };
 
+    let path = key.path();
+
     let snapshot = session
         .take_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 db = session.project_db_or_default(&path);
+    let db = session.project_db_or_default(path);
 
     let Some(diagnostics) = compute_diagnostics(db, &snapshot) else {
         return Ok(());

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

@@ -8,7 +8,7 @@ use crate::server::api::diagnostics::publish_diagnostics;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::session::Session;
 use crate::session::client::Client;
-use crate::system::{AnySystemPath, url_to_any_system_path};
+use crate::system::AnySystemPath;
 use ty_project::watch::ChangeEvent;
 
 pub(crate) struct DidChangeTextDocumentHandler;
@@ -28,30 +28,32 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
             content_changes,
         } = params;
 
-        let Ok(path) = url_to_any_system_path(&uri) else {
+        let Ok(key) = session.key_from_url(uri.clone()) else {
+            tracing::debug!("Failed to create document key from URI: {}", uri);
             return Ok(());
         };
 
-        let key = session.key_from_url(uri.clone());
-
         session
             .update_text_document(&key, content_changes, version)
             .with_failure_code(ErrorCode::InternalError)?;
 
-        match path {
+        match key.path() {
             AnySystemPath::System(path) => {
                 let db = match session.project_db_for_path_mut(path.as_std_path()) {
                     Some(db) => db,
                     None => session.default_project_db_mut(),
                 };
-                db.apply_changes(vec![ChangeEvent::file_content_changed(path)], None);
+                db.apply_changes(vec![ChangeEvent::file_content_changed(path.clone())], None);
             }
             AnySystemPath::SystemVirtual(virtual_path) => {
                 let db = session.default_project_db_mut();
-                db.apply_changes(vec![ChangeEvent::ChangedVirtual(virtual_path)], None);
+                db.apply_changes(
+                    vec![ChangeEvent::ChangedVirtual(virtual_path.clone())],
+                    None,
+                );
             }
         }
 
-        publish_diagnostics(session, uri, client)
+        publish_diagnostics(session, &key, client)
     }
 }

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

@@ -4,7 +4,7 @@ use crate::server::api::diagnostics::publish_diagnostics;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::session::Session;
 use crate::session::client::Client;
-use crate::system::{AnySystemPath, url_to_any_system_path};
+use crate::system::AnySystemPath;
 use lsp_types as types;
 use lsp_types::{FileChangeType, notification as notif};
 use rustc_hash::FxHashMap;
@@ -26,7 +26,7 @@ impl SyncNotificationHandler for DidChangeWatchedFiles {
         let mut events_by_db: FxHashMap<_, Vec<ChangeEvent>> = FxHashMap::default();
 
         for change in params.changes {
-            let path = match url_to_any_system_path(&change.uri) {
+            let path = match AnySystemPath::try_from_url(&change.uri) {
                 Ok(path) => path,
                 Err(err) => {
                     tracing::warn!(
@@ -111,8 +111,8 @@ impl SyncNotificationHandler for DidChangeWatchedFiles {
                     )
                     .with_failure_code(lsp_server::ErrorCode::InternalError)?;
             } else {
-                for url in session.text_document_urls() {
-                    publish_diagnostics(session, url.clone(), client)?;
+                for key in session.text_document_keys() {
+                    publish_diagnostics(session, &key, client)?;
                 }
             }
 

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

@@ -4,7 +4,7 @@ use crate::server::api::diagnostics::clear_diagnostics;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::session::Session;
 use crate::session::client::Client;
-use crate::system::{AnySystemPath, url_to_any_system_path};
+use crate::system::AnySystemPath;
 use lsp_server::ErrorCode;
 use lsp_types::DidCloseTextDocumentParams;
 use lsp_types::notification::DidCloseTextDocument;
@@ -22,22 +22,25 @@ impl SyncNotificationHandler for DidCloseTextDocumentHandler {
         client: &Client,
         params: DidCloseTextDocumentParams,
     ) -> Result<()> {
-        let Ok(path) = url_to_any_system_path(&params.text_document.uri) else {
+        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
+            );
             return Ok(());
         };
-
-        let key = session.key_from_url(params.text_document.uri);
         session
             .close_document(&key)
             .with_failure_code(ErrorCode::InternalError)?;
 
-        if let AnySystemPath::SystemVirtual(virtual_path) = path {
+        if let AnySystemPath::SystemVirtual(virtual_path) = key.path() {
             let db = session.default_project_db_mut();
-            db.apply_changes(vec![ChangeEvent::DeletedVirtual(virtual_path)], None);
+            db.apply_changes(
+                vec![ChangeEvent::DeletedVirtual(virtual_path.clone())],
+                None,
+            );
         }
 
-        clear_diagnostics(key.url(), client)?;
-
-        Ok(())
+        clear_diagnostics(&key, client)
     }
 }

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

@@ -6,7 +6,7 @@ use crate::server::api::LSPResult;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::session::Session;
 use crate::session::client::Client;
-use crate::system::{AnySystemPath, url_to_any_system_path};
+use crate::system::AnySystemPath;
 use ty_project::watch::ChangeEvent;
 
 pub(crate) struct DidCloseNotebookHandler;
@@ -21,18 +21,23 @@ impl SyncNotificationHandler for DidCloseNotebookHandler {
         _client: &Client,
         params: DidCloseNotebookDocumentParams,
     ) -> Result<()> {
-        let Ok(path) = url_to_any_system_path(&params.notebook_document.uri) else {
+        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
+            );
             return Ok(());
         };
-
-        let key = session.key_from_url(params.notebook_document.uri);
         session
             .close_document(&key)
             .with_failure_code(lsp_server::ErrorCode::InternalError)?;
 
-        if let AnySystemPath::SystemVirtual(virtual_path) = path {
+        if let AnySystemPath::SystemVirtual(virtual_path) = key.path() {
             let db = session.default_project_db_mut();
-            db.apply_changes(vec![ChangeEvent::DeletedVirtual(virtual_path)], None);
+            db.apply_changes(
+                vec![ChangeEvent::DeletedVirtual(virtual_path.clone())],
+                None,
+            );
         }
 
         Ok(())

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

@@ -7,7 +7,7 @@ use crate::server::api::diagnostics::publish_diagnostics;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::session::Session;
 use crate::session::client::Client;
-use crate::system::{AnySystemPath, url_to_any_system_path};
+use crate::system::AnySystemPath;
 use ruff_db::Db;
 use ty_project::watch::ChangeEvent;
 
@@ -31,20 +31,21 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
                 },
         }: DidOpenTextDocumentParams,
     ) -> Result<()> {
-        let Ok(path) = url_to_any_system_path(&uri) else {
+        let Ok(key) = session.key_from_url(uri.clone()) else {
+            tracing::debug!("Failed to create document key from URI: {}", uri);
             return Ok(());
         };
 
         let document = TextDocument::new(text, version).with_language_id(&language_id);
-        session.open_text_document(uri.clone(), document);
+        session.open_text_document(key.path(), document);
 
-        match &path {
-            AnySystemPath::System(path) => {
-                let db = match session.project_db_for_path_mut(path.as_std_path()) {
+        match key.path() {
+            AnySystemPath::System(system_path) => {
+                let db = match session.project_db_for_path_mut(system_path.as_std_path()) {
                     Some(db) => db,
                     None => session.default_project_db_mut(),
                 };
-                db.apply_changes(vec![ChangeEvent::Opened(path.clone())], None);
+                db.apply_changes(vec![ChangeEvent::Opened(system_path.clone())], None);
             }
             AnySystemPath::SystemVirtual(virtual_path) => {
                 let db = session.default_project_db_mut();
@@ -52,6 +53,6 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
             }
         }
 
-        publish_diagnostics(session, uri, client)
+        publish_diagnostics(session, &key, client)
     }
 }

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

@@ -11,7 +11,7 @@ use crate::server::api::LSPResult;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::session::Session;
 use crate::session::client::Client;
-use crate::system::{AnySystemPath, url_to_any_system_path};
+use crate::system::AnySystemPath;
 
 pub(crate) struct DidOpenNotebookHandler;
 
@@ -25,7 +25,7 @@ impl SyncNotificationHandler for DidOpenNotebookHandler {
         _client: &Client,
         params: DidOpenNotebookDocumentParams,
     ) -> Result<()> {
-        let Ok(path) = url_to_any_system_path(&params.notebook_document.uri) else {
+        let Ok(path) = AnySystemPath::try_from_url(&params.notebook_document.uri) else {
             return Ok(());
         };
 
@@ -36,19 +36,19 @@ impl SyncNotificationHandler for DidOpenNotebookHandler {
             params.cell_text_documents,
         )
         .with_failure_code(ErrorCode::InternalError)?;
-        session.open_notebook_document(params.notebook_document.uri, notebook);
+        session.open_notebook_document(&path, notebook);
 
-        match path {
-            AnySystemPath::System(path) => {
-                let db = match session.project_db_for_path_mut(path.as_std_path()) {
+        match &path {
+            AnySystemPath::System(system_path) => {
+                let db = match session.project_db_for_path_mut(system_path.as_std_path()) {
                     Some(db) => db,
                     None => session.default_project_db_mut(),
                 };
-                db.apply_changes(vec![ChangeEvent::Opened(path)], None);
+                db.apply_changes(vec![ChangeEvent::Opened(system_path.clone())], None);
             }
             AnySystemPath::SystemVirtual(virtual_path) => {
                 let db = session.default_project_db_mut();
-                db.files().virtual_file(db, &virtual_path);
+                db.files().virtual_file(db, virtual_path);
             }
         }
 

crates/ty_server/src/session.rs

@@ -19,7 +19,7 @@ pub use self::settings::ClientSettings;
 pub(crate) use self::settings::Experimental;
 use crate::document::{DocumentKey, DocumentVersion, NotebookDocument};
 use crate::session::request_queue::RequestQueue;
-use crate::system::{AnySystemPath, LSPSystem, url_to_any_system_path};
+use crate::system::{AnySystemPath, LSPSystem};
 use crate::{PositionEncoding, TextDocument};
 
 mod capabilities;
@@ -158,35 +158,43 @@ impl Session {
             .unwrap()
     }
 
-    pub fn key_from_url(&self, url: Url) -> DocumentKey {
+    pub(crate) fn key_from_url(&self, url: Url) -> crate::Result<DocumentKey> {
         self.index().key_from_url(url)
     }
 
     /// 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 fn take_snapshot(&self, url: Url) -> Option<DocumentSnapshot> {
-        let key = self.key_from_url(url);
+        let key = self.key_from_url(url).ok()?;
         Some(DocumentSnapshot {
             resolved_client_capabilities: self.resolved_client_capabilities.clone(),
-            document_ref: self.index().make_document_ref(key)?,
+            document_ref: self.index().make_document_ref(&key)?,
             position_encoding: self.position_encoding,
         })
     }
 
-    /// Iterates over the LSP URLs for all open text documents. These URLs are valid file paths.
-    pub(super) fn text_document_urls(&self) -> impl Iterator<Item = &Url> + '_ {
-        self.index().text_document_urls()
+    /// Iterates over the document keys for all open text documents.
+    pub(super) fn text_document_keys(&self) -> impl Iterator<Item = DocumentKey> + '_ {
+        self.index()
+            .text_document_paths()
+            .map(|path| DocumentKey::Text(path.clone()))
     }
 
-    /// Registers a notebook document at the provided `url`.
+    /// Registers a notebook document at the provided `path`.
     /// If a document is already open here, it will be overwritten.
-    pub fn open_notebook_document(&mut self, url: Url, document: NotebookDocument) {
-        self.index_mut().open_notebook_document(url, document);
+    pub(crate) fn open_notebook_document(
+        &mut self,
+        path: &AnySystemPath,
+        document: NotebookDocument,
+    ) {
+        self.index_mut().open_notebook_document(path, document);
     }
 
-    /// Registers a text document at the provided `url`.
+    /// Registers a text document at the provided `path`.
     /// If a document is already open here, it will be overwritten.
-    pub(crate) fn open_text_document(&mut self, url: Url, document: TextDocument) {
-        self.index_mut().open_text_document(url, document);
+    pub(crate) fn open_text_document(&mut self, path: &AnySystemPath, document: TextDocument) {
+        self.index_mut().open_text_document(path, document);
     }
 
     /// Updates a text document at the associated `key`.
@@ -314,7 +322,7 @@ impl DocumentSnapshot {
     }
 
     pub(crate) fn file(&self, db: &dyn Db) -> Option<File> {
-        match url_to_any_system_path(self.document_ref.file_url()).ok()? {
+        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()

crates/ty_server/src/session/index.rs

@@ -1,4 +1,3 @@
-use std::path::Path;
 use std::sync::Arc;
 
 use lsp_types::Url;
@@ -7,6 +6,7 @@ use rustc_hash::FxHashMap;
 use crate::{
     PositionEncoding, TextDocument,
     document::{DocumentKey, DocumentVersion, NotebookDocument},
+    system::AnySystemPath,
 };
 
 use super::ClientSettings;
@@ -14,11 +14,11 @@ use super::ClientSettings;
 /// Stores and tracks all open documents in a session, along with their associated settings.
 #[derive(Default, Debug)]
 pub(crate) struct Index {
-    /// Maps all document file URLs to the associated document controller
-    documents: FxHashMap<Url, DocumentController>,
+    /// Maps all document file paths to the associated document controller
+    documents: FxHashMap<AnySystemPath, DocumentController>,
 
-    /// Maps opaque cell URLs to a notebook URL (document)
-    notebook_cells: FxHashMap<Url, Url>,
+    /// Maps opaque cell URLs to a notebook path (document)
+    notebook_cells: FxHashMap<Url, AnySystemPath>,
 
     /// Global settings provided by the client.
     #[expect(dead_code)]
@@ -34,18 +34,18 @@ impl Index {
         }
     }
 
-    pub(super) fn text_document_urls(&self) -> impl Iterator<Item = &Url> + '_ {
+    pub(super) fn text_document_paths(&self) -> impl Iterator<Item = &AnySystemPath> + '_ {
         self.documents
             .iter()
-            .filter_map(|(url, doc)| doc.as_text().and(Some(url)))
+            .filter_map(|(path, doc)| doc.as_text().and(Some(path)))
     }
 
     #[expect(dead_code)]
-    pub(super) fn notebook_document_urls(&self) -> impl Iterator<Item = &Url> + '_ {
+    pub(super) fn notebook_document_paths(&self) -> impl Iterator<Item = &AnySystemPath> + '_ {
         self.documents
             .iter()
             .filter(|(_, doc)| doc.as_notebook().is_some())
-            .map(|(url, _)| url)
+            .map(|(path, _)| path)
     }
 
     pub(super) fn update_text_document(
@@ -57,7 +57,7 @@ impl Index {
     ) -> crate::Result<()> {
         let controller = self.document_controller_for_key(key)?;
         let Some(document) = controller.as_text_mut() else {
-            anyhow::bail!("Text document URI does not point to a text document");
+            anyhow::bail!("Text document path does not point to a text document");
         };
 
         if content_changes.is_empty() {
@@ -70,16 +70,24 @@ impl Index {
         Ok(())
     }
 
-    pub(crate) fn key_from_url(&self, url: Url) -> DocumentKey {
-        if self.notebook_cells.contains_key(&url) {
-            DocumentKey::NotebookCell(url)
-        } else if Path::new(url.path())
-            .extension()
-            .is_some_and(|ext| ext.eq_ignore_ascii_case("ipynb"))
-        {
-            DocumentKey::Notebook(url)
+    pub(crate) fn key_from_url(&self, url: Url) -> crate::Result<DocumentKey> {
+        if let Some(notebook_path) = self.notebook_cells.get(&url) {
+            Ok(DocumentKey::NotebookCell {
+                cell_url: url,
+                notebook_path: notebook_path.clone(),
+            })
         } else {
-            DocumentKey::Text(url)
+            let path = AnySystemPath::try_from_url(&url)
+                .map_err(|()| anyhow::anyhow!("Failed to convert URL to system path: {}", url))?;
+
+            if path
+                .extension()
+                .is_some_and(|ext| ext.eq_ignore_ascii_case("ipynb"))
+            {
+                Ok(DocumentKey::Notebook(path))
+            } else {
+                Ok(DocumentKey::Text(path))
+            }
         }
     }
 
@@ -98,48 +106,55 @@ impl Index {
             ..
         }) = 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}`");
-            };
+            let notebook_path = key.path().clone();
 
             for opened_cell in did_open {
                 self.notebook_cells
-                    .insert(opened_cell.uri.clone(), path.clone());
+                    .insert(opened_cell.uri.clone(), notebook_path.clone());
             }
             // deleted notebook cells are closed via textDocument/didClose - we don't close them here.
         }
 
         let controller = self.document_controller_for_key(key)?;
         let Some(notebook) = controller.as_notebook_mut() else {
-            anyhow::bail!("Notebook document URI does not point to a notebook document");
+            anyhow::bail!("Notebook document path does not point to a notebook document");
         };
 
         notebook.update(cells, metadata, new_version, encoding)?;
         Ok(())
     }
 
-    pub(crate) fn make_document_ref(&self, key: DocumentKey) -> Option<DocumentQuery> {
-        let url = self.url_for_key(&key)?.clone();
-        let controller = self.documents.get(&url)?;
-        let cell_url = match key {
-            DocumentKey::NotebookCell(cell_url) => Some(cell_url),
-            _ => None,
+    pub(crate) fn make_document_ref(&self, key: &DocumentKey) -> Option<DocumentQuery> {
+        let path = key.path();
+        let controller = self.documents.get(path)?;
+        let (cell_url, file_url) = 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(controller.make_ref(cell_url, url))
+        Some(controller.make_ref(cell_url, file_url))
     }
 
-    pub(super) fn open_text_document(&mut self, url: Url, document: TextDocument) {
+    pub(super) fn open_text_document(&mut self, path: &AnySystemPath, document: TextDocument) {
         self.documents
-            .insert(url, DocumentController::new_text(document));
+            .insert(path.clone(), DocumentController::new_text(document));
     }
 
-    pub(super) fn open_notebook_document(&mut self, notebook_url: Url, document: NotebookDocument) {
+    pub(super) fn open_notebook_document(
+        &mut self,
+        notebook_path: &AnySystemPath,
+        document: NotebookDocument,
+    ) {
         for cell_url in document.cell_urls() {
             self.notebook_cells
-                .insert(cell_url.clone(), notebook_url.clone());
+                .insert(cell_url.clone(), notebook_path.clone());
         }
-        self.documents
-            .insert(notebook_url, DocumentController::new_notebook(document));
+        self.documents.insert(
+            notebook_path.clone(),
+            DocumentController::new_notebook(document),
+        );
     }
 
     pub(super) fn close_document(&mut self, key: &DocumentKey) -> crate::Result<()> {
@@ -148,18 +163,16 @@ impl Index {
         // 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}",);
+        if let DocumentKey::NotebookCell { cell_url, .. } = key {
+            if self.notebook_cells.remove(cell_url).is_none() {
+                tracing::warn!("Tried to remove a notebook cell that does not exist: {cell_url}",);
             }
             return Ok(());
         }
-        let Some(url) = self.url_for_key(key).cloned() else {
-            anyhow::bail!("Tried to close unavailable document `{key}`");
-        };
+        let path = key.path();
 
-        let Some(_) = self.documents.remove(&url) else {
-            anyhow::bail!("tried to close document that didn't exist at {}", url)
+        let Some(_) = self.documents.remove(path) else {
+            anyhow::bail!("tried to close document that didn't exist at {}", key)
         };
         Ok(())
     }
@@ -168,21 +181,12 @@ impl Index {
         &mut self,
         key: &DocumentKey,
     ) -> crate::Result<&mut DocumentController> {
-        let Some(url) = self.url_for_key(key).cloned() else {
-            anyhow::bail!("Tried to open unavailable document `{key}`");
-        };
-        let Some(controller) = self.documents.get_mut(&url) else {
-            anyhow::bail!("Document controller not available at `{}`", url);
+        let path = key.path();
+        let Some(controller) = self.documents.get_mut(path) else {
+            anyhow::bail!("Document controller not available at `{}`", key);
         };
         Ok(controller)
     }
-
-    fn url_for_key<'a>(&'a self, key: &'a DocumentKey) -> Option<&'a Url> {
-        match key {
-            DocumentKey::Notebook(path) | DocumentKey::Text(path) => Some(path),
-            DocumentKey::NotebookCell(uri) => self.notebook_cells.get(uri),
-        }
-    }
 }
 
 /// A mutable handler to an underlying document.
@@ -263,19 +267,6 @@ pub enum DocumentQuery {
 }
 
 impl DocumentQuery {
-    /// Retrieve the original key that describes this document query.
-    #[expect(dead_code)]
-    pub(crate) fn make_key(&self) -> DocumentKey {
-        match self {
-            Self::Text { file_url, .. } => DocumentKey::Text(file_url.clone()),
-            Self::Notebook {
-                cell_url: Some(cell_uri),
-                ..
-            } => DocumentKey::NotebookCell(cell_uri.clone()),
-            Self::Notebook { file_url, .. } => DocumentKey::Notebook(file_url.clone()),
-        }
-    }
-
     /// Attempts to access the underlying notebook document that this query is selecting.
     pub fn as_notebook(&self) -> Option<&NotebookDocument> {
         match self {

crates/ty_server/src/system.rs

@@ -14,28 +14,9 @@ use ruff_notebook::{Notebook, NotebookError};
 use ty_python_semantic::Db;
 
 use crate::DocumentQuery;
+use crate::document::DocumentKey;
 use crate::session::index::Index;
 
-/// Converts the given [`Url`] to an [`AnySystemPath`].
-///
-/// If the URL scheme is `file`, then the path is converted to a [`SystemPathBuf`]. Otherwise, the
-/// URL is converted to a [`SystemVirtualPathBuf`].
-///
-/// This fails in the following cases:
-/// * The URL cannot be converted to a file path (refer to [`Url::to_file_path`]).
-/// * If the URL is not a valid UTF-8 string.
-pub(crate) fn url_to_any_system_path(url: &Url) -> std::result::Result<AnySystemPath, ()> {
-    if url.scheme() == "file" {
-        Ok(AnySystemPath::System(
-            SystemPathBuf::from_path_buf(url.to_file_path()?).map_err(|_| ())?,
-        ))
-    } else {
-        Ok(AnySystemPath::SystemVirtual(
-            SystemVirtualPath::new(url.as_str()).to_path_buf(),
-        ))
-    }
-}
-
 pub(crate) fn file_to_url(db: &dyn Db, file: File) -> Option<Url> {
     match file.path(db) {
         FilePath::System(system) => Url::from_file_path(system.as_std_path()).ok(),
@@ -47,19 +28,57 @@ pub(crate) fn file_to_url(db: &dyn Db, file: File) -> Option<Url> {
 }
 
 /// Represents either a [`SystemPath`] or a [`SystemVirtualPath`].
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Hash, PartialEq, Eq)]
 pub(crate) enum AnySystemPath {
     System(SystemPathBuf),
     SystemVirtual(SystemVirtualPathBuf),
 }
 
 impl AnySystemPath {
+    /// Converts the given [`Url`] to an [`AnySystemPath`].
+    ///
+    /// If the URL scheme is `file`, then the path is converted to a [`SystemPathBuf`]. Otherwise, the
+    /// URL is converted to a [`SystemVirtualPathBuf`].
+    ///
+    /// This fails in the following cases:
+    /// * The URL cannot be converted to a file path (refer to [`Url::to_file_path`]).
+    /// * If the URL is not a valid UTF-8 string.
+    pub(crate) fn try_from_url(url: &Url) -> std::result::Result<Self, ()> {
+        if url.scheme() == "file" {
+            Ok(AnySystemPath::System(
+                SystemPathBuf::from_path_buf(url.to_file_path()?).map_err(|_| ())?,
+            ))
+        } else {
+            Ok(AnySystemPath::SystemVirtual(
+                SystemVirtualPath::new(url.as_str()).to_path_buf(),
+            ))
+        }
+    }
+
     pub(crate) const fn as_system(&self) -> Option<&SystemPathBuf> {
         match self {
             AnySystemPath::System(system_path_buf) => Some(system_path_buf),
             AnySystemPath::SystemVirtual(_) => None,
         }
     }
+
+    /// Returns the extension of the path, if any.
+    pub(crate) fn extension(&self) -> Option<&str> {
+        match self {
+            AnySystemPath::System(system_path) => system_path.extension(),
+            AnySystemPath::SystemVirtual(virtual_path) => virtual_path.extension(),
+        }
+    }
+
+    /// Converts the path to a URL.
+    pub(crate) fn to_url(&self) -> Option<Url> {
+        match self {
+            AnySystemPath::System(system_path) => {
+                Url::from_file_path(system_path.as_std_path()).ok()
+            }
+            AnySystemPath::SystemVirtual(virtual_path) => Url::parse(virtual_path.as_str()).ok(),
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -107,39 +126,29 @@ impl LSPSystem {
         self.index.as_ref().unwrap()
     }
 
-    fn make_document_ref(&self, url: Url) -> Option<DocumentQuery> {
+    fn make_document_ref(&self, path: AnySystemPath) -> Option<DocumentQuery> {
         let index = self.index();
-        let key = index.key_from_url(url);
-        index.make_document_ref(key)
+        let key = DocumentKey::from_path(path);
+        index.make_document_ref(&key)
     }
 
-    fn system_path_to_document_ref(&self, path: &SystemPath) -> Result<Option<DocumentQuery>> {
-        let url = Url::from_file_path(path.as_std_path()).map_err(|()| {
-            std::io::Error::new(
-                std::io::ErrorKind::InvalidInput,
-                format!("Failed to convert system path to URL: {path:?}"),
-            )
-        })?;
-        Ok(self.make_document_ref(url))
+    fn system_path_to_document_ref(&self, path: &SystemPath) -> Option<DocumentQuery> {
+        let any_path = AnySystemPath::System(path.to_path_buf());
+        self.make_document_ref(any_path)
     }
 
     fn system_virtual_path_to_document_ref(
         &self,
         path: &SystemVirtualPath,
-    ) -> Result<Option<DocumentQuery>> {
-        let url = Url::parse(path.as_str()).map_err(|_| {
-            std::io::Error::new(
-                std::io::ErrorKind::InvalidInput,
-                format!("Failed to convert virtual path to URL: {path:?}"),
-            )
-        })?;
-        Ok(self.make_document_ref(url))
+    ) -> Option<DocumentQuery> {
+        let any_path = AnySystemPath::SystemVirtual(path.to_path_buf());
+        self.make_document_ref(any_path)
     }
 }
 
 impl System for LSPSystem {
     fn path_metadata(&self, path: &SystemPath) -> Result<Metadata> {
-        let document = self.system_path_to_document_ref(path)?;
+        let document = self.system_path_to_document_ref(path);
 
         if let Some(document) = document {
             Ok(Metadata::new(
@@ -161,7 +170,7 @@ impl System for LSPSystem {
     }
 
     fn read_to_string(&self, path: &SystemPath) -> Result<String> {
-        let document = self.system_path_to_document_ref(path)?;
+        let document = self.system_path_to_document_ref(path);
 
         match document {
             Some(DocumentQuery::Text { document, .. }) => Ok(document.contents().to_string()),
@@ -170,7 +179,7 @@ impl System for LSPSystem {
     }
 
     fn read_to_notebook(&self, path: &SystemPath) -> std::result::Result<Notebook, NotebookError> {
-        let document = self.system_path_to_document_ref(path)?;
+        let document = self.system_path_to_document_ref(path);
 
         match document {
             Some(DocumentQuery::Text { document, .. }) => {
@@ -183,7 +192,7 @@ impl System for LSPSystem {
 
     fn read_virtual_path_to_string(&self, path: &SystemVirtualPath) -> Result<String> {
         let document = self
-            .system_virtual_path_to_document_ref(path)?
+            .system_virtual_path_to_document_ref(path)
             .ok_or_else(|| virtual_path_not_found(path))?;
 
         if let DocumentQuery::Text { document, .. } = &document {
@@ -198,7 +207,7 @@ impl System for LSPSystem {
         path: &SystemVirtualPath,
     ) -> std::result::Result<Notebook, NotebookError> {
         let document = self
-            .system_virtual_path_to_document_ref(path)?
+            .system_virtual_path_to_document_ref(path)
             .ok_or_else(|| virtual_path_not_found(path))?;
 
         match document {