Fix document colors not showing on file reopen (#33009)

Closes https://github.com/zed-industries/zed/issues/32989

Release Notes:

- Fixed document colors not showing on file reopen

crates/collab/src/tests/editor_tests.rs

@@ -2028,7 +2028,6 @@ async fn test_lsp_document_color(cx_a: &mut TestAppContext, cx_b: &mut TestAppCo
         .unwrap();
 
     let (workspace_a, cx_a) = client_a.build_workspace(&project_a, cx_a);
-    executor.start_waiting();
 
     // The host opens a rust file.
     let _buffer_a = project_a

crates/editor/src/editor.rs

@@ -1802,7 +1802,7 @@ impl Editor {
                                 editor.tasks_update_task =
                                     Some(editor.refresh_runnables(window, cx));
                             }
-                            editor.update_lsp_data(false, Some(*server_id), None, window, cx);
+                            editor.update_lsp_data(Some(*server_id), None, window, cx);
                         }
                         project::Event::SnippetEdit(id, snippet_edits) => {
                             if let Some(buffer) = editor.buffer.read(cx).buffer(*id) {
@@ -2240,7 +2240,7 @@ impl Editor {
             editor.minimap =
                 editor.create_minimap(EditorSettings::get_global(cx).minimap, window, cx);
             editor.colors = Some(LspColorData::new(cx));
-            editor.update_lsp_data(false, None, None, window, cx);
+            editor.update_lsp_data(None, None, window, cx);
         }
 
         editor.report_editor_event("Editor Opened", None, cx);
@@ -19194,7 +19194,7 @@ impl Editor {
                 cx.emit(SearchEvent::MatchesInvalidated);
 
                 if let Some(buffer) = edited_buffer {
-                    self.update_lsp_data(true, None, Some(buffer.read(cx).remote_id()), window, cx);
+                    self.update_lsp_data(None, Some(buffer.read(cx).remote_id()), window, cx);
                 }
 
                 if *singleton_buffer_edited {
@@ -19259,7 +19259,7 @@ impl Editor {
                         .detach();
                     }
                 }
-                self.update_lsp_data(false, None, Some(buffer_id), window, cx);
+                self.update_lsp_data(None, Some(buffer_id), window, cx);
                 cx.emit(EditorEvent::ExcerptsAdded {
                     buffer: buffer.clone(),
                     predecessor: *predecessor,
@@ -19442,7 +19442,7 @@ impl Editor {
             if !inlay_splice.to_insert.is_empty() || !inlay_splice.to_remove.is_empty() {
                 self.splice_inlays(&inlay_splice.to_remove, inlay_splice.to_insert, cx);
             }
-            self.refresh_colors(true, None, None, window, cx);
+            self.refresh_colors(None, None, window, cx);
         }
 
         cx.notify();
@@ -20331,14 +20331,13 @@ impl Editor {
 
     fn update_lsp_data(
         &mut self,
-        update_on_edit: bool,
         for_server_id: Option<LanguageServerId>,
         for_buffer: Option<BufferId>,
         window: &mut Window,
         cx: &mut Context<'_, Self>,
     ) {
         self.pull_diagnostics(for_buffer, window, cx);
-        self.refresh_colors(update_on_edit, for_server_id, for_buffer, window, cx);
+        self.refresh_colors(for_server_id, for_buffer, window, cx);
     }
 }
 

crates/editor/src/editor_tests.rs

@@ -15,7 +15,7 @@ use crate::{
 use buffer_diff::{BufferDiff, DiffHunkSecondaryStatus, DiffHunkStatus, DiffHunkStatusKind};
 use futures::StreamExt;
 use gpui::{
-    BackgroundExecutor, DismissEvent, SemanticVersion, TestAppContext, UpdateGlobal,
+    BackgroundExecutor, DismissEvent, Rgba, SemanticVersion, TestAppContext, UpdateGlobal,
     VisualTestContext, WindowBounds, WindowOptions, div,
 };
 use indoc::indoc;
@@ -22269,3 +22269,185 @@ async fn test_add_selection_after_moving_with_multiple_cursors(cx: &mut TestAppC
         "Should have 4 cursors after moving and adding another"
     );
 }
+
+#[gpui::test]
+async fn test_mtime_and_document_colors(cx: &mut TestAppContext) {
+    let expected_color = Rgba {
+        r: 0.33,
+        g: 0.33,
+        b: 0.33,
+        a: 0.33,
+    };
+
+    init_test(cx, |_| {});
+
+    let fs = FakeFs::new(cx.executor());
+    fs.insert_tree(
+        path!("/a"),
+        json!({
+            "first.rs": "fn main() { let a = 5; }",
+        }),
+    )
+    .await;
+
+    let project = Project::test(fs, [path!("/a").as_ref()], cx).await;
+    let workspace = cx.add_window(|window, cx| Workspace::test_new(project.clone(), window, cx));
+    let cx = &mut VisualTestContext::from_window(*workspace, cx);
+
+    let language_registry = project.read_with(cx, |project, _| project.languages().clone());
+    language_registry.add(rust_lang());
+    let mut fake_servers = language_registry.register_fake_lsp(
+        "Rust",
+        FakeLspAdapter {
+            capabilities: lsp::ServerCapabilities {
+                color_provider: Some(lsp::ColorProviderCapability::Simple(true)),
+                ..lsp::ServerCapabilities::default()
+            },
+            ..FakeLspAdapter::default()
+        },
+    );
+
+    let editor = workspace
+        .update(cx, |workspace, window, cx| {
+            workspace.open_abs_path(
+                PathBuf::from(path!("/a/first.rs")),
+                OpenOptions::default(),
+                window,
+                cx,
+            )
+        })
+        .unwrap()
+        .await
+        .unwrap()
+        .downcast::<Editor>()
+        .unwrap();
+    let fake_language_server = fake_servers.next().await.unwrap();
+    let requests_made = Arc::new(AtomicUsize::new(0));
+    let closure_requests_made = Arc::clone(&requests_made);
+    let mut color_request_handle = fake_language_server
+        .set_request_handler::<lsp::request::DocumentColor, _, _>(move |params, _| {
+            let requests_made = Arc::clone(&closure_requests_made);
+            async move {
+                assert_eq!(
+                    params.text_document.uri,
+                    lsp::Url::from_file_path(path!("/a/first.rs")).unwrap()
+                );
+                requests_made.fetch_add(1, atomic::Ordering::Release);
+                Ok(vec![lsp::ColorInformation {
+                    range: lsp::Range {
+                        start: lsp::Position {
+                            line: 0,
+                            character: 0,
+                        },
+                        end: lsp::Position {
+                            line: 0,
+                            character: 1,
+                        },
+                    },
+                    color: lsp::Color {
+                        red: 0.33,
+                        green: 0.33,
+                        blue: 0.33,
+                        alpha: 0.33,
+                    },
+                }])
+            }
+        });
+    color_request_handle.next().await.unwrap();
+    cx.run_until_parked();
+    color_request_handle.next().await.unwrap();
+    cx.run_until_parked();
+    assert_eq!(
+        2,
+        requests_made.load(atomic::Ordering::Acquire),
+        "Should query for colors once per editor open and once after the language server startup"
+    );
+
+    cx.executor().advance_clock(Duration::from_millis(500));
+    let save = editor.update_in(cx, |editor, window, cx| {
+        assert_eq!(
+            vec![expected_color],
+            extract_color_inlays(editor, cx),
+            "Should have an initial inlay"
+        );
+
+        editor.move_to_end(&MoveToEnd, window, cx);
+        editor.handle_input("dirty", window, cx);
+        editor.save(
+            SaveOptions {
+                format: true,
+                autosave: true,
+            },
+            project.clone(),
+            window,
+            cx,
+        )
+    });
+    save.await.unwrap();
+
+    color_request_handle.next().await.unwrap();
+    cx.run_until_parked();
+    color_request_handle.next().await.unwrap();
+    cx.run_until_parked();
+    assert_eq!(
+        4,
+        requests_made.load(atomic::Ordering::Acquire),
+        "Should query for colors once per save and once per formatting after save"
+    );
+
+    drop(editor);
+    let close = workspace
+        .update(cx, |workspace, window, cx| {
+            workspace.active_pane().update(cx, |pane, cx| {
+                pane.close_active_item(&CloseActiveItem::default(), window, cx)
+            })
+        })
+        .unwrap();
+    close.await.unwrap();
+    assert_eq!(
+        4,
+        requests_made.load(atomic::Ordering::Acquire),
+        "After saving and closing the editor, no extra requests should be made"
+    );
+
+    workspace
+        .update(cx, |workspace, window, cx| {
+            workspace.active_pane().update(cx, |pane, cx| {
+                pane.navigate_backward(window, cx);
+            })
+        })
+        .unwrap();
+    color_request_handle.next().await.unwrap();
+    cx.run_until_parked();
+    assert_eq!(
+        5,
+        requests_made.load(atomic::Ordering::Acquire),
+        "After navigating back to an editor and reopening it, another color request should be made"
+    );
+    let editor = workspace
+        .update(cx, |workspace, _, cx| {
+            workspace
+                .active_item(cx)
+                .expect("Should have reopened the editor again after navigating back")
+                .downcast::<Editor>()
+                .expect("Should be an editor")
+        })
+        .unwrap();
+    editor.update(cx, |editor, cx| {
+        assert_eq!(
+            vec![expected_color],
+            extract_color_inlays(editor, cx),
+            "Should have an initial inlay"
+        );
+    });
+}
+
+#[track_caller]
+fn extract_color_inlays(editor: &Editor, cx: &App) -> Vec<Rgba> {
+    editor
+        .all_inlays(cx)
+        .into_iter()
+        .filter_map(|inlay| inlay.get_color())
+        .map(Rgba::from)
+        .collect()
+}

crates/editor/src/lsp_colors.rs

@@ -122,7 +122,6 @@ impl LspColorData {
 impl Editor {
     pub(super) fn refresh_colors(
         &mut self,
-        update_on_edit: bool,
         for_server_id: Option<LanguageServerId>,
         buffer_id: Option<BufferId>,
         _: &Window,
@@ -158,8 +157,7 @@ impl Editor {
                 .into_iter()
                 .filter_map(|buffer| {
                     let buffer_id = buffer.read(cx).remote_id();
-                    let colors_task =
-                        lsp_store.document_colors(update_on_edit, for_server_id, buffer, cx)?;
+                    let colors_task = lsp_store.document_colors(for_server_id, buffer, cx)?;
                     Some(async move { (buffer_id, colors_task.await) })
                 })
                 .collect::<Vec<_>>()

crates/project/src/lsp_store.rs

@@ -4169,7 +4169,7 @@ impl LspStore {
                         // and reused later in the invisible worktrees.
                         plain_text_buffers.sort_by_key(|buffer| {
                             Reverse(
-                                crate::File::from_dyn(buffer.read(cx).file())
+                                File::from_dyn(buffer.read(cx).file())
                                     .map(|file| file.worktree.read(cx).is_visible()),
                             )
                         });
@@ -4513,7 +4513,7 @@ impl LspStore {
                 let mut rebase = lsp_tree.rebase();
                 for buffer_handle in buffer_store.read(cx).buffers().sorted_by_key(|buffer| {
                     Reverse(
-                        crate::File::from_dyn(buffer.read(cx).file())
+                        File::from_dyn(buffer.read(cx).file())
                             .map(|file| file.worktree.read(cx).is_visible()),
                     )
                 }) {
@@ -4914,7 +4914,7 @@ impl LspStore {
         } else {
             let path = match buffer
                 .update(cx, |buffer, cx| {
-                    Some(crate::File::from_dyn(buffer.file())?.abs_path(cx))
+                    Some(File::from_dyn(buffer.file())?.abs_path(cx))
                 })
                 .context("buffer with the missing path")
             {
@@ -6077,32 +6077,30 @@ impl LspStore {
 
     pub fn document_colors(
         &mut self,
-        update_on_edit: bool,
         for_server_id: Option<LanguageServerId>,
         buffer: Entity<Buffer>,
         cx: &mut Context<Self>,
     ) -> Option<DocumentColorTask> {
         let buffer_mtime = buffer.read(cx).saved_mtime()?;
-        let abs_path = crate::File::from_dyn(buffer.read(cx).file())?.abs_path(cx);
         let buffer_version = buffer.read(cx).version();
-        let ignore_existing_mtime = update_on_edit
-            && self.lsp_data.as_ref().is_none_or(|lsp_data| {
-                lsp_data.last_version_queried.get(&abs_path) != Some(&buffer_version)
-            });
+        let abs_path = File::from_dyn(buffer.read(cx).file())?.abs_path(cx);
 
-        let mut has_other_versions = false;
         let mut received_colors_data = false;
-        let mut outdated_lsp_data = false;
         let buffer_lsp_data = self
             .lsp_data
             .as_ref()
             .into_iter()
             .filter(|lsp_data| {
-                if ignore_existing_mtime {
-                    return false;
+                if buffer_mtime == lsp_data.mtime {
+                    lsp_data
+                        .last_version_queried
+                        .get(&abs_path)
+                        .is_none_or(|version_queried| {
+                            !buffer_version.changed_since(version_queried)
+                        })
+                } else {
+                    !buffer_mtime.bad_is_greater_than(lsp_data.mtime)
                 }
-                has_other_versions |= lsp_data.mtime != buffer_mtime;
-                lsp_data.mtime == buffer_mtime
             })
             .flat_map(|lsp_data| lsp_data.buffer_lsp_data.values())
             .filter_map(|buffer_data| buffer_data.get(&abs_path))
@@ -6118,16 +6116,22 @@ impl LspStore {
         if buffer_lsp_data.is_empty() || for_server_id.is_some() {
             if received_colors_data && for_server_id.is_none() {
                 return None;
-            } else if has_other_versions && !ignore_existing_mtime {
-                return None;
             }
 
-            if ignore_existing_mtime
-                || self.lsp_data.is_none()
-                || self
-                    .lsp_data
-                    .as_ref()
-                    .is_some_and(|lsp_data| buffer_mtime != lsp_data.mtime)
+            let mut outdated_lsp_data = false;
+            if self.lsp_data.is_none()
+                || self.lsp_data.as_ref().is_some_and(|lsp_data| {
+                    if buffer_mtime == lsp_data.mtime {
+                        lsp_data
+                            .last_version_queried
+                            .get(&abs_path)
+                            .is_none_or(|version_queried| {
+                                buffer_version.changed_since(version_queried)
+                            })
+                    } else {
+                        buffer_mtime.bad_is_greater_than(lsp_data.mtime)
+                    }
+                })
             {
                 self.lsp_data = Some(LspData {
                     mtime: buffer_mtime,
@@ -6207,6 +6211,7 @@ impl LspStore {
             lsp_data
                 .last_version_queried
                 .insert(abs_path, buffer_version);
+            lsp_data.mtime = buffer_mtime;
             Some(new_task)
         } else {
             Some(Task::ready(Ok(buffer_lsp_data)).shared())