tenant.openagents/omega

1use collections::HashMap;
2use futures::future::join_all;
3use gpui::{App, Entity, Task};
4use itertools::Itertools;
5use language::{Buffer, BufferSnapshot};
6use lsp::LanguageServerId;
7use project::lsp_store::{BufferDocumentLinks, LspDocumentLink, ResolvedDocumentLink};
8use settings::Settings;
9use text::BufferId;
10use ui::Context;
11
12use crate::{Editor, LSP_REQUEST_DEBOUNCE_TIMEOUT, editor_settings::EditorSettings};
13
14pub(super) struct LspDocumentLinks {
15    pub(super) enabled: bool,
16    pub(super) per_buffer: HashMap<BufferId, BufferDocumentLinks>,
17    pub(super) refresh_task: Task<()>,
18}
19
20impl LspDocumentLinks {
21    pub(super) fn new(cx: &App) -> Self {
22        Self {
23            enabled: EditorSettings::get_global(cx).lsp_document_links,
24            per_buffer: HashMap::default(),
25            refresh_task: Task::ready(()),
26        }
27    }
28}
29
30impl Editor {
31    pub(super) fn refresh_document_links(
32        &mut self,
33        for_buffer: Option<BufferId>,
34        cx: &mut Context<Self>,
35    ) {
36        if !self.lsp_data_enabled() || !self.lsp_document_links.enabled {
37            return;
38        }
39        let Some(project) = self.project.as_ref().map(|p| p.downgrade()) else {
40            return;
41        };
42
43        let buffers_to_query = self
44            .visible_buffers(cx)
45            .into_iter()
46            .filter(|buffer| self.is_lsp_relevant(buffer.read(cx).file(), cx))
47            .chain(for_buffer.and_then(|id| self.buffer.read(cx).buffer(id)))
48            .filter(|buffer| {
49                let id = buffer.read(cx).remote_id();
50                for_buffer.is_none_or(|target| target == id)
51                    && self.registered_buffers.contains_key(&id)
52            })
53            .unique_by(|buffer| buffer.read(cx).remote_id())
54            .collect::<Vec<_>>();
55        if buffers_to_query.is_empty() {
56            self.lsp_document_links.refresh_task = Task::ready(());
57            return;
58        }
59
60        self.lsp_document_links.refresh_task = cx.spawn(async move |editor, cx| {
61            cx.background_executor()
62                .timer(LSP_REQUEST_DEBOUNCE_TIMEOUT)
63                .await;
64
65            let Some(tasks_for_buffers) = project
66                .update(cx, |project, cx| {
67                    project.lsp_store().update(cx, |lsp_store, cx| {
68                        buffers_to_query
69                            .into_iter()
70                            .map(|buffer| {
71                                let buffer_id = buffer.read(cx).remote_id();
72                                let task = lsp_store.fetch_document_links(&buffer, cx);
73                                async move { (buffer_id, task.await) }
74                            })
75                            .collect::<Vec<_>>()
76                    })
77                })
78                .ok()
79            else {
80                return;
81            };
82
83            let new_links_for_buffers = join_all(tasks_for_buffers).await;
84            editor
85                .update(cx, |editor, _| {
86                    for (buffer_id, links) in new_links_for_buffers {
87                        let Some(links) = links else {
88                            continue;
89                        };
90                        if links.is_empty() {
91                            editor.lsp_document_links.per_buffer.remove(&buffer_id);
92                        } else {
93                            editor
94                                .lsp_document_links
95                                .per_buffer
96                                .insert(buffer_id, links);
97                        }
98                    }
99                })
100                .ok();
101        });
102    }
103
104    /// Returns a task yielding the resolved document links covering `position`
105    /// in `buffer`, paired with the language server that owns each link.
106    /// Resolution is deduplicated through `LspStore`'s per-`(server_id,
107    /// link_id)` `Shared` task; the editor's mirror is updated when the
108    /// resolves complete so subsequent renders/hovers find resolved data
109    /// without re-issuing requests.
110    ///
111    /// Returns `None` when nothing is cached at `position` so callers can skip
112    /// spawning anything.
113    pub fn document_links_at(
114        &mut self,
115        buffer: Entity<Buffer>,
116        position: text::Anchor,
117        cx: &mut Context<Self>,
118    ) -> Option<Task<Vec<(LanguageServerId, LspDocumentLink)>>> {
119        let buffer_id = buffer.read(cx).remote_id();
120        let snapshot = buffer.read(cx).snapshot();
121        let matches = self
122            .lsp_document_links
123            .per_buffer
124            .get(&buffer_id)?
125            .iter()
126            .flat_map(|(server_id, per_server)| {
127                per_server
128                    .iter()
129                    .map(move |(link_id, link)| (*server_id, *link_id, link))
130            })
131            .filter(|(_, _, link)| link_contains(link, &position, &snapshot))
132            .map(|(server_id, link_id, link)| (server_id, link_id, link.clone()))
133            .collect::<Vec<_>>();
134        if matches.is_empty() {
135            return None;
136        }
137
138        let project = self.project.clone()?;
139        let mut resolved_links = Vec::with_capacity(matches.len());
140        let mut pending = Vec::new();
141        project.update(cx, |project, cx| {
142            project.lsp_store().update(cx, |lsp_store, cx| {
143                for (server_id, link_id, _link) in matches {
144                    match lsp_store.resolved_document_link(&buffer, server_id, link_id, cx) {
145                        Some(ResolvedDocumentLink::Resolved(resolved)) => {
146                            resolved_links.push((server_id, link_id, resolved));
147                        }
148                        Some(ResolvedDocumentLink::Resolving(task)) => {
149                            pending.push((server_id, task));
150                        }
151                        None => {
152                            // Cache no longer holds the link (likely a version
153                            // bump between the mirror snapshot and now); skip.
154                        }
155                    }
156                }
157            })
158        });
159
160        if pending.is_empty() {
161            return Some(Task::ready(
162                resolved_links
163                    .into_iter()
164                    .map(|(server_id, _, link)| (server_id, link))
165                    .collect(),
166            ));
167        }
168
169        Some(cx.spawn(async move |editor, cx| {
170            let pending_results =
171                join_all(pending.into_iter().map(|(server_id, task)| async move {
172                    task.await.map(|(link_id, link)| (server_id, link_id, link))
173                }))
174                .await;
175            resolved_links.extend(pending_results.into_iter().flatten());
176            editor
177                .update(cx, |editor, cx| {
178                    if let Some(by_server) =
179                        editor.lsp_document_links.per_buffer.get_mut(&buffer_id)
180                    {
181                        for (server_id, link_id, resolved) in &resolved_links {
182                            if let Some(slot) = by_server
183                                .get_mut(server_id)
184                                .and_then(|links| links.get_mut(link_id))
185                            {
186                                *slot = resolved.clone();
187                            }
188                        }
189                    }
190                    cx.notify();
191                })
192                .ok();
193
194            resolved_links
195                .into_iter()
196                .map(|(server_id, _, link)| (server_id, link))
197                .collect()
198        }))
199    }
200}
201
202fn link_contains(
203    link: &LspDocumentLink,
204    position: &text::Anchor,
205    snapshot: &BufferSnapshot,
206) -> bool {
207    link.range.start.cmp(position, snapshot).is_le()
208        && link.range.end.cmp(position, snapshot).is_ge()
209}
210