mcpls_core/bridge/translator/
enclosing.rs

1//! Opt-in enclosing-symbol context for locations and diagnostics.
2//!
3//! Callers that pass [`ResultContext::EnclosingSymbol`] get, per item, the
4//! innermost symbol of the item's file that contains it. The default
5//! ([`ResultContext::None`]) issues no extra request and adds no field, so
6//! existing output stays byte-identical.
7
8use std::cmp::Reverse;
9use std::collections::HashMap;
10use std::ops::Deref;
11use std::time::Duration;
12
13use schemars::JsonSchema;
14use serde::{Deserialize, Serialize};
15use tokio::time::Instant;
16
17use super::Translator;
18use super::dto::{Diagnostic, Location, Position2D, PositionDegradation, Range, Symbol};
19use super::routing::{FileLocalCapability, IndexingGate};
20use super::symbols::{DocumentSymbolTree, FetchedSymbols, FlatSymbol, symbol_tree};
21use crate::config::DocumentLimit;
22use crate::error::Error;
23use crate::redaction::{Redactions, ServerText};
24
25/// Upper bound on the number of distinct files one call enriches.
26const MAX_ENRICHED_FILES: usize = 16;
27
28/// Overall time budget for the `documentSymbol` lookups of one call, so a
29/// slow server cannot multiply the call's latency by the file cap.
30const ENRICHMENT_DEADLINE: Duration = Duration::from_secs(30);
31
32/// Extra context a tool can attach to each returned item.
33///
34/// Only tools whose input schema declares `context` act on it.
35///
36/// # Examples
37///
38/// ```
39/// use mcpls_core::bridge::ResultContext;
40///
41/// assert_eq!(ResultContext::default(), ResultContext::None);
42/// let parsed: ResultContext = serde_json::from_str("\"enclosing_symbol\"").unwrap();
43/// assert_eq!(parsed, ResultContext::EnclosingSymbol);
44/// ```
45#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema)]
46#[serde(rename_all = "snake_case")]
47pub enum ResultContext {
48    /// Return the plain result with no extra lookups.
49    #[default]
50    None,
51    /// Attach the innermost enclosing symbol (name path, kind, range) to each
52    /// item, at the cost of one `textDocument/documentSymbol` request per
53    /// distinct file, capped per call.
54    EnclosingSymbol,
55}
56
57/// How much ancestry an [`EnclosingSymbol`] name path carries.
58#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
59#[serde(rename_all = "snake_case")]
60pub enum SymbolFidelity {
61    /// Built from the server's `DocumentSymbol` ancestor chain: complete.
62    Hierarchical,
63    /// Built from a flat `SymbolInformation` list: at most one container
64    /// segment from `containerName`, so ancestors may be missing.
65    Flat,
66}
67
68/// The innermost symbol containing an item.
69#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
70pub struct EnclosingSymbol {
71    /// Symbol names from the outermost ancestor down to the symbol itself.
72    pub name_path: Vec<String>,
73    /// LSP numeric symbol kind, the same value `get_document_symbols`
74    /// reports (usable as a `kind_filter`).
75    pub kind: u32,
76    /// Range of the symbol (1-based).
77    pub range: Range,
78    /// How complete `name_path` is.
79    pub fidelity: SymbolFidelity,
80}
81
82/// Why an item's enclosing symbol was not looked up.
83#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
84#[serde(rename_all = "snake_case")]
85pub enum NotComputedReason {
86    /// The call already enriched the maximum number of distinct files.
87    FileCap,
88    /// The file is outside every workspace root, so it is not opened.
89    OutOfWorkspace,
90    /// The document tracker is at its limit and could not open the file.
91    TrackerLimit,
92    /// The call's overall enrichment time budget was spent.
93    Deadline,
94}
95
96/// Why a lookup was attempted but yielded no symbol tree.
97#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
98#[serde(rename_all = "snake_case")]
99pub enum UnavailableReason {
100    /// The routed server does not advertise `documentSymbolProvider`.
101    CapabilityAbsent,
102    /// The `documentSymbol` request or file open failed.
103    RequestFailed,
104    /// The `documentSymbol` request timed out.
105    TimedOut,
106}
107
108/// What is known about the enclosing symbol of one item.
109///
110/// `top_level` means the file's symbols were read and none contains the
111/// item. `not_computed` and `unavailable` mean nothing is known and must
112/// never be read as top level.
113#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
114#[serde(tag = "status", rename_all = "snake_case")]
115pub enum EnclosingSymbolOutcome {
116    /// The innermost containing symbol.
117    Resolved(EnclosingSymbol),
118    /// No symbol of the file contains the item.
119    TopLevel,
120    /// The lookup was skipped.
121    NotComputed {
122        /// Why it was skipped.
123        reason: NotComputedReason,
124    },
125    /// The lookup was attempted and failed.
126    Unavailable {
127        /// Why it failed.
128        reason: UnavailableReason,
129    },
130}
131
132/// Result-level statement about how complete the enrichment is.
133#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
134pub struct EnrichmentSummary {
135    /// Distinct files whose symbol tree was obtained.
136    pub files_enriched: usize,
137    /// Distinct files for which no symbol tree was obtained, for any reason.
138    pub files_skipped: usize,
139    /// Whether the file cap or the time budget skipped some files, so a
140    /// retry with fewer items may enrich more.
141    pub cut_short: bool,
142}
143
144/// An item optionally carrying its enclosing symbol.
145///
146/// Serializes as the item's own fields with `enclosing_symbol` alongside, so
147/// the default output of a tool is unchanged.
148#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
149#[schemars(rename = "Contextual{T}")]
150pub struct Contextual<T> {
151    /// The item itself.
152    #[serde(flatten)]
153    pub inner: T,
154    /// Present only when `context: "enclosing_symbol"` was requested.
155    #[serde(default, skip_serializing_if = "Option::is_none")]
156    #[schemars(with = "EnclosingSymbolOutcome")]
157    pub enclosing_symbol: Option<EnclosingSymbolOutcome>,
158}
159
160/// A [`Location`] optionally carrying its enclosing symbol.
161pub type ContextualLocation = Contextual<Location>;
162
163/// A [`Diagnostic`] optionally carrying its enclosing symbol.
164pub type ContextualDiagnostic = Contextual<Diagnostic>;
165
166impl<T> From<T> for Contextual<T> {
167    fn from(inner: T) -> Self {
168        Self {
169            inner,
170            enclosing_symbol: None,
171        }
172    }
173}
174
175impl<T> Deref for Contextual<T> {
176    type Target = T;
177
178    fn deref(&self) -> &T {
179        &self.inner
180    }
181}
182
183/// Items after the optional enrichment pass.
184pub(super) struct Contextualized<T> {
185    pub(super) items: Vec<Contextual<T>>,
186    pub(super) enrichment: Option<EnrichmentSummary>,
187    pub(super) positions_degraded: Option<PositionDegradation>,
188}
189
190/// Number of distinct files one call may enrich: 16 when the tracker limit is
191/// disabled, else a quarter of the limit clamped to `1..=16`, so one
192/// call cannot evict most of the documents other handlers depend on.
193pub(super) fn enrichment_file_cap(max_documents: DocumentLimit) -> usize {
194    max_documents.get().map_or(MAX_ENRICHED_FILES, |max| {
195        (max.get() / 4).clamp(1, MAX_ENRICHED_FILES)
196    })
197}
198
199/// How strictly a symbol must contain an item to enclose it.
200#[derive(Clone, Copy)]
201enum Containment {
202    /// The symbol contains the whole item range.
203    Whole,
204    /// The symbol contains the item's start; used when none contains it whole.
205    Start,
206}
207
208const fn key(position: &Position2D) -> (u32, u32) {
209    (position.line, position.character)
210}
211
212impl Containment {
213    fn admits(self, symbol: &Range, item: &Range) -> bool {
214        let starts_inside = key(&symbol.start) <= key(&item.start);
215        match self {
216            Self::Whole => starts_inside && key(&item.end) <= key(&symbol.end),
217            Self::Start => starts_inside && key(&item.start) <= key(&symbol.end),
218        }
219    }
220}
221
222const CONTAINMENT_RULES: [Containment; 2] = [Containment::Whole, Containment::Start];
223
224pub(super) fn innermost_hierarchical(symbols: &[Symbol], item: &Range) -> Option<EnclosingSymbol> {
225    CONTAINMENT_RULES
226        .into_iter()
227        .find_map(|rule| descend(symbols, item, rule))
228}
229
230fn descend(symbols: &[Symbol], item: &Range, rule: Containment) -> Option<EnclosingSymbol> {
231    let mut name_path = Vec::new();
232    let mut level = symbols;
233    let mut found: Option<&Symbol> = None;
234    while let Some(symbol) = level
235        .iter()
236        .filter(|symbol| rule.admits(&symbol.range, item))
237        .max_by_key(|symbol| {
238            (
239                key(&symbol.range.start),
240                Reverse(key(&symbol.range.end)),
241                Reverse(symbol.name.as_str()),
242                Reverse(symbol.kind),
243            )
244        })
245    {
246        name_path.push(symbol.name.clone());
247        found = Some(symbol);
248        level = symbol.children.as_deref().unwrap_or_default();
249    }
250    found.map(|symbol| EnclosingSymbol {
251        name_path,
252        kind: symbol.kind,
253        range: symbol.range.clone(),
254        fidelity: SymbolFidelity::Hierarchical,
255    })
256}
257
258pub(super) fn innermost_flat(symbols: &[FlatSymbol], item: &Range) -> Option<EnclosingSymbol> {
259    CONTAINMENT_RULES.into_iter().find_map(|rule| {
260        symbols
261            .iter()
262            .filter(|flat| rule.admits(&flat.symbol.range, item))
263            .max_by_key(|flat| {
264                let symbol = &flat.symbol;
265                (
266                    key(&symbol.range.start),
267                    Reverse(key(&symbol.range.end)),
268                    Reverse(symbol.name.as_str()),
269                    Reverse(symbol.kind),
270                )
271            })
272            .map(|flat| EnclosingSymbol {
273                name_path: flat
274                    .container_name
275                    .iter()
276                    .filter(|container| !container.is_empty())
277                    .cloned()
278                    .chain([flat.symbol.name.clone()])
279                    .collect(),
280                kind: flat.symbol.kind,
281                range: flat.symbol.range.clone(),
282                fidelity: SymbolFidelity::Flat,
283            })
284    })
285}
286
287fn outcome_for(tree: &DocumentSymbolTree, item: &Range) -> EnclosingSymbolOutcome {
288    let found = match tree {
289        DocumentSymbolTree::Hierarchical(symbols) => innermost_hierarchical(symbols, item),
290        DocumentSymbolTree::Flat(symbols) => innermost_flat(symbols, item),
291    };
292    found.map_or(
293        EnclosingSymbolOutcome::TopLevel,
294        EnclosingSymbolOutcome::Resolved,
295    )
296}
297
298/// An item to resolve: the URI string of its file and its normalized range.
299struct Hit<'a> {
300    uri: &'a str,
301    range: &'a Range,
302}
303
304/// What one file contributed to a call's enrichment.
305enum FileResolution {
306    Tree(DocumentSymbolTree),
307    Skipped(EnclosingSymbolOutcome),
308}
309
310/// Outcomes for every hit, in input order, plus the call-level summary.
311struct Resolution {
312    outcomes: Vec<EnclosingSymbolOutcome>,
313    summary: EnrichmentSummary,
314    positions_degraded: Option<PositionDegradation>,
315}
316
317const fn skipped_outcome_for(error: &Error) -> EnclosingSymbolOutcome {
318    match error {
319        Error::CapabilityNotSupported { .. } => EnclosingSymbolOutcome::Unavailable {
320            reason: UnavailableReason::CapabilityAbsent,
321        },
322        Error::DocumentLimitExceeded { .. } => EnclosingSymbolOutcome::NotComputed {
323            reason: NotComputedReason::TrackerLimit,
324        },
325        Error::Timeout(_) => EnclosingSymbolOutcome::Unavailable {
326            reason: UnavailableReason::TimedOut,
327        },
328        _ => EnclosingSymbolOutcome::Unavailable {
329            reason: UnavailableReason::RequestFailed,
330        },
331    }
332}
333
334const fn not_computed(reason: NotComputedReason) -> FileResolution {
335    FileResolution::Skipped(EnclosingSymbolOutcome::NotComputed { reason })
336}
337
338impl Translator {
339    /// Attaches enclosing symbols to `locations` when `context` asks for it.
340    pub(super) async fn contextualize_locations(
341        &self,
342        locations: Vec<Location>,
343        context: ResultContext,
344        positions_degraded: Option<PositionDegradation>,
345    ) -> Contextualized<Location> {
346        self.contextualize_items(
347            locations,
348            &(),
349            context,
350            positions_degraded,
351            |(), location| Hit {
352                uri: &location.uri,
353                range: &location.range,
354            },
355        )
356        .await
357    }
358
359    /// Attaches enclosing symbols to the diagnostics of the file `uri` when
360    /// `context` asks for it.
361    pub(super) async fn contextualize_diagnostics(
362        &self,
363        uri: &str,
364        diagnostics: Vec<Diagnostic>,
365        context: ResultContext,
366        positions_degraded: Option<PositionDegradation>,
367    ) -> Contextualized<Diagnostic> {
368        self.contextualize_items(
369            diagnostics,
370            uri,
371            context,
372            positions_degraded,
373            |uri, diagnostic| Hit {
374                uri,
375                range: &diagnostic.range,
376            },
377        )
378        .await
379    }
380
381    /// `shared` is data every hit borrows from besides its own item, such as
382    /// the one URI of a file's diagnostics.
383    async fn contextualize_items<T: Send, S: Sync + ?Sized>(
384        &self,
385        items: Vec<T>,
386        shared: &S,
387        context: ResultContext,
388        positions_degraded: Option<PositionDegradation>,
389        hit: impl for<'a> Fn(&'a S, &'a T) -> Hit<'a> + Sync,
390    ) -> Contextualized<T> {
391        if context == ResultContext::None || items.is_empty() {
392            return Contextualized {
393                items: items.into_iter().map(Contextual::from).collect(),
394                enrichment: None,
395                positions_degraded,
396            };
397        }
398        let hits: Vec<Hit<'_>> = items.iter().map(|item| hit(shared, item)).collect();
399        let resolution = self.resolve_enclosing(&hits).await;
400        let items = items
401            .into_iter()
402            .zip(resolution.outcomes)
403            .map(|(inner, outcome)| Contextual {
404                inner,
405                enclosing_symbol: Some(outcome),
406            })
407            .collect();
408        Contextualized {
409            items,
410            enrichment: Some(resolution.summary),
411            positions_degraded: positions_degraded.max(resolution.positions_degraded),
412        }
413    }
414
415    async fn resolve_enclosing(&self, hits: &[Hit<'_>]) -> Resolution {
416        let cap = enrichment_file_cap(self.document_tracker.limits().max_documents);
417        let started = Instant::now();
418        let mut attempted = 0_usize;
419        let mut positions_degraded = None;
420        let mut files: HashMap<&str, FileResolution> = HashMap::new();
421
422        for hit in hits {
423            if files.contains_key(hit.uri) {
424                continue;
425            }
426            let resolution = self
427                .resolve_file(
428                    hit.uri,
429                    cap,
430                    &mut attempted,
431                    started,
432                    &mut positions_degraded,
433                )
434                .await;
435            files.insert(hit.uri, resolution);
436        }
437
438        let outcomes = hits
439            .iter()
440            .map(|hit| match files.get(hit.uri) {
441                Some(FileResolution::Tree(tree)) => outcome_for(tree, hit.range),
442                Some(FileResolution::Skipped(outcome)) => outcome.clone(),
443                None => EnclosingSymbolOutcome::Unavailable {
444                    reason: UnavailableReason::RequestFailed,
445                },
446            })
447            .collect();
448
449        let files_enriched = files
450            .values()
451            .filter(|file| matches!(file, FileResolution::Tree(_)))
452            .count();
453        let cut_short = files.values().any(|file| {
454            matches!(
455                file,
456                FileResolution::Skipped(EnclosingSymbolOutcome::NotComputed {
457                    reason: NotComputedReason::FileCap | NotComputedReason::Deadline
458                })
459            )
460        });
461        Resolution {
462            outcomes,
463            summary: EnrichmentSummary {
464                files_enriched,
465                files_skipped: files.len().saturating_sub(files_enriched),
466                cut_short,
467            },
468            positions_degraded,
469        }
470    }
471
472    /// Resolves one file's symbol tree. The URI is the server's text, so it
473    /// is validated against the workspace roots before the file is opened.
474    async fn resolve_file(
475        &self,
476        uri: &str,
477        cap: usize,
478        attempted: &mut usize,
479        started: Instant,
480        positions_degraded: &mut Option<PositionDegradation>,
481    ) -> FileResolution {
482        if *attempted >= cap {
483            return not_computed(NotComputedReason::FileCap);
484        }
485        let remaining = ENRICHMENT_DEADLINE.saturating_sub(started.elapsed());
486        if remaining.is_zero() {
487            return not_computed(NotComputedReason::Deadline);
488        }
489        let Ok(path) = self.parse_file_uri(&lsp_types::Uri::from(uri)).await else {
490            return not_computed(NotComputedReason::OutOfWorkspace);
491        };
492        *attempted = attempted.saturating_add(1);
493
494        let lookup = async {
495            let doc = self
496                .prepare_gated_document_for_path(
497                    &path,
498                    IndexingGate::FileLocal(FileLocalCapability::DocumentSymbols),
499                )
500                .await?;
501            let FetchedSymbols { doc, ctx, response } = self.fetch_document_symbols(doc).await?;
502            let tree = symbol_tree(response, &ctx, doc.uri()).await;
503            Ok((tree, ctx.positions_degraded()))
504        };
505        match tokio::time::timeout(remaining, lookup).await {
506            Err(_elapsed) => FileResolution::Skipped(EnclosingSymbolOutcome::Unavailable {
507                reason: UnavailableReason::TimedOut,
508            }),
509            Ok(Err(error)) => FileResolution::Skipped(skipped_outcome_for(&error)),
510            Ok(Ok((tree, degraded))) => {
511                *positions_degraded = (*positions_degraded).max(degraded);
512                FileResolution::Tree(tree)
513            }
514        }
515    }
516}
517
518impl ServerText for EnclosingSymbol {
519    fn redact_server_text(&mut self, redactions: &Redactions) {
520        let Self {
521            name_path,
522            kind: _,
523            range: _,
524            fidelity: _,
525        } = self;
526        for segment in name_path {
527            redactions.note_payload(segment);
528        }
529    }
530}
531
532impl ServerText for EnclosingSymbolOutcome {
533    fn redact_server_text(&mut self, redactions: &Redactions) {
534        match self {
535            Self::Resolved(symbol) => symbol.redact_server_text(redactions),
536            Self::TopLevel | Self::NotComputed { reason: _ } | Self::Unavailable { reason: _ } => {}
537        }
538    }
539}
540
541impl<T: ServerText> ServerText for Contextual<T> {
542    fn redact_server_text(&mut self, redactions: &Redactions) {
543        let Self {
544            inner,
545            enclosing_symbol,
546        } = self;
547        inner.redact_server_text(redactions);
548        enclosing_symbol.redact_server_text(redactions);
549    }
550}
551
552#[cfg(test)]
553mod contextual_tests {
554    use super::*;
555
556    #[test]
557    fn test_contextual_serializes_flat_for_any_inner_type() {
558        let item = Contextual {
559            inner: serde_json::json!({"name": "f"}),
560            enclosing_symbol: Some(EnclosingSymbolOutcome::TopLevel),
561        };
562        let value = serde_json::to_value(&item).unwrap();
563        assert_eq!(value["name"], "f");
564        assert_eq!(value["enclosing_symbol"]["status"], "top_level");
565
566        let plain = Contextual::from(serde_json::json!({"name": "f"}));
567        assert!(
568            serde_json::to_value(&plain)
569                .unwrap()
570                .get("enclosing_symbol")
571                .is_none()
572        );
573    }
574}
575
576#[cfg(test)]
577mod tests {
578    use super::*;
579
580    #[test]
581    fn test_enrichment_file_cap() {
582        assert_eq!(enrichment_file_cap(DocumentLimit::UNLIMITED), 16);
583        assert_eq!(enrichment_file_cap(DocumentLimit::new(1)), 1);
584        assert_eq!(enrichment_file_cap(DocumentLimit::new(3)), 1);
585        assert_eq!(enrichment_file_cap(DocumentLimit::new(8)), 2);
586        assert_eq!(enrichment_file_cap(DocumentLimit::new(64)), 16);
587        assert_eq!(enrichment_file_cap(DocumentLimit::new(10_000)), 16);
588    }
589}
590
591#[cfg(test)]
592mod server_text_tests {
593    use super::*;
594
595    #[test]
596    fn test_enclosing_symbol_name_path_passes_through() {
597        let secret = "SuperSecretValue123";
598        let set = Redactions::new([("API_TOKEN".to_owned(), secret.to_owned())]);
599        let mut outcome = EnclosingSymbolOutcome::Resolved(EnclosingSymbol {
600            name_path: vec![format!("mod_{secret}"), "f".to_owned()],
601            kind: 12,
602            range: Range {
603                start: Position2D {
604                    line: 1,
605                    character: 1,
606                },
607                end: Position2D {
608                    line: 1,
609                    character: 2,
610                },
611            },
612            fidelity: SymbolFidelity::Hierarchical,
613        });
614        let before = outcome.clone();
615
616        outcome.redact_server_text(&set);
617
618        assert_eq!(outcome, before);
619    }
620}
mcpls_core/bridge/translator/enclosing.rs

mcpls_core/bridge/translator/
enclosing.rs
