tenant.openagents/omega

1use std::ops::Range;
2use std::sync::Arc;
3use std::time::Duration;
4
5use anyhow::Context as _;
6use collections::{HashMap, HashSet};
7use futures::FutureExt as _;
8use futures::future::{Shared, join_all};
9use gpui::{AppContext as _, AsyncApp, Context, Entity, Task};
10use itertools::Itertools;
11use language::{Buffer, BufferSnapshot, OutlineItem};
12use lsp::LanguageServerId;
13use rpc::{TypedEnvelope, proto};
14use settings::Settings as _;
15use text::{Anchor, Bias, PointUtf16};
16use util::ResultExt;
17
18use crate::DocumentSymbol;
19use crate::lsp_command::{GetDocumentSymbols, LspCommand as _};
20use crate::lsp_store::{
21    LspStore, LspStoreEvent, RunningFetch, missing_servers_to_query, next_lsp_fetch_id,
22    upstream_lsp_query_server_filter,
23};
24use crate::project_settings::ProjectSettings;
25
26pub(super) type DocumentSymbolsTask =
27    Shared<Task<std::result::Result<Vec<OutlineItem<Anchor>>, Arc<anyhow::Error>>>>;
28
29#[derive(Debug, Default)]
30pub(super) struct DocumentSymbolsData {
31    symbols: HashMap<LanguageServerId, Vec<OutlineItem<Anchor>>>,
32    fetched_servers: HashSet<LanguageServerId>,
33    symbols_update: Option<RunningFetch<DocumentSymbolsTask>>,
34}
35
36impl DocumentSymbolsData {
37    pub(super) fn remove_server_data(&mut self, for_server: LanguageServerId) {
38        self.symbols.remove(&for_server);
39        self.fetched_servers.remove(&for_server);
40        RunningFetch::discard_if_queried(&mut self.symbols_update, for_server);
41    }
42
43    fn evict(&mut self, for_server: Option<LanguageServerId>) {
44        match for_server {
45            Some(server_id) => self.remove_server_data(server_id),
46            None => {
47                self.symbols.clear();
48                self.fetched_servers.clear();
49            }
50        }
51        self.symbols_update = None;
52    }
53}
54
55impl LspStore {
56    pub(super) fn refresh_document_symbols(
57        &mut self,
58        for_server: Option<LanguageServerId>,
59        cx: &mut Context<Self>,
60    ) {
61        for lsp_data in self.lsp_data.values_mut() {
62            if let Some(document_symbols) = &mut lsp_data.document_symbols {
63                document_symbols.evict(for_server);
64            }
65        }
66
67        cx.emit(LspStoreEvent::RefreshDocumentSymbols {
68            server_id: for_server,
69        });
70        if let Some((downstream_client, project_id)) = self.downstream_client.as_ref() {
71            downstream_client
72                .send(proto::RefreshDocumentSymbols {
73                    project_id: *project_id,
74                    server_id: for_server.map(|server_id| server_id.to_proto()),
75                })
76                .context("sending refresh document symbols downstream")
77                .log_err();
78        }
79    }
80
81    pub(super) async fn handle_refresh_document_symbols(
82        lsp_store: Entity<Self>,
83        envelope: TypedEnvelope<proto::RefreshDocumentSymbols>,
84        mut cx: AsyncApp,
85    ) -> anyhow::Result<proto::Ack> {
86        lsp_store.update(&mut cx, |lsp_store, cx| {
87            let server_id = envelope.payload.server_id.map(LanguageServerId::from_proto);
88            lsp_store.refresh_document_symbols(server_id, cx);
89        });
90        Ok(proto::Ack {})
91    }
92
93    /// Returns a task that resolves to the document symbol outline items for
94    /// the given buffer.
95    ///
96    /// Caches results per buffer version so repeated calls for the same version
97    /// return immediately. Deduplicates concurrent in-flight requests.
98    ///
99    /// The returned items contain text and ranges but no syntax highlights.
100    /// Callers (e.g. the editor) are responsible for applying highlights
101    /// via the buffer's tree-sitter data and the active theme.
102    pub fn fetch_document_symbols(
103        &mut self,
104        buffer: &Entity<Buffer>,
105        cx: &mut Context<Self>,
106    ) -> Task<Vec<OutlineItem<Anchor>>> {
107        let version_queried_for = buffer.read(cx).version();
108        let buffer_id = buffer.read(cx).remote_id();
109
110        let current_servers = self.relevant_server_ids_for_capability_check(buffer, cx);
111
112        let mut servers_to_query = None;
113        if let Some(lsp_data) = self.current_lsp_data(buffer_id) {
114            if !version_queried_for.changed_since(&lsp_data.buffer_version)
115                && let Some(cached) = &mut lsp_data.document_symbols
116            {
117                match missing_servers_to_query(
118                    &mut cached.symbols,
119                    &mut cached.fetched_servers,
120                    &current_servers,
121                ) {
122                    Some(missing_servers) => servers_to_query = Some(missing_servers),
123                    None => {
124                        let snapshot = buffer.read(cx).snapshot();
125                        return Task::ready(
126                            cached
127                                .symbols
128                                .values()
129                                .flatten()
130                                .unique()
131                                .cloned()
132                                .sorted_by(|a, b| a.range.start.cmp(&b.range.start, &snapshot))
133                                .collect(),
134                        );
135                    }
136                }
137            }
138            if let Some(document_symbols) = &lsp_data.document_symbols
139                && let Some(running) = &document_symbols.symbols_update
140                && !version_queried_for.changed_since(&running.version)
141                && servers_to_query
142                    .as_ref()
143                    .is_none_or(|missing| missing.is_subset(&running.servers))
144            {
145                let running = running.task.clone();
146                return cx
147                    .background_spawn(async move { running.await.log_err().unwrap_or_default() });
148            }
149        }
150
151        let doc_symbols_data = self
152            .latest_lsp_data(buffer, cx)
153            .document_symbols
154            .get_or_insert_default();
155        let fetch_id = next_lsp_fetch_id();
156        let queried_servers = servers_to_query
157            .clone()
158            .unwrap_or_else(|| current_servers.clone());
159        let buffer = buffer.clone();
160        let query_version = version_queried_for.clone();
161        let new_task = cx
162            .spawn({
163                let queried_servers = queried_servers.clone();
164                async move |lsp_store, cx| {
165                    cx.background_executor()
166                        .timer(Duration::from_millis(30))
167                        .await;
168
169                    let fetched = lsp_store
170                        .update(cx, |lsp_store, cx| {
171                            lsp_store.fetch_document_symbols_for_buffer(
172                                &buffer,
173                                servers_to_query,
174                                cx,
175                            )
176                        })
177                        .map_err(Arc::new)?
178                        .await
179                        .context("fetching document symbols")
180                        .map_err(Arc::new);
181
182                    let fetched = match fetched {
183                        Ok(fetched) => fetched,
184                        Err(e) => {
185                            lsp_store
186                                .update(cx, |lsp_store, _| {
187                                    if let Some(lsp_data) = lsp_store.lsp_data.get_mut(&buffer_id)
188                                        && let Some(document_symbols) =
189                                            &mut lsp_data.document_symbols
190                                    {
191                                        RunningFetch::take_finished(
192                                            &mut document_symbols.symbols_update,
193                                            fetch_id,
194                                        );
195                                    }
196                                })
197                                .ok();
198                            return Err(e);
199                        }
200                    };
201
202                    lsp_store
203                        .update(cx, |lsp_store, cx| {
204                            let snapshot = buffer.read(cx).snapshot();
205                            let lsp_data = lsp_store.latest_lsp_data(&buffer, cx);
206                            let doc_symbols = lsp_data.document_symbols.get_or_insert_default();
207
208                            if RunningFetch::take_finished(
209                                &mut doc_symbols.symbols_update,
210                                fetch_id,
211                            ) && let Some(fetched_symbols) = fetched
212                            {
213                                let converted = fetched_symbols
214                                    .iter()
215                                    .map(|(&server_id, symbols)| {
216                                        let mut items = Vec::new();
217                                        flatten_document_symbols(symbols, &snapshot, 0, &mut items);
218                                        (server_id, items)
219                                    })
220                                    .collect();
221                                if lsp_data.buffer_version == query_version {
222                                    doc_symbols.symbols.extend(converted);
223                                    doc_symbols.fetched_servers.extend(queried_servers);
224                                } else if !lsp_data.buffer_version.changed_since(&query_version) {
225                                    lsp_data.buffer_version = query_version;
226                                    doc_symbols.symbols = converted;
227                                    doc_symbols.fetched_servers = queried_servers;
228                                }
229                            }
230                            doc_symbols
231                                .symbols
232                                .values()
233                                .flatten()
234                                .unique()
235                                .cloned()
236                                .sorted_by(|a, b| a.range.start.cmp(&b.range.start, &snapshot))
237                                .collect()
238                        })
239                        .map_err(Arc::new)
240                }
241            })
242            .shared();
243
244        doc_symbols_data.symbols_update = Some(RunningFetch {
245            id: fetch_id,
246            version: version_queried_for,
247            servers: queried_servers,
248            task: new_task.clone(),
249        });
250
251        cx.background_spawn(async move { new_task.await.log_err().unwrap_or_default() })
252    }
253
254    fn fetch_document_symbols_for_buffer(
255        &mut self,
256        buffer: &Entity<Buffer>,
257        for_servers: Option<HashSet<LanguageServerId>>,
258        cx: &mut Context<Self>,
259    ) -> Task<anyhow::Result<Option<HashMap<LanguageServerId, Vec<DocumentSymbol>>>>> {
260        if let Some((client, project_id)) = self.upstream_client() {
261            let request = GetDocumentSymbols;
262            if !self.is_capable_for_proto_request(buffer, &request, cx) {
263                return Task::ready(Ok(None));
264            }
265
266            let request_timeout = ProjectSettings::get_global(cx)
267                .global_lsp_settings
268                .get_request_timeout();
269            let request_task = client.request_lsp(
270                project_id,
271                upstream_lsp_query_server_filter(for_servers.as_ref()),
272                request_timeout,
273                cx.background_executor().clone(),
274                request.to_proto(project_id, buffer.read(cx)),
275            );
276            let buffer = buffer.clone();
277            cx.spawn(async move |weak_lsp_store, cx| {
278                let Some(lsp_store) = weak_lsp_store.upgrade() else {
279                    return Ok(None);
280                };
281                let Some(responses) = request_task.await? else {
282                    return Ok(None);
283                };
284
285                let document_symbols = join_all(responses.payload.into_iter().map(|response| {
286                    let lsp_store = lsp_store.clone();
287                    let buffer = buffer.clone();
288                    let cx = cx.clone();
289                    async move {
290                        (
291                            LanguageServerId::from_proto(response.server_id),
292                            GetDocumentSymbols
293                                .response_from_proto(response.response, lsp_store, buffer, cx)
294                                .await,
295                        )
296                    }
297                }))
298                .await;
299
300                let mut has_errors = false;
301                let result = document_symbols
302                    .into_iter()
303                    .filter_map(|(server_id, symbols)| match symbols {
304                        Ok(symbols) => Some((server_id, symbols)),
305                        Err(e) => {
306                            has_errors = true;
307                            log::error!("Failed to fetch document symbols: {e:#}");
308                            None
309                        }
310                    })
311                    .collect::<HashMap<_, _>>();
312                anyhow::ensure!(
313                    !has_errors || !result.is_empty(),
314                    "Failed to fetch document symbols"
315                );
316                Ok(Some(result))
317            })
318        } else {
319            let symbols_task = self.request_filtered_lsp_locally(
320                buffer,
321                None::<usize>,
322                GetDocumentSymbols,
323                for_servers.as_ref(),
324                cx,
325            );
326            cx.background_spawn(async move { Ok(Some(symbols_task.await.into_iter().collect())) })
327        }
328    }
329}
330
331fn flatten_document_symbols(
332    symbols: &[DocumentSymbol],
333    snapshot: &BufferSnapshot,
334    depth: usize,
335    output: &mut Vec<OutlineItem<Anchor>>,
336) {
337    for symbol in symbols {
338        let name = super::collapse_newlines(&symbol.name, " ");
339
340        let start = snapshot.clip_point_utf16(symbol.range.start, Bias::Right);
341        let end = snapshot.clip_point_utf16(symbol.range.end, Bias::Left);
342        let selection_start = snapshot.clip_point_utf16(symbol.selection_range.start, Bias::Right);
343        let selection_end = snapshot.clip_point_utf16(symbol.selection_range.end, Bias::Left);
344
345        let range = snapshot.anchor_after(start)..snapshot.anchor_before(end);
346        let selection_range =
347            snapshot.anchor_after(selection_start)..snapshot.anchor_before(selection_end);
348
349        let (text, name_ranges, source_range_for_text) =
350            enriched_symbol_text(&name, start, selection_start, selection_end, snapshot)
351                .unwrap_or_else(|| {
352                    let name_len = name.len();
353                    (name.clone(), vec![0..name_len], selection_range.clone())
354                });
355
356        output.push(OutlineItem {
357            depth,
358            range,
359            selection_range: selection_range.clone(),
360            source_range_for_text,
361            text: text.into(),
362            highlight_ranges: Vec::new(),
363            name_ranges,
364            body_range: None,
365            annotation_range: None,
366        });
367
368        if !symbol.children.is_empty() {
369            flatten_document_symbols(&symbol.children, snapshot, depth + 1, output);
370        }
371    }
372}
373
374/// Tries to build an enriched label by including buffer text from the symbol
375/// range start to the selection range end (e.g., "struct Foo" instead of just "Foo").
376/// Only uses same-line prefix to avoid pulling in attributes/decorators.
377fn enriched_symbol_text(
378    name: &str,
379    range_start: PointUtf16,
380    selection_start: PointUtf16,
381    selection_end: PointUtf16,
382    snapshot: &BufferSnapshot,
383) -> Option<(String, Vec<Range<usize>>, Range<Anchor>)> {
384    let text_start = if range_start.row == selection_start.row {
385        range_start
386    } else {
387        PointUtf16::new(selection_start.row, 0)
388    };
389
390    let start_offset = snapshot.point_utf16_to_offset(text_start);
391    let end_offset = snapshot.point_utf16_to_offset(selection_end);
392    if start_offset >= end_offset {
393        return None;
394    }
395
396    let raw: String = snapshot.text_for_range(start_offset..end_offset).collect();
397    let trimmed = raw.trim_start();
398    if trimmed.len() <= name.len() || !trimmed.ends_with(name) {
399        return None;
400    }
401
402    let name_start = trimmed.len() - name.len();
403    let leading_ws = raw.len() - trimmed.len();
404    let adjusted_start = start_offset + leading_ws;
405
406    Some((
407        trimmed.to_string(),
408        vec![name_start..trimmed.len()],
409        snapshot.anchor_after(adjusted_start)..snapshot.anchor_before(end_offset),
410    ))
411}
412
413#[cfg(test)]
414mod tests {
415    use super::*;
416    use gpui::TestAppContext;
417    use language::lsp_to_symbol_kind;
418    use text::{OffsetRangeExt, Point, Unclipped};
419
420    fn make_symbol(
421        name: &str,
422        kind: lsp::SymbolKind,
423        range: std::ops::Range<(u32, u32)>,
424        selection_range: std::ops::Range<(u32, u32)>,
425        children: Vec<DocumentSymbol>,
426    ) -> DocumentSymbol {
427        use text::PointUtf16;
428        DocumentSymbol {
429            name: name.to_string(),
430            kind: lsp_to_symbol_kind(kind),
431            range: Unclipped(PointUtf16::new(range.start.0, range.start.1))
432                ..Unclipped(PointUtf16::new(range.end.0, range.end.1)),
433            selection_range: Unclipped(PointUtf16::new(
434                selection_range.start.0,
435                selection_range.start.1,
436            ))
437                ..Unclipped(PointUtf16::new(
438                    selection_range.end.0,
439                    selection_range.end.1,
440                )),
441            children,
442        }
443    }
444
445    #[gpui::test]
446    async fn test_flatten_document_symbols(cx: &mut TestAppContext) {
447        let buffer = cx.new(|cx| {
448            Buffer::local(
449                concat!(
450                    "struct Foo {\n",
451                    "    bar: u32,\n",
452                    "    baz: String,\n",
453                    "}\n",
454                    "\n",
455                    "impl Foo {\n",
456                    "    fn new() -> Self {\n",
457                    "        Foo { bar: 0, baz: String::new() }\n",
458                    "    }\n",
459                    "}\n",
460                ),
461                cx,
462            )
463        });
464
465        let symbols = vec![
466            make_symbol(
467                "Foo",
468                lsp::SymbolKind::STRUCT,
469                (0, 0)..(3, 1),
470                (0, 7)..(0, 10),
471                vec![
472                    make_symbol(
473                        "bar",
474                        lsp::SymbolKind::FIELD,
475                        (1, 4)..(1, 13),
476                        (1, 4)..(1, 7),
477                        vec![],
478                    ),
479                    make_symbol(
480                        "baz",
481                        lsp::SymbolKind::FIELD,
482                        (2, 4)..(2, 15),
483                        (2, 4)..(2, 7),
484                        vec![],
485                    ),
486                ],
487            ),
488            make_symbol(
489                "Foo",
490                lsp::SymbolKind::STRUCT,
491                (5, 0)..(9, 1),
492                (5, 5)..(5, 8),
493                vec![make_symbol(
494                    "new",
495                    lsp::SymbolKind::FUNCTION,
496                    (6, 4)..(8, 5),
497                    (6, 7)..(6, 10),
498                    vec![],
499                )],
500            ),
501        ];
502
503        let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot());
504
505        let mut items = Vec::new();
506        flatten_document_symbols(&symbols, &snapshot, 0, &mut items);
507
508        assert_eq!(items.len(), 5);
509
510        assert_eq!(items[0].depth, 0);
511        assert_eq!(items[0].text, "struct Foo");
512        assert_eq!(items[0].name_ranges, vec![7..10]);
513        assert_eq!(
514            items[0].selection_range.to_point(&snapshot),
515            Point::new(0, 7)..Point::new(0, 10)
516        );
517
518        assert_eq!(items[1].depth, 1);
519        assert_eq!(items[1].text, "bar");
520        assert_eq!(items[1].name_ranges, vec![0..3]);
521        assert_eq!(
522            items[1].selection_range.to_point(&snapshot),
523            Point::new(1, 4)..Point::new(1, 7)
524        );
525
526        assert_eq!(items[2].depth, 1);
527        assert_eq!(items[2].text, "baz");
528        assert_eq!(items[2].name_ranges, vec![0..3]);
529        assert_eq!(
530            items[2].selection_range.to_point(&snapshot),
531            Point::new(2, 4)..Point::new(2, 7)
532        );
533
534        assert_eq!(items[3].depth, 0);
535        assert_eq!(items[3].text, "impl Foo");
536        assert_eq!(items[3].name_ranges, vec![5..8]);
537        assert_eq!(
538            items[3].selection_range.to_point(&snapshot),
539            Point::new(5, 5)..Point::new(5, 8)
540        );
541
542        assert_eq!(items[4].depth, 1);
543        assert_eq!(items[4].text, "fn new");
544        assert_eq!(items[4].name_ranges, vec![3..6]);
545        assert_eq!(
546            items[4].selection_range.to_point(&snapshot),
547            Point::new(6, 7)..Point::new(6, 10)
548        );
549    }
550
551    #[gpui::test]
552    async fn test_empty_symbols(cx: &mut TestAppContext) {
553        let buffer = cx.new(|cx| Buffer::local("", cx));
554        let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot());
555
556        let symbols: Vec<DocumentSymbol> = Vec::new();
557        let mut items = Vec::new();
558        flatten_document_symbols(&symbols, &snapshot, 0, &mut items);
559        assert!(items.is_empty());
560    }
561
562    #[gpui::test]
563    async fn test_newlines_collapsed_in_name(cx: &mut TestAppContext) {
564        let buffer = cx.new(|cx| Buffer::local("x = 1\ny = 2\n", cx));
565
566        let symbols = vec![
567            make_symbol(
568                "line1\nline2",
569                lsp::SymbolKind::VARIABLE,
570                (0, 0)..(0, 5),
571                (0, 0)..(0, 1),
572                vec![],
573            ),
574            make_symbol(
575                "  a  \n  b  ",
576                lsp::SymbolKind::VARIABLE,
577                (1, 0)..(1, 5),
578                (1, 0)..(1, 1),
579                vec![],
580            ),
581            make_symbol(
582                "a\r\nb",
583                lsp::SymbolKind::VARIABLE,
584                (0, 0)..(1, 5),
585                (0, 0)..(0, 1),
586                vec![],
587            ),
588            make_symbol(
589                "a\n\nb",
590                lsp::SymbolKind::VARIABLE,
591                (0, 0)..(1, 5),
592                (0, 0)..(0, 1),
593                vec![],
594            ),
595        ];
596
597        let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot());
598        let mut items = Vec::new();
599        flatten_document_symbols(&symbols, &snapshot, 0, &mut items);
600
601        assert_eq!(items.len(), 4);
602        assert_eq!(items[0].text, "line1 line2");
603        assert_eq!(items[1].text, "a b");
604        assert_eq!(items[2].text, "a b");
605        assert_eq!(items[3].text, "a b");
606    }
607}
608