tenant.openagents/omega

1use std::{sync::Arc, time::Duration};
2
3use anyhow::{Context as _, Result};
4use collections::{HashMap, HashSet};
5use futures::{
6    FutureExt as _,
7    future::{Shared, join_all},
8};
9use gpui::{AppContext as _, AsyncApp, Context, Entity, SharedString, Task};
10use language::{
11    Buffer, LocalFile as _, PointUtf16, point_to_lsp,
12    proto::{deserialize_lsp_edit, serialize_lsp_edit},
13};
14use lsp::LanguageServerId;
15use rpc::{TypedEnvelope, proto};
16use settings::Settings as _;
17use text::BufferId;
18use util::ResultExt as _;
19use worktree::File;
20
21use crate::{
22    ColorPresentation, DocumentColor, LspStore,
23    lsp_command::{GetDocumentColor, LspCommand as _, make_text_document_identifier},
24    lsp_store::{
25        LspStoreEvent, RunningFetch, missing_servers_to_query, next_lsp_fetch_id,
26        upstream_lsp_query_server_filter,
27    },
28    project_settings::ProjectSettings,
29};
30
31#[derive(Debug, Default, Clone)]
32pub struct DocumentColors {
33    pub colors: HashSet<DocumentColor>,
34}
35
36pub(super) type DocumentColorTask =
37    Shared<Task<std::result::Result<DocumentColors, Arc<anyhow::Error>>>>;
38
39#[derive(Debug, Default)]
40pub(super) struct DocumentColorData {
41    pub(super) colors: HashMap<LanguageServerId, HashSet<DocumentColor>>,
42    fetched_servers: HashSet<LanguageServerId>,
43    pub(super) colors_update: Option<RunningFetch<DocumentColorTask>>,
44}
45
46impl DocumentColorData {
47    pub(super) fn remove_server_data(&mut self, server_id: LanguageServerId) {
48        self.colors.remove(&server_id);
49        self.fetched_servers.remove(&server_id);
50        RunningFetch::discard_if_queried(&mut self.colors_update, server_id);
51    }
52
53    fn evict(&mut self, for_server: Option<LanguageServerId>) {
54        match for_server {
55            Some(server_id) => self.remove_server_data(server_id),
56            None => {
57                self.colors.clear();
58                self.fetched_servers.clear();
59            }
60        }
61        self.colors_update = None;
62    }
63}
64
65impl LspStore {
66    pub(super) fn refresh_document_colors(
67        &mut self,
68        for_server: Option<LanguageServerId>,
69        cx: &mut Context<Self>,
70    ) {
71        for lsp_data in self.lsp_data.values_mut() {
72            if let Some(document_colors) = &mut lsp_data.document_colors {
73                document_colors.evict(for_server);
74            }
75        }
76
77        cx.emit(LspStoreEvent::RefreshDocumentColors {
78            server_id: for_server,
79        });
80        if let Some((downstream_client, project_id)) = self.downstream_client.as_ref() {
81            downstream_client
82                .send(proto::RefreshDocumentColors {
83                    project_id: *project_id,
84                    server_id: for_server.map(|server_id| server_id.to_proto()),
85                })
86                .context("sending refresh document colors downstream")
87                .log_err();
88        }
89    }
90
91    pub(super) async fn handle_refresh_document_colors(
92        lsp_store: Entity<Self>,
93        envelope: TypedEnvelope<proto::RefreshDocumentColors>,
94        mut cx: AsyncApp,
95    ) -> Result<proto::Ack> {
96        lsp_store.update(&mut cx, |lsp_store, cx| {
97            let server_id = envelope.payload.server_id.map(LanguageServerId::from_proto);
98            lsp_store.refresh_document_colors(server_id, cx);
99        });
100        Ok(proto::Ack {})
101    }
102
103    pub fn document_colors(
104        &mut self,
105        buffer: Entity<Buffer>,
106        cx: &mut Context<Self>,
107    ) -> Option<DocumentColorTask> {
108        let version_queried_for = buffer.read(cx).version();
109        let buffer_id = buffer.read(cx).remote_id();
110
111        let current_servers = self.relevant_server_ids_for_capability_check(&buffer, cx);
112
113        let mut servers_to_query = None;
114        if let Some(lsp_data) = self.current_lsp_data(buffer_id) {
115            if !version_queried_for.changed_since(&lsp_data.buffer_version)
116                && let Some(cached_colors) = &mut lsp_data.document_colors
117            {
118                match missing_servers_to_query(
119                    &mut cached_colors.colors,
120                    &mut cached_colors.fetched_servers,
121                    &current_servers,
122                ) {
123                    Some(missing_servers) => servers_to_query = Some(missing_servers),
124                    None => {
125                        return Some(
126                            Task::ready(Ok(DocumentColors {
127                                colors: cached_colors.colors.values().flatten().cloned().collect(),
128                            }))
129                            .shared(),
130                        );
131                    }
132                }
133            }
134            if let Some(document_colors) = &lsp_data.document_colors
135                && let Some(running) = &document_colors.colors_update
136                && !version_queried_for.changed_since(&running.version)
137                && servers_to_query
138                    .as_ref()
139                    .is_none_or(|missing| missing.is_subset(&running.servers))
140            {
141                return Some(running.task.clone());
142            }
143        }
144
145        let color_lsp_data = self
146            .latest_lsp_data(&buffer, cx)
147            .document_colors
148            .get_or_insert_default();
149        let fetch_id = next_lsp_fetch_id();
150        let queried_servers = servers_to_query
151            .clone()
152            .unwrap_or_else(|| current_servers.clone());
153        let buffer_version_queried_for = version_queried_for.clone();
154        let new_task = cx
155            .spawn({
156                let queried_servers = queried_servers.clone();
157                async move |lsp_store, cx| {
158                    cx.background_executor()
159                        .timer(Duration::from_millis(30))
160                        .await;
161                    let fetched_colors = lsp_store
162                        .update(cx, |lsp_store, cx| {
163                            lsp_store.fetch_document_colors_for_buffer(
164                                &buffer,
165                                servers_to_query,
166                                cx,
167                            )
168                        })?
169                        .await
170                        .context("fetching document colors")
171                        .map_err(Arc::new);
172                    let fetched_colors = match fetched_colors {
173                        Ok(fetched_colors) => {
174                            if buffer.update(cx, |buffer, _| {
175                                buffer.version() != buffer_version_queried_for
176                            }) {
177                                return Ok(DocumentColors::default());
178                            }
179                            fetched_colors
180                        }
181                        Err(e) => {
182                            lsp_store
183                                .update(cx, |lsp_store, _| {
184                                    if let Some(lsp_data) = lsp_store.lsp_data.get_mut(&buffer_id)
185                                        && let Some(document_colors) = &mut lsp_data.document_colors
186                                    {
187                                        RunningFetch::take_finished(
188                                            &mut document_colors.colors_update,
189                                            fetch_id,
190                                        );
191                                    }
192                                })
193                                .ok();
194                            return Err(e);
195                        }
196                    };
197
198                    lsp_store
199                        .update(cx, |lsp_store, cx| {
200                            let lsp_data = lsp_store.latest_lsp_data(&buffer, cx);
201                            let lsp_colors = lsp_data.document_colors.get_or_insert_default();
202
203                            if RunningFetch::take_finished(&mut lsp_colors.colors_update, fetch_id)
204                                && let Some(fetched_colors) = fetched_colors
205                            {
206                                if lsp_data.buffer_version == buffer_version_queried_for {
207                                    lsp_colors.colors.extend(fetched_colors);
208                                    lsp_colors.fetched_servers.extend(queried_servers);
209                                } else if !lsp_data
210                                    .buffer_version
211                                    .changed_since(&buffer_version_queried_for)
212                                {
213                                    lsp_data.buffer_version = buffer_version_queried_for;
214                                    lsp_colors.colors = fetched_colors;
215                                    lsp_colors.fetched_servers = queried_servers;
216                                }
217                            }
218                            let colors = lsp_colors
219                                .colors
220                                .values()
221                                .flatten()
222                                .cloned()
223                                .collect::<HashSet<_>>();
224                            DocumentColors { colors }
225                        })
226                        .map_err(Arc::new)
227                }
228            })
229            .shared();
230        color_lsp_data.colors_update = Some(RunningFetch {
231            id: fetch_id,
232            version: version_queried_for,
233            servers: queried_servers,
234            task: new_task.clone(),
235        });
236        Some(new_task)
237    }
238
239    pub fn resolve_color_presentation(
240        &mut self,
241        mut color: DocumentColor,
242        buffer: Entity<Buffer>,
243        server_id: LanguageServerId,
244        cx: &mut Context<Self>,
245    ) -> Task<Result<DocumentColor>> {
246        if color.resolved {
247            return Task::ready(Ok(color));
248        }
249
250        if let Some((upstream_client, project_id)) = self.upstream_client() {
251            let start = color.lsp_range.start;
252            let end = color.lsp_range.end;
253            let request = proto::GetColorPresentation {
254                project_id,
255                server_id: server_id.to_proto(),
256                buffer_id: buffer.read(cx).remote_id().into(),
257                color: Some(proto::ColorInformation {
258                    red: color.color.red,
259                    green: color.color.green,
260                    blue: color.color.blue,
261                    alpha: color.color.alpha,
262                    lsp_range_start: Some(proto::PointUtf16 {
263                        row: start.line,
264                        column: start.character,
265                    }),
266                    lsp_range_end: Some(proto::PointUtf16 {
267                        row: end.line,
268                        column: end.character,
269                    }),
270                }),
271            };
272            cx.background_spawn(async move {
273                let response = upstream_client
274                    .request(request)
275                    .await
276                    .context("color presentation proto request")?;
277                color.resolved = true;
278                color.color_presentations = response
279                    .presentations
280                    .into_iter()
281                    .map(|presentation| ColorPresentation {
282                        label: SharedString::from(presentation.label),
283                        text_edit: presentation.text_edit.and_then(deserialize_lsp_edit),
284                        additional_text_edits: presentation
285                            .additional_text_edits
286                            .into_iter()
287                            .filter_map(deserialize_lsp_edit)
288                            .collect(),
289                    })
290                    .collect();
291                Ok(color)
292            })
293        } else {
294            let path = match buffer
295                .update(cx, |buffer, cx| {
296                    Some(File::from_dyn(buffer.file())?.abs_path(cx))
297                })
298                .context("buffer with the missing path")
299            {
300                Ok(path) => path,
301                Err(e) => return Task::ready(Err(e)),
302            };
303            let Some(lang_server) = buffer.update(cx, |buffer, cx| {
304                self.language_server_for_local_buffer(buffer, server_id, cx)
305                    .map(|(_, server)| server.clone())
306            }) else {
307                return Task::ready(Ok(color));
308            };
309
310            let request_timeout = ProjectSettings::get_global(cx)
311                .global_lsp_settings
312                .get_request_timeout();
313            cx.background_spawn(async move {
314                let resolve_task = lang_server.request::<lsp::request::ColorPresentationRequest>(
315                    lsp::ColorPresentationParams {
316                        text_document: make_text_document_identifier(&path)?,
317                        color: color.color,
318                        range: color.lsp_range,
319                        work_done_progress_params: Default::default(),
320                        partial_result_params: Default::default(),
321                    },
322                    request_timeout,
323                );
324                color.color_presentations = resolve_task
325                    .await
326                    .into_response()
327                    .context("color presentation resolve LSP request")?
328                    .into_iter()
329                    .map(|presentation| ColorPresentation {
330                        label: SharedString::from(presentation.label),
331                        text_edit: presentation.text_edit,
332                        additional_text_edits: presentation
333                            .additional_text_edits
334                            .unwrap_or_default(),
335                    })
336                    .collect();
337                color.resolved = true;
338                Ok(color)
339            })
340        }
341    }
342
343    pub(super) fn fetch_document_colors_for_buffer(
344        &mut self,
345        buffer: &Entity<Buffer>,
346        for_servers: Option<HashSet<LanguageServerId>>,
347        cx: &mut Context<Self>,
348    ) -> Task<anyhow::Result<Option<HashMap<LanguageServerId, HashSet<DocumentColor>>>>> {
349        if let Some((client, project_id)) = self.upstream_client() {
350            let request = GetDocumentColor {};
351            if !self.is_capable_for_proto_request(buffer, &request, cx) {
352                return Task::ready(Ok(None));
353            }
354
355            let request_timeout = ProjectSettings::get_global(cx)
356                .global_lsp_settings
357                .get_request_timeout();
358            let request_task = client.request_lsp(
359                project_id,
360                upstream_lsp_query_server_filter(for_servers.as_ref()),
361                request_timeout,
362                cx.background_executor().clone(),
363                request.to_proto(project_id, buffer.read(cx)),
364            );
365            let buffer = buffer.clone();
366            cx.spawn(async move |lsp_store, cx| {
367                let Some(lsp_store) = lsp_store.upgrade() else {
368                    return Ok(None);
369                };
370                let colors: HashMap<LanguageServerId, HashSet<DocumentColor>> = join_all(
371                    request_task
372                        .await
373                        .log_err()
374                        .flatten()
375                        .map(|response| response.payload)
376                        .unwrap_or_default()
377                        .into_iter()
378                        .map(|color_response| {
379                            let response = request.response_from_proto(
380                                color_response.response,
381                                lsp_store.clone(),
382                                buffer.clone(),
383                                cx.clone(),
384                            );
385                            async move {
386                                let colors: Vec<DocumentColor> =
387                                    response.await.log_err().unwrap_or_default();
388                                (
389                                    LanguageServerId::from_proto(color_response.server_id),
390                                    colors,
391                                )
392                            }
393                        }),
394                )
395                .await
396                .into_iter()
397                .fold(HashMap::default(), |mut acc, (server_id, colors)| {
398                    acc.entry(server_id)
399                        .or_insert_with(HashSet::default)
400                        .extend(colors);
401                    acc
402                });
403                Ok(Some(colors))
404            })
405        } else {
406            let document_colors_task = self.request_filtered_lsp_locally(
407                buffer,
408                None::<usize>,
409                GetDocumentColor,
410                for_servers.as_ref(),
411                cx,
412            );
413            cx.background_spawn(async move {
414                Ok(Some(
415                    document_colors_task
416                        .await
417                        .into_iter()
418                        .fold(HashMap::default(), |mut acc, (server_id, colors)| {
419                            acc.entry(server_id)
420                                .or_insert_with(HashSet::default)
421                                .extend(colors);
422                            acc
423                        })
424                        .into_iter()
425                        .collect(),
426                ))
427            })
428        }
429    }
430
431    pub(super) async fn handle_get_color_presentation(
432        lsp_store: Entity<Self>,
433        envelope: TypedEnvelope<proto::GetColorPresentation>,
434        mut cx: AsyncApp,
435    ) -> Result<proto::GetColorPresentationResponse> {
436        let buffer_id = BufferId::new(envelope.payload.buffer_id)?;
437        let buffer = lsp_store.update(&mut cx, |lsp_store, cx| {
438            lsp_store.buffer_store.read(cx).get_existing(buffer_id)
439        })?;
440
441        let color = envelope
442            .payload
443            .color
444            .context("invalid color resolve request")?;
445        let start = color
446            .lsp_range_start
447            .context("invalid color resolve request")?;
448        let end = color
449            .lsp_range_end
450            .context("invalid color resolve request")?;
451
452        let color = DocumentColor {
453            lsp_range: lsp::Range {
454                start: point_to_lsp(PointUtf16::new(start.row, start.column)),
455                end: point_to_lsp(PointUtf16::new(end.row, end.column)),
456            },
457            color: lsp::Color {
458                red: color.red,
459                green: color.green,
460                blue: color.blue,
461                alpha: color.alpha,
462            },
463            resolved: false,
464            color_presentations: Vec::new(),
465        };
466        let resolved_color = lsp_store
467            .update(&mut cx, |lsp_store, cx| {
468                lsp_store.resolve_color_presentation(
469                    color,
470                    buffer.clone(),
471                    LanguageServerId(envelope.payload.server_id as usize),
472                    cx,
473                )
474            })
475            .await
476            .context("resolving color presentation")?;
477
478        Ok(proto::GetColorPresentationResponse {
479            presentations: resolved_color
480                .color_presentations
481                .into_iter()
482                .map(|presentation| proto::ColorPresentation {
483                    label: presentation.label.to_string(),
484                    text_edit: presentation.text_edit.map(serialize_lsp_edit),
485                    additional_text_edits: presentation
486                        .additional_text_edits
487                        .into_iter()
488                        .map(serialize_lsp_edit)
489                        .collect(),
490                })
491                .collect(),
492        })
493    }
494}
495