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

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 {