fix

Author: asukaminato0721

pyrefly/lib/lsp/non_wasm/server.rs

@@ -746,6 +746,7 @@ pub struct Server {
     /// should be mapped through here in case they correspond to a cell.
     open_notebook_cells: RwLock<HashMap<Url, PathBuf>>,
     open_files: RwLock<HashMap<PathBuf, Arc<LspFile>>>,
+    open_files_with_unsaved_changes: Mutex<HashSet<PathBuf>>,
     /// Last published fingerprint for unversioned file-backed workspace diagnostics.
     published_workspace_diagnostics: Mutex<HashMap<Url, u64>>,
     /// Tracks URIs (including virtual/untitled ones) to synthetic on-disk paths so we can
@@ -2416,6 +2417,7 @@ impl Server {
             state: State::new(config_finder, thread_count),
             open_notebook_cells: RwLock::new(HashMap::new()),
             open_files: RwLock::new(HashMap::new()),
+            open_files_with_unsaved_changes: Mutex::new(HashSet::new()),
             published_workspace_diagnostics: Mutex::new(HashMap::new()),
             unsaved_file_tracker: UnsavedFileTracker::new(),
             indexed_configs: Mutex::new(HashSet::new()),
@@ -3285,6 +3287,7 @@ impl Server {
 
     fn did_save(&self, url: Url) {
         if let Some(path) = self.path_for_uri(&url) {
+            self.open_files_with_unsaved_changes.lock().remove(&path);
             self.invalidate(TelemetryEventKind::InvalidateDisk, move |t| {
                 t.invalidate_disk(&[path])
             })
@@ -3322,6 +3325,7 @@ impl Server {
         } else {
             None
         };
+        self.open_files_with_unsaved_changes.lock().remove(&path);
         self.version_info.lock().insert(path.clone(), version);
         self.open_files.write().insert(path.clone(), contents);
         self.queue_source_db_rebuild_and_recheck(telemetry, telemetry_event, false);
@@ -3387,6 +3391,9 @@ impl Server {
             params.content_changes,
         )));
         drop(lock);
+        self.open_files_with_unsaved_changes
+            .lock()
+            .insert(file_path.clone());
         // Update version_info only after the mutation has fully succeeded.
         self.version_info.lock().insert(file_path.clone(), version);
         if !subsequent_mutation {
@@ -3549,6 +3556,9 @@ impl Server {
         let new_notebook = Arc::new(LspNotebook::new(ruff_notebook, notebook_document));
         *original = Arc::new(LspFile::Notebook(new_notebook));
         drop(lock);
+        self.open_files_with_unsaved_changes
+            .lock()
+            .insert(file_path.clone());
         // Update version_info only after the mutation has fully succeeded, so
         // that on error the version stays at the old value and subsequent
         // notifications operate against consistent state.
@@ -3586,6 +3596,28 @@ impl Server {
         config_changed || files_added_or_removed
     }
 
+    fn refresh_clean_open_files_from_disk(&self, events: &CategorizedEvents) {
+        let unsaved = self.open_files_with_unsaved_changes.lock().clone();
+        let mut open_files = self.open_files.write();
+        for path in events.modified.iter() {
+            if unsaved.contains(path) {
+                continue;
+            }
+            let Some(open_file) = open_files.get_mut(path) else {
+                continue;
+            };
+            let LspFile::Source(current_contents) = open_file.as_ref() else {
+                continue;
+            };
+            let Ok(updated_contents) = std::fs::read_to_string(path) else {
+                continue;
+            };
+            if current_contents.as_str() != updated_contents {
+                *open_file = Arc::new(LspFile::from_source(updated_contents));
+            }
+        }
+    }
+
     fn did_change_watched_files(
         &self,
         params: DidChangeWatchedFilesParams,
@@ -3621,6 +3653,8 @@ impl Server {
 
         let should_requery_build_system = should_requery_build_system(&events);
 
+        self.refresh_clean_open_files_from_disk(&events);
+
         // Rewatch files if necessary (config changed, files added/removed, etc.)
         if Self::should_rewatch(&events) {
             info!("[Pyrefly] Re-registering file watchers");
@@ -3710,6 +3744,7 @@ impl Server {
             },
         }
         drop(open_files);
+        self.open_files_with_unsaved_changes.lock().remove(&path);
         self.unsaved_file_tracker.forget_uri_path(&url);
         self.queue_source_db_rebuild_and_recheck(telemetry, telemetry_event, false);
         self.recheck_queue.queue_task(

pyrefly/lib/state/load.rs

@@ -81,13 +81,33 @@ pub struct Load {
 }
 
 impl Load {
+    fn load_open_filesystem_path(
+        path: &Path,
+        memory_lookup: &MemoryFilesLookup,
+    ) -> Option<Arc<FileContents>> {
+        if let Some(contents) = memory_lookup.get(path) {
+            return Some(Arc::clone(contents));
+        }
+        let canonical_path = path.canonicalize().ok()?;
+        if canonical_path == path {
+            return None;
+        }
+        memory_lookup.get(&canonical_path).map(Arc::clone)
+    }
+
     /// Return the code for this module, optional notebook cell mapping, and whether there was an error while loading (a self-error).
     pub fn load_from_path(
         path: &ModulePath,
         memory_lookup: &MemoryFilesLookup,
     ) -> (FileContents, Option<anyhow::Error>) {
         let res = match path.details() {
-            ModulePathDetails::FileSystem(path) => Self::load_from_filesystem(path),
+            ModulePathDetails::FileSystem(path) => {
+                if let Some(contents) = Self::load_open_filesystem_path(path, memory_lookup) {
+                    Ok((*contents).clone())
+                } else {
+                    Self::load_from_filesystem(path)
+                }
+            }
             ModulePathDetails::Namespace(_) => Ok(FileContents::from_source("".to_owned())),
             ModulePathDetails::Memory(path) => memory_lookup
                 .get(path)

pyrefly/lib/state/state.rs

@@ -948,13 +948,7 @@ impl<'a> Transaction<'a> {
                 .get_cached_loader(&self.get_module(handle).config.read())
                 .find_import(module, Some(handle.path())),
         };
-        path.map(|path| {
-            Handle::new(
-                module,
-                self.prefer_memory_path(path),
-                handle.sys_info().dupe(),
-            )
-        })
+        path.map(|path| Handle::new(module, path, handle.sys_info().dupe()))
     }
 
     /// Create a handle for import `module` within the handle `handle`, preferring `.py` over `.pyi`
@@ -970,13 +964,7 @@ impl<'a> Transaction<'a> {
                 .get_cached_loader(&self.get_module(handle).config.read())
                 .find_import_prefer_executable(module, Some(handle.path())),
         };
-        path.map(|path| {
-            Handle::new(
-                module,
-                self.prefer_memory_path(path),
-                handle.sys_info().dupe(),
-            )
-        })
+        path.map(|path| Handle::new(module, path, handle.sys_info().dupe()))
     }
 
     /// Create a handle for import `module` within the handle `handle`
@@ -1576,23 +1564,6 @@ impl<'a> Transaction<'a> {
         MemoryFilesLookup::new(&self.readable.memory, &self.data.memory_overlay)
     }
 
-    fn prefer_memory_path(&self, path: ModulePath) -> ModulePath {
-        let ModulePathDetails::FileSystem(filesystem_path) = path.details() else {
-            return path;
-        };
-        let memory_lookup = self.memory_lookup();
-        if memory_lookup.get(filesystem_path).is_some() {
-            return ModulePath::memory((**filesystem_path).clone());
-        }
-        if let Ok(canonical_path) = filesystem_path.canonicalize()
-            && canonical_path != **filesystem_path
-            && memory_lookup.get(&canonical_path).is_some()
-        {
-            return ModulePath::memory(canonical_path);
-        }
-        path
-    }
-
     fn get_cached_loader(&self, loader: &ArcId<ConfigFile>) -> Arc<LoaderFindCache> {
         self.data
             .updated_loaders
@@ -2069,7 +2040,13 @@ impl<'a> Transaction<'a> {
         for (path, contents) in files {
             if self.memory_lookup().get(&path) != contents.as_ref() {
                 self.data.memory_overlay.set(path.clone(), contents);
-                changed.insert(ModulePath::memory(path));
+                changed.insert(ModulePath::memory(path.clone()));
+                changed.insert(ModulePath::filesystem(path.clone()));
+                if let Ok(canonical_path) = path.canonicalize()
+                    && canonical_path != path
+                {
+                    changed.insert(ModulePath::filesystem(canonical_path));
+                }
             }
         }
         self.stats.lock().set_memory_dirty = changed.len();
@@ -2333,7 +2310,7 @@ impl<'a> TransactionHandle<'a> {
                 // Explicit path — already resolved. Bypass imports entirely.
                 FindingOrError::new_finding(Handle::new(
                     module,
-                    self.transaction.prefer_memory_path(path.dupe()),
+                    path.dupe(),
                     self.module_data.handle.sys_info().dupe(),
                 ))
             }
@@ -2358,7 +2335,7 @@ impl<'a> TransactionHandle<'a> {
                 path.map(|path| {
                     Handle::new(
                         module,
-                        self.transaction.prefer_memory_path(path.dupe()),
+                        path.dupe(),
                         self.module_data.handle.sys_info().dupe(),
                     )
                 })