fix

Author: asukaminato0721

pyrefly/lib/lsp/non_wasm/server.rs

@@ -201,10 +201,12 @@ use pyrefly_python::module::TextRangeWithModule;
 use pyrefly_python::module_name::ModuleName;
 use pyrefly_python::module_name::ModuleNameWithKind;
 use pyrefly_python::module_path::ModulePath;
+use pyrefly_python::module_path::ModulePathDetails;
 use pyrefly_python::short_identifier::ShortIdentifier;
 use pyrefly_util::absolutize::Absolutize as _;
 use pyrefly_util::arc_id::ArcId;
 use pyrefly_util::events::CategorizedEvents;
+use pyrefly_util::fs_anyhow;
 use pyrefly_util::globs::FilteredGlobs;
 use pyrefly_util::globs::HiddenDirFilter;
 use pyrefly_util::includes::Includes as _;
@@ -313,6 +315,7 @@ use crate::lsp::wasm::provide_type::ProvideTypeParams;
 use crate::lsp::wasm::provide_type::ProvideTypeResponse;
 use crate::lsp::wasm::provide_type::provide_type;
 use crate::module::bundled::BundledStub;
+use crate::state::load::FileContents;
 use crate::state::load::LspFile;
 use crate::state::lsp::DisplayTypeErrors;
 use crate::state::lsp::FindDefinitionItemWithDocstring;
@@ -1528,6 +1531,41 @@ impl Server {
         }
     }
 
+    fn stale_definition_request_disk_contents(
+        &self,
+        transaction: &Transaction<'_>,
+        handle: &Handle,
+        path: &Path,
+    ) -> Option<String> {
+        let disk_contents = fs_anyhow::read_to_string(path).ok()?;
+        if let Some(open_file) = self.open_files.read().get(path) {
+            return match &**open_file {
+                LspFile::Source(contents) if contents.as_str() != disk_contents.as_str() => {
+                    Some(disk_contents)
+                }
+                LspFile::Source(_) | LspFile::Notebook(_) => None,
+            };
+        }
+        let module_info = transaction.get_module_info(handle)?;
+        if module_info.is_notebook() || module_info.contents().as_str() == disk_contents.as_str() {
+            None
+        } else {
+            Some(disk_contents)
+        }
+    }
+
+    fn invalidate_all_definition_files(&self, transaction: &mut Transaction<'_>) {
+        let files = transaction
+            .handles()
+            .into_iter()
+            .filter_map(|handle| match handle.path().details() {
+                ModulePathDetails::FileSystem(path) => Some(path.as_path().to_path_buf()),
+                _ => None,
+            })
+            .collect::<Vec<_>>();
+        transaction.invalidate_disk(&files);
+    }
+
     /// Returns a snapshot of all currently open notebooks, keyed by their filesystem path.
     /// Used to remap file paths to notebook cell URIs in closures that don't have
     /// access to `self`.
