Skip to main content

flymcprust_core/
completion.rs

1//! Argument completion (`completion/complete`) for prompt arguments and
2//! resource-template variables.
3//!
4//! With the macros, annotate a parameter with `#[complete(path)]` (a function
5//! `fn(&str) -> impl IntoIterator<Item = impl Into<String>>` that receives the
6//! partial value) or `#[complete(values = ["a", "b"])]` (a fixed list filtered
7//! by prefix). By hand, attach a [`CompleteFn`] with
8//! [`PromptDef::complete`](crate::PromptDef::complete) or
9//! [`ResourceDef::complete`](crate::ResourceDef::complete).
10
11use std::collections::BTreeMap;
12
13use crate::spec;
14
15/// Maximum number of values in one completion result (specification limit).
16pub const MAX_COMPLETION_VALUES: usize = 100;
17
18/// A completion provider. Completion runs synchronously and must be fast.
19pub type CompleteFn = fn(&CompletionRequest<'_>) -> Completion;
20
21/// What the client asked to complete.
22#[derive(Debug, Clone, Copy)]
23pub struct CompletionRequest<'a> {
24    pub(crate) argument: &'a str,
25    pub(crate) value: &'a str,
26    pub(crate) context: &'a BTreeMap<String, String>,
27}
28
29impl<'a> CompletionRequest<'a> {
30    /// The argument (or template variable) being completed.
31    #[must_use]
32    pub fn argument(&self) -> &'a str {
33        self.argument
34    }
35
36    /// The partial value typed so far.
37    #[must_use]
38    pub fn value(&self) -> &'a str {
39        self.value
40    }
41
42    /// Values already chosen for other arguments, when the client sent them.
43    #[must_use]
44    pub fn context(&self) -> &'a BTreeMap<String, String> {
45        self.context
46    }
47}
48
49/// Completion values. Results longer than [`MAX_COMPLETION_VALUES`] are
50/// truncated, with `total` and `hasMore` set accordingly.
51#[derive(Debug, Clone, Default, PartialEq, Eq)]
52pub struct Completion {
53    values: Vec<String>,
54    total: Option<i64>,
55    has_more: Option<bool>,
56}
57
58impl Completion {
59    /// Completion from a list of values.
60    pub fn new<I>(values: I) -> Self
61    where
62        I: IntoIterator,
63        I::Item: Into<String>,
64    {
65        Self {
66            values: values.into_iter().map(Into::into).collect(),
67            total: None,
68            has_more: None,
69        }
70    }
71
72    /// The values of `candidates` that start with `prefix`, ignoring ASCII
73    /// case.
74    #[must_use]
75    pub fn filter_prefix(candidates: &[&str], prefix: &str) -> Self {
76        Self::new(
77            candidates
78                .iter()
79                .filter(|c| {
80                    c.get(..prefix.len())
81                        .is_some_and(|head| head.eq_ignore_ascii_case(prefix))
82                })
83                .copied(),
84        )
85    }
86
87    /// Sets the total number of matches, when known.
88    #[must_use]
89    pub fn with_total(mut self, total: u64) -> Self {
90        self.total = Some(i64::try_from(total).unwrap_or(i64::MAX));
91        self
92    }
93
94    /// Declares that more matches exist than the values returned.
95    #[must_use]
96    pub fn with_more(mut self, has_more: bool) -> Self {
97        self.has_more = Some(has_more);
98        self
99    }
100
101    pub(crate) fn into_spec(mut self) -> spec::CompleteResultCompletion {
102        if self.values.len() > MAX_COMPLETION_VALUES {
103            let total = self.values.len();
104            self.values.truncate(MAX_COMPLETION_VALUES);
105            self.total = Some(
106                self.total
107                    .unwrap_or_else(|| i64::try_from(total).unwrap_or(i64::MAX)),
108            );
109            self.has_more = Some(true);
110        }
111        spec::CompleteResultCompletion {
112            has_more: self.has_more,
113            total: self.total,
114            values: self.values,
115        }
116    }
117}
118
119impl<S: Into<String>> FromIterator<S> for Completion {
120    fn from_iter<I: IntoIterator<Item = S>>(iter: I) -> Self {
121        Self::new(iter)
122    }
123}
124
125/// A completion provider attached to one argument.
126#[derive(Debug, Clone)]
127pub(crate) struct Completer {
128    pub(crate) argument: String,
129    pub(crate) complete: CompleteFn,
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    #[test]
137    fn prefix_filter_ignores_case() {
138        let c = Completion::filter_prefix(&["Rust", "Ruby", "Python"], "ru");
139        assert_eq!(c.into_spec().values, ["Rust", "Ruby"]);
140    }
141
142    #[test]
143    fn long_results_are_truncated() {
144        let c = Completion::new((0..150).map(|i| i.to_string())).into_spec();
145        assert_eq!(c.values.len(), 100);
146        assert_eq!(c.total, Some(150));
147        assert_eq!(c.has_more, Some(true));
148    }
149}