tenant.openagents/omega

1use std::{collections::hash_map, ops::Range, sync::Arc};
2
3use anyhow::{Context as _, Result};
4use collections::{HashMap, HashSet};
5use futures::future::Shared;
6use gpui::{App, AppContext as _, AsyncApp, Context, Entity, Task};
7use language::{
8    Buffer,
9    row_chunk::{RowChunk, RowChunks},
10};
11use lsp::LanguageServerId;
12use rpc::{TypedEnvelope, proto};
13use settings::Settings as _;
14use text::{BufferId, Point};
15use util::ResultExt as _;
16
17use crate::{
18    InlayHint, InlayId, LspStore, LspStoreEvent, ResolveState, lsp_command::InlayHints,
19    project_settings::ProjectSettings,
20};
21
22pub type CacheInlayHints = HashMap<LanguageServerId, Vec<(InlayId, InlayHint)>>;
23pub type CacheInlayHintsTask = Shared<Task<Result<CacheInlayHints, Arc<anyhow::Error>>>>;
24
25/// A logic to apply when querying for new inlay hints and deciding what to do with the old entries in the cache in case of conflicts.
26#[derive(Debug, Clone, Copy)]
27pub enum InvalidationStrategy {
28    /// Language servers reset hints via <a href="https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#workspace_inlayHint_refresh">request</a>.
29    /// Demands to re-query all inlay hints needed and invalidate all cached entries, but does not require instant update with invalidation.
30    ///
31    /// Despite nothing forbids language server from sending this request on every edit, it is expected to be sent only when certain internal server state update, invisible for the editor otherwise.
32    RefreshRequested { server_id: LanguageServerId },
33    /// Multibuffer excerpt(s) and/or singleton buffer(s) were edited at least on one place.
34    /// Neither editor nor LSP is able to tell which open file hints' are not affected, so all of them have to be invalidated, re-queried and do that fast enough to avoid being slow, but also debounce to avoid loading hints on every fast keystroke sequence.
35    BufferEdited,
36    /// A new file got opened/new excerpt was added to a multibuffer/a [multi]buffer was scrolled to a new position.
37    /// No invalidation should be done at all, all new hints are added to the cache.
38    ///
39    /// A special case is the editor toggles and settings change:
40    /// in addition to LSP capabilities, Zed allows omitting certain hint kinds (defined by the corresponding LSP part: type/parameter/other) and toggling hints.
41    /// This does not lead to cache invalidation, but would require cache usage for determining which hints are not displayed and issuing an update to inlays on the screen.
42    None,
43}
44
45impl InvalidationStrategy {
46    pub fn should_invalidate(&self) -> bool {
47        matches!(
48            self,
49            InvalidationStrategy::RefreshRequested { .. } | InvalidationStrategy::BufferEdited
50        )
51    }
52}
53
54pub struct BufferInlayHints {
55    chunks: RowChunks,
56    hints_by_chunks: Vec<Option<CacheInlayHints>>,
57    fetches_by_chunks: Vec<Option<CacheInlayHintsTask>>,
58    hints_by_id: HashMap<InlayId, HintForId>,
59    pending_refreshes: HashSet<LanguageServerId>,
60    work_end_refreshes: HashSet<LanguageServerId>,
61    pub(super) fetched_servers: HashSet<LanguageServerId>,
62    pub(super) hint_resolves: HashMap<InlayId, Shared<Task<()>>>,
63}
64
65#[derive(Debug, Clone, Copy)]
66struct HintForId {
67    chunk_id: usize,
68    server_id: LanguageServerId,
69    position: usize,
70}
71
72impl std::fmt::Debug for BufferInlayHints {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        f.debug_struct("BufferInlayHints")
75            .field("buffer_chunks", &self.chunks)
76            .field("hints_by_chunks", &self.hints_by_chunks)
77            .field("fetches_by_chunks", &self.fetches_by_chunks)
78            .field("hints_by_id", &self.hints_by_id)
79            .finish_non_exhaustive()
80    }
81}
82
83const MAX_ROWS_IN_A_CHUNK: u32 = 50;
84
85impl BufferInlayHints {
86    pub fn new(buffer: &Entity<Buffer>, cx: &mut App) -> Self {
87        let chunks = RowChunks::new(buffer.read(cx).as_text_snapshot(), MAX_ROWS_IN_A_CHUNK);
88
89        Self {
90            hints_by_chunks: vec![None; chunks.len()],
91            fetches_by_chunks: vec![None; chunks.len()],
92            pending_refreshes: HashSet::default(),
93            work_end_refreshes: HashSet::default(),
94            fetched_servers: HashSet::default(),
95            hints_by_id: HashMap::default(),
96            hint_resolves: HashMap::default(),
97            chunks,
98        }
99    }
100
101    pub fn applicable_chunks(&self, ranges: &[Range<Point>]) -> impl Iterator<Item = RowChunk> {
102        self.chunks.applicable_chunks(ranges)
103    }
104
105    pub fn cached_hints(&mut self, chunk: &RowChunk) -> Option<&CacheInlayHints> {
106        self.hints_by_chunks[chunk.id].as_ref()
107    }
108
109    pub fn fetched_hints(&mut self, chunk: &RowChunk) -> &mut Option<CacheInlayHintsTask> {
110        &mut self.fetches_by_chunks[chunk.id]
111    }
112
113    #[cfg(any(test, feature = "test-support"))]
114    pub fn all_cached_hints(&self) -> Vec<InlayHint> {
115        self.hints_by_chunks
116            .iter()
117            .filter_map(|hints| hints.as_ref())
118            .flat_map(|hints| hints.values().cloned())
119            .flatten()
120            .map(|(_, hint)| hint)
121            .collect()
122    }
123
124    #[cfg(any(test, feature = "test-support"))]
125    pub fn all_fetched_hints(&self) -> Vec<CacheInlayHintsTask> {
126        self.fetches_by_chunks
127            .iter()
128            .filter_map(|fetches| fetches.clone())
129            .collect()
130    }
131
132    pub fn remove_server_data(&mut self, for_server: LanguageServerId) {
133        for (chunk_index, hints) in self.hints_by_chunks.iter_mut().enumerate() {
134            if let Some(hints) = hints {
135                if hints.remove(&for_server).is_some() {
136                    self.fetches_by_chunks[chunk_index] = None;
137                }
138            }
139        }
140        self.pending_refreshes.remove(&for_server);
141        self.work_end_refreshes.remove(&for_server);
142        self.fetched_servers.remove(&for_server);
143    }
144
145    fn mark_refresh_pending(&mut self, server_id: LanguageServerId) {
146        self.pending_refreshes.insert(server_id);
147        self.work_end_refreshes.insert(server_id);
148    }
149
150    /// A server finishing a long-running operation may have produced new hints without
151    /// requesting an explicit refresh, but servers like rust-analyzer report work (e.g.
152    /// flycheck) after every save: refresh at most once per server until the buffer
153    /// changes. Explicit refresh requests also count against this limit, as they make a
154    /// subsequent work-end refresh redundant.
155    fn mark_refresh_pending_on_work_end(&mut self, server_id: LanguageServerId) -> bool {
156        if self.work_end_refreshes.insert(server_id) {
157            self.pending_refreshes.insert(server_id);
158            true
159        } else {
160            false
161        }
162    }
163
164    pub fn clear(&mut self) {
165        self.hints_by_chunks = vec![None; self.chunks.len()];
166        self.fetches_by_chunks = vec![None; self.chunks.len()];
167        self.hints_by_id.clear();
168        self.hint_resolves.clear();
169        self.pending_refreshes.clear();
170        self.work_end_refreshes.clear();
171        self.fetched_servers.clear();
172    }
173
174    pub fn insert_new_hints(
175        &mut self,
176        chunk: RowChunk,
177        server_id: LanguageServerId,
178        new_hints: Vec<(InlayId, InlayHint)>,
179    ) {
180        let chunk_hints = self.hints_by_chunks[chunk.id].get_or_insert_default();
181
182        // A response always covers the entire chunk, so it supersedes the server's previously
183        // cached hints for this chunk: a concurrent fetch for the same chunk and server
184        // (e.g. a server refresh arriving mid-fetch) would otherwise append the same hints
185        // again under fresh ids, duplicating them.
186        for (stale_id, _) in chunk_hints.remove(&server_id).into_iter().flatten() {
187            self.hints_by_id.remove(&stale_id);
188            self.hint_resolves.remove(&stale_id);
189        }
190        let mut inserted_hints = Vec::with_capacity(new_hints.len());
191        for (id, new_hint) in new_hints {
192            if let hash_map::Entry::Vacant(vacant_entry) = self.hints_by_id.entry(id) {
193                vacant_entry.insert(HintForId {
194                    chunk_id: chunk.id,
195                    server_id,
196                    position: inserted_hints.len(),
197                });
198                inserted_hints.push((id, new_hint));
199            }
200        }
201        chunk_hints.insert(server_id, inserted_hints);
202        *self.fetched_hints(&chunk) = None;
203    }
204
205    pub fn hint_for_id(&mut self, id: InlayId) -> Option<&mut InlayHint> {
206        let hint_for_id = self.hints_by_id.get(&id)?;
207        let (hint_id, hint) = self
208            .hints_by_chunks
209            .get_mut(hint_for_id.chunk_id)?
210            .as_mut()?
211            .get_mut(&hint_for_id.server_id)?
212            .get_mut(hint_for_id.position)?;
213        debug_assert_eq!(*hint_id, id, "Invalid pointer {hint_for_id:?}");
214        Some(hint)
215    }
216
217    /// Consumes a pending refresh for the given server, if any, invalidating its cached
218    /// hints: only the first query after [`LspStore::refresh_inlay_hints`] invalidates,
219    /// while concurrent queriers share the re-fetch it has started.
220    pub(crate) fn invalidate_for_server_refresh(&mut self, for_server: LanguageServerId) -> bool {
221        if !self.pending_refreshes.remove(&for_server) {
222            return false;
223        }
224
225        for (chunk_id, chunk_data) in self.hints_by_chunks.iter_mut().enumerate() {
226            if let Some(removed_hints) = chunk_data
227                .as_mut()
228                .and_then(|chunk_data| chunk_data.remove(&for_server))
229            {
230                for (id, _) in removed_hints {
231                    self.hints_by_id.remove(&id);
232                    self.hint_resolves.remove(&id);
233                }
234                self.fetches_by_chunks[chunk_id] = None;
235            }
236        }
237        self.fetched_servers.remove(&for_server);
238
239        true
240    }
241
242    pub(crate) fn invalidate_for_chunk(&mut self, chunk: RowChunk) {
243        self.fetches_by_chunks[chunk.id] = None;
244        if let Some(hints_by_server) = self.hints_by_chunks[chunk.id].take() {
245            for (hint_id, _) in hints_by_server.into_values().flatten() {
246                self.hints_by_id.remove(&hint_id);
247                self.hint_resolves.remove(&hint_id);
248            }
249        }
250    }
251}
252
253impl LspStore {
254    pub(super) fn resolve_inlay_hint(
255        &self,
256        mut hint: InlayHint,
257        buffer: Entity<Buffer>,
258        server_id: LanguageServerId,
259        cx: &mut Context<Self>,
260    ) -> Task<anyhow::Result<InlayHint>> {
261        if let Some((upstream_client, project_id)) = self.upstream_client() {
262            if !self.check_if_capable_for_proto_request(&buffer, InlayHints::can_resolve_inlays, cx)
263            {
264                hint.resolve_state = ResolveState::Resolved;
265                return Task::ready(Ok(hint));
266            }
267            let request = proto::ResolveInlayHint {
268                project_id,
269                buffer_id: buffer.read(cx).remote_id().into(),
270                language_server_id: server_id.0 as u64,
271                hint: Some(InlayHints::project_to_proto_hint(hint.clone())),
272            };
273            cx.background_spawn(async move {
274                let response = upstream_client
275                    .request(request)
276                    .await
277                    .context("inlay hints proto request")?;
278                match response.hint {
279                    Some(resolved_hint) => InlayHints::proto_to_project_hint(resolved_hint)
280                        .context("inlay hints proto resolve response conversion"),
281                    None => Ok(hint),
282                }
283            })
284        } else {
285            let Some(lang_server) = buffer.update(cx, |buffer, cx| {
286                self.language_server_for_local_buffer(buffer, server_id, cx)
287                    .map(|(_, server)| server.clone())
288            }) else {
289                return Task::ready(Ok(hint));
290            };
291            if !InlayHints::can_resolve_inlays(&lang_server.capabilities()) {
292                return Task::ready(Ok(hint));
293            }
294            let buffer_snapshot = buffer.read(cx).snapshot();
295            let request_timeout = ProjectSettings::get_global(cx)
296                .global_lsp_settings
297                .get_request_timeout();
298            cx.spawn(async move |_, cx| {
299                let resolve_task = lang_server.request::<lsp::request::InlayHintResolveRequest>(
300                    InlayHints::project_to_lsp_hint(hint, &buffer_snapshot),
301                    request_timeout,
302                );
303                let resolved_hint = resolve_task
304                    .await
305                    .into_response()
306                    .context("inlay hint resolve LSP request")?;
307                let resolved_hint = InlayHints::lsp_to_project_hint(
308                    resolved_hint,
309                    &buffer,
310                    server_id,
311                    ResolveState::Resolved,
312                    false,
313                    cx,
314                )
315                .await?;
316                Ok(resolved_hint)
317            })
318        }
319    }
320
321    /// Marks the server's inlay hints as refresh-pending in every buffer, to be
322    /// invalidated by the next query, and notifies the observers.
323    pub(super) fn refresh_inlay_hints(
324        &mut self,
325        server_id: LanguageServerId,
326        cx: &mut Context<Self>,
327    ) {
328        self.mark_inlay_hints_refresh_pending(server_id, cx);
329        if let Some((downstream_client, project_id)) = self.downstream_client.as_ref() {
330            downstream_client
331                .send(proto::RefreshInlayHints {
332                    project_id: *project_id,
333                    server_id: server_id.to_proto(),
334                    request_id: Some(super::next_wire_refresh_request_id()),
335                })
336                .context("sending refresh inlay hints downstream")
337                .log_err();
338        }
339    }
340
341    pub(super) fn mark_inlay_hints_refresh_pending(
342        &mut self,
343        server_id: LanguageServerId,
344        cx: &mut Context<Self>,
345    ) {
346        for lsp_data in self.lsp_data.values_mut() {
347            lsp_data.inlay_hints.mark_refresh_pending(server_id);
348        }
349        cx.emit(LspStoreEvent::RefreshInlayHints { server_id });
350    }
351
352    /// Same as [`Self::mark_inlay_hints_refresh_pending`], but rate-limited per server
353    /// via [`BufferInlayHints::mark_refresh_pending_on_work_end`].
354    pub(super) fn refresh_inlay_hints_on_work_end(
355        &mut self,
356        server_id: LanguageServerId,
357        cx: &mut Context<Self>,
358    ) {
359        let mut marked = false;
360        for lsp_data in self.lsp_data.values_mut() {
361            marked |= lsp_data
362                .inlay_hints
363                .mark_refresh_pending_on_work_end(server_id);
364        }
365        if marked {
366            cx.emit(LspStoreEvent::RefreshInlayHints { server_id });
367        }
368    }
369
370    pub(super) async fn handle_refresh_inlay_hints(
371        lsp_store: Entity<Self>,
372        envelope: TypedEnvelope<proto::RefreshInlayHints>,
373        mut cx: AsyncApp,
374    ) -> Result<proto::Ack> {
375        lsp_store.update(&mut cx, |lsp_store, cx| {
376            lsp_store
377                .refresh_inlay_hints(LanguageServerId::from_proto(envelope.payload.server_id), cx);
378        });
379        Ok(proto::Ack {})
380    }
381
382    pub(super) async fn handle_resolve_inlay_hint(
383        lsp_store: Entity<Self>,
384        envelope: TypedEnvelope<proto::ResolveInlayHint>,
385        mut cx: AsyncApp,
386    ) -> Result<proto::ResolveInlayHintResponse> {
387        let proto_hint = envelope
388            .payload
389            .hint
390            .expect("incorrect protobuf resolve inlay hint message: missing the inlay hint");
391        let hint = InlayHints::proto_to_project_hint(proto_hint)
392            .context("resolved proto inlay hint conversion")?;
393        let buffer = lsp_store.update(&mut cx, |lsp_store, cx| {
394            let buffer_id = BufferId::new(envelope.payload.buffer_id)?;
395            lsp_store.buffer_store.read(cx).get_existing(buffer_id)
396        })?;
397        let response_hint = lsp_store
398            .update(&mut cx, |lsp_store, cx| {
399                lsp_store.resolve_inlay_hint(
400                    hint,
401                    buffer,
402                    LanguageServerId(envelope.payload.language_server_id as usize),
403                    cx,
404                )
405            })
406            .await
407            .context("inlay hints fetch")?;
408        Ok(proto::ResolveInlayHintResponse {
409            hint: Some(InlayHints::project_to_proto_hint(response_hint)),
410        })
411    }
412}
413