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