Fix document colors issues with other inlays and multi buffers (#33598)

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

* Fixes inlay colors spoiled after document color displayed
* Optimizes the query pattern for large multi buffers

Release Notes:

- Fixed document colors issues with other inlays and multi buffers

crates/editor/src/display_map/inlay_map.rs

@@ -327,9 +327,9 @@ impl<'a> Iterator for InlayChunks<'a> {
                     InlayId::DebuggerValue(_) => self.highlight_styles.inlay_hint,
                     InlayId::Color(_) => match inlay.color {
                         Some(color) => {
-                            let style = self.highlight_styles.inlay_hint.get_or_insert_default();
+                            let mut style = self.highlight_styles.inlay_hint.unwrap_or_default();
                             style.color = Some(color);
-                            Some(*style)
+                            Some(style)
                         }
                         None => self.highlight_styles.inlay_hint,
                     },

crates/editor/src/editor.rs

@@ -1845,13 +1845,13 @@ impl Editor {
                             editor
                                 .refresh_inlay_hints(InlayHintRefreshReason::RefreshRequested, cx);
                         }
-                        project::Event::LanguageServerAdded(server_id, ..)
+                        project::Event::LanguageServerAdded(..)
-                        | project::Event::LanguageServerRemoved(server_id) => {
+                        | project::Event::LanguageServerRemoved(..) => {
                             if editor.tasks_update_task.is_none() {
                                 editor.tasks_update_task =
                                     Some(editor.refresh_runnables(window, cx));
                             }
-                            editor.update_lsp_data(Some(*server_id), None, window, cx);
+                            editor.update_lsp_data(true, None, window, cx);
                         }
                         project::Event::SnippetEdit(id, snippet_edits) => {
                             if let Some(buffer) = editor.buffer.read(cx).buffer(*id) {
@@ -2291,7 +2291,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(None, None, window, cx);
+            editor.update_lsp_data(false, None, window, cx);
         }
 
         editor.report_editor_event("Editor Opened", None, cx);
@@ -5103,7 +5103,7 @@ impl Editor {
             to_insert,
         }) = self.inlay_hint_cache.spawn_hint_refresh(
             reason_description,
-            self.excerpts_for_inlay_hints_query(required_languages.as_ref(), cx),
+            self.visible_excerpts(required_languages.as_ref(), cx),
             invalidate_cache,
             ignore_debounce,
             cx,
@@ -5121,7 +5121,7 @@ impl Editor {
             .collect()
     }
 
-    pub fn excerpts_for_inlay_hints_query(
+    pub fn visible_excerpts(
         &self,
         restrict_to_languages: Option<&HashSet<Arc<Language>>>,
         cx: &mut Context<Editor>,
@@ -19562,7 +19562,7 @@ impl Editor {
                 cx.emit(SearchEvent::MatchesInvalidated);
 
                 if let Some(buffer) = edited_buffer {
-                    self.update_lsp_data(None, Some(buffer.read(cx).remote_id()), window, cx);
+                    self.update_lsp_data(false, Some(buffer.read(cx).remote_id()), window, cx);
                 }
 
                 if *singleton_buffer_edited {
@@ -19627,7 +19627,7 @@ impl Editor {
                         .detach();
                     }
                 }
-                self.update_lsp_data(None, Some(buffer_id), window, cx);
+                self.update_lsp_data(false, Some(buffer_id), window, cx);
                 cx.emit(EditorEvent::ExcerptsAdded {
                     buffer: buffer.clone(),
                     predecessor: *predecessor,
@@ -19813,7 +19813,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(None, None, window, cx);
+            self.refresh_colors(false, None, window, cx);
         }
 
         cx.notify();
@@ -20714,13 +20714,13 @@ impl Editor {
 
     fn update_lsp_data(
         &mut self,
-        for_server_id: Option<LanguageServerId>,
+        ignore_cache: bool,
         for_buffer: Option<BufferId>,
         window: &mut Window,
         cx: &mut Context<'_, Self>,
     ) {
         self.pull_diagnostics(for_buffer, window, cx);
-        self.refresh_colors(for_server_id, for_buffer, window, cx);
+        self.refresh_colors(ignore_cache, for_buffer, window, cx);
     }
 }
 

crates/editor/src/editor_tests.rs

@@ -55,7 +55,8 @@ use util::{
     uri,
 };
 use workspace::{
-    CloseActiveItem, CloseAllItems, CloseInactiveItems, NavigationEntry, OpenOptions, ViewId,
+    CloseActiveItem, CloseAllItems, CloseInactiveItems, MoveItemToPaneInDirection, NavigationEntry,
+    OpenOptions, ViewId,
     item::{FollowEvent, FollowableItem, Item, ItemHandle, SaveOptions},
 };
 
@@ -22601,8 +22602,8 @@ async fn test_add_selection_after_moving_with_multiple_cursors(cx: &mut TestAppC
     );
 }
 
-#[gpui::test]
+#[gpui::test(iterations = 10)]
-async fn test_mtime_and_document_colors(cx: &mut TestAppContext) {
+async fn test_document_colors(cx: &mut TestAppContext) {
     let expected_color = Rgba {
         r: 0.33,
         g: 0.33,
@@ -22723,24 +22724,73 @@ async fn test_mtime_and_document_colors(cx: &mut TestAppContext) {
         .set_request_handler::<lsp::request::DocumentColor, _, _>(move |_, _| async move {
             panic!("Should not be called");
         });
-    color_request_handle.next().await.unwrap();
+    cx.executor().advance_clock(Duration::from_millis(100));
-    cx.run_until_parked();
     color_request_handle.next().await.unwrap();
     cx.run_until_parked();
     assert_eq!(
-        3,
+        1,
         requests_made.load(atomic::Ordering::Acquire),
-        "Should query for colors once per editor open (1) and once after the language server startup (2)"
+        "Should query for colors once per editor open"
     );
-
+    editor.update_in(cx, |editor, _, cx| {
-    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"
         );
+    });
 
+    // opening another file in a split should not influence the LSP query counter
+    workspace
+        .update(cx, |workspace, window, cx| {
+            assert_eq!(
+                workspace.panes().len(),
+                1,
+                "Should have one pane with one editor"
+            );
+            workspace.move_item_to_pane_in_direction(
+                &MoveItemToPaneInDirection {
+                    direction: SplitDirection::Right,
+                    focus: false,
+                    clone: true,
+                },
+                window,
+                cx,
+            );
+        })
+        .unwrap();
+    cx.run_until_parked();
+    workspace
+        .update(cx, |workspace, _, cx| {
+            let panes = workspace.panes();
+            assert_eq!(panes.len(), 2, "Should have two panes after splitting");
+            for pane in panes {
+                let editor = pane
+                    .read(cx)
+                    .active_item()
+                    .and_then(|item| item.downcast::<Editor>())
+                    .expect("Should have opened an editor in each split");
+                let editor_file = editor
+                    .read(cx)
+                    .buffer()
+                    .read(cx)
+                    .as_singleton()
+                    .expect("test deals with singleton buffers")
+                    .read(cx)
+                    .file()
+                    .expect("test buffese should have a file")
+                    .path();
+                assert_eq!(
+                    editor_file.as_ref(),
+                    Path::new("first.rs"),
+                    "Both editors should be opened for the same file"
+                )
+            }
+        })
+        .unwrap();
+
+    cx.executor().advance_clock(Duration::from_millis(500));
+    let save = editor.update_in(cx, |editor, window, cx| {
         editor.move_to_end(&MoveToEnd, window, cx);
         editor.handle_input("dirty", window, cx);
         editor.save(
@@ -22755,12 +22805,10 @@ async fn test_mtime_and_document_colors(cx: &mut TestAppContext) {
     });
     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!(
-        5,
+        3,
         requests_made.load(atomic::Ordering::Acquire),
         "Should query for colors once per save and once per formatting after save"
     );
@@ -22774,11 +22822,27 @@ async fn test_mtime_and_document_colors(cx: &mut TestAppContext) {
         })
         .unwrap();
     close.await.unwrap();
+    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!(
-        5,
+        3,
         requests_made.load(atomic::Ordering::Acquire),
-        "After saving and closing the editor, no extra requests should be made"
+        "After saving and closing all editors, no extra requests should be made"
     );
+    workspace
+        .update(cx, |workspace, _, cx| {
+            assert!(
+                workspace.active_item(cx).is_none(),
+                "Should close all editors"
+            )
+        })
+        .unwrap();
 
     workspace
         .update(cx, |workspace, window, cx| {
@@ -22788,13 +22852,7 @@ async fn test_mtime_and_document_colors(cx: &mut TestAppContext) {
         })
         .unwrap();
     cx.executor().advance_clock(Duration::from_millis(100));
-    color_request_handle.next().await.unwrap();
     cx.run_until_parked();
-    assert_eq!(
-        6,
-        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
@@ -22804,6 +22862,12 @@ async fn test_mtime_and_document_colors(cx: &mut TestAppContext) {
                 .expect("Should be an editor")
         })
         .unwrap();
+    color_request_handle.next().await.unwrap();
+    assert_eq!(
+        3,
+        requests_made.load(atomic::Ordering::Acquire),
+        "Cache should be reused on buffer close and reopen"
+    );
     editor.update(cx, |editor, cx| {
         assert_eq!(
             vec![expected_color],

crates/editor/src/inlay_hint_cache.rs

@@ -956,7 +956,7 @@ fn fetch_and_update_hints(
             .update(cx, |editor, cx| {
                 if got_throttled {
                     let query_not_around_visible_range = match editor
-                        .excerpts_for_inlay_hints_query(None, cx)
+                        .visible_excerpts(None, cx)
                         .remove(&query.excerpt_id)
                     {
                         Some((_, _, current_visible_range)) => {
@@ -2525,9 +2525,7 @@ pub mod tests {
         cx: &mut gpui::TestAppContext,
     ) -> Range<Point> {
         let ranges = editor
-            .update(cx, |editor, _window, cx| {
+            .update(cx, |editor, _window, cx| editor.visible_excerpts(None, cx))
-                editor.excerpts_for_inlay_hints_query(None, cx)
-            })
             .unwrap();
         assert_eq!(
             ranges.len(),

crates/editor/src/lsp_colors.rs

@@ -3,10 +3,10 @@ use std::{cmp, ops::Range};
 use collections::HashMap;
 use futures::future::join_all;
 use gpui::{Hsla, Rgba};
+use itertools::Itertools;
 use language::point_from_lsp;
-use lsp::LanguageServerId;
 use multi_buffer::Anchor;
-use project::DocumentColor;
+use project::{DocumentColor, lsp_store::ColorFetchStrategy};
 use settings::Settings as _;
 use text::{Bias, BufferId, OffsetRangeExt as _};
 use ui::{App, Context, Window};
@@ -19,6 +19,7 @@ use crate::{
 
 #[derive(Debug)]
 pub(super) struct LspColorData {
+    cache_version_used: usize,
     colors: Vec<(Range<Anchor>, DocumentColor, InlayId)>,
     inlay_colors: HashMap<InlayId, usize>,
     render_mode: DocumentColorsRenderMode,
@@ -27,6 +28,7 @@ pub(super) struct LspColorData {
 impl LspColorData {
     pub fn new(cx: &App) -> Self {
         Self {
+            cache_version_used: 0,
             colors: Vec::new(),
             inlay_colors: HashMap::default(),
             render_mode: EditorSettings::get_global(cx).lsp_document_colors,
@@ -122,7 +124,7 @@ impl LspColorData {
 impl Editor {
     pub(super) fn refresh_colors(
         &mut self,
-        for_server_id: Option<LanguageServerId>,
+        ignore_cache: bool,
         buffer_id: Option<BufferId>,
         _: &Window,
         cx: &mut Context<Self>,
@@ -141,29 +143,41 @@ impl Editor {
             return;
         }
 
+        let visible_buffers = self
+            .visible_excerpts(None, cx)
+            .into_values()
+            .map(|(buffer, ..)| buffer)
+            .filter(|editor_buffer| {
+                buffer_id.is_none_or(|buffer_id| buffer_id == editor_buffer.read(cx).remote_id())
+            })
+            .unique_by(|buffer| buffer.read(cx).remote_id())
+            .collect::<Vec<_>>();
+
         let all_colors_task = project.read(cx).lsp_store().update(cx, |lsp_store, cx| {
-            self.buffer()
+            visible_buffers
-                .update(cx, |multi_buffer, cx| {
-                    multi_buffer
-                        .all_buffers()
-                        .into_iter()
-                        .filter(|editor_buffer| {
-                            buffer_id.is_none_or(|buffer_id| {
-                                buffer_id == editor_buffer.read(cx).remote_id()
-                            })
-                        })
-                        .collect::<Vec<_>>()
-                })
                 .into_iter()
                 .filter_map(|buffer| {
                     let buffer_id = buffer.read(cx).remote_id();
-                    let colors_task = lsp_store.document_colors(for_server_id, buffer, cx)?;
+                    let fetch_strategy = if ignore_cache {
+                        ColorFetchStrategy::IgnoreCache
+                    } else {
+                        ColorFetchStrategy::UseCache {
+                            known_cache_version: self
+                                .colors
+                                .as_ref()
+                                .map(|colors| colors.cache_version_used),
+                        }
+                    };
+                    let colors_task = lsp_store.document_colors(fetch_strategy, buffer, cx)?;
                     Some(async move { (buffer_id, colors_task.await) })
                 })
                 .collect::<Vec<_>>()
         });
         cx.spawn(async move |editor, cx| {
             let all_colors = join_all(all_colors_task).await;
+            if all_colors.is_empty() {
+                return;
+            }
             let Ok((multi_buffer_snapshot, editor_excerpts)) = editor.update(cx, |editor, cx| {
                 let multi_buffer_snapshot = editor.buffer().read(cx).snapshot(cx);
                 let editor_excerpts = multi_buffer_snapshot.excerpts().fold(
@@ -187,6 +201,7 @@ impl Editor {
                 return;
             };
 
+            let mut cache_version = None;
             let mut new_editor_colors = Vec::<(Range<Anchor>, DocumentColor)>::new();
             for (buffer_id, colors) in all_colors {
                 let Some(excerpts) = editor_excerpts.get(&buffer_id) else {
@@ -194,7 +209,8 @@ impl Editor {
                 };
                 match colors {
                     Ok(colors) => {
-                        for color in colors {
+                        cache_version = colors.cache_version;
+                        for color in colors.colors {
                             let color_start = point_from_lsp(color.lsp_range.start);
                             let color_end = point_from_lsp(color.lsp_range.end);
 
@@ -337,6 +353,9 @@ impl Editor {
                     }
 
                     let mut updated = colors.set_colors(new_color_inlays);
+                    if let Some(cache_version) = cache_version {
+                        colors.cache_version_used = cache_version;
+                    }
                     if colors.render_mode == DocumentColorsRenderMode::Inlay
                         && (!colors_splice.to_insert.is_empty()
                             || !colors_splice.to_remove.is_empty())

crates/editor/src/scroll.rs

@@ -487,8 +487,9 @@ impl Editor {
         if opened_first_time {
             cx.spawn_in(window, async move |editor, cx| {
                 editor
-                    .update(cx, |editor, cx| {
+                    .update_in(cx, |editor, window, cx| {
-                        editor.refresh_inlay_hints(InlayHintRefreshReason::NewLinesShown, cx)
+                        editor.refresh_inlay_hints(InlayHintRefreshReason::NewLinesShown, cx);
+                        editor.refresh_colors(false, None, window, cx);
                     })
                     .ok()
             })
@@ -599,6 +600,7 @@ impl Editor {
         );
 
         self.refresh_inlay_hints(InlayHintRefreshReason::NewLinesShown, cx);
+        self.refresh_colors(false, None, window, cx);
     }
 
     pub fn scroll_position(&self, cx: &mut Context<Self>) -> gpui::Point<f32> {

crates/project/src/lsp_store.rs

@@ -3542,23 +3542,29 @@ pub struct LspStore {
     _maintain_buffer_languages: Task<()>,
     diagnostic_summaries:
         HashMap<WorktreeId, HashMap<Arc<Path>, HashMap<LanguageServerId, DiagnosticSummary>>>,
-    lsp_data: Option<LspData>,
+    lsp_data: HashMap<BufferId, DocumentColorData>,
 }
 
-type DocumentColorTask =
+#[derive(Debug, Default, Clone)]
-    Shared<Task<std::result::Result<HashSet<DocumentColor>, Arc<anyhow::Error>>>>;
+pub struct DocumentColors {
-
+    pub colors: HashSet<DocumentColor>,
-#[derive(Debug)]
+    pub cache_version: Option<usize>,
-struct LspData {
-    mtime: MTime,
-    buffer_lsp_data: HashMap<LanguageServerId, HashMap<PathBuf, BufferLspData>>,
-    colors_update: HashMap<PathBuf, DocumentColorTask>,
-    last_version_queried: HashMap<PathBuf, Global>,
 }
 
+type DocumentColorTask = Shared<Task<std::result::Result<DocumentColors, Arc<anyhow::Error>>>>;
+
 #[derive(Debug, Default)]
-struct BufferLspData {
+struct DocumentColorData {
-    colors: Option<HashSet<DocumentColor>>,
+    colors_for_version: Global,
+    colors: HashMap<LanguageServerId, HashSet<DocumentColor>>,
+    cache_version: usize,
+    colors_update: Option<(Global, DocumentColorTask)>,
+}
+
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
+pub enum ColorFetchStrategy {
+    IgnoreCache,
+    UseCache { known_cache_version: Option<usize> },
 }
 
 #[derive(Debug)]
@@ -3792,7 +3798,7 @@ impl LspStore {
             language_server_statuses: Default::default(),
             nonce: StdRng::from_entropy().r#gen(),
             diagnostic_summaries: HashMap::default(),
-            lsp_data: None,
+            lsp_data: HashMap::default(),
             active_entry: None,
             _maintain_workspace_config,
             _maintain_buffer_languages: Self::maintain_buffer_languages(languages, cx),
@@ -3849,7 +3855,7 @@ impl LspStore {
             language_server_statuses: Default::default(),
             nonce: StdRng::from_entropy().r#gen(),
             diagnostic_summaries: HashMap::default(),
-            lsp_data: None,
+            lsp_data: HashMap::default(),
             active_entry: None,
             toolchain_store,
             _maintain_workspace_config,
@@ -4138,15 +4144,20 @@ impl LspStore {
                 local.register_buffer_with_language_servers(buffer, only_register_servers, cx);
             }
             if !ignore_refcounts {
-                cx.observe_release(&handle, move |this, buffer, cx| {
+                cx.observe_release(&handle, move |lsp_store, buffer, cx| {
-                    let local = this.as_local_mut().unwrap();
+                    let refcount = {
-                    let Some(refcount) = local.registered_buffers.get_mut(&buffer_id) else {
+                        let local = lsp_store.as_local_mut().unwrap();
-                        debug_panic!("bad refcounting");
+                        let Some(refcount) = local.registered_buffers.get_mut(&buffer_id) else {
-                        return;
+                            debug_panic!("bad refcounting");
-                    };
+                            return;
+                        };
 
-                    *refcount -= 1;
+                        *refcount -= 1;
-                    if *refcount == 0 {
+                        *refcount
+                    };
+                    if refcount == 0 {
+                        lsp_store.lsp_data.remove(&buffer_id);
+                        let local = lsp_store.as_local_mut().unwrap();
                         local.registered_buffers.remove(&buffer_id);
                         if let Some(file) = File::from_dyn(buffer.read(cx).file()).cloned() {
                             local.unregister_old_buffer_from_language_servers(&buffer, &file, cx);
@@ -5012,7 +5023,7 @@ impl LspStore {
                     .presentations
                     .into_iter()
                     .map(|presentation| ColorPresentation {
-                        label: presentation.label,
+                        label: SharedString::from(presentation.label),
                         text_edit: presentation.text_edit.and_then(deserialize_lsp_edit),
                         additional_text_edits: presentation
                             .additional_text_edits
@@ -5055,7 +5066,7 @@ impl LspStore {
                     .context("color presentation resolve LSP request")?
                     .into_iter()
                     .map(|presentation| ColorPresentation {
-                        label: presentation.label,
+                        label: SharedString::from(presentation.label),
                         text_edit: presentation.text_edit,
                         additional_text_edits: presentation
                             .additional_text_edits
@@ -6210,135 +6221,127 @@ impl LspStore {
 
     pub fn document_colors(
         &mut self,
-        for_server_id: Option<LanguageServerId>,
+        fetch_strategy: ColorFetchStrategy,
         buffer: Entity<Buffer>,
         cx: &mut Context<Self>,
     ) -> Option<DocumentColorTask> {
-        let buffer_mtime = buffer.read(cx).saved_mtime()?;
+        let version_queried_for = buffer.read(cx).version();
-        let buffer_version = buffer.read(cx).version();
+        let buffer_id = buffer.read(cx).remote_id();
-        let abs_path = File::from_dyn(buffer.read(cx).file())?.abs_path(cx);
 
-        let mut received_colors_data = false;
+        match fetch_strategy {
-        let buffer_lsp_data = self
+            ColorFetchStrategy::IgnoreCache => {}
-            .lsp_data
+            ColorFetchStrategy::UseCache {
-            .as_ref()
+                known_cache_version,
-            .into_iter()
+            } => {
-            .filter(|lsp_data| {
+                if let Some(cached_data) = self.lsp_data.get(&buffer_id) {
-                if buffer_mtime == lsp_data.mtime {
+                    if !version_queried_for.changed_since(&cached_data.colors_for_version) {
-                    lsp_data
+                        if Some(cached_data.cache_version) == known_cache_version {
-                        .last_version_queried
+                            return None;
-                        .get(&abs_path)
+                        } else {
-                        .is_none_or(|version_queried| {
+                            return Some(
-                            !buffer_version.changed_since(version_queried)
+                                Task::ready(Ok(DocumentColors {
-                        })
+                                    colors: cached_data
-                } else {
+                                        .colors
-                    !buffer_mtime.bad_is_greater_than(lsp_data.mtime)
+                                        .values()
-                }
+                                        .flatten()
-            })
+                                        .cloned()
-            .flat_map(|lsp_data| lsp_data.buffer_lsp_data.values())
+                                        .collect(),
-            .filter_map(|buffer_data| buffer_data.get(&abs_path))
+                                    cache_version: Some(cached_data.cache_version),
-            .filter_map(|buffer_data| {
+                                }))
-                let colors = buffer_data.colors.as_ref()?;
+                                .shared(),
-                received_colors_data = true;
+                            );
-                Some(colors)
-            })
-            .flatten()
-            .cloned()
-            .collect::<HashSet<_>>();
-
-        if buffer_lsp_data.is_empty() || for_server_id.is_some() {
-            if received_colors_data && for_server_id.is_none() {
-                return None;
-            }
-
-            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,
-                    buffer_lsp_data: HashMap::default(),
-                    colors_update: HashMap::default(),
-                    last_version_queried: HashMap::default(),
-                });
-                outdated_lsp_data = true;
-            }
-
-            {
-                let lsp_data = self.lsp_data.as_mut()?;
-                match for_server_id {
-                    Some(for_server_id) if !outdated_lsp_data => {
-                        lsp_data.buffer_lsp_data.remove(&for_server_id);
-                    }
-                    None | Some(_) => {
-                        let existing_task = lsp_data.colors_update.get(&abs_path).cloned();
-                        if !outdated_lsp_data && existing_task.is_some() {
-                            return existing_task;
-                        }
-                        for buffer_data in lsp_data.buffer_lsp_data.values_mut() {
-                            if let Some(buffer_data) = buffer_data.get_mut(&abs_path) {
-                                buffer_data.colors = None;
-                            }
                         }
                     }
                 }
             }
-
-            let task_abs_path = abs_path.clone();
-            let new_task = cx
-                .spawn(async move |lsp_store, cx| {
-                    match fetch_document_colors(
-                        lsp_store.clone(),
-                        buffer,
-                        task_abs_path.clone(),
-                        cx,
-                    )
-                    .await
-                    {
-                        Ok(colors) => Ok(colors),
-                        Err(e) => {
-                            lsp_store
-                                .update(cx, |lsp_store, _| {
-                                    if let Some(lsp_data) = lsp_store.lsp_data.as_mut() {
-                                        lsp_data.colors_update.remove(&task_abs_path);
-                                    }
-                                })
-                                .ok();
-                            Err(Arc::new(e))
-                        }
-                    }
-                })
-                .shared();
-            let lsp_data = self.lsp_data.as_mut()?;
-            lsp_data
-                .colors_update
-                .insert(abs_path.clone(), new_task.clone());
-            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())
         }
+
+        let lsp_data = self.lsp_data.entry(buffer_id).or_default();
+        if let Some((updating_for, running_update)) = &lsp_data.colors_update {
+            if !version_queried_for.changed_since(&updating_for) {
+                return Some(running_update.clone());
+            }
+        }
+        let query_version_queried_for = version_queried_for.clone();
+        let new_task = cx
+            .spawn(async move |lsp_store, cx| {
+                cx.background_executor()
+                    .timer(Duration::from_millis(30))
+                    .await;
+                let fetched_colors = lsp_store
+                    .update(cx, |lsp_store, cx| {
+                        lsp_store.fetch_document_colors_for_buffer(buffer.clone(), cx)
+                    })?
+                    .await
+                    .context("fetching document colors")
+                    .map_err(Arc::new);
+                let fetched_colors = match fetched_colors {
+                    Ok(fetched_colors) => {
+                        if fetch_strategy != ColorFetchStrategy::IgnoreCache
+                            && Some(true)
+                                == buffer
+                                    .update(cx, |buffer, _| {
+                                        buffer.version() != query_version_queried_for
+                                    })
+                                    .ok()
+                        {
+                            return Ok(DocumentColors::default());
+                        }
+                        fetched_colors
+                    }
+                    Err(e) => {
+                        lsp_store
+                            .update(cx, |lsp_store, _| {
+                                lsp_store
+                                    .lsp_data
+                                    .entry(buffer_id)
+                                    .or_default()
+                                    .colors_update = None;
+                            })
+                            .ok();
+                        return Err(e);
+                    }
+                };
+
+                lsp_store
+                    .update(cx, |lsp_store, _| {
+                        let lsp_data = lsp_store.lsp_data.entry(buffer_id).or_default();
+
+                        if lsp_data.colors_for_version == query_version_queried_for {
+                            lsp_data.colors.extend(fetched_colors.clone());
+                            lsp_data.cache_version += 1;
+                        } else if !lsp_data
+                            .colors_for_version
+                            .changed_since(&query_version_queried_for)
+                        {
+                            lsp_data.colors_for_version = query_version_queried_for;
+                            lsp_data.colors = fetched_colors.clone();
+                            lsp_data.cache_version += 1;
+                        }
+                        lsp_data.colors_update = None;
+                        let colors = lsp_data
+                            .colors
+                            .values()
+                            .flatten()
+                            .cloned()
+                            .collect::<HashSet<_>>();
+                        DocumentColors {
+                            colors,
+                            cache_version: Some(lsp_data.cache_version),
+                        }
+                    })
+                    .map_err(Arc::new)
+            })
+            .shared();
+        lsp_data.colors_update = Some((version_queried_for, new_task.clone()));
+        Some(new_task)
     }
 
     fn fetch_document_colors_for_buffer(
         &mut self,
         buffer: Entity<Buffer>,
         cx: &mut Context<Self>,
-    ) -> Task<anyhow::Result<Vec<(LanguageServerId, HashSet<DocumentColor>)>>> {
+    ) -> Task<anyhow::Result<HashMap<LanguageServerId, HashSet<DocumentColor>>>> {
         if let Some((client, project_id)) = self.upstream_client() {
             let request_task = client.request(proto::MultiLspQuery {
                 project_id,
@@ -6353,7 +6356,7 @@ impl LspStore {
             });
             cx.spawn(async move |project, cx| {
                 let Some(project) = project.upgrade() else {
-                    return Ok(Vec::new());
+                    return Ok(HashMap::default());
                 };
                 let colors = join_all(
                     request_task
@@ -6391,9 +6394,7 @@ impl LspStore {
                         .or_insert_with(HashSet::default)
                         .extend(colors);
                     acc
-                })
+                });
-                .into_iter()
-                .collect();
                 Ok(colors)
             })
         } else {
@@ -8942,7 +8943,7 @@ impl LspStore {
                 .color_presentations
                 .into_iter()
                 .map(|presentation| proto::ColorPresentation {
-                    label: presentation.label,
+                    label: presentation.label.to_string(),
                     text_edit: presentation.text_edit.map(serialize_lsp_edit),
                     additional_text_edits: presentation
                         .additional_text_edits
@@ -10605,8 +10606,9 @@ impl LspStore {
     }
 
     fn cleanup_lsp_data(&mut self, for_server: LanguageServerId) {
-        if let Some(lsp_data) = &mut self.lsp_data {
+        for buffer_lsp_data in self.lsp_data.values_mut() {
-            lsp_data.buffer_lsp_data.remove(&for_server);
+            buffer_lsp_data.colors.remove(&for_server);
+            buffer_lsp_data.cache_version += 1;
         }
         if let Some(local) = self.as_local_mut() {
             local.buffer_pull_diagnostics_result_ids.remove(&for_server);
@@ -10679,53 +10681,6 @@ impl LspStore {
     }
 }
 
-async fn fetch_document_colors(
-    lsp_store: WeakEntity<LspStore>,
-    buffer: Entity<Buffer>,
-    task_abs_path: PathBuf,
-    cx: &mut AsyncApp,
-) -> anyhow::Result<HashSet<DocumentColor>> {
-    cx.background_executor()
-        .timer(Duration::from_millis(50))
-        .await;
-    let Some(buffer_mtime) = buffer.update(cx, |buffer, _| buffer.saved_mtime())? else {
-        return Ok(HashSet::default());
-    };
-    let fetched_colors = lsp_store
-        .update(cx, |lsp_store, cx| {
-            lsp_store.fetch_document_colors_for_buffer(buffer, cx)
-        })?
-        .await
-        .with_context(|| {
-            format!("Fetching document colors for buffer with path {task_abs_path:?}")
-        })?;
-
-    lsp_store.update(cx, |lsp_store, _| {
-        let lsp_data = lsp_store.lsp_data.as_mut().with_context(|| {
-            format!(
-                "Document lsp data got updated between fetch and update for path {task_abs_path:?}"
-            )
-        })?;
-        let mut lsp_colors = HashSet::default();
-        anyhow::ensure!(
-            lsp_data.mtime == buffer_mtime,
-            "Buffer lsp data got updated between fetch and update for path {task_abs_path:?}"
-        );
-        for (server_id, colors) in fetched_colors {
-            let colors_lsp_data = &mut lsp_data
-                .buffer_lsp_data
-                .entry(server_id)
-                .or_default()
-                .entry(task_abs_path.clone())
-                .or_default()
-                .colors;
-            *colors_lsp_data = Some(colors.clone());
-            lsp_colors.extend(colors);
-        }
-        Ok(lsp_colors)
-    })?
-}
-
 fn subscribe_to_binary_statuses(
     languages: &Arc<LanguageRegistry>,
     cx: &mut Context<'_, LspStore>,

crates/project/src/project.rs

@@ -795,7 +795,7 @@ impl std::hash::Hash for DocumentColor {
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct ColorPresentation {
-    pub label: String,
+    pub label: SharedString,
     pub text_edit: Option<lsp::TextEdit>,
     pub additional_text_edits: Vec<lsp::TextEdit>,
 }