@@ -1861,7 +1899,7 @@ impl Server {
                             params, &x.id,
                         )
                     {
-                        let response = match self.goto_definition(&transaction, params) {
+                        let response = match self.goto_definition(&mut transaction, params) {
                             Ok(response) => response,
                             Err(reason) => {
                                 telemetry_event.set_empty_response_reason(reason);
@@ -3933,19 +3971,18 @@ impl Server {
             .map(|(handle, _)| handle)
     }
 
-    fn goto_definition(
+    fn goto_definition_response(
         &self,
         transaction: &Transaction<'_>,
-        params: GotoDefinitionParams,
+        handle: &Handle,
+        uri: &Url,
+        position: Position,
     ) -> Result<Option<GotoDefinitionResponse>, EmptyResponseReason> {
-        let uri = &params.text_document_position_params.text_document.uri;
-        let handle = self.make_handle_if_enabled(uri, Some(GotoDefinition::METHOD))?;
         let info = transaction
-            .get_module_info(&handle)
+            .get_module_info(handle)
             .ok_or(EmptyResponseReason::ModuleInfoNotFound)?;
-        let range =
-            self.from_lsp_position(uri, &info, params.text_document_position_params.position);
-        let targets = transaction.goto_definition(&handle, range)?;
+        let range = self.from_lsp_position(uri, &info, position);
+        let targets = transaction.goto_definition(handle, range)?;
         let mut lsp_targets = targets
             .iter()
             .filter_map(|x| self.to_lsp_location(x))
@@ -3967,6 +4004,55 @@ impl Server {
         }
     }
 
+    fn goto_definition(
+        &self,
+        transaction: &mut Transaction<'_>,
+        params: GotoDefinitionParams,
+    ) -> Result<Option<GotoDefinitionResponse>, EmptyResponseReason> {
+        let uri = &params.text_document_position_params.text_document.uri;
+        let handle = self.make_handle_if_enabled(uri, Some(GotoDefinition::METHOD))?;
+        let path = self.path_for_uri(uri);
+        let stale_disk_contents = path.as_deref().and_then(|path| {
+            self.stale_definition_request_disk_contents(transaction, &handle, path)
+        });
+        if stale_disk_contents.is_some() {
+            self.invalidate_all_definition_files(transaction);
+            transaction.run(
+                std::slice::from_ref(&handle),
+                Require::Everything,
+                Some(&self.lsp_thread_pool),
+            );
+        }
+        let position = params.text_document_position_params.position;
+        let response = self.goto_definition_response(transaction, &handle, uri, position);
+        if matches!(response, Ok(Some(_))) {
+            return response;
+        }
+        let Some(path) = path else {
+            return response;
+        };
+        let Some(disk_contents) = stale_disk_contents else {
+            return response;
+        };
+        let should_retry = matches!(
+            self.open_files.read().get(&path).map(|file| &**file),
+            Some(LspFile::Source(_))
+        );
+        if !should_retry {
+            return response;
+        }
+        transaction.set_memory(vec![(
+            path,
+            Some(Arc::new(FileContents::from_source(disk_contents))),
+        )]);
+        transaction.run(
+            std::slice::from_ref(&handle),
+            Require::Everything,
+            Some(&self.lsp_thread_pool),
+        );
+        self.goto_definition_response(transaction, &handle, uri, position)
+    }
+
     fn goto_declaration(
         &self,
         transaction: &Transaction<'_>,

pyrefly/lib/test/lsp/lsp_interaction/definition.rs

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+use std::fs;
 use std::path::PathBuf;
 use std::sync::Arc;
 
@@ -141,6 +142,47 @@ fn test_go_to_def_basic(root: &TempDir, workspace_folders: Option<Vec<(String, U
         .unwrap();
 }
 
+#[test]
+fn goto_definition_recovers_from_stale_disk_source_file() {
+    let root = TempDir::with_prefix("pyrefly_goto_def_stale_disk").unwrap();
+    fs::write(
+        root.path().join("foo.py"),
+        "from bar import target\n\nvalue = target\n",
+    )
+    .unwrap();
+    fs::write(root.path().join("bar.py"), "target = 1\n").unwrap();
+
+    let mut interaction = LspInteraction::new();
+    interaction.set_root(root.path().to_path_buf());
+    interaction
+        .initialize(InitializeSettings {
+            ..Default::default()
+        })
+        .unwrap();
+    interaction.client.did_open("foo.py");
+
+    interaction
+        .client
+        .definition("foo.py", 2, 10)
+        .expect_definition_response_from_root("bar.py", 0, 0, 0, 6)
+        .unwrap();
+
+    fs::write(
+        root.path().join("foo.py"),
+        "# rebased\nfrom bar import target\n\nvalue = target\n",
+    )
+    .unwrap();
+    fs::write(root.path().join("bar.py"), "# moved\ntarget = 1\n").unwrap();
+
+    interaction
+        .client
+        .definition("foo.py", 3, 10)
+        .expect_definition_response_from_root("bar.py", 1, 0, 1, 6)
+        .unwrap();
+
+    interaction.shutdown().unwrap();
+}
+
 #[test]
 fn test_go_to_def_single_root() {
     let root = get_test_files_root();