prism_mcp_rs/core/
completion.rs

1//! Completion handling for MCP servers
2//!
3//! This module provides the foundation for implementing autocompletion features
4//! in MCP servers, supporting the completion API introduced in MCP 2025-11-25.
5
6use crate::core::error::McpResult;
7use crate::protocol::messages::{CompletionArgument, CompletionReference};
8use crate::protocol::types::*;
9use async_trait::async_trait;
10use std::collections::HashMap;
11
12/// Trait for handling completion requests
13///
14/// Implement Trait to provide smart autocompletion for prompts,
15/// tools, and resources in your MCP server.
16#[async_trait]
17pub trait CompletionHandler: Send + Sync {
18    /// Generate completion suggestions
19    ///
20    /// Method is called when a client requests autocompletion for an
21    /// argument value. The implementation should analyze the context and
22    /// return relevant completion suggestions.
23    ///
24    /// # Arguments
25    /// * `reference` - The item being completed (prompt, tool, or resource)
26    /// * `argument` - The argument being completed with current value
27    /// * `context` - Optional additional context for completion
28    ///
29    /// # Returns
30    /// List of completion suggestions
31    async fn complete(
32        &self,
33        reference: &CompletionReference,
34        argument: &CompletionArgument,
35        context: Option<&CompletionContext>,
36    ) -> McpResult<Vec<String>>;
37}
38
39/// Default prompt completion handler with fuzzy matching
40///
41/// This handler provides basic completion for prompt names by matching
42/// against a predefined list of available prompts.
43///
44/// # Example
45/// ```rust
46/// use prism_mcp_rs::core::completion::PromptCompletionHandler;
47///
48/// let prompts = vec![
49/// "analyze_data".to_string(),
50/// "analyze_text".to_string(),
51/// "create_report".to_string(),
52/// ];
53/// let handler = PromptCompletionHandler::new(prompts);
54/// ```
55pub struct PromptCompletionHandler {
56    prompts: Vec<String>,
57    argument_completions: HashMap<String, HashMap<String, Vec<String>>>,
58}
59
60impl PromptCompletionHandler {
61    /// Create a new prompt completion handler
62    ///
63    /// # Arguments
64    /// * `prompts` - List of available prompt names
65    pub fn new(prompts: Vec<String>) -> Self {
66        Self {
67            prompts,
68            argument_completions: HashMap::new(),
69        }
70    }
71
72    /// Add a prompt to the completion list
73    pub fn add_prompt<S: Into<String>>(&mut self, name: S) {
74        self.prompts.push(name.into());
75    }
76
77    /// Remove a prompt from the completion list
78    pub fn remove_prompt(&mut self, name: &str) {
79        self.prompts.retain(|p| p != name);
80    }
81
82    /// Get all available prompts
83    pub fn prompts(&self) -> &[String] {
84        &self.prompts
85    }
86
87    /// Add argument completions for a specific prompt
88    pub fn add_argument_completions(
89        &mut self,
90        prompt_name: &str,
91        argument_name: &str,
92        completions: Vec<String>,
93    ) {
94        let prompt_args = self
95            .argument_completions
96            .entry(prompt_name.to_string())
97            .or_default();
98        prompt_args.insert(argument_name.to_string(), completions);
99    }
100
101    /// Perform fuzzy matching on a list of strings
102    pub fn fuzzy_match(&self, items: &[String], query: &str) -> Vec<String> {
103        let mut matches: Vec<(String, f32)> = items
104            .iter()
105            .filter_map(|item| {
106                let score = self.calculate_match_score(item, query);
107                if score > 0.0 {
108                    Some((item.clone(), score))
109                } else {
110                    None
111                }
112            })
113            .collect();
114
115        // Sort by score (descending)
116        matches.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
117
118        matches.into_iter().map(|(item, _)| item).collect()
119    }
120
121    /// Calculate fuzzy match score
122    fn calculate_match_score(&self, item: &str, query: &str) -> f32 {
123        if item == query {
124            return 1.0;
125        }
126        if item.starts_with(query) {
127            return 0.8;
128        }
129        if item.contains(query) {
130            return 0.6;
131        }
132        0.0
133    }
134
135    /// Get supported reference types
136    pub fn supported_reference_types(&self) -> Vec<&str> {
137        vec!["ref/prompt"]
138    }
139}
140
141#[async_trait]
142impl CompletionHandler for PromptCompletionHandler {
143    async fn complete(
144        &self,
145        reference: &CompletionReference,
146        argument: &CompletionArgument,
147        _context: Option<&CompletionContext>,
148    ) -> McpResult<Vec<String>> {
149        match reference {
150            CompletionReference::Prompt { name } => {
151                if argument.name == "name" {
152                    // Return prompt names that start with the current value
153                    Ok(self
154                        .prompts
155                        .iter()
156                        .filter(|prompt_name| prompt_name.starts_with(&argument.value))
157                        .take(10) // Limit to 10 suggestions
158                        .cloned()
159                        .collect())
160                } else {
161                    // Check for argument completions
162                    if let Some(prompt_args) = self.argument_completions.get(name) {
163                        if let Some(values) = prompt_args.get(&argument.name) {
164                            return Ok(values
165                                .iter()
166                                .filter(|value| value.starts_with(&argument.value))
167                                .take(10)
168                                .cloned()
169                                .collect());
170                        }
171                    }
172                    Ok(vec![])
173                }
174            }
175            _ => Ok(vec![]), // Only handle prompt references
176        }
177    }
178}
179
180/// Resource URI completion handler
181///
182/// This handler provides completion for resource URIs based on resource
183/// templates. It can suggest URI completions by analyzing URI patterns
184/// and available resources.
185///
186/// # Example
187/// ```rust
188/// use prism_mcp_rs::core::completion::ResourceUriCompletionHandler;
189/// use prism_mcp_rs::protocol::types::ResourceTemplate;
190///
191/// let templates = vec![
192/// ResourceTemplate::new("file:///docs/{category}/{filename}".to_string(), "Documentation".to_string()),
193/// ResourceTemplate::new("db:/// {table}/{id}".to_string(), "Database Records".to_string()),
194/// ];
195/// let handler = ResourceUriCompletionHandler::new(templates);
196/// ```
197pub struct ResourceUriCompletionHandler {
198    templates: Vec<ResourceTemplate>,
199    static_resources: Vec<String>,
200}
201
202impl ResourceUriCompletionHandler {
203    /// Create a new resource URI completion handler
204    ///
205    /// # Arguments
206    /// * `templates` - List of resource templates to use for completion
207    pub fn new(templates: Vec<ResourceTemplate>) -> Self {
208        Self {
209            templates,
210            static_resources: Vec::new(),
211        }
212    }
213
214    /// Add a static resource URI for completion
215    ///
216    /// # Arguments
217    /// * `uri` - Static resource URI to add to completions
218    pub fn add_static_resource<S: Into<String>>(&mut self, uri: S) {
219        self.static_resources.push(uri.into());
220    }
221
222    /// Generate URI completions based on templates
223    async fn generate_uri_completions(
224        &self,
225        uri_template: &str,
226        current_value: &str,
227        _context: Option<&CompletionContext>,
228    ) -> McpResult<Vec<String>> {
229        let mut completions = Vec::new();
230
231        // Simple template variable completion
232        if uri_template.contains("{category}") && current_value.contains("/docs/") {
233            // Example categories for documentation
234            let categories = ["api", "guides", "tutorials", "reference"];
235            for category in categories {
236                let suggestion = uri_template.replace("{category}", category);
237                if suggestion.starts_with(current_value) {
238                    completions.push(suggestion);
239                }
240            }
241        }
242
243        if uri_template.contains("{filename}") {
244            // Example filenames
245            let filenames = [
246                "overview.md",
247                "getting-started.md",
248                "reference.md",
249                "examples.md",
250            ];
251            for filename in filenames {
252                let suggestion = uri_template.replace("{filename}", filename);
253                if suggestion.starts_with(current_value) {
254                    completions.push(suggestion);
255                }
256            }
257        }
258
259        if uri_template.contains("{table}") {
260            // Example database tables
261            let tables = ["users", "products", "orders", "categories"];
262            for table in tables {
263                let suggestion = uri_template.replace("{table}", table);
264                if suggestion.starts_with(current_value) {
265                    completions.push(suggestion);
266                }
267            }
268        }
269
270        if uri_template.contains("{id}") {
271            // Example IDs
272            let ids = ["1", "2", "3", "latest", "featured"];
273            for id in ids {
274                let suggestion = uri_template.replace("{id}", id);
275                if suggestion.starts_with(current_value) {
276                    completions.push(suggestion);
277                }
278            }
279        }
280
281        Ok(completions)
282    }
283
284    /// Extract template variables from a URI template
285    pub fn extract_template_variables(&self, template: &str) -> Vec<String> {
286        let mut variables = Vec::new();
287        let mut chars = template.chars().peekable();
288
289        while let Some(ch) = chars.next() {
290            if ch == '{' {
291                let mut var_name = String::new();
292                for ch in chars.by_ref() {
293                    if ch == '}' {
294                        break;
295                    }
296                    var_name.push(ch);
297                }
298                if !var_name.is_empty() {
299                    variables.push(var_name);
300                }
301            }
302        }
303
304        variables
305    }
306
307    /// Get supported reference types
308    pub fn supported_reference_types(&self) -> Vec<&str> {
309        vec!["ref/resource"]
310    }
311}
312
313#[async_trait]
314impl CompletionHandler for ResourceUriCompletionHandler {
315    async fn complete(
316        &self,
317        reference: &CompletionReference,
318        argument: &CompletionArgument,
319        context: Option<&CompletionContext>,
320    ) -> McpResult<Vec<String>> {
321        match reference {
322            CompletionReference::Resource { uri: _ } => {
323                // Match against URI templates and suggest completions
324                let mut suggestions = Vec::new();
325
326                for template in &self.templates {
327                    if let Ok(completions) = self
328                        .generate_uri_completions(&template.uri_template, &argument.value, context)
329                        .await
330                    {
331                        suggestions.extend(completions);
332                    }
333                }
334
335                // Add static resource completions
336                for resource_uri in &self.static_resources {
337                    if resource_uri.starts_with(&argument.value) {
338                        suggestions.push(resource_uri.clone());
339                    }
340                }
341
342                Ok(suggestions.into_iter().take(10).collect())
343            }
344            _ => Ok(vec![]), // Only handle resource references
345        }
346    }
347}
348
349/// Tool argument completion handler
350///
351/// This handler provides completion for tool arguments based on the tool's
352/// input schema and available values.
353///
354/// # Example
355/// ```rust
356/// use prism_mcp_rs::core::completion::ToolCompletionHandler;
357/// use std::collections::HashMap;
358///
359/// let mut completions = HashMap::new();
360/// completions.insert("file_reader".to_string(), vec![
361/// ("path".to_string(), vec!["/home/user/file1.txt".to_string(), "/home/user/file2.txt".to_string()]),
362/// ]);
363/// let handler = ToolCompletionHandler::new(completions);
364/// ```
365pub struct ToolCompletionHandler {
366    tool_completions: HashMap<String, Vec<(String, Vec<String>)>>,
367}
368
369impl ToolCompletionHandler {
370    /// Create a new tool completion handler
371    ///
372    /// # Arguments
373    /// * `tool_completions` - Map of tool names to argument completions
374    pub fn new(tool_completions: HashMap<String, Vec<(String, Vec<String>)>>) -> Self {
375        Self { tool_completions }
376    }
377
378    /// Add completion values for a tool argument
379    ///
380    /// # Arguments
381    /// * `tool_name` - Name of the tool
382    /// * `argument_name` - Name of the argument
383    /// * `values` - List of possible values for completion
384    pub fn add_tool_argument_completions<S: Into<String>>(
385        &mut self,
386        tool_name: S,
387        argument_name: S,
388        values: Vec<String>,
389    ) {
390        let tool_name = tool_name.into();
391        let argument_name = argument_name.into();
392
393        let tool_entry = self.tool_completions.entry(tool_name).or_default();
394
395        // Remove existing entry for this argument if it exists
396        tool_entry.retain(|(name, _)| name != &argument_name);
397
398        // Add new entry
399        tool_entry.push((argument_name, values));
400    }
401
402    /// Remove completions for a tool
403    pub fn remove_tool_completions(&mut self, tool_name: &str) {
404        self.tool_completions.remove(tool_name);
405    }
406
407    /// Get supported reference types
408    pub fn supported_reference_types(&self) -> Vec<&str> {
409        vec!["ref/tool"]
410    }
411}
412
413#[async_trait]
414impl CompletionHandler for ToolCompletionHandler {
415    async fn complete(
416        &self,
417        reference: &CompletionReference,
418        argument: &CompletionArgument,
419        _context: Option<&CompletionContext>,
420    ) -> McpResult<Vec<String>> {
421        match reference {
422            CompletionReference::Tool { name } => {
423                if let Some(tool_args) = self.tool_completions.get(name) {
424                    for (arg_name, values) in tool_args {
425                        if arg_name == &argument.name {
426                            // Return values that start with the current value
427                            return Ok(values
428                                .iter()
429                                .filter(|value| value.starts_with(&argument.value))
430                                .take(10)
431                                .cloned()
432                                .collect());
433                        }
434                    }
435                }
436                Ok(vec![])
437            }
438            _ => Ok(vec![]), // Only handle tool references
439        }
440    }
441}
442
443/// Combined completion handler that delegates to different handlers based on reference type
444///
445/// This handler allows you to combine multiple completion handlers and route
446/// completion requests to the appropriate handler based on the reference type.
447///
448/// # Example
449/// ```rust
450/// use prism_mcp_rs::core::completion::*;
451/// use std::collections::HashMap;
452///
453/// let prompt_handler = PromptCompletionHandler::new(vec!["analyze".to_string()]);
454/// let resource_handler = ResourceUriCompletionHandler::new(vec![]);
455/// let tool_handler = ToolCompletionHandler::new(HashMap::new());
456///
457/// let combined = CompositeCompletionHandler::new()
458/// .with_prompt_handler(prompt_handler)
459/// .with_resource_handler(resource_handler)
460/// .with_tool_handler(tool_handler);
461/// ```
462pub struct CompositeCompletionHandler {
463    prompt_handler: Option<Box<dyn CompletionHandler>>,
464    resource_handler: Option<Box<dyn CompletionHandler>>,
465    tool_handler: Option<Box<dyn CompletionHandler>>,
466}
467
468impl CompositeCompletionHandler {
469    /// Create a new combined completion handler
470    pub fn new() -> Self {
471        Self {
472            prompt_handler: None,
473            resource_handler: None,
474            tool_handler: None,
475        }
476    }
477
478    /// Add a handler for a specific reference type
479    ///
480    /// # Arguments
481    /// * `reference_type` - The reference type this handler supports (e.g., "ref/prompt")
482    /// * `handler` - The completion handler
483    pub fn add_handler<H>(&mut self, reference_type: String, handler: H) -> &mut Self
484    where
485        H: CompletionHandler + 'static,
486    {
487        match reference_type.as_str() {
488            "ref/prompt" => self.prompt_handler = Some(Box::new(handler)),
489            "ref/resource" => self.resource_handler = Some(Box::new(handler)),
490            "ref/tool" => self.tool_handler = Some(Box::new(handler)),
491            _ => {} // Ignore unknown types
492        }
493        self
494    }
495
496    /// Add a handler for a specific reference type (builder pattern)
497    ///
498    /// # Arguments
499    /// * `reference_type` - The reference type this handler supports (e.g., "ref/prompt")
500    /// * `handler` - The completion handler
501    pub fn with_handler<H>(mut self, reference_type: impl Into<String>, handler: H) -> Self
502    where
503        H: CompletionHandler + 'static,
504    {
505        let ref_type = reference_type.into();
506        match ref_type.as_str() {
507            "ref/prompt" => self.prompt_handler = Some(Box::new(handler)),
508            "ref/resource" => self.resource_handler = Some(Box::new(handler)),
509            "ref/tool" => self.tool_handler = Some(Box::new(handler)),
510            _ => {} // Ignore unknown reference types
511        }
512        self
513    }
514
515    /// Set the prompt completion handler
516    pub fn with_prompt_handler<H: CompletionHandler + 'static>(mut self, handler: H) -> Self {
517        self.prompt_handler = Some(Box::new(handler));
518        self
519    }
520
521    /// Set the resource completion handler
522    pub fn with_resource_handler<H: CompletionHandler + 'static>(mut self, handler: H) -> Self {
523        self.resource_handler = Some(Box::new(handler));
524        self
525    }
526
527    /// Set the tool completion handler
528    pub fn with_tool_handler<H: CompletionHandler + 'static>(mut self, handler: H) -> Self {
529        self.tool_handler = Some(Box::new(handler));
530        self
531    }
532}
533
534impl std::fmt::Debug for CompositeCompletionHandler {
535    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
536        f.debug_struct("CompositeCompletionHandler")
537            .field("prompt_handler", &self.prompt_handler.is_some())
538            .field("resource_handler", &self.resource_handler.is_some())
539            .field("tool_handler", &self.tool_handler.is_some())
540            .finish()
541    }
542}
543
544impl Default for CompositeCompletionHandler {
545    fn default() -> Self {
546        Self::new()
547    }
548}
549
550#[async_trait]
551impl CompletionHandler for CompositeCompletionHandler {
552    async fn complete(
553        &self,
554        reference: &CompletionReference,
555        argument: &CompletionArgument,
556        context: Option<&CompletionContext>,
557    ) -> McpResult<Vec<String>> {
558        match reference {
559            CompletionReference::Prompt { .. } => {
560                if let Some(handler) = &self.prompt_handler {
561                    handler.complete(reference, argument, context).await
562                } else {
563                    Ok(vec![])
564                }
565            }
566            CompletionReference::Resource { .. } => {
567                if let Some(handler) = &self.resource_handler {
568                    handler.complete(reference, argument, context).await
569                } else {
570                    Ok(vec![])
571                }
572            }
573            CompletionReference::Tool { .. } => {
574                if let Some(handler) = &self.tool_handler {
575                    handler.complete(reference, argument, context).await
576                } else {
577                    Ok(vec![])
578                }
579            }
580        }
581    }
582}
583
584/// Completion context for additional information
585///
586/// Struct can be extended to provide additional context for completion
587/// requests, such as current file paths, user preferences, etc.
588#[derive(Debug, Clone, Default)]
589pub struct CompletionContext {
590    /// Additional context arguments
591    pub arguments: Option<HashMap<String, String>>,
592    /// Current working directory or context path
593    pub context_path: Option<String>,
594    /// User preferences for completion
595    pub preferences: Option<HashMap<String, serde_json::Value>>,
596}
597
598impl CompletionContext {
599    /// Create a new completion context
600    pub fn new() -> Self {
601        Self::default()
602    }
603
604    /// Set context arguments
605    pub fn with_arguments(mut self, arguments: HashMap<String, String>) -> Self {
606        self.arguments = Some(arguments);
607        self
608    }
609
610    /// Set context path
611    pub fn with_context_path<S: Into<String>>(mut self, path: S) -> Self {
612        self.context_path = Some(path.into());
613        self
614    }
615
616    /// Set user preferences
617    pub fn with_preferences(mut self, preferences: HashMap<String, serde_json::Value>) -> Self {
618        self.preferences = Some(preferences);
619        self
620    }
621}
622
623#[cfg(test)]
624mod tests {
625    use super::*;
626
627    #[tokio::test]
628    async fn test_prompt_completion() {
629        let handler = PromptCompletionHandler::new(vec![
630            "analyze_data".to_string(),
631            "analyze_text".to_string(),
632            "create_report".to_string(),
633        ]);
634
635        let reference = CompletionReference::Prompt {
636            name: "test".to_string(),
637        };
638
639        let argument = CompletionArgument {
640            name: "name".to_string(),
641            value: "ana".to_string(),
642        };
643
644        let results = handler.complete(&reference, &argument, None).await.unwrap();
645        assert_eq!(results.len(), 2);
646        assert!(results.contains(&"analyze_data".to_string()));
647        assert!(results.contains(&"analyze_text".to_string()));
648    }
649
650    #[tokio::test]
651    async fn test_prompt_completion_with_arguments() {
652        let mut handler = PromptCompletionHandler::new(vec!["analyze".to_string()]);
653        handler.add_argument_completions(
654            "analyze",
655            "format",
656            vec!["json".to_string(), "xml".to_string(), "yaml".to_string()],
657        );
658
659        let reference = CompletionReference::Prompt {
660            name: "analyze".to_string(),
661        };
662
663        let argument = CompletionArgument {
664            name: "format".to_string(),
665            value: "j".to_string(),
666        };
667
668        let results = handler.complete(&reference, &argument, None).await.unwrap();
669        assert_eq!(results.len(), 1);
670        assert!(results.contains(&"json".to_string()));
671    }
672
673    #[tokio::test]
674    async fn test_resource_uri_completion_handler() {
675        let templates = vec![ResourceTemplate::new(
676            "file:///docs/{category}/{filename}".to_string(),
677            "Documentation".to_string(),
678        )];
679        let handler = ResourceUriCompletionHandler::new(templates);
680
681        let reference = CompletionReference::Resource {
682            uri: "file:///docs/".to_string(),
683        };
684
685        let argument = CompletionArgument {
686            name: "uri".to_string(),
687            value: "file:///docs/".to_string(),
688        };
689
690        let results = handler.complete(&reference, &argument, None).await.unwrap();
691        // Results should contain expanded template examples
692        assert!(!results.is_empty());
693    }
694
695    #[tokio::test]
696    async fn test_composite_completion_handler() {
697        let mut composite = CompositeCompletionHandler::new();
698
699        let prompt_handler = PromptCompletionHandler::new(vec!["test_prompt".to_string()]);
700
701        composite.add_handler("ref/prompt".to_string(), prompt_handler);
702
703        let reference = CompletionReference::Prompt {
704            name: "test".to_string(),
705        };
706
707        let argument = CompletionArgument {
708            name: "name".to_string(),
709            value: "test".to_string(),
710        };
711
712        let results = composite
713            .complete(&reference, &argument, None)
714            .await
715            .unwrap();
716        assert_eq!(results.len(), 1);
717        assert!(results.contains(&"test_prompt".to_string()));
718    }
719
720    #[test]
721    fn test_fuzzy_matching() {
722        let handler = PromptCompletionHandler::new(vec![
723            "analyze_data".to_string(),
724            "create_report".to_string(),
725            "data_analysis".to_string(),
726        ]);
727
728        // Test exact match
729        let results = handler.fuzzy_match(&handler.prompts, "analyze_data");
730        assert_eq!(results[0], "analyze_data");
731
732        // Test prefix match
733        let results = handler.fuzzy_match(&handler.prompts, "ana");
734        assert!(results.contains(&"analyze_data".to_string()));
735
736        // Test substring match
737        let results = handler.fuzzy_match(&handler.prompts, "data");
738        assert!(results.len() >= 2); // Should match both "analyze_data" and "data_analysis"
739    }
740
741    #[test]
742    fn test_template_variable_extraction() {
743        let handler = ResourceUriCompletionHandler::new(vec![]);
744
745        let template = "file:///docs/{category}/{filename}";
746        let variables = handler.extract_template_variables(template);
747
748        assert_eq!(variables.len(), 2);
749        assert!(variables.contains(&"category".to_string()));
750        assert!(variables.contains(&"filename".to_string()));
751    }
752
753    #[test]
754    fn test_supported_reference_types() {
755        let prompt_handler = PromptCompletionHandler::new(vec![]);
756        assert_eq!(
757            prompt_handler.supported_reference_types(),
758            vec!["ref/prompt"]
759        );
760
761        let resource_handler = ResourceUriCompletionHandler::new(vec![]);
762        assert_eq!(
763            resource_handler.supported_reference_types(),
764            vec!["ref/resource"]
765        );
766
767        let tool_handler = ToolCompletionHandler::new(HashMap::new());
768        assert_eq!(tool_handler.supported_reference_types(), vec!["ref/tool"]);
769    }
770}
prism_mcp_rs/core/completion.rs

prism_mcp_rs/core/
completion.rs
