Respect server capabilities on queries (#33538)

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

Turns out a bunch of Zed requests were not checking their capabilities
correctly, due to odd copy-paste and due to default that assumed that
the capabilities are met.

Adjust the code, which includes the document colors, add the test on the
colors case.

Release Notes:

- Fixed excessive document colors requests for unrelated files
This commit is contained in:
Kirill Bulatov 2025-06-27 19:31:40 +03:00 committed by GitHub
parent f9987a1141
commit 01dfb6fa82
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 259 additions and 74 deletions

View file

@ -3545,7 +3545,8 @@ pub struct LspStore {
lsp_data: Option<LspData>,
}
type DocumentColorTask = Shared<Task<std::result::Result<Vec<DocumentColor>, Arc<anyhow::Error>>>>;
type DocumentColorTask =
Shared<Task<std::result::Result<HashSet<DocumentColor>, Arc<anyhow::Error>>>>;
#[derive(Debug)]
struct LspData {
@ -3557,7 +3558,7 @@ struct LspData {
#[derive(Debug, Default)]
struct BufferLspData {
colors: Option<Vec<DocumentColor>>,
colors: Option<HashSet<DocumentColor>>,
}
#[derive(Debug)]
@ -6237,13 +6238,13 @@ impl LspStore {
.flat_map(|lsp_data| lsp_data.buffer_lsp_data.values())
.filter_map(|buffer_data| buffer_data.get(&abs_path))
.filter_map(|buffer_data| {
let colors = buffer_data.colors.as_deref()?;
let colors = buffer_data.colors.as_ref()?;
received_colors_data = true;
Some(colors)
})
.flatten()
.cloned()
.collect::<Vec<_>>();
.collect::<HashSet<_>>();
if buffer_lsp_data.is_empty() || for_server_id.is_some() {
if received_colors_data && for_server_id.is_none() {
@ -6297,42 +6298,25 @@ impl LspStore {
let task_abs_path = abs_path.clone();
let new_task = cx
.spawn(async move |lsp_store, cx| {
cx.background_executor().timer(Duration::from_millis(50)).await;
let fetched_colors = match lsp_store
.update(cx, |lsp_store, cx| {
lsp_store.fetch_document_colors(buffer, cx)
}) {
Ok(fetch_task) => fetch_task.await
.with_context(|| {
format!(
"Fetching document colors for buffer with path {task_abs_path:?}"
)
}),
Err(e) => return Err(Arc::new(e)),
};
let fetched_colors = match fetched_colors {
Ok(fetched_colors) => fetched_colors,
Err(e) => return Err(Arc::new(e)),
};
let lsp_colors = 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 = Vec::new();
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);
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))
}
Ok(lsp_colors)
});
match lsp_colors {
Ok(Ok(lsp_colors)) => Ok(lsp_colors),
Ok(Err(e)) => Err(Arc::new(e)),
Err(e) => Err(Arc::new(e)),
}
})
.shared();
@ -6350,11 +6334,11 @@ impl LspStore {
}
}
fn fetch_document_colors(
fn fetch_document_colors_for_buffer(
&mut self,
buffer: Entity<Buffer>,
cx: &mut Context<Self>,
) -> Task<anyhow::Result<Vec<(LanguageServerId, Vec<DocumentColor>)>>> {
) -> Task<anyhow::Result<Vec<(LanguageServerId, HashSet<DocumentColor>)>>> {
if let Some((client, project_id)) = self.upstream_client() {
let request_task = client.request(proto::MultiLspQuery {
project_id,
@ -6403,7 +6387,9 @@ impl LspStore {
.await
.into_iter()
.fold(HashMap::default(), |mut acc, (server_id, colors)| {
acc.entry(server_id).or_insert_with(Vec::new).extend(colors);
acc.entry(server_id)
.or_insert_with(HashSet::default)
.extend(colors);
acc
})
.into_iter()
@ -6418,7 +6404,9 @@ impl LspStore {
.await
.into_iter()
.fold(HashMap::default(), |mut acc, (server_id, colors)| {
acc.entry(server_id).or_insert_with(Vec::new).extend(colors);
acc.entry(server_id)
.or_insert_with(HashSet::default)
.extend(colors);
acc
})
.into_iter()
@ -10691,6 +10679,53 @@ 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>,