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