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;
17
18use crate::{
19 CodeAction, LspAction, LspStore, LspStoreEvent,
20 lsp_command::{GetCodeLens, LspCommand as _},
21 project_settings::ProjectSettings,
22};
23
24pub(super) type CodeLensTask =
25 Shared<Task<std::result::Result<Option<Vec<CodeAction>>, Arc<anyhow::Error>>>>;
26
27#[derive(Debug, Default)]
28pub(super) struct CodeLensData {
29 pub(super) lens: HashMap<LanguageServerId, Vec<CodeAction>>,
30 pub(super) update: Option<(Global, CodeLensTask)>,
31}
32
33impl CodeLensData {
34 pub(super) fn remove_server_data(&mut self, server_id: LanguageServerId) {
35 self.lens.remove(&server_id);
36 }
37}
38
39impl LspStore {
40 pub(super) fn invalidate_code_lens(&mut self) {
41 for lsp_data in self.lsp_data.values_mut() {
42 lsp_data.code_lens = None;
43 }
44 }
45
46 pub fn code_lens_actions(
47 &mut self,
48 buffer: &Entity<Buffer>,
49 cx: &mut Context<Self>,
50 ) -> CodeLensTask {
51 let version_queried_for = buffer.read(cx).version();
52 let buffer_id = buffer.read(cx).remote_id();
53 let existing_servers = self.as_local().map(|local| {
54 local
55 .buffers_opened_in_servers
56 .get(&buffer_id)
57 .cloned()
58 .unwrap_or_default()
59 });
60
61 if let Some(lsp_data) = self.current_lsp_data(buffer_id) {
62 if let Some(cached_lens) = &lsp_data.code_lens {
63 if !version_queried_for.changed_since(&lsp_data.buffer_version) {
64 let has_different_servers = existing_servers.is_some_and(|existing_servers| {
65 existing_servers != cached_lens.lens.keys().copied().collect()
66 });
67 if !has_different_servers {
68 return Task::ready(Ok(Some(
69 cached_lens.lens.values().flatten().cloned().collect(),
70 )))
71 .shared();
72 }
73 } else if let Some((updating_for, running_update)) = cached_lens.update.as_ref() {
74 if !version_queried_for.changed_since(updating_for) {
75 return running_update.clone();
76 }
77 }
78 }
79 }
80
81 let lens_lsp_data = self
82 .latest_lsp_data(buffer, cx)
83 .code_lens
84 .get_or_insert_default();
85 let buffer = buffer.clone();
86 let query_version_queried_for = version_queried_for.clone();
87 let new_task = cx
88 .spawn(async move |lsp_store, cx| {
89 cx.background_executor()
90 .timer(Duration::from_millis(30))
91 .await;
92 let fetched_lens = lsp_store
93 .update(cx, |lsp_store, cx| lsp_store.fetch_code_lens(&buffer, cx))
94 .map_err(Arc::new)?
95 .await
96 .context("fetching code lens")
97 .map_err(Arc::new);
98 let fetched_lens = match fetched_lens {
99 Ok(fetched_lens) => fetched_lens,
100 Err(e) => {
101 lsp_store
102 .update(cx, |lsp_store, _| {
103 if let Some(lens_lsp_data) = lsp_store
104 .lsp_data
105 .get_mut(&buffer_id)
106 .and_then(|lsp_data| lsp_data.code_lens.as_mut())
107 {
108 lens_lsp_data.update = None;
109 }
110 })
111 .ok();
112 return Err(e);
113 }
114 };
115
116 lsp_store
117 .update(cx, |lsp_store, cx| {
118 let lsp_data = lsp_store.current_lsp_data(buffer_id)?;
119 let code_lens = lsp_data.code_lens.as_mut()?;
120 if let Some(fetched_lens) = fetched_lens {
121 if lsp_data.buffer_version == query_version_queried_for {
122 code_lens.lens.extend(fetched_lens);
123 } else if !lsp_data
124 .buffer_version
125 .changed_since(&query_version_queried_for)
126 {
127 lsp_data.buffer_version = query_version_queried_for;
128 code_lens.lens = fetched_lens;
129 }
130 let snapshot = buffer.read(cx).snapshot();
131 for actions in code_lens.lens.values_mut() {
132 actions
133 .sort_by(|a, b| a.range.start.cmp(&b.range.start, &snapshot));
134 }
135 }
136 code_lens.update = None;
137 Some(code_lens.lens.values().flatten().cloned().collect())
138 })
139 .map_err(Arc::new)
140 })
141 .shared();
142 lens_lsp_data.update = Some((version_queried_for, new_task.clone()));
143 new_task
144 }
145
146 pub(super) fn fetch_code_lens(
147 &mut self,
148 buffer: &Entity<Buffer>,
149 cx: &mut Context<Self>,
150 ) -> Task<Result<Option<HashMap<LanguageServerId, Vec<CodeAction>>>>> {
151 if let Some((upstream_client, project_id)) = self.upstream_client() {
152 let request = GetCodeLens;
153 if !self.is_capable_for_proto_request(buffer, &request, cx) {
154 return Task::ready(Ok(None));
155 }
156 let request_timeout = ProjectSettings::get_global(cx)
157 .global_lsp_settings
158 .get_request_timeout();
159 let request_task = upstream_client.request_lsp(
160 project_id,
161 None,
162 request_timeout,
163 cx.background_executor().clone(),
164 request.to_proto(project_id, buffer.read(cx)),
165 );
166 let buffer = buffer.clone();
167 cx.spawn(async move |weak_lsp_store, cx| {
168 let Some(lsp_store) = weak_lsp_store.upgrade() else {
169 return Ok(None);
170 };
171 let Some(responses) = request_task.await? else {
172 return Ok(None);
173 };
174
175 let code_lens_actions = join_all(responses.payload.into_iter().map(|response| {
176 let lsp_store = lsp_store.clone();
177 let buffer = buffer.clone();
178 let cx = cx.clone();
179 async move {
180 (
181 LanguageServerId::from_proto(response.server_id),
182 GetCodeLens
183 .response_from_proto(response.response, lsp_store, buffer, cx)
184 .await,
185 )
186 }
187 }))
188 .await;
189
190 let mut has_errors = false;
191 let code_lens_actions = code_lens_actions
192 .into_iter()
193 .filter_map(|(server_id, code_lens)| match code_lens {
194 Ok(code_lens) => Some((server_id, code_lens)),
195 Err(e) => {
196 has_errors = true;
197 log::error!("{e:#}");
198 None
199 }
200 })
201 .collect::<HashMap<_, _>>();
202 anyhow::ensure!(
203 !has_errors || !code_lens_actions.is_empty(),
204 "Failed to fetch code lens"
205 );
206 Ok(Some(code_lens_actions))
207 })
208 } else {
209 let code_lens_actions_task =
210 self.request_multiple_lsp_locally(buffer, None::<usize>, GetCodeLens, cx);
211 cx.background_spawn(async move {
212 Ok(Some(code_lens_actions_task.await.into_iter().collect()))
213 })
214 }
215 }
216
217 pub fn resolve_visible_code_lenses(
218 &mut self,
219 buffer: &Entity<Buffer>,
220 visible_range: Range<Anchor>,
221 cx: &mut Context<Self>,
222 ) -> Task<Vec<CodeAction>> {
223 let buffer_id = buffer.read(cx).remote_id();
224 let snapshot = buffer.read(cx).snapshot();
225 let visible_start = visible_range.start.to_offset(&snapshot);
226 let visible_end = visible_range.end.to_offset(&snapshot);
227
228 let Some(code_lens) = self
229 .lsp_data
230 .get(&buffer_id)
231 .and_then(|data| data.code_lens.as_ref())
232 else {
233 return Task::ready(Vec::new());
234 };
235
236 let capable_servers = code_lens
237 .lens
238 .keys()
239 .filter_map(|server_id| {
240 let server = self.language_server_for_id(*server_id)?;
241 GetCodeLens::can_resolve_lens(&server.capabilities())
242 .then_some((*server_id, server))
243 })
244 .collect::<HashMap<_, _>>();
245 if capable_servers.is_empty() {
246 return Task::ready(Vec::new());
247 }
248
249 let to_resolve = code_lens
250 .lens
251 .iter()
252 .flat_map(|(server_id, actions)| {
253 let start_idx =
254 actions.partition_point(|a| a.range.start.to_offset(&snapshot) < visible_start);
255 let end_idx = start_idx
256 + actions[start_idx..]
257 .partition_point(|a| a.range.start.to_offset(&snapshot) <= visible_end);
258 actions[start_idx..end_idx].iter().enumerate().filter_map(
259 move |(local_idx, action)| {
260 let LspAction::CodeLens(lens) = &action.lsp_action else {
261 return None;
262 };
263 if lens.command.is_some() {
264 return None;
265 }
266 Some((*server_id, start_idx + local_idx, lens.clone()))
267 },
268 )
269 })
270 .collect::<Vec<_>>();
271 if to_resolve.is_empty() {
272 return Task::ready(Vec::new());
273 }
274
275 let request_timeout = ProjectSettings::get_global(cx)
276 .global_lsp_settings
277 .get_request_timeout();
278
279 cx.spawn(async move |lsp_store, cx| {
280 let mut resolved = Vec::new();
281 for (server_id, index, lens) in to_resolve {
282 let Some(server) = capable_servers.get(&server_id) else {
283 continue;
284 };
285 match server
286 .request::<lsp::request::CodeLensResolve>(lens, request_timeout)
287 .await
288 .into_response()
289 {
290 Ok(resolved_lens) => resolved.push((server_id, index, resolved_lens)),
291 Err(e) => log::warn!("Failed to resolve code lens: {e:#}"),
292 }
293 }
294 if resolved.is_empty() {
295 return Vec::new();
296 }
297
298 lsp_store
299 .update(cx, |lsp_store, _| {
300 let Some(code_lens) = lsp_store
301 .lsp_data
302 .get_mut(&buffer_id)
303 .and_then(|data| data.code_lens.as_mut())
304 else {
305 return Vec::new();
306 };
307 let mut newly_resolved = Vec::new();
308 for (server_id, index, resolved_lens) in resolved {
309 if let Some(actions) = code_lens.lens.get_mut(&server_id) {
310 if let Some(action) = actions.get_mut(index) {
311 action.resolved = true;
312 action.lsp_action = LspAction::CodeLens(resolved_lens);
313 newly_resolved.push(action.clone());
314 }
315 }
316 }
317 newly_resolved
318 })
319 .unwrap_or_default()
320 })
321 }
322
323 #[cfg(any(test, feature = "test-support"))]
324 pub fn forget_code_lens_task(&mut self, buffer_id: text::BufferId) -> Option<CodeLensTask> {
325 Some(
326 self.lsp_data
327 .get_mut(&buffer_id)?
328 .code_lens
329 .take()?
330 .update
331 .take()?
332 .1,
333 )
334 }
335
336 pub(super) async fn handle_refresh_code_lens(
337 this: Entity<Self>,
338 _: TypedEnvelope<proto::RefreshCodeLens>,
339 mut cx: AsyncApp,
340 ) -> Result<proto::Ack> {
341 this.update(&mut cx, |this, cx| {
342 this.invalidate_code_lens();
343 cx.emit(LspStoreEvent::RefreshCodeLens);
344 });
345 Ok(proto::Ack {})
346 }
347}