1use std::sync::Arc;
2
3use anyhow::{Context as _, Result};
4use clock::Global;
5use collections::HashMap;
6use futures::{
7 FutureExt as _,
8 future::{Shared, join_all},
9};
10use gpui::{AppContext as _, AsyncApp, Context, Entity, Task};
11use language::Buffer;
12use lsp::LanguageServerId;
13use rpc::{TypedEnvelope, proto};
14use settings::Settings as _;
15use std::time::Duration;
16
17use crate::{
18 CodeAction, LspStore, LspStoreEvent,
19 lsp_command::{GetCodeLens, LspCommand as _},
20 project_settings::ProjectSettings,
21};
22
23pub(super) type CodeLensTask =
24 Shared<Task<std::result::Result<Option<Vec<CodeAction>>, Arc<anyhow::Error>>>>;
25
26#[derive(Debug, Default)]
27pub(super) struct CodeLensData {
28 pub(super) lens: HashMap<LanguageServerId, Vec<CodeAction>>,
29 pub(super) update: Option<(Global, CodeLensTask)>,
30}
31
32impl CodeLensData {
33 pub(super) fn remove_server_data(&mut self, server_id: LanguageServerId) {
34 self.lens.remove(&server_id);
35 }
36}
37
38impl LspStore {
39 pub(super) fn invalidate_code_lens(&mut self) {
40 for lsp_data in self.lsp_data.values_mut() {
41 lsp_data.code_lens = None;
42 }
43 }
44
45 pub fn code_lens_actions(
46 &mut self,
47 buffer: &Entity<Buffer>,
48 cx: &mut Context<Self>,
49 ) -> CodeLensTask {
50 let version_queried_for = buffer.read(cx).version();
51 let buffer_id = buffer.read(cx).remote_id();
52 let existing_servers = self.as_local().map(|local| {
53 local
54 .buffers_opened_in_servers
55 .get(&buffer_id)
56 .cloned()
57 .unwrap_or_default()
58 });
59
60 if let Some(lsp_data) = self.current_lsp_data(buffer_id) {
61 if let Some(cached_lens) = &lsp_data.code_lens {
62 if !version_queried_for.changed_since(&lsp_data.buffer_version) {
63 let has_different_servers = existing_servers.is_some_and(|existing_servers| {
64 existing_servers != cached_lens.lens.keys().copied().collect()
65 });
66 if !has_different_servers {
67 return Task::ready(Ok(Some(
68 cached_lens.lens.values().flatten().cloned().collect(),
69 )))
70 .shared();
71 }
72 } else if let Some((updating_for, running_update)) = cached_lens.update.as_ref() {
73 if !version_queried_for.changed_since(updating_for) {
74 return running_update.clone();
75 }
76 }
77 }
78 }
79
80 let lens_lsp_data = self
81 .latest_lsp_data(buffer, cx)
82 .code_lens
83 .get_or_insert_default();
84 let buffer = buffer.clone();
85 let query_version_queried_for = version_queried_for.clone();
86 let new_task = cx
87 .spawn(async move |lsp_store, cx| {
88 cx.background_executor()
89 .timer(Duration::from_millis(30))
90 .await;
91 let fetched_lens = lsp_store
92 .update(cx, |lsp_store, cx| lsp_store.fetch_code_lens(&buffer, cx))
93 .map_err(Arc::new)?
94 .await
95 .context("fetching code lens")
96 .map_err(Arc::new);
97 let fetched_lens = match fetched_lens {
98 Ok(fetched_lens) => fetched_lens,
99 Err(e) => {
100 lsp_store
101 .update(cx, |lsp_store, _| {
102 if let Some(lens_lsp_data) = lsp_store
103 .lsp_data
104 .get_mut(&buffer_id)
105 .and_then(|lsp_data| lsp_data.code_lens.as_mut())
106 {
107 lens_lsp_data.update = None;
108 }
109 })
110 .ok();
111 return Err(e);
112 }
113 };
114
115 lsp_store
116 .update(cx, |lsp_store, _| {
117 let lsp_data = lsp_store.current_lsp_data(buffer_id)?;
118 let code_lens = lsp_data.code_lens.as_mut()?;
119 if let Some(fetched_lens) = fetched_lens {
120 if lsp_data.buffer_version == query_version_queried_for {
121 code_lens.lens.extend(fetched_lens);
122 } else if !lsp_data
123 .buffer_version
124 .changed_since(&query_version_queried_for)
125 {
126 lsp_data.buffer_version = query_version_queried_for;
127 code_lens.lens = fetched_lens;
128 }
129 }
130 code_lens.update = None;
131 Some(code_lens.lens.values().flatten().cloned().collect())
132 })
133 .map_err(Arc::new)
134 })
135 .shared();
136 lens_lsp_data.update = Some((version_queried_for, new_task.clone()));
137 new_task
138 }
139
140 pub(super) fn fetch_code_lens(
141 &mut self,
142 buffer: &Entity<Buffer>,
143 cx: &mut Context<Self>,
144 ) -> Task<Result<Option<HashMap<LanguageServerId, Vec<CodeAction>>>>> {
145 if let Some((upstream_client, project_id)) = self.upstream_client() {
146 let request = GetCodeLens;
147 if !self.is_capable_for_proto_request(buffer, &request, cx) {
148 return Task::ready(Ok(None));
149 }
150 let request_timeout = ProjectSettings::get_global(cx)
151 .global_lsp_settings
152 .get_request_timeout();
153 let request_task = upstream_client.request_lsp(
154 project_id,
155 None,
156 request_timeout,
157 cx.background_executor().clone(),
158 request.to_proto(project_id, buffer.read(cx)),
159 );
160 let buffer = buffer.clone();
161 cx.spawn(async move |weak_lsp_store, cx| {
162 let Some(lsp_store) = weak_lsp_store.upgrade() else {
163 return Ok(None);
164 };
165 let Some(responses) = request_task.await? else {
166 return Ok(None);
167 };
168
169 let code_lens_actions = join_all(responses.payload.into_iter().map(|response| {
170 let lsp_store = lsp_store.clone();
171 let buffer = buffer.clone();
172 let cx = cx.clone();
173 async move {
174 (
175 LanguageServerId::from_proto(response.server_id),
176 GetCodeLens
177 .response_from_proto(response.response, lsp_store, buffer, cx)
178 .await,
179 )
180 }
181 }))
182 .await;
183
184 let mut has_errors = false;
185 let code_lens_actions = code_lens_actions
186 .into_iter()
187 .filter_map(|(server_id, code_lens)| match code_lens {
188 Ok(code_lens) => Some((server_id, code_lens)),
189 Err(e) => {
190 has_errors = true;
191 log::error!("{e:#}");
192 None
193 }
194 })
195 .collect::<HashMap<_, _>>();
196 anyhow::ensure!(
197 !has_errors || !code_lens_actions.is_empty(),
198 "Failed to fetch code lens"
199 );
200 Ok(Some(code_lens_actions))
201 })
202 } else {
203 let code_lens_actions_task =
204 self.request_multiple_lsp_locally(buffer, None::<usize>, GetCodeLens, cx);
205 cx.background_spawn(async move {
206 Ok(Some(code_lens_actions_task.await.into_iter().collect()))
207 })
208 }
209 }
210
211 #[cfg(any(test, feature = "test-support"))]
212 pub fn forget_code_lens_task(&mut self, buffer_id: text::BufferId) -> Option<CodeLensTask> {
213 Some(
214 self.lsp_data
215 .get_mut(&buffer_id)?
216 .code_lens
217 .take()?
218 .update
219 .take()?
220 .1,
221 )
222 }
223
224 pub(super) async fn handle_refresh_code_lens(
225 this: Entity<Self>,
226 _: TypedEnvelope<proto::RefreshCodeLens>,
227 mut cx: AsyncApp,
228 ) -> Result<proto::Ack> {
229 this.update(&mut cx, |this, cx| {
230 this.invalidate_code_lens();
231 cx.emit(LspStoreEvent::RefreshCodeLens);
232 });
233 Ok(proto::Ack {})
234 }
235}