1use std::ops::Range;
2use std::sync::Arc;
3
4use anyhow::{Context as _, Result};
5use clock::Global;
6use collections::HashMap;
7use futures::{
8 FutureExt as _,
9 future::{Shared, join_all},
10};
11use gpui::{AppContext as _, AsyncApp, Context, Entity, Task};
12use language::{Anchor, Buffer, ToOffset as _};
13use lsp::LanguageServerId;
14use rpc::{TypedEnvelope, proto};
15use settings::Settings as _;
16use std::time::Duration;
17use text::OffsetRangeExt as _;
18
19use crate::{
20 CodeAction, LspAction, LspStore, LspStoreEvent, Project,
21 lsp_command::{GetCodeLens, LspCommand as _},
22 project_settings::ProjectSettings,
23};
24
25pub(super) type CodeLensTask =
26 Shared<Task<std::result::Result<Option<Vec<CodeAction>>, Arc<anyhow::Error>>>>;
27
28#[derive(Debug, Default)]
29pub(super) struct CodeLensData {
30 pub(super) lens: HashMap<LanguageServerId, Vec<CodeAction>>,
31 pub(super) update: Option<(Global, CodeLensTask)>,
32}
33
34impl CodeLensData {
35 pub(super) fn remove_server_data(&mut self, server_id: LanguageServerId) {
36 self.lens.remove(&server_id);
37 }
38}
39
40impl LspStore {
41 pub(super) fn invalidate_code_lens(&mut self) {
42 for lsp_data in self.lsp_data.values_mut() {
43 lsp_data.code_lens = None;
44 }
45 }
46
47 /// Fetches and returns all code lenses for the buffer.
48 ///
49 /// Resolution of individual lenses is the caller's responsibility; see
50 /// [`LspStore::resolve_visible_code_lenses`].
51 pub fn code_lens_actions(
52 &mut self,
53 buffer: &Entity<Buffer>,
54 cx: &mut Context<Self>,
55 ) -> Task<Result<Option<Vec<CodeAction>>>> {
56 let buffer_id = buffer.read(cx).remote_id();
57 let fetch_task = self.fetch_code_lenses(buffer, cx);
58
59 cx.spawn(async move |lsp_store, cx| {
60 fetch_task
61 .await
62 .map_err(|e| anyhow::anyhow!("code lens fetch failed: {e:#}"))?;
63
64 let actions = lsp_store.read_with(cx, |lsp_store, _| {
65 lsp_store
66 .lsp_data
67 .get(&buffer_id)
68 .and_then(|data| data.code_lens.as_ref())
69 .map(|code_lens| code_lens.lens.values().flatten().cloned().collect())
70 })?;
71 Ok(actions)
72 })
73 }
74
75 fn fetch_code_lenses(
76 &mut self,
77 buffer: &Entity<Buffer>,
78 cx: &mut Context<Self>,
79 ) -> CodeLensTask {
80 let version_queried_for = buffer.read(cx).version();
81 let buffer_id = buffer.read(cx).remote_id();
82 let existing_servers = self.as_local().map(|local| {
83 local
84 .buffers_opened_in_servers
85 .get(&buffer_id)
86 .cloned()
87 .unwrap_or_default()
88 });
89
90 if let Some(lsp_data) = self.current_lsp_data(buffer_id) {
91 if let Some(cached_lens) = &lsp_data.code_lens {
92 if !version_queried_for.changed_since(&lsp_data.buffer_version) {
93 let has_different_servers = existing_servers.is_some_and(|existing_servers| {
94 existing_servers != cached_lens.lens.keys().copied().collect()
95 });
96 if !has_different_servers {
97 return Task::ready(Ok(Some(
98 cached_lens.lens.values().flatten().cloned().collect(),
99 )))
100 .shared();
101 }
102 } else if let Some((updating_for, running_update)) = cached_lens.update.as_ref() {
103 if !version_queried_for.changed_since(updating_for) {
104 return running_update.clone();
105 }
106 }
107 }
108 }
109
110 let lens_lsp_data = self
111 .latest_lsp_data(buffer, cx)
112 .code_lens
113 .get_or_insert_default();
114 let buffer = buffer.clone();
115 let query_version_queried_for = version_queried_for.clone();
116 let new_task = cx
117 .spawn(async move |lsp_store, cx| {
118 cx.background_executor()
119 .timer(Duration::from_millis(30))
120 .await;
121 let fetched_lens = lsp_store
122 .update(cx, |lsp_store, cx| {
123 lsp_store.fetch_code_lens_for_buffer(&buffer, cx)
124 })
125 .map_err(Arc::new)?
126 .await
127 .context("fetching code lens")
128 .map_err(Arc::new);
129 let fetched_lens = match fetched_lens {
130 Ok(fetched_lens) => fetched_lens,
131 Err(e) => {
132 lsp_store
133 .update(cx, |lsp_store, _| {
134 if let Some(lens_lsp_data) = lsp_store
135 .lsp_data
136 .get_mut(&buffer_id)
137 .and_then(|lsp_data| lsp_data.code_lens.as_mut())
138 {
139 lens_lsp_data.update = None;
140 }
141 })
142 .ok();
143 return Err(e);
144 }
145 };
146
147 lsp_store
148 .update(cx, |lsp_store, cx| {
149 let lsp_data = lsp_store.current_lsp_data(buffer_id)?;
150 let code_lens = lsp_data.code_lens.as_mut()?;
151 if let Some(fetched_lens) = fetched_lens {
152 if lsp_data.buffer_version == query_version_queried_for {
153 code_lens.lens.extend(fetched_lens);
154 } else if !lsp_data
155 .buffer_version
156 .changed_since(&query_version_queried_for)
157 {
158 lsp_data.buffer_version = query_version_queried_for;
159 code_lens.lens = fetched_lens;
160 }
161 let snapshot = buffer.read(cx).snapshot();
162 for actions in code_lens.lens.values_mut() {
163 actions
164 .sort_by(|a, b| a.range.start.cmp(&b.range.start, &snapshot));
165 }
166 }
167 code_lens.update = None;
168 Some(code_lens.lens.values().flatten().cloned().collect())
169 })
170 .map_err(Arc::new)
171 })
172 .shared();
173 lens_lsp_data.update = Some((version_queried_for, new_task.clone()));
174 new_task
175 }
176
177 fn fetch_code_lens_for_buffer(
178 &mut self,
179 buffer: &Entity<Buffer>,
180 cx: &mut Context<Self>,
181 ) -> Task<Result<Option<HashMap<LanguageServerId, Vec<CodeAction>>>>> {
182 if let Some((upstream_client, project_id)) = self.upstream_client() {
183 let request = GetCodeLens;
184 if !self.is_capable_for_proto_request(buffer, &request, cx) {
185 return Task::ready(Ok(None));
186 }
187 let request_timeout = ProjectSettings::get_global(cx)
188 .global_lsp_settings
189 .get_request_timeout();
190 let request_task = upstream_client.request_lsp(
191 project_id,
192 None,
193 request_timeout,
194 cx.background_executor().clone(),
195 request.to_proto(project_id, buffer.read(cx)),
196 );
197 let buffer = buffer.clone();
198 cx.spawn(async move |weak_lsp_store, cx| {
199 let Some(lsp_store) = weak_lsp_store.upgrade() else {
200 return Ok(None);
201 };
202 let Some(responses) = request_task.await? else {
203 return Ok(None);
204 };
205
206 let code_lens_actions = join_all(responses.payload.into_iter().map(|response| {
207 let lsp_store = lsp_store.clone();
208 let buffer = buffer.clone();
209 let cx = cx.clone();
210 async move {
211 (
212 LanguageServerId::from_proto(response.server_id),
213 GetCodeLens
214 .response_from_proto(response.response, lsp_store, buffer, cx)
215 .await,
216 )
217 }
218 }))
219 .await;
220
221 let mut has_errors = false;
222 let code_lens_actions = code_lens_actions
223 .into_iter()
224 .filter_map(|(server_id, code_lens)| match code_lens {
225 Ok(code_lens) => Some((server_id, code_lens)),
226 Err(e) => {
227 has_errors = true;
228 log::error!("{e:#}");
229 None
230 }
231 })
232 .collect::<HashMap<_, _>>();
233 anyhow::ensure!(
234 !has_errors || !code_lens_actions.is_empty(),
235 "Failed to fetch code lens"
236 );
237 Ok(Some(code_lens_actions))
238 })
239 } else {
240 let code_lens_actions_task =
241 self.request_multiple_lsp_locally(buffer, None::<usize>, GetCodeLens, cx);
242 cx.background_spawn(async move {
243 Ok(Some(code_lens_actions_task.await.into_iter().collect()))
244 })
245 }
246 }
247
248 pub fn resolve_visible_code_lenses(
249 &mut self,
250 buffer: &Entity<Buffer>,
251 visible_range: Range<Anchor>,
252 cx: &mut Context<Self>,
253 ) -> Task<Vec<CodeAction>> {
254 let buffer_id = buffer.read(cx).remote_id();
255 let snapshot = buffer.read(cx).snapshot();
256 let visible_start = visible_range.start.to_offset(&snapshot);
257 let visible_end = visible_range.end.to_offset(&snapshot);
258
259 let Some(code_lens) = self
260 .lsp_data
261 .get(&buffer_id)
262 .and_then(|data| data.code_lens.as_ref())
263 else {
264 return Task::ready(Vec::new());
265 };
266
267 let capable_servers = code_lens
268 .lens
269 .keys()
270 .filter_map(|server_id| {
271 let server = self.language_server_for_id(*server_id)?;
272 GetCodeLens::can_resolve_lens(&server.capabilities())
273 .then_some((*server_id, server))
274 })
275 .collect::<HashMap<_, _>>();
276 if capable_servers.is_empty() {
277 return Task::ready(Vec::new());
278 }
279
280 let to_resolve = code_lens
281 .lens
282 .iter()
283 .flat_map(|(server_id, actions)| {
284 let start_idx =
285 actions.partition_point(|a| a.range.start.to_offset(&snapshot) < visible_start);
286 let end_idx = start_idx
287 + actions[start_idx..]
288 .partition_point(|a| a.range.start.to_offset(&snapshot) <= visible_end);
289 actions[start_idx..end_idx].iter().enumerate().filter_map(
290 move |(local_idx, action)| {
291 let LspAction::CodeLens(lens) = &action.lsp_action else {
292 return None;
293 };
294 if lens.command.is_some() {
295 return None;
296 }
297 Some((*server_id, start_idx + local_idx, lens.clone()))
298 },
299 )
300 })
301 .collect::<Vec<_>>();
302 if to_resolve.is_empty() {
303 return Task::ready(Vec::new());
304 }
305
306 let request_timeout = ProjectSettings::get_global(cx)
307 .global_lsp_settings
308 .get_request_timeout();
309
310 cx.spawn(async move |lsp_store, cx| {
311 let mut resolved = Vec::new();
312 for (server_id, index, lens) in to_resolve {
313 let Some(server) = capable_servers.get(&server_id) else {
314 continue;
315 };
316 match server
317 .request::<lsp::request::CodeLensResolve>(lens, request_timeout)
318 .await
319 .into_response()
320 {
321 Ok(resolved_lens) => resolved.push((server_id, index, resolved_lens)),
322 Err(e) => log::warn!("Failed to resolve code lens: {e:#}"),
323 }
324 }
325 if resolved.is_empty() {
326 return Vec::new();
327 }
328
329 lsp_store
330 .update(cx, |lsp_store, _| {
331 let Some(code_lens) = lsp_store
332 .lsp_data
333 .get_mut(&buffer_id)
334 .and_then(|data| data.code_lens.as_mut())
335 else {
336 return Vec::new();
337 };
338 let mut newly_resolved = Vec::new();
339 for (server_id, index, resolved_lens) in resolved {
340 if let Some(actions) = code_lens.lens.get_mut(&server_id) {
341 if let Some(action) = actions.get_mut(index) {
342 action.resolved = true;
343 action.lsp_action = LspAction::CodeLens(resolved_lens);
344 newly_resolved.push(action.clone());
345 }
346 }
347 }
348 newly_resolved
349 })
350 .unwrap_or_default()
351 })
352 }
353
354 #[cfg(any(test, feature = "test-support"))]
355 pub fn forget_code_lens_task(&mut self, buffer_id: text::BufferId) -> Option<CodeLensTask> {
356 Some(
357 self.lsp_data
358 .get_mut(&buffer_id)?
359 .code_lens
360 .take()?
361 .update
362 .take()?
363 .1,
364 )
365 }
366
367 pub(super) async fn handle_refresh_code_lens(
368 lsp_store: Entity<Self>,
369 _: TypedEnvelope<proto::RefreshCodeLens>,
370 mut cx: AsyncApp,
371 ) -> Result<proto::Ack> {
372 lsp_store.update(&mut cx, |lsp_store, cx| {
373 lsp_store.invalidate_code_lens();
374 cx.emit(LspStoreEvent::RefreshCodeLens);
375 });
376 Ok(proto::Ack {})
377 }
378}
379
380impl Project {
381 pub fn code_lens_actions(
382 &mut self,
383 buffer: &Entity<Buffer>,
384 range: Range<Anchor>,
385 cx: &mut Context<Self>,
386 ) -> Task<Result<Option<Vec<CodeAction>>>> {
387 let snapshot = buffer.read(cx).snapshot();
388 let range = range.to_point(&snapshot);
389 let range_start = snapshot.anchor_before(range.start);
390 let range_end = if range.start == range.end {
391 range_start
392 } else {
393 snapshot.anchor_after(range.end)
394 };
395 let range = range_start..range_end;
396 let lsp_store = self.lsp_store();
397 let fetch_task =
398 lsp_store.update(cx, |lsp_store, cx| lsp_store.code_lens_actions(buffer, cx));
399 let buffer = buffer.clone();
400 cx.spawn(async move |_, cx| {
401 let mut actions = fetch_task.await?;
402 if let Some(actions) = &mut actions {
403 let resolve_task = lsp_store.update(cx, |lsp_store, cx| {
404 lsp_store.resolve_visible_code_lenses(&buffer, range.clone(), cx)
405 });
406 let resolved = resolve_task.await;
407 for resolved_action in resolved {
408 if let Some(action) = actions.iter_mut().find(|a| {
409 a.server_id == resolved_action.server_id && a.range == resolved_action.range
410 }) {
411 *action = resolved_action;
412 }
413 }
414 let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot());
415 actions.retain(|action| {
416 range.start.cmp(&action.range.start, &snapshot).is_ge()
417 && range.end.cmp(&action.range.end, &snapshot).is_le()
418 });
419 }
420 Ok(actions)
421 })
422 }
423